diff --git a/packages/ai-sdk/src/activities.ts b/packages/ai-sdk/src/activities.ts index 397b76922..20c2c0f27 100644 --- a/packages/ai-sdk/src/activities.ts +++ b/packages/ai-sdk/src/activities.ts @@ -5,8 +5,9 @@ import type { SharedV3ProviderOptions, SharedV3Headers, ProviderV3, + JSONSchema7, } from '@ai-sdk/provider'; -import type { FlexibleSchema, ToolExecutionOptions } from 'ai'; +import type { ToolExecutionOptions } from 'ai'; import { ApplicationFailure } from '@temporalio/common'; import type { McpClientFactories, McpClientFactory } from './mcp'; @@ -42,9 +43,15 @@ export interface InvokeEmbeddingModelArgs { */ export type InvokeEmbeddingModelResult = EmbeddingModelV3Result; +/** + * Result from listing MCP tools, returned by the listTools activity. + * Note: inputSchema is the serialized form of the AI SDK's schema wrapper. + * After JSON serialization through Temporal's activity boundary, the schema + * wrapper becomes { jsonSchema: JSONSchema7 }. + */ export interface ListToolResult { description?: string; - inputSchema: FlexibleSchema; + inputSchema: { jsonSchema: JSONSchema7 }; } export interface ListToolArgs { @@ -101,12 +108,14 @@ function activitiesForName(name: string, mcpClientFactory: McpClientFactory): ob try { const tools = await mcpClient.tools(); + // The activity returns FlexibleSchema objects, but ListToolResult describes the + // post-serialization form that workflows receive. The cast bridges this gap. return Object.fromEntries( Object.entries(tools).map(([k, v]) => [ k, { description: v.description, - inputSchema: v.inputSchema, + inputSchema: v.inputSchema as ListToolResult['inputSchema'], }, ]) ); diff --git a/packages/ai-sdk/src/mcp.ts b/packages/ai-sdk/src/mcp.ts index 988bb6246..9c4d883b1 100644 --- a/packages/ai-sdk/src/mcp.ts +++ b/packages/ai-sdk/src/mcp.ts @@ -1,5 +1,4 @@ import { jsonSchema, type ToolSet } from 'ai'; -import type { JSONSchema7 } from '@ai-sdk/provider'; import type { experimental_MCPClient as MCPClient } from '@ai-sdk/mcp'; import * as workflow from '@temporalio/workflow'; import type { ActivityOptions } from '@temporalio/workflow'; @@ -52,7 +51,8 @@ export class TemporalMCPClient { const callActivity = activities[this.options.name + '-callTool']!; return await callActivity({ name: toolName, input, options, clientArgs: this.options.clientArgs }); }, - inputSchema: jsonSchema(toolResult.inputSchema as JSONSchema7), + // toolResult.inputSchema is the serialized form { jsonSchema: JSONSchema7 } + inputSchema: jsonSchema(toolResult.inputSchema.jsonSchema), type: 'dynamic', }, ]) diff --git a/packages/test/src/test-ai-sdk.ts b/packages/test/src/test-ai-sdk.ts index 50cc726c8..04be292b9 100644 --- a/packages/test/src/test-ai-sdk.ts +++ b/packages/test/src/test-ai-sdk.ts @@ -22,6 +22,9 @@ import * as opentelemetry from '@opentelemetry/sdk-node'; import { SEMRESATTRS_SERVICE_NAME } from '@opentelemetry/semantic-conventions'; import { ExportResultCode } from '@opentelemetry/core'; import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'; +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js'; +import { ListToolsRequestSchema, CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js'; import { experimental_createMCPClient as createMCPClient } from '@ai-sdk/mcp'; import { AiSdkPlugin, createActivities } from '@temporalio/ai-sdk'; import { temporal } from '@temporalio/proto'; @@ -39,6 +42,7 @@ import { embeddingWorkflow, generateObjectWorkflow, helloWorldAgent, + mcpSchemaTestWorkflow, mcpWorkflow, middlewareWorkflow, telemetryWorkflow, @@ -544,6 +548,65 @@ test('callToolActivity awaits tool.execute before closing MCP client', async (t) ); }); +test('MCP tool schema survives activity serialization', async (t) => { + const { createWorker, executeWorkflow } = helpers(t); + + // Create in-memory MCP server with test tool + const createTestMcpClient = async () => { + const server = new Server( + { name: 'test-server', version: '1.0.0' }, + { capabilities: { tools: {} } } + ); + + // Register tools/list handler + server.setRequestHandler(ListToolsRequestSchema, async () => ({ + tools: [ + { + name: 'testTool', + description: 'A test tool', + inputSchema: { + type: 'object', + properties: { + testParam: { type: 'string', description: 'Test parameter' }, + }, + required: ['testParam'], + }, + }, + ], + })); + + // Register tools/call handler + server.setRequestHandler(CallToolRequestSchema, async () => ({ + content: [{ type: 'text', text: 'ok' }], + })); + + // Create linked in-memory transports + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + await server.connect(serverTransport); + + // Return real AI SDK MCP client connected to in-memory server + return createMCPClient({ transport: clientTransport }); + }; + + const worker = await createWorker({ + plugins: [ + new AiSdkPlugin({ + modelProvider: new TestProvider(helloWorkflowGenerator()), + mcpClientFactories: { testServer: createTestMcpClient }, + }), + ], + }); + + await worker.runUntil(async () => { + const result = await executeWorkflow(mcpSchemaTestWorkflow); + + // These would FAIL with the bug (undefined), PASS with fix + t.is(result.schemaType, 'object', 'Schema type should be "object", not undefined'); + t.true(result.hasProperties, 'Schema should have properties'); + t.is(result.propertyDescription, 'Test parameter', 'Property description should be preserved'); + }); +}); + // Currently fails in CI due to invalid server response but passes locally test.skip('MCP Use', async (t) => { if (remoteTests) { diff --git a/packages/test/src/workflows/ai-sdk.ts b/packages/test/src/workflows/ai-sdk.ts index 32bc28fbd..43840091f 100644 --- a/packages/test/src/workflows/ai-sdk.ts +++ b/packages/test/src/workflows/ai-sdk.ts @@ -113,6 +113,31 @@ export async function mcpWorkflow(prompt: string): Promise { return result.text; } +/** + * Test workflow that returns MCP tool schema structure for assertion. + * Used to verify inputSchema survives activity serialization correctly. + */ +export async function mcpSchemaTestWorkflow(): Promise<{ + toolName: string; + schemaType: string | undefined; + hasProperties: boolean; + propertyDescription: string | undefined; +}> { + const mcpClient = new TemporalMCPClient({ name: 'testServer' }); + const tools = await mcpClient.tools(); + + const [toolName, tool] = Object.entries(tools)[0]; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const schema = (tool as any).inputSchema.jsonSchema; + + return { + toolName, + schemaType: schema.type, + hasProperties: schema.properties !== undefined, + propertyDescription: schema.properties?.testParam?.description, + }; +} + /** * Workflow that demonstrates embedding model support. * Uses the temporalProvider to generate embeddings for multiple text values.