@mcp-ts/sdk 1.5.2 → 1.6.0

package/dist/{tool-router-DsKhRmJm.d.ts → tool-router-Bn9R0KWr.d.ts}

@@ -80,6 +80,39 @@ interface ToolSummary {
     /** Estimated token cost of the full inputSchema */
     estimatedTokens: number;
 }
+/** Server-level summary derived from indexed tools. */
+interface ToolServerSummary {
+    /** Human-readable server name */
+    serverName: string;
+    /** Stable server identifier */
+    serverId: string;
+    /** Session the server belongs to */
+    sessionId: string;
+    /** Number of indexed tools for this server */
+    toolCount: number;
+}
+/** Optional filters for search and listing. */
+interface ToolSearchOptions {
+    /** Restrict results to this server ID. */
+    serverId?: string;
+    /** Restrict results to servers whose name or ID matches this value. */
+    serverName?: string;
+}
+/** Paginated tool listing result. */
+interface ToolListResult {
+    tools: ToolSummary[];
+    totalCount: number;
+    returnedCount: number;
+    nextCursor?: string;
+    servers: ToolServerSummary[];
+}
+interface ToolLookupOptions {
+    /**
+     * Allow namespace to match a fragment of serverName after exact
+     * sessionId/serverId matching fails.
+     */
+    allowServerNameFragment?: boolean;
+}
 /** A tool with routing metadata attached during indexing. */
 interface IndexedTool extends Tool {
     sessionId: string;
@@ -147,7 +180,7 @@ declare class ToolIndex {
      *
      *   `score = keywordWeight × keyword_score + (1 - keywordWeight) × cosine_score`
      */
-    search(query: string, topK?: number): Promise<ToolSummary[]>;
+    search(query: string, topK?: number, options?: ToolSearchOptions): Promise<ToolSummary[]>;
     /**
      * Search tools using a regex pattern.
      * Matches against name, description, and parameter metadata.
@@ -155,11 +188,19 @@ declare class ToolIndex {
     searchRegex(pattern: string, topK?: number): ToolSummary[];
     /**
      * Get tool definition(s) by name.
-     * If namespace is provided, it tries to match sessionId or serverName.
+     * If namespace is provided, exact sessionId/serverId matches take precedence.
+     * Falls back to serverName fragment matching only when explicitly allowed.
      */
-    getTool(name: string, namespace?: string): IndexedTool[];
+    getTool(name: string, namespace?: string, options?: ToolLookupOptions): IndexedTool[];
     /** All indexed tool names. */
     getToolNames(): string[];
+    /** List indexed servers with tool counts. */
+    listServers(options?: ToolSearchOptions): ToolServerSummary[];
+    /** List tools deterministically, optionally scoped to a server. */
+    listTools(options?: ToolSearchOptions & {
+        limit?: number;
+        cursor?: string;
+    }): ToolListResult;
     /** Number of indexed tools (including duplicates). */
     get size(): number;
     /** Total estimated token cost of all indexed tool schemas. */
@@ -174,6 +215,7 @@ declare class ToolIndex {
     /** Build a single searchable string from tool metadata. */
     private buildSearchableText;
     private getDocumentKey;
+    private matchesServer;
     /** Simple whitespace + camelCase + snake_case tokenizer. */
     private tokenize;
     /** Cosine similarity between two vectors. */
@@ -262,7 +304,7 @@ declare class ToolRouter {
      * This is the main method adapters should call.
      *
      * - `all`    → returns all tools (unchanged behavior)
-     * - `search` → returns only meta-tools (mcp_search_tool_bm25, mcp_get_tool_schema, mcp_execute_tool)
+     * - `search` → returns only meta-tools (mcp_search_tools, mcp_get_tool_schema, mcp_execute_tool)
      * - `groups` → returns tools from active groups only
      */
     getFilteredTools(): Promise<Tool[]>;
@@ -270,17 +312,24 @@ declare class ToolRouter {
      * Search tools by natural-language query.
      * Works regardless of strategy.
      */
-    searchTools(query: string, topK?: number): Promise<ToolSummary[]>;
+    searchTools(query: string, topK?: number, options?: ToolSearchOptions): Promise<ToolSummary[]>;
     /**
      * Search tools by regex pattern.
      * Matches against name, description, and parameter metadata.
      */
     searchToolsRegex(pattern: string, topK?: number): Promise<ToolSummary[]>;
+    /** List connected MCP servers with indexed tool counts. */
+    listServers(options?: ToolSearchOptions): Promise<ToolServerSummary[]>;
+    /** List tools deterministically, optionally scoped to a server. */
+    listTools(options?: ToolSearchOptions & {
+        limit?: number;
+        cursor?: string;
+    }): Promise<ToolListResult>;
     /**
      * Get the full tool definition by name.
      * If tool name is ambiguous, use namespace to specify the server.
      */
-    getToolSchema(toolName: string, namespace?: string): IndexedTool | undefined;
+    getToolSchema(toolName: string, namespace?: string, options?: ToolLookupOptions): IndexedTool | undefined;
     /**
      * Get compact (schema-less) summaries for all tools.
      */
@@ -325,4 +374,4 @@ declare class ToolRouter {
     private getMetaToolDefinitions;
 }
 
-export { type CompactTool as C, type EmbedFn as E, type IndexedTool as I, SchemaCompressor as S, type ToolGroupInfo as T, type CompressionStats as a, ToolIndex as b, type ToolIndexOptions as c, ToolRouter as d, type ToolRouterClientInput as e, type ToolRouterOptions as f, type ToolRouterStrategy as g, type ToolSummary as h };
+export { type CompactTool as C, type EmbedFn as E, type IndexedTool as I, SchemaCompressor as S, type ToolGroupInfo as T, type CompressionStats as a, ToolIndex as b, type ToolIndexOptions as c, type ToolListResult as d, ToolRouter as e, type ToolRouterClientInput as f, type ToolRouterOptions as g, type ToolRouterStrategy as h, type ToolSearchOptions as i, type ToolServerSummary as j, type ToolSummary as k };

package/dist/{tool-router-DK0RJblO.d.mts → tool-router-_O2tIwf7.d.mts}

@@ -80,6 +80,39 @@ interface ToolSummary {
     /** Estimated token cost of the full inputSchema */
     estimatedTokens: number;
 }
+/** Server-level summary derived from indexed tools. */
+interface ToolServerSummary {
+    /** Human-readable server name */
+    serverName: string;
+    /** Stable server identifier */
+    serverId: string;
+    /** Session the server belongs to */
+    sessionId: string;
+    /** Number of indexed tools for this server */
+    toolCount: number;
+}
+/** Optional filters for search and listing. */
+interface ToolSearchOptions {
+    /** Restrict results to this server ID. */
+    serverId?: string;
+    /** Restrict results to servers whose name or ID matches this value. */
+    serverName?: string;
+}
+/** Paginated tool listing result. */
+interface ToolListResult {
+    tools: ToolSummary[];
+    totalCount: number;
+    returnedCount: number;
+    nextCursor?: string;
+    servers: ToolServerSummary[];
+}
+interface ToolLookupOptions {
+    /**
+     * Allow namespace to match a fragment of serverName after exact
+     * sessionId/serverId matching fails.
+     */
+    allowServerNameFragment?: boolean;
+}
 /** A tool with routing metadata attached during indexing. */
 interface IndexedTool extends Tool {
     sessionId: string;
@@ -147,7 +180,7 @@ declare class ToolIndex {
      *
      *   `score = keywordWeight × keyword_score + (1 - keywordWeight) × cosine_score`
      */
-    search(query: string, topK?: number): Promise<ToolSummary[]>;
+    search(query: string, topK?: number, options?: ToolSearchOptions): Promise<ToolSummary[]>;
     /**
      * Search tools using a regex pattern.
      * Matches against name, description, and parameter metadata.
@@ -155,11 +188,19 @@ declare class ToolIndex {
     searchRegex(pattern: string, topK?: number): ToolSummary[];
     /**
      * Get tool definition(s) by name.
-     * If namespace is provided, it tries to match sessionId or serverName.
+     * If namespace is provided, exact sessionId/serverId matches take precedence.
+     * Falls back to serverName fragment matching only when explicitly allowed.
      */
-    getTool(name: string, namespace?: string): IndexedTool[];
+    getTool(name: string, namespace?: string, options?: ToolLookupOptions): IndexedTool[];
     /** All indexed tool names. */
     getToolNames(): string[];
+    /** List indexed servers with tool counts. */
+    listServers(options?: ToolSearchOptions): ToolServerSummary[];
+    /** List tools deterministically, optionally scoped to a server. */
+    listTools(options?: ToolSearchOptions & {
+        limit?: number;
+        cursor?: string;
+    }): ToolListResult;
     /** Number of indexed tools (including duplicates). */
     get size(): number;
     /** Total estimated token cost of all indexed tool schemas. */
@@ -174,6 +215,7 @@ declare class ToolIndex {
     /** Build a single searchable string from tool metadata. */
     private buildSearchableText;
     private getDocumentKey;
+    private matchesServer;
     /** Simple whitespace + camelCase + snake_case tokenizer. */
     private tokenize;
     /** Cosine similarity between two vectors. */
@@ -262,7 +304,7 @@ declare class ToolRouter {
      * This is the main method adapters should call.
      *
      * - `all`    → returns all tools (unchanged behavior)
-     * - `search` → returns only meta-tools (mcp_search_tool_bm25, mcp_get_tool_schema, mcp_execute_tool)
+     * - `search` → returns only meta-tools (mcp_search_tools, mcp_get_tool_schema, mcp_execute_tool)
      * - `groups` → returns tools from active groups only
      */
     getFilteredTools(): Promise<Tool[]>;
@@ -270,17 +312,24 @@ declare class ToolRouter {
      * Search tools by natural-language query.
      * Works regardless of strategy.
      */
-    searchTools(query: string, topK?: number): Promise<ToolSummary[]>;
+    searchTools(query: string, topK?: number, options?: ToolSearchOptions): Promise<ToolSummary[]>;
     /**
      * Search tools by regex pattern.
      * Matches against name, description, and parameter metadata.
      */
     searchToolsRegex(pattern: string, topK?: number): Promise<ToolSummary[]>;
+    /** List connected MCP servers with indexed tool counts. */
+    listServers(options?: ToolSearchOptions): Promise<ToolServerSummary[]>;
+    /** List tools deterministically, optionally scoped to a server. */
+    listTools(options?: ToolSearchOptions & {
+        limit?: number;
+        cursor?: string;
+    }): Promise<ToolListResult>;
     /**
      * Get the full tool definition by name.
      * If tool name is ambiguous, use namespace to specify the server.
      */
-    getToolSchema(toolName: string, namespace?: string): IndexedTool | undefined;
+    getToolSchema(toolName: string, namespace?: string, options?: ToolLookupOptions): IndexedTool | undefined;
     /**
      * Get compact (schema-less) summaries for all tools.
      */
@@ -325,4 +374,4 @@ declare class ToolRouter {
     private getMetaToolDefinitions;
 }
 
-export { type CompactTool as C, type EmbedFn as E, type IndexedTool as I, SchemaCompressor as S, type ToolGroupInfo as T, type CompressionStats as a, ToolIndex as b, type ToolIndexOptions as c, ToolRouter as d, type ToolRouterClientInput as e, type ToolRouterOptions as f, type ToolRouterStrategy as g, type ToolSummary as h };
+export { type CompactTool as C, type EmbedFn as E, type IndexedTool as I, SchemaCompressor as S, type ToolGroupInfo as T, type CompressionStats as a, ToolIndex as b, type ToolIndexOptions as c, type ToolListResult as d, ToolRouter as e, type ToolRouterClientInput as f, type ToolRouterOptions as g, type ToolRouterStrategy as h, type ToolSearchOptions as i, type ToolServerSummary as j, type ToolSummary as k };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mcp-ts/sdk",
-  "version": "1.5.2",
+  "version": "1.6.0",
   "publishConfig": {
     "access": "public"
   },
@@ -84,6 +84,8 @@
     "test": "npx playwright test",
     "test:ui": "npx playwright test --ui",
     "test:debug": "npx playwright test --debug",
+    "benchmark:toolrouter": "npm run build && node benchmarks/toolrouter-efficiency.mjs",
+    "benchmark:toolrouter:live": "npm run build && node benchmarks/toolrouter-live.mjs",
     "supabase:push": "supabase db push",
     "supabase:reset": "supabase db reset",
     "supabase:status": "supabase status",
@@ -160,7 +162,7 @@
     "nanoid": "^5.1.6"
   },
   "devDependencies": {
-    "@ag-ui/client": "^0.0.42",
+    "@ag-ui/client": "^0.0.52",
     "@langchain/core": "^1.1.39",
     "@playwright/test": "^1.58.0",
     "@supabase/supabase-js": "^2.48.0",
@@ -185,4 +187,4 @@
   "engines": {
     "node": ">=18.0.0"
   }
-}
+}

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 };