From 2abfa3d54e8449043bdb4d7d44e863e04ee6d36a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=98=93=E8=89=AF?= <1204183885@qq.com> Date: Thu, 30 Jul 2026 09:33:40 +0800 Subject: [PATCH] fix(core): preserve active Todo context across tool turns (#7919) * fix(core): preserve active Todo context across tool turns * test(cli): update automatic turn prompt expectation * fix(core): preserve Todo ownership across automatic turns * fix(core): preserve Todo ownership at prompt boundaries * test(todo): cover automatic reminder boundaries * fix(core): throttle active Todo reminder re-injection to bound history growth Every injected reminder copy lands permanently in chat history, so per-turn injection grew the live context linearly with tool turns. Tool-turn injection now re-issues the reminder only every third tool turn since the state was last presented; turn-start injections always fire and reset the cadence. The payload becomes a compact status/content line list capped at 800 characters. History stays append-only, so provider prefix caching is unaffected. Also: cover the new-ordinary-prompt-clears-stale-reminders invariant on the real Config, add TUI coverage for the work-chain notification batch split, cover todoWorkChainId continuation forwarding, and document the deliberate enterWith binding in the daemon tool runner. * fix(core): keep todo reminder before drained input --------- Co-authored-by: qwen-code-dev-bot Co-authored-by: Qwen-Coder --- docs/design/active-todo-context.md | 43 ++++ .../session/Session.review-lease.test.ts | 6 + .../acp-integration/session/Session.test.ts | 190 +++++++++++++++++- .../src/acp-integration/session/Session.ts | 116 +++++++++-- .../session/Session.worktree.test.ts | 6 + packages/cli/src/nonInteractiveCli.test.ts | 115 ++++++++++- packages/cli/src/nonInteractiveCli.ts | 19 +- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 73 +++++++ packages/cli/src/ui/hooks/useGeminiStream.ts | 18 +- .../core/src/agents/background-tasks.test.ts | 10 + packages/core/src/agents/background-tasks.ts | 4 + packages/core/src/agents/tasks/types.ts | 2 + packages/core/src/config/config.test.ts | 129 ++++++++++++ packages/core/src/config/config.ts | 125 ++++++++++++ packages/core/src/core/client-goal.test.ts | 4 + packages/core/src/core/client.test.ts | 152 ++++++++++++++ packages/core/src/core/client.ts | 88 +++++++- .../core/src/core/coreToolScheduler.test.ts | 24 ++- packages/core/src/core/coreToolScheduler.ts | 43 +++- .../services/backgroundShellRegistry.test.ts | 10 + .../src/services/backgroundShellRegistry.ts | 4 + packages/core/src/services/cronScheduler.ts | 12 +- .../core/src/services/monitorRegistry.test.ts | 9 + packages/core/src/services/monitorRegistry.ts | 5 + packages/core/src/tools/loop-wakeup.test.ts | 6 +- packages/core/src/tools/loop-wakeup.ts | 7 +- packages/core/src/tools/todoWrite.test.ts | 68 ++++++- packages/core/src/tools/todoWrite.ts | 24 +++ packages/core/src/utils/promptIdContext.ts | 1 + 29 files changed, 1260 insertions(+), 53 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 0000000000..d50b809d7a --- /dev/null +++ b/docs/design/active-todo-context.md @@ -0,0 +1,43 @@ +# 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 under a stable work-chain owner. Prompt IDs used by retries and related +automatic turns resolve to that owner, so concurrent notification branches do +not move or overwrite the foreground reminder. Background tasks and loop +wakeups capture the owner when they are created and carry it back with their +automatic turn; unrelated cron and notification turns use an isolated owner +that is removed when the turn ends. Inject the reminder on the first request of +a retry or related automatic turn and after function responses on later tool +turns. Clear it when all todos complete, a new ordinary work chain starts, or +the session changes. + +Every injected copy is recorded permanently in chat history, so per-turn +injection would grow the live context linearly with tool turns. Tool-turn +injection therefore re-issues the reminder only every third tool turn since +the last time the state was presented (the `todo_write` result itself counts); +turn-start injections always fire and reset that cadence. The payload is a +compact `- [status] content` line list capped at 800 characters. History stays +append-only, so provider prefix caching is unaffected. + +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. +- Terminal automatic turns release their temporary ownership state. diff --git a/packages/cli/src/acp-integration/session/Session.review-lease.test.ts b/packages/cli/src/acp-integration/session/Session.review-lease.test.ts index e580cf54b4..8edd8ca882 100644 --- a/packages/cli/src/acp-integration/session/Session.review-lease.test.ts +++ b/packages/cli/src/acp-integration/session/Session.review-lease.test.ts @@ -102,6 +102,12 @@ describe('Session review-worktree lease sweep', () => { switchModel: vi.fn(), getModel: vi.fn().mockReturnValue('qwen3'), getSessionId: vi.fn().mockReturnValue(SESSION_ID), + takeActiveTodoReminder: vi.fn().mockReturnValue(undefined), + getActiveTodoWorkChainOwner: vi.fn((promptId: string) => promptId), + setActiveTodoReminder: vi.fn(), + startActiveTodoWorkChain: vi.fn(), + startAutomaticActiveTodoWorkChain: vi.fn(), + endAutomaticActiveTodoWorkChain: vi.fn(), assertCanStartTurn: vi.fn().mockResolvedValue(undefined), getWorkingDir: vi.fn().mockReturnValue('/tmp'), getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(false), diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index aaaaf083b0..aefd1b475b 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -569,6 +569,12 @@ describe('Session', () => { switchModel: switchModelSpy, getModel: vi.fn().mockImplementation(() => currentModel), getSessionId: vi.fn().mockReturnValue('test-session-id'), + takeActiveTodoReminder: vi.fn().mockReturnValue(undefined), + setActiveTodoReminder: vi.fn(), + startActiveTodoWorkChain: vi.fn(), + startAutomaticActiveTodoWorkChain: vi.fn(), + endAutomaticActiveTodoWorkChain: vi.fn(), + getActiveTodoWorkChainOwner: vi.fn((promptId: string) => promptId), assertCanStartTurn: vi.fn().mockResolvedValue(undefined), getWorkingDir: vi.fn().mockReturnValue(process.cwd()), getProjectRoot: vi.fn().mockReturnValue('/repo'), @@ -765,6 +771,152 @@ 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('includes active Todo context on the first retry request', async () => { + const reminder = + 'unfinished todo: run tests'; + vi.mocked(mockConfig.takeActiveTodoReminder).mockReturnValue(reminder); + mockChat.sendMessageStream = vi + .fn() + .mockImplementation(async () => createEmptyStream()); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'start work' }], + }); + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'start work' }], + retry: true, + } as PromptRequest); + + const retryCall = vi + .mocked(mockChat.sendMessageStream) + .mock.calls.at(-1)?.[1] as { + message: Part[]; + }; + expect(textParts(retryCall.message)).toContain(reminder); + }); + + 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 reminder = + 'unfinished todo: wait for agent'; + vi.mocked(mockConfig.takeActiveTodoReminder).mockReturnValue(reminder); + 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; + todoWorkChainId?: string; + }, + ) => void; + + callback('Background task completed.', '', { + agentId: 'related-agent', + status: 'completed', + todoWorkChainId: 'test-session-id########1', + }); + + await vi.waitFor(() => + expect(mockConfig.startAutomaticActiveTodoWorkChain).toHaveBeenCalledWith( + expect.stringContaining('########notification'), + 'test-session-id########1', + ), + ); + await vi.waitFor(() => + expect(mockConfig.endAutomaticActiveTodoWorkChain).toHaveBeenCalledWith( + expect.stringContaining('########notification'), + ), + ); + const notificationCall = vi + .mocked(mockChat.sendMessageStream) + .mock.calls.at(-1)?.[1] as { message: Part[] }; + expect(textParts(notificationCall.message)).toContain(reminder); + }); + + it('does not infer Todo ownership from Todo Stop Guard lineage', 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('guard-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: 'guard-related-agent', + status: 'completed', + }); + + await vi.waitFor(() => + expect(mockConfig.startAutomaticActiveTodoWorkChain).toHaveBeenCalledWith( + expect.stringContaining('########notification'), + undefined, + ), + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'start work' }], + retry: true, + } as PromptRequest); + expect(mockConfig.startActiveTodoWorkChain).toHaveBeenLastCalledWith( + 'test-session-id########2', + 'test-session-id########1', + ); + }); + it('holds the close gate until active turns settle', async () => { let resolveTurn!: () => void; const turnCompletion = new Promise((resolve) => { @@ -6062,9 +6214,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.takeActiveTodoReminder).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', @@ -6116,9 +6285,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 '); @@ -10478,6 +10657,7 @@ describe('Session', () => { expect(mockGeminiClient.tryCompressChat).not.toHaveBeenCalled(); expect(mockChat.sendMessageStream).not.toHaveBeenCalled(); + expect(mockConfig.startActiveTodoWorkChain).not.toHaveBeenCalled(); expect(mockClient.sessionUpdate).toHaveBeenCalledWith({ sessionId: 'test-session-id', update: { diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 585e2483ef..71194357f1 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -144,6 +144,7 @@ import { setGoalTerminalObserver, sessionIdContext, promptIdContext, + todoWorkChainContext, dedupeToolCallsById, getProviderToolCallId, parsePositiveIntegerEnv, @@ -820,6 +821,7 @@ interface BackgroundNotificationQueueItem { kind: 'agent' | 'monitor' | 'shell'; continuesTodoStopGuardWorkChain: boolean; toolUseId?: string; + todoWorkChainId?: string; } /** The slice of `CronJob` a fire delivers to this session. Structural, not the @@ -834,6 +836,7 @@ interface CronFire { * identifies this fire's entry in `runs[]`. */ lastFiredAt?: number; delivery?: CronTaskDelivery; + todoWorkChainId?: string; } interface CronQueueItem { @@ -842,6 +845,7 @@ interface CronQueueItem { taskId?: string; firedAt?: number; delivery?: CronTaskDelivery; + todoWorkChainId?: string; } interface PromptChannelDelivery { @@ -1237,6 +1241,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 @@ -1346,14 +1351,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); @@ -2635,6 +2634,12 @@ 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; // 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 @@ -2854,6 +2859,17 @@ export class Session implements SessionContext { } } + if (!continuesCurrentWorkChain && !this.todoStopGuard.enabled) { + this.#resetTodoStopGuardBackgroundLineage(); + } + this.config.startActiveTodoWorkChain( + promptId, + continuesCurrentWorkChain + ? this.activeTodoWorkChainPromptId + : undefined, + ); + this.activeTodoWorkChainPromptId = promptId; + // Snapshot file state before this turn (mirrors the makeSnapshot // block in GeminiClient.sendMessageStream). Placed after // slash-command and hook early-returns so locally handled commands @@ -2932,6 +2948,19 @@ export class Session implements SessionContext { this.pendingRecoveredAgentsNotice = null; } + const activeTodoReminder = this.config.takeActiveTodoReminder( + promptId, + true, + ); + if ( + activeTodoReminder && + !parts.some((part) => part.text === activeTodoReminder) + ) { + parts = insertAfterFunctionResponses(parts, [ + { text: activeTodoReminder }, + ]); + } + let nextMessage: Content | null = { role: 'user', parts }; let turnCount = 0; const toolLoopState = createDaemonToolLoopState(); @@ -3192,6 +3221,7 @@ export class Session implements SessionContext { await this.#buildNextMessageAfterToolRun( toolRun, pendingSend.signal, + promptId, onFullTurnModel, ); nextMessage = nextAfterTools.message; @@ -4136,6 +4166,7 @@ export class Session implements SessionContext { const nextAfterTools = await this.#buildNextMessageAfterToolRun( toolRun, pendingSend.signal, + toolPromptId, options.onFullTurnModel, ); nextMessage = nextAfterTools.message; @@ -4539,6 +4570,7 @@ export class Session implements SessionContext { async #buildNextMessageAfterToolRun( toolRun: RunToolResult, abortSignal: AbortSignal, + promptId: string, onFullTurnModel?: (model: string) => boolean, ): Promise { if (toolRun.loopDetected) { @@ -4559,7 +4591,12 @@ export class Session implements SessionContext { if (hadMidTurnUserInput) { this.todoStopGuard.acceptMidTurnUserInput(); } - const parts = [...toolRun.parts, ...drained.parts]; + const activeTodoReminder = this.config.takeActiveTodoReminder(promptId); + const parts = [ + ...toolRun.parts, + ...(activeTodoReminder ? [{ text: activeTodoReminder }] : []), + ...drained.parts, + ]; return { message: { role: 'user', parts }, hadMidTurnUserInput, @@ -4948,6 +4985,9 @@ export class Session implements SessionContext { ...(job.id ? { taskId: job.id } : {}), ...(job.lastFiredAt !== undefined ? { firedAt: job.lastFiredAt } : {}), ...(job.delivery ? { delivery: job.delivery } : {}), + ...(job.todoWorkChainId + ? { todoWorkChainId: job.todoWorkChainId } + : {}), }); void this.#drainCronQueue(); }); @@ -5140,9 +5180,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; @@ -5162,6 +5202,10 @@ export class Session implements SessionContext { try { await this.assertCanStartTurn(); if (ac.signal.aborted) return; + this.config.startAutomaticActiveTodoWorkChain( + promptId, + item.todoWorkChainId, + ); // 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 @@ -5294,9 +5338,17 @@ export class Session implements SessionContext { // Prepend session-level system reminders (same rationale as the // user-query path in #executePrompt). const cronReminders = await this.#buildInitialSystemReminders(); + const activeTodoReminder = this.config.takeActiveTodoReminder( + promptId, + true, + ); let nextMessage: Content | null = { role: 'user', - parts: [...cronReminders, { text: modelText }], + parts: [ + ...cronReminders, + ...(activeTodoReminder ? [{ text: activeTodoReminder }] : []), + { text: modelText }, + ], }; const toolLoopState = createDaemonToolLoopState(); @@ -5464,6 +5516,7 @@ export class Session implements SessionContext { await this.#buildNextMessageAfterToolRun( toolRun, ac.signal, + promptId, ); nextMessage = nextAfterTools.message; if (toolRun.loopDetected) { @@ -5504,6 +5557,7 @@ export class Session implements SessionContext { `[${item.source} error] ${msg}`, ); } finally { + this.config.endAutomaticActiveTodoWorkChain(promptId); if (this.cronAbortController === ac) { this.cronAbortController = null; } @@ -5572,6 +5626,7 @@ export class Session implements SessionContext { continuesTodoStopGuardWorkChain: this.#agentContinuesTodoStopGuardWorkChain(meta.agentId), toolUseId: meta.toolUseId, + todoWorkChainId: meta.todoWorkChainId, }); }, ); @@ -5594,6 +5649,7 @@ export class Session implements SessionContext { meta.ownerAgentId, ), toolUseId: meta.toolUseId, + todoWorkChainId: meta.todoWorkChainId, }); }); @@ -5607,6 +5663,7 @@ export class Session implements SessionContext { kind: 'shell', continuesTodoStopGuardWorkChain: !this.todoStopGuardBackgroundBaseline.shells.has(meta.shellId), + todoWorkChainId: meta.todoWorkChainId, }); }); @@ -5782,14 +5839,18 @@ 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, + item.todoWorkChainId, + ); await this.#emitBackgroundNotificationDisplay(item); const notificationParts: Part[] = [{ text: item.modelText }]; @@ -5804,9 +5865,17 @@ export class Session implements SessionContext { const notificationReminders = await this.#buildInitialSystemReminders(); + const activeTodoReminder = this.config.takeActiveTodoReminder( + promptId, + true, + ); let nextMessage: Content | null = { role: 'user', - parts: [...notificationReminders, ...notificationParts], + parts: [ + ...notificationReminders, + ...(activeTodoReminder ? [{ text: activeTodoReminder }] : []), + ...notificationParts, + ], }; const toolLoopState = createDaemonToolLoopState(); @@ -5962,6 +6031,7 @@ export class Session implements SessionContext { const nextAfterTools = await this.#buildNextMessageAfterToolRun( toolRun, ac.signal, + promptId, ); nextMessage = nextAfterTools.message; if (toolRun.loopDetected) { @@ -6016,6 +6086,7 @@ export class Session implements SessionContext { await this.#emitBackgroundNotificationEndTurn('end_turn'); } } finally { + this.config.endAutomaticActiveTodoWorkChain(promptId); if (this.notificationAbortController === ac) { this.notificationAbortController = null; } @@ -6386,6 +6457,17 @@ export class Session implements SessionContext { toolLoopState?: DaemonToolLoopState, onFullTurnModel?: (model: string) => boolean, ): Promise { + // The daemon executes tools directly rather than through + // CoreToolScheduler, so the ALS bindings the scheduler would provide must + // happen here. `enterWith` (not `run`) is deliberate: background + // task/shell/monitor registration can occur in async continuations of + // this batch after runToolCalls resolves, and those must still observe + // this prompt's work-chain owner. The turn loop rebinds on the next + // runToolCalls, and turn starts re-enter via #executePrompt. + promptIdContext.enterWith(promptId); + todoWorkChainContext.enterWith( + this.config.getActiveTodoWorkChainOwner(promptId), + ); const dedupedFunctionCalls = dedupeToolCallsById(functionCalls); const generatedCallIdBase = randomUUID(); const executionCallIds = new Map( diff --git a/packages/cli/src/acp-integration/session/Session.worktree.test.ts b/packages/cli/src/acp-integration/session/Session.worktree.test.ts index c23bae2833..92a0ef3a3d 100644 --- a/packages/cli/src/acp-integration/session/Session.worktree.test.ts +++ b/packages/cli/src/acp-integration/session/Session.worktree.test.ts @@ -108,6 +108,12 @@ describe('Session.pendingWorktreeNotice', () => { switchModel: vi.fn(), getModel: vi.fn().mockReturnValue('qwen3'), getSessionId: vi.fn().mockReturnValue(SESSION_ID), + takeActiveTodoReminder: vi.fn().mockReturnValue(undefined), + getActiveTodoWorkChainOwner: vi.fn((promptId: string) => promptId), + setActiveTodoReminder: vi.fn(), + startActiveTodoWorkChain: vi.fn(), + startAutomaticActiveTodoWorkChain: vi.fn(), + endAutomaticActiveTodoWorkChain: vi.fn(), assertCanStartTurn: vi.fn().mockResolvedValue(undefined), getWorkingDir: vi.fn().mockReturnValue('/tmp'), getTelemetryLogPromptsEnabled: vi.fn().mockReturnValue(false), diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index b24401a855..d0f9d943fb 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -1547,7 +1547,7 @@ describe('runNonInteractive', () => { 4, [{ text: 'drain image a' }, { text: 'drain image b' }], expect.any(AbortSignal), - 'prompt-drain', + 'prompt-drain/automatic/3', { type: SendMessageType.ToolResult, modelOverride: first }, ); }); @@ -2275,6 +2275,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[]; @@ -3468,7 +3474,7 @@ describe('runNonInteractive', () => { 2, [{ text: notificationXml }], expect.any(AbortSignal), - 'prompt-monitor', + 'prompt-monitor/automatic/2', { type: SendMessageType.Notification, modelOverride: undefined, @@ -3506,6 +3512,111 @@ describe('runNonInteractive', () => { }); }); + it('keeps notifications from different Todo work chains in separate batches', async () => { + setupMetricsMock(); + + const firstNotificationXml = + '\n' + + 'mon_1\n' + + 'monitor\n' + + 'running\n' + + 'Monitor emitted event #1.\n' + + 'ready\n' + + ''; + const secondNotificationXml = + '\n' + + 'mon_2\n' + + 'monitor\n' + + 'running\n' + + 'Monitor emitted event #2.\n' + + 'also ready\n' + + ''; + + mockMonitorRegistry.setNotificationCallback.mockImplementation((cb) => { + if (!cb) { + return; + } + cb('Monitor "logs" event #1: ready', firstNotificationXml, { + monitorId: 'mon_1', + toolUseId: 'tool_mon_1', + status: 'running', + eventCount: 1, + todoWorkChainId: 'chain-1', + }); + cb('Monitor "build" event #1: ready', secondNotificationXml, { + monitorId: 'mon_2', + toolUseId: 'tool_mon_2', + status: 'running', + eventCount: 1, + todoWorkChainId: 'chain-2', + }); + }); + mockGeminiClient.sendMessageStream + .mockReturnValueOnce( + createStreamFromEvents([ + { type: GeminiEventType.Content, value: 'Started.' }, + { + type: GeminiEventType.Finished, + value: { + reason: undefined, + usageMetadata: { totalTokenCount: 1 }, + }, + }, + ]), + ) + .mockReturnValueOnce( + createStreamFromEvents([ + { type: GeminiEventType.Content, value: 'First notification.' }, + { + type: GeminiEventType.Finished, + value: { + reason: undefined, + usageMetadata: { totalTokenCount: 1 }, + }, + }, + ]), + ) + .mockReturnValueOnce( + createStreamFromEvents([ + { type: GeminiEventType.Content, value: 'Second notification.' }, + { + type: GeminiEventType.Finished, + value: { + reason: undefined, + usageMetadata: { totalTokenCount: 1 }, + }, + }, + ]), + ); + + await runNonInteractive( + mockConfig, + mockSettings, + 'Watch the logs', + 'prompt-monitor-work-chains', + ); + + expect(mockGeminiClient.sendMessageStream).toHaveBeenCalledTimes(3); + expect(mockGeminiClient.sendMessageStream).toHaveBeenNthCalledWith( + 2, + [{ text: firstNotificationXml }], + expect.any(AbortSignal), + 'prompt-monitor-work-chains/automatic/2', + expect.objectContaining({ + todoWorkChainId: 'chain-1', + }), + ); + expect(mockGeminiClient.sendMessageStream).toHaveBeenNthCalledWith( + 3, + [{ text: secondNotificationXml }], + expect.any(AbortSignal), + 'prompt-monitor-work-chains/automatic/3', + expect.objectContaining({ + todoWorkChainId: 'chain-2', + }), + ); + }); + it.skip('should emit a single user envelope when userEnvelope is provided', async () => { (mockConfig.getOutputFormat as Mock).mockReturnValue('stream-json'); (mockConfig.getIncludePartialMessages as Mock).mockReturnValue(false); diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 85afcfe85f..5231321e2d 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -456,6 +456,7 @@ export async function runNonInteractive( displayText: string; modelText: string; sendMessageType: SendMessageType; + todoWorkChainId?: string; monitorId?: string; sdkNotification?: { task_id: string; @@ -901,6 +902,7 @@ export async function runNonInteractive( displayText, modelText, sendMessageType: SendMessageType.Notification, + todoWorkChainId: meta.todoWorkChainId, sdkNotification: { task_id: meta.agentId, tool_use_id: meta.toolUseId, @@ -946,6 +948,7 @@ export async function runNonInteractive( displayText, modelText, sendMessageType: SendMessageType.Notification, + todoWorkChainId: meta.todoWorkChainId, monitorId: meta.monitorId, sdkNotification: { task_id: meta.monitorId, @@ -1629,6 +1632,7 @@ export async function runNonInteractive( }; }; + let currentPromptId = prompt_id; while (true) { // Drain pending teammate messages into the conversation. // sendMessageStream only reads currentMessages[0].parts, @@ -1682,13 +1686,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, @@ -1898,7 +1905,9 @@ export async function runNonInteractive( if (splitIdx === 0) { while ( splitIdx < localQueue.length && - localQueue[splitIdx]!.sendMessageType === targetType + localQueue[splitIdx]!.sendMessageType === targetType && + localQueue[splitIdx]!.todoWorkChainId === + localQueue[0]!.todoWorkChainId ) { splitIdx++; } @@ -1917,6 +1926,7 @@ export async function runNonInteractive( displayText: batch.map((i) => i.displayText).join('; '), modelText: batch.map((i) => i.modelText).join('\n\n'), sendMessageType: targetType, + todoWorkChainId: batch[0]?.todoWorkChainId, }; turnCount++; @@ -1933,6 +1943,7 @@ export async function runNonInteractive( ]; let itemIsFirstTurn = true; let itemModelOverride: string | undefined; + const itemPromptId = `${prompt_id}/automatic/${turnCount}`; while (true) { const itemToolCallRequests: ToolCallRequestInfo[] = []; @@ -1940,7 +1951,7 @@ export async function runNonInteractive( const itemStream = geminiClient.sendMessageStream( itemMessages[0]?.parts || [], abortController.signal, - prompt_id, + itemPromptId, { type: itemIsFirstTurn ? item.sendMessageType @@ -1948,6 +1959,7 @@ export async function runNonInteractive( modelOverride: itemModelOverride, ...(itemIsFirstTurn && { notificationDisplayText: item.displayText, + todoWorkChainId: item.todoWorkChainId, }), }, ); @@ -2183,6 +2195,7 @@ export async function runNonInteractive( displayText: `${job.cronExpr === '@wakeup' ? 'Loop' : 'Cron'}: ${label}`, modelText: job.prompt, sendMessageType: SendMessageType.Cron, + todoWorkChainId: job.todoWorkChainId, }); drainLocalQueue().then(checkCronDone, onDrainError); }); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 70cc35a1f8..aa89113af4 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -7159,6 +7159,79 @@ describe('useGeminiStream', () => { ); }); + it('keeps notifications from different Todo work chains in separate turns', async () => { + renderTestHook(); + + const callback = mockMonitorRegistry.setNotificationCallback.mock + .calls[0][0] as ( + displayText: string, + modelText: string, + meta: { + monitorId: string; + status: string; + todoWorkChainId?: string; + }, + ) => void; + mockSendMessageStream.mockClear(); + + await act(async () => { + callback( + 'Monitor "logs" event #1: ready', + 'first', + { + monitorId: 'mon_1', + status: 'completed', + todoWorkChainId: 'chain-1', + }, + ); + callback( + 'Monitor "build" event #1: ready', + 'second', + { + monitorId: 'mon_2', + status: 'completed', + todoWorkChainId: 'chain-2', + }, + ); + }); + + await waitFor(() => + expect(mockSendMessageStream).toHaveBeenCalledTimes(1), + ); + const firstCall = mockSendMessageStream.mock.calls[0]; + expect(JSON.stringify(firstCall[0])).toContain('first'); + expect(JSON.stringify(firstCall[0])).not.toContain('second'); + expect(firstCall[3]).toMatchObject({ + type: SendMessageType.Notification, + todoWorkChainId: 'chain-1', + }); + + // A further chain-2 event re-triggers the drain; the two queued + // chain-2 items batch into a single turn. + await act(async () => { + callback( + 'Monitor "build" event #2: done', + 'third', + { + monitorId: 'mon_2', + status: 'completed', + todoWorkChainId: 'chain-2', + }, + ); + }); + + await waitFor(() => + expect(mockSendMessageStream).toHaveBeenCalledTimes(2), + ); + const secondCall = mockSendMessageStream.mock.calls[1]; + expect(JSON.stringify(secondCall[0])).toContain('second'); + expect(JSON.stringify(secondCall[0])).toContain('third'); + expect(secondCall[3]).toMatchObject({ + type: SendMessageType.Notification, + todoWorkChainId: 'chain-2', + }); + }); + // Regression for #7156: progress setState calls issued from inside a // background subagent's AsyncLocalStorage frame can batch with the // notification trigger into one React commit, so the drain effect diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index d2948685dd..df8b380200 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -2662,6 +2662,7 @@ export const useGeminiStream = ( prompt_id?: string, metadata?: { notificationDisplayText?: string; + todoWorkChainId?: string; onDelivered?: () => void; onDeliveryFailed?: () => void; steerInput?: SteerInput; @@ -2912,6 +2913,7 @@ export const useGeminiStream = ( const sendOptions = { type: submitType, notificationDisplayText: metadata?.notificationDisplayText, + todoWorkChainId: metadata?.todoWorkChainId, modelOverride: modelOverrideRef.current, steerInput: metadata?.steerInput, ...(submittedPrompt !== undefined ? { submittedPrompt } : {}), @@ -3823,6 +3825,7 @@ export const useGeminiStream = ( modelText: string; sendMessageType: SendMessageType; monitor?: { id: string; status: string }; + todoWorkChainId?: string; onDelivered?: () => void; onDeliveryFailed?: () => void; }> @@ -3898,6 +3901,7 @@ export const useGeminiStream = ( prompt: string; cronExpr?: string; missed?: boolean; + todoWorkChainId?: string; }) => { const source = job.cronExpr === '@wakeup' ? 'Loop' : 'Cron'; const autonomousMode = detectAutonomousSentinel(job.prompt); @@ -3913,6 +3917,7 @@ export const useGeminiStream = ( displayText: `${job.missed ? 'Missed' : source}: ${label}`, modelText, sendMessageType: SendMessageType.Cron, + todoWorkChainId: job.todoWorkChainId, onDelivered: () => resolver.markDelivered(), }); setNotificationTrigger((n) => n + 1); @@ -3922,6 +3927,7 @@ export const useGeminiStream = ( displayText: `${job.missed ? 'Missed' : source}: ${label}`, modelText, sendMessageType: SendMessageType.Cron, + todoWorkChainId: job.todoWorkChainId, }); setNotificationTrigger((n) => n + 1); }, @@ -3941,11 +3947,12 @@ export const useGeminiStream = ( // Register background agent notification callback onto the shared queue. useEffect(() => { const registry = config.getBackgroundTaskRegistry(); - registry.setNotificationCallback((displayText, modelText) => { + registry.setNotificationCallback((displayText, modelText, meta) => { notificationQueueRef.current.push({ displayText, modelText, sendMessageType: SendMessageType.Notification, + todoWorkChainId: meta?.todoWorkChainId, }); setNotificationTrigger((n) => n + 1); }); @@ -3957,11 +3964,12 @@ export const useGeminiStream = ( // Register background shell terminal notification callback onto the shared queue. useEffect(() => { const registry = config.getBackgroundShellRegistry(); - registry.setNotificationCallback((displayText, modelText) => { + registry.setNotificationCallback((displayText, modelText, meta) => { notificationQueueRef.current.push({ displayText, modelText, sendMessageType: SendMessageType.Notification, + todoWorkChainId: meta?.todoWorkChainId, }); setNotificationTrigger((n) => n + 1); }); @@ -3983,6 +3991,7 @@ export const useGeminiStream = ( modelText, sendMessageType: SendMessageType.Notification, monitor: { id: meta.monitorId, status: meta.status }, + todoWorkChainId: meta.todoWorkChainId, }); setNotificationTrigger((n) => n + 1); }); @@ -4040,6 +4049,7 @@ export const useGeminiStream = ( ); submitQuery(item.modelText, item.sendMessageType, undefined, { notificationDisplayText: item.displayText, + todoWorkChainId: item.todoWorkChainId, onDelivered: item.onDelivered, onDeliveryFailed: item.onDeliveryFailed, }); @@ -4050,7 +4060,8 @@ export const useGeminiStream = ( let splitIdx = 0; while ( splitIdx < queue.length && - queue[splitIdx]!.sendMessageType === targetType + queue[splitIdx]!.sendMessageType === targetType && + queue[splitIdx]!.todoWorkChainId === queue[0]!.todoWorkChainId ) { splitIdx++; } @@ -4068,6 +4079,7 @@ export const useGeminiStream = ( const combinedDisplayText = batch.map((e) => e.displayText).join('; '); submitQuery(combinedModelText, targetType, undefined, { notificationDisplayText: combinedDisplayText, + todoWorkChainId: batch[0]?.todoWorkChainId, }); }); } diff --git a/packages/core/src/agents/background-tasks.test.ts b/packages/core/src/agents/background-tasks.test.ts index 60f19680f1..29a4a3085c 100644 --- a/packages/core/src/agents/background-tasks.test.ts +++ b/packages/core/src/agents/background-tasks.test.ts @@ -26,6 +26,7 @@ import { import * as transcript from './agent-transcript.js'; import { AgentEventEmitter, AgentEventType } from './runtime/agent-events.js'; import { ToolConfirmationOutcome } from '../tools/tools.js'; +import { todoWorkChainContext } from '../utils/promptIdContext.js'; function makeApproval( callId: string, @@ -60,6 +61,15 @@ function makeRegistration( } describe('notification emission and agent context (#7156)', () => { + it('captures the Todo work-chain owner at registration', () => { + const registry = new BackgroundTaskRegistry(); + const entry = todoWorkChainContext.run('work-chain-1', () => + registry.register(makeRegistration('bg-owner')), + ); + + expect(entry.todoWorkChainId).toBe('work-chain-1'); + }); + // A background agent's terminal transition fires inside its own // AsyncLocalStorage frame, and ALS context follows every async // continuation the notification callback starts (React state updates, diff --git a/packages/core/src/agents/background-tasks.ts b/packages/core/src/agents/background-tasks.ts index c6e8c41add..ccda093b66 100644 --- a/packages/core/src/agents/background-tasks.ts +++ b/packages/core/src/agents/background-tasks.ts @@ -24,6 +24,7 @@ import { ToolConfirmationOutcome } from '../tools/tools.js'; import { createDebugLogger } from '../utils/debugLogger.js'; import { parsePositiveIntegerEnv } from '../utils/env.js'; +import { todoWorkChainContext } from '../utils/promptIdContext.js'; import { escapeXml } from '../utils/xml.js'; import { patchAgentMeta } from './agent-transcript.js'; import { runOutsideAgentContext } from './runtime/agent-context.js'; @@ -385,6 +386,7 @@ export interface NotificationMeta { status: TaskStatus; stats?: AgentCompletionStats; toolUseId?: string; + todoWorkChainId?: string; } export type BackgroundNotificationCallback = ( @@ -676,6 +678,7 @@ export class BackgroundTaskRegistry { entry.notified = options.preserveNotificationState ? ((registration as AgentTask).notified ?? false) : false; + entry.todoWorkChainId ??= todoWorkChainContext.getStore(); entry.pendingMessages = registration.pendingMessages ?? []; // Resolve the parent's display name at registration time — before the // parent can evict — so the UI's orphan annotation survives it. Owned @@ -1601,6 +1604,7 @@ export class BackgroundTaskRegistry { status: entry.status, stats: entry.stats, toolUseId: entry.toolUseId, + todoWorkChainId: entry.todoWorkChainId, }; try { diff --git a/packages/core/src/agents/tasks/types.ts b/packages/core/src/agents/tasks/types.ts index ea3a2a9798..554a921923 100644 --- a/packages/core/src/agents/tasks/types.ts +++ b/packages/core/src/agents/tasks/types.ts @@ -98,6 +98,8 @@ export interface TaskBase { outputOffset: number; /** True once the kind's terminal notification has fired. */ notified: boolean; + /** Todo work chain that created this task, when it was model-launched. */ + todoWorkChainId?: string; /** Unified cancellation handle. */ abortController: AbortController; } diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 5c049d26a7..6a35d1a86c 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -8959,4 +8959,133 @@ 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')).toBe( + 'unfinished user work', + ); + expect(config.getActiveTodoReminder('prompt-cron')).toBeUndefined(); + }); + + it('clears stale Todo reminders when a new ordinary work chain starts', () => { + const config = Object.create(Config.prototype) as Config; + config.startActiveTodoWorkChain('prompt-old'); + config.setActiveTodoReminder('prompt-old', 'old work'); + + config.startActiveTodoWorkChain('prompt-new'); + + expect(config.getActiveTodoReminder('prompt-new')).toBeUndefined(); + expect(config.getActiveTodoReminder('prompt-old')).toBeUndefined(); + }); + + it('re-issues the active Todo reminder only every third tool turn', () => { + const config = Object.create(Config.prototype) as Config; + config.startActiveTodoWorkChain('prompt-user'); + config.setActiveTodoReminder('prompt-user', 'unfinished work'); + + expect(config.takeActiveTodoReminder('prompt-user')).toBeUndefined(); + expect(config.takeActiveTodoReminder('prompt-user')).toBeUndefined(); + expect(config.takeActiveTodoReminder('prompt-user')).toBe( + 'unfinished work', + ); + expect(config.takeActiveTodoReminder('prompt-user')).toBeUndefined(); + + expect(config.takeActiveTodoReminder('prompt-user', true)).toBe( + 'unfinished work', + ); + expect(config.takeActiveTodoReminder('prompt-user')).toBeUndefined(); + expect(config.takeActiveTodoReminder('prompt-user')).toBeUndefined(); + expect(config.takeActiveTodoReminder('prompt-user')).toBe( + 'unfinished work', + ); + + config.setActiveTodoReminder('prompt-user', 'updated work'); + expect(config.takeActiveTodoReminder('prompt-user')).toBeUndefined(); + }); + + it('moves related automatic work without clearing unrelated reminders', () => { + const config = Object.create(Config.prototype) as Config; + config.startActiveTodoWorkChain('prompt-user'); + config.setActiveTodoReminder('prompt-user', 'unfinished user work'); + config.startAutomaticActiveTodoWorkChain('prompt-unrelated'); + config.setActiveTodoReminder('prompt-unrelated', 'other work'); + + config.startAutomaticActiveTodoWorkChain('prompt-cron'); + config.startAutomaticActiveTodoWorkChain( + 'prompt-related-notification', + 'prompt-user', + ); + + expect(config.getActiveTodoReminder('prompt-user')).toBe( + 'unfinished user work', + ); + expect(config.getActiveTodoReminder('prompt-cron')).toBeUndefined(); + expect(config.getActiveTodoReminder('prompt-related-notification')).toBe( + 'unfinished user work', + ); + expect( + config.getActiveTodoWorkChainOwner( + 'prompt-related-notification', + 'stale-owner', + ), + ).toBe('prompt-user'); + expect( + config.getActiveTodoWorkChainOwner('prompt-unmapped', 'inherited-owner'), + ).toBe('inherited-owner'); + expect(config.getActiveTodoReminder('prompt-unrelated')).toBe('other work'); + + config.endAutomaticActiveTodoWorkChain('prompt-cron'); + config.endAutomaticActiveTodoWorkChain('prompt-related-notification'); + + expect(config.getActiveTodoReminder('prompt-cron')).toBeUndefined(); + expect(config.getActiveTodoReminder('prompt-user')).toBe( + 'unfinished user work', + ); + + config.startAutomaticActiveTodoWorkChain( + 'prompt-stale-notification', + 'prompt-stale-owner', + ); + config.setActiveTodoReminder( + 'prompt-stale-notification', + 'stale automatic work', + ); + config.endAutomaticActiveTodoWorkChain('prompt-stale-notification'); + + expect(config.getActiveTodoReminder('prompt-stale-owner')).toBeUndefined(); + }); + + 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.startActiveTodoWorkChain('old-retry', 'old-prompt'); + + config.startNewSession('new-session-id'); + + expect(config.getActiveTodoReminder('old-prompt')).toBeUndefined(); + expect(config.getActiveTodoWorkChainOwner('old-retry')).toBe('old-retry'); + }); }); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 68c013ca4b..031859c02f 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -251,6 +251,9 @@ const memoryPressureConfigLogger = createDebugLogger('MEMORY_PRESSURE'); const MEMORY_CONTEXT_WARNING_RATIO = 0.15; +/** Re-inject the active Todo reminder every Nth tool turn, not every turn. */ +const ACTIVE_TODO_REMINDER_REFRESH_TURNS = 3; + // Default `tools.toolSearch.threshold` (percent of the context window): // mirrors the settings-schema default in packages/cli. const DEFAULT_TOOL_SEARCH_THRESHOLD = 10; @@ -1857,6 +1860,9 @@ export class Config { private readonly gitCoAuthor: GitCoAuthorSettings; private readonly usageStatisticsEnabled: boolean; private readonly fileReadCacheDisabled: boolean; + private activeTodoReminders = new Map(); + private activeTodoWorkChainOwners = new Map(); + private activeTodoReminderTurns = new Map(); private geminiClient!: GeminiClient; private baseLlmClient!: BaseLlmClient; private cronScheduler: CronScheduler | null = null; @@ -3668,6 +3674,9 @@ export class Config { } this.sessionData = sessionData; this.pendingRecoveredAgentsNotice = null; + this.getOwnActiveTodoReminders().clear(); + this.getOwnActiveTodoWorkChainOwners().clear(); + this.getOwnActiveTodoReminderTurns().clear(); setDebugLogSession(this); this.debugLogger = createDebugLogger(); this.chatRecordingService = this.chatRecordingEnabled @@ -6072,6 +6081,122 @@ export class Config { return this.geminiClient; } + private getOwnActiveTodoReminders(): Map { + if (!Object.prototype.hasOwnProperty.call(this, 'activeTodoReminders')) { + this.activeTodoReminders = new Map(); + } + return this.activeTodoReminders; + } + + private getOwnActiveTodoWorkChainOwners(): Map { + if ( + !Object.prototype.hasOwnProperty.call(this, 'activeTodoWorkChainOwners') + ) { + this.activeTodoWorkChainOwners = new Map(); + } + return this.activeTodoWorkChainOwners; + } + + private getOwnActiveTodoReminderTurns(): Map { + if ( + !Object.prototype.hasOwnProperty.call(this, 'activeTodoReminderTurns') + ) { + this.activeTodoReminderTurns = new Map(); + } + return this.activeTodoReminderTurns; + } + + getActiveTodoWorkChainOwner( + promptId: string, + fallbackOwner = promptId, + ): string { + return ( + this.getOwnActiveTodoWorkChainOwners().get(promptId) ?? fallbackOwner + ); + } + + getActiveTodoReminder(promptId: string): string | undefined { + return this.getOwnActiveTodoReminders().get( + this.getActiveTodoWorkChainOwner(promptId), + ); + } + + /** + * Reads the reminder for injection, re-issuing it only every + * ACTIVE_TODO_REMINDER_REFRESH_TURNS tool turns: each injected copy lands in + * chat history permanently, so per-turn injection would grow the context + * linearly with tool turns. `force` is for turn-start injections (retry / + * related automatic turns), which always need the context and reset the + * cadence. + */ + takeActiveTodoReminder(promptId: string, force = false): string | undefined { + const owner = this.getActiveTodoWorkChainOwner(promptId); + const reminder = this.getOwnActiveTodoReminders().get(owner); + if (!reminder) return undefined; + const turns = this.getOwnActiveTodoReminderTurns(); + const elapsed = (turns.get(owner) ?? 0) + 1; + if (!force && elapsed < ACTIVE_TODO_REMINDER_REFRESH_TURNS) { + turns.set(owner, elapsed); + return undefined; + } + turns.set(owner, 0); + return reminder; + } + + setActiveTodoReminder(promptId: string, reminder: string | undefined): void { + const reminders = this.getOwnActiveTodoReminders(); + const owner = this.getActiveTodoWorkChainOwner(promptId); + if (reminder) { + reminders.set(owner, reminder); + // The todo_write result itself just presented the full state. + this.getOwnActiveTodoReminderTurns().set(owner, 0); + } else { + reminders.delete(owner); + this.getOwnActiveTodoReminderTurns().delete(owner); + } + } + + startActiveTodoWorkChain(promptId: string, continuedFrom?: string): void { + const reminders = this.getOwnActiveTodoReminders(); + const owners = this.getOwnActiveTodoWorkChainOwners(); + if (!continuedFrom) { + reminders.clear(); + owners.clear(); + this.getOwnActiveTodoReminderTurns().clear(); + owners.set(promptId, promptId); + return; + } + + const owner = this.getActiveTodoWorkChainOwner(continuedFrom); + for (const reminderOwner of reminders.keys()) { + if (reminderOwner !== owner) reminders.delete(reminderOwner); + } + owners.clear(); + owners.set(promptId, owner); + } + + startAutomaticActiveTodoWorkChain( + promptId: string, + continuedFrom?: string, + ): void { + const reminders = this.getOwnActiveTodoReminders(); + const owners = this.getOwnActiveTodoWorkChainOwners(); + const owner = continuedFrom + ? this.getActiveTodoWorkChainOwner(continuedFrom) + : promptId; + owners.set(promptId, owner); + if (owner === promptId) reminders.delete(owner); + } + + endAutomaticActiveTodoWorkChain(promptId: string): void { + const owners = this.getOwnActiveTodoWorkChainOwners(); + const owner = this.getActiveTodoWorkChainOwner(promptId); + owners.delete(promptId); + if (![...owners.values()].includes(owner)) { + this.getOwnActiveTodoReminders().delete(owner); + } + } + /** * 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-goal.test.ts b/packages/core/src/core/client-goal.test.ts index 8e57bdcf19..1570b6a19c 100644 --- a/packages/core/src/core/client-goal.test.ts +++ b/packages/core/src/core/client-goal.test.ts @@ -207,6 +207,10 @@ function setupGoalClient() { getModel: vi.fn(() => 'test-model'), getSkipNextSpeakerCheck: vi.fn(() => false), getSkipLoopDetection: vi.fn(() => false), + startActiveTodoWorkChain: vi.fn(), + startAutomaticActiveTodoWorkChain: vi.fn(), + endAutomaticActiveTodoWorkChain: vi.fn(), + takeActiveTodoReminder: vi.fn(() => undefined), getContentGeneratorConfig: vi.fn(() => undefined), hasHooksForEvent: vi.fn(() => false), getStopHookBlockingCap: vi.fn(() => 8), diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 51c8214913..e219c8d2cb 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -564,6 +564,11 @@ describe('Gemini Client (client.ts)', () => { setStaticSystemPrefix: vi.fn(), getFullContext: vi.fn().mockReturnValue(false), getSessionId: vi.fn().mockReturnValue('test-session-id'), + takeActiveTodoReminder: vi.fn().mockReturnValue(undefined), + getActiveTodoWorkChainOwner: vi.fn((promptId: string) => promptId), + startActiveTodoWorkChain: vi.fn(), + startAutomaticActiveTodoWorkChain: vi.fn(), + endAutomaticActiveTodoWorkChain: vi.fn(), getProxy: vi.fn().mockReturnValue(undefined), getWorkingDir: vi.fn().mockReturnValue('/test/dir'), getFileService: vi.fn().mockReturnValue(fileService), @@ -1808,6 +1813,153 @@ 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.takeActiveTodoReminder).mockReturnValue(reminder); + + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'response' }; + })(), + ); + const stream = client.sendMessageStream( + [ + { functionResponse: { name: 'read_file', response: { ok: true } } }, + 'user changed priority mid-turn', + ], + 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(request.indexOf(reminder)).toBeLessThan( + request.indexOf('user changed priority mid-turn'), + ); + expect(mockConfig.takeActiveTodoReminder).toHaveBeenCalledWith( + 'prompt-tool-result', + ); + + await runTurn(SendMessageType.UserQuery); + + expect(mockConfig.startActiveTodoWorkChain).toHaveBeenCalledWith( + 'prompt-userQuery', + ); + + await runTurn(SendMessageType.Cron); + + expect(mockConfig.startAutomaticActiveTodoWorkChain).toHaveBeenCalledWith( + 'prompt-cron', + undefined, + ); + expect(mockConfig.endAutomaticActiveTodoWorkChain).toHaveBeenCalledWith( + 'prompt-cron', + ); + + await runTurn(SendMessageType.Retry); + + expect(mockConfig.startActiveTodoWorkChain).toHaveBeenCalledWith( + 'prompt-retry', + 'prompt-userQuery', + ); + }); + + it('includes active Todo context on the first retry request', async () => { + const reminder = + 'unfinished todo: run tests'; + vi.mocked(mockConfig.takeActiveTodoReminder).mockReturnValue(reminder); + + await runTurn(SendMessageType.UserQuery); + await runTurn(SendMessageType.Retry); + + const request = mockTurnRunFn.mock.lastCall?.[1] as unknown[]; + expect(request).toContain(reminder); + }); + + it('continues the carried Todo work chain for related notifications', async () => { + mockTurnRunFn.mockReturnValue( + (async function* () { + yield { type: GeminiEventType.Content, value: 'response' }; + })(), + ); + + const stream = client.sendMessageStream( + [{ text: 'related notification' }], + new AbortController().signal, + 'prompt-related-notification', + { + type: SendMessageType.Notification, + todoWorkChainId: 'prompt-owner', + }, + ); + for await (const _ of stream) { + // drain + } + + expect(mockConfig.startAutomaticActiveTodoWorkChain).toHaveBeenCalledWith( + 'prompt-related-notification', + 'prompt-owner', + ); + }); + + it('keeps automatic Todo ownership through its tool-result turns', async () => { + const reminder = + 'unfinished todo: finish automatic work'; + vi.mocked(mockConfig.takeActiveTodoReminder).mockReturnValue(reminder); + mockTurnRunFn + .mockReturnValueOnce( + (async function* () { + yield { + type: GeminiEventType.ToolCallRequest, + value: { callId: 'call-1', name: 'read_file', args: {} }, + }; + })(), + ) + .mockReturnValueOnce( + (async function* () { + yield { type: GeminiEventType.Content, value: 'done' }; + })(), + ); + + for await (const _ of client.sendMessageStream( + [{ text: 'automatic work' }], + new AbortController().signal, + 'prompt-automatic', + { type: SendMessageType.Notification }, + )) { + // drain + } + expect(mockConfig.endAutomaticActiveTodoWorkChain).not.toHaveBeenCalled(); + + for await (const _ of client.sendMessageStream( + [{ functionResponse: { name: 'read_file', response: { ok: true } } }], + new AbortController().signal, + 'prompt-automatic', + { type: SendMessageType.ToolResult }, + )) { + // drain + } + + expect(mockConfig.takeActiveTodoReminder).toHaveBeenCalledWith( + 'prompt-automatic', + ); + expect(mockConfig.endAutomaticActiveTodoWorkChain).toHaveBeenCalledWith( + 'prompt-automatic', + ); + }); + 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 0f91f2f41a..f2af5cd93d 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -200,6 +200,8 @@ export interface SendMessageOptions { }; /** Display text for notification messages (persisted for session resume). */ notificationDisplayText?: string; + /** Todo work chain that owns this automatic turn, when it is related. */ + todoWorkChainId?: string; /** Model override from skill execution. When present, overrides the session model for this turn. */ modelOverride?: string; /** Exact runtime permit authorizing this Goal-bound turn. */ @@ -325,6 +327,8 @@ export class GeminiClient { private readonly loopDetector: LoopDetectionService; private lastPromptId: string | undefined = undefined; + private activeTodoWorkChainPromptId: string | undefined; + private readonly activeAutomaticTodoWorkChainPromptIds = new Set(); private lastSentIdeContext: IdeContext | undefined; private forceFullIdeContext = true; private recentCompletedToolNames: string[] = []; @@ -2447,6 +2451,30 @@ 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.activeAutomaticTodoWorkChainPromptIds.clear(); + 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.startAutomaticActiveTodoWorkChain( + prompt_id, + options?.todoWorkChainId ?? + (messageType === SendMessageType.Teammate + ? this.activeTodoWorkChainPromptId + : undefined), + ); + this.activeAutomaticTodoWorkChainPromptIds.add(prompt_id); + } const isTopLevelInteraction = messageType === SendMessageType.UserQuery || messageType === SendMessageType.Cron || @@ -2486,6 +2514,7 @@ export class GeminiClient { // early-return) leaves this `false`, and the `finally` block aborts the // prefetch as a safety net. let normalCompletion = false; + let hasToolCalls = false; // Declared outside the try so the finally block can close it out on // uncaught-exception exits too; created (when the hook is registered) // right before the turn's streaming loop below. @@ -2901,6 +2930,39 @@ export class GeminiClient { requestToSend = [...systemReminders, ...requestToSend]; } + if ( + messageType === SendMessageType.Retry || + messageType === SendMessageType.Cron || + messageType === SendMessageType.Notification || + messageType === SendMessageType.Teammate + ) { + const activeTodoReminder = this.config.takeActiveTodoReminder( + prompt_id, + true, + ); + const alreadyHasActiveTodoReminder = requestToSend.some( + (part) => + part === activeTodoReminder || + (typeof part === 'object' && + part !== null && + 'text' in part && + part.text === activeTodoReminder), + ); + if (activeTodoReminder && !alreadyHasActiveTodoReminder) { + const insertAt = requestToSend.findIndex( + (part) => + typeof part !== 'object' || + part === null || + !('functionResponse' in part), + ); + requestToSend.splice( + insertAt < 0 ? requestToSend.length : insertAt, + 0, + activeTodoReminder, + ); + } + } + if (messageType === SendMessageType.ToolResult) { const toolResultMemory = await this.tryConsumeMemoryPrefetch('tool_result'); @@ -2915,6 +2977,21 @@ export class GeminiClient { // text as a separate user message after the tool messages. requestToSend = [...requestToSend, toolResultMemory.prompt]; } + const activeTodoReminder = + this.config.takeActiveTodoReminder(prompt_id); + if (activeTodoReminder) { + const insertAt = requestToSend.findIndex( + (part) => + typeof part !== 'object' || + part === null || + !('functionResponse' in part), + ); + requestToSend.splice( + insertAt < 0 ? requestToSend.length : insertAt, + 0, + activeTodoReminder, + ); + } await this.microcompactHistoryBeforeSend(null, { sizeOnly: true, pendingContent: createUserContent(requestToSend), @@ -2977,7 +3054,6 @@ export class GeminiClient { const resultStream = turn.run(model, requestToSend, signal); let didUpdateIdeContextState = false; let steerInputSettled = false; - let hasToolCalls = false; try { for await (const event of resultStream) { if (!steerInputSettled) { @@ -3193,6 +3269,7 @@ export class GeminiClient { } if (isTopLevelInteraction) endInteractionSpan(signal.aborted ? 'cancelled' : 'ok'); + hasToolCalls = steeredTurn.pendingToolCalls.length > 0; normalCompletion = true; return steeredTurn; } @@ -3525,6 +3602,7 @@ export class GeminiClient { } if (isTopLevelInteraction) endInteractionSpan(signal.aborted ? 'cancelled' : 'ok'); + hasToolCalls = hookTurn.pendingToolCalls.length > 0; // Preserve the pending prefetch: the inner Hook turn we just // yielded may have produced tool calls, and the caller's next // ToolResult turn still needs to consume the recall result. @@ -3640,6 +3718,7 @@ export class GeminiClient { } if (isTopLevelInteraction) endInteractionSpan(signal.aborted ? 'cancelled' : 'ok'); + hasToolCalls = continueTurn.pendingToolCalls.length > 0; // Preserve the pending prefetch: same reasoning as the // `return hookTurn` site above — the recursive Hook turn may // have produced tool calls whose ToolResult turn still needs @@ -3685,6 +3764,13 @@ export class GeminiClient { } throw error; } finally { + if ( + this.activeAutomaticTodoWorkChainPromptIds.has(prompt_id) && + (!normalCompletion || !hasToolCalls) + ) { + this.activeAutomaticTodoWorkChainPromptIds.delete(prompt_id); + this.config.endAutomaticActiveTodoWorkChain(prompt_id); + } if (!goalPermitReleased && (callerSignal.aborted || !normalCompletion)) { await releaseGoalPermitOnInterruptedExit(); } diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 92cacff053..9c21fe4a83 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -84,6 +84,10 @@ 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, + todoWorkChainContext, +} from '../utils/promptIdContext.js'; type ToolSpanRecord = { name: string; @@ -737,6 +741,10 @@ describe('CoreToolScheduler', () => { visionBridge?: boolean; visionAgent?: boolean; onToolResultFullTurnModel?: (model: string) => boolean; + getActiveTodoWorkChainOwner?: ( + promptId: string, + fallbackOwner?: string, + ) => string; }) { const ensureTool = vi.fn( async (name: string) => @@ -829,6 +837,7 @@ describe('CoreToolScheduler', () => { isInteractive: () => true, getInputFormat: () => undefined, getExperimentalZedIntegration: () => false, + getActiveTodoWorkChainOwner: options.getActiveTodoWorkChainOwner, } as unknown as Config, onAllToolCallsComplete, onToolCallsUpdate, @@ -858,6 +867,8 @@ describe('CoreToolScheduler', () => { promptId: 'unrelated-prompt', }; let observedContext: InvocationContextV1 | undefined; + let observedPromptId: string | undefined; + let observedTodoWorkChainId: string | undefined; const tool = new MockTool({ name: 'approval-context-tool', getDefaultPermission: async () => 'ask', @@ -869,12 +880,15 @@ describe('CoreToolScheduler', () => { }), execute: async () => { observedContext = getInvocationContext(); + observedPromptId = promptIdContext.getStore(); + observedTodoWorkChainId = todoWorkChainContext.getStore(); return { llmContent: 'ok', returnDisplay: 'ok' }; }, }); const { scheduler, onToolCallsUpdate } = createSchedulerForLegacyToolTests({ toolsByName: new Map([[tool.name, tool]]), approvalMode: ApprovalMode.DEFAULT, + getActiveTodoWorkChainOwner: () => 'mapped-work-chain', }); await runWithInvocationContext(invocationContext, () => @@ -896,13 +910,17 @@ describe('CoreToolScheduler', () => { 'awaiting_approval', )) as WaitingToolCall; - await runWithInvocationContext(unrelatedContext, () => - waiting.confirmationDetails.onConfirm( - ToolConfirmationOutcome.ProceedOnce, + await todoWorkChainContext.run('stale-work-chain', () => + runWithInvocationContext(unrelatedContext, () => + waiting.confirmationDetails.onConfirm( + ToolConfirmationOutcome.ProceedOnce, + ), ), ); expect(observedContext).toEqual(invocationContext); + expect(observedPromptId).toBe(invocationContext.promptId); + expect(observedTodoWorkChainId).toBe('mapped-work-chain'); }); 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 53823c89b1..54f53613c1 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -70,6 +70,10 @@ import { type AvailableSkillEntry, } from '../tools/skill-utils.js'; import { escapeSystemReminderTags } from '../utils/xml.js'; +import { + promptIdContext, + todoWorkChainContext, +} 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'; @@ -4217,6 +4221,15 @@ export class CoreToolScheduler { } } + const inheritedTodoWorkChainId = todoWorkChainContext.getStore(); + const todoWorkChainId = + this.config.getActiveTodoWorkChainOwner?.( + scheduledCall.request.prompt_id, + inheritedTodoWorkChainId, + ) ?? + inheritedTodoWorkChainId ?? + scheduledCall.request.prompt_id; + if (invocation instanceof ShellToolInvocation) { const setPidCallback = (pid: number) => { this.toolCalls = this.toolCalls.map((tc) => @@ -4250,20 +4263,28 @@ export class CoreToolScheduler { ); }; this.safelyAddToolArgumentsAttributes(span, invocation.params); - promise = invocation.execute( - execSignal, - liveOutputCallback, - shellExecutionConfig, - setPidCallback, - setPromoteAbortControllerCallback, - canPromoteForegroundShell, + promise = todoWorkChainContext.run(todoWorkChainId, () => + 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 = todoWorkChainContext.run(todoWorkChainId, () => + promptIdContext.run(scheduledCall.request.prompt_id, () => + invocation.execute( + execSignal, + liveOutputCallback, + shellExecutionConfig, + ), + ), ); } diff --git a/packages/core/src/services/backgroundShellRegistry.test.ts b/packages/core/src/services/backgroundShellRegistry.test.ts index c64f1ae5d9..ba07b80f64 100644 --- a/packages/core/src/services/backgroundShellRegistry.test.ts +++ b/packages/core/src/services/backgroundShellRegistry.test.ts @@ -26,6 +26,7 @@ import { statusFilePathFor, type ShellTaskRegistration, } from './backgroundShellRegistry.js'; +import { todoWorkChainContext } from '../utils/promptIdContext.js'; let tmpDirs: string[] = []; @@ -67,6 +68,15 @@ function makeEntry( describe('BackgroundShellRegistry', () => { describe('register / get / getAll', () => { + it('captures the Todo work-chain owner at registration', () => { + const reg = new BackgroundShellRegistry(); + const entry = todoWorkChainContext.run('work-chain-1', () => + reg.register(makeEntry()), + ); + + expect(entry.todoWorkChainId).toBe('work-chain-1'); + }); + it('round-trips a registered entry by id', () => { const reg = new BackgroundShellRegistry(); const e = makeEntry({ shellId: 'a' }); diff --git a/packages/core/src/services/backgroundShellRegistry.ts b/packages/core/src/services/backgroundShellRegistry.ts index fb3e61328d..094f647c70 100644 --- a/packages/core/src/services/backgroundShellRegistry.ts +++ b/packages/core/src/services/backgroundShellRegistry.ts @@ -23,6 +23,7 @@ import * as fs from 'node:fs'; import type { TaskBase, TaskRegistration } from '../agents/tasks/types.js'; import { atomicWriteFileSync } from '../utils/atomicFileWrite.js'; import { createDebugLogger } from '../utils/debugLogger.js'; +import { todoWorkChainContext } from '../utils/promptIdContext.js'; import { escapeXml } from '../utils/xml.js'; const debugLogger = createDebugLogger('BACKGROUND_SHELLS'); @@ -244,6 +245,7 @@ export interface ShellNotificationMeta { shellId: string; status: BackgroundShellStatus; exitCode?: number; + todoWorkChainId?: string; } export type BackgroundShellNotificationCallback = ( @@ -309,6 +311,7 @@ export class BackgroundShellRegistry { entry.outputFile = registration.outputPath; entry.outputOffset = 0; entry.notified = false; + entry.todoWorkChainId ??= todoWorkChainContext.getStore(); this.entries.set(entry.shellId, entry); this.writeStatusFile(entry); this.fireRegister(entry); @@ -550,6 +553,7 @@ export class BackgroundShellRegistry { shellId: entry.shellId, status: entry.status, exitCode: entry.exitCode, + todoWorkChainId: entry.todoWorkChainId, }; try { diff --git a/packages/core/src/services/cronScheduler.ts b/packages/core/src/services/cronScheduler.ts index 4e0841f08b..60c0579ddb 100644 --- a/packages/core/src/services/cronScheduler.ts +++ b/packages/core/src/services/cronScheduler.ts @@ -92,6 +92,7 @@ export interface CronJob { delivery?: CronTaskDelivery; /** One-shot that was due while no owning session ran — fired late. */ missed?: boolean; + todoWorkChainId?: string; } /** @@ -104,6 +105,7 @@ interface SessionWakeup { fireAtMs: number; prompt: string; createdAt: number; + todoWorkChainId?: string; } /** @@ -219,6 +221,7 @@ function wakeupToJob(wakeup: SessionWakeup): CronJob { expiresAt: Infinity, fireAtMs: wakeup.fireAtMs, jitterMs: 0, + todoWorkChainId: wakeup.todoWorkChainId, }; } @@ -426,6 +429,7 @@ export class CronScheduler { scheduleWakeup( delaySeconds: number, prompt: string, + todoWorkChainId?: string, ): { id: string; scheduledFor: string; @@ -472,7 +476,13 @@ export class CronScheduler { if (replacedId) { debugLogger.debug(`Replacing pending wakeup ${replacedId}`); } - this.wakeups.set(id, { id, fireAtMs, prompt, createdAt: now }); + this.wakeups.set(id, { + id, + fireAtMs, + prompt, + createdAt: now, + todoWorkChainId, + }); debugLogger.debug( `Wakeup ${id} scheduled for ${new Date(fireAtMs).toISOString()} ` + `(delay=${clampedDelaySeconds}s)`, diff --git a/packages/core/src/services/monitorRegistry.test.ts b/packages/core/src/services/monitorRegistry.test.ts index 4526e2b4d7..054eafa7ac 100644 --- a/packages/core/src/services/monitorRegistry.test.ts +++ b/packages/core/src/services/monitorRegistry.test.ts @@ -10,6 +10,7 @@ import { MonitorRegistry, type MonitorTaskRegistration, } from './monitorRegistry.js'; +import { todoWorkChainContext } from '../utils/promptIdContext.js'; function createEntry( overrides: Partial = {}, @@ -51,6 +52,14 @@ describe('MonitorRegistry', () => { expect(registry.get('mon-1')).toBe(entry); }); + it('captures the Todo work-chain owner at registration', () => { + const entry = todoWorkChainContext.run('work-chain-1', () => + registry.register(createEntry()), + ); + + expect(entry.todoWorkChainId).toBe('work-chain-1'); + }); + it('emits event notification via callback', () => { const callback = vi.fn(); registry.setNotificationCallback(callback); diff --git a/packages/core/src/services/monitorRegistry.ts b/packages/core/src/services/monitorRegistry.ts index 01e4fdcd49..cde406fa6f 100644 --- a/packages/core/src/services/monitorRegistry.ts +++ b/packages/core/src/services/monitorRegistry.ts @@ -18,6 +18,7 @@ import * as path from 'node:path'; import { sanitizeFilenameComponent } from '../agents/agent-transcript.js'; import { createDebugLogger } from '../utils/debugLogger.js'; +import { todoWorkChainContext } from '../utils/promptIdContext.js'; import { stripDisplayControlChars } from '../utils/terminalSafe.js'; import { escapeXml } from '../utils/xml.js'; import type { TaskBase, TaskRegistration } from '../agents/tasks/types.js'; @@ -118,6 +119,7 @@ export interface MonitorNotificationMeta { eventCount: number; toolUseId?: string; ownerAgentId?: string; + todoWorkChainId?: string; } export type MonitorNotificationCallback = ( @@ -184,6 +186,7 @@ export class MonitorRegistry { entry.kind = 'monitor'; entry.outputOffset = 0; entry.notified = false; + entry.todoWorkChainId ??= todoWorkChainContext.getStore(); this.monitors.set(entry.monitorId, entry); debugLogger.info(`Registered monitor: ${entry.monitorId}`); this.resetIdleTimer(entry); @@ -524,6 +527,7 @@ export class MonitorRegistry { eventCount: entry.eventCount, toolUseId: entry.toolUseId, ownerAgentId: entry.ownerAgentId, + todoWorkChainId: entry.todoWorkChainId, }; this.dispatchNotification(entry, displayLine, xmlParts.join('\n'), meta); @@ -577,6 +581,7 @@ export class MonitorRegistry { eventCount: entry.eventCount, toolUseId: entry.toolUseId, ownerAgentId: entry.ownerAgentId, + todoWorkChainId: entry.todoWorkChainId, }; this.dispatchNotification(entry, displayLine, xmlParts.join('\n'), meta); diff --git a/packages/core/src/tools/loop-wakeup.test.ts b/packages/core/src/tools/loop-wakeup.test.ts index e792d957ea..e3fb478fe0 100644 --- a/packages/core/src/tools/loop-wakeup.test.ts +++ b/packages/core/src/tools/loop-wakeup.test.ts @@ -12,6 +12,7 @@ import { Storage } from '../config/storage.js'; import { CronScheduler } from '../services/cronScheduler.js'; import type { Config } from '../config/config.js'; import { LoopWakeupTool } from './loop-wakeup.js'; +import { todoWorkChainContext } from '../utils/promptIdContext.js'; // The scheduling math (clamp / wasClamped / second-precise fire time) is // covered in cronScheduler.test.ts `session wakeups`. These tests cover the @@ -103,7 +104,9 @@ describe('LoopWakeupTool', () => { reason: 'CI is still running', }); - const result = await invocation.execute(new AbortController().signal); + const result = await todoWorkChainContext.run('work-chain-1', () => + invocation.execute(new AbortController().signal), + ); expect(result.error).toBeUndefined(); expect(result.llmContent).toContain('Session-only one-shot'); @@ -114,6 +117,7 @@ describe('LoopWakeupTool', () => { expect(scheduler.list()[0]).toMatchObject({ cronExpr: '@wakeup', prompt: 'continue loop', + todoWorkChainId: 'work-chain-1', }); expect(scheduler.size).toBe(1); }); diff --git a/packages/core/src/tools/loop-wakeup.ts b/packages/core/src/tools/loop-wakeup.ts index 73fb0f728d..097976f456 100644 --- a/packages/core/src/tools/loop-wakeup.ts +++ b/packages/core/src/tools/loop-wakeup.ts @@ -15,6 +15,7 @@ import { clampWakeupSeconds, } from '../services/cronScheduler.js'; import { getErrorMessage } from '../utils/errors.js'; +import { todoWorkChainContext } from '../utils/promptIdContext.js'; export interface LoopWakeupParams { delaySeconds: number; @@ -83,7 +84,11 @@ class LoopWakeupInvocation extends BaseToolInvocation< }; } const { id, scheduledFor, clampedDelaySeconds, wasClamped, replacedId } = - scheduler.scheduleWakeup(this.params.delaySeconds, prompt); + scheduler.scheduleWakeup( + this.params.delaySeconds, + prompt, + todoWorkChainContext.getStore(), + ); const reason = this.params.reason?.trim(); const llmContent = [ diff --git a/packages/core/src/tools/todoWrite.test.ts b/packages/core/src/tools/todoWrite.test.ts index 2b796f7345..59d69d1493 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,48 @@ 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(1100); + }); + + it('skips active Todo reminder when no prompt id is active', async () => { + const params: TodoWriteParams = { + todos: [{ id: '1', content: 'Task 1', status: 'pending' }], + }; + const enoentError = new Error('ENOENT') as Error & { code: string }; + enoentError.code = 'ENOENT'; + mockFs.readFile.mockRejectedValue(enoentError); + mockFs.mkdir.mockResolvedValue(undefined); + mockAtomicWrite.mockResolvedValue(undefined); + + const result = await tool.build(params).execute(mockAbortSignal); + + expect(result.llmContent).toContain( + 'Todos have been modified successfully', + ); + expect(mockConfig.setActiveTodoReminder).not.toHaveBeenCalled(); }); it('should replace todos with new ones', async () => { @@ -194,7 +240,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 +262,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 +308,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 +325,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 +852,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 e58501b97b..03c3edeff4 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 = 800; export interface TodoWriteParams { todos: TodoItem[]; @@ -250,6 +253,27 @@ 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( + unfinishedTodos + .map((todo) => `- [${todo.status}] ${todo.content}`) + .join('\n'), + ); + 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 diff --git a/packages/core/src/utils/promptIdContext.ts b/packages/core/src/utils/promptIdContext.ts index 6344bd0b83..45ac19fdad 100644 --- a/packages/core/src/utils/promptIdContext.ts +++ b/packages/core/src/utils/promptIdContext.ts @@ -7,3 +7,4 @@ import { AsyncLocalStorage } from 'node:async_hooks'; export const promptIdContext = new AsyncLocalStorage(); +export const todoWorkChainContext = new AsyncLocalStorage();