mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-19 09:33:53 +00:00
457 lines
13 KiB
TypeScript
457 lines
13 KiB
TypeScript
/**
|
|
* @license
|
|
* Copyright 2025 Qwen Team
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*/
|
|
|
|
/**
|
|
* E2E tests for SDK-embedded MCP servers
|
|
*
|
|
* Tests that the SDK can create and manage MCP servers running in the SDK process
|
|
* using the tool() and createSdkMcpServer() APIs.
|
|
*/
|
|
|
|
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
|
|
import { z } from 'zod';
|
|
import {
|
|
query,
|
|
tool,
|
|
createSdkMcpServer,
|
|
isSDKAssistantMessage,
|
|
isSDKResultMessage,
|
|
isSDKSystemMessage,
|
|
type SDKMessage,
|
|
type SDKSystemMessage,
|
|
} from '@qwen-code/sdk';
|
|
import {
|
|
SDKTestHelper,
|
|
extractText,
|
|
findToolUseBlocks,
|
|
createSharedTestOptions,
|
|
} from './test-helper.js';
|
|
|
|
const SHARED_TEST_OPTIONS = {
|
|
...createSharedTestOptions(),
|
|
permissionMode: 'yolo' as const,
|
|
};
|
|
|
|
describe('SDK MCP Server Integration (E2E)', () => {
|
|
let helper: SDKTestHelper;
|
|
let testDir: string;
|
|
|
|
beforeEach(async () => {
|
|
helper = new SDKTestHelper();
|
|
testDir = await helper.setup('sdk-mcp-server-integration');
|
|
});
|
|
|
|
afterEach(async () => {
|
|
await helper.cleanup();
|
|
});
|
|
|
|
describe('Basic SDK MCP Tool Usage', () => {
|
|
it('should use SDK MCP tool to perform a simple calculation', async () => {
|
|
// Define a simple calculator tool using the tool() API with Zod schema
|
|
const calculatorTool = tool(
|
|
'calculate_sum',
|
|
'Calculate the sum of two numbers',
|
|
z.object({
|
|
a: z.number().describe('First number'),
|
|
b: z.number().describe('Second number'),
|
|
}).shape,
|
|
async (args) => ({
|
|
content: [{ type: 'text', text: String(args.a + args.b) }],
|
|
}),
|
|
);
|
|
|
|
// Create SDK MCP server with the tool
|
|
const serverConfig = createSdkMcpServer({
|
|
name: 'sdk-calculator',
|
|
version: '1.0.0',
|
|
tools: [calculatorTool],
|
|
});
|
|
|
|
const q = query({
|
|
prompt:
|
|
'Use the calculate_sum tool to add 25 and 17. Output the result of tool only.',
|
|
options: {
|
|
...SHARED_TEST_OPTIONS,
|
|
cwd: testDir,
|
|
mcpServers: {
|
|
'sdk-calculator': serverConfig,
|
|
},
|
|
},
|
|
});
|
|
|
|
const messages: SDKMessage[] = [];
|
|
let assistantText = '';
|
|
let foundToolUse = false;
|
|
|
|
try {
|
|
for await (const message of q) {
|
|
messages.push(message);
|
|
|
|
if (isSDKAssistantMessage(message)) {
|
|
const toolUseBlocks = findToolUseBlocks(message, 'calculate_sum');
|
|
if (toolUseBlocks.length > 0) {
|
|
foundToolUse = true;
|
|
}
|
|
assistantText += extractText(message.message.content);
|
|
}
|
|
}
|
|
|
|
// Validate tool was called
|
|
expect(foundToolUse).toBe(true);
|
|
|
|
// Validate result contains expected answer: 25 + 17 = 42
|
|
expect(assistantText).toMatch(/42/);
|
|
|
|
// Validate successful completion
|
|
const lastMessage = messages[messages.length - 1];
|
|
expect(isSDKResultMessage(lastMessage)).toBe(true);
|
|
if (isSDKResultMessage(lastMessage)) {
|
|
expect(lastMessage.subtype).toBe('success');
|
|
}
|
|
} finally {
|
|
await q.close();
|
|
}
|
|
});
|
|
|
|
it('should use SDK MCP tool with string operations', async () => {
|
|
// Define a string manipulation tool with Zod schema
|
|
const stringTool = tool(
|
|
'reverse_string',
|
|
'Reverse a string',
|
|
{
|
|
text: z.string().describe('The text to reverse'),
|
|
},
|
|
async (args) => ({
|
|
content: [
|
|
{ type: 'text', text: args.text.split('').reverse().join('') },
|
|
],
|
|
}),
|
|
);
|
|
|
|
const serverConfig = createSdkMcpServer({
|
|
name: 'sdk-string-utils',
|
|
version: '1.0.0',
|
|
tools: [stringTool],
|
|
});
|
|
|
|
const q = query({
|
|
prompt: `Use the 'reverse_string' tool to process the word "hello world". Output the tool result only.`,
|
|
options: {
|
|
...SHARED_TEST_OPTIONS,
|
|
cwd: testDir,
|
|
mcpServers: {
|
|
'sdk-string-utils': serverConfig,
|
|
},
|
|
},
|
|
});
|
|
|
|
const messages: SDKMessage[] = [];
|
|
let assistantText = '';
|
|
let foundToolUse = false;
|
|
|
|
try {
|
|
for await (const message of q) {
|
|
messages.push(message);
|
|
|
|
if (isSDKAssistantMessage(message)) {
|
|
const toolUseBlocks = findToolUseBlocks(message, 'reverse_string');
|
|
if (toolUseBlocks.length > 0) {
|
|
foundToolUse = true;
|
|
}
|
|
assistantText += extractText(message.message.content);
|
|
}
|
|
}
|
|
|
|
// Validate tool was called
|
|
expect(foundToolUse).toBe(true);
|
|
|
|
// Validate result contains reversed string: "olleh"
|
|
expect(assistantText.toLowerCase()).toMatch(/olleh/);
|
|
|
|
// Validate successful completion
|
|
const lastMessage = messages[messages.length - 1];
|
|
expect(isSDKResultMessage(lastMessage)).toBe(true);
|
|
} finally {
|
|
await q.close();
|
|
}
|
|
});
|
|
});
|
|
|
|
describe('Multiple SDK MCP Tools', () => {
|
|
it('should use multiple tools from the same SDK MCP server', async () => {
|
|
// Define the Zod schema shape for two numbers
|
|
const twoNumbersSchema = {
|
|
a: z.number().describe('First number'),
|
|
b: z.number().describe('Second number'),
|
|
};
|
|
|
|
// Define multiple tools
|
|
const addTool = tool(
|
|
'sdk_add',
|
|
'Add two numbers',
|
|
twoNumbersSchema,
|
|
async (args) => ({
|
|
content: [{ type: 'text', text: String(args.a + args.b) }],
|
|
}),
|
|
);
|
|
|
|
const multiplyTool = tool(
|
|
'sdk_multiply',
|
|
'Multiply two numbers',
|
|
twoNumbersSchema,
|
|
async (args) => ({
|
|
content: [{ type: 'text', text: String(args.a * args.b) }],
|
|
}),
|
|
);
|
|
|
|
const serverConfig = createSdkMcpServer({
|
|
name: 'sdk-math',
|
|
version: '1.0.0',
|
|
tools: [addTool, multiplyTool],
|
|
});
|
|
|
|
const q = query({
|
|
prompt:
|
|
'First use sdk_add to calculate 10 + 5, then use sdk_multiply to multiply the result by 3. Give me the final answer.',
|
|
options: {
|
|
...SHARED_TEST_OPTIONS,
|
|
cwd: testDir,
|
|
debug: false,
|
|
mcpServers: {
|
|
'sdk-math': serverConfig,
|
|
},
|
|
},
|
|
});
|
|
|
|
const messages: SDKMessage[] = [];
|
|
let assistantText = '';
|
|
const toolCalls: string[] = [];
|
|
|
|
try {
|
|
for await (const message of q) {
|
|
messages.push(message);
|
|
|
|
if (isSDKAssistantMessage(message)) {
|
|
const toolUseBlocks = findToolUseBlocks(message);
|
|
toolUseBlocks.forEach((block) => {
|
|
toolCalls.push(block.name);
|
|
});
|
|
assistantText += extractText(message.message.content);
|
|
}
|
|
}
|
|
|
|
// Validate both tools were called
|
|
expect(toolCalls).toContain('sdk_add');
|
|
expect(toolCalls).toContain('sdk_multiply');
|
|
|
|
// Validate result: (10 + 5) * 3 = 45
|
|
expect(assistantText).toMatch(/45/);
|
|
|
|
// Validate successful completion
|
|
const lastMessage = messages[messages.length - 1];
|
|
expect(isSDKResultMessage(lastMessage)).toBe(true);
|
|
} finally {
|
|
await q.close();
|
|
}
|
|
});
|
|
});
|
|
|
|
describe('SDK MCP Server Discovery', () => {
|
|
it('should list SDK MCP servers in system init message', async () => {
|
|
// Define echo tool with Zod schema
|
|
const echoTool = tool(
|
|
'echo',
|
|
'Echo a message',
|
|
{
|
|
message: z.string().describe('Message to echo'),
|
|
},
|
|
async (args) => ({
|
|
content: [{ type: 'text', text: args.message }],
|
|
}),
|
|
);
|
|
|
|
const serverConfig = createSdkMcpServer({
|
|
name: 'sdk-echo',
|
|
version: '1.0.0',
|
|
tools: [echoTool],
|
|
});
|
|
|
|
const q = query({
|
|
prompt: 'Hello',
|
|
options: {
|
|
...SHARED_TEST_OPTIONS,
|
|
cwd: testDir,
|
|
debug: false,
|
|
mcpServers: {
|
|
'sdk-echo': serverConfig,
|
|
},
|
|
},
|
|
});
|
|
|
|
let systemMessage: SDKSystemMessage | null = null;
|
|
|
|
try {
|
|
for await (const message of q) {
|
|
if (isSDKSystemMessage(message) && message.subtype === 'init') {
|
|
systemMessage = message;
|
|
break;
|
|
}
|
|
}
|
|
|
|
// Validate MCP server is listed
|
|
expect(systemMessage).not.toBeNull();
|
|
expect(systemMessage!.mcp_servers).toBeDefined();
|
|
expect(Array.isArray(systemMessage!.mcp_servers)).toBe(true);
|
|
|
|
// Find our SDK MCP server
|
|
const sdkServer = systemMessage!.mcp_servers?.find(
|
|
(server) => server.name === 'sdk-echo',
|
|
);
|
|
expect(sdkServer).toBeDefined();
|
|
} finally {
|
|
await q.close();
|
|
}
|
|
});
|
|
});
|
|
|
|
describe('SDK MCP Tool Error Handling', () => {
|
|
it('should handle tool errors gracefully', async () => {
|
|
// Define a tool that throws an error with Zod schema
|
|
const errorTool = tool(
|
|
'maybe_fail',
|
|
'A tool that may fail based on input',
|
|
{
|
|
shouldFail: z.boolean().describe('If true, the tool will fail'),
|
|
},
|
|
async (args) => {
|
|
if (args.shouldFail) {
|
|
throw new Error('Tool intentionally failed');
|
|
}
|
|
return { content: [{ type: 'text', text: 'Success!' }] };
|
|
},
|
|
);
|
|
|
|
const serverConfig = createSdkMcpServer({
|
|
name: 'sdk-error-test',
|
|
version: '1.0.0',
|
|
tools: [errorTool],
|
|
});
|
|
|
|
const q = query({
|
|
prompt:
|
|
'Use the maybe_fail tool with shouldFail set to true. Tell me what happens.',
|
|
options: {
|
|
...SHARED_TEST_OPTIONS,
|
|
cwd: testDir,
|
|
debug: false,
|
|
mcpServers: {
|
|
'sdk-error-test': serverConfig,
|
|
},
|
|
},
|
|
});
|
|
|
|
const messages: SDKMessage[] = [];
|
|
let foundToolUse = false;
|
|
|
|
try {
|
|
for await (const message of q) {
|
|
messages.push(message);
|
|
|
|
if (isSDKAssistantMessage(message)) {
|
|
const toolUseBlocks = findToolUseBlocks(message, 'maybe_fail');
|
|
if (toolUseBlocks.length > 0) {
|
|
foundToolUse = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Tool should be called
|
|
expect(foundToolUse).toBe(true);
|
|
|
|
// Query should complete (even with tool error)
|
|
const lastMessage = messages[messages.length - 1];
|
|
expect(isSDKResultMessage(lastMessage)).toBe(true);
|
|
} finally {
|
|
await q.close();
|
|
}
|
|
});
|
|
});
|
|
|
|
describe('Async Tool Handlers', () => {
|
|
it('should handle async tool handlers with delays', async () => {
|
|
// Define a tool with async delay using Zod schema
|
|
const delayedTool = tool(
|
|
'delayed_response',
|
|
'Returns a value after a delay',
|
|
{
|
|
delay: z.number().describe('Delay in milliseconds (max 100)'),
|
|
value: z.string().describe('Value to return'),
|
|
},
|
|
async (args) => {
|
|
// Cap delay at 100ms for test performance
|
|
const actualDelay = Math.min(args.delay, 100);
|
|
await new Promise((resolve) => setTimeout(resolve, actualDelay));
|
|
return {
|
|
content: [{ type: 'text', text: `Delayed result: ${args.value}` }],
|
|
};
|
|
},
|
|
);
|
|
|
|
const serverConfig = createSdkMcpServer({
|
|
name: 'sdk-async',
|
|
version: '1.0.0',
|
|
tools: [delayedTool],
|
|
});
|
|
|
|
const q = query({
|
|
prompt:
|
|
'Use the delayed_response tool with delay=50 and value="test_async". Tell me the result.',
|
|
options: {
|
|
...SHARED_TEST_OPTIONS,
|
|
cwd: testDir,
|
|
debug: false,
|
|
mcpServers: {
|
|
'sdk-async': serverConfig,
|
|
},
|
|
},
|
|
});
|
|
|
|
const messages: SDKMessage[] = [];
|
|
let assistantText = '';
|
|
let foundToolUse = false;
|
|
|
|
try {
|
|
for await (const message of q) {
|
|
messages.push(message);
|
|
|
|
if (isSDKAssistantMessage(message)) {
|
|
const toolUseBlocks = findToolUseBlocks(
|
|
message,
|
|
'delayed_response',
|
|
);
|
|
if (toolUseBlocks.length > 0) {
|
|
foundToolUse = true;
|
|
}
|
|
assistantText += extractText(message.message.content);
|
|
}
|
|
}
|
|
|
|
// Validate tool was called
|
|
expect(foundToolUse).toBe(true);
|
|
|
|
// Validate result contains the delayed response
|
|
expect(assistantText.toLowerCase()).toMatch(/test_async/i);
|
|
|
|
// Validate successful completion
|
|
const lastMessage = messages[messages.length - 1];
|
|
expect(isSDKResultMessage(lastMessage)).toBe(true);
|
|
} finally {
|
|
await q.close();
|
|
}
|
|
});
|
|
});
|
|
});
|