diff --git a/integration-tests/sdk-typescript/sdk-mcp-server.test.ts b/integration-tests/sdk-typescript/sdk-mcp-server.test.ts index 8cb28f0b88..bfc15114f0 100644 --- a/integration-tests/sdk-typescript/sdk-mcp-server.test.ts +++ b/integration-tests/sdk-typescript/sdk-mcp-server.test.ts @@ -207,6 +207,112 @@ describe('SDK MCP Server Integration (E2E)', () => { await q.close(); } }); + + it('keeps previously used MCP tools available when resuming a session', async () => { + // Resume needs a persisted transcript; the rest of this suite keeps + // recording disabled so enable it only for this case. + testDir = await helper.setup('sdk-mcp-server-integration', { + chatRecording: true, + }); + + 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) }], + }), + ); + const serverConfig = createSdkMcpServer({ + name: 'sdk-calculator', + version: '1.0.0', + tools: [calculatorTool], + }); + let streamingRequestIndex = 0; + fakeResponse = ({ body }) => { + if (body['stream'] !== true) { + return { content: '{"selected_memories":[]}' }; + } + const requestIndex = streamingRequestIndex++; + if (requestIndex === 0) { + return { + toolCalls: [ + fakeToolCall('tool_search', { + query: `select:${MCP_CALCULATE_SUM}`, + }), + ], + }; + } + if (requestIndex === 1) { + return { + toolCalls: [fakeToolCall(MCP_CALCULATE_SUM, { a: 25, b: 17 })], + }; + } + if (requestIndex === 3) { + // The resumed model calls the historical tool directly, without a + // second tool_search request. + return { + toolCalls: [fakeToolCall(MCP_CALCULATE_SUM, { a: 8, b: 5 })], + }; + } + return { content: 'Done.' }; + }; + + const firstQuery = query({ + prompt: 'Calculate 25 + 17.', + options: { + ...SHARED_TEST_OPTIONS, + ...fakeModelOptions(fakeServer.baseUrl), + cwd: testDir, + mcpServers: { 'sdk-calculator': serverConfig }, + }, + }); + const firstMessages: SDKMessage[] = []; + const sessionId = firstQuery.getSessionId(); + try { + for await (const message of firstQuery) { + firstMessages.push(message); + } + expect( + findToolResults(firstMessages, MCP_CALCULATE_SUM)[0]?.content, + ).toContain('42'); + assertSuccessfulCompletion(firstMessages); + } finally { + await firstQuery.close(); + } + + const resumedQuery = query({ + prompt: 'Now calculate 8 + 5 with the same tool.', + options: { + ...SHARED_TEST_OPTIONS, + ...fakeModelOptions(fakeServer.baseUrl), + cwd: testDir, + resume: sessionId, + mcpServers: { 'sdk-calculator': serverConfig }, + }, + }); + const resumedMessages: SDKMessage[] = []; + try { + for await (const message of resumedQuery) { + resumedMessages.push(message); + } + + expect(advertisedToolNames(fakeServer, 3)).toContain(MCP_CALCULATE_SUM); + const resumedResults = findToolResults( + resumedMessages, + MCP_CALCULATE_SUM, + ); + expect(resumedResults).toHaveLength(1); + expect(resumedResults[0]?.isError).toBe(false); + expect(resumedResults[0]?.content).toContain('13'); + assertSuccessfulCompletion(resumedMessages); + } finally { + await resumedQuery.close(); + } + }); }); describe('SDK MCP Tool Error Handling', () => { diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index c70dbd2e97..f46ad2d0ae 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -1335,6 +1335,9 @@ describe('Gemini Client (client.ts)', () => { { functionCall: { name: 'cron_create', args: {} }, } as never, + { + functionCall: { name: 'removed_deferred_tool', args: {} }, + } as never, ], }, ]); @@ -1342,6 +1345,33 @@ describe('Gemini Client (client.ts)', () => { expect(reg.revealDeferredTool).toHaveBeenCalledWith('cron_create'); // cron_list NOT in history → must NOT be revealed by the resume scan. expect(reg.revealDeferredTool).not.toHaveBeenCalledWith('cron_list'); + // A historical call whose tool is no longer registered must stay absent. + expect(reg.revealDeferredTool).not.toHaveBeenCalledWith( + 'removed_deferred_tool', + ); + expect(mockClientDebugLogger.debug).toHaveBeenCalledWith( + '[DEFERRED_TOOLS] revealed from history: cron_create', + ); + }); + + it('does not scan resumed history again from startChat setTools', async () => { + const reg = getRegistryMock(); + reg.getDeferredToolSummary.mockReturnValue([ + { name: 'cron_create', description: 'schedule' }, + ]); + reg.getTool.mockImplementation((name: string) => + name === 'tool_search' ? ({} as never) : null, + ); + const getHistorySpy = vi.spyOn(client, 'getHistoryShallow'); + + await client.startChat([ + { + role: 'model', + parts: [{ functionCall: { name: 'cron_create', args: {} } } as never], + }, + ]); + + expect(getHistorySpy).not.toHaveBeenCalled(); }); it('eagerly reveals every deferred tool when ToolSearch is unavailable', async () => { @@ -1870,6 +1900,19 @@ describe('Gemini Client (client.ts)', () => { } } + it('avoids reading history without hidden deferred tools and resolves one summary', async () => { + const reg = getRegistryMock(); + reg.getDeferredToolSummary.mockReturnValue([]); + const getHistorySpy = vi.spyOn(client, 'getHistoryShallow'); + vi.spyOn(client.getChat(), 'setTools').mockImplementation(() => {}); + reg.getDeferredToolSummary.mockClear(); + + await client.setTools(); + + expect(getHistorySpy).not.toHaveBeenCalled(); + expect(reg.getDeferredToolSummary).toHaveBeenCalledTimes(1); + }); + it('carries active todos after tool results and clears them for new work', async () => { const reminder = 'unfinished todo: run tests'; @@ -2193,6 +2236,88 @@ describe('Gemini Client (client.ts)', () => { }); }); + it('does not announce a still-registered tool as removed after history reveals it', async () => { + const reg = getRegistryMock(); + const tool = { + name: 'mcp__calculator__add', + description: 'Add two numbers', + serverName: 'calculator', + }; + let revealed = false; + let registered = true; + reg.getTool.mockImplementation((name: string) => + name === 'tool_search' || (name === tool.name && registered) + ? ({} as never) + : null, + ); + reg.getDeferredToolSummary.mockImplementation(() => + registered ? [tool] : [], + ); + reg.isDeferredToolRevealed.mockImplementation( + (name: string) => name === tool.name && revealed, + ); + reg.revealDeferredTool.mockImplementation((name: string) => { + if (name === tool.name) revealed = true; + }); + vi.spyOn(client.getChat(), 'setTools').mockImplementation(() => {}); + const addHistorySpy = vi.spyOn(client.getChat(), 'addHistory'); + const reminderState = client as unknown as { + announcedDeferredToolNames: Set; + announcedMcpToolNames: Set; + }; + reminderState.announcedDeferredToolNames = new Set([tool.name]); + reminderState.announcedMcpToolNames = new Set([tool.name]); + + client.setHistory([ + { + role: 'model', + parts: [ + { + functionCall: { name: tool.name, args: { a: 1, b: 2 } }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + name: tool.name, + response: { output: '3' }, + }, + }, + ], + }, + ]); + + await client.setTools(); + await runTurn(); + + expect(revealed).toBe(true); + expect(buildChangedMcpToolsReminder).not.toHaveBeenCalled(); + expect(addHistorySpy).not.toHaveBeenCalled(); + + registered = false; + vi.mocked(buildChangedMcpToolsReminder).mockClear(); + addHistorySpy.mockClear(); + + await client.setTools(); + await runTurn(); + + expect(buildChangedMcpToolsReminder).toHaveBeenCalledWith( + [], + [tool.name], + ); + expect(addHistorySpy).toHaveBeenCalledWith({ + role: 'user', + parts: [ + { + text: '\nchanged mcp: added= removed=mcp__calculator__add\n', + }, + ], + }); + }); + it('keeps queued MCP changes when the reminder builder returns null', () => { const priv = client as unknown as { pendingAddedMcpTools: Map< @@ -2212,6 +2337,58 @@ describe('Gemini Client (client.ts)', () => { ); }); + it('re-reveals MCP tools from resumed history after progressive discovery', async () => { + const reg = getRegistryMock(); + reg.getTool.mockImplementation((name: string) => + name === 'tool_search' ? ({} as never) : null, + ); + + // The resumed chat is constructed before progressive MCP discovery, so + // startChat() cannot match this historical call until the server's tools + // are registered. setTools() is the common refresh path once they are. + client.setHistory([ + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call-resumed-mcp', + name: 'mcp__calculator__add', + args: { a: 1, b: 2 }, + }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'call-resumed-mcp', + name: 'mcp__calculator__add', + response: { output: '3' }, + }, + }, + ], + }, + ]); + reg.getDeferredToolSummary.mockReturnValue([ + { + name: 'mcp__calculator__add', + description: 'Add two numbers', + serverName: 'calculator', + }, + ]); + reg.revealDeferredTool.mockClear(); + vi.spyOn(client.getChat(), 'setTools').mockImplementation(() => {}); + + await client.setTools(); + + expect(reg.revealDeferredTool).toHaveBeenCalledWith( + 'mcp__calculator__add', + ); + }); + it('eagerly reveals every deferred tool when ToolSearch is unavailable', async () => { // Mirrors startChat's silent-disappearance guard: without ToolSearch // a deferred MCP tool can't be reached, so the only safe option is diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index d093a54bdb..18cb53216d 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -738,14 +738,24 @@ export class GeminiClient { this.forceFullIdeContext = true; } - async setTools(): Promise { + async setTools(options: { skipHistoryReveal?: boolean } = {}): Promise { if (!this.isInitialized()) { return; } const toolRegistry = this.config.getToolRegistry(); await toolRegistry.warmAll(); - const deferredTools = this.resolveDeferredToolsForReminder(); + const deferredSummary = toolRegistry.getDeferredToolSummary(); + // Progressive MCP discovery registers tools after a resumed chat has + // already been constructed. Re-scan the live history here so historical + // MCP calls reveal their newly registered schemas before declarations are + // refreshed. setTools() is shared by interactive and headless refreshes. + if (!options.skipHistoryReveal) { + this.revealDeferredToolsReferencedInHistory(deferredSummary, () => + this.getHistoryShallow(), + ); + } + const deferredTools = this.resolveDeferredToolsForReminder(deferredSummary); const toolDeclarations = toolRegistry.getFunctionDeclarations(); const tools: Tool[] = [{ functionDeclarations: toolDeclarations }]; this.getChat().setTools(tools); @@ -1116,6 +1126,54 @@ export class GeminiClient { ); } + /** + * Reveals deferred tools referenced by function calls in existing history. + * + * On resume this runs once before startup reminders are built. It also runs + * from setTools() because progressive MCP discovery can register deferred + * tools only after the resumed chat and its initial declarations exist. + */ + private revealDeferredToolsReferencedInHistory( + deferredSummary: readonly DeferredToolSummary[], + getHistory: () => readonly Content[] | undefined, + ): void { + const toolRegistry = this.config.getToolRegistry(); + const deferredNames = new Set( + deferredSummary + .filter((tool) => !toolRegistry.isDeferredToolRevealed(tool.name)) + .map((tool) => tool.name), + ); + if (deferredNames.size === 0) { + return; + } + + // Reading live history is O(history), so defer it until the registry proves + // there is at least one hidden deferred tool that could be matched. + const history = getHistory(); + if (!history || history.length === 0) { + return; + } + + const revealedNames: string[] = []; + for (const entry of history) { + for (const part of entry.parts ?? []) { + const callName = part.functionCall?.name; + if (callName && deferredNames.delete(callName)) { + toolRegistry.revealDeferredTool(callName); + revealedNames.push(callName); + } + } + if (deferredNames.size === 0) { + break; + } + } + if (revealedNames.length > 0) { + debugLogger.debug( + `[DEFERRED_TOOLS] revealed from history: ${revealedNames.join(', ')}`, + ); + } + } + /** * Computes the deferred-tools list that should be announced through * user-role system reminders. @@ -1134,9 +1192,10 @@ export class GeminiClient { * Returns `undefined` when ToolSearch is unavailable: reminders must not * advertise tools the model has no way to load on demand. */ - private resolveDeferredToolsForReminder(): DeferredToolSummary[] | undefined { + private resolveDeferredToolsForReminder( + deferredSummary: readonly DeferredToolSummary[], + ): DeferredToolSummary[] | undefined { const toolRegistry = this.config.getToolRegistry(); - const deferredSummary = toolRegistry.getDeferredToolSummary(); const toolSearchAvailable = !!toolRegistry.getTool(ToolNames.TOOL_SEARCH); if (!toolSearchAvailable) { if (deferredSummary.length > 0) { @@ -1169,6 +1228,7 @@ export class GeminiClient { private queueAddedMcpToolsReminder( deferredTools: readonly DeferredToolSummary[], ): void { + const toolRegistry = this.config.getToolRegistry(); const currentDeferredNames = new Set( deferredTools.map((tool) => tool.name), ); @@ -1181,7 +1241,7 @@ export class GeminiClient { } } for (const name of this.pendingRemovedMcpToolNames) { - if (currentMcpToolNames.has(name)) { + if (currentMcpToolNames.has(name) || toolRegistry.getTool(name)) { this.pendingRemovedMcpToolNames.delete(name); } } @@ -1197,7 +1257,14 @@ export class GeminiClient { } } for (const name of this.announcedMcpToolNames) { - if (!currentMcpToolNames.has(name)) { + if (currentMcpToolNames.has(name)) { + continue; + } + // A revealed or newly-visible tool is absent from the deferred reminder + // summary but still present in the registry. Keep tracking it as + // model-visible so a later real disconnect can still be announced; only + // a tool actually removed from the registry is unavailable now. + if (!toolRegistry.getTool(name)) { this.pendingRemovedMcpToolNames.add(name); } } @@ -1481,6 +1548,7 @@ export class GeminiClient { // calling us. const toolRegistry = this.config.getToolRegistry(); await profiler.time('tool_registry_warm', () => toolRegistry.warmAll()); + const deferredSummary = toolRegistry.getDeferredToolSummary(); // Resume support: when a transcript contains prior calls to a deferred // tool, re-reveal that tool so `setTools()` below sends its schema in // the declaration list. Without this, the model sees history like @@ -1489,21 +1557,10 @@ export class GeminiClient { // BEFORE `resolveDeferredToolsForReminder()` runs so the resumed tools // are correctly filtered out of the startup reminder built below. profiler.timeSync('resume_deferred_tool_reveal', () => { - if (extraHistory && extraHistory.length > 0) { - const deferredNames = new Set( - toolRegistry.getDeferredToolSummary().map((t) => t.name), - ); - if (deferredNames.size > 0) { - for (const entry of extraHistory) { - for (const part of entry.parts ?? []) { - const callName = part.functionCall?.name; - if (callName && deferredNames.has(callName)) { - toolRegistry.revealDeferredTool(callName); - } - } - } - } - } + this.revealDeferredToolsReferencedInHistory( + deferredSummary, + () => extraHistory, + ); }); // Budget-based deferred-tool preload runs BEFORE the deferred // reminder is resolved so preloaded tools are filtered out of the @@ -1512,7 +1569,7 @@ export class GeminiClient { this.preloadDeferredToolsWithinBudget(); }); const deferredTools = profiler.timeSync('deferred_reminder_setup', () => { - const resolved = this.resolveDeferredToolsForReminder(); + const resolved = this.resolveDeferredToolsForReminder(deferredSummary); this.rememberAnnouncedDeferredTools(resolved); return resolved; }); @@ -1584,7 +1641,9 @@ export class GeminiClient { // setTools() intentionally keeps its own warmAll() guard, so this stage // overlaps with tool_registry_warm while preserving the startup path. - await profiler.time('set_tools', () => this.setTools()); + await profiler.time('set_tools', () => + this.setTools({ skipHistoryReveal: true }), + ); finishProfile(true); return this.chat;