mirror of
https://github.com/QwenLM/qwen-code.git
synced 2026-08-22 07:04:58 +00:00
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 <qwen-code-dev-bot@users.noreply.github.com> Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>
This commit is contained in:
易良
GitHub
parent
eb28b3038e
commit
2abfa3d54e
29 changed files with 1260 additions and 53 deletions
43
docs/design/active-todo-context.md
Normal file
43
docs/design/active-todo-context.md
Normal file
|
|
@ -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.
|
||||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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 =
|
||||
'<system-reminder>unfinished todo: run tests</system-reminder>';
|
||||
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 =
|
||||
'<system-reminder>unfinished todo: wait for agent</system-reminder>';
|
||||
vi.mocked(mockConfig.takeActiveTodoReminder).mockReturnValue(reminder);
|
||||
const internals = session as unknown as {
|
||||
relatedAgentIds: Set<string>;
|
||||
};
|
||||
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.', '<task-notification/>', {
|
||||
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<string>;
|
||||
};
|
||||
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.', '<task-notification/>', {
|
||||
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<void>((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 =
|
||||
'<system-reminder>unfinished todo: check tests</system-reminder>';
|
||||
const activeTodoReminders = new Map<string, string>();
|
||||
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: {
|
||||
|
|
|
|||
|
|
@ -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<string, unknown> })
|
||||
._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<NextMessageAfterToolRun> {
|
||||
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 `<<loop.md>>` / `<<loop.md-dynamic>>` 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<RunToolResult> {
|
||||
// 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(
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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 =
|
||||
'<task-notification>\n' +
|
||||
'<task-id>mon_1</task-id>\n' +
|
||||
'<kind>monitor</kind>\n' +
|
||||
'<status>running</status>\n' +
|
||||
'<summary>Monitor emitted event #1.</summary>\n' +
|
||||
'<result>ready</result>\n' +
|
||||
'</task-notification>';
|
||||
const secondNotificationXml =
|
||||
'<task-notification>\n' +
|
||||
'<task-id>mon_2</task-id>\n' +
|
||||
'<kind>monitor</kind>\n' +
|
||||
'<status>running</status>\n' +
|
||||
'<summary>Monitor emitted event #2.</summary>\n' +
|
||||
'<result>also ready</result>\n' +
|
||||
'</task-notification>';
|
||||
|
||||
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);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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',
|
||||
'<task-notification>first</task-notification>',
|
||||
{
|
||||
monitorId: 'mon_1',
|
||||
status: 'completed',
|
||||
todoWorkChainId: 'chain-1',
|
||||
},
|
||||
);
|
||||
callback(
|
||||
'Monitor "build" event #1: ready',
|
||||
'<task-notification>second</task-notification>',
|
||||
{
|
||||
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',
|
||||
'<task-notification>third</task-notification>',
|
||||
{
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
});
|
||||
});
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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');
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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<string, string>();
|
||||
private activeTodoWorkChainOwners = new Map<string, string>();
|
||||
private activeTodoReminderTurns = new Map<string, number>();
|
||||
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<string, string> {
|
||||
if (!Object.prototype.hasOwnProperty.call(this, 'activeTodoReminders')) {
|
||||
this.activeTodoReminders = new Map();
|
||||
}
|
||||
return this.activeTodoReminders;
|
||||
}
|
||||
|
||||
private getOwnActiveTodoWorkChainOwners(): Map<string, string> {
|
||||
if (
|
||||
!Object.prototype.hasOwnProperty.call(this, 'activeTodoWorkChainOwners')
|
||||
) {
|
||||
this.activeTodoWorkChainOwners = new Map();
|
||||
}
|
||||
return this.activeTodoWorkChainOwners;
|
||||
}
|
||||
|
||||
private getOwnActiveTodoReminderTurns(): Map<string, number> {
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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 =
|
||||
'<system-reminder>unfinished todo: run tests</system-reminder>';
|
||||
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 =
|
||||
'<system-reminder>unfinished todo: run tests</system-reminder>';
|
||||
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 =
|
||||
'<system-reminder>unfinished todo: finish automatic work</system-reminder>';
|
||||
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) =>
|
||||
|
|
|
|||
|
|
@ -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<string>();
|
||||
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();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
),
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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' });
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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)`,
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import {
|
|||
MonitorRegistry,
|
||||
type MonitorTaskRegistration,
|
||||
} from './monitorRegistry.js';
|
||||
import { todoWorkChainContext } from '../utils/promptIdContext.js';
|
||||
|
||||
function createEntry(
|
||||
overrides: Partial<MonitorTaskRegistration> = {},
|
||||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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 = [
|
||||
|
|
|
|||
|
|
@ -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('<system-reminder>');
|
||||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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
|
||||
? `<system-reminder>\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</system-reminder>`
|
||||
: 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
|
||||
|
|
|
|||
|
|
@ -7,3 +7,4 @@
|
|||
import { AsyncLocalStorage } from 'node:async_hooks';
|
||||
|
||||
export const promptIdContext = new AsyncLocalStorage<string>();
|
||||
export const todoWorkChainContext = new AsyncLocalStorage<string>();
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue