diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 61beb142c10..8cc540b90c2 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -33,6 +33,8 @@ import { SendMessageType, ToolErrorType, ToolConfirmationOutcome, + getRuntimeContentGenerator, + runWithRuntimeContentGenerator, } from '@qwen-code/qwen-code-core'; import type { Part, PartListUnion } from '@google/genai'; import type { UseHistoryManagerReturn } from './useHistoryManager.js'; @@ -6435,6 +6437,50 @@ describe('useGeminiStream', () => { mockSendMessageStream.mock.calls[0][3].modelOverride, ).toBeUndefined(); }); + + // 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 + // executes on a stack that still carries the subagent's frame. The + // drained turn — and every async continuation it starts — then + // resolves Config.getModel() to the subagent's runtime view and the + // main session switches onto the subagent's model. The drain effect + // must therefore run via runOutsideAgentContext. This test drives the + // notification callback from inside an agent frame (bypassing the + // producer-side guard in BackgroundTaskRegistry.emitNotification) and + // fails if the consumer-side wrapping is removed. + it('drains a notification outside a background agent ALS frame', async () => { + renderTestHook(); + const callback = mockBackgroundShellRegistry.setNotificationCallback + .mock.calls[0][0] as (displayText: string, modelText: string) => void; + + let capturedRuntimeView: unknown = 'unset'; + mockSendMessageStream.mockImplementationOnce(() => { + capturedRuntimeView = getRuntimeContentGenerator(); + return (async function* () {})(); + }); + + const subagentView = { + contentGenerator: {}, + contentGeneratorConfig: { model: 'small-default' }, + } as never; + // The whole act() flush runs inside the agent frame, mirroring the + // contaminated React commit from the issue. + await runWithRuntimeContentGenerator(subagentView, async () => { + await act(async () => { + callback( + 'Background shell "npm test" completed.', + 'completed', + ); + }); + }); + + await waitFor(() => expect(mockSendMessageStream).toHaveBeenCalled()); + expect(mockSendMessageStream.mock.calls[0][3]).toMatchObject({ + type: SendMessageType.Notification, + }); + expect(capturedRuntimeView).toBeUndefined(); + }); }); }); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 2be5bb40dbf..190fda936f3 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -13,6 +13,7 @@ import { useLayoutEffect, } from 'react'; import { + runOutsideAgentContext, type Config, type EditorType, type GeminiClient, @@ -3760,45 +3761,59 @@ export const useGeminiStream = ( !isSubmittingQueryRef.current && notificationQueueRef.current.length > 0 ) { - const queue = notificationQueueRef.current; - const targetType = queue[0]!.sendMessageType; - - // Cron prompts must run as individual turns — each needs its own - // slash/shell/@ preprocessing and approval cycle. Only batch - // Notification items (which pass through without preprocessing). - if (targetType === SendMessageType.Cron) { - const item = queue.shift()!; - addItem( - { type: 'notification' as const, text: item.displayText }, - Date.now(), - ); - submitQuery(item.modelText, item.sendMessageType, undefined, { - notificationDisplayText: item.displayText, - onDelivered: item.onDelivered, - onDeliveryFailed: item.onDeliveryFailed, - }); - return; - } + // Consumer-side guard for #7156: this effect can run on a render pass + // that React batched together with progress setState calls issued from + // INSIDE a subagent's AsyncLocalStorage frame, in which case the whole + // synchronous effect stack — and every async continuation submitQuery + // starts — inherits the subagent's runtime view, and the notification + // turn resolves Config.getModel() to the SUBAGENT's model. Exiting the + // frame here guarantees the drained turn always runs on the main + // session's configuration, regardless of which producer's setState + // triggered the commit. + runOutsideAgentContext(() => { + const queue = notificationQueueRef.current; + const targetType = queue[0]!.sendMessageType; + + // Cron prompts must run as individual turns — each needs its own + // slash/shell/@ preprocessing and approval cycle. Only batch + // Notification items (which pass through without preprocessing). + if (targetType === SendMessageType.Cron) { + const item = queue.shift()!; + addItem( + { type: 'notification' as const, text: item.displayText }, + Date.now(), + ); + submitQuery(item.modelText, item.sendMessageType, undefined, { + notificationDisplayText: item.displayText, + onDelivered: item.onDelivered, + onDeliveryFailed: item.onDeliveryFailed, + }); + return; + } - // Drain contiguous leading Notification items into one batch. - let splitIdx = 0; - while ( - splitIdx < queue.length && - queue[splitIdx]!.sendMessageType === targetType - ) { - splitIdx++; - } - const batch = queue.splice(0, splitIdx); + // Drain contiguous leading Notification items into one batch. + let splitIdx = 0; + while ( + splitIdx < queue.length && + queue[splitIdx]!.sendMessageType === targetType + ) { + splitIdx++; + } + const batch = queue.splice(0, splitIdx); - const now = Date.now(); - for (const item of batch) { - addItem({ type: 'notification' as const, text: item.displayText }, now); - } + const now = Date.now(); + for (const item of batch) { + addItem( + { type: 'notification' as const, text: item.displayText }, + now, + ); + } - const combinedModelText = batch.map((e) => e.modelText).join('\n\n'); - const combinedDisplayText = batch.map((e) => e.displayText).join('; '); - submitQuery(combinedModelText, targetType, undefined, { - notificationDisplayText: combinedDisplayText, + const combinedModelText = batch.map((e) => e.modelText).join('\n\n'); + const combinedDisplayText = batch.map((e) => e.displayText).join('; '); + submitQuery(combinedModelText, targetType, undefined, { + notificationDisplayText: combinedDisplayText, + }); }); } }, [streamingState, submitQuery, notificationTrigger, addItem]); diff --git a/packages/core/src/agents/background-tasks.test.ts b/packages/core/src/agents/background-tasks.test.ts index 061934fb2aa..11ae0e63c09 100644 --- a/packages/core/src/agents/background-tasks.test.ts +++ b/packages/core/src/agents/background-tasks.test.ts @@ -16,6 +16,12 @@ import { type BackgroundApproval, type BackgroundTaskEntry, } from './background-tasks.js'; +import { + getCurrentAgentId, + getRuntimeContentGenerator, + runWithAgentContext, + runWithRuntimeContentGenerator, +} from './runtime/agent-context.js'; import * as transcript from './agent-transcript.js'; import { AgentEventEmitter, AgentEventType } from './runtime/agent-events.js'; import { ToolConfirmationOutcome } from '../tools/tools.js'; @@ -52,6 +58,56 @@ function makeRegistration( }; } +describe('notification emission and agent context (#7156)', () => { + // 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, + // the drain effect, the next conversation turn). The callback must + // therefore run with NO agent frame, or Config.getModel() resolves to + // the subagent's model for the notification turn and the main session's + // history can overflow the smaller context window. + it('invokes the notification callback outside the agent ALS frame', async () => { + const registry = new BackgroundTaskRegistry(); + const seen: Array<{ + agentId: string | null; + runtimeView: unknown; + }> = []; + registry.setNotificationCallback(() => { + seen.push({ + agentId: getCurrentAgentId(), + runtimeView: getRuntimeContentGenerator(), + }); + }); + + registry.register({ + agentId: 'bg-1', + description: 'bg agent', + status: 'running', + startTime: Date.now(), + abortController: new AbortController(), + isBackgrounded: true, + outputFile: '/tmp/bg.jsonl', + }); + + const fakeView = { + contentGenerator: {} as never, + contentGeneratorConfig: { model: 'subagent-model' } as never, + }; + await runWithAgentContext('bg-1', () => + runWithRuntimeContentGenerator(fakeView, async () => { + // Sanity: we ARE inside the subagent frame here. + expect(getCurrentAgentId()).toBe('bg-1'); + expect(getRuntimeContentGenerator()).toBe(fakeView); + registry.complete('bg-1', 'done'); + }), + ); + + expect(seen).toHaveLength(1); + expect(seen[0]!.agentId).toBeNull(); + expect(seen[0]!.runtimeView).toBeUndefined(); + }); +}); + describe('BackgroundTaskRegistry', () => { let registry: BackgroundTaskRegistry; diff --git a/packages/core/src/agents/background-tasks.ts b/packages/core/src/agents/background-tasks.ts index cca7d28f005..a1be9bff549 100644 --- a/packages/core/src/agents/background-tasks.ts +++ b/packages/core/src/agents/background-tasks.ts @@ -26,6 +26,7 @@ import { createDebugLogger } from '../utils/debugLogger.js'; import { parsePositiveIntegerEnv } from '../utils/env.js'; import { escapeXml } from '../utils/xml.js'; import { patchAgentMeta } from './agent-transcript.js'; +import { runOutsideAgentContext } from './runtime/agent-context.js'; import { AgentEventType, type AgentApprovalRequestEvent, @@ -1428,7 +1429,18 @@ export class BackgroundTaskRegistry { }; try { - this.notificationCallback(displayLine, xmlParts.join('\n'), meta); + // The terminal transition (complete/fail/cancel) that reaches this + // point runs inside the finished agent's AsyncLocalStorage frame, and + // ALS context follows every async continuation the callback starts — + // including the React state update that drains the notification into + // a new conversation turn. Without exiting the frame here, that turn + // resolves Config.getModel() to the SUBAGENT's model and the main + // session's history can overflow its smaller context window (#7156). + // A notification is main-session-owned, so emit it with no agent + // frame at all. + runOutsideAgentContext(() => + this.notificationCallback!(displayLine, xmlParts.join('\n'), meta), + ); } catch (error) { debugLogger.error('Failed to emit background notification:', error); } diff --git a/packages/core/src/agents/runtime/agent-context.ts b/packages/core/src/agents/runtime/agent-context.ts index 2a7bc4276f5..433be1ba66d 100644 --- a/packages/core/src/agents/runtime/agent-context.ts +++ b/packages/core/src/agents/runtime/agent-context.ts @@ -98,6 +98,23 @@ export function getRuntimeContentGenerator(): return storage.getStore()?.runtimeView; } +/** + * Runs `fn` with NO agent frame on the async-local stack, so + * `Config.getModel()` / `getContentGeneratorConfig()` resolve to the main + * session's configuration and `getCurrentAgentId()` returns null. + * + * AsyncLocalStorage context propagates through every async continuation + * started inside `fn` — React state updates, queued microtasks, timers — + * which is exactly how a background agent's runtime view leaked into the + * notification drain and switched the main session onto the subagent's + * model (#7156). Wrap main-session-owned work that can be triggered from + * inside an agent frame (notification emission, completion bookkeeping) + * with this helper. + */ +export function runOutsideAgentContext(fn: () => T): T { + return storage.exit(fn); +} + /** * True when there is no active agent frame — i.e. we are in the top-level * user session, not inside a sub-agent. The canonical "top-level only" diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 0ed26ae50a4..8516e302bf2 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -62,6 +62,7 @@ export { getRuntimeContentGenerator, runWithRuntimeContentGenerator, type RuntimeContentGeneratorView, + runOutsideAgentContext, } from './agents/runtime/agent-context.js'; export * from './core/reasoning-effort.js'; export * from './core/coreToolScheduler.js';