@mcp-ts/sdk 1.5.2 → 1.5.3

package/src/adapters/agui-middleware.ts

@@ -8,7 +8,7 @@
  * @requires rxjs - Uses RxJS Observables for event streaming
  */
 
-import { Observable, Subscriber } from 'rxjs';
+import { Observable, Subscriber, Subscription } from 'rxjs';
 import {
     Middleware,
     EventType,
@@ -28,11 +28,19 @@ interface ToolResult {
     messageId: string;
 }
 
+interface SnapshotToolCall {
+    id: string;
+    name: string;
+    arguments: string;
+}
+
 /** State for tracking tool calls during a run */
 interface RunState {
     toolCallArgsBuffer: Map<string, string>;
     toolCallNames: Map<string, string>;
     pendingMcpCalls: Set<string>;
+    completedMcpCalls: Set<string>;
+    assistantMessageId?: string;
     textContent?: string;
     error: boolean;
 }
@@ -87,6 +95,10 @@ export class McpMiddleware extends Middleware {
         }
     }
 
+    private normalizeArgsString(argsString: string): string {
+        return JSON.stringify(this.parseArgs(argsString));
+    }
+
     private async executeTool(toolName: string, args: Record<string, any>): Promise<string> {
         const tool = this.tools.find(t => t.name === toolName);
         if (!tool?.handler) {
@@ -116,13 +128,82 @@ export class McpMiddleware extends Middleware {
         if (!anyInput.runId) anyInput.runId = this.generateId('mcp_run');
     }
 
+    private getSnapshotToolCalls(message: any): SnapshotToolCall[] {
+        const rawToolCalls = Array.isArray(message?.toolCalls)
+            ? message.toolCalls
+            : (Array.isArray(message?.tool_calls) ? message.tool_calls : []);
+
+        return rawToolCalls.flatMap((toolCall: any) => {
+            if (!toolCall || typeof toolCall !== 'object') return [];
+
+            const id = typeof toolCall.id === 'string' ? toolCall.id : undefined;
+            const fn = toolCall.function && typeof toolCall.function === 'object'
+                ? toolCall.function
+                : undefined;
+            const name = typeof fn?.name === 'string'
+                ? fn.name
+                : (typeof toolCall.name === 'string' ? toolCall.name : undefined);
+
+            if (!id || !name) return [];
+
+            const rawArgs = fn?.arguments ?? toolCall.arguments ?? toolCall.args;
+            const argsString = typeof rawArgs === 'string'
+                ? rawArgs
+                : (rawArgs && typeof rawArgs === 'object' ? JSON.stringify(rawArgs) : '{}');
+
+            return [{ id, name, arguments: argsString }];
+        });
+    }
+
+    private getResolvedToolCallIds(messages: any[]): Set<string> {
+        const resolved = new Set<string>();
+        for (const message of messages) {
+            if (message?.role !== 'tool') continue;
+            const toolCallId = message.toolCallId ?? message.tool_call_id;
+            if (typeof toolCallId === 'string') {
+                resolved.add(toolCallId);
+            }
+        }
+        return resolved;
+    }
+
+    private shouldFilterMessagesSnapshot(event: BaseEvent, state: RunState): boolean {
+        if (state.completedMcpCalls.size === 0 || state.pendingMcpCalls.size > 0) {
+            return false;
+        }
+
+        // Only suppress snapshots that replay an assistant answer we already streamed.
+        // Snapshot-only final answers still need to reach the UI.
+        if (!state.assistantMessageId && !state.textContent) {
+            return false;
+        }
+
+        const messages = (event as any).messages;
+        if (!Array.isArray(messages)) {
+            return false;
+        }
+
+        const lastAssistant = [...messages].reverse().find((message) => message?.role === 'assistant');
+        if (!lastAssistant) {
+            return false;
+        }
+
+        return (
+            (typeof lastAssistant.id === 'string' && lastAssistant.id === state.assistantMessageId) ||
+            (typeof lastAssistant.content === 'string' && lastAssistant.content === state.textContent)
+        );
+    }
+
     /** Process tool call events and update state */
     private handleToolCallEvent(event: BaseEvent, state: RunState): void {
-        const { toolCallArgsBuffer, toolCallNames, pendingMcpCalls } = state;
+        const { toolCallArgsBuffer, toolCallNames, pendingMcpCalls, completedMcpCalls } = state;
 
         // Accumulate text content for reconstruction
-        if (event.type === EventType.TEXT_MESSAGE_CHUNK) {
+        if (event.type === EventType.TEXT_MESSAGE_CHUNK || event.type === EventType.TEXT_MESSAGE_CONTENT) {
             const e = event as any;
+            if (typeof e.messageId === 'string') {
+                state.assistantMessageId = e.messageId;
+            }
             if (e.delta) {
                 state.textContent = (state.textContent || '') + e.delta;
             }
@@ -132,7 +213,7 @@ export class McpMiddleware extends Middleware {
             const e = event as any;
             if (e.toolCallId && e.toolCallName) {
                 toolCallNames.set(e.toolCallId, e.toolCallName);
-                if (this.isMcpTool(e.toolCallName)) {
+                if (this.isMcpTool(e.toolCallName) && !completedMcpCalls.has(e.toolCallId)) {
                     pendingMcpCalls.add(e.toolCallId);
                 }
                 console.log(`[McpMiddleware] TOOL_CALL_START: ${e.toolCallName} (id: ${e.toolCallId}, isMCP: ${this.isMcpTool(e.toolCallName)})`);
@@ -152,33 +233,52 @@ export class McpMiddleware extends Middleware {
             console.log(`[McpMiddleware] TOOL_CALL_END: ${toolCallNames.get(e.toolCallId) ?? 'unknown'} (id: ${e.toolCallId})`);
         }
 
+        if (event.type === EventType.TOOL_CALL_RESULT) {
+            const e = event as any;
+            if (typeof e.toolCallId === 'string') {
+                pendingMcpCalls.delete(e.toolCallId);
+                completedMcpCalls.add(e.toolCallId);
+            }
+        }
+
         // Workaround: Extract parallel tool calls from MESSAGES_SNAPSHOT
         if (event.type === EventType.MESSAGES_SNAPSHOT) {
             const messages = (event as any).messages || [];
             if (messages.length > 0) {
-                const lastMsg = messages[messages.length - 1];
-                // Update text content from snapshot if available (often more reliable)
-                if (lastMsg.role === 'assistant' && lastMsg.content) {
-                    state.textContent = lastMsg.content;
+                const resolvedToolCallIds = this.getResolvedToolCallIds(messages);
+
+                for (const toolCallId of [...pendingMcpCalls]) {
+                    if (resolvedToolCallIds.has(toolCallId) || completedMcpCalls.has(toolCallId)) {
+                        pendingMcpCalls.delete(toolCallId);
+                    }
                 }
 
-                // Discover tools
+                // Discover unresolved tools. LangGraph snapshots include historical
+                // assistant tool calls, so only calls without a tool result are pending.
                 for (let i = messages.length - 1; i >= 0; i--) {
                     const msg = messages[i];
-                    const tools = Array.isArray(msg.toolCalls) ? msg.toolCalls :
-                        (Array.isArray(msg.tool_calls) ? msg.tool_calls : []);
+                    const tools = this.getSnapshotToolCalls(msg)
+                        .filter(tc =>
+                            !resolvedToolCallIds.has(tc.id) &&
+                            !completedMcpCalls.has(tc.id) &&
+                            this.isMcpTool(tc.name)
+                        );
 
                     if (msg.role === 'assistant' && tools.length > 0) {
                         for (const tc of tools) {
-                            if (tc.id && tc.function?.name && !toolCallNames.has(tc.id)) {
-                                toolCallNames.set(tc.id, tc.function.name);
-                                toolCallArgsBuffer.set(tc.id, tc.function.arguments || '{}');
-                                if (this.isMcpTool(tc.function.name)) {
-                                    pendingMcpCalls.add(tc.id);
-                                    console.log(`[McpMiddleware] MESSAGES_SNAPSHOT: Discovered ${tc.function.name} (id: ${tc.id})`);
-                                }
+                            if (!toolCallNames.has(tc.id)) {
+                                toolCallNames.set(tc.id, tc.name);
+                                toolCallArgsBuffer.set(tc.id, tc.arguments || '{}');
+                                pendingMcpCalls.add(tc.id);
+                                console.log(`[McpMiddleware] MESSAGES_SNAPSHOT: Discovered ${tc.name} (id: ${tc.id})`);
                             }
                         }
+                        if (typeof msg.id === 'string') {
+                            state.assistantMessageId = msg.id;
+                        }
+                        if (typeof msg.content === 'string') {
+                            state.textContent = msg.content;
+                        }
                         break;
                     }
                 }
@@ -188,10 +288,15 @@ export class McpMiddleware extends Middleware {
 
     /** Execute pending MCP tools and return results */
     private async executeTools(state: RunState): Promise<ToolResult[]> {
-        const { toolCallArgsBuffer, toolCallNames, pendingMcpCalls } = state;
+        const { toolCallArgsBuffer, toolCallNames, pendingMcpCalls, completedMcpCalls } = state;
         const results: ToolResult[] = [];
 
         const promises = [...pendingMcpCalls].map(async (toolCallId) => {
+            if (completedMcpCalls.has(toolCallId)) {
+                pendingMcpCalls.delete(toolCallId);
+                return;
+            }
+
             const toolName = toolCallNames.get(toolCallId);
             if (!toolName) return;
 
@@ -206,6 +311,7 @@ export class McpMiddleware extends Middleware {
                 messageId: this.generateId('mcp_result'),
             });
             pendingMcpCalls.delete(toolCallId);
+            completedMcpCalls.add(toolCallId);
         });
 
         await Promise.all(promises);
@@ -229,10 +335,13 @@ export class McpMiddleware extends Middleware {
 
     run(input: RunAgentInput, next: AbstractAgent): Observable<BaseEvent> {
         return new Observable<BaseEvent>((observer: Subscriber<BaseEvent>) => {
+            const subscriptions = new Set<Subscription>();
+            let continuationInProgress = false;
             const state: RunState = {
                 toolCallArgsBuffer: new Map(),
                 toolCallNames: new Map(),
                 pendingMcpCalls: new Set(),
+                completedMcpCalls: new Set(),
                 textContent: '',
                 error: false,
             };
@@ -265,12 +374,13 @@ export class McpMiddleware extends Middleware {
                 }
 
                 console.log(`[McpMiddleware] RUN_FINISHED with ${state.pendingMcpCalls.size} pending calls`);
+                continuationInProgress = true;
 
                 // Reconstruct the Assistant Message that triggered these tools
                 const toolCalls = [];
                 for (const toolCallId of state.pendingMcpCalls) {
                     const name = state.toolCallNames.get(toolCallId);
-                    const args = state.toolCallArgsBuffer.get(toolCallId) || '{}';
+                    const args = this.normalizeArgsString(state.toolCallArgsBuffer.get(toolCallId) || '{}');
                     if (name) {
                         toolCalls.push({
                             id: toolCallId,
@@ -283,10 +393,10 @@ export class McpMiddleware extends Middleware {
                 // Add the Assistant Message to history FIRST
                 if (toolCalls.length > 0 || state.textContent) {
                     const assistantMsg = {
-                        id: this.generateId('msg_ast'),
+                        id: state.assistantMessageId || this.generateId('msg_ast'),
                         role: 'assistant',
                         content: state.textContent || null, // Ensure null if empty string for strict LLMs
-                        tool_calls: toolCalls.length > 0 ? toolCalls : undefined
+                        toolCalls: toolCalls.length > 0 ? toolCalls : undefined
                     };
                     input.messages.push(assistantMsg as any);
                     console.log(`[McpMiddleware] Added assistant message to history before tools: ${state.textContent?.slice(0, 50)}... [${toolCalls.length} tools]`);
@@ -304,7 +414,7 @@ export class McpMiddleware extends Middleware {
                     input.messages.push({
                         id: messageId,
                         role: 'tool',
-                        tool_call_id: toolCallId,
+                        toolCallId,
                         content: result,
                     } as any);
                 }
@@ -312,35 +422,45 @@ export class McpMiddleware extends Middleware {
                 // Reset state for next turn
                 state.toolCallArgsBuffer.clear();
                 state.toolCallNames.clear();
+                state.assistantMessageId = undefined;
                 state.textContent = ''; // Clear text content for next turn
 
-                anyInput.runId = this.generateId('mcp_run');
                 console.log(`[McpMiddleware] === CONTINUATION RUN === messages: ${input.messages.length}`);
 
-                // Subscribe to continuation
-                next.run(input).subscribe({
+                subscribeToRun(true);
+            };
+
+            const subscribeToRun = (isContinuation: boolean): Subscription => {
+                let subscription = new Subscription();
+                subscription = next.run(input).subscribe({
                     next: (event) => {
                         if (state.error) return;
 
                         this.handleToolCallEvent(event, state);
 
                         if (event.type === EventType.RUN_ERROR) {
-                            console.log(`[McpMiddleware] RUN_ERROR received in continuation`);
+                            console.log(`[McpMiddleware] RUN_ERROR received${isContinuation ? ' in continuation' : ''}`);
                             state.error = true;
                             observer.next(event);
                             observer.complete();
                             return;
                         }
 
-                        if (event.type === EventType.RUN_STARTED) {
+                        if (isContinuation && event.type === EventType.RUN_STARTED) {
                             console.log(`[McpMiddleware] Filtering RUN_STARTED from continuation`);
                             return;
                         }
 
+                        if (event.type === EventType.MESSAGES_SNAPSHOT && this.shouldFilterMessagesSnapshot(event, state)) {
+                            console.log(`[McpMiddleware] Filtering completed MCP MESSAGES_SNAPSHOT to preserve streamed message order`);
+                            return;
+                        }
+
                         if (event.type === EventType.RUN_FINISHED) {
                             if (state.pendingMcpCalls.size > 0) {
                                 handleRunFinished();
                             } else {
+                                continuationInProgress = false;
                                 observer.next(event);
                                 observer.complete();
                             }
@@ -350,44 +470,28 @@ export class McpMiddleware extends Middleware {
                     },
                     error: (err) => {
                         state.error = true;
+                        continuationInProgress = false;
                         observer.error(err);
                     },
                     complete: () => {
-                        if (!state.error && state.pendingMcpCalls.size === 0) observer.complete();
+                        subscriptions.delete(subscription);
+                        if (!state.error && state.pendingMcpCalls.size === 0 && !continuationInProgress) observer.complete();
                     },
                 });
+                if (!subscription.closed) {
+                    subscriptions.add(subscription);
+                }
+                return subscription;
             };
 
-            const subscription = next.run(input).subscribe({
-                next: (event) => {
-                    if (state.error) return;
-
-                    this.handleToolCallEvent(event, state);
+            subscribeToRun(false);
 
-                    if (event.type === EventType.RUN_ERROR) {
-                        console.log(`[McpMiddleware] RUN_ERROR received`);
-                        state.error = true;
-                        observer.next(event);
-                        observer.complete();
-                        return;
-                    }
-
-                    if (event.type === EventType.RUN_FINISHED) {
-                        handleRunFinished();
-                        return;
-                    }
-                    observer.next(event);
-                },
-                error: (err) => {
-                    state.error = true;
-                    observer.error(err);
-                },
-                complete: () => {
-                    if (!state.error && state.pendingMcpCalls.size === 0) observer.complete();
-                },
-            });
-
-            return () => subscription.unsubscribe();
+            return () => {
+                for (const subscription of subscriptions) {
+                    subscription.unsubscribe();
+                }
+                subscriptions.clear();
+            };
         });
     }
 }
@@ -410,4 +514,4 @@ export { createMcpMiddleware as createMcpToolMiddleware };
 
 // Re-exports
 export { Middleware, EventType };
-export type { RunAgentInput, BaseEvent, AbstractAgent, ToolCallEndEvent, Tool };
+export type { RunAgentInput, BaseEvent, AbstractAgent, ToolCallEndEvent, Tool };

package/src/shared/meta-tools.ts

@@ -17,6 +17,7 @@
 
 import type { Tool, CallToolResult } from '@modelcontextprotocol/sdk/types.js';
 import type { ToolRouter } from './tool-router.js';
+import type { IndexedTool } from './tool-index.js';
 
 // ---------------------------------------------------------------------------
 // Tool Definitions
@@ -99,7 +100,8 @@ export function createGetSchemaToolDefinition(): Tool {
     description:
       'Get the full input schema (parameters) for a specific tool. ' +
       'Call this after mcp_search_tool_bm25 to get the parameter details ' +
-      'needed to call a tool correctly.',
+      'needed to call a tool correctly. ' +
+      'Do NOT call the discovered tool directly; after reading the schema, call mcp_execute_tool.',
     inputSchema: {
       type: 'object' as const,
       properties: {
@@ -189,7 +191,7 @@ export async function executeMetaTool(
   router: ToolRouter,
   callToolFn?: CallToolFn
 ): Promise<CallToolResult | null> {
-  const resolveToolSchema = (name: string, namespace?: string): { tool?: Tool; error?: CallToolResult } => {
+  const resolveToolSchema = (name: string, namespace?: string): { tool?: IndexedTool; error?: CallToolResult } => {
     try {
       return { tool: router.getToolSchema(name, namespace) };
     } catch (err) {
@@ -322,6 +324,13 @@ export async function executeMetaTool(
         name: tool.name,
         description: tool.description,
         inputSchema: tool.inputSchema,
+        executionInstructions: {
+          nextTool: 'mcp_execute_tool',
+          toolName: tool.name,
+          serverId: tool.serverId,
+          note:
+            'Do not call this discovered tool directly unless it was explicitly registered as a runtime tool. Execute it via mcp_execute_tool and pass these parameters inside args.',
+        },
       };
 
       return {