mirror of
https://github.com/QwenLM/qwen-code.git
synced 2026-08-13 02:35:57 +00:00
fix(core): restore deferred MCP tools on resumed sessions (#8475)
* fix(core): restore deferred MCP tools on resume (#8433) * fix(core): reconcile resumed deferred MCP tools --------- Co-authored-by: Shaojin Wen <shaojin.wensj@alibaba-inc.com>
This commit is contained in:
Nothing Chan
GitHub
parent
6c330c17bb
commit
4ac606dd18
3 changed files with 365 additions and 23 deletions
|
|
@ -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', () => {
|
||||
|
|
|
|||
|
|
@ -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 =
|
||||
'<system-reminder>unfinished todo: run tests</system-reminder>';
|
||||
|
|
@ -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<string>;
|
||||
announcedMcpToolNames: Set<string>;
|
||||
};
|
||||
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: '<system-reminder>\nchanged mcp: added= removed=mcp__calculator__add\n</system-reminder>',
|
||||
},
|
||||
],
|
||||
});
|
||||
});
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -738,14 +738,24 @@ export class GeminiClient {
|
|||
this.forceFullIdeContext = true;
|
||||
}
|
||||
|
||||
async setTools(): Promise<void> {
|
||||
async setTools(options: { skipHistoryReveal?: boolean } = {}): Promise<void> {
|
||||
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;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue