Skip to content

Commit 1772fd5

Browse files
bizob2828bizob2828
authored
test: Updated mcp streaming tests to generate sessions to reuse transport for all streaming tests (newrelic#3729)
1 parent 961f13d commit 1772fd5

File tree

3 files changed

+134
-140
lines changed

3 files changed

+134
-140
lines changed

‎test/versioned/mcp-sdk/client-streaming.test.js‎

Lines changed: 131 additions & 134 deletions
Original file line numberDiff line numberDiff line change
@@ -5,169 +5,166 @@
55

66
'use strict'
77

8-
const { describe, test } = require('node:test')
8+
const test = require('node:test')
99
const assert = require('node:assert')
10-
const semver = require('semver')
1110

1211
const { removeModules } = require('../../lib/cache-buster')
1312
const helper = require('../../lib/agent_helper')
1413
const { assertPackageMetrics, assertSegments, assertSpanKind } = require('../../lib/custom-assertions')
1514
const {
1615
MCP
1716
} = require('../../../lib/metrics/names')
18-
const version = helper.readPackageVersion(__dirname, '@modelcontextprotocol/sdk')
19-
20-
describe('MCP streaming tests', { skip: semver.gte(version, '1.26.0') }, () => {
21-
test.beforeEach(async (ctx) => {
22-
ctx.nr = {}
23-
ctx.nr.agent = helper.instrumentMockedAgent({
24-
ai_monitoring: {
25-
enabled: ctx.name.includes('disabled') ? false : true
26-
}
27-
})
2817

29-
const { Client } = require('@modelcontextprotocol/sdk/client/index.js')
30-
const { StreamableHTTPClientTransport } = require('@modelcontextprotocol/sdk/client/streamableHttp.js')
31-
// Set up server
32-
const McpTestServer = require('./streaming-server')
33-
ctx.nr.mcpServer = new McpTestServer()
34-
const port = await ctx.nr.mcpServer.start()
35-
36-
// Set up client
37-
ctx.nr.transport = new StreamableHTTPClientTransport(
38-
new URL(`http://localhost:${port}/mcp`)
39-
)
40-
ctx.nr.client = new Client(
41-
{
42-
name: 'test-client',
43-
version: '1.0.0'
44-
}
45-
)
46-
await ctx.nr.client.connect(ctx.nr.transport)
18+
test.beforeEach(async (ctx) => {
19+
ctx.nr = {}
20+
ctx.nr.agent = helper.instrumentMockedAgent({
21+
ai_monitoring: {
22+
enabled: ctx.name.includes('disabled') ? false : true
23+
}
4724
})
4825

49-
test.afterEach(async (ctx) => {
50-
await ctx.nr.client.close()
51-
await ctx.nr.transport.close()
52-
await ctx.nr.mcpServer.stop()
53-
helper.unloadAgent(ctx.nr.agent)
54-
removeModules([
55-
'@modelcontextprotocol/sdk/client/index.js',
56-
'@modelcontextprotocol/sdk/client/streamableHttp.js'
57-
])
26+
const { Client } = require('@modelcontextprotocol/sdk/client/index.js')
27+
const { StreamableHTTPClientTransport } = require('@modelcontextprotocol/sdk/client/streamableHttp.js')
28+
// Set up server
29+
const McpTestServer = require('./streaming-server')
30+
ctx.nr.mcpServer = new McpTestServer()
31+
const port = await ctx.nr.mcpServer.start()
32+
33+
// Set up client
34+
ctx.nr.transport = new StreamableHTTPClientTransport(
35+
new URL(`http://localhost:${port}/mcp`)
36+
)
37+
ctx.nr.client = new Client(
38+
{
39+
name: 'test-client',
40+
version: '1.0.0'
41+
}
42+
)
43+
await ctx.nr.client.connect(ctx.nr.transport)
44+
})
45+
46+
test.afterEach(async (ctx) => {
47+
await ctx.nr.client.close()
48+
await ctx.nr.transport.close()
49+
await ctx.nr.mcpServer.stop()
50+
helper.unloadAgent(ctx.nr.agent)
51+
removeModules([
52+
'@modelcontextprotocol/sdk/client/index.js',
53+
'@modelcontextprotocol/sdk/client/streamableHttp.js'
54+
])
55+
})
56+
57+
test('should log package tracking metrics', (t) => {
58+
const { agent } = t.nr
59+
const version = helper.readPackageVersion(__dirname, '@modelcontextprotocol/sdk')
60+
assertPackageMetrics({
61+
agent,
62+
pkg: '@modelcontextprotocol/sdk',
63+
version,
64+
subscriberType: true
5865
})
66+
})
5967

60-
test('should log package tracking metrics', (t) => {
61-
const { agent } = t.nr
62-
assertPackageMetrics({
68+
test('should create span for callTool', (t, end) => {
69+
const { agent, client } = t.nr
70+
helper.runInTransaction(agent, async (tx) => {
71+
const result = await client.callTool({
72+
name: 'echo',
73+
arguments: {
74+
message: 'example message'
75+
}
76+
})
77+
assert.ok(result, 'should return a result from the tool call')
78+
const name = `${MCP.TOOL}/callTool/echo`
79+
assertSegments(tx.trace, tx.trace.root, [name], { exact: false })
80+
tx.end()
81+
assertSpanKind({
6382
agent,
64-
pkg: '@modelcontextprotocol/sdk',
65-
version,
66-
subscriberType: true
83+
segments: [
84+
{ name, kind: 'internal' }
85+
]
6786
})
87+
88+
end()
6889
})
90+
})
6991

70-
test('should create span for callTool', (t, end) => {
71-
const { agent, client } = t.nr
72-
helper.runInTransaction(agent, async (tx) => {
73-
const result = await client.callTool({
74-
name: 'echo',
75-
arguments: {
76-
message: 'example message'
77-
}
78-
})
79-
assert.ok(result, 'should return a result from the tool call')
80-
const name = `${MCP.TOOL}/callTool/echo`
81-
assertSegments(tx.trace, tx.trace.root, [name], { exact: false })
82-
tx.end()
83-
assertSpanKind({
84-
agent,
85-
segments: [
86-
{ name, kind: 'internal' }
87-
]
88-
})
89-
90-
end()
92+
test('should create span for readResource', (t, end) => {
93+
const { agent, client } = t.nr
94+
helper.runInTransaction(agent, async (tx) => {
95+
const resource = await client.readResource({
96+
uri: 'echo://hello-world',
9197
})
92-
})
9398

94-
test('should create span for readResource', (t, end) => {
95-
const { agent, client } = t.nr
96-
helper.runInTransaction(agent, async (tx) => {
97-
const resource = await client.readResource({
98-
uri: 'echo://hello-world',
99-
})
99+
assert.ok(resource, 'should return a resource from readResource')
100+
101+
const name = `${MCP.RESOURCE}/readResource/echo`
102+
assertSegments(tx.trace, tx.trace.root, [name], { exact: false })
103+
104+
tx.end()
105+
assertSpanKind({
106+
agent,
107+
segments: [
108+
{ name, kind: 'internal' }
109+
]
110+
})
111+
112+
end()
113+
})
114+
})
100115

101-
assert.ok(resource, 'should return a resource from readResource')
116+
test('should create span for getPrompt', (t, end) => {
117+
const { agent, client } = t.nr
118+
helper.runInTransaction(agent, async (tx) => {
119+
const prompt = await client.getPrompt({
120+
name: 'echo',
121+
arguments: {
122+
message: 'example message'
123+
}
124+
})
102125

103-
const name = `${MCP.RESOURCE}/readResource/echo`
104-
assertSegments(tx.trace, tx.trace.root, [name], { exact: false })
126+
assert.ok(prompt, 'should return a prompt from getPrompt')
105127

106-
tx.end()
107-
assertSpanKind({
108-
agent,
109-
segments: [
110-
{ name, kind: 'internal' }
111-
]
112-
})
128+
const name = `${MCP.PROMPT}/getPrompt/echo`
129+
assertSegments(tx.trace, tx.trace.root, [name], { exact: false })
113130

114-
end()
131+
tx.end()
132+
assertSpanKind({
133+
agent,
134+
segments: [
135+
{ name, kind: 'internal' }
136+
]
115137
})
138+
139+
end()
116140
})
141+
})
117142

118-
test('should create span for getPrompt', (t, end) => {
119-
const { agent, client } = t.nr
120-
helper.runInTransaction(agent, async (tx) => {
121-
const prompt = await client.getPrompt({
122-
name: 'echo',
123-
arguments: {
124-
message: 'example message'
125-
}
126-
})
127-
128-
assert.ok(prompt, 'should return a prompt from getPrompt')
129-
130-
const name = `${MCP.PROMPT}/getPrompt/echo`
131-
assertSegments(tx.trace, tx.trace.root, [name], { exact: false })
132-
133-
tx.end()
134-
assertSpanKind({
135-
agent,
136-
segments: [
137-
{ name, kind: 'internal' }
138-
]
139-
})
140-
141-
end()
143+
test('should not instrument if ai_monitoring is disabled', (t, end) => {
144+
const { agent, client } = t.nr
145+
146+
helper.runInTransaction(agent, async (tx) => {
147+
const result = await client.callTool({
148+
name: 'echo',
149+
arguments: {
150+
message: 'example message'
151+
}
142152
})
143-
})
144153

145-
test('should not instrument if ai_monitoring is disabled', (t, end) => {
146-
const { agent, client } = t.nr
147-
148-
helper.runInTransaction(agent, async (tx) => {
149-
const result = await client.callTool({
150-
name: 'echo',
151-
arguments: {
152-
message: 'example message'
153-
}
154-
})
155-
156-
assert.ok(result, 'should still return a result from the tool call')
157-
158-
const name = `${MCP.TOOL}/callTool/echo`
159-
const root = tx?.trace?.segments?.root
160-
assert.ok(root)
161-
function assertNoMcpSegment(node) {
162-
assert.notEqual(node?.segment?.name, name, 'should not create MCP segment')
163-
for (const child of node?.children) {
164-
assertNoMcpSegment(child)
165-
}
154+
assert.ok(result, 'should still return a result from the tool call')
155+
156+
const name = `${MCP.TOOL}/callTool/echo`
157+
const root = tx?.trace?.segments?.root
158+
assert.ok(root)
159+
function assertNoMcpSegment(node) {
160+
assert.notEqual(node?.segment?.name, name, 'should not create MCP segment')
161+
for (const child of node?.children) {
162+
assertNoMcpSegment(child)
166163
}
167-
assertNoMcpSegment(root)
164+
}
165+
assertNoMcpSegment(root)
168166

169-
tx.end()
170-
end()
171-
})
167+
tx.end()
168+
end()
172169
})
173170
})

‎test/versioned/mcp-sdk/package.json‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,5 @@
2424
"client-streaming.test.js"
2525
]
2626
}
27-
],
28-
"dependencies": {
29-
"express": "^5.1.0",
30-
"zod": "^3.25.76"
31-
}
27+
]
3228
}

‎test/versioned/mcp-sdk/streaming-server.js‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ const express = require('express')
88
const { McpServer, ResourceTemplate } = require('@modelcontextprotocol/sdk/server/mcp.js')
99
const { StreamableHTTPServerTransport } = require('@modelcontextprotocol/sdk/server/streamableHttp.js')
1010
const { z } = require('zod')
11+
const { randomUUID } = require('node:crypto')
1112

1213
class McpTestServer {
1314
constructor() {
@@ -22,7 +23,7 @@ class McpTestServer {
2223
})
2324

2425
this.transport = new StreamableHTTPServerTransport({
25-
sessionIdGenerator: undefined,
26+
sessionIdGenerator: () => randomUUID(),
2627
enableDnsRebindingProtection: true,
2728
})
2829

0 commit comments

Comments
 (0)