From 48865472657725b75075b5ad3b846d7b28de136f Mon Sep 17 00:00:00 2001 From: yiliang114 Date: Tue, 28 Jul 2026 15:48:18 +0800 Subject: [PATCH] fix(core): preserve active Todo context across tool turns --- docs/design/active-todo-context.md | 33 ++++++ .../acp-integration/session/Session.test.ts | 101 +++++++++++++++++- .../src/acp-integration/session/Session.ts | 70 +++++++++--- packages/cli/src/nonInteractiveCli.test.ts | 8 +- packages/cli/src/nonInteractiveCli.ts | 9 +- packages/core/src/config/config.test.ts | 56 ++++++++++ packages/core/src/config/config.ts | 40 +++++++ packages/core/src/core/client.test.ts | 56 ++++++++++ packages/core/src/core/client.ts | 22 ++++ .../core/src/core/coreToolScheduler.test.ts | 4 + packages/core/src/core/coreToolScheduler.ts | 27 +++-- packages/core/src/tools/todoWrite.test.ts | 50 ++++++++- packages/core/src/tools/todoWrite.ts | 22 ++++ 13 files changed, 459 insertions(+), 39 deletions(-) create mode 100644 docs/design/active-todo-context.md diff --git a/docs/design/active-todo-context.md b/docs/design/active-todo-context.md new file mode 100644 index 00000000000..31a53d901da --- /dev/null +++ b/docs/design/active-todo-context.md @@ -0,0 +1,33 @@ +# Active Todo Context + +## Problem + +`todo_write` presents the current list as a reminder only in its own tool +result. After more tool calls, that reminder loses salience and the model may +end the turn with unfinished items. The persisted todo file is unsuitable as +live control state because it can outlive the work chain that created it. + +## Design + +After a successful `todo_write`, keep a reminder containing only unfinished +items, keyed by the prompt ID that owns the work chain. Append it after function +responses on subsequent tool-result turns with the same owner in both the core +and ACP loops. This isolates ordinary prompts, cron jobs, and background +notifications even when they share a session. Clear the reminder when all todos +complete, a new work chain starts, or the session changes. Retry, continue, and +explicitly related automatic requests move the reminder to the new prompt ID +because they resume the same work chain. The repeated reminder is capped at +4,000 characters. + +This does not change stop semantics or enable `todoStopGuard`. The guard remains +an optional bounded recovery after a model has already tried to stop; this +change instead preserves task context before that decision. + +## Verification + +- A successful write with unfinished items updates the session reminder. +- A completed list clears it. +- Core and ACP tool-result messages append the reminder after function results. +- ACP mid-turn user input remains last and therefore keeps precedence. +- An ordinary new prompt clears stale state while retry/continue retains it. +- Independent automatic turns are isolated; related automatic turns inherit. diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 927781ed5c5..65a91e838da 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -563,6 +563,10 @@ describe('Session', () => { switchModel: switchModelSpy, getModel: vi.fn().mockImplementation(() => currentModel), getSessionId: vi.fn().mockReturnValue('test-session-id'), + getActiveTodoReminder: vi.fn().mockReturnValue(undefined), + setActiveTodoReminder: vi.fn(), + startActiveTodoWorkChain: vi.fn(), + startAutomaticActiveTodoWorkChain: vi.fn(), assertCanStartTurn: vi.fn().mockResolvedValue(undefined), getWorkingDir: vi.fn().mockReturnValue(process.cwd()), getProjectRoot: vi.fn().mockReturnValue('/repo'), @@ -759,6 +763,66 @@ describe('Session', () => { expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); }); + it('clears active todo context when an ordinary prompt starts', async () => { + mockChat.sendMessageStream = vi + .fn() + .mockImplementation(async () => createEmptyStream()); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'start different work' }], + }); + + expect(mockConfig.startActiveTodoWorkChain).toHaveBeenCalledWith( + 'test-session-id########1', + undefined, + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'start different work' }], + retry: true, + } as PromptRequest); + + expect(mockConfig.startActiveTodoWorkChain).toHaveBeenCalledWith( + 'test-session-id########2', + 'test-session-id########1', + ); + }); + + it('continues active Todo context for related automatic turns', async () => { + mockChat.sendMessageStream = vi + .fn() + .mockImplementation(async () => createEmptyStream()); + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'start work' }], + }); + vi.mocked(mockConfig.startAutomaticActiveTodoWorkChain).mockClear(); + const internals = session as unknown as { + relatedAgentIds: Set; + }; + internals.relatedAgentIds.add('related-agent'); + const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock + .calls[0][0] as ( + displayText: string, + modelText: string, + meta: { agentId: string; status: string }, + ) => void; + + callback('Background task completed.', '', { + agentId: 'related-agent', + status: 'completed', + }); + + await vi.waitFor(() => + expect(mockConfig.startAutomaticActiveTodoWorkChain).toHaveBeenCalledWith( + expect.stringContaining('########notification'), + 'test-session-id########1', + ), + ); + }); + it('holds the close gate until active turns settle', async () => { let resolveTurn!: () => void; const turnCompletion = new Promise((resolve) => { @@ -6056,9 +6120,26 @@ describe('Session', () => { }); it('injects drained mid-turn user messages with tool responses', async () => { - const executeSpy = vi.fn().mockResolvedValue({ - llmContent: 'file contents', - returnDisplay: 'file contents', + const todoReminder = + 'unfinished todo: check tests'; + const activeTodoReminders = new Map(); + vi.mocked(mockConfig.getActiveTodoReminder).mockImplementation( + (promptId) => activeTodoReminders.get(promptId), + ); + vi.mocked(mockConfig.setActiveTodoReminder).mockImplementation( + (promptId, reminder) => { + if (reminder) activeTodoReminders.set(promptId, reminder); + }, + ); + const executeSpy = vi.fn().mockImplementation(async () => { + const promptId = core.promptIdContext.getStore(); + if (promptId) { + mockConfig.setActiveTodoReminder(promptId, todoReminder); + } + return { + llmContent: 'file contents', + returnDisplay: 'file contents', + }; }); const tool = { name: 'read_file', @@ -6110,9 +6191,19 @@ describe('Session', () => { const midTurnPart = { text: '\n[User message received during tool execution]: please also check tests ', }; - expect(secondCall?.[1].message).toEqual( - expect.arrayContaining([midTurnPart]), + const nextMessage = secondCall?.[1].message as Part[]; + const functionResponseIndex = nextMessage.findIndex( + (part) => part.functionResponse !== undefined, + ); + const reminderIndex = nextMessage.findIndex( + (part) => part.text === todoReminder, + ); + const midTurnIndex = nextMessage.findIndex( + (part) => part.text === midTurnPart.text, ); + expect(functionResponseIndex).toBeGreaterThanOrEqual(0); + expect(reminderIndex).toBeGreaterThan(functionResponseIndex); + expect(midTurnIndex).toBeGreaterThan(reminderIndex); expect( mockChatRecordingService.recordMidTurnUserMessage, ).toHaveBeenCalledWith([midTurnPart], ' please also check tests '); diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 1d0b35f987a..a3d9516ef67 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -1236,6 +1236,7 @@ export class Session implements SessionContext { */ private followupAbort: AbortController | null = null; private turn: number = 0; + private activeTodoWorkChainPromptId: string | undefined; private readonly createdAt: number = Date.now(); /** * Running cumulative usage for this session, snapshotted onto each todo/plan @@ -1345,14 +1346,8 @@ export class Session implements SessionContext { !this.config.getBareMode() && !this.config.isSafeMode(); this.todoStopGuard = new DaemonTodoStopGuard(todoStopGuardEnabled); - this.todoStopGuardBackgroundBaseline = todoStopGuardEnabled - ? this.#captureTodoStopGuardBackgroundBaseline() - : { - agents: new Set(), - shells: new Set(), - monitors: new Set(), - wakeups: new Set(), - }; + this.todoStopGuardBackgroundBaseline = + this.#captureTodoStopGuardBackgroundBaseline(); // Initialize modular components with this session as context this.toolCallEmitter = new ToolCallEmitter(this); @@ -2634,6 +2629,22 @@ export class Session implements SessionContext { this.turn += 1; const promptId = this.config.getSessionId() + '########' + this.turn; + const promptMetadata = (params as { _meta?: Record }) + ._meta; + const continuesCurrentWorkChain = + (params as { retry?: boolean }).retry === true || + promptMetadata?.[DAEMON_RETRY_META_KEY] === true || + promptMetadata?.[DAEMON_CONTINUE_META_KEY] === true; + if (!continuesCurrentWorkChain && !this.todoStopGuard.enabled) { + this.#resetTodoStopGuardBackgroundLineage(); + } + this.config.startActiveTodoWorkChain( + promptId, + continuesCurrentWorkChain + ? this.activeTodoWorkChainPromptId + : undefined, + ); + this.activeTodoWorkChainPromptId = promptId; // Bind the prompt ID for the remainder of this turn, mirroring the // sessionIdContext.run wrapper in #executePrompt. Shell subprocesses // read it via getShellContextEnvVars (QWEN_CODE_PROMPT_ID) — without @@ -3187,6 +3198,7 @@ export class Session implements SessionContext { await this.#buildNextMessageAfterToolRun( toolRun, pendingSend.signal, + promptId, onFullTurnModel, ); nextMessage = nextAfterTools.message; @@ -4127,6 +4139,7 @@ export class Session implements SessionContext { const nextAfterTools = await this.#buildNextMessageAfterToolRun( toolRun, pendingSend.signal, + toolPromptId, options.onFullTurnModel, ); nextMessage = nextAfterTools.message; @@ -4530,6 +4543,7 @@ export class Session implements SessionContext { async #buildNextMessageAfterToolRun( toolRun: RunToolResult, abortSignal: AbortSignal, + promptId: string, onFullTurnModel?: (model: string) => boolean, ): Promise { if (toolRun.loopDetected) { @@ -4550,7 +4564,12 @@ export class Session implements SessionContext { if (hadMidTurnUserInput) { this.todoStopGuard.acceptMidTurnUserInput(); } - const parts = [...toolRun.parts, ...drained.parts]; + const activeTodoReminder = this.config.getActiveTodoReminder(promptId); + const parts = [ + ...toolRun.parts, + ...(activeTodoReminder ? [{ text: activeTodoReminder }] : []), + ...drained.parts, + ]; return { message: { role: 'user', parts }, hadMidTurnUserInput, @@ -5131,9 +5150,9 @@ export class Session implements SessionContext { async () => { const ac = new AbortController(); this.cronAbortController = ac; - this.#prepareTodoStopGuardForAutomaticTurn( - this.#cronContinuesTodoStopGuardWorkChain(item), - ); + const continuesCurrentWorkChain = + this.#cronContinuesTodoStopGuardWorkChain(item); + this.#prepareTodoStopGuardForAutomaticTurn(continuesCurrentWorkChain); const promptId = this.config.getSessionId() + '########cron' + Date.now(); let cronHadError = false; @@ -5153,6 +5172,15 @@ export class Session implements SessionContext { try { await this.assertCanStartTurn(); if (ac.signal.aborted) return; + this.config.startAutomaticActiveTodoWorkChain( + promptId, + continuesCurrentWorkChain + ? this.activeTodoWorkChainPromptId + : undefined, + ); + if (continuesCurrentWorkChain) { + this.activeTodoWorkChainPromptId = promptId; + } // A `<>` / `<>` sentinel is expanded at // fire time into the loop.md task block — full on the first or a // changed fire, a short reminder when unchanged. Non-sentinel @@ -5455,6 +5483,7 @@ export class Session implements SessionContext { await this.#buildNextMessageAfterToolRun( toolRun, ac.signal, + promptId, ); nextMessage = nextAfterTools.message; if (toolRun.loopDetected) { @@ -5773,14 +5802,23 @@ export class Session implements SessionContext { async () => { const ac = new AbortController(); this.notificationAbortController = ac; - this.#prepareTodoStopGuardForAutomaticTurn( - this.#notificationContinuesTodoStopGuardWorkChain(item), - ); + const continuesCurrentWorkChain = + this.#notificationContinuesTodoStopGuardWorkChain(item); + this.#prepareTodoStopGuardForAutomaticTurn(continuesCurrentWorkChain); const promptId = this.config.getSessionId() + '########notification' + Date.now(); try { await this.assertCanStartTurn(); if (ac.signal.aborted) return; + this.config.startAutomaticActiveTodoWorkChain( + promptId, + continuesCurrentWorkChain + ? this.activeTodoWorkChainPromptId + : undefined, + ); + if (continuesCurrentWorkChain) { + this.activeTodoWorkChainPromptId = promptId; + } await this.#emitBackgroundNotificationDisplay(item); const notificationParts: Part[] = [{ text: item.modelText }]; @@ -5953,6 +5991,7 @@ export class Session implements SessionContext { const nextAfterTools = await this.#buildNextMessageAfterToolRun( toolRun, ac.signal, + promptId, ); nextMessage = nextAfterTools.message; if (toolRun.loopDetected) { @@ -6376,6 +6415,7 @@ export class Session implements SessionContext { functionCalls: FunctionCall[], toolLoopState?: DaemonToolLoopState, ): Promise { + promptIdContext.enterWith(promptId); const dedupedFunctionCalls = dedupeToolCallsById(functionCalls); const generatedCallIdBase = randomUUID(); const executionCallIds = new Map( diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index fc15cc551e9..6d727777f9d 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -1995,6 +1995,12 @@ describe('runNonInteractive', () => { expect(exitCode).toBe(1); expect(mockGeminiClient.sendMessageStream).toHaveBeenCalledTimes(3); expect(mockCoreExecuteToolCall).not.toHaveBeenCalled(); + const drainPromptIds = mockGeminiClient.sendMessageStream.mock.calls + .slice(1) + .map((call) => call[2]); + expect(new Set(drainPromptIds)).toEqual( + new Set(['prompt-id-drain-dup-loop/automatic/2']), + ); const duplicateParts = mockGeminiClient.sendMessageStream.mock .calls[2][0] as Part[]; @@ -3188,7 +3194,7 @@ describe('runNonInteractive', () => { 2, [{ text: notificationXml }], expect.any(AbortSignal), - 'prompt-monitor', + 'prompt-monitor/automatic/2', { type: SendMessageType.Notification, modelOverride: undefined, diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 60b448bf61a..614456b7851 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -1625,6 +1625,7 @@ export async function runNonInteractive( }; }; + let currentPromptId = prompt_id; while (true) { // Drain pending teammate messages into the conversation. // sendMessageStream only reads currentMessages[0].parts, @@ -1678,13 +1679,16 @@ export async function runNonInteractive( } else { sendType = SendMessageType.ToolResult; } + if (isTeammateTurn) { + currentPromptId = `${prompt_id}/teammate/${turnCount}`; + } const toolCallRequests: ToolCallRequestInfo[] = []; const apiStartTime = Date.now(); const responseStream = geminiClient.sendMessageStream( currentMessages[0]?.parts || [], abortController.signal, - prompt_id, + currentPromptId, { type: sendType, modelOverride, @@ -1926,6 +1930,7 @@ export async function runNonInteractive( ]; let itemIsFirstTurn = true; let itemModelOverride: string | undefined; + const itemPromptId = `${prompt_id}/automatic/${turnCount}`; while (true) { const itemToolCallRequests: ToolCallRequestInfo[] = []; @@ -1933,7 +1938,7 @@ export async function runNonInteractive( const itemStream = geminiClient.sendMessageStream( itemMessages[0]?.parts || [], abortController.signal, - prompt_id, + itemPromptId, { type: itemIsFirstTurn ? item.sendMessageType diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index da384624b67..5f5adda90a0 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -8347,4 +8347,60 @@ describe('Model Switching and Config Updates', () => { expect(response.success).toBe(true); }); }); + + it('moves only the continued work chain Todo reminder', () => { + const config = Object.create(Config.prototype) as Config; + config.setActiveTodoReminder('prompt-user', 'unfinished user work'); + config.setActiveTodoReminder('prompt-cron', 'unfinished cron work'); + + config.startActiveTodoWorkChain('prompt-retry', 'prompt-user'); + + expect(config.getActiveTodoReminder('prompt-retry')).toBe( + 'unfinished user work', + ); + expect(config.getActiveTodoReminder('prompt-user')).toBeUndefined(); + expect(config.getActiveTodoReminder('prompt-cron')).toBeUndefined(); + }); + + it('moves related automatic work without clearing unrelated reminders', () => { + const config = Object.create(Config.prototype) as Config; + config.setActiveTodoReminder('prompt-user', 'unfinished user work'); + config.setActiveTodoReminder('prompt-unrelated', 'other work'); + + config.startAutomaticActiveTodoWorkChain('prompt-cron'); + config.startAutomaticActiveTodoWorkChain( + 'prompt-related-notification', + 'prompt-user', + ); + + expect(config.getActiveTodoReminder('prompt-user')).toBeUndefined(); + expect(config.getActiveTodoReminder('prompt-cron')).toBeUndefined(); + expect(config.getActiveTodoReminder('prompt-related-notification')).toBe( + 'unfinished user work', + ); + expect(config.getActiveTodoReminder('prompt-unrelated')).toBe('other work'); + }); + + it('isolates active Todo reminders inherited through child Configs', () => { + const parent = Object.create(Config.prototype) as Config; + const child = Object.create(parent) as Config; + parent.setActiveTodoReminder('parent-prompt', 'parent work'); + + child.setActiveTodoReminder('child-prompt', 'child work'); + child.startActiveTodoWorkChain('child-retry', 'child-prompt'); + + expect(parent.getActiveTodoReminder('parent-prompt')).toBe('parent work'); + expect(parent.getActiveTodoReminder('child-retry')).toBeUndefined(); + expect(child.getActiveTodoReminder('parent-prompt')).toBeUndefined(); + expect(child.getActiveTodoReminder('child-retry')).toBe('child work'); + }); + + it('clears active Todo reminders for a new session', () => { + const config = new Config(baseParams); + config.setActiveTodoReminder('old-prompt', 'unfinished old work'); + + config.startNewSession('new-session-id'); + + expect(config.getActiveTodoReminder('old-prompt')).toBeUndefined(); + }); }); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 8c99912f27d..da82acd61e1 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -1806,6 +1806,7 @@ export class Config { private readonly gitCoAuthor: GitCoAuthorSettings; private readonly usageStatisticsEnabled: boolean; private readonly fileReadCacheDisabled: boolean; + private activeTodoReminders = new Map(); private geminiClient!: GeminiClient; private baseLlmClient!: BaseLlmClient; private cronScheduler: CronScheduler | null = null; @@ -3531,6 +3532,7 @@ export class Config { } this.sessionData = sessionData; this.pendingRecoveredAgentsNotice = null; + this.getOwnActiveTodoReminders().clear(); setDebugLogSession(this); this.debugLogger = createDebugLogger(); this.chatRecordingService = this.chatRecordingEnabled @@ -5840,6 +5842,44 @@ export class Config { return this.geminiClient; } + private getOwnActiveTodoReminders(): Map { + if (!Object.prototype.hasOwnProperty.call(this, 'activeTodoReminders')) { + this.activeTodoReminders = new Map(); + } + return this.activeTodoReminders; + } + + getActiveTodoReminder(promptId: string): string | undefined { + return this.getOwnActiveTodoReminders().get(promptId); + } + + setActiveTodoReminder(promptId: string, reminder: string | undefined): void { + const reminders = this.getOwnActiveTodoReminders(); + if (reminder) { + reminders.set(promptId, reminder); + } else { + reminders.delete(promptId); + } + } + + startActiveTodoWorkChain(promptId: string, continuedFrom?: string): void { + const reminders = this.getOwnActiveTodoReminders(); + const reminder = continuedFrom ? reminders.get(continuedFrom) : undefined; + reminders.clear(); + if (reminder) reminders.set(promptId, reminder); + } + + startAutomaticActiveTodoWorkChain( + promptId: string, + continuedFrom?: string, + ): void { + const reminders = this.getOwnActiveTodoReminders(); + const reminder = continuedFrom ? reminders.get(continuedFrom) : undefined; + if (continuedFrom) reminders.delete(continuedFrom); + reminders.delete(promptId); + if (reminder) reminders.set(promptId, reminder); + } + /** * Session-scoped memory pressure monitor. Child Configs created with * `Object.create(parent)` inherit the parent's monitor through the prototype diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 26cdb85b9ca..38be11527b1 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -559,6 +559,9 @@ describe('Gemini Client (client.ts)', () => { getAppendSystemPrompt: vi.fn().mockReturnValue(undefined), getFullContext: vi.fn().mockReturnValue(false), getSessionId: vi.fn().mockReturnValue('test-session-id'), + getActiveTodoReminder: vi.fn().mockReturnValue(undefined), + startActiveTodoWorkChain: vi.fn(), + startAutomaticActiveTodoWorkChain: vi.fn(), getProxy: vi.fn().mockReturnValue(undefined), getWorkingDir: vi.fn().mockReturnValue('/test/dir'), getFileService: vi.fn().mockReturnValue(fileService), @@ -1709,6 +1712,59 @@ describe('Gemini Client (client.ts)', () => { } } + it('carries active todos after tool results and clears them for new work', async () => { + const reminder = + 'unfinished todo: run tests'; + vi.mocked(mockConfig.getActiveTodoReminder).mockReturnValue(reminder); + + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'response' }; + })(), + ); + const stream = client.sendMessageStream( + [{ functionResponse: { name: 'read_file', response: { ok: true } } }], + new AbortController().signal, + 'prompt-tool-result', + { type: SendMessageType.ToolResult }, + ); + for await (const _ of stream) { + // drain + } + + const request = mockTurnRunFn.mock.lastCall?.[1] as unknown[]; + const functionResponseIndex = request.findIndex( + (part) => + typeof part === 'object' && + part !== null && + 'functionResponse' in part, + ); + expect(functionResponseIndex).toBeGreaterThanOrEqual(0); + expect(request.indexOf(reminder)).toBeGreaterThan(functionResponseIndex); + expect(mockConfig.getActiveTodoReminder).toHaveBeenCalledWith( + 'prompt-tool-result', + ); + + await runTurn(SendMessageType.UserQuery); + + expect(mockConfig.startActiveTodoWorkChain).toHaveBeenCalledWith( + 'prompt-userQuery', + ); + + await runTurn(SendMessageType.Retry); + + expect(mockConfig.startActiveTodoWorkChain).toHaveBeenCalledWith( + 'prompt-retry', + 'prompt-userQuery', + ); + + await runTurn(SendMessageType.Cron); + + expect(mockConfig.startActiveTodoWorkChain).toHaveBeenCalledWith( + 'prompt-cron', + ); + }); + it('queues and drains a reminder for newly registered MCP deferred tools', async () => { const reg = getRegistryMock(); reg.getTool.mockImplementation((n: string) => diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 16cfcb7db0f..6a0f66d9820 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -256,6 +256,7 @@ export class GeminiClient { private readonly loopDetector: LoopDetectionService; private lastPromptId: string | undefined = undefined; + private activeTodoWorkChainPromptId: string | undefined; private lastSentIdeContext: IdeContext | undefined; private forceFullIdeContext = true; private recentCompletedToolNames: string[] = []; @@ -2085,6 +2086,23 @@ export class GeminiClient { // Notifications start a fresh Turn with a new prompt_id, so the loop // detector must reset — otherwise a prior turn's count can trip // LoopDetected early on the notification turn. + if (messageType === SendMessageType.UserQuery) { + this.config.startActiveTodoWorkChain(prompt_id); + this.activeTodoWorkChainPromptId = prompt_id; + } else if (messageType === SendMessageType.Retry) { + this.config.startActiveTodoWorkChain( + prompt_id, + this.activeTodoWorkChainPromptId, + ); + this.activeTodoWorkChainPromptId = prompt_id; + } else if ( + messageType === SendMessageType.Cron || + messageType === SendMessageType.Notification || + messageType === SendMessageType.Teammate + ) { + this.config.startActiveTodoWorkChain(prompt_id); + this.activeTodoWorkChainPromptId = prompt_id; + } const isTopLevelInteraction = messageType === SendMessageType.UserQuery || messageType === SendMessageType.Cron || @@ -2524,6 +2542,10 @@ export class GeminiClient { // text as a separate user message after the tool messages. requestToSend = [...requestToSend, toolResultMemory.prompt]; } + const activeTodoReminder = this.config.getActiveTodoReminder(prompt_id); + if (activeTodoReminder) { + requestToSend = [...requestToSend, activeTodoReminder]; + } await this.microcompactHistoryBeforeSend(null, { sizeOnly: true, pendingContent: createUserContent(requestToSend), diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index bcab894c2c0..518b705edb0 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -84,6 +84,7 @@ import { } from '../utils/invocation-context.js'; import { getPlanModeSystemReminder } from './prompts.js'; import { PLAN_MODE_ENTRY_SIBLING_SKIP_MESSAGE } from './plan-mode-entry-policy.js'; +import { promptIdContext } from '../utils/promptIdContext.js'; type ToolSpanRecord = { name: string; @@ -845,6 +846,7 @@ describe('CoreToolScheduler', () => { promptId: 'unrelated-prompt', }; let observedContext: InvocationContextV1 | undefined; + let observedPromptId: string | undefined; const tool = new MockTool({ name: 'approval-context-tool', getDefaultPermission: async () => 'ask', @@ -856,6 +858,7 @@ describe('CoreToolScheduler', () => { }), execute: async () => { observedContext = getInvocationContext(); + observedPromptId = promptIdContext.getStore(); return { llmContent: 'ok', returnDisplay: 'ok' }; }, }); @@ -890,6 +893,7 @@ describe('CoreToolScheduler', () => { ); expect(observedContext).toEqual(invocationContext); + expect(observedPromptId).toBe(invocationContext.promptId); }); it('isolates enter_plan_mode as a batch boundary and preserves its full reminder', async () => { diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index 250d48e3f1c..afe61c4e278 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -70,6 +70,7 @@ import { type AvailableSkillEntry, } from '../tools/skill-utils.js'; import { escapeSystemReminderTags } from '../utils/xml.js'; +import { promptIdContext } from '../utils/promptIdContext.js'; import { unescapePath, PATH_ARG_KEYS } from '../utils/paths.js'; import type { MemoryPressureMonitor } from '../services/memoryPressureMonitor.js'; import { CONCURRENCY_SAFE_KINDS, isShellProgressData } from '../tools/tools.js'; @@ -4136,20 +4137,24 @@ export class CoreToolScheduler { ); }; this.safelyAddToolArgumentsAttributes(span, invocation.params); - promise = invocation.execute( - execSignal, - liveOutputCallback, - shellExecutionConfig, - setPidCallback, - setPromoteAbortControllerCallback, - canPromoteForegroundShell, + promise = promptIdContext.run(scheduledCall.request.prompt_id, () => + invocation.execute( + execSignal, + liveOutputCallback, + shellExecutionConfig, + setPidCallback, + setPromoteAbortControllerCallback, + canPromoteForegroundShell, + ), ); } else { this.safelyAddToolArgumentsAttributes(span, invocation.params); - promise = invocation.execute( - execSignal, - liveOutputCallback, - shellExecutionConfig, + promise = promptIdContext.run(scheduledCall.request.prompt_id, () => + invocation.execute( + execSignal, + liveOutputCallback, + shellExecutionConfig, + ), ); } diff --git a/packages/core/src/tools/todoWrite.test.ts b/packages/core/src/tools/todoWrite.test.ts index 2b796f7345c..1c502eeb3e0 100644 --- a/packages/core/src/tools/todoWrite.test.ts +++ b/packages/core/src/tools/todoWrite.test.ts @@ -15,6 +15,7 @@ import type { Config } from '../config/config.js'; import type { AggregatedHookResult } from '../hooks/hookAggregator.js'; import { Storage } from '../config/storage.js'; import { atomicWriteFile } from '../utils/atomicFileWrite.js'; +import { promptIdContext } from '../utils/promptIdContext.js'; // Mock fs modules vi.mock('fs/promises'); @@ -37,7 +38,8 @@ describe('TodoWriteTool', () => { mockConfig = { getSessionId: () => 'test-session-123', getHookSystem: () => undefined, - } as Config; + setActiveTodoReminder: vi.fn(), + } as unknown as Config; tool = new TodoWriteTool(mockConfig); mockAbortSignal = new AbortController().signal; vi.clearAllMocks(); @@ -152,7 +154,9 @@ describe('TodoWriteTool', () => { mockAtomicWrite.mockResolvedValue(undefined); const invocation = tool.build(params); - const result = await invocation.execute(mockAbortSignal); + const result = await promptIdContext.run('todo-prompt', () => + invocation.execute(mockAbortSignal), + ); expect(result.llmContent).toContain( 'Todos have been modified successfully', @@ -172,6 +176,30 @@ describe('TodoWriteTool', () => { expect.stringContaining('"todos"'), { encoding: 'utf-8' }, ); + expect(mockConfig.setActiveTodoReminder).toHaveBeenCalledWith( + 'todo-prompt', + expect.stringContaining('Task 1'), + ); + }); + + it('bounds the active Todo reminder', async () => { + const params: TodoWriteParams = { + todos: [{ id: '1', content: 'x'.repeat(5000), status: 'in_progress' }], + }; + const enoentError = new Error('ENOENT') as Error & { code: string }; + enoentError.code = 'ENOENT'; + mockFs.readFile.mockRejectedValue(enoentError); + mockFs.mkdir.mockResolvedValue(undefined); + mockAtomicWrite.mockResolvedValue(undefined); + + await promptIdContext.run('todo-prompt', () => + tool.build(params).execute(mockAbortSignal), + ); + + const reminder = vi.mocked(mockConfig.setActiveTodoReminder).mock + .lastCall?.[1]; + expect(reminder).toContain('[truncated]'); + expect(reminder?.length).toBeLessThan(4300); }); it('should replace todos with new ones', async () => { @@ -194,7 +222,9 @@ describe('TodoWriteTool', () => { mockAtomicWrite.mockResolvedValue(undefined); const invocation = tool.build(params); - const result = await invocation.execute(mockAbortSignal); + const result = await promptIdContext.run('todo-prompt', () => + invocation.execute(mockAbortSignal), + ); expect(result.llmContent).toContain( 'Todos have been modified successfully', @@ -214,6 +244,10 @@ describe('TodoWriteTool', () => { expect.stringMatching(/"Updated Task"/), { encoding: 'utf-8' }, ); + const reminder = vi.mocked(mockConfig.setActiveTodoReminder).mock + .lastCall?.[1]; + expect(reminder).toContain('New Task'); + expect(reminder).not.toContain('Updated Task'); }); it('should handle file write errors', async () => { @@ -256,7 +290,9 @@ describe('TodoWriteTool', () => { ); const invocation = tool.build(params); - const result = await invocation.execute(mockAbortSignal); + const result = await promptIdContext.run('todo-prompt', () => + invocation.execute(mockAbortSignal), + ); expect(result.llmContent).toContain('Todo list has been cleared'); expect(result.llmContent).toContain(''); @@ -271,6 +307,10 @@ describe('TodoWriteTool', () => { expect.stringContaining('"todos"'), { encoding: 'utf-8' }, ); + expect(mockConfig.setActiveTodoReminder).toHaveBeenCalledWith( + 'todo-prompt', + undefined, + ); }); it('should block todo creation when validation hook returns block', async () => { @@ -794,7 +834,7 @@ describe('TodoWriteTool – runtime output directory', () => { mockConfig = { getSessionId: () => 'runtime-session', getHookSystem: () => undefined, - } as Config; + } as unknown as Config; tool = new TodoWriteTool(mockConfig); mockAbortSignal = new AbortController().signal; Storage.setRuntimeBaseDir(null); diff --git a/packages/core/src/tools/todoWrite.ts b/packages/core/src/tools/todoWrite.ts index e58501b97b1..69089a6226a 100644 --- a/packages/core/src/tools/todoWrite.ts +++ b/packages/core/src/tools/todoWrite.ts @@ -17,9 +17,12 @@ import { ToolDisplayNames, ToolNames } from './tool-names.js'; import { atomicWriteFile } from '../utils/atomicFileWrite.js'; import { createDebugLogger } from '../utils/debugLogger.js'; import { detectTodoChanges, HookPhase, type TodoItem } from '../hooks/types.js'; +import { escapeSystemReminderTags } from '../utils/xml.js'; +import { promptIdContext } from '../utils/promptIdContext.js'; export type { TodoItem } from '../hooks/types.js'; const debugLogger = createDebugLogger('TODO_WRITE'); +const MAX_ACTIVE_TODO_CONTEXT_CHARS = 4000; export interface TodoWriteParams { todos: TodoItem[]; @@ -250,6 +253,25 @@ class TodoWriteToolInvocation extends BaseToolInvocation< // 4. Write new todos AFTER all validation passes await writeTodosToFile(finalTodos, sessionId); + const unfinishedTodos = finalTodos.filter( + (todo) => todo.status !== 'completed', + ); + const promptId = promptIdContext.getStore(); + if (promptId) { + const serializedTodos = escapeSystemReminderTags( + JSON.stringify(unfinishedTodos), + ); + const todoContext = serializedTodos.slice( + 0, + MAX_ACTIVE_TODO_CONTEXT_CHARS, + ); + this.config.setActiveTodoReminder( + promptId, + unfinishedTodos.length > 0 + ? `\nThe current task still has unfinished todo items:\n${todoContext}${serializedTodos.length > todoContext.length ? '\n[truncated]' : ''}\nKeep the todo list current and continue the task. Do not treat a successful intermediate tool call as task completion.\n` + : undefined, + ); + } // 5. POST-WRITE PHASE: Execute hooks for side effects (logging, HTTP sync, etc.) // These hooks can now safely perform side effects knowing data is persisted