diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 4b469a5102b..d52b8c3b51a 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -28,6 +28,9 @@ import * as nonInteractiveCliCommands from '../../nonInteractiveCliCommands.js'; import { CommandKind } from '../../ui/commands/types.js'; const debugLoggerWarnSpy = vi.hoisted(() => vi.fn()); +const mockSetAgentNotificationCallback = vi.hoisted(() => vi.fn()); +const mockSetMonitorNotificationCallback = vi.hoisted(() => vi.fn()); +const mockSetShellNotificationCallback = vi.hoisted(() => vi.fn()); vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { const actual = @@ -40,6 +43,9 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { warn: debugLoggerWarnSpy, error: vi.fn(), }), + setAgentNotificationCallback: mockSetAgentNotificationCallback, + setMonitorNotificationCallback: mockSetMonitorNotificationCallback, + setShellNotificationCallback: mockSetShellNotificationCallback, }; }); @@ -196,15 +202,6 @@ describe('Session', () => { getChat: ReturnType; tryCompressChat: ReturnType; }; - let mockBackgroundTaskRegistry: { - setNotificationCallback: ReturnType; - }; - let mockMonitorRegistry: { - setNotificationCallback: ReturnType; - }; - let mockBackgroundShellRegistry: { - setNotificationCallback: ReturnType; - }; let mockToolRegistry: { getTool: ReturnType; ensureTool: ReturnType; @@ -237,15 +234,9 @@ describe('Session', () => { compressionStatus: core.CompressionStatus.NOOP, }), }; - mockBackgroundTaskRegistry = { - setNotificationCallback: vi.fn(), - }; - mockMonitorRegistry = { - setNotificationCallback: vi.fn(), - }; - mockBackgroundShellRegistry = { - setNotificationCallback: vi.fn(), - }; + mockSetAgentNotificationCallback.mockReset(); + mockSetMonitorNotificationCallback.mockReset(); + mockSetShellNotificationCallback.mockReset(); mockChatRecordingService = { recordUserMessage: vi.fn(), @@ -289,13 +280,10 @@ describe('Session', () => { getSessionTokenLimit: vi.fn().mockReturnValue(0), getStopHookBlockingCap: vi.fn().mockReturnValue(8), getGeminiClient: vi.fn().mockReturnValue(mockGeminiClient), - getBackgroundTaskRegistry: vi - .fn() - .mockReturnValue(mockBackgroundTaskRegistry), - getBackgroundShellRegistry: vi - .fn() - .mockReturnValue(mockBackgroundShellRegistry), - getMonitorRegistry: vi.fn().mockReturnValue(mockMonitorRegistry), + // Background-notification callbacks are keyed by the session's + // TaskRegistry, so the constructor reads this. The task-module setters + // are mocked above, so a stub instance is sufficient. + getTaskRegistry: vi.fn().mockReturnValue({}), } as unknown as Config; mockClient = { @@ -888,8 +876,7 @@ describe('Session', () => { prompt: [{ type: 'text', text: 'start background work' }], }); - const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock - .calls[0][0] as ( + const callback = mockSetAgentNotificationCallback.mock.calls[0][1] as ( displayText: string, modelText: string, meta: { agentId: string; status: string; toolUseId?: string }, @@ -1013,8 +1000,7 @@ describe('Session', () => { prompt: [{ type: 'text', text: 'start background work' }], }); - const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock - .calls[0][0] as ( + const callback = mockSetAgentNotificationCallback.mock.calls[0][1] as ( displayText: string, modelText: string, meta: { agentId: string; status: string; toolUseId?: string }, @@ -1075,8 +1061,7 @@ describe('Session', () => { prompt: [{ type: 'text', text: 'start background work' }], }); - const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock - .calls[0][0] as ( + const callback = mockSetAgentNotificationCallback.mock.calls[0][1] as ( displayText: string, modelText: string, meta: { agentId: string; status: string; toolUseId?: string }, @@ -1108,8 +1093,7 @@ describe('Session', () => { } ).pendingPrompt = new AbortController(); - const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock - .calls[0][0] as ( + const callback = mockSetAgentNotificationCallback.mock.calls[0][1] as ( displayText: string, modelText: string, meta: { agentId: string; status: string; toolUseId?: string }, @@ -1157,8 +1141,7 @@ describe('Session', () => { prompt: [{ type: 'text', text: 'start background work' }], }); - const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock - .calls[0][0] as ( + const callback = mockSetAgentNotificationCallback.mock.calls[0][1] as ( displayText: string, modelText: string, meta: { agentId: string; status: string; toolUseId?: string }, @@ -1223,8 +1206,7 @@ describe('Session', () => { prompt: [{ type: 'text', text: 'start background work' }], }); - const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock - .calls[0][0] as ( + const callback = mockSetAgentNotificationCallback.mock.calls[0][1] as ( displayText: string, modelText: string, meta: { agentId: string; status: string; toolUseId?: string }, @@ -1250,8 +1232,7 @@ describe('Session', () => { prompt: [{ type: 'text', text: 'start monitor' }], }); - const callback = mockMonitorRegistry.setNotificationCallback.mock - .calls[0][0] as ( + const callback = mockSetMonitorNotificationCallback.mock.calls[0][1] as ( displayText: string, modelText: string, meta: { monitorId: string; status: string; toolUseId?: string }, @@ -1310,8 +1291,7 @@ describe('Session', () => { prompt: [{ type: 'text', text: 'start background shell' }], }); - const callback = mockBackgroundShellRegistry.setNotificationCallback.mock - .calls[0][0] as ( + const callback = mockSetShellNotificationCallback.mock.calls[0][1] as ( displayText: string, modelText: string, meta: { shellId: string; status: string }, @@ -4239,15 +4219,18 @@ describe('Session', () => { expect(internals.notificationQueue).toHaveLength(0); expect(internals.cronQueue).toHaveLength(0); expect(internals.notificationProcessing).toBe(false); - expect( - mockBackgroundTaskRegistry.setNotificationCallback, - ).toHaveBeenLastCalledWith(undefined); - expect( - mockMonitorRegistry.setNotificationCallback, - ).toHaveBeenLastCalledWith(undefined); - expect( - mockBackgroundShellRegistry.setNotificationCallback, - ).toHaveBeenLastCalledWith(undefined); + expect(mockSetAgentNotificationCallback).toHaveBeenLastCalledWith( + expect.anything(), + undefined, + ); + expect(mockSetMonitorNotificationCallback).toHaveBeenLastCalledWith( + expect.anything(), + undefined, + ); + expect(mockSetShellNotificationCallback).toHaveBeenLastCalledWith( + expect.anything(), + undefined, + ); }); it('aborts an active notificationAbortController and nulls the reference', () => { @@ -4288,7 +4271,7 @@ describe('Session', () => { const internals = session as unknown as SessionInternals; session.dispose(); const callsAfterFirst = - mockBackgroundTaskRegistry.setNotificationCallback.mock.calls.length; + mockSetAgentNotificationCallback.mock.calls.length; expect(() => session.dispose()).not.toThrow(); expect(internals.disposed).toBe(true); @@ -4296,11 +4279,10 @@ describe('Session', () => { expect(internals.cronQueue).toHaveLength(0); // The second dispose still unregisters (passes undefined again), which // is harmless. We only care that no surprise re-registration occurs. - const last = - mockBackgroundTaskRegistry.setNotificationCallback.mock.calls.at(-1); - expect(last?.[0]).toBeUndefined(); + const last = mockSetAgentNotificationCallback.mock.calls.at(-1); + expect(last?.[1]).toBeUndefined(); expect( - mockBackgroundTaskRegistry.setNotificationCallback.mock.calls.length, + mockSetAgentNotificationCallback.mock.calls.length, ).toBeGreaterThanOrEqual(callsAfterFirst); }); diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 6d38c6ce5fc..14dd6bf55ab 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -76,6 +76,9 @@ import { shouldFirePermissionDeniedForAutoMode, shouldRunAutoModeForCall, acquireSleepInhibitor, + setAgentNotificationCallback, + setMonitorNotificationCallback, + setShellNotificationCallback, } from '@qwen-code/qwen-code-core'; import { getCommandSubcommandNames } from '../../services/commandMetadata.js'; import { getEffectiveSupportedModes } from '../../services/commandUtils.js'; @@ -408,9 +411,10 @@ export class Session implements SessionContext { this.cronProcessing = false; this.cronCompletion = null; - this.config.getBackgroundTaskRegistry().setNotificationCallback(undefined); - this.config.getMonitorRegistry().setNotificationCallback(undefined); - this.config.getBackgroundShellRegistry().setNotificationCallback(undefined); + const taskRegistry = this.config.getTaskRegistry(); + setAgentNotificationCallback(taskRegistry, undefined); + setMonitorNotificationCallback(taskRegistry, undefined); + setShellNotificationCallback(taskRegistry, undefined); } /** @@ -1649,8 +1653,9 @@ export class Session implements SessionContext { } #registerBackgroundNotificationCallbacks(): void { - const backgroundRegistry = this.config.getBackgroundTaskRegistry(); - backgroundRegistry.setNotificationCallback( + const taskRegistry = this.config.getTaskRegistry(); + setAgentNotificationCallback( + taskRegistry, (displayText, modelText, meta) => { this.#enqueueBackgroundNotification({ displayText, @@ -1663,32 +1668,36 @@ export class Session implements SessionContext { }, ); - const monitorRegistry = this.config.getMonitorRegistry(); - monitorRegistry.setNotificationCallback((displayText, modelText, meta) => { - if (meta.status === 'running') { - return; - } + setMonitorNotificationCallback( + taskRegistry, + (displayText, modelText, meta) => { + if (meta.status === 'running') { + return; + } - this.#enqueueBackgroundNotification({ - displayText, - modelText, - taskId: meta.monitorId, - status: meta.status, - kind: 'monitor', - toolUseId: meta.toolUseId, - }); - }); + this.#enqueueBackgroundNotification({ + displayText, + modelText, + taskId: meta.monitorId, + status: meta.status, + kind: 'monitor', + toolUseId: meta.toolUseId, + }); + }, + ); - const shellRegistry = this.config.getBackgroundShellRegistry(); - shellRegistry.setNotificationCallback((displayText, modelText, meta) => { - this.#enqueueBackgroundNotification({ - displayText, - modelText, - taskId: meta.shellId, - status: meta.status, - kind: 'shell', - }); - }); + setShellNotificationCallback( + taskRegistry, + (displayText, modelText, meta) => { + this.#enqueueBackgroundNotification({ + displayText, + modelText, + taskId: meta.shellId, + status: meta.status, + kind: 'shell', + }); + }, + ); } #enqueueBackgroundNotification(item: BackgroundNotificationQueueItem): void { diff --git a/packages/cli/src/nonInteractive/session.test.ts b/packages/cli/src/nonInteractive/session.test.ts index 6cc12775303..d81548cf56d 100644 --- a/packages/cli/src/nonInteractive/session.test.ts +++ b/packages/cli/src/nonInteractive/session.test.ts @@ -6,6 +6,50 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { SendMessageType, type Config } from '@qwen-code/qwen-code-core'; + +// Hoisted mocks for the kind-local helpers session.ts now imports +// from core. The setNotificationCallback/setRegisterCallback mocks +// store the callback so tests can invoke it with synthetic events. +const monitorRegisterCb = vi.hoisted(() => ({ + current: undefined as ((entry: unknown) => void) | undefined, +})); +const monitorNotificationCb = vi.hoisted(() => ({ + current: undefined as + | ((displayText: string, modelText: string, meta: unknown) => void) + | undefined, +})); +const mockSetMonitorRegisterCallback = vi.hoisted(() => + vi.fn((cb: ((entry: unknown) => void) | undefined) => { + monitorRegisterCb.current = cb ?? undefined; + }), +); +const mockSetMonitorNotificationCallback = vi.hoisted(() => + vi.fn( + ( + cb: + | ((displayText: string, modelText: string, meta: unknown) => void) + | undefined, + ) => { + monitorNotificationCb.current = cb ?? undefined; + }, + ), +); +const mockMonitorAbortAll = vi.hoisted(() => vi.fn()); +const mockShellAbortAll = vi.hoisted(() => vi.fn()); +const mockAgentAbortAll = vi.hoisted(() => vi.fn()); + +vi.mock('@qwen-code/qwen-code-core', async () => { + const actual = await vi.importActual('@qwen-code/qwen-code-core'); + return { + ...actual, + setMonitorRegisterCallback: mockSetMonitorRegisterCallback, + setMonitorNotificationCallback: mockSetMonitorNotificationCallback, + monitorAbortAll: mockMonitorAbortAll, + shellAbortAll: mockShellAbortAll, + agentAbortAll: mockAgentAbortAll, + }; +}); + import { runNonInteractiveStreamJson } from './session.js'; import type { CLIUserMessage, @@ -56,16 +100,17 @@ interface ConfigOverrides { [key: string]: unknown; } -let mockMonitorRegistry: { - setNotificationCallback: ReturnType; - setRegisterCallback: ReturnType; - abortAll: ReturnType; -}; -let mockBackgroundShellRegistry: { - abortAll: ReturnType; -}; -let mockBackgroundTaskRegistry: { - abortAll: ReturnType; +// Stub TaskRegistry exposing the narrow surface session.ts uses (none +// directly — it only invokes the kind-local abort helpers from core). +let mockTaskRegistry: { + getAll: () => unknown[]; + getByKind: () => unknown[]; + get: () => unknown; + register: () => unknown; + update: () => unknown; + evict: () => unknown; + kill: () => unknown; + subscribe: () => () => void; }; function createConfig(overrides: ConfigOverrides = {}): Config { @@ -78,9 +123,7 @@ function createConfig(overrides: ConfigOverrides = {}): Config { getOutputFormat: () => 'stream-json', initialize: vi.fn(), waitForMcpReady: vi.fn().mockResolvedValue(undefined), - getMonitorRegistry: () => mockMonitorRegistry, - getBackgroundShellRegistry: () => mockBackgroundShellRegistry, - getBackgroundTaskRegistry: () => mockBackgroundTaskRegistry, + getTaskRegistry: () => mockTaskRegistry, }; return { ...base, ...overrides } as unknown as Config; } @@ -174,17 +217,23 @@ describe('runNonInteractiveStreamJson', () => { }; }; beforeEach(() => { - mockMonitorRegistry = { - setNotificationCallback: vi.fn(), - setRegisterCallback: vi.fn(), - abortAll: vi.fn(), - }; - mockBackgroundShellRegistry = { - abortAll: vi.fn(), - }; - mockBackgroundTaskRegistry = { - abortAll: vi.fn(), + mockTaskRegistry = { + getAll: () => [], + getByKind: () => [], + get: () => undefined, + register: () => undefined, + update: () => undefined, + evict: () => undefined, + kill: () => undefined, + subscribe: () => () => {}, }; + mockSetMonitorRegisterCallback.mockClear(); + mockSetMonitorNotificationCallback.mockClear(); + mockMonitorAbortAll.mockClear(); + mockShellAbortAll.mockClear(); + mockAgentAbortAll.mockClear(); + monitorRegisterCb.current = undefined; + monitorNotificationCb.current = undefined; config = createConfig(); runNonInteractiveMock.mockReset(); @@ -315,10 +364,10 @@ describe('runNonInteractiveStreamJson', () => { }, ) => void) | undefined; - mockMonitorRegistry.setRegisterCallback.mockImplementation((cb) => { + mockSetMonitorRegisterCallback.mockImplementation((_registry, cb) => { registerCallback = cb; }); - mockMonitorRegistry.setNotificationCallback.mockImplementation((cb) => { + mockSetMonitorNotificationCallback.mockImplementation((_registry, cb) => { monitorCallback = cb; }); @@ -421,10 +470,10 @@ describe('runNonInteractiveStreamJson', () => { ) => void) | undefined; - mockMonitorRegistry.setRegisterCallback.mockImplementation((cb) => { + mockSetMonitorRegisterCallback.mockImplementation((_registry, cb) => { registerCallback = cb; }); - mockMonitorRegistry.setNotificationCallback.mockImplementation((cb) => { + mockSetMonitorNotificationCallback.mockImplementation((_registry, cb) => { notificationCallback = cb; }); @@ -480,10 +529,12 @@ describe('runNonInteractiveStreamJson', () => { closeInput?.(); await vi.waitFor(() => { - expect( - mockMonitorRegistry.setNotificationCallback, - ).toHaveBeenLastCalledWith(undefined); - expect(mockMonitorRegistry.setRegisterCallback).toHaveBeenLastCalledWith( + expect(mockSetMonitorNotificationCallback).toHaveBeenLastCalledWith( + expect.anything(), + undefined, + ); + expect(mockSetMonitorRegisterCallback).toHaveBeenLastCalledWith( + expect.anything(), undefined, ); }); @@ -517,11 +568,12 @@ describe('runNonInteractiveStreamJson', () => { ); expect(runNonInteractiveMock).toHaveBeenCalledTimes(2); - const clearCalls = mockMonitorRegistry.setNotificationCallback.mock.calls - .map(([cb]) => cb) + const clearCalls = mockSetMonitorNotificationCallback.mock.calls + .map(([, cb]) => cb) .filter((cb) => cb === undefined); expect(clearCalls).toHaveLength(1); - expect(mockMonitorRegistry.setRegisterCallback).toHaveBeenLastCalledWith( + expect(mockSetMonitorRegisterCallback).toHaveBeenLastCalledWith( + expect.anything(), undefined, ); }); @@ -843,9 +895,9 @@ describe('runNonInteractiveStreamJson', () => { await runNonInteractiveStreamJson(config, ''); - expect(mockMonitorRegistry.abortAll).toHaveBeenCalledTimes(2); - expect(mockBackgroundShellRegistry.abortAll).toHaveBeenCalledTimes(2); - expect(mockBackgroundTaskRegistry.abortAll).toHaveBeenCalledTimes(2); + expect(mockMonitorAbortAll).toHaveBeenCalledTimes(2); + expect(mockShellAbortAll).toHaveBeenCalledTimes(2); + expect(mockAgentAbortAll).toHaveBeenCalledTimes(2); }); it('aborts background registries on error shutdown', async () => { @@ -859,9 +911,9 @@ describe('runNonInteractiveStreamJson', () => { 'Stream error', ); - expect(mockMonitorRegistry.abortAll).toHaveBeenCalledTimes(2); - expect(mockBackgroundShellRegistry.abortAll).toHaveBeenCalledTimes(2); - expect(mockBackgroundTaskRegistry.abortAll).toHaveBeenCalledTimes(2); + expect(mockMonitorAbortAll).toHaveBeenCalledTimes(2); + expect(mockShellAbortAll).toHaveBeenCalledTimes(2); + expect(mockAgentAbortAll).toHaveBeenCalledTimes(2); }); it('runs final background cleanup after in-flight processing drains', async () => { @@ -870,13 +922,13 @@ describe('runNonInteractiveStreamJson', () => { let releaseProcessing: (() => void) | undefined; const callOrder: string[] = []; - mockMonitorRegistry.abortAll.mockImplementation(() => { + mockMonitorAbortAll.mockImplementation(() => { callOrder.push('monitor:abortAll'); }); - mockBackgroundShellRegistry.abortAll.mockImplementation(() => { + mockShellAbortAll.mockImplementation(() => { callOrder.push('background:abortAll'); }); - mockBackgroundTaskRegistry.abortAll.mockImplementation(() => { + mockAgentAbortAll.mockImplementation(() => { callOrder.push('agent:abortAll'); }); @@ -901,9 +953,9 @@ describe('runNonInteractiveStreamJson', () => { expect(releaseProcessing).toBeDefined(); }); - expect(mockMonitorRegistry.abortAll).toHaveBeenCalledTimes(1); - expect(mockBackgroundShellRegistry.abortAll).toHaveBeenCalledTimes(1); - expect(mockBackgroundTaskRegistry.abortAll).toHaveBeenCalledTimes(1); + expect(mockMonitorAbortAll).toHaveBeenCalledTimes(1); + expect(mockShellAbortAll).toHaveBeenCalledTimes(1); + expect(mockAgentAbortAll).toHaveBeenCalledTimes(1); expect(callOrder).toContain('run:start'); expect(callOrder).toContain('monitor:abortAll'); expect(callOrder).toContain('background:abortAll'); @@ -912,9 +964,9 @@ describe('runNonInteractiveStreamJson', () => { releaseProcessing?.(); await sessionPromise; - expect(mockMonitorRegistry.abortAll).toHaveBeenCalledTimes(2); - expect(mockBackgroundShellRegistry.abortAll).toHaveBeenCalledTimes(2); - expect(mockBackgroundTaskRegistry.abortAll).toHaveBeenCalledTimes(2); + expect(mockMonitorAbortAll).toHaveBeenCalledTimes(2); + expect(mockShellAbortAll).toHaveBeenCalledTimes(2); + expect(mockAgentAbortAll).toHaveBeenCalledTimes(2); expect(callOrder.slice(-4)).toEqual([ 'run:end', 'monitor:abortAll', @@ -930,13 +982,13 @@ describe('runNonInteractiveStreamJson', () => { const callOrder: string[] = []; const streamError = new Error('Stream error'); - mockMonitorRegistry.abortAll.mockImplementation(() => { + mockMonitorAbortAll.mockImplementation(() => { callOrder.push('monitor:abortAll'); }); - mockBackgroundShellRegistry.abortAll.mockImplementation(() => { + mockShellAbortAll.mockImplementation(() => { callOrder.push('background:abortAll'); }); - mockBackgroundTaskRegistry.abortAll.mockImplementation(() => { + mockAgentAbortAll.mockImplementation(() => { callOrder.push('agent:abortAll'); }); @@ -962,17 +1014,17 @@ describe('runNonInteractiveStreamJson', () => { expect(releaseProcessing).toBeDefined(); }); - expect(mockMonitorRegistry.abortAll).toHaveBeenCalledTimes(1); - expect(mockBackgroundShellRegistry.abortAll).toHaveBeenCalledTimes(1); - expect(mockBackgroundTaskRegistry.abortAll).toHaveBeenCalledTimes(1); + expect(mockMonitorAbortAll).toHaveBeenCalledTimes(1); + expect(mockShellAbortAll).toHaveBeenCalledTimes(1); + expect(mockAgentAbortAll).toHaveBeenCalledTimes(1); expect(callOrder).toContain('run:start'); releaseProcessing?.(); await expect(sessionPromise).rejects.toThrow('Stream error'); - expect(mockMonitorRegistry.abortAll).toHaveBeenCalledTimes(2); - expect(mockBackgroundShellRegistry.abortAll).toHaveBeenCalledTimes(2); - expect(mockBackgroundTaskRegistry.abortAll).toHaveBeenCalledTimes(2); + expect(mockMonitorAbortAll).toHaveBeenCalledTimes(2); + expect(mockShellAbortAll).toHaveBeenCalledTimes(2); + expect(mockAgentAbortAll).toHaveBeenCalledTimes(2); expect(callOrder.slice(-4)).toEqual([ 'run:end', 'monitor:abortAll', diff --git a/packages/cli/src/nonInteractive/session.ts b/packages/cli/src/nonInteractive/session.ts index f5e32aebbff..851d79797be 100644 --- a/packages/cli/src/nonInteractive/session.ts +++ b/packages/cli/src/nonInteractive/session.ts @@ -8,7 +8,15 @@ import type { Config, ConfigInitializeOptions, } from '@qwen-code/qwen-code-core'; -import { createDebugLogger, SendMessageType } from '@qwen-code/qwen-code-core'; +import { + agentAbortAll, + createDebugLogger, + monitorAbortAll, + SendMessageType, + setMonitorNotificationCallback, + setMonitorRegisterCallback, + shellAbortAll, +} from '@qwen-code/qwen-code-core'; import { StreamJsonInputReader } from './io/StreamJsonInputReader.js'; import { StreamJsonOutputAdapter } from './io/StreamJsonOutputAdapter.js'; import { ControlContext } from './control/ControlContext.js'; @@ -193,21 +201,23 @@ class Session { return; } - const registry = this.config.getMonitorRegistry(); - registry.setNotificationCallback((displayText, modelText, meta) => { - if (this.isShuttingDown || this.abortController.signal.aborted) { - return; - } - this.enqueueMonitorNotification({ - displayText, - modelText, - sdkNotification: { - task_id: meta.monitorId, - tool_use_id: meta.toolUseId, - status: meta.status, - }, - }); - }); + setMonitorNotificationCallback( + this.config.getTaskRegistry(), + (displayText, modelText, meta) => { + if (this.isShuttingDown || this.abortController.signal.aborted) { + return; + } + this.enqueueMonitorNotification({ + displayText, + modelText, + sdkNotification: { + task_id: meta.monitorId, + tool_use_id: meta.toolUseId, + status: meta.status, + }, + }); + }, + ); this.monitorNotificationsRegistered = true; } @@ -216,8 +226,7 @@ class Session { return; } - const registry = this.config.getMonitorRegistry(); - registry.setRegisterCallback((entry) => { + setMonitorRegisterCallback(this.config.getTaskRegistry(), (entry) => { if (this.isShuttingDown || this.abortController.signal.aborted) { return; } @@ -664,9 +673,10 @@ class Session { } private abortTaskRegistries(): void { - this.config.getMonitorRegistry().abortAll({ notify: false }); - this.config.getBackgroundShellRegistry().abortAll(); - this.config.getBackgroundTaskRegistry().abortAll(); + const registry = this.config.getTaskRegistry(); + monitorAbortAll(registry, { notify: false }); + shellAbortAll(registry); + agentAbortAll(registry); } private finishShutdown(): void { @@ -682,13 +692,13 @@ class Session { return; } - const registry = this.config.getMonitorRegistry(); + const taskRegistry = this.config.getTaskRegistry(); if (this.monitorNotificationsRegistered) { - registry.setNotificationCallback(undefined); + setMonitorNotificationCallback(taskRegistry, undefined); this.monitorNotificationsRegistered = false; } if (this.monitorRegistrationsRegistered) { - registry.setRegisterCallback(undefined); + setMonitorRegisterCallback(taskRegistry, undefined); this.monitorRegistrationsRegistered = false; } } diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index 01c642121fa..3bec40b406b 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -11,6 +11,71 @@ import type { SessionMetrics, } from '@qwen-code/qwen-code-core'; import type { CLIUserMessage } from './nonInteractive/types.js'; + +// Hoisted mocks for the kind-local helpers nonInteractiveCli now imports +// from core. The set*Callback mocks store the callback so tests can +// invoke it with synthetic notification/register events. The abortAll +// mocks let tests assert call counts and order across the three kinds. +const monitorRegisterCb = vi.hoisted(() => ({ + current: undefined as ((entry: unknown) => void) | undefined, +})); +const monitorNotificationCb = vi.hoisted(() => ({ + current: undefined as + | ((displayText: string, modelText: string, meta: unknown) => void) + | undefined, +})); +const agentRegisterCb = vi.hoisted(() => ({ + current: undefined as ((entry: unknown) => void) | undefined, +})); +const agentNotificationCb = vi.hoisted(() => ({ + current: undefined as + | ((displayText: string, modelText: string, meta: unknown) => void) + | undefined, +})); +const mockSetMonitorRegisterCallback = vi.hoisted(() => + vi.fn((cb: ((entry: unknown) => void) | undefined) => { + monitorRegisterCb.current = cb ?? undefined; + }), +); +const mockSetMonitorNotificationCallback = vi.hoisted(() => + vi.fn( + ( + cb: + | ((displayText: string, modelText: string, meta: unknown) => void) + | undefined, + ) => { + monitorNotificationCb.current = cb ?? undefined; + }, + ), +); +const mockSetAgentRegisterCallback = vi.hoisted(() => + vi.fn((cb: ((entry: unknown) => void) | undefined) => { + agentRegisterCb.current = cb ?? undefined; + }), +); +const mockSetAgentNotificationCallback = vi.hoisted(() => + vi.fn( + ( + cb: + | ((displayText: string, modelText: string, meta: unknown) => void) + | undefined, + ) => { + agentNotificationCb.current = cb ?? undefined; + }, + ), +); +const mockMonitorAbortAll = vi.hoisted(() => vi.fn()); +const mockShellAbortAll = vi.hoisted(() => vi.fn()); +const mockAgentAbortAll = vi.hoisted(() => vi.fn()); +const mockAgentHasUnfinalizedTasks = vi.hoisted(() => + vi.fn().mockReturnValue(false), +); +const mockGetRunningMonitorTasks = vi.hoisted(() => + vi.fn().mockReturnValue([]), +); + +import { vi, type Mock, type MockInstance } from 'vitest'; + import { executeToolCall, ToolErrorType, @@ -24,7 +89,6 @@ import { } from '@qwen-code/qwen-code-core'; import type { Part } from '@google/genai'; import { runNonInteractive } from './nonInteractiveCli.js'; -import { vi, type Mock, type MockInstance } from 'vitest'; import * as fs from 'node:fs/promises'; import * as os from 'node:os'; import * as path from 'node:path'; @@ -59,6 +123,15 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { uiTelemetryService: { getMetrics: vi.fn(), }, + setMonitorRegisterCallback: mockSetMonitorRegisterCallback, + setMonitorNotificationCallback: mockSetMonitorNotificationCallback, + setAgentRegisterCallback: mockSetAgentRegisterCallback, + setAgentNotificationCallback: mockSetAgentNotificationCallback, + monitorAbortAll: mockMonitorAbortAll, + shellAbortAll: mockShellAbortAll, + agentAbortAll: mockAgentAbortAll, + agentHasUnfinalizedTasks: mockAgentHasUnfinalizedTasks, + getRunningMonitorTasks: mockGetRunningMonitorTasks, }; }); @@ -75,18 +148,16 @@ describe('runNonInteractive', () => { let mockConfig: Config; let mockSettings: LoadedSettings; let mockToolRegistry: ToolRegistry; - let mockBackgroundTaskRegistry: { - setNotificationCallback: ReturnType; - setRegisterCallback: ReturnType; + let mockTaskRegistry: { getAll: ReturnType; - hasUnfinalizedTasks: ReturnType; - abortAll: ReturnType; - }; - let mockMonitorRegistry: { - setNotificationCallback: ReturnType; - setRegisterCallback: ReturnType; - getRunning: ReturnType; - abortAll: ReturnType; + getByKind: ReturnType; + get: ReturnType; + register: ReturnType; + update: ReturnType; + mutateSilent: ReturnType; + evict: ReturnType; + kill: ReturnType; + subscribe: ReturnType; }; let mockCoreExecuteToolCall: Mock; let mockShutdownTelemetry: Mock; @@ -134,20 +205,32 @@ describe('runNonInteractive', () => { getAllToolNames: vi.fn().mockReturnValue([]), } as unknown as ToolRegistry; - mockBackgroundTaskRegistry = { - setNotificationCallback: vi.fn(), - setRegisterCallback: vi.fn(), + mockTaskRegistry = { getAll: vi.fn().mockReturnValue([]), - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - abortAll: vi.fn(), - }; - - mockMonitorRegistry = { - setNotificationCallback: vi.fn(), - setRegisterCallback: vi.fn(), - getRunning: vi.fn().mockReturnValue([]), - abortAll: vi.fn(), + getByKind: vi.fn().mockReturnValue([]), + get: vi.fn(), + register: vi.fn(), + update: vi.fn(), + mutateSilent: vi.fn(), + evict: vi.fn(), + kill: vi.fn(), + subscribe: vi.fn(() => () => {}), }; + // Reset module-level helpers between tests so callback storage and + // call counts don't leak across the suite. + mockSetMonitorRegisterCallback.mockClear(); + mockSetMonitorNotificationCallback.mockClear(); + mockSetAgentRegisterCallback.mockClear(); + mockSetAgentNotificationCallback.mockClear(); + mockMonitorAbortAll.mockClear(); + mockShellAbortAll.mockClear(); + mockAgentAbortAll.mockClear(); + mockAgentHasUnfinalizedTasks.mockReset().mockReturnValue(false); + mockGetRunningMonitorTasks.mockReset().mockReturnValue([]); + monitorRegisterCb.current = undefined; + monitorNotificationCb.current = undefined; + agentRegisterCb.current = undefined; + agentNotificationCb.current = undefined; mockGetDebugResponses = vi.fn(() => []); @@ -206,10 +289,7 @@ describe('runNonInteractive', () => { setModelInvocableCommandsExecutor: vi.fn(), getAutoSkillEnabled: vi.fn().mockReturnValue(false), getDisabledSlashCommands: vi.fn().mockReturnValue([]), - getBackgroundTaskRegistry: vi - .fn() - .mockReturnValue(mockBackgroundTaskRegistry), - getMonitorRegistry: vi.fn().mockReturnValue(mockMonitorRegistry), + getTaskRegistry: vi.fn().mockReturnValue(mockTaskRegistry), // Phase C: headless --resume reads the resumed session + sidecar to // restore worktree context. These tests don't exercise resume, so // return undefined to short-circuit the helper. @@ -1318,7 +1398,7 @@ describe('runNonInteractive', () => { ) => void) | undefined; - mockMonitorRegistry.setNotificationCallback.mockImplementation((cb) => { + mockSetMonitorNotificationCallback.mockImplementation((_registry, cb) => { monitorNotificationCallback = cb ?? undefined; if (!cb) { return; @@ -1330,7 +1410,7 @@ describe('runNonInteractive', () => { eventCount: 1, }); }); - mockMonitorRegistry.abortAll.mockImplementation(() => { + mockMonitorAbortAll.mockImplementation(() => { monitorNotificationCallback?.( 'Monitor "logs" was cancelled.', cancelledXml, @@ -1412,7 +1492,7 @@ describe('runNonInteractive', () => { const resultIndex = envelopes.findIndex((env) => env.type === 'result'); expect(cancelledNotificationIndex).toBeGreaterThanOrEqual(0); expect(resultIndex).toBeGreaterThan(cancelledNotificationIndex); - expect(mockMonitorRegistry.abortAll).toHaveBeenCalledTimes(1); + expect(mockMonitorAbortAll).toHaveBeenCalledTimes(1); expect(envelopes.at(-1)).toMatchObject({ type: 'result', is_error: false, @@ -1530,7 +1610,7 @@ describe('runNonInteractive', () => { ) => void) | undefined; - mockMonitorRegistry.setNotificationCallback.mockImplementation((cb) => { + mockSetMonitorNotificationCallback.mockImplementation((_registry, cb) => { monitorNotificationCallback = cb ?? undefined; if (!cb) { return; @@ -1542,7 +1622,7 @@ describe('runNonInteractive', () => { eventCount: 1, }); }); - mockMonitorRegistry.abortAll.mockImplementation(() => { + mockMonitorAbortAll.mockImplementation(() => { monitorNotificationCallback?.( 'Monitor "logs" was cancelled.', cancelledXml, @@ -1642,7 +1722,7 @@ describe('runNonInteractive', () => { let keepBackgroundTaskOpen = true; let lateMonitorEventEmitted = false; - mockBackgroundTaskRegistry.hasUnfinalizedTasks.mockImplementation(() => { + mockAgentHasUnfinalizedTasks.mockImplementation(() => { if (keepBackgroundTaskOpen && !lateMonitorEventEmitted) { lateMonitorEventEmitted = true; monitorNotificationCallback?.( @@ -1696,7 +1776,7 @@ describe('runNonInteractive', () => { ) => void) | undefined; - mockMonitorRegistry.setNotificationCallback.mockImplementation((cb) => { + mockSetMonitorNotificationCallback.mockImplementation((_registry, cb) => { monitorNotificationCallback = cb ?? undefined; if (!cb) { return; @@ -1708,7 +1788,7 @@ describe('runNonInteractive', () => { eventCount: 1, }); }); - mockMonitorRegistry.abortAll.mockImplementation(() => { + mockMonitorAbortAll.mockImplementation(() => { monitorNotificationCallback?.( 'Monitor "logs" was cancelled.', cancelledXml, @@ -2478,17 +2558,10 @@ describe('runNonInteractive', () => { (mockConfig.getOutputFormat as Mock).mockReturnValue(OutputFormat.JSON); setupMetricsMock(); - // Spy on the registry returned by getBackgroundTaskRegistry so we can - // assert abortAll() is called as part of the deterministic shutdown - // contract for structured-output mode. - const abortAllSpy = vi.fn(); - (mockConfig.getBackgroundTaskRegistry as Mock).mockReturnValue({ - setNotificationCallback: vi.fn(), - setRegisterCallback: vi.fn(), - getAll: vi.fn().mockReturnValue([]), - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - abortAll: abortAllSpy, - }); + // Spy on agentAbortAll so we can assert it is called as part of the + // deterministic shutdown contract for structured-output mode. + const abortAllSpy = mockAgentAbortAll; + abortAllSpy.mockClear(); const writes: string[] = []; processStdoutSpy.mockImplementation((chunk: string | Uint8Array) => { @@ -2597,14 +2670,8 @@ describe('runNonInteractive', () => { (mockConfig.getOutputFormat as Mock).mockReturnValue(OutputFormat.JSON); setupMetricsMock(); - const abortAllSpy = vi.fn(); - (mockConfig.getBackgroundTaskRegistry as Mock).mockReturnValue({ - setNotificationCallback: vi.fn(), - setRegisterCallback: vi.fn(), - getAll: vi.fn().mockReturnValue([]), - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - abortAll: abortAllSpy, - }); + const abortAllSpy = mockAgentAbortAll; + abortAllSpy.mockClear(); const writes: string[] = []; processStdoutSpy.mockImplementation((chunk: string | Uint8Array) => { @@ -2721,14 +2788,8 @@ describe('runNonInteractive', () => { (mockConfig.getOutputFormat as Mock).mockReturnValue(OutputFormat.JSON); setupMetricsMock(); - const abortAllSpy = vi.fn(); - (mockConfig.getBackgroundTaskRegistry as Mock).mockReturnValue({ - setNotificationCallback: vi.fn(), - setRegisterCallback: vi.fn(), - getAll: vi.fn().mockReturnValue([]), - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - abortAll: abortAllSpy, - }); + const abortAllSpy = mockAgentAbortAll; + abortAllSpy.mockClear(); const writes: string[] = []; processStdoutSpy.mockImplementation((chunk: string | Uint8Array) => { @@ -3142,7 +3203,7 @@ describe('runNonInteractive', () => { 'Monitor emitted event #1.\n' + 'ready\n' + ''; - mockMonitorRegistry.setNotificationCallback.mockImplementation((cb) => { + mockSetMonitorNotificationCallback.mockImplementation((_registry, cb) => { if (!cb) return; cb('Monitor "logs" event #1: ready', notificationXml, { monitorId: 'mon_1', @@ -3244,12 +3305,13 @@ describe('runNonInteractive', () => { ); setupMetricsMock(); - const abortAllSpy = vi.fn(); - // Returns true once, then false. After abortAll() is called the + const abortAllSpy = mockAgentAbortAll; + abortAllSpy.mockClear(); + // Returns true once, then false. After agentAbortAll() is called the // holdback's `while` body executes one iteration of `setTimeout(50)` // and re-checks; on the second call we report tasks finalized. let unfinalizedCalls = 0; - const hasUnfinalizedTasksSpy = vi.fn(() => { + mockAgentHasUnfinalizedTasks.mockReset().mockImplementation(() => { unfinalizedCalls++; return unfinalizedCalls === 1; }); @@ -3269,29 +3331,23 @@ describe('runNonInteractive', () => { }, ) => void) | null = null; - (mockConfig.getBackgroundTaskRegistry as Mock).mockReturnValue({ - setNotificationCallback: vi.fn((cb) => { - notificationCallback = cb; - }), - setRegisterCallback: vi.fn(), - getAll: vi.fn().mockReturnValue([]), - hasUnfinalizedTasks: hasUnfinalizedTasksSpy, - abortAll: vi.fn(() => { - abortAllSpy(); - // The natural cancel-handler enqueues the terminal - // task_notification synchronously when abortAll is invoked. - // Fire the captured callback immediately so it lands in - // localQueue before the holdback flush runs. - notificationCallback?.( - 'Agent cancelled: bg-task-1', - 'Agent bg-task-1 was cancelled', - { - agentId: 'bg-task-1', - toolUseId: 'tool-bg-1', - status: 'cancelled' as never, - }, - ); - }), + mockSetAgentNotificationCallback.mockImplementation((_registry, cb) => { + notificationCallback = cb ?? null; + }); + abortAllSpy.mockImplementation(() => { + // The natural cancel-handler enqueues the terminal + // task_notification synchronously when agentAbortAll is invoked. + // Fire the captured callback immediately so it lands in + // localQueue before the holdback flush runs. + notificationCallback?.( + 'Agent cancelled: bg-task-1', + 'Agent bg-task-1 was cancelled', + { + agentId: 'bg-task-1', + toolUseId: 'tool-bg-1', + status: 'cancelled' as never, + }, + ); }); const writes: string[] = []; @@ -3382,13 +3438,7 @@ describe('runNonInteractive', () => { (mockConfig.getOutputFormat as Mock).mockReturnValue(OutputFormat.TEXT); setupMetricsMock(); - (mockConfig.getBackgroundTaskRegistry as Mock).mockReturnValue({ - setNotificationCallback: vi.fn(), - setRegisterCallback: vi.fn(), - getAll: vi.fn().mockReturnValue([]), - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - abortAll: vi.fn(), - }); + mockAgentAbortAll.mockClear(); const writes: string[] = []; processStdoutSpy.mockImplementation((chunk: string | Uint8Array) => { diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 72389d6c0ac..6d8adb95a5f 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -12,7 +12,14 @@ import type { import { isSlashCommand } from './ui/utils/commandUtils.js'; import type { LoadedSettings } from './config/settings.js'; import { + agentAbortAll, + agentHasUnfinalizedTasks, executeToolCall, + monitorAbortAll, + setAgentNotificationCallback, + setAgentRegisterCallback, + setMonitorNotificationCallback, + setMonitorRegisterCallback, shutdownTelemetry, isTelemetrySdkInitialized, GeminiEventType, @@ -299,7 +306,7 @@ export async function runNonInteractive( return; oneShotMonitorsFinalized = true; captureMonitorTurnsInLocalQueue = false; - config.getMonitorRegistry().abortAll(); + monitorAbortAll(config.getTaskRegistry()); flushQueuedNotificationsToSdk(sdkOnlyMonitorQueue); }; @@ -468,8 +475,8 @@ export async function runNonInteractive( // Register the callback early so background agents launched during the main // tool-call chain can push completions onto the queue. - const registry = config.getBackgroundTaskRegistry(); - registry.setNotificationCallback((displayText, modelText, meta) => { + const registry = config.getTaskRegistry(); + setAgentNotificationCallback(registry, (displayText, modelText, meta) => { localQueue.push({ displayText, modelText, @@ -489,7 +496,7 @@ export async function runNonInteractive( }); }); - registry.setRegisterCallback((entry) => { + setAgentRegisterCallback(registry, (entry) => { adapter.emitSystemMessage('task_started', { task_id: entry.agentId, tool_use_id: entry.toolUseId, @@ -498,14 +505,14 @@ export async function runNonInteractive( }); }); - const monitorRegistry = config.getMonitorRegistry(); if (options.captureMonitorNotifications !== false) { // One-shot headless runs capture monitor notifications locally so any // events already emitted before exit can be surfaced to the SDK/model. // Persistent stream-json sessions own this callback at the Session // layer instead, so future monitor events can continue after the // originating turn has already completed. - monitorRegistry.setNotificationCallback( + setMonitorNotificationCallback( + registry, (displayText, modelText, meta) => { const queueItem = { displayText, @@ -529,7 +536,7 @@ export async function runNonInteractive( } if (options.captureMonitorRegistrations !== false) { - monitorRegistry.setRegisterCallback((entry) => { + setMonitorRegisterCallback(registry, (entry) => { adapter.emitSystemMessage('task_started', { task_id: entry.monitorId, tool_use_id: entry.toolUseId, @@ -564,8 +571,8 @@ export async function runNonInteractive( // no-op), so unconditional invocation is safe even when the drain // path already finalized monitors before reaching here. const emitStructuredSuccess = async (): Promise<0> => { - registry.abortAll(); - // `abortAll()` marks each task `cancelled` synchronously, but + agentAbortAll(registry); + // `agentAbortAll()` marks each task `cancelled` synchronously, but // the matching `task_notification` is emitted later by the // task's natural handler. Hold back briefly (capped at // STRUCTURED_SHUTDOWN_HOLDBACK_MS) so consumers see every @@ -575,7 +582,7 @@ export async function runNonInteractive( const holdbackDeadline = Date.now() + STRUCTURED_SHUTDOWN_HOLDBACK_MS; while ( Date.now() < holdbackDeadline && - registry.hasUnfinalizedTasks() + agentHasUnfinalizedTasks(registry) ) { await new Promise((r) => setTimeout(r, 50)); } @@ -1119,7 +1126,7 @@ export async function runNonInteractive( // silently convert a cancellation into a completion. while (true) { if (abortController.signal.aborted) { - registry.abortAll(); + agentAbortAll(registry); // Flush queued terminal notifications before routeAbort // exits so stream-json consumers always see a task_notification // paired with every task_started. @@ -1143,7 +1150,7 @@ export async function runNonInteractive( // paired with one. Monitors are different: they intentionally // continue in the background, so final result emission is not // gated on monitor lifetime. - if (!registry.hasUnfinalizedTasks() && localQueue.length === 0) + if (!agentHasUnfinalizedTasks(registry) && localQueue.length === 0) break; await new Promise((r) => setTimeout(r, 100)); } @@ -1305,21 +1312,20 @@ export async function runNonInteractive( // daemon, SDK) that reuse a single process across many runs. budgetEnforcer.stop(); - const reg = config.getBackgroundTaskRegistry(); - reg.setNotificationCallback(undefined); - reg.setRegisterCallback(undefined); - const monReg = config.getMonitorRegistry(); + const taskRegistry = config.getTaskRegistry(); + setAgentNotificationCallback(taskRegistry, undefined); + setAgentRegisterCallback(taskRegistry, undefined); // In one-shot (non-Session) runs, abort all running monitors so their // piped stdio refs don't keep the Node event loop alive after the result // is emitted. Session runs manage monitor lifecycle independently. if (options.captureMonitorNotifications !== false) { if (!oneShotMonitorsFinalized) { - monReg.abortAll({ notify: false }); + monitorAbortAll(taskRegistry, { notify: false }); } - monReg.setNotificationCallback(undefined); + setMonitorNotificationCallback(taskRegistry, undefined); } if (options.captureMonitorRegistrations !== false) { - monReg.setRegisterCallback(undefined); + setMonitorRegisterCallback(taskRegistry, undefined); } process.stdout.removeListener('error', stdoutErrorHandler); diff --git a/packages/cli/src/ui/commands/clearCommand.test.ts b/packages/cli/src/ui/commands/clearCommand.test.ts index 92699e514cb..47de1f09e1e 100644 --- a/packages/cli/src/ui/commands/clearCommand.test.ts +++ b/packages/cli/src/ui/commands/clearCommand.test.ts @@ -5,22 +5,43 @@ */ import { vi, describe, it, expect, beforeEach } from 'vitest'; -import { clearCommand } from './clearCommand.js'; import { type CommandContext } from './types.js'; import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; import { SessionEndReason } from '@qwen-code/qwen-code-core'; -// Mock the telemetry service +// Hoisted spies for the kind-local abort helpers and the +// resetBackgroundStateForSessionSwitch helper. The mock factory below +// runs before the test body, so any spy referenced inside it must be +// hoisted via vi.hoisted(). The clear flow calls each of these in a +// fixed order; tests assert that order via invocationCallOrder. +const mockAgentAbortAll = vi.hoisted(() => vi.fn()); +const mockMonitorAbortAll = vi.hoisted(() => vi.fn()); +const mockShellAbortAll = vi.hoisted(() => vi.fn()); + vi.mock('@qwen-code/qwen-code-core', async () => { const actual = await vi.importActual('@qwen-code/qwen-code-core'); return { ...actual, + agentAbortAll: mockAgentAbortAll, + monitorAbortAll: mockMonitorAbortAll, + shellAbortAll: mockShellAbortAll, uiTelemetryService: { reset: vi.fn(), }, }; }); +const mockResetBackgroundStateForSessionSwitch = vi.hoisted(() => vi.fn()); +const mockHasBlockingBackgroundWork = vi.hoisted(() => + vi.fn().mockReturnValue(false), +); +vi.mock('../utils/backgroundWorkUtils.js', () => ({ + hasBlockingBackgroundWork: mockHasBlockingBackgroundWork, + resetBackgroundStateForSessionSwitch: + mockResetBackgroundStateForSessionSwitch, +})); + +import { clearCommand } from './clearCommand.js'; import type { GeminiClient } from '@qwen-code/qwen-code-core'; describe('clearCommand', () => { @@ -30,12 +51,6 @@ describe('clearCommand', () => { let mockFireSessionEndEvent: ReturnType; let mockFireSessionStartEvent: ReturnType; let mockGetHookSystem: ReturnType; - let mockAbortBackgroundTasks: ReturnType; - let mockAbortMonitors: ReturnType; - let mockAbortBackgroundShells: ReturnType; - let mockResetBackgroundTasks: ReturnType; - let mockResetMonitors: ReturnType; - let mockResetBackgroundShells: ReturnType; beforeEach(() => { mockResetChat = vi.fn().mockResolvedValue(undefined); @@ -46,14 +61,20 @@ describe('clearCommand', () => { fireSessionEndEvent: mockFireSessionEndEvent, fireSessionStartEvent: mockFireSessionStartEvent, }); - mockAbortBackgroundTasks = vi.fn(); - mockAbortMonitors = vi.fn(); - mockAbortBackgroundShells = vi.fn(); - mockResetBackgroundTasks = vi.fn(); - mockResetMonitors = vi.fn(); - mockResetBackgroundShells = vi.fn(); vi.clearAllMocks(); + const stubRegistry = { + getAll: () => [], + getByKind: () => [], + get: vi.fn(), + register: vi.fn(), + update: vi.fn(), + mutateSilent: vi.fn(), + evict: vi.fn(), + kill: vi.fn(), + subscribe: vi.fn(() => () => {}), + }; + mockContext = createMockCommandContext({ services: { config: { @@ -61,17 +82,7 @@ describe('clearCommand', () => { ({ resetChat: mockResetChat, }) as unknown as GeminiClient, - getBackgroundTaskRegistry: vi.fn().mockReturnValue({ - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - reset: mockResetBackgroundTasks, - abortAll: mockAbortBackgroundTasks, - }), - getBackgroundShellRegistry: vi.fn().mockReturnValue({ - getAll: vi.fn().mockReturnValue([]), - hasRunningEntries: vi.fn().mockReturnValue(false), - reset: mockResetBackgroundShells, - abortAll: mockAbortBackgroundShells, - }), + getTaskRegistry: vi.fn().mockReturnValue(stubRegistry), startNewSession: mockStartNewSession, getHookSystem: mockGetHookSystem, getDebugLogger: () => ({ @@ -80,11 +91,6 @@ describe('clearCommand', () => { getModel: () => 'test-model', getToolRegistry: () => undefined, getApprovalMode: () => 'default', - getMonitorRegistry: () => ({ - getRunning: vi.fn().mockReturnValue([]), - abortAll: mockAbortMonitors, - reset: mockResetMonitors, - }), }, }, session: { @@ -111,13 +117,6 @@ describe('clearCommand', () => { ); expect(mockResetChat).toHaveBeenCalledTimes(1); expect(mockContext.ui.clear).toHaveBeenCalledTimes(1); - - // Check that all expected operations were called - expect(mockContext.ui.setDebugMessage).toHaveBeenCalled(); - expect(mockStartNewSession).toHaveBeenCalled(); - expect(mockContext.session.startNewSession).toHaveBeenCalled(); - expect(mockResetChat).toHaveBeenCalled(); - expect(mockContext.ui.clear).toHaveBeenCalled(); }); it('should fire SessionEnd event before clearing', async () => { @@ -131,6 +130,9 @@ describe('clearCommand', () => { expect(mockFireSessionEndEvent).toHaveBeenCalledWith( SessionEndReason.Clear, ); + // PR #4115 moved SessionStart firing out of `/clear` so the hook's + // additionalContext lands inside the new chat session rather than + // before it. The clear command now only fires SessionEnd. expect(mockFireSessionStartEvent).not.toHaveBeenCalled(); }); @@ -141,33 +143,26 @@ describe('clearCommand', () => { await clearCommand.action(mockContext, ''); - expect(mockAbortBackgroundTasks).toHaveBeenCalledWith({ notify: false }); - expect(mockAbortMonitors).toHaveBeenCalledWith({ notify: false }); - expect(mockAbortBackgroundShells).toHaveBeenCalledTimes(1); - expect(mockAbortBackgroundTasks.mock.invocationCallOrder[0]).toBeLessThan( - mockStartNewSession.mock.invocationCallOrder[0], - ); - expect(mockAbortMonitors.mock.invocationCallOrder[0]).toBeLessThan( - mockStartNewSession.mock.invocationCallOrder[0], - ); - expect(mockAbortBackgroundShells.mock.invocationCallOrder[0]).toBeLessThan( - mockResetBackgroundShells.mock.invocationCallOrder[0], - ); - expect(mockAbortBackgroundTasks.mock.invocationCallOrder[0]).toBeLessThan( - mockResetBackgroundTasks.mock.invocationCallOrder[0], - ); - expect(mockAbortMonitors.mock.invocationCallOrder[0]).toBeLessThan( - mockResetMonitors.mock.invocationCallOrder[0], - ); - expect(mockResetBackgroundShells.mock.invocationCallOrder[0]).toBeLessThan( - mockStartNewSession.mock.invocationCallOrder[0], - ); - expect(mockResetBackgroundTasks.mock.invocationCallOrder[0]).toBeLessThan( - mockStartNewSession.mock.invocationCallOrder[0], - ); - expect(mockResetMonitors.mock.invocationCallOrder[0]).toBeLessThan( - mockStartNewSession.mock.invocationCallOrder[0], - ); + expect(mockAgentAbortAll).toHaveBeenCalledWith(expect.anything(), { + notify: false, + }); + expect(mockMonitorAbortAll).toHaveBeenCalledWith(expect.anything(), { + notify: false, + }); + expect(mockShellAbortAll).toHaveBeenCalledTimes(1); + expect(mockResetBackgroundStateForSessionSwitch).toHaveBeenCalledTimes(1); + + const agentAbort = mockAgentAbortAll.mock.invocationCallOrder[0]; + const monitorAbort = mockMonitorAbortAll.mock.invocationCallOrder[0]; + const shellAbort = mockShellAbortAll.mock.invocationCallOrder[0]; + const reset = + mockResetBackgroundStateForSessionSwitch.mock.invocationCallOrder[0]; + const newSession = mockStartNewSession.mock.invocationCallOrder[0]; + + expect(agentAbort).toBeLessThan(newSession); + expect(monitorAbort).toBeLessThan(newSession); + expect(shellAbort).toBeLessThan(newSession); + expect(reset).toBeLessThan(newSession); }); it('should handle hook errors gracefully and continue execution', async () => { @@ -190,139 +185,74 @@ describe('clearCommand', () => { expect(mockContext.ui.clear).toHaveBeenCalledTimes(1); }); - it('should clear UI before resetChat for immediate responsiveness', async () => { + it('should handle missing hook system gracefully', async () => { if (!clearCommand.action) { throw new Error('clearCommand must have an action.'); } - const callOrder: string[] = []; - (mockContext.ui.clear as ReturnType).mockImplementation( - () => { - callOrder.push('ui.clear'); - }, - ); - mockResetChat.mockImplementation(async () => { - callOrder.push('resetChat'); - }); + mockGetHookSystem.mockReturnValue(undefined); await clearCommand.action(mockContext, ''); - // ui.clear should be called before resetChat for immediate UI feedback - const clearIndex = callOrder.indexOf('ui.clear'); - const resetIndex = callOrder.indexOf('resetChat'); - expect(clearIndex).toBeGreaterThanOrEqual(0); - expect(resetIndex).toBeGreaterThanOrEqual(0); - expect(clearIndex).toBeLessThan(resetIndex); - }); - - it('should not await hook events (fire-and-forget)', async () => { - if (!clearCommand.action) { - throw new Error('clearCommand must have an action.'); - } - - // Make hooks take a long time - they should not block - let sessionEndResolved = false; - let sessionStartResolved = false; - mockFireSessionEndEvent.mockImplementation( - () => - new Promise((resolve) => { - setTimeout(() => { - sessionEndResolved = true; - resolve(undefined); - }, 5000); - }), - ); - mockFireSessionStartEvent.mockImplementation( - () => - new Promise((resolve) => { - setTimeout(() => { - sessionStartResolved = true; - resolve(undefined); - }, 5000); - }), - ); - - await clearCommand.action(mockContext, ''); - - // The action should complete immediately without waiting for hooks - expect(mockContext.ui.clear).toHaveBeenCalledTimes(1); - expect(mockResetChat).toHaveBeenCalledTimes(1); - // SessionEnd hook should have been called but not necessarily resolved - expect(mockFireSessionEndEvent).toHaveBeenCalled(); + expect(mockFireSessionEndEvent).not.toHaveBeenCalled(); expect(mockFireSessionStartEvent).not.toHaveBeenCalled(); - // SessionEnd hook should NOT have resolved yet since it has a 5s timeout - expect(sessionEndResolved).toBe(false); - expect(sessionStartResolved).toBe(false); + expect(mockStartNewSession).toHaveBeenCalledTimes(1); + expect(mockResetChat).toHaveBeenCalledTimes(1); }); - it('should not attempt to reset chat if config service is not available', async () => { + it('should handle missing config gracefully', async () => { if (!clearCommand.action) { throw new Error('clearCommand must have an action.'); } - const nullConfigContext = createMockCommandContext({ - services: { - config: null, - }, - session: { - startNewSession: vi.fn(), - }, + const ctxNoConfig = createMockCommandContext({ + services: {}, }); - await clearCommand.action(nullConfigContext, ''); + const result = await clearCommand.action(ctxNoConfig, ''); - expect(nullConfigContext.ui.setDebugMessage).toHaveBeenCalledWith( - 'Starting a new session and clearing.', - ); - expect(mockResetChat).not.toHaveBeenCalled(); - expect(nullConfigContext.ui.clear).toHaveBeenCalledTimes(1); + expect(result).toBeUndefined(); + expect(mockStartNewSession).not.toHaveBeenCalled(); }); describe('non-interactive mode', () => { let nonInteractiveContext: ReturnType; beforeEach(() => { + const stubRegistry = { + getAll: () => [], + getByKind: () => [], + get: vi.fn(), + register: vi.fn(), + update: vi.fn(), + mutateSilent: vi.fn(), + evict: vi.fn(), + kill: vi.fn(), + subscribe: vi.fn(() => () => {}), + }; nonInteractiveContext = createMockCommandContext({ executionMode: 'non_interactive', services: { config: { - getHookSystem: mockGetHookSystem, - getBackgroundTaskRegistry: vi.fn().mockReturnValue({ - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - reset: mockResetBackgroundTasks, - abortAll: mockAbortBackgroundTasks, - }), - getBackgroundShellRegistry: vi.fn().mockReturnValue({ - getAll: vi.fn().mockReturnValue([]), - hasRunningEntries: vi.fn().mockReturnValue(false), - reset: mockResetBackgroundShells, - abortAll: mockAbortBackgroundShells, - }), + getGeminiClient: () => + ({ resetChat: mockResetChat }) as unknown as GeminiClient, + getTaskRegistry: vi.fn().mockReturnValue(stubRegistry), startNewSession: mockStartNewSession, - getGeminiClient: vi.fn().mockReturnValue({ - resetChat: mockResetChat, - } as unknown as GeminiClient), - getModel: vi.fn().mockReturnValue('test-model'), - getApprovalMode: vi.fn().mockReturnValue('default'), - getToolRegistry: vi.fn().mockReturnValue({ - getAllTools: vi.fn().mockReturnValue([]), - }), - getMonitorRegistry: () => ({ - getRunning: vi.fn().mockReturnValue([]), - abortAll: mockAbortMonitors, - reset: mockResetMonitors, - }), + getHookSystem: mockGetHookSystem, + getDebugLogger: () => ({ warn: vi.fn() }), + getModel: () => 'test-model', + getToolRegistry: () => undefined, + getApprovalMode: () => 'default', }, }, - session: { - startNewSession: vi.fn(), - }, + session: { startNewSession: vi.fn() }, }); }); - it('should return context boundary message in non-interactive mode', async () => { + it('returns the context-cleared message and resets chat when not blocked', async () => { if (!clearCommand.action) throw new Error('clearCommand must have an action.'); + mockHasBlockingBackgroundWork.mockReturnValue(false); const result = await clearCommand.action(nonInteractiveContext, ''); @@ -331,23 +261,7 @@ describe('clearCommand', () => { messageType: 'info', content: 'Context cleared. Previous messages are no longer in context.', }); - }); - - it('should still call resetChat in non-interactive mode', async () => { - if (!clearCommand.action) - throw new Error('clearCommand must have an action.'); - - await clearCommand.action(nonInteractiveContext, ''); - expect(mockResetChat).toHaveBeenCalledTimes(1); - }); - - it('should still fire SessionEnd in non-interactive mode', async () => { - if (!clearCommand.action) - throw new Error('clearCommand must have an action.'); - - await clearCommand.action(nonInteractiveContext, ''); - expect(mockFireSessionEndEvent).toHaveBeenCalledWith( SessionEndReason.Clear, ); @@ -357,155 +271,9 @@ describe('clearCommand', () => { it('blocks session clearing while background work is still running', async () => { if (!clearCommand.action) throw new Error('clearCommand must have an action.'); + mockHasBlockingBackgroundWork.mockReturnValue(true); - const blockedContext = createMockCommandContext({ - executionMode: 'non_interactive', - services: { - config: { - getBackgroundTaskRegistry: vi.fn().mockReturnValue({ - hasUnfinalizedTasks: vi.fn().mockReturnValue(true), - reset: vi.fn(), - }), - getBackgroundShellRegistry: vi.fn().mockReturnValue({ - getAll: vi.fn().mockReturnValue([]), - hasRunningEntries: vi.fn().mockReturnValue(false), - reset: vi.fn(), - }), - getMonitorRegistry: vi.fn().mockReturnValue({ - getRunning: vi.fn().mockReturnValue([]), - reset: vi.fn(), - }), - getHookSystem: mockGetHookSystem, - startNewSession: mockStartNewSession, - getGeminiClient: vi.fn().mockReturnValue({ - resetChat: mockResetChat, - } as unknown as GeminiClient), - getModel: vi.fn().mockReturnValue('test-model'), - getApprovalMode: vi.fn().mockReturnValue('default'), - getToolRegistry: vi.fn().mockReturnValue({ - getAllTools: vi.fn().mockReturnValue([]), - }), - getDebugLogger: vi.fn().mockReturnValue({ warn: vi.fn() }), - }, - }, - session: { - startNewSession: vi.fn(), - }, - }); - - const result = await clearCommand.action(blockedContext, ''); - - expect(result).toEqual({ - type: 'message', - messageType: 'error', - content: - "Stop the current session's running background tasks before starting a new session.", - }); - expect(mockStartNewSession).not.toHaveBeenCalled(); - expect(mockResetChat).not.toHaveBeenCalled(); - }); - - it('blocks session clearing while a monitor is still running', async () => { - if (!clearCommand.action) - throw new Error('clearCommand must have an action.'); - - const blockedContext = createMockCommandContext({ - executionMode: 'non_interactive', - services: { - config: { - getBackgroundTaskRegistry: vi.fn().mockReturnValue({ - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - reset: vi.fn(), - }), - getBackgroundShellRegistry: vi.fn().mockReturnValue({ - getAll: vi.fn().mockReturnValue([]), - hasRunningEntries: vi.fn().mockReturnValue(false), - reset: vi.fn(), - }), - getMonitorRegistry: vi.fn().mockReturnValue({ - getRunning: vi.fn().mockReturnValue([ - { - monitorId: 'mon_123', - status: 'running', - }, - ]), - reset: vi.fn(), - }), - getHookSystem: mockGetHookSystem, - startNewSession: mockStartNewSession, - getGeminiClient: vi.fn().mockReturnValue({ - resetChat: mockResetChat, - } as unknown as GeminiClient), - getModel: vi.fn().mockReturnValue('test-model'), - getApprovalMode: vi.fn().mockReturnValue('default'), - getToolRegistry: vi.fn().mockReturnValue({ - getAllTools: vi.fn().mockReturnValue([]), - }), - getDebugLogger: vi.fn().mockReturnValue({ warn: vi.fn() }), - }, - }, - session: { - startNewSession: vi.fn(), - }, - }); - - const result = await clearCommand.action(blockedContext, ''); - - expect(result).toEqual({ - type: 'message', - messageType: 'error', - content: - "Stop the current session's running background tasks before starting a new session.", - }); - expect(mockStartNewSession).not.toHaveBeenCalled(); - expect(mockResetChat).not.toHaveBeenCalled(); - }); - - it('blocks session clearing while a background shell is still running', async () => { - if (!clearCommand.action) - throw new Error('clearCommand must have an action.'); - - const blockedContext = createMockCommandContext({ - executionMode: 'non_interactive', - services: { - config: { - getBackgroundTaskRegistry: vi.fn().mockReturnValue({ - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - reset: vi.fn(), - }), - getBackgroundShellRegistry: vi.fn().mockReturnValue({ - getAll: vi.fn().mockReturnValue([ - { - shellId: 'shell_123', - status: 'running', - }, - ]), - hasRunningEntries: vi.fn().mockReturnValue(true), - reset: vi.fn(), - }), - getMonitorRegistry: vi.fn().mockReturnValue({ - getRunning: vi.fn().mockReturnValue([]), - reset: vi.fn(), - }), - getHookSystem: mockGetHookSystem, - startNewSession: mockStartNewSession, - getGeminiClient: vi.fn().mockReturnValue({ - resetChat: mockResetChat, - } as unknown as GeminiClient), - getModel: vi.fn().mockReturnValue('test-model'), - getApprovalMode: vi.fn().mockReturnValue('default'), - getToolRegistry: vi.fn().mockReturnValue({ - getAllTools: vi.fn().mockReturnValue([]), - }), - getDebugLogger: vi.fn().mockReturnValue({ warn: vi.fn() }), - }, - }, - session: { - startNewSession: vi.fn(), - }, - }); - - const result = await clearCommand.action(blockedContext, ''); + const result = await clearCommand.action(nonInteractiveContext, ''); expect(result).toEqual({ type: 'message', diff --git a/packages/cli/src/ui/commands/clearCommand.ts b/packages/cli/src/ui/commands/clearCommand.ts index 41bfb695cfd..f801499d079 100644 --- a/packages/cli/src/ui/commands/clearCommand.ts +++ b/packages/cli/src/ui/commands/clearCommand.ts @@ -8,6 +8,9 @@ import type { SlashCommand } from './types.js'; import { CommandKind } from './types.js'; import { t } from '../../i18n/index.js'; import { + agentAbortAll, + monitorAbortAll, + shellAbortAll, uiTelemetryService, SessionEndReason, ToolNames, @@ -53,9 +56,10 @@ export const clearCommand: SlashCommand = { // Abort old-session async work before creating the new session so // cancellation notifications cannot leak across the reset boundary. - config.getBackgroundTaskRegistry().abortAll({ notify: false }); - config.getMonitorRegistry().abortAll({ notify: false }); - config.getBackgroundShellRegistry().abortAll(); + const taskRegistry = config.getTaskRegistry(); + agentAbortAll(taskRegistry, { notify: false }); + monitorAbortAll(taskRegistry, { notify: false }); + shellAbortAll(taskRegistry); resetBackgroundStateForSessionSwitch(config); const newSessionId = config.startNewSession(); diff --git a/packages/cli/src/ui/commands/tasksCommand.test.ts b/packages/cli/src/ui/commands/tasksCommand.test.ts index fecddfeb28e..c80815a5998 100644 --- a/packages/cli/src/ui/commands/tasksCommand.test.ts +++ b/packages/cli/src/ui/commands/tasksCommand.test.ts @@ -97,9 +97,9 @@ describe('tasksCommand', () => { executionMode: 'non_interactive', services: { config: { - getBackgroundShellRegistry: () => ({ getAll: getShells }), - getBackgroundTaskRegistry: () => ({ getAll: getAgents }), - getMonitorRegistry: () => ({ getAll: getMonitors }), + getTaskRegistry: () => ({ + getAll: () => [...getAgents(), ...getShells(), ...getMonitors()], + }), }, }, } as unknown as Parameters[0]); @@ -304,9 +304,9 @@ describe('tasksCommand', () => { executionMode: 'interactive', services: { config: { - getBackgroundShellRegistry: () => ({ getAll: getShells }), - getBackgroundTaskRegistry: () => ({ getAll: getAgents }), - getMonitorRegistry: () => ({ getAll: getMonitors }), + getTaskRegistry: () => ({ + getAll: () => [...getAgents(), ...getShells(), ...getMonitors()], + }), }, }, } as unknown as Parameters[0]); @@ -332,9 +332,9 @@ describe('tasksCommand', () => { executionMode: 'acp', services: { config: { - getBackgroundShellRegistry: () => ({ getAll: getShells }), - getBackgroundTaskRegistry: () => ({ getAll: getAgents }), - getMonitorRegistry: () => ({ getAll: getMonitors }), + getTaskRegistry: () => ({ + getAll: () => [...getAgents(), ...getShells(), ...getMonitors()], + }), }, }, } as unknown as Parameters[0]); diff --git a/packages/cli/src/ui/commands/tasksCommand.ts b/packages/cli/src/ui/commands/tasksCommand.ts index 3096be24fef..e02b378466a 100644 --- a/packages/cli/src/ui/commands/tasksCommand.ts +++ b/packages/cli/src/ui/commands/tasksCommand.ts @@ -176,22 +176,11 @@ export const tasksCommand: SlashCommand = { }; } - // Each registry already tags entries with `kind`, so no per-entry - // mapping is needed here — just spread into a single sorted list. - const agentEntries: AgentTask[] = [ - ...config.getBackgroundTaskRegistry().getAll(), - ]; - const shellEntries: ShellTask[] = [ - ...config.getBackgroundShellRegistry().getAll(), - ]; - const monitorEntries: MonitorTask[] = [ - ...config.getMonitorRegistry().getAll(), - ]; - const entries: TaskEntry[] = [ - ...agentEntries, - ...shellEntries, - ...monitorEntries, - ].sort((a, b) => a.startTime - b.startTime); + // Every entry in the registry carries `kind` already; one + // `getAll()` is enough to render the full mix in launch order. + const entries: TaskEntry[] = [...config.getTaskRegistry().getAll()].sort( + (a, b) => a.startTime - b.startTime, + ); if (entries.length === 0) { return { diff --git a/packages/cli/src/ui/components/background-view/BackgroundTasksDialog.test.tsx b/packages/cli/src/ui/components/background-view/BackgroundTasksDialog.test.tsx index 69bc50f681a..ac31dc997a3 100644 --- a/packages/cli/src/ui/components/background-view/BackgroundTasksDialog.test.tsx +++ b/packages/cli/src/ui/components/background-view/BackgroundTasksDialog.test.tsx @@ -24,6 +24,34 @@ import { } from '../../hooks/useBackgroundTaskView.js'; import { useKeypress } from '../../hooks/useKeypress.js'; +// Hoisted spies for the kind-local kill dispatch. The dialog's +// cancelSelected path now calls `getTaskByType(target.kind).kill(...)`; +// each mock here lets the test assert against a kind without wiring +// the real per-kind module. +const mockAgentKill = vi.hoisted(() => vi.fn()); +const mockShellKill = vi.hoisted(() => vi.fn()); +const mockMonitorKill = vi.hoisted(() => vi.fn()); +const mockDreamKill = vi.hoisted(() => vi.fn()); + +vi.mock('@qwen-code/qwen-code-core', async () => { + const actual = await vi.importActual('@qwen-code/qwen-code-core'); + return { + ...actual, + getTaskByType: (kind: 'agent' | 'shell' | 'monitor' | 'dream') => ({ + kind, + name: `${kind} (mock)`, + kill: + kind === 'agent' + ? mockAgentKill + : kind === 'shell' + ? mockShellKill + : kind === 'monitor' + ? mockMonitorKill + : mockDreamKill, + }), + }; +}); + vi.mock('../../hooks/useBackgroundTaskView.js', () => ({ useBackgroundTaskView: vi.fn(), // Re-export the helper so Dialog renderers can still resolve it under the @@ -141,27 +169,33 @@ function setup(initial: readonly DialogEntry[]): Harness { // live activity/stats mutations the snapshot misses. let currentEntries: readonly DialogEntry[] = initial; const config = { - getBackgroundTaskRegistry: () => ({ - cancel, - setActivityChangeCallback: vi.fn(), - get: (id: string) => { - const match = currentEntries.find( - (e) => e.kind === 'agent' && e.agentId === id, - ); - return match; - }, - }), - getMonitorRegistry: () => ({ - cancel: monitorCancel, - // Resolve `.get(monitorId)` against the snapshot so the dialog's - // `selectedEntry` re-resolution path works for monitor kind too. + getTaskRegistry: () => ({ get: (id: string) => { + // Tests reach in for both agent and monitor kinds via the + // unified `get(id)` shape; resolve from the snapshot using + // each kind's id field. Shells aren't exercised in this dialog + // suite, so they fall through. const match = currentEntries.find( - (e) => e.kind === 'monitor' && e.monitorId === id, + (e) => + (e.kind === 'agent' && e.agentId === id) || + (e.kind === 'monitor' && e.monitorId === id), ); return match; }, + getAll: () => currentEntries, + getByKind: () => [], + register: () => undefined, + update: () => undefined, + evict: () => undefined, + kill: () => undefined, + subscribe: () => () => {}, }), + // Replaces the old per-kind cancel mocks. The dialog now dispatches + // through `getTaskByType(target.kind).kill(...)`, but the dispatcher + // is exercised in the registry's own tests; here we just verify the + // dialog routes through ConfigContext correctly. + _agentCancel: cancel, + _monitorCancel: monitorCancel, getMemoryManager: () => ({ cancelTask: dreamCancelTask, }), @@ -265,7 +299,7 @@ describe('BackgroundTasksDialog', () => { expect(h.probe.current!.state.dialogMode).toBe('detail'); h.pressKey({ sequence: 'x' }); - expect(h.cancel).toHaveBeenCalledWith('a'); + expect(mockAgentKill).toHaveBeenCalledWith('a', expect.anything()); // Registry would push the cancelled status; simulate that update. h.setEntries([{ ...running, status: 'cancelled' }]); @@ -286,10 +320,10 @@ describe('BackgroundTasksDialog', () => { expect(h.probe.current!.state.dialogMode).toBe('detail'); h.pressKey({ sequence: 'x' }); - expect(h.monitorCancel).toHaveBeenCalledWith('mon-zzz'); + expect(mockMonitorKill).toHaveBeenCalledWith('mon-zzz', expect.anything()); // Agent registry's cancel must NOT be called for a monitor entry — // belt-and-braces guard against the kind switch falling through. - expect(h.cancel).not.toHaveBeenCalled(); + expect(mockAgentKill).not.toHaveBeenCalled(); }); it('keeps detail mode when an already-terminal entry is opened (no spurious fallback)', () => { @@ -322,10 +356,10 @@ describe('BackgroundTasksDialog', () => { h.call(() => h.probe.current!.actions.openDialog()); h.pressKey({ sequence: 'x' }); - expect(h.cancel).not.toHaveBeenCalled(); + expect(mockAgentKill).not.toHaveBeenCalled(); h.pressKey({ sequence: 'x' }); - expect(h.cancel).toHaveBeenCalledWith('fg-1'); + expect(mockAgentKill).toHaveBeenCalledWith('fg-1', expect.anything()); }); it('background cancel still fires on the first `x` press (no confirm)', () => { @@ -342,7 +376,7 @@ describe('BackgroundTasksDialog', () => { h.call(() => h.probe.current!.actions.openDialog()); h.pressKey({ sequence: 'x' }); - expect(h.cancel).toHaveBeenCalledWith('bg-1'); + expect(mockAgentKill).toHaveBeenCalledWith('bg-1', expect.anything()); }); it('ignores `x` on a terminal foreground entry (no arm, no cancel call)', () => { @@ -365,7 +399,7 @@ describe('BackgroundTasksDialog', () => { expect(h.lastFrame()).not.toContain('x again to confirm stop'); h.pressKey({ sequence: 'x' }); - expect(h.cancel).not.toHaveBeenCalled(); + expect(mockAgentKill).not.toHaveBeenCalled(); }); it('detail-mode left clears any armed foreground cancel before exiting', () => { @@ -391,7 +425,7 @@ describe('BackgroundTasksDialog', () => { // Back in list mode, the next `x` arms again rather than confirming // a stale armed state inherited from detail mode. h.pressKey({ sequence: 'x' }); - expect(h.cancel).not.toHaveBeenCalled(); + expect(mockAgentKill).not.toHaveBeenCalled(); }); it('Esc backs out of an armed foreground cancel without closing the dialog', () => { @@ -412,7 +446,7 @@ describe('BackgroundTasksDialog', () => { // After the Esc reset, the next `x` arms again rather than confirming. h.pressKey({ sequence: 'x' }); - expect(h.cancel).not.toHaveBeenCalled(); + expect(mockAgentKill).not.toHaveBeenCalled(); }); it('clamps selectedIndex when entries shrink', () => { @@ -661,12 +695,12 @@ describe('BackgroundTasksDialog', () => { const h = setup([dreamEntry({ dreamId: 'd-zzz', status: 'running' })]); h.call(() => h.probe.current!.actions.openDialog()); h.pressKey({ sequence: 'x' }); - expect(h.dreamCancelTask).toHaveBeenCalledWith('d-zzz'); + expect(mockDreamKill).toHaveBeenCalledWith('d-zzz', expect.anything()); // Belt-and-braces — the registry-side cancel paths must not fire // for a dream entry, otherwise the wrong AbortController gets // signalled. - expect(h.cancel).not.toHaveBeenCalled(); - expect(h.monitorCancel).not.toHaveBeenCalled(); + expect(mockAgentKill).not.toHaveBeenCalled(); + expect(mockMonitorKill).not.toHaveBeenCalled(); }); it('omits the topics block entirely while the dream is still running', () => { diff --git a/packages/cli/src/ui/components/background-view/BackgroundTasksDialog.tsx b/packages/cli/src/ui/components/background-view/BackgroundTasksDialog.tsx index b240f037d27..db9d61a4d1b 100644 --- a/packages/cli/src/ui/components/background-view/BackgroundTasksDialog.tsx +++ b/packages/cli/src/ui/components/background-view/BackgroundTasksDialog.tsx @@ -25,9 +25,10 @@ import { theme } from '../../semantic-colors.js'; import { useConfig } from '../../contexts/ConfigContext.js'; import { buildBackgroundEntryLabel, + getAgentTask, + getMonitorTask, ToolDisplayNames, ToolNames, - type AgentTask, type MonitorTask, } from '@qwen-code/qwen-code-core'; import { formatDuration, formatTokenCount } from '../../utils/formatters.js'; @@ -924,21 +925,25 @@ export const BackgroundTasksDialog: React.FC = ({ // statusChange (so the pill / AppContainer don't churn under heavy // tool / event traffic), so for the detail view we have to re-resolve // explicitly: - // - agent: `recentActivities` is reassigned by `appendActivity`, - // which fires `activityChange` (subscribed below). - // - monitor: `eventCount` / `droppedLines` are mutated by - // `emitEvent`, which intentionally does NOT fire `statusChange` - // to avoid per-event refresh churn. The 1s wall-clock tick below - // drives the recompute instead. + // - agent: `recentActivities` is reassigned by `agentAppendActivity`, + // which uses `mutateSilent` and intentionally does NOT fire the + // subscription (to avoid per-activity refresh churn), so the 1s + // wall-clock tick below drives the recompute instead. + // - monitor: `eventCount` / `droppedLines` are mutated by `emitEvent` + // via `mutateSilent`, which likewise does NOT fire the subscription; + // the same 1s wall-clock tick drives the recompute. // Shells don't mutate detail-visible fields between statusChange // events, so the snapshot stays correct for them. if (fromSnapshot.kind === 'agent') { - const live = config.getBackgroundTaskRegistry().get(fromSnapshot.agentId); - return live ? { ...live, kind: 'agent' as const } : fromSnapshot; + const live = getAgentTask(config.getTaskRegistry(), fromSnapshot.agentId); + return live ?? fromSnapshot; } if (fromSnapshot.kind === 'monitor') { - const live = config.getMonitorRegistry().get(fromSnapshot.monitorId); - return live ? { ...live, kind: 'monitor' as const } : fromSnapshot; + const live = getMonitorTask( + config.getTaskRegistry(), + fromSnapshot.monitorId, + ); + return live ?? fromSnapshot; } return fromSnapshot; // activityTick is a dep on purpose: the registry mutation is invisible @@ -947,25 +952,38 @@ export const BackgroundTasksDialog: React.FC = ({ }, [entries, selectedIndex, config, activityTick]); const selectedEntryId = selectedEntry ? entryId(selectedEntry) : undefined; - // Activity callback is agent-only — shells don't emit per-tool events. - const selectedAgentIdForActivity = - selectedEntry?.kind === 'agent' ? selectedEntry.agentId : undefined; + // Live in-place updates (agent activity bursts, monitor event count) + // don't change the snapshot's shape signature, so useBackgroundTaskView + // skips its setEntries fan-out for them. The detail view subscribes + // directly to the selected agent / monitor so its per-entry display + // still reflects them. + const selectedLiveEntityId = + selectedEntry?.kind === 'agent' + ? selectedEntry.agentId + : selectedEntry?.kind === 'monitor' + ? selectedEntry.monitorId + : undefined; useEffect(() => { - if (!dialogOpen || !isDetailMode || !selectedAgentIdForActivity) return; - const registry = config.getBackgroundTaskRegistry(); - const onActivity = (entry: AgentTask) => { - if (entry.agentId !== selectedAgentIdForActivity) return; + // `isDetailMode` covers both `dialogMode === 'detail'` (entered from + // the list) and `dialogMode === 'detail-from-panel'` (entered from + // the LiveAgentPanel). The detail body renders in both modes and + // every sibling effect (below) uses the same predicate; narrowing + // to just `'detail'` would silently drop live-refresh on the + // panel-entered path. + if (!dialogOpen || !isDetailMode || !selectedLiveEntityId) return; + const registry = config.getTaskRegistry(); + return registry.subscribe((entry) => { + if (!entry) return; + const id = + entry.kind === 'agent' + ? entry.agentId + : entry.kind === 'monitor' + ? entry.monitorId + : undefined; + if (id !== selectedLiveEntityId) return; setActivityTick((n) => n + 1); - }; - registry.setActivityChangeCallback(onActivity); - return () => registry.setActivityChangeCallback(undefined); - }, [ - dialogOpen, - dialogMode, - isDetailMode, - config, - selectedAgentIdForActivity, - ]); + }); + }, [dialogOpen, isDetailMode, config, selectedLiveEntityId]); // Wall-clock tick for the running agent's duration. Activity callbacks // fire when tools run, but duration needs to advance even when the agent diff --git a/packages/cli/src/ui/components/background-view/LiveAgentPanel.test.tsx b/packages/cli/src/ui/components/background-view/LiveAgentPanel.test.tsx index 8534b8ca770..7fe77739167 100644 --- a/packages/cli/src/ui/components/background-view/LiveAgentPanel.test.tsx +++ b/packages/cli/src/ui/components/background-view/LiveAgentPanel.test.tsx @@ -98,8 +98,18 @@ function makeRegistryConfig(agents: readonly AgentDialogEntry[]): { const store = new Map(); for (const a of agents) store.set(a.agentId, a); const config = { - getBackgroundTaskRegistry: () => ({ + getTaskRegistry: () => ({ get: (id: string) => store.get(id), + // The panel uses `getAgentTask(registry, id)` which calls + // `registry.get(id)` and narrows by kind. The stub entries + // already carry `kind: 'agent'`, so the narrowing succeeds. + getAll: () => Array.from(store.values()), + getByKind: () => Array.from(store.values()), + register: () => undefined, + update: () => undefined, + evict: () => undefined, + kill: () => undefined, + subscribe: () => () => {}, }), } as unknown as Config; return { config, store }; diff --git a/packages/cli/src/ui/components/background-view/LiveAgentPanel.tsx b/packages/cli/src/ui/components/background-view/LiveAgentPanel.tsx index 84b6835c40c..274f19c8f47 100644 --- a/packages/cli/src/ui/components/background-view/LiveAgentPanel.tsx +++ b/packages/cli/src/ui/components/background-view/LiveAgentPanel.tsx @@ -31,6 +31,7 @@ import { useContext, useEffect, useMemo, useRef, useState } from 'react'; import { Box, Text } from 'ink'; import { DEFAULT_BUILTIN_SUBAGENT_TYPE as CORE_DEFAULT_SUBAGENT_TYPE, + getAgentTask, ToolDisplayNames, ToolNames, } from '@qwen-code/qwen-code-core'; @@ -281,7 +282,7 @@ export const LiveAgentPanel: React.FC = ({ const liveAgentSnapshots: AgentDialogEntry[] = useMemo(() => { const snapshots = entries.filter(isAgentEntry); if (!config) return snapshots; - const registry = config.getBackgroundTaskRegistry(); + const registry = config.getTaskRegistry(); // `now` participates in the dependency array so the memo recomputes // each tick and picks up `recentActivities` the registry mutated in // place via appendActivity. Reading it here makes the dependency @@ -293,13 +294,13 @@ export const LiveAgentPanel: React.FC = ({ const next = snapshots .map((snap) => { seenIds.add(snap.agentId); - const live = registry.get(snap.agentId); + const live = getAgentTask(registry, snap.agentId); if (live) { // Recovered (or never went missing) — drop any stale // missing-since record so a future re-disappearance // gets a fresh timestamp. missingSinceRef.current.delete(snap.agentId); - return { ...live, kind: 'agent' as const }; + return live; } if (snap.status === 'running' || snap.status === 'paused') { // Pin the disappearance time on first observation so diff --git a/packages/cli/src/ui/components/messages/InlineParallelAgentsDisplay.test.tsx b/packages/cli/src/ui/components/messages/InlineParallelAgentsDisplay.test.tsx index 4b823f6cc15..f58853c861e 100644 --- a/packages/cli/src/ui/components/messages/InlineParallelAgentsDisplay.test.tsx +++ b/packages/cli/src/ui/components/messages/InlineParallelAgentsDisplay.test.tsx @@ -60,7 +60,7 @@ function makeRegistryConfig(entries: Array>): { } } const config = { - getBackgroundTaskRegistry: () => ({ + getTaskRegistry: () => ({ get: (id: string) => store.get(id), }), } as unknown as Config; diff --git a/packages/cli/src/ui/components/messages/InlineParallelAgentsDisplay.tsx b/packages/cli/src/ui/components/messages/InlineParallelAgentsDisplay.tsx index a444a48dc5f..b296d62e863 100644 --- a/packages/cli/src/ui/components/messages/InlineParallelAgentsDisplay.tsx +++ b/packages/cli/src/ui/components/messages/InlineParallelAgentsDisplay.tsx @@ -221,14 +221,16 @@ export const InlineParallelAgentsDisplay: React.FC< // mutates `recentActivities` in place, so without a tick the // component would freeze on the first row of activity. const rows: RowData[] = useMemo(() => { - const registry = config?.getBackgroundTaskRegistry(); + const registry = config?.getTaskRegistry(); // Touch `now` so a future "remove dead dep" cleanup can't silently // freeze the panel — the registry mutates in place and we need to // re-read on every tick to surface fresh activity. void now; return agentEntries.map(({ toolCall, result }) => { const agentId = deriveAgentId(toolCall, result); - const live = registry?.get(agentId); + const liveEntry = registry?.get(agentId); + const live = + liveEntry?.kind === 'agent' ? liveEntry : undefined; const recent = live?.recentActivities?.at(-1); return { agentId, diff --git a/packages/cli/src/ui/contexts/BackgroundTaskViewContext.tsx b/packages/cli/src/ui/contexts/BackgroundTaskViewContext.tsx index 43007342eff..e9af8065562 100644 --- a/packages/cli/src/ui/contexts/BackgroundTaskViewContext.tsx +++ b/packages/cli/src/ui/contexts/BackgroundTaskViewContext.tsx @@ -19,8 +19,13 @@ import { useMemo, useState, } from 'react'; -import { type Config, createDebugLogger } from '@qwen-code/qwen-code-core'; import { + getTaskByType, + type Config, + createDebugLogger, +} from '@qwen-code/qwen-code-core'; +import { + entryId, type DialogEntry, useBackgroundTaskView, } from '../hooks/useBackgroundTaskView.js'; @@ -220,51 +225,36 @@ export function BackgroundTaskViewProvider({ config.abandonBackgroundAgent(target.agentId); return; } - // All three registries' cancel paths are no-ops on non-running - // entries, so no pre-check here. Shell cancel goes through - // requestCancel — it triggers the AbortController only and lets the - // spawn's settle path record the real terminal moment + outcome - // (mirrors the task_stop tool path in #3687). Monitor cancel is - // synchronous: settle + abort happen inside the registry's cancel(), - // matching its own task_stop path. - switch (target.kind) { - case 'agent': - config.getBackgroundTaskRegistry().cancel(target.agentId); - break; - case 'shell': - config.getBackgroundShellRegistry().requestCancel(target.shellId); - break; - case 'monitor': - config.getMonitorRegistry().cancel(target.monitorId); - break; - case 'dream': { - // Aborts the dream fork-agent via MemoryManager.cancelTask; - // the manager flips status to 'cancelled' before aborting, and - // the runDream finally block releases the consolidation lock as - // the agent unwinds. Same one-shot fire-and-forget shape as - // shell.requestCancel above. - // - // cancelTask returns false in the contract-violation path - // (running record without an AbortController). Today this is - // unreachable because the controller is registered before - // storeWith fires the notify, but if a future refactor - // breaks the invariant a silent ignore here would let the - // user think the cancel took. Log + leave the dialog open. - const ok = config.getMemoryManager().cancelTask(target.dreamId); - if (!ok) { + // Fire-and-forget: the dialog re-renders from the registry/dream + // subscription once kill updates state. + const ctx = { + registry: config.getTaskRegistry(), + memoryManager: config.getMemoryManager(), + }; + // `getTaskByType(kind)` throws synchronously for unregistered kinds + // (defensive — should not happen, but our React tree shouldn't + // crash if it does). Catch the sync throw from the lookup with a + // try/catch, then dispatch the (async) kill call through a Promise + // chain so any later async rejection is also captured. Keeping + // `kill()` synchronous-at-call-site preserves the test contract + // "pressing 'x' invokes kill before the next tick." + try { + const taskKind = getTaskByType(target.kind); + void Promise.resolve(taskKind.kill(entryId(target), ctx)).catch( + (err) => { debugLogger.warn( - `cancelSelected: dream task ${target.dreamId} could not be cancelled ` + - `(internal state inconsistency — see MemoryManager.cancelTask warn).`, + `cancelSelected: kill for ${target.kind} task ${entryId(target)} failed: ${ + err instanceof Error ? err.message : String(err) + }`, ); - } - break; - } - default: { - const _exhaustive: never = target; - throw new Error( - `cancelSelected: unknown DialogEntry kind: ${JSON.stringify(_exhaustive)}`, - ); - } + }, + ); + } catch (err) { + debugLogger.warn( + `cancelSelected: kind dispatch for ${target.kind} threw: ${ + err instanceof Error ? err.message : String(err) + }`, + ); } }, [config, entries, selectedIndex]); diff --git a/packages/cli/src/ui/hooks/useBackgroundTaskView.test.ts b/packages/cli/src/ui/hooks/useBackgroundTaskView.test.ts index 36a00a681f3..fa028eb0c39 100644 --- a/packages/cli/src/ui/hooks/useBackgroundTaskView.test.ts +++ b/packages/cli/src/ui/hooks/useBackgroundTaskView.test.ts @@ -6,36 +6,29 @@ import { describe, it, expect, vi } from 'vitest'; import { renderHook, act } from '@testing-library/react'; -import type { Config } from '@qwen-code/qwen-code-core'; +import { + TaskRegistry, + agentRegister, + monitorRegister, + shellRegister, + type Config, +} from '@qwen-code/qwen-code-core'; import { useBackgroundTaskView, entryId } from './useBackgroundTaskView.js'; -interface FakeRegistry { - setStatusChangeCallback: ReturnType; - /** Test helper — invokes the currently-set callback. */ - fire: () => void; -} - -function makeFakeRegistry(): FakeRegistry { - let cb: (() => void) | undefined; - return { - setStatusChangeCallback: vi.fn((next: (() => void) | undefined) => { - cb = next; - }), - fire: () => cb?.(), - }; -} - interface FakeMemoryManager { subscribe: ReturnType; unsubscribe: ReturnType; - /** Captured opts from the most recent subscribe() call (the hook - * passes `{ taskType: 'dream' }` to skip per-extract notifies). */ + /** Captured opts from the most recent subscribe() call (the dream + * adapter passes `{ taskType: 'dream' }` to skip per-extract + * notifies). */ lastSubscribeOpts: { taskType?: 'extract' | 'dream' } | undefined; /** Test helper — invokes the currently-subscribed listener. */ fire: () => void; } -function makeFakeMemoryManager(): FakeMemoryManager { +function makeFakeMemoryManager( + listTasksByType: () => unknown[], +): FakeMemoryManager { let listener: (() => void) | undefined; const ref: { lastSubscribeOpts: FakeMemoryManager['lastSubscribeOpts'] } = { lastSubscribeOpts: undefined, @@ -50,115 +43,60 @@ function makeFakeMemoryManager(): FakeMemoryManager { return unsubscribe; }, ); - return { - subscribe, - unsubscribe, - get lastSubscribeOpts() { - return ref.lastSubscribeOpts; + return Object.assign( + { + subscribe, + unsubscribe, + get lastSubscribeOpts() { + return ref.lastSubscribeOpts; + }, + fire: () => listener?.(), }, - fire: () => listener?.(), - }; + { + listTasksByType, + }, + ) as FakeMemoryManager & { listTasksByType: typeof listTasksByType }; } -function makeConfig(opts: { - agents: () => unknown[]; - shells: () => unknown[]; - monitors: () => unknown[]; - dreams?: () => unknown[]; -}) { - const agentReg = makeFakeRegistry(); - const shellReg = makeFakeRegistry(); - const monitorReg = makeFakeRegistry(); - const memoryMgr = makeFakeMemoryManager(); - const dreams = opts.dreams ?? (() => []); - +function makeConfig( + registry: TaskRegistry, + dreams: () => unknown[] = () => [], +) { + const memoryMgr = makeFakeMemoryManager(dreams); const config = { - getBackgroundTaskRegistry: () => ({ - ...agentReg, - getAll: opts.agents, - }), - getBackgroundShellRegistry: () => ({ - ...shellReg, - getAll: opts.shells, - }), - getMonitorRegistry: () => ({ - ...monitorReg, - getAll: opts.monitors, - }), - getMemoryManager: () => ({ - subscribe: memoryMgr.subscribe, - // Hook only ever requests dream-typed records; ignore the type arg - // and return whatever the test provided. - listTasksByType: (_type: string, _projectRoot?: string) => dreams(), - }), + getTaskRegistry: () => registry, + getMemoryManager: () => memoryMgr, getProjectRoot: () => '/test/project', } as unknown as Config; - - return { config, agentReg, shellReg, monitorReg, memoryMgr }; + return { config, memoryMgr }; } -type StatusOverride = { - status?: 'running' | 'paused' | 'completed' | 'failed' | 'cancelled'; - endTime?: number; -}; - -const agent = ( - id: string, - startTime: number, - overrides: StatusOverride = {}, -) => ({ - id, - kind: 'agent' as const, +const agentReg = (id: string, startTime: number) => ({ agentId: id, description: 'desc', isBackgrounded: true, - status: overrides.status ?? ('running' as const), + status: 'running' as const, startTime, - endTime: overrides.endTime, abortController: new AbortController(), outputFile: '/tmp/agent.jsonl', - outputOffset: 0, - notified: false, }); -const shell = ( - id: string, - startTime: number, - overrides: Omit & { - status?: 'running' | 'completed' | 'failed' | 'cancelled'; - } = {}, -) => ({ - id, - kind: 'shell' as const, +const shellReg = (id: string, startTime: number) => ({ shellId: id, command: 'sleep 60', - description: 'sleep 60', cwd: '/tmp', - status: overrides.status ?? ('running' as const), + status: 'running' as const, startTime, - endTime: overrides.endTime, outputPath: '/tmp/x.out', - outputFile: '/tmp/x.out', - outputOffset: 0, - notified: false, abortController: new AbortController(), }); -const monitor = ( - id: string, - startTime: number, - overrides: Omit & { - status?: 'running' | 'completed' | 'failed' | 'cancelled'; - } = {}, -) => ({ - id, - kind: 'monitor' as const, +const monitorReg = (id: string, startTime: number) => ({ monitorId: id, - command: 'tail -f log', description: 'watch logs', - status: overrides.status ?? ('running' as const), + command: 'tail -f log', + status: 'running' as const, startTime, - endTime: overrides.endTime, abortController: new AbortController(), eventCount: 0, lastEventTime: 0, @@ -166,15 +104,8 @@ const monitor = ( idleTimeoutMs: 300_000, droppedLines: 0, outputFile: '/tmp/monitor.log', - outputOffset: 0, - notified: false, }); -// Mirror the MemoryTaskRecord shape that MemoryManager.listTasksByType -// returns. Status defaults to 'running'; tests override to exercise the -// filter (`pending` / `skipped` records must be excluded; `cancelled` -// flows through the same terminal-cap path as `completed` / `failed` -// once the task_stop / dialog cancel keystroke lands one). const dream = ( id: string, startTimeMs: number, @@ -208,346 +139,115 @@ describe('useBackgroundTaskView', () => { expect(result.current.entries).toEqual([]); }); - it('merges entries from all three registries on mount', () => { - const { config } = makeConfig({ - agents: () => [agent('a1', 100)], - shells: () => [shell('s1', 50)], - monitors: () => [monitor('m1', 200)], - }); + it('merges agent, shell, and monitor entries from the unified registry on mount', () => { + const registry = new TaskRegistry(); + agentRegister(registry, agentReg('a1', 100)); + shellRegister(registry, shellReg('s1', 50)); + monitorRegister(registry, monitorReg('m1', 200)); + const { config } = makeConfig(registry); const { result } = renderHook(() => useBackgroundTaskView(config)); expect(result.current.entries).toHaveLength(3); - // Sort order is by startTime descending — newest first: monitor - // (200) → agent (100) → shell (50). The dialog opens with the - // cursor on row 0, so the most recently launched task is the one - // immediately selected. + // `buildMerged` puts running entries first, sorted by startTime DESC + // (most recent launch on top). startTimes: a1=100, s1=50, m1=200. expect(result.current.entries.map(entryId)).toEqual(['m1', 'a1', 's1']); }); - it('orders entries newest-first across all kinds', () => { - // Pin the descending sort so a future refactor that flips the - // comparator silently re-introduces the "new task buried at the - // bottom of a long list" UX. Mix all four kinds at varying - // startTimes to exercise the merge path end-to-end. - const { config } = makeConfig({ - agents: () => [agent('a-old', 10), agent('a-new', 400)], - shells: () => [shell('s-mid', 200)], - monitors: () => [monitor('m-second-newest', 300)], - dreams: () => [dream('d-oldest', 5)], - }); - const { result } = renderHook(() => useBackgroundTaskView(config)); - expect(result.current.entries.map(entryId)).toEqual([ - 'a-new', - 'm-second-newest', - 's-mid', - 'a-old', - 'd-oldest', - ]); - }); - - it('puts active (running + paused) entries above terminal entries even when terminals are newer', () => { - // The literal phrasing of the issue is "new OR running tasks - // should appear at the top". A pure startTime DESC sort handles - // the "new" half but lets a long-running entry get buried under a - // batch of newer terminals (a quick agent that started AND - // finished after the long one). Pin the bucket order so the user - // opening the dialog to check on running work doesn't have to - // scroll past stale completed rows to find it. - const { config } = makeConfig({ - agents: () => [ - // Old running agent — must NOT be pushed below newer terminals. - agent('a-running-old', 100), - // Recently-completed agent — newer startTime than the running - // one, but should still sort below it because it's terminal. - agent('a-done-fresh', 500, { status: 'completed', endTime: 600 }), - // Paused agent — same bucket as running (user can resume / - // abandon), ranks by startTime DESC inside the bucket. - agent('a-paused', 300, { status: 'paused' }), - ], - shells: () => [ - // Failed shell launched in between the two active agents — - // belongs in the terminal bucket regardless of startTime. - shell('s-failed', 400, { status: 'failed', endTime: 450 }), - ], - monitors: () => [], - }); - const { result } = renderHook(() => useBackgroundTaskView(config)); - expect(result.current.entries.map(entryId)).toEqual([ - // Active bucket (startTime DESC): paused (300), running (100). - 'a-paused', - 'a-running-old', - // Terminal bucket (endTime DESC): a-done-fresh (600), s-failed (450). - 'a-done-fresh', - 's-failed', - ]); - }); - - it('orders the terminal bucket by endTime DESC (not startTime)', () => { - // A long-running task that just settled is more "interesting" to - // a returning user than an old quick task that finished hours - // ago, even if the latter has a higher startTime. - const { config } = makeConfig({ - agents: () => [ - // Started early, just finished — most recent terminal event. - agent('a-just-finished', 100, { - status: 'completed', - endTime: 1_000, - }), - // Started later, finished early — older terminal event. - agent('a-quick-and-old', 500, { - status: 'completed', - endTime: 600, - }), - ], - shells: () => [], - monitors: () => [], - }); - const { result } = renderHook(() => useBackgroundTaskView(config)); - expect(result.current.entries.map(entryId)).toEqual([ - 'a-just-finished', - 'a-quick-and-old', - ]); - }); - it('tags each merged entry with the right `kind` discriminator', () => { - const { config } = makeConfig({ - agents: () => [agent('a1', 0)], - shells: () => [shell('s1', 0)], - monitors: () => [monitor('m1', 0)], - }); + const registry = new TaskRegistry(); + agentRegister(registry, agentReg('a1', 0)); + shellRegister(registry, shellReg('s1', 0)); + monitorRegister(registry, monitorReg('m1', 0)); + const { config } = makeConfig(registry); const { result } = renderHook(() => useBackgroundTaskView(config)); const kinds = result.current.entries.map((e) => e.kind).sort(); expect(kinds).toEqual(['agent', 'monitor', 'shell']); }); - it('subscribes to all three registries on mount', () => { - const { config, agentReg, shellReg, monitorReg } = makeConfig({ - agents: () => [], - shells: () => [], - monitors: () => [], - }); - renderHook(() => useBackgroundTaskView(config)); - expect(agentReg.setStatusChangeCallback).toHaveBeenCalledWith( - expect.any(Function), - ); - expect(shellReg.setStatusChangeCallback).toHaveBeenCalledWith( - expect.any(Function), - ); - expect(monitorReg.setStatusChangeCallback).toHaveBeenCalledWith( - expect.any(Function), - ); - }); - - it('refreshes entries when any registry fires statusChange', () => { - const agents: Array> = []; - const monitors: Array> = []; - const { config, agentReg, monitorReg } = makeConfig({ - agents: () => agents, - shells: () => [], - monitors: () => monitors, - }); + it('refreshes entries when the registry fires a change', () => { + const registry = new TaskRegistry(); + const { config } = makeConfig(registry); const { result } = renderHook(() => useBackgroundTaskView(config)); expect(result.current.entries).toEqual([]); - // Simulate registry mutation + statusChange fire from each registry. - agents.push(agent('a1', 100)); - act(() => agentReg.fire()); + act(() => { + agentRegister(registry, agentReg('a1', 100)); + }); expect(result.current.entries.map(entryId)).toEqual(['a1']); - monitors.push(monitor('m1', 50)); - act(() => monitorReg.fire()); - // Sort is descending by startTime: agent (100) sits above monitor - // (50) because the user wants the newest entry on top. + act(() => { + monitorRegister(registry, monitorReg('m1', 50)); + }); + // startTime DESC: a1=100 sorts before m1=50. expect(result.current.entries.map(entryId)).toEqual(['a1', 'm1']); }); - it('clears all three subscriptions on unmount', () => { - const { config, agentReg, shellReg, monitorReg, memoryMgr } = makeConfig({ - agents: () => [], - shells: () => [], - monitors: () => [], - }); + it('clears the registry subscription and the dream subscription on unmount', () => { + const registry = new TaskRegistry(); + const { config, memoryMgr } = makeConfig(registry); const { unmount } = renderHook(() => useBackgroundTaskView(config)); unmount(); - // Each setStatusChangeCallback should have been called twice — once - // with the refresh function on mount, once with `undefined` on - // cleanup. Failing this check would mean stale subscribers can fire - // into an unmounted component (warning + state-update on unmounted - // tree, sometimes crashes the next render). - expect(agentReg.setStatusChangeCallback.mock.calls).toEqual([ - [expect.any(Function)], - [undefined], - ]); - expect(shellReg.setStatusChangeCallback.mock.calls).toEqual([ - [expect.any(Function)], - [undefined], - ]); - expect(monitorReg.setStatusChangeCallback.mock.calls).toEqual([ - [expect.any(Function)], - [undefined], - ]); - // MemoryManager uses subscribe()/unsubscribe rather than the - // setCallback pattern; the unsubscribe returned from subscribe must - // run on cleanup or stale dream listeners leak across remounts. expect(memoryMgr.subscribe).toHaveBeenCalledTimes(1); expect(memoryMgr.unsubscribe).toHaveBeenCalledTimes(1); - }); - - it('surfaces dream tasks with kind=dream and skips pending/skipped records', () => { - const { config } = makeConfig({ - agents: () => [], - shells: () => [], - monitors: () => [], - // Three dream records covering: a pre-fire pending record (must - // not surface — would flood the dialog with one row per - // UserQuery), a running fire (must surface), and a skipped - // gate-miss (must not surface — same flood concern). - dreams: () => [ - dream('d-pending', 100, { status: 'pending' }), - dream('d-running', 200), - dream('d-skipped', 300, { status: 'skipped' }), - ], - }); + // After unmount, registry mutations must not throw or update state. + // The hook's cleanup unregisters the listener; if it didn't, this + // would log a "setState on unmounted component" warning. + expect(() => agentRegister(registry, agentReg('a-late', 0))).not.toThrow(); + }); + + it('surfaces dream tasks and skips pending/skipped records', () => { + const registry = new TaskRegistry(); + const dreams = () => [ + dream('d-pending', 100, { status: 'pending' }), + dream('d-running', 200), + dream('d-skipped', 300, { status: 'skipped' }), + ]; + const { config } = makeConfig(registry, dreams); const { result } = renderHook(() => useBackgroundTaskView(config)); expect(result.current.entries).toHaveLength(1); const [only] = result.current.entries; expect(only.kind).toBe('dream'); - expect(only.status).toBe('running'); expect(entryId(only)).toBe('d-running'); }); - it('caps retained terminal dream entries at 3 most-recent (by updatedAt) plus all running', () => { - // MemoryManager has no eviction; without the cap, accumulating - // completed dreams across a long session would blow up the dialog. - // The cap keeps the dialog glanceable while still surfacing the - // most recent outcomes (mirrors MonitorRegistry's terminal cap). - const baseMs = Date.parse('2026-05-04T12:00:00.000Z'); - const completed = (id: string, mtime: number) => ({ - id, - taskType: 'dream' as const, - projectRoot: '/test/project', - status: 'completed' as const, - createdAt: new Date(baseMs + mtime - 1000).toISOString(), - updatedAt: new Date(baseMs + mtime).toISOString(), - }); - const { config } = makeConfig({ - agents: () => [], - shells: () => [], - monitors: () => [], - dreams: () => [ - completed('d-old-1', 1_000), - completed('d-old-2', 2_000), - completed('d-mid', 3_000), - completed('d-recent', 4_000), - completed('d-newest', 5_000), - // Plus a running entry that must always survive the cap (caps - // only trim terminals; running dreams are uncapped). - dream('d-running-now', baseMs + 6_000, { status: 'running' }), - ], - }); - const { result } = renderHook(() => useBackgroundTaskView(config)); - const ids = result.current.entries.map(entryId).sort(); - // Surviving terminal entries: d-newest, d-recent, d-mid (top 3 by - // updatedAt desc). The two oldest (d-old-1, d-old-2) get dropped. - // The running dream survives unconditionally. - expect(ids).toEqual( - ['d-mid', 'd-newest', 'd-recent', 'd-running-now'].sort(), - ); - }); - - it('surfaces a cancelled dream with kind=dream so the dialog can render the terminal status', () => { - // `'cancelled'` arrives via the dialog `x stop` / `task_stop` path - // which routes through `MemoryManager.cancelTask`. The view-model - // must accept it the same way it accepts `'completed'` / `'failed'`, - // because the dialog's terminal-cap window depends on showing the - // user the outcome of the abort they just triggered. - const { config } = makeConfig({ - agents: () => [], - shells: () => [], - monitors: () => [], - dreams: () => [dream('d-stopped', 100, { status: 'cancelled' })], - }); + it('caps retained terminal dream entries at MAX_RETAINED_TERMINAL_DREAMS = 3', () => { + const registry = new TaskRegistry(); + const dreams = () => [ + dream('d-1', 100, { status: 'completed' }), + dream('d-2', 200, { status: 'completed' }), + dream('d-3', 300, { status: 'completed' }), + dream('d-4', 400, { status: 'completed' }), + dream('d-5', 500, { status: 'completed' }), + ]; + const { config } = makeConfig(registry, dreams); const { result } = renderHook(() => useBackgroundTaskView(config)); - expect(result.current.entries).toHaveLength(1); - const [only] = result.current.entries; - expect(only.kind).toBe('dream'); - expect(only.status).toBe('cancelled'); + expect(result.current.entries).toHaveLength(3); + // Most-recent-first by updatedAt → d-5, d-4, d-3 (sorted by + // startTime ascending in the merged output). + expect(result.current.entries.map(entryId).sort()).toEqual([ + 'd-3', + 'd-4', + 'd-5', + ]); }); - it('subscribes to MemoryManager with a dream taskType filter so extract notifies are skipped at the source', () => { - // The taskType filter on MemoryManager.subscribe() is the - // primary perf guard — it prevents the per-UserQuery extract - // notify from waking the bg-tasks UI listener at all (avoids the - // O(n) dream-snapshot fetch + signature compare that would - // otherwise run on every extract transition). Pin the filter so - // a future refactor that drops the opts arg fails the test - // rather than silently re-introducing the wakeups. - const { config, memoryMgr } = makeConfig({ - agents: () => [], - shells: () => [], - monitors: () => [], - }); + it('subscribes to MemoryManager dream events with `{ taskType: dream }`', () => { + const registry = new TaskRegistry(); + const { config, memoryMgr } = makeConfig(registry); renderHook(() => useBackgroundTaskView(config)); - expect(memoryMgr.subscribe).toHaveBeenCalledTimes(1); - expect(memoryMgr.lastSubscribeOpts).toEqual({ taskType: 'dream' }); - }); - - it('skips setEntries when the memory listener fires with unchanged dream content', () => { - // MemoryManager.subscribe() fires for ALL task transitions, including - // extract task records that have no dialog surface. Without the - // dream-signature dedup, every extract notify would trigger a full - // re-merge + a fresh array reference into setEntries — re-rendering - // the dialog and pill on entries that are byte-identical to the - // previous snapshot. This test pins the dedup by firing the memory - // listener while the dream snapshot stays unchanged and asserting - // that the entries reference is preserved. - const dreams: Array> = [dream('d-only', 100)]; - const { config, memoryMgr } = makeConfig({ - agents: () => [], - shells: () => [], - monitors: () => [], - dreams: () => dreams, + expect(memoryMgr.subscribe).toHaveBeenCalledWith(expect.any(Function), { + taskType: 'dream', }); - const { result } = renderHook(() => useBackgroundTaskView(config)); - const before = result.current.entries; - expect(before.map(entryId)).toEqual(['d-only']); - - // Fire the memory listener without mutating `dreams`. With the - // signature-dedup in place, this must NOT call setEntries; React - // will then preserve the existing array reference. - act(() => memoryMgr.fire()); - expect(result.current.entries).toBe(before); - - // Sanity check the inverse path: when dreams DO change, the - // listener must propagate. A flipped status should change the - // signature and force a fresh setEntries. - dreams.splice(0, 1, dream('d-only', 100, { status: 'completed' })); - act(() => memoryMgr.fire()); - expect(result.current.entries).not.toBe(before); - expect(result.current.entries[0]?.status).toBe('completed'); }); - it('refreshes entries when the memory manager fires its subscribe listener', () => { - const dreams: Array> = []; - const { config, memoryMgr } = makeConfig({ - agents: () => [], - shells: () => [], - monitors: () => [], - dreams: () => dreams, - }); + it('refreshes entries when MemoryManager dream subscription fires', () => { + const registry = new TaskRegistry(); + let dreamRecords: Array> = []; + const { config, memoryMgr } = makeConfig(registry, () => dreamRecords); const { result } = renderHook(() => useBackgroundTaskView(config)); expect(result.current.entries).toEqual([]); - dreams.push(dream('d-1', 100)); + dreamRecords = [dream('d-1', 100)]; act(() => memoryMgr.fire()); expect(result.current.entries.map(entryId)).toEqual(['d-1']); - - // A subsequent terminal state update must propagate the new status - // (running → completed) and survive the filter (only pending / - // skipped get dropped). - dreams.splice(0, dreams.length, dream('d-1', 100, { status: 'completed' })); - act(() => memoryMgr.fire()); - const [only] = result.current.entries; - expect(only.kind).toBe('dream'); - expect(only.status).toBe('completed'); }); }); diff --git a/packages/cli/src/ui/hooks/useBackgroundTaskView.ts b/packages/cli/src/ui/hooks/useBackgroundTaskView.ts index 5c82d58af11..93a8e02290e 100644 --- a/packages/cli/src/ui/hooks/useBackgroundTaskView.ts +++ b/packages/cli/src/ui/hooks/useBackgroundTaskView.ts @@ -5,44 +5,31 @@ */ /** - * useBackgroundTaskView — subscribes to the three background-task - * registries (background subagents, managed shells, and event monitors) - * AND to `MemoryManager` for dream consolidation tasks, merging them - * into a single ordered snapshot of `DialogEntry`s. Each registry fires - * `statusChange` on register too, so a single subscription per registry - * is enough to keep the snapshot fresh for new + transitioning entries. - * The `MemoryManager.subscribe({ taskType: 'dream' })` filter routes - * dream-task transitions to the same refresh path while skipping the - * per-UserQuery extract notifies that have no dialog surface. + * useBackgroundTaskView — subscribes to the unified `TaskRegistry` + * (agents + shells + monitors) AND to `MemoryManager` via the dream + * adapter for dream consolidation tasks. Merges them into a single + * ordered snapshot of `DialogEntry`s. + * + * The registry's `subscribe` fires on every change to any task entry — + * register, status transition, activity append. The dream adapter + * surfaces dream task changes through a separate subscription against + * `MemoryManager`. Both feed the same `refresh` path. * * Surfaces that only care about live work (the footer pill, the * composer's Down-arrow route) filter for `running` themselves. - * - * Intentionally ignores activity updates (appendActivity). Tool-call - * traffic from a running background agent would otherwise churn the - * Footer pill and the AppContainer every few hundred ms. The detail - * dialog subscribes to the activity callback directly when it needs - * live Progress updates. */ import { useState, useEffect } from 'react'; import { + dreamSnapshotSignature, + listDreamTasks, + subscribeDreams, type AgentTask, type Config, - type MemoryTaskRecord, - type MonitorTask, - type ShellTask, + type DreamTask, type TaskState, } from '@qwen-code/qwen-code-core'; -// Cap on retained terminal dream entries surfaced via the dialog. -// `MemoryManager.tasks` has no eviction; without this cap the list -// grows unboundedly with completed dreams over the project's lifetime. -// 3 is small enough to stay glanceable yet keeps the most recent -// outcomes visible across rapid succession (e.g. the user opening the -// dialog right after two dreams completed). -const MAX_RETAINED_TERMINAL_DREAMS = 3; - /** * @deprecated Use {@link AgentTask} from `@qwen-code/qwen-code-core` * directly. Kept as a one-release alias while UI consumers migrate. @@ -50,58 +37,24 @@ const MAX_RETAINED_TERMINAL_DREAMS = 3; export type AgentDialogEntry = AgentTask; /** - * Dream-task adapter. MemoryManager owns its own task records - * (MemoryTaskRecord) and intentionally lives outside the registry trio; - * this view-model wraps the subset of fields the dialog needs and - * narrows status to the four values that ever appear in the dialog - * (skipped/pending records are filtered out at the source). + * @deprecated Renamed to {@link DreamTask}; kept as a one-release alias + * for UI consumers migrating off the previous local view-model. */ -export type DreamDialogEntry = { - kind: 'dream'; - /** MemoryTaskRecord.id — used as React key + lookup. */ - dreamId: string; - status: 'running' | 'completed' | 'failed' | 'cancelled'; - startTime: number; - /** - * Wall-clock instant the record's `status` last changed. For - * `completed` / `failed` this is when the dream actually finished; - * for `cancelled` this is the moment `cancelTask` ran (NOT when - * the fork agent finishes unwinding — that can lag by seconds for - * agents mid-tool-call). The dialog renders elapsed from this - * value, so a freshly-cancelled record snaps to "Stopped · Ns" - * even while the underlying fork is still releasing the lock. - */ - endTime?: number; - progressText?: string; - error?: string; - /** Number of sessions the dream is reviewing — populated on schedule. */ - sessionCount?: number; - /** Memory topic files written — populated on completion. */ - touchedTopics?: readonly string[]; - /** - * Best-effort warnings populated by `runDream` when post-fork - * housekeeping fails (gating-metadata write or consolidation-lock - * release). The dream itself completed successfully — these are - * informational so the user can explain why subsequent dreams may - * be silently skipped as `'locked'` or why the scheduler gate - * isn't seeing the most recent dream's timestamp. - */ - lockReleaseError?: string; - metadataWriteError?: string; -}; +export type DreamDialogEntry = DreamTask; /** * A unified view-model entry the dialog/pill/context render against. * Discriminated by `kind`; per-kind fields are inlined verbatim so * renderer code can stay mechanical (`entry.kind === 'agent'` / - * `'shell'` / `'monitor'` / `'dream'` guard, then access fields directly). + * `'shell'` / `'monitor'` / `'dream'` guard, then access fields + * directly). * * The `agent`/`shell`/`monitor` arms are the core `TaskState` union - * member — `kind` lives on the core entry, so the merge step here no - * longer tags it. `dream` remains adapted from `MemoryManager` and is - * unioned in here while the dream task placement is decided in PR 2. + * member, held by `TaskRegistry`. The `dream` arm comes from the + * dream-task adapter (`tasks/dream-task.ts`), which synthesizes its + * view-model from `MemoryManager` records. */ -export type DialogEntry = TaskState | DreamDialogEntry; +export type DialogEntry = TaskState | DreamTask; export interface UseBackgroundTaskViewResult { entries: readonly DialogEntry[]; @@ -127,6 +80,28 @@ export function entryId(entry: DialogEntry): string { } } +/** + * Signature of the registry fields the dialog list / pill renderers + * depend on: id, kind, status. Activity bursts and event-count bumps + * mutate in place and don't change this signature, so the hook can + * suppress both the merge work and the setEntries fan-out for them. + * Per-entry surfaces (LiveAgentPanel, the dialog's selected-entry + * tick) carry their own subscriptions. + * + * Computed directly from `TaskState`s (not merged + dreams) so the + * cheap probe inside the registry subscription doesn't pay the + * `listDreamTasks` + MemoryManager cost on every `fireChange`. Dreams + * have their own subscription with a separate signature + * (`dreamSnapshotSignature`). + */ +function registrySnapshotShape(entries: readonly TaskState[]): string { + let sig = ''; + for (const e of entries) { + sig += `${e.kind}:${e.id}:${e.status}|`; + } + return sig; +} + export function useBackgroundTaskView( config: Config | null, ): UseBackgroundTaskViewResult { @@ -134,174 +109,74 @@ export function useBackgroundTaskView( useEffect(() => { if (!config) return; - const agentRegistry = config.getBackgroundTaskRegistry(); - const shellRegistry = config.getBackgroundShellRegistry(); - const monitorRegistry = config.getMonitorRegistry(); + const registry = config.getTaskRegistry(); const memoryManager = config.getMemoryManager(); const projectRoot = config.getProjectRoot(); - // Dream snapshot signature, kept as a defense-in-depth dedup for - // the dream-filtered memory listener below. The taskType filter - // already skips the listener entirely on extract notifies; this - // signature additionally absorbs the rare case where dream - // metadata is updated without an observable dialog change. + let lastRegistryShape = ''; let lastDreamSig = ''; - // Declared before `refresh` so the function ordering can't trip - // the temporal-dead-zone if a future refactor adds a synchronous - // call to refresh between the two `const` bindings. - const computeDreamSig = (dreams: readonly MemoryTaskRecord[]): string => - dreams.map((t) => `${t.id}:${t.status}:${t.updatedAt}`).join('|'); - - // refresh accepts a pre-fetched dream snapshot so the memory - // listener can reuse the same array it computed for its dedup - // check — avoids a second listTasksByType call AND eliminates the - // race window where the listener's gate sig and the entries it - // builds would otherwise come from two separate snapshots. - const refresh = (dreamSnapshot?: readonly MemoryTaskRecord[]) => { - const agentEntries: AgentTask[] = [...agentRegistry.getAll()]; - const shellEntries: ShellTask[] = [...shellRegistry.getAll()]; - const monitorEntries: MonitorTask[] = [...monitorRegistry.getAll()]; - // Dream entries: only surface tasks that actually fired. - // `pending` is a sub-second transition state and `skipped` - // records arise from the rare race where the schedule-time - // lock check passed but `acquireDreamLock` then hit EEXIST in - // runDream — these never reflect user-visible work, so filter - // them out. (Most gate misses don't create a record at all; - // scheduleDream returns `{status: 'skipped'}` early without - // touching the task map.) Extract tasks also intentionally - // stay out of this view — they fire on every UserQuery and - // their completion is already covered by the `memory_saved` - // toast in useGeminiStream. - // - // Cap retained terminal entries — MemoryManager.tasks Map has no - // eviction path, so completed/failed dreams accumulate forever - // (every fired dream over the project's lifetime). Without this - // cap the dialog would grow unbounded; with it the user sees all - // running dreams plus the most recent few terminal results - // (mirrors MonitorRegistry.MAX_RETAINED_TERMINAL_MONITORS). - const allDreams = - dreamSnapshot ?? memoryManager.listTasksByType('dream', projectRoot); - const runningDreams = allDreams.filter((t) => t.status === 'running'); - const terminalDreams = allDreams - .filter( - (t) => - t.status === 'completed' || - t.status === 'failed' || - t.status === 'cancelled', - ) - .sort((a, b) => b.updatedAt.localeCompare(a.updatedAt)) - .slice(0, MAX_RETAINED_TERMINAL_DREAMS); - const dreamEntries: DialogEntry[] = [ - ...runningDreams, - ...terminalDreams, - ].map((t) => { - const sessionCount = t.metadata?.['sessionCount']; - const touchedTopics = t.metadata?.['touchedTopics']; - const lockReleaseError = t.metadata?.['lockReleaseError']; - const metadataWriteError = t.metadata?.['metadataWriteError']; - return { - kind: 'dream' as const, - dreamId: t.id, - status: t.status as 'running' | 'completed' | 'failed' | 'cancelled', - startTime: Date.parse(t.createdAt), - endTime: t.status === 'running' ? undefined : Date.parse(t.updatedAt), - progressText: t.progressText, - error: t.error, - sessionCount: - typeof sessionCount === 'number' ? sessionCount : undefined, - touchedTopics: Array.isArray(touchedTopics) - ? (touchedTopics.filter((s) => typeof s === 'string') as string[]) - : undefined, - lockReleaseError: - typeof lockReleaseError === 'string' ? lockReleaseError : undefined, - metadataWriteError: - typeof metadataWriteError === 'string' - ? metadataWriteError - : undefined, - }; - }); - // Two-bucket merge so "new OR running tasks should appear at the - // top" (the literal phrasing of the issue this view-model serves). - // A pure startTime DESC sort surfaces the newest LAUNCH but lets - // an older long-running / paused entry fall below a batch of - // newer terminal entries — the user opens the dialog wanting to - // check the running work, and finds it buried under noise. - // - // bucket 1 — active (running + paused), sorted by startTime DESC - // so the most recent launch sits at the very top. - // bucket 2 — terminal (completed / failed / cancelled), sorted - // by endTime DESC so the most recently FINISHED entry - // is the first terminal row (matches "what changed - // while I wasn't looking" intuition; startTime would - // put a long-running task that just settled below an - // old quick task that finished hours ago). - // - // Entries falling out the bottom of bucket 2 are eventually - // pruned by each registry's terminal-entry cap (see - // `MAX_RETAINED_TERMINAL_AGENTS` / `MAX_RETAINED_TERMINAL_SHELLS` - // / `MAX_RETAINED_TERMINAL_MONITORS`). - const isActive = (entry: DialogEntry): boolean => - entry.status === 'running' || entry.status === 'paused'; - const merged = [ - ...agentEntries, - ...shellEntries, - ...monitorEntries, - ...dreamEntries, + // Two-bucket merge so "new OR running tasks should appear at the + // top" (the literal phrasing of the issue this view-model serves). + // A pure startTime sort surfaces the newest LAUNCH but lets an older + // long-running / paused entry fall below a batch of newer terminal + // entries — the user opens the dialog wanting the running work and + // finds it buried under noise. + // + // bucket 1 — active (running + paused), startTime DESC so the most + // recent launch sits at the very top. + // bucket 2 — terminal (completed / failed / cancelled), endTime + // DESC so the most recently FINISHED entry is the first + // terminal row. + // + // Entries falling out the bottom of bucket 2 are pruned by each + // kind's terminal-entry cap (MAX_RETAINED_TERMINAL_AGENTS / + // _SHELLS / _MONITORS) and, for dreams, by listDreamTasks' + // MAX_RETAINED_TERMINAL_DREAMS. + const isActive = (entry: DialogEntry): boolean => + entry.status === 'running' || entry.status === 'paused'; + const buildMerged = (): DialogEntry[] => + [ + ...registry.getAll(), + ...listDreamTasks(memoryManager, projectRoot), ].sort((a, b) => { const aActive = isActive(a); const bActive = isActive(b); if (aActive !== bActive) return aActive ? -1 : 1; if (aActive) return b.startTime - a.startTime; // Terminal bucket: fall back to startTime when an entry has no - // endTime yet (defensive — the registries stamp endTime on - // every running → terminal transition, so this only matters - // for synthetic / partially-restored entries). + // endTime yet (defensive — the registries stamp endTime on every + // running → terminal transition). return (b.endTime ?? b.startTime) - (a.endTime ?? a.startTime); }); - // Cache the dream signature derived from the freshly-built - // entries — the memory listener uses this to skip redundant - // setEntries calls when an extract notify fires (extract has no - // dialog surface, so the merged result is identical). Computed - // from the same `allDreams` snapshot used to build dreamEntries - // so the gate value can never desync from what's on screen. - lastDreamSig = computeDreamSig(allDreams); - setEntries(merged); - }; - - // Wrap registry callbacks in a thunk so React's setStatusChange - // signature (no-arg) doesn't accidentally pass an entry into - // refresh's `dreamSnapshot` parameter. - const refreshFromRegistry = () => refresh(); - refresh(); - - agentRegistry.setStatusChangeCallback(refreshFromRegistry); - shellRegistry.setStatusChangeCallback(refreshFromRegistry); - monitorRegistry.setStatusChangeCallback(refreshFromRegistry); + const initial = buildMerged(); + setEntries(initial); + lastRegistryShape = registrySnapshotShape(registry.getAll()); + lastDreamSig = dreamSnapshotSignature( + memoryManager.listTasksByType('dream', projectRoot), + ); + + const unsubscribeRegistry = registry.subscribe(() => { + // Cheap probe BEFORE buildMerged to avoid the `listDreamTasks` + + // sort cost on activity bursts and monitor event bumps, which + // mutate in place and don't change `id+kind+status`. + const shape = registrySnapshotShape(registry.getAll()); + if (shape === lastRegistryShape) return; + lastRegistryShape = shape; + setEntries(buildMerged()); + }); - // Memory listener fires only on dream-task transitions — - // `subscribe({ taskType: 'dream' })` skips the per-extract notify - // entirely so we don't pay the per-UserQuery O(n) signature cost - // for transitions we have no surface for. The dream-content - // signature dedup remains as a second-line guard against the rare - // case where dream metadata is updated without observable changes - // to the dialog (e.g. a future progressText-only patch on the - // same status). The fetched snapshot is forwarded to refresh so - // both the gate and the rendered dreamEntries come from one read. - const memoryListener = () => { - const dreams = memoryManager.listTasksByType('dream', projectRoot); - const sig = computeDreamSig(dreams); + const unsubscribeMemory = subscribeDreams(memoryManager, () => { + const sig = dreamSnapshotSignature( + memoryManager.listTasksByType('dream', projectRoot), + ); if (sig === lastDreamSig) return; - refresh(dreams); - }; - const unsubscribeMemory = memoryManager.subscribe(memoryListener, { - taskType: 'dream', + lastDreamSig = sig; + setEntries(buildMerged()); }); return () => { - agentRegistry.setStatusChangeCallback(undefined); - shellRegistry.setStatusChangeCallback(undefined); - monitorRegistry.setStatusChangeCallback(undefined); + unsubscribeRegistry(); unsubscribeMemory(); }; }, [config]); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index f077712b407..7660757fd70 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -92,6 +92,7 @@ const mockGetActiveGoal = vi.hoisted(() => vi.fn()); const mockActiveGoalEquals = vi.hoisted(() => vi.fn()); const mockSetActiveGoal = vi.hoisted(() => vi.fn()); const mockClearActiveGoal = vi.hoisted(() => vi.fn()); +const mockSetShellNotificationCallback = vi.hoisted(() => vi.fn()); vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { const actualCoreModule = (await importOriginal()) as any; @@ -107,6 +108,7 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { activeGoalEquals: mockActiveGoalEquals, setActiveGoal: mockSetActiveGoal, clearActiveGoal: mockClearActiveGoal, + setShellNotificationCallback: mockSetShellNotificationCallback, }; }); @@ -165,7 +167,6 @@ describe('useGeminiStream', () => { let mockScheduleToolCalls: Mock; let mockCancelAllToolCalls: Mock; let mockMarkToolsAsSubmitted: Mock; - let mockBackgroundShellRegistry: { setNotificationCallback: Mock }; let handleAtCommandSpy: MockInstance; beforeEach(() => { @@ -194,9 +195,6 @@ describe('useGeminiStream', () => { vertexai: false, authType: AuthType.USE_GEMINI, }; - mockBackgroundShellRegistry = { - setNotificationCallback: vi.fn(), - }; mockConfig = { apiKey: 'test-api-key', @@ -242,12 +240,15 @@ describe('useGeminiStream', () => { getCronScheduler: vi.fn(() => null), getEmitToolUseSummaries: vi.fn(() => false), getFastModel: vi.fn(() => undefined), - getBackgroundTaskRegistry: vi.fn(() => ({ - setNotificationCallback: vi.fn(), - })), - getBackgroundShellRegistry: vi.fn(() => mockBackgroundShellRegistry), - getMonitorRegistry: vi.fn(() => ({ - setNotificationCallback: vi.fn(), + getTaskRegistry: vi.fn(() => ({ + getAll: () => [], + getByKind: () => [], + get: () => undefined, + register: () => undefined, + update: () => undefined, + evict: () => undefined, + kill: () => undefined, + subscribe: () => () => {}, })), } as unknown as Config; mockOnDebugMessage = vi.fn(); @@ -379,13 +380,16 @@ describe('useGeminiStream', () => { '\nshell\ncompleted\n'; await waitFor(() => { - expect( - mockBackgroundShellRegistry.setNotificationCallback, - ).toHaveBeenCalledWith(expect.any(Function)); + expect(mockSetShellNotificationCallback).toHaveBeenCalledWith( + expect.anything(), + expect.any(Function), + ); }); - const callback = mockBackgroundShellRegistry.setNotificationCallback.mock - .calls[0][0] as (displayText: string, modelText: string) => void; + const callback = mockSetShellNotificationCallback.mock.calls[0][1] as ( + displayText: string, + modelText: string, + ) => void; act(() => { callback(displayText, modelText); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index f7d53a500fd..ef89f3068a3 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -56,6 +56,9 @@ import { activeGoalEquals, setActiveGoal, clearActiveGoal, + setAgentNotificationCallback, + setMonitorNotificationCallback, + setShellNotificationCallback, } from '@qwen-code/qwen-code-core'; import { type Part, type PartListUnion, FinishReason } from '@google/genai'; import type { @@ -2558,8 +2561,8 @@ export const useGeminiStream = ( // Register background agent notification callback onto the shared queue. useEffect(() => { - const registry = config.getBackgroundTaskRegistry(); - registry.setNotificationCallback((displayText, modelText) => { + const taskRegistry = config.getTaskRegistry(); + setAgentNotificationCallback(taskRegistry, (displayText, modelText) => { notificationQueueRef.current.push({ displayText, modelText, @@ -2568,14 +2571,14 @@ export const useGeminiStream = ( setNotificationTrigger((n) => n + 1); }); return () => { - registry.setNotificationCallback(undefined); + setAgentNotificationCallback(taskRegistry, undefined); }; }, [config]); // Register background shell terminal notification callback onto the shared queue. useEffect(() => { - const registry = config.getBackgroundShellRegistry(); - registry.setNotificationCallback((displayText, modelText) => { + const taskRegistry = config.getTaskRegistry(); + setShellNotificationCallback(taskRegistry, (displayText, modelText) => { notificationQueueRef.current.push({ displayText, modelText, @@ -2584,14 +2587,14 @@ export const useGeminiStream = ( setNotificationTrigger((n) => n + 1); }); return () => { - registry.setNotificationCallback(undefined); + setShellNotificationCallback(taskRegistry, undefined); }; }, [config]); // Register monitor notification callback onto the shared queue. useEffect(() => { - const registry = config.getMonitorRegistry(); - registry.setNotificationCallback((displayText, modelText) => { + const taskRegistry = config.getTaskRegistry(); + setMonitorNotificationCallback(taskRegistry, (displayText, modelText) => { notificationQueueRef.current.push({ displayText, modelText, @@ -2600,7 +2603,7 @@ export const useGeminiStream = ( setNotificationTrigger((n) => n + 1); }); return () => { - registry.setNotificationCallback(undefined); + setMonitorNotificationCallback(taskRegistry, undefined); }; }, [config]); diff --git a/packages/cli/src/ui/hooks/useResumeCommand.test.ts b/packages/cli/src/ui/hooks/useResumeCommand.test.ts index eefbe3ea98c..c5e6f86c13c 100644 --- a/packages/cli/src/ui/hooks/useResumeCommand.test.ts +++ b/packages/cli/src/ui/hooks/useResumeCommand.test.ts @@ -66,6 +66,16 @@ vi.mock('@qwen-code/qwen-code-core', () => { return { SessionService, + // The unified-registry helpers backgroundWorkUtils.ts now imports + // from core. Default to "no work" for the resume tests that don't + // exercise the blocking path; the resume-blocking tests below + // override these via vi.mocked(...). + agentHasUnfinalizedTasks: vi.fn(() => false), + agentReset: vi.fn(), + monitorReset: vi.fn(), + shellReset: vi.fn(), + getRunningMonitorTasks: vi.fn(() => []), + shellHasRunningEntries: vi.fn(() => false), }; }); @@ -156,24 +166,24 @@ describe('useResumeCommand', () => { const geminiClient = { initialize: vi.fn().mockResolvedValue(undefined), }; - const resetMonitorRegistry = vi.fn(); + // resetBackgroundStateForSessionSwitch dispatches to the kind-local + // reset helpers in core; we assert against those module-level mocks. + const core = await import('@qwen-code/qwen-code-core'); + vi.mocked(core.monitorReset).mockClear(); const config = { getTargetDir: () => '/tmp', getGeminiClient: () => geminiClient, startNewSession: vi.fn(), - getBackgroundTaskRegistry: () => ({ - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - reset: vi.fn(), - }), - getBackgroundShellRegistry: () => ({ - getAll: vi.fn().mockReturnValue([]), - hasRunningEntries: vi.fn().mockReturnValue(false), - reset: vi.fn(), - }), - getMonitorRegistry: () => ({ - getRunning: vi.fn().mockReturnValue([]), - reset: resetMonitorRegistry, + getTaskRegistry: () => ({ + getAll: () => [], + getByKind: () => [], + get: () => undefined, + register: () => undefined, + update: () => undefined, + evict: () => undefined, + kill: () => undefined, + subscribe: () => () => {}, }), loadPausedBackgroundAgents: vi.fn().mockResolvedValue([]), getBackgroundAgentResumeService: () => ({ @@ -228,7 +238,7 @@ describe('useResumeCommand', () => { expect(geminiClient.initialize).toHaveBeenCalledWith(); expect(historyManager.clearItems).toHaveBeenCalledTimes(1); expect(historyManager.loadHistory).toHaveBeenCalledTimes(1); - expect(resetMonitorRegistry).toHaveBeenCalledTimes(1); + expect(core.monitorReset).toHaveBeenCalledTimes(1); // Goal must be re-armed under the resumed sessionId so the in-memory // activeGoalStore entry (potentially stale across /new + /resume) gets // a fresh setAt / hookId / observer — otherwise the footer pill ticks @@ -258,18 +268,15 @@ describe('useResumeCommand', () => { getTargetDir: () => '/tmp', getGeminiClient: () => geminiClient, startNewSession: vi.fn(), - getBackgroundTaskRegistry: () => ({ - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - reset: vi.fn(), - }), - getBackgroundShellRegistry: () => ({ - getAll: vi.fn().mockReturnValue([]), - hasRunningEntries: vi.fn().mockReturnValue(false), - reset: vi.fn(), - }), - getMonitorRegistry: () => ({ - getRunning: vi.fn().mockReturnValue([]), - reset: vi.fn(), + getTaskRegistry: () => ({ + getAll: () => [], + getByKind: () => [], + get: () => undefined, + register: () => undefined, + update: () => undefined, + evict: () => undefined, + kill: () => undefined, + subscribe: () => () => {}, }), loadPausedBackgroundAgents: vi .fn() @@ -309,6 +316,9 @@ describe('useResumeCommand', () => { }); it('blocks resume when the current session still has running background work', async () => { + const core = await import('@qwen-code/qwen-code-core'); + vi.mocked(core.agentHasUnfinalizedTasks).mockReturnValueOnce(true); + const historyManager = { addItem: vi.fn(), clearItems: vi.fn(), @@ -317,18 +327,15 @@ describe('useResumeCommand', () => { const startNewSession = vi.fn(); const config = { - getBackgroundTaskRegistry: () => ({ - hasUnfinalizedTasks: vi.fn().mockReturnValue(true), - reset: vi.fn(), - }), - getBackgroundShellRegistry: () => ({ - getAll: vi.fn().mockReturnValue([]), - hasRunningEntries: vi.fn().mockReturnValue(false), - reset: vi.fn(), - }), - getMonitorRegistry: () => ({ - getRunning: vi.fn().mockReturnValue([]), - reset: vi.fn(), + getTaskRegistry: () => ({ + getAll: () => [], + getByKind: () => [], + get: () => undefined, + register: () => undefined, + update: () => undefined, + evict: () => undefined, + kill: () => undefined, + subscribe: () => () => {}, }), getTargetDir: () => '/tmp', getDebugLogger: () => ({ @@ -368,6 +375,14 @@ describe('useResumeCommand', () => { }); it('blocks resume when the current session still has a running monitor', async () => { + const core = await import('@qwen-code/qwen-code-core'); + vi.mocked(core.getRunningMonitorTasks).mockReturnValueOnce([ + // Just need length > 0; the helper's other fields aren't read by + // hasBlockingBackgroundWork. + // eslint-disable-next-line @typescript-eslint/no-explicit-any + { monitorId: 'mon_123', status: 'running' } as any, + ]); + const historyManager = { addItem: vi.fn(), clearItems: vi.fn(), @@ -376,23 +391,15 @@ describe('useResumeCommand', () => { const startNewSession = vi.fn(); const config = { - getBackgroundTaskRegistry: () => ({ - hasUnfinalizedTasks: vi.fn().mockReturnValue(false), - reset: vi.fn(), - }), - getBackgroundShellRegistry: () => ({ - getAll: vi.fn().mockReturnValue([]), - hasRunningEntries: vi.fn().mockReturnValue(false), - reset: vi.fn(), - }), - getMonitorRegistry: () => ({ - getRunning: vi.fn().mockReturnValue([ - { - monitorId: 'mon_123', - status: 'running', - }, - ]), - reset: vi.fn(), + getTaskRegistry: () => ({ + getAll: () => [], + getByKind: () => [], + get: () => undefined, + register: () => undefined, + update: () => undefined, + evict: () => undefined, + kill: () => undefined, + subscribe: () => () => {}, }), getTargetDir: () => '/tmp', getDebugLogger: () => ({ diff --git a/packages/cli/src/ui/utils/backgroundWorkUtils.test.ts b/packages/cli/src/ui/utils/backgroundWorkUtils.test.ts index cee3841248a..85c3998aa34 100644 --- a/packages/cli/src/ui/utils/backgroundWorkUtils.test.ts +++ b/packages/cli/src/ui/utils/backgroundWorkUtils.test.ts @@ -4,113 +4,131 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { describe, it, expect, vi } from 'vitest'; -import type { Config } from '@qwen-code/qwen-code-core'; +import { describe, it, expect } from 'vitest'; +import { + agentRegister, + monitorRegister, + shellRegister, + TaskRegistry, + type Config, +} from '@qwen-code/qwen-code-core'; import { hasBlockingBackgroundWork, resetBackgroundStateForSessionSwitch, } from './backgroundWorkUtils.js'; -function createMockConfig(overrides?: { - hasUnfinalizedTasks?: boolean; - runningMonitors?: unknown[]; - hasRunningEntries?: boolean; -}): Config { +function createMockConfig(registry: TaskRegistry): Config { return { - getBackgroundTaskRegistry: () => ({ - hasUnfinalizedTasks: () => overrides?.hasUnfinalizedTasks ?? false, - reset: vi.fn(), - }), - getMonitorRegistry: () => ({ - getRunning: () => overrides?.runningMonitors ?? [], - reset: vi.fn(), - }), - getBackgroundShellRegistry: () => ({ - hasRunningEntries: () => overrides?.hasRunningEntries ?? false, - reset: vi.fn(), - }), + getTaskRegistry: () => registry, } as unknown as Config; } describe('hasBlockingBackgroundWork', () => { it('returns false when nothing is running', () => { - expect(hasBlockingBackgroundWork(createMockConfig())).toBe(false); - }); - - it('returns true when background tasks are unfinalized', () => { - expect( - hasBlockingBackgroundWork( - createMockConfig({ hasUnfinalizedTasks: true }), - ), - ).toBe(true); + const registry = new TaskRegistry(); + expect(hasBlockingBackgroundWork(createMockConfig(registry))).toBe(false); }); - it('returns true when monitors are running', () => { - expect( - hasBlockingBackgroundWork( - createMockConfig({ runningMonitors: [{ id: 'm1' }] }), - ), - ).toBe(true); + it('returns true when a backgrounded agent is still running (unfinalized)', () => { + const registry = new TaskRegistry(); + agentRegister(registry, { + agentId: 'a1', + description: 'agent', + isBackgrounded: true, + status: 'running', + startTime: Date.now(), + abortController: new AbortController(), + outputFile: '/tmp/a.jsonl', + }); + expect(hasBlockingBackgroundWork(createMockConfig(registry))).toBe(true); }); - it('returns true when shell entries are running', () => { - expect( - hasBlockingBackgroundWork(createMockConfig({ hasRunningEntries: true })), - ).toBe(true); + it('returns true when a monitor is running', () => { + const registry = new TaskRegistry(); + monitorRegister(registry, { + monitorId: 'm1', + description: 'monitor', + command: 'tail -f log', + status: 'running', + startTime: Date.now(), + abortController: new AbortController(), + outputFile: '/tmp/m.log', + eventCount: 0, + lastEventTime: 0, + maxEvents: 100, + idleTimeoutMs: 60000, + droppedLines: 0, + }); + expect(hasBlockingBackgroundWork(createMockConfig(registry))).toBe(true); }); - it('short-circuits: does not check monitors or shells when tasks are unfinalized', () => { - const config = { - getBackgroundTaskRegistry: () => ({ - hasUnfinalizedTasks: () => true, - reset: vi.fn(), - }), - getMonitorRegistry: () => { - throw new Error('should not be called'); - }, - getBackgroundShellRegistry: () => { - throw new Error('should not be called'); - }, - } as unknown as Config; - - expect(hasBlockingBackgroundWork(config)).toBe(true); + it('returns true when a shell is running', () => { + const registry = new TaskRegistry(); + shellRegister(registry, { + shellId: 's1', + command: 'sleep 30', + cwd: '/tmp', + status: 'running', + startTime: Date.now(), + outputPath: '/tmp/s.out', + abortController: new AbortController(), + }); + expect(hasBlockingBackgroundWork(createMockConfig(registry))).toBe(true); }); - it('short-circuits: does not check shells when monitors are running', () => { - const config = { - getBackgroundTaskRegistry: () => ({ - hasUnfinalizedTasks: () => false, - reset: vi.fn(), - }), - getMonitorRegistry: () => ({ - getRunning: () => [{ id: 'm1' }], - reset: vi.fn(), - }), - getBackgroundShellRegistry: () => { - throw new Error('should not be called'); - }, - } as unknown as Config; - - expect(hasBlockingBackgroundWork(config)).toBe(true); + it('returns false once the only running entry has settled', () => { + const registry = new TaskRegistry(); + shellRegister(registry, { + shellId: 's1', + command: 'sleep 30', + cwd: '/tmp', + status: 'completed', + startTime: Date.now(), + outputPath: '/tmp/s.out', + abortController: new AbortController(), + }); + expect(hasBlockingBackgroundWork(createMockConfig(registry))).toBe(false); }); }); describe('resetBackgroundStateForSessionSwitch', () => { - it('calls reset on all three registries', () => { - const resetTasks = vi.fn(); - const resetMonitors = vi.fn(); - const resetShells = vi.fn(); - - const config = { - getBackgroundTaskRegistry: () => ({ reset: resetTasks }), - getMonitorRegistry: () => ({ reset: resetMonitors }), - getBackgroundShellRegistry: () => ({ reset: resetShells }), - } as unknown as Config; + it('clears every kind from the registry', () => { + const registry = new TaskRegistry(); + agentRegister(registry, { + agentId: 'a1', + description: 'agent', + isBackgrounded: false, + status: 'paused', + startTime: Date.now(), + abortController: new AbortController(), + outputFile: '/tmp/a.jsonl', + }); + shellRegister(registry, { + shellId: 's1', + command: 'sleep 30', + cwd: '/tmp', + status: 'completed', + startTime: Date.now(), + outputPath: '/tmp/s.out', + abortController: new AbortController(), + }); + monitorRegister(registry, { + monitorId: 'm1', + description: 'monitor', + command: 'tail -f log', + status: 'completed', + startTime: Date.now(), + abortController: new AbortController(), + outputFile: '/tmp/m.log', + eventCount: 0, + lastEventTime: 0, + maxEvents: 100, + idleTimeoutMs: 60000, + droppedLines: 0, + }); - resetBackgroundStateForSessionSwitch(config); + resetBackgroundStateForSessionSwitch(createMockConfig(registry)); - expect(resetTasks).toHaveBeenCalledOnce(); - expect(resetMonitors).toHaveBeenCalledOnce(); - expect(resetShells).toHaveBeenCalledOnce(); + expect(registry.getAll()).toEqual([]); }); }); diff --git a/packages/cli/src/ui/utils/backgroundWorkUtils.ts b/packages/cli/src/ui/utils/backgroundWorkUtils.ts index 9e7b4d71c9f..106934fdf1f 100644 --- a/packages/cli/src/ui/utils/backgroundWorkUtils.ts +++ b/packages/cli/src/ui/utils/backgroundWorkUtils.ts @@ -4,18 +4,28 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { Config } from '@qwen-code/qwen-code-core'; +import { + agentHasUnfinalizedTasks, + agentReset, + getRunningMonitorTasks, + monitorReset, + shellHasRunningEntries, + shellReset, + type Config, +} from '@qwen-code/qwen-code-core'; export function hasBlockingBackgroundWork(config: Config): boolean { + const registry = config.getTaskRegistry(); return ( - config.getBackgroundTaskRegistry().hasUnfinalizedTasks() || - config.getMonitorRegistry().getRunning().length > 0 || - config.getBackgroundShellRegistry().hasRunningEntries() + agentHasUnfinalizedTasks(registry) || + getRunningMonitorTasks(registry).length > 0 || + shellHasRunningEntries(registry) ); } export function resetBackgroundStateForSessionSwitch(config: Config): void { - config.getBackgroundTaskRegistry().reset(); - config.getMonitorRegistry().reset(); - config.getBackgroundShellRegistry().reset(); + const registry = config.getTaskRegistry(); + agentReset(registry); + monitorReset(registry); + shellReset(registry); } diff --git a/packages/core/src/agents/background-agent-resume.test.ts b/packages/core/src/agents/background-agent-resume.test.ts index 4899a2f2961..bc8e2512fb3 100644 --- a/packages/core/src/agents/background-agent-resume.test.ts +++ b/packages/core/src/agents/background-agent-resume.test.ts @@ -9,7 +9,17 @@ import * as fs from 'node:fs'; import * as os from 'node:os'; import * as path from 'node:path'; import type { Config } from '../config/config.js'; -import { BackgroundTaskRegistry } from './background-tasks.js'; +import { TaskRegistry } from '../tasks/registry.js'; +import { + agentAbortAll, + agentCancel, + agentDrainMessages, + agentRegister, + agentWaitForMessages, + getAgentTask, + setAgentBackgroundCapForTest, +} from '../tasks/agent-task.js'; +import * as monitorTaskModule from '../tasks/monitor-task.js'; import { BackgroundAgentResumeService } from './background-agent-resume.js'; import { getAgentJsonlPath, @@ -26,14 +36,16 @@ import { describe('BackgroundAgentResumeService', () => { let tempDir: string; - let registry: BackgroundTaskRegistry; + let registry: TaskRegistry; beforeEach(() => { tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'bg-agent-resume-')); - registry = new BackgroundTaskRegistry(); + registry = new TaskRegistry(); }); afterEach(() => { + setAgentBackgroundCapForTest(undefined); + vi.restoreAllMocks(); fs.rmSync(tempDir, { recursive: true, force: true, @@ -69,17 +81,28 @@ describe('BackgroundAgentResumeService', () => { getAllToolNames: vi.fn().mockReturnValue([]), stop: vi.fn().mockResolvedValue(undefined), }; + // The collapsed task-registry architecture exposes + // `setMonitorAgentNotificationCallback` / `setMonitorAgentLifecycleCallback` + // / `monitorCancelRunningForOwner` as module-level free functions in + // `tasks/monitor-task.ts`, replacing the old `MonitorRegistry` methods. + // Spy on the namespace so assertions framed in terms of the old + // registry handle keep reading naturally. const monitorRegistry = { - setAgentNotificationCallback: vi.fn(), - setAgentLifecycleCallback: vi.fn(), - cancelRunningForOwner: vi.fn(), + setAgentNotificationCallback: vi + .spyOn(monitorTaskModule, 'setMonitorAgentNotificationCallback') + .mockImplementation(() => {}), + setAgentLifecycleCallback: vi + .spyOn(monitorTaskModule, 'setMonitorAgentLifecycleCallback') + .mockImplementation(() => {}), + cancelRunningForOwner: vi + .spyOn(monitorTaskModule, 'monitorCancelRunningForOwner') + .mockImplementation(() => {}), }; const config = { storage: { getProjectDir: () => tempDir, }, - getBackgroundTaskRegistry: () => registry, - getMonitorRegistry: () => monitorRegistry, + getTaskRegistry: () => registry, getSubagentManager: () => subagentManager, getHookSystem: () => hookSystem, getStopHookBlockingCap: () => options.stopHookBlockingCap ?? 8, @@ -396,7 +419,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Resume with hooks', subagentType: 'researcher', @@ -456,9 +479,8 @@ describe('BackgroundAgentResumeService', () => { }); it('can resume into the final background concurrency slot', async () => { - registry = new BackgroundTaskRegistry({ - maxConcurrentBackgroundAgents: 1, - }); + setAgentBackgroundCapForTest(1); + registry = new TaskRegistry(); const sessionId = 'session-resume-cap'; const agentId = 'agent-resume-cap'; const metaPath = getAgentMetaPath(tempDir, sessionId, agentId); @@ -488,7 +510,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Resume at cap', subagentType: 'researcher', @@ -526,15 +548,14 @@ describe('BackgroundAgentResumeService', () => { }); it('keeps a paused agent paused when resume cannot claim a background slot', async () => { - registry = new BackgroundTaskRegistry({ - maxConcurrentBackgroundAgents: 1, - }); + setAgentBackgroundCapForTest(1); + registry = new TaskRegistry(); const sessionId = 'session-resume-full'; const agentId = 'agent-resume-full'; const metaPath = getAgentMetaPath(tempDir, sessionId, agentId); const outputFile = getAgentJsonlPath(tempDir, sessionId, agentId); - registry.register({ + agentRegister(registry, { agentId: 'already-running', description: 'Already running', subagentType: 'researcher', @@ -568,7 +589,7 @@ describe('BackgroundAgentResumeService', () => { }) + '\n', 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Resume while full', subagentType: 'researcher', @@ -623,7 +644,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Resume stop hook path', subagentType: 'researcher', @@ -697,7 +718,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Resume cap path', subagentType: 'researcher', @@ -741,7 +762,7 @@ describe('BackgroundAgentResumeService', () => { }); expect(hookSystem.fireSubagentStopEvent).toHaveBeenCalledTimes(2); expect(subagent.execute).toHaveBeenCalledTimes(2); - expect(registry.get(agentId)?.result).toContain( + expect(getAgentTask(registry, agentId)?.result).toContain( 'SubagentStop hook blocked continuation 2 consecutive times; overriding and ending the turn.', ); }); @@ -783,7 +804,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Resume after trust revoked', subagentType: 'researcher', @@ -853,7 +874,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Resume once', subagentType: 'researcher', @@ -937,7 +958,7 @@ describe('BackgroundAgentResumeService', () => { }) + '\n', 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Resume monitor owner', subagentType: 'researcher', @@ -981,7 +1002,7 @@ describe('BackgroundAgentResumeService', () => { callback('Monitor "logs" event #1: ready', ''); - expect(registry.get(agentId)?.pendingMessages).toContainEqual({ + expect(getAgentTask(registry, agentId)?.pendingMessages).toContainEqual({ kind: 'notification', text: '', }); @@ -989,8 +1010,9 @@ describe('BackgroundAgentResumeService', () => { expect(subagent.setExternalMessageWaitPredicate).toHaveBeenCalled(); const lifecycleCallback = monitorRegistry.setAgentLifecycleCallback.mock .calls[0][1] as () => void; - registry.drainMessages(agentId); - const waitPromise = registry.waitForMessages( + agentDrainMessages(registry, agentId); + const waitPromise = agentWaitForMessages( + registry, agentId, new AbortController().signal, ); @@ -1010,6 +1032,7 @@ describe('BackgroundAgentResumeService', () => { undefined, ); expect(monitorRegistry.cancelRunningForOwner).toHaveBeenCalledWith( + registry, agentId, { notify: false, @@ -1050,7 +1073,7 @@ describe('BackgroundAgentResumeService', () => { }) + '\n', 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Resume monitor setup failure', subagentType: 'researcher', @@ -1100,6 +1123,7 @@ describe('BackgroundAgentResumeService', () => { undefined, ); expect(monitorRegistry.cancelRunningForOwner).toHaveBeenCalledWith( + registry, agentId, { notify: false, @@ -1179,7 +1203,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: launchPrompt, subagentType: FORK_SUBAGENT_TYPE, @@ -1275,7 +1299,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Legacy fork task', subagentType: FORK_SUBAGENT_TYPE, @@ -1294,7 +1318,7 @@ describe('BackgroundAgentResumeService', () => { expect(resumed).toBeUndefined(); expect(registry.get(agentId)?.status).toBe('paused'); - expect(registry.get(agentId)?.resumeBlockedReason).toContain( + expect(getAgentTask(registry, agentId)?.resumeBlockedReason).toContain( 'bootstrap transcript is missing', ); expect(registry.get(agentId)?.error).toBeUndefined(); @@ -1357,7 +1381,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Legacy fork task without capabilities', subagentType: FORK_SUBAGENT_TYPE, @@ -1376,7 +1400,7 @@ describe('BackgroundAgentResumeService', () => { expect(resumed).toBeUndefined(); expect(registry.get(agentId)?.status).toBe('paused'); - expect(registry.get(agentId)?.resumeBlockedReason).toContain( + expect(getAgentTask(registry, agentId)?.resumeBlockedReason).toContain( 'runtime constraints are missing', ); expect(createSpy).not.toHaveBeenCalled(); @@ -1399,7 +1423,7 @@ describe('BackgroundAgentResumeService', () => { resolvedApprovalMode: 'default', }); - registry.register({ + agentRegister(registry, { agentId, description: 'Interrupted by shutdown', subagentType: 'researcher', @@ -1412,7 +1436,7 @@ describe('BackgroundAgentResumeService', () => { isBackgrounded: true, }); - registry.abortAll(); + agentAbortAll(registry); expect(readMetaStatus(metaPath)).toBe('running'); }); @@ -1447,7 +1471,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Resume then shutdown', subagentType: 'researcher', @@ -1484,7 +1508,7 @@ describe('BackgroundAgentResumeService', () => { const resumed = await service.resumeBackgroundAgent(agentId, 'continue'); expect(resumed).toBeDefined(); - registry.abortAll(); + agentAbortAll(registry); releaseExecute?.(); await vi.waitFor(() => { expect(registry.get(agentId)?.status).toBe('cancelled'); @@ -1522,7 +1546,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Resume then cancel', subagentType: 'researcher', @@ -1559,7 +1583,7 @@ describe('BackgroundAgentResumeService', () => { const resumed = await service.resumeBackgroundAgent(agentId, 'continue'); expect(resumed).toBeDefined(); - registry.cancel(agentId); + agentCancel(registry, agentId); releaseExecute?.(); await vi.waitFor(() => { expect(registry.get(agentId)?.status).toBe('cancelled'); @@ -1615,7 +1639,7 @@ describe('BackgroundAgentResumeService', () => { 'utf8', ); - registry.register({ + agentRegister(registry, { agentId, description: 'Pending user tail', subagentType: 'researcher', diff --git a/packages/core/src/agents/background-agent-resume.ts b/packages/core/src/agents/background-agent-resume.ts index ce98722f396..6641459a094 100644 --- a/packages/core/src/agents/background-agent-resume.ts +++ b/packages/core/src/agents/background-agent-resume.ts @@ -43,7 +43,27 @@ import type { AgentCompletionStats, AgentTask, AgentTaskRegistration, -} from './background-tasks.js'; +} from '../tasks/agent-task.js'; +import { + agentAppendActivity, + agentAbandon, + agentComplete, + agentDrainMessages, + agentFail, + agentFinalizeCancelled, + agentQueueExternalInput, + agentQueueMessage, + agentRegister, + agentWaitForMessages, + agentWakeExternalInputWaiters, + getAgentTask, +} from '../tasks/agent-task.js'; +import { + monitorCancelRunningForOwner, + monitorHasRunningForOwner, + setMonitorAgentLifecycleCallback, + setMonitorAgentNotificationCallback, +} from '../tasks/monitor-task.js'; import type { SubagentConfig } from '../subagents/types.js'; import type { PromptConfig, @@ -377,7 +397,7 @@ export class BackgroundAgentResumeService { throw error; } - const registry = this.config.getBackgroundTaskRegistry(); + const registry = this.config.getTaskRegistry(); const recovered: AgentTask[] = []; for (const fileName of files) { @@ -386,7 +406,7 @@ export class BackgroundAgentResumeService { try { const meta = readAgentMeta(metaPath); if (!meta || meta.status !== 'running') continue; - if (registry.get(meta.agentId)) continue; + if (getAgentTask(registry, meta.agentId)) continue; const subagentName = meta.subagentName ?? meta.agentType; if (!subagentName) continue; const target = await this.resolveResumeTarget(subagentName); @@ -426,7 +446,7 @@ export class BackgroundAgentResumeService { meta.lastError === resumeBlockedReason ? undefined : meta.lastError, resumeBlockedReason, }; - const entry = registry.register(registration); + const entry = agentRegister(registry, registration); recovered.push(entry); } catch (error) { debugLogger.warn( @@ -447,8 +467,8 @@ export class BackgroundAgentResumeService { const existingOperation = this.resumeOperations.get(agentId); if (existingOperation) { if (trimmedMessage) { - const registry = this.config.getBackgroundTaskRegistry(); - if (!registry.queueMessage(agentId, trimmedMessage)) { + const registry = this.config.getTaskRegistry(); + if (!agentQueueMessage(registry, agentId, trimmedMessage)) { existingOperation.continuationMessages.push(trimmedMessage); } } @@ -473,8 +493,8 @@ export class BackgroundAgentResumeService { agentId: string, operation: ResumeOperation, ): Promise { - const registry = this.config.getBackgroundTaskRegistry(); - const existing = registry.get(agentId); + const registry = this.config.getTaskRegistry(); + const existing = getAgentTask(registry, agentId); if (!existing || existing.status !== 'paused') { return existing; } @@ -493,7 +513,7 @@ export class BackgroundAgentResumeService { const bgAbortController = new AbortController(); try { - registry.register({ + agentRegister(registry, { ...existing, status: 'running', abortController: bgAbortController, @@ -658,7 +678,7 @@ export class BackgroundAgentResumeService { }); const pendingMessages = [ - ...(registry.get(meta.agentId)?.pendingMessages ?? []), + ...(getAgentTask(registry, meta.agentId)?.pendingMessages ?? []), ]; const registration: AgentTaskRegistration = { ...existing, @@ -675,44 +695,43 @@ export class BackgroundAgentResumeService { recentActivities: [], pendingMessages, }; - const entry = registry.register(registration); + const entry = agentRegister(registry, registration); const lateContinuationMessages = operation.continuationMessages.slice( promptMessages.length, ); for (const message of lateContinuationMessages) { - registry.queueMessage(meta.agentId, message); + agentQueueMessage(registry, meta.agentId, message); } subagent.setExternalMessageProvider(() => - registry.drainMessages(meta.agentId), + agentDrainMessages(registry, meta.agentId), ); subagent.setExternalMessageWaiter?.((waitSignal) => - registry.waitForMessages(meta.agentId, waitSignal), + agentWaitForMessages(registry, meta.agentId, waitSignal), ); - const monitorRegistry = this.config.getMonitorRegistry(); subagent.setExternalMessageWaitPredicate?.(() => - monitorRegistry.hasRunningForOwner(meta.agentId), + monitorHasRunningForOwner(registry, meta.agentId), ); - monitorRegistry.setAgentNotificationCallback( + setMonitorAgentNotificationCallback( meta.agentId, - (_displayText, modelText) => - void registry.queueExternalInput(meta.agentId, { + (_displayText: string, modelText: string) => + void agentQueueExternalInput(registry, meta.agentId, { kind: 'notification', text: modelText, }), ); - monitorRegistry.setAgentLifecycleCallback(meta.agentId, () => - registry.wakeExternalInputWaiters(meta.agentId), + setMonitorAgentLifecycleCallback(meta.agentId, () => + agentWakeExternalInputWaiters(registry, meta.agentId), ); let cleanedUpOwnedMonitorNotifications = false; cleanupOwnedMonitorNotifications = () => { if (cleanedUpOwnedMonitorNotifications) return; cleanedUpOwnedMonitorNotifications = true; - monitorRegistry.cancelRunningForOwner(meta.agentId, { + monitorCancelRunningForOwner(registry, meta.agentId, { notify: false, }); - monitorRegistry.setAgentNotificationCallback(meta.agentId, undefined); - monitorRegistry.setAgentLifecycleCallback(meta.agentId, undefined); + setMonitorAgentNotificationCallback(meta.agentId, undefined); + setMonitorAgentLifecycleCallback(meta.agentId, undefined); }; const hookSystem = this.config.getHookSystem(); @@ -729,14 +748,14 @@ export class BackgroundAgentResumeService { let liveToolCallCount = 0; const refreshLiveStats = () => { - const target = registry.get(meta.agentId); + const target = getAgentTask(registry, meta.agentId); if (!target || target.status !== 'running') return; target.stats = getCompletionStats(subagent, liveToolCallCount); }; const onToolCall = (event: AgentToolCallEvent) => { liveToolCallCount += 1; refreshLiveStats(); - registry.appendActivity(meta.agentId, { + agentAppendActivity(registry, meta.agentId, { name: event.name, description: event.description, at: event.timestamp, @@ -771,23 +790,23 @@ export class BackgroundAgentResumeService { ); const stats = getCompletionStats(subagent, liveToolCallCount); if (terminateMode === AgentTerminateMode.GOAL) { - registry.complete(meta.agentId, finalText, stats); + agentComplete(registry, meta.agentId, finalText, stats); patchAgentMeta(metaPath, { status: 'completed', lastUpdatedAt: new Date().toISOString(), lastError: undefined, }); } else if (terminateMode === AgentTerminateMode.CANCELLED) { - registry.finalizeCancelled(meta.agentId, finalText, stats); + agentFinalizeCancelled(registry, meta.agentId, finalText, stats); persistBackgroundCancellation( metaPath, - registry.get(meta.agentId)?.persistedCancellationStatus ?? - 'cancelled', + getAgentTask(registry, meta.agentId) + ?.persistedCancellationStatus ?? 'cancelled', ); } else { const failureText = finalText || `Agent terminated with mode: ${terminateMode}`; - registry.fail(meta.agentId, failureText, stats); + agentFail(registry, meta.agentId, failureText, stats); patchAgentMeta(metaPath, { status: 'failed', lastUpdatedAt: new Date().toISOString(), @@ -801,18 +820,20 @@ export class BackgroundAgentResumeService { `[BackgroundAgentResume] Background agent failed: ${errorMessage}`, ); if (bgAbortController.signal.aborted) { - registry.finalizeCancelled( + agentFinalizeCancelled( + registry, meta.agentId, errorMessage, getCompletionStats(subagent, liveToolCallCount), ); persistBackgroundCancellation( metaPath, - registry.get(meta.agentId)?.persistedCancellationStatus ?? - 'cancelled', + getAgentTask(registry, meta.agentId) + ?.persistedCancellationStatus ?? 'cancelled', ); } else { - registry.fail( + agentFail( + registry, meta.agentId, errorMessage, getCompletionStats(subagent, liveToolCallCount), @@ -860,10 +881,10 @@ export class BackgroundAgentResumeService { lastError: errorMessage, lastUpdatedAt: new Date().toISOString(), }); - const latest = registry.get(agentId); + const latest = getAgentTask(registry, agentId); if (latest?.status === 'running') { if (latest.abortController.signal.aborted) { - registry.finalizeCancelled(agentId, errorMessage); + agentFinalizeCancelled(registry, agentId, errorMessage); } else { this.restorePausedEntry(agentId, { error: errorMessage }); } @@ -873,8 +894,8 @@ export class BackgroundAgentResumeService { } abandonBackgroundAgent(agentId: string): boolean { - const registry = this.config.getBackgroundTaskRegistry(); - const entry = registry.get(agentId); + const registry = this.config.getTaskRegistry(); + const entry = getAgentTask(registry, agentId); if (!entry || entry.status !== 'paused' || !entry.metaPath) { return false; } @@ -884,7 +905,7 @@ export class BackgroundAgentResumeService { lastUpdatedAt: new Date().toISOString(), lastError: undefined, }); - registry.abandon(agentId); + agentAbandon(registry, agentId); return true; } @@ -925,8 +946,8 @@ export class BackgroundAgentResumeService { agentId: string, options: RestorePausedEntryOptions = {}, ): AgentTask | undefined { - const registry = this.config.getBackgroundTaskRegistry(); - const latest = registry.get(agentId); + const registry = this.config.getTaskRegistry(); + const latest = getAgentTask(registry, agentId); if (!latest) return undefined; const registration: AgentTaskRegistration = { @@ -942,7 +963,7 @@ export class BackgroundAgentResumeService { recentActivities: [], pendingMessages: [...(latest.pendingMessages ?? [])], }; - return registry.register(registration); + return agentRegister(registry, registration); } private async createResumedForkSubagent( diff --git a/packages/core/src/agents/background-tasks.test.ts b/packages/core/src/agents/background-tasks.test.ts deleted file mode 100644 index fbb7319ba12..00000000000 --- a/packages/core/src/agents/background-tasks.test.ts +++ /dev/null @@ -1,1561 +0,0 @@ -/** - * @license - * Copyright 2025 Qwen - * SPDX-License-Identifier: Apache-2.0 - */ - -import { describe, it, expect, vi, beforeEach } from 'vitest'; -import { - BACKGROUND_AGENT_CONCURRENCY_ENV, - BackgroundTaskRegistry, - DEFAULT_MAX_CONCURRENT_BACKGROUND_AGENTS, - MAX_CONCURRENT_BACKGROUND_AGENTS, - MAX_RETAINED_TERMINAL_AGENTS, - resolveMaxConcurrentBackgroundAgents, - type AgentTaskRegistration, - type BackgroundTaskEntry, -} from './background-tasks.js'; -import * as transcript from './agent-transcript.js'; - -function makeRegistration( - agentId: string, - overrides: Partial = {}, -): AgentTaskRegistration { - return { - agentId, - description: agentId, - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: `/tmp/${agentId}.jsonl`, - ...overrides, - }; -} - -describe('BackgroundTaskRegistry', () => { - let registry: BackgroundTaskRegistry; - - beforeEach(() => { - registry = new BackgroundTaskRegistry(); - }); - - it('registers and retrieves a background agent', () => { - const entry = { - agentId: 'test-1', - description: 'test agent', - status: 'running' as const, - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }; - - registry.register(entry); - expect(registry.get('test-1')).toBe(entry); - }); - - it('completes a background agent and sends notification', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.complete('test-1', 'The result text'); - - const entry = registry.get('test-1')!; - expect(entry.status).toBe('completed'); - expect(entry.result).toBe('The result text'); - expect(entry.endTime).toBeDefined(); - expect(callback).toHaveBeenCalledOnce(); - const [displayText, modelText] = callback.mock.calls[0] as [string, string]; - // Display text: short summary without the full result - expect(displayText).toContain('completed'); - expect(displayText).toContain('test agent'); - expect(displayText).not.toContain('The result text'); - // Model text: full details including result for the LLM - expect(modelText).toContain('The result text'); - }); - - it('fails a background agent and sends notification', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.fail('test-1', 'Something went wrong'); - - const entry = registry.get('test-1')!; - expect(entry.status).toBe('failed'); - expect(entry.error).toBe('Something went wrong'); - expect(callback).toHaveBeenCalledOnce(); - const [displayText] = callback.mock.calls[0] as [string, string]; - expect(displayText).toContain('failed'); - }); - - it('cancels a running background agent without emitting a notification', () => { - // cancel() is intent-only: it aborts the signal and marks the entry - // cancelled, but does not emit a task-notification. The natural - // completion handler (bgBody) emits the terminal notification with - // the agent's real partial/final result via complete()/fail(). - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - const abortController = new AbortController(); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController, - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.cancel('test-1'); - - expect(registry.get('test-1')!.status).toBe('cancelled'); - expect(abortController.signal.aborted).toBe(true); - expect(callback).not.toHaveBeenCalled(); - }); - - it('persists explicit cancellations as cancelled sidecar state', () => { - const patchSpy = vi - .spyOn(transcript, 'patchAgentMeta') - .mockImplementation(() => undefined); - try { - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - metaPath: '/tmp/test-1.meta.json', - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.cancel('test-1'); - - expect(patchSpy).toHaveBeenCalledWith( - '/tmp/test-1.meta.json', - expect.objectContaining({ - status: 'cancelled', - lastError: undefined, - }), - ); - } finally { - patchSpy.mockRestore(); - } - }); - - it('emits a fallback cancelled notification after the grace period when the natural handler never runs', () => { - vi.useFakeTimers(); - try { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.cancel('test-1'); - expect(callback).not.toHaveBeenCalled(); - - // Pathological tool case: bgBody never emits. After the grace period - // the fallback fires so hasUnfinalizedTasks() stops reporting true - // and the headless wait loop can exit. - vi.runAllTimers(); - - expect(callback).toHaveBeenCalledOnce(); - const [, modelText] = callback.mock.calls[0] as [string, string]; - expect(modelText).toContain('cancelled'); - expect(registry.hasUnfinalizedTasks()).toBe(false); - } finally { - vi.useRealTimers(); - } - }); - - it('skips the fallback notification when the natural handler finalizes first', () => { - vi.useFakeTimers(); - try { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.cancel('test-1'); - // Natural handler wins the race with the partial result. - registry.finalizeCancelled('test-1', 'partial output'); - expect(callback).toHaveBeenCalledOnce(); - callback.mockClear(); - - vi.runAllTimers(); - - // Fallback lands on a notified entry and no-ops. - expect(callback).not.toHaveBeenCalled(); - } finally { - vi.useRealTimers(); - } - }); - - it('finalizeCancellationIfPending emits a fallback cancelled notification', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.cancel('test-1'); - registry.finalizeCancellationIfPending('test-1'); - - expect(callback).toHaveBeenCalledOnce(); - const [, modelText] = callback.mock.calls[0] as [string, string]; - expect(modelText).toContain('cancelled'); - }); - - it('complete() after the cancellation has already been notified is a no-op', () => { - // Once finalizeCancelled has emitted the terminal notification, a - // late-arriving complete() must not double-fire — the SDK contract - // is one notification per task_started. - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.cancel('test-1'); - registry.finalizeCancelled('test-1', 'partial'); - expect(callback).toHaveBeenCalledOnce(); - callback.mockClear(); - - registry.complete('test-1', 'late result'); - - expect(callback).not.toHaveBeenCalled(); - // Status stays cancelled — the notified terminal state wins. - expect(registry.get('test-1')!.status).toBe('cancelled'); - expect(registry.get('test-1')!.result).toBe('partial'); - }); - - it('does not cancel a non-running agent', () => { - const abortController = new AbortController(); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController, - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.complete('test-1', 'done'); - registry.cancel('test-1'); // should be a no-op - - expect(registry.get('test-1')!.status).toBe('completed'); - expect(abortController.signal.aborted).toBe(false); - }); - - it('abandons a paused agent without emitting a notification', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'paused-1', - description: 'paused agent', - status: 'paused', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.abandon('paused-1'); - - expect(registry.get('paused-1')!.status).toBe('cancelled'); - expect(registry.get('paused-1')!.notified).toBe(true); - expect(callback).not.toHaveBeenCalled(); - }); - - it('does not treat paused entries as unfinalized work', () => { - registry.register({ - agentId: 'paused-1', - description: 'paused agent', - status: 'paused', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - expect(registry.hasUnfinalizedTasks()).toBe(false); - }); - - it('lists running agents', () => { - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - registry.register({ - agentId: 'b', - description: 'agent b', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.complete('a', 'done'); - - const running = registry.getAll().filter((e) => e.status === 'running'); - expect(running).toHaveLength(1); - expect(running[0].agentId).toBe('b'); - }); - - describe('background concurrency limit', () => { - it('resolves the default and env override for the background agent cap', () => { - expect(resolveMaxConcurrentBackgroundAgents({})).toBe( - DEFAULT_MAX_CONCURRENT_BACKGROUND_AGENTS, - ); - expect( - resolveMaxConcurrentBackgroundAgents({ - [BACKGROUND_AGENT_CONCURRENCY_ENV]: '3', - }), - ).toBe(3); - expect( - resolveMaxConcurrentBackgroundAgents({ - [BACKGROUND_AGENT_CONCURRENCY_ENV]: '0', - }), - ).toBe(DEFAULT_MAX_CONCURRENT_BACKGROUND_AGENTS); - expect(MAX_CONCURRENT_BACKGROUND_AGENTS).toBeGreaterThanOrEqual(1); - }); - - it('rejects new running background agents once the cap is reached', () => { - registry = new BackgroundTaskRegistry({ - maxConcurrentBackgroundAgents: 2, - }); - - registry.register(makeRegistration('bg-1')); - registry.register(makeRegistration('bg-2')); - - expect(() => registry.register(makeRegistration('bg-3'))).toThrow( - 'Cannot start background agent: maximum concurrent background agents ' + - '(2) reached. Stop an existing agent first.', - ); - expect(registry.get('bg-3')).toBeUndefined(); - }); - - it('allows replacing the same running background agent at the cap', () => { - registry = new BackgroundTaskRegistry({ - maxConcurrentBackgroundAgents: 1, - }); - - registry.register(makeRegistration('bg-1')); - - expect(() => - registry.register( - makeRegistration('bg-1', { - prompt: 'resumed continuation', - }), - ), - ).not.toThrow(); - expect(registry.get('bg-1')?.prompt).toBe('resumed continuation'); - }); - - it('does not count foreground, paused, or terminal entries toward the cap', () => { - registry = new BackgroundTaskRegistry({ - maxConcurrentBackgroundAgents: 1, - }); - - registry.register( - makeRegistration('fg-1', { - isBackgrounded: false, - }), - ); - registry.register( - makeRegistration('paused-1', { - status: 'paused', - }), - ); - - registry.register(makeRegistration('bg-1')); - expect(() => registry.register(makeRegistration('bg-2'))).toThrow( - 'maximum concurrent background agents (1) reached', - ); - - registry.complete('bg-1', 'done'); - registry.register(makeRegistration('bg-2')); - - expect(registry.get('fg-1')).toBeDefined(); - expect(registry.get('paused-1')).toBeDefined(); - expect(registry.get('bg-2')?.status).toBe('running'); - }); - }); - - it('aborts all running agents and emits fallback notifications', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - const ac1 = new AbortController(); - const ac2 = new AbortController(); - - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: ac1, - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - registry.register({ - agentId: 'b', - description: 'agent b', - status: 'running', - startTime: Date.now(), - abortController: ac2, - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.abortAll(); - - expect(ac1.signal.aborted).toBe(true); - expect(ac2.signal.aborted).toBe(true); - expect(registry.get('a')!.status).toBe('cancelled'); - expect(registry.get('b')!.status).toBe('cancelled'); - // abortAll is a shutdown path — no natural handler will fire, so - // finalizeCancellationIfPending emits one cancelled notification per - // agent to keep the SDK contract intact. - expect(callback).toHaveBeenCalledTimes(2); - }); - - it('abortAll({ notify: false }) suppresses terminal notifications from old tasks', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.abortAll({ notify: false }); - - expect(registry.get('a')!.status).toBe('cancelled'); - expect(registry.hasUnfinalizedTasks()).toBe(false); - expect(callback).not.toHaveBeenCalled(); - - registry.complete('a', 'late result'); - registry.finalizeCancelled('a', 'late partial'); - - expect(callback).not.toHaveBeenCalled(); - expect(registry.get('a')!.status).toBe('cancelled'); - expect(registry.get('a')!.result).toBeUndefined(); - }); - - it('abortAll({ notify: false }) suppresses pending fallback notifications', () => { - vi.useFakeTimers(); - try { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.cancel('a'); - registry.abortAll({ notify: false }); - vi.runAllTimers(); - - expect(callback).not.toHaveBeenCalled(); - expect(registry.hasUnfinalizedTasks()).toBe(false); - } finally { - vi.useRealTimers(); - } - }); - - it('hasUnfinalizedTasks reports cancelled-but-not-notified entries', () => { - // Headless runs rely on this to keep the event loop alive after a - // task_stop until the agent's natural handler has emitted the - // terminal task-notification — otherwise the matching notification - // can be dropped before stream-json/SDK consumers observe it. - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - expect(registry.hasUnfinalizedTasks()).toBe(true); - - registry.cancel('test-1'); - expect(registry.get('test-1')!.status).toBe('cancelled'); - expect(registry.hasUnfinalizedTasks()).toBe(true); - - registry.finalizeCancelled('test-1', ''); - expect(registry.hasUnfinalizedTasks()).toBe(false); - }); - - it('hasUnfinalizedTasks clears once every entry has been notified', () => { - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - registry.register({ - agentId: 'b', - description: 'agent b', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - expect(registry.hasUnfinalizedTasks()).toBe(true); - registry.complete('a', 'done'); - expect(registry.hasUnfinalizedTasks()).toBe(true); - registry.fail('b', 'boom'); - expect(registry.hasUnfinalizedTasks()).toBe(false); - }); - - it('complete after cancellation surfaces the real result', () => { - // When cancel races with the natural completion handler, the agent's - // reasoning loop may have finished with a real result before the abort - // landed. complete() transitions cancelled → completed and emits the - // terminal notification carrying that real result, instead of letting - // the bare "cancelled" notification discard it. - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.cancel('test-1'); - registry.complete('test-1', 'real result after cancel race'); - - expect(registry.get('test-1')!.status).toBe('completed'); - expect(registry.get('test-1')!.result).toBe( - 'real result after cancel race', - ); - expect(callback).toHaveBeenCalledTimes(1); - const [, modelText] = callback.mock.calls[0]; - expect(modelText).toContain('completed'); - expect(modelText).toContain('real result after cancel race'); - }); - - it('fail after cancellation surfaces the real error', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.cancel('test-1'); - registry.fail('test-1', 'real error after cancel race'); - - expect(registry.get('test-1')!.status).toBe('failed'); - expect(registry.get('test-1')!.error).toBe('real error after cancel race'); - expect(callback).toHaveBeenCalledTimes(1); - const [, modelText] = callback.mock.calls[0]; - expect(modelText).toContain('failed'); - }); - - it('second terminal call does not double-notify', () => { - // Once a terminal notification has fired, subsequent terminal calls - // (from late fire-and-forget paths) must not produce a duplicate. - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.complete('test-1', 'first'); - registry.fail('test-1', 'late error'); - - expect(callback).toHaveBeenCalledTimes(1); - expect(registry.get('test-1')!.status).toBe('completed'); - }); - - it('does not send notification without callback', () => { - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - // Should not throw - registry.complete('test-1', 'done'); - expect(registry.get('test-1')!.status).toBe('completed'); - }); - - it('propagates toolUseId through XML and notification meta', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - toolUseId: 'call-abc-123', - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.complete('test-1', 'done'); - - expect(callback).toHaveBeenCalledOnce(); - const [, modelText, meta] = callback.mock.calls[0]; - expect(modelText).toContain('call-abc-123'); - expect(meta.toolUseId).toBe('call-abc-123'); - }); - - it('omits tool-use-id XML tag when toolUseId is absent', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.complete('test-1', 'done'); - - const [, modelText, meta] = callback.mock.calls[0]; - expect(modelText).not.toContain(''); - expect(meta.toolUseId).toBeUndefined(); - }); - - it('getAll returns every entry regardless of status', () => { - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - registry.register({ - agentId: 'b', - description: 'agent b', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - registry.register({ - agentId: 'c', - description: 'agent c', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.complete('a', 'done'); - registry.fail('b', 'boom'); - - const all = registry.getAll(); - expect(all).toHaveLength(3); - expect(all.map((e) => e.status).sort()).toEqual([ - 'completed', - 'failed', - 'running', - ]); - // Callers that need only running entries filter getAll() themselves. - expect( - registry - .getAll() - .filter((e) => e.status === 'running') - .map((e) => e.agentId), - ).toEqual(['c']); - }); - - it('statusChange callback fires on register and every state transition', () => { - const seen: Array<{ id: string; status: string }> = []; - registry.setStatusChangeCallback((entry) => { - if (entry) { - seen.push({ id: entry.agentId, status: entry.status }); - } - }); - - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - registry.register({ - agentId: 'b', - description: 'agent b', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - registry.complete('a', 'ok'); - registry.fail('b', 'err'); - - expect(seen).toEqual([ - { id: 'a', status: 'running' }, - { id: 'b', status: 'running' }, - { id: 'a', status: 'completed' }, - { id: 'b', status: 'failed' }, - ]); - }); - - it('statusChange callback errors do not break registry operations', () => { - registry.setStatusChangeCallback(() => { - throw new Error('listener broke'); - }); - - // Should not throw even though the callback does. - expect(() => - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }), - ).not.toThrow(); - expect(registry.get('a')?.status).toBe('running'); - }); - - it('statusChange callback can be cleared with undefined', () => { - const cb = vi.fn(); - registry.setStatusChangeCallback(cb); - registry.setStatusChangeCallback(undefined); - - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - expect(cb).not.toHaveBeenCalled(); - }); - - it('appendActivity builds a rolling buffer capped at 5', () => { - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - for (let i = 0; i < 7; i++) { - registry.appendActivity('a', { - name: `Tool${i}`, - description: `call ${i}`, - at: i, - }); - } - - const activities = registry.get('a')!.recentActivities ?? []; - expect(activities.map((a) => a.name)).toEqual([ - 'Tool2', - 'Tool3', - 'Tool4', - 'Tool5', - 'Tool6', - ]); - }); - - it('appendActivity no-ops after the agent terminates', () => { - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.complete('a', 'done'); - registry.appendActivity('a', { name: 'Late', description: 'x', at: 99 }); - - expect(registry.get('a')!.recentActivities ?? []).toHaveLength(0); - }); - - it('appendActivity fires activityChange, not statusChange', () => { - const statusCb = vi.fn(); - const activityCb = vi.fn(); - registry.setStatusChangeCallback(statusCb); - registry.setActivityChangeCallback(activityCb); - - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - statusCb.mockClear(); - activityCb.mockClear(); - - registry.appendActivity('a', { name: 'T', description: 'd', at: 0 }); - - expect(statusCb).not.toHaveBeenCalled(); - expect(activityCb).toHaveBeenCalledOnce(); - expect(activityCb.mock.calls[0][0].agentId).toBe('a'); - }); - - it('stores prompt verbatim on the entry', () => { - registry.register({ - agentId: 'a', - description: 'agent a', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - prompt: 'Run sleep 30 and report done.', - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - expect(registry.get('a')!.prompt).toBe('Run sleep 30 and report done.'); - }); - - it('escapes XML metacharacters in interpolated fields', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'summarize & ', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.complete('test-1', 'here is bold & '); - - const [, modelText] = callback.mock.calls[0]; - // No injected closing tags — subagent text is escaped so the - // parent envelope stays a single task-notification element. - expect(modelText.match(/<\/task-notification>/g)!.length).toBe(1); - expect(modelText).toContain('</result>'); - expect(modelText).toContain('</task-notification>'); - expect(modelText).toContain('<b>bold</b>'); - expect(modelText).toContain('&'); - }); - - describe('terminal-entry retention cap', () => { - function makeRegisteredEntry(id: string, startTime: number) { - return { - agentId: id, - description: id, - status: 'running' as const, - startTime, - abortController: new AbortController(), - outputFile: `/tmp/${id}.jsonl`, - isBackgrounded: true, - }; - } - - it('retains only a bounded number of fully-finalized terminal entries', () => { - // Register and complete one more entry than the cap allows so - // the prune kicks in. Use strictly increasing startTimes so the - // synthetic endTimes (Date.now() inside complete) preserve a - // deterministic eviction order via the startTime tiebreaker. - for (let i = 0; i < MAX_RETAINED_TERMINAL_AGENTS + 2; i++) { - registry.register(makeRegisteredEntry(`a-${i}`, i * 1000)); - registry.complete(`a-${i}`, 'done'); - } - expect(registry.getAll()).toHaveLength(MAX_RETAINED_TERMINAL_AGENTS); - // The two oldest (`a-0`, `a-1`) get pruned; the newest survives. - expect(registry.get('a-0')).toBeUndefined(); - expect(registry.get('a-1')).toBeUndefined(); - expect( - registry.get(`a-${MAX_RETAINED_TERMINAL_AGENTS + 1}`), - ).toBeDefined(); - }); - - it('never evicts running entries even when terminal entries blow past the cap', () => { - // The user's only handle on a live subagent is its row in the - // dialog; a prune that drops a running entry would silently - // strand work in progress. - registry.register(makeRegisteredEntry('live', 1)); - for (let i = 0; i < MAX_RETAINED_TERMINAL_AGENTS + 1; i++) { - registry.register(makeRegisteredEntry(`done-${i}`, 100 + i * 1000)); - registry.complete(`done-${i}`, 'done'); - } - // Cap-of-32 terminals + 1 running survivor = 33 entries kept. - expect(registry.getAll()).toHaveLength(MAX_RETAINED_TERMINAL_AGENTS + 1); - expect(registry.get('live')?.status).toBe('running'); - // The oldest terminal entry is the one evicted. - expect(registry.get('done-0')).toBeUndefined(); - }); - - it('never evicts paused entries (recoverable, awaiting resume/abandon)', () => { - // Manually plant a paused entry — the registry exposes - // abandon/resume but no public "transition to paused" call; - // resume restoration on Config init writes paused entries - // directly via register(). - registry.register({ - agentId: 'paused-1', - description: 'paused', - status: 'paused', - startTime: 1, - abortController: new AbortController(), - outputFile: '/tmp/paused-1.jsonl', - isBackgrounded: true, - }); - // Push terminal entries past the cap so prune is forced to choose - // an eviction set. - for (let i = 0; i < MAX_RETAINED_TERMINAL_AGENTS + 1; i++) { - registry.register(makeRegisteredEntry(`done-${i}`, 100 + i * 1000)); - registry.complete(`done-${i}`, 'done'); - } - expect(registry.get('paused-1')?.status).toBe('paused'); - }); - - it('never evicts cancelled-but-not-yet-notified entries', () => { - // cancel() flips the entry to cancelled immediately but defers - // the terminal task-notification to the natural handler / grace - // timer. Pruning here would break the SDK contract that every - // register pairs with exactly one terminal task-notification. - registry.setNotificationCallback(() => {}); - registry.register(makeRegisteredEntry('pending-cancel', 1)); - registry.cancel('pending-cancel'); - // Push terminal entries past the cap. - for (let i = 0; i < MAX_RETAINED_TERMINAL_AGENTS + 1; i++) { - registry.register(makeRegisteredEntry(`done-${i}`, 100 + i * 1000)); - registry.complete(`done-${i}`, 'done'); - } - // pending-cancel survives because it has notified=false; it's - // still owed a terminal notification. - expect(registry.get('pending-cancel')?.status).toBe('cancelled'); - expect(registry.get('pending-cancel')?.notified).toBeFalsy(); - }); - - it('prunes an abandoned (paused → cancelled) entry the same as any other terminal', () => { - // abandon() is the only path that flips notified=true on a - // previously-paused entry. Make sure the resulting terminal - // counts toward the cap so a session that abandons many - // paused agents doesn't bypass the retention bound. - registry.register({ - agentId: 'paused-overflow', - description: 'paused', - status: 'paused', - startTime: 1, - abortController: new AbortController(), - outputFile: '/tmp/paused-overflow.jsonl', - isBackgrounded: true, - }); - registry.abandon('paused-overflow'); - for (let i = 0; i < MAX_RETAINED_TERMINAL_AGENTS; i++) { - registry.register(makeRegisteredEntry(`done-${i}`, 100 + i * 1000)); - registry.complete(`done-${i}`, 'done'); - } - // After the loop, terminal count = 1 (abandon) + 32 (complete) = - // 33, exceeds the cap → oldest evicted. The abandoned entry - // (startTime=1, endTime=earliest) is the one evicted. - expect(registry.getAll()).toHaveLength(MAX_RETAINED_TERMINAL_AGENTS); - expect(registry.get('paused-overflow')).toBeUndefined(); - }); - }); - - describe('queueMessage', () => { - it('queues a message for a running agent', () => { - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - const result = registry.queueMessage('test-1', 'hello'); - expect(result).toBe(true); - expect(registry.get('test-1')!.pendingMessages).toEqual(['hello']); - }); - - it('returns false for non-existent agent', () => { - expect(registry.queueMessage('nope', 'hello')).toBe(false); - }); - - it('returns false for non-running agent', () => { - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - registry.complete('test-1', 'done'); - - expect(registry.queueMessage('test-1', 'hello')).toBe(false); - }); - }); - - describe('drainMessages', () => { - it('drains all messages and clears the queue', () => { - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.queueMessage('test-1', 'msg-1'); - registry.queueMessage('test-1', 'msg-2'); - - const messages = registry.drainMessages('test-1'); - expect(messages).toEqual(['msg-1', 'msg-2']); - expect(registry.get('test-1')!.pendingMessages).toEqual([]); - }); - - it('returns empty array when no messages queued', () => { - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - expect(registry.drainMessages('test-1')).toEqual([]); - }); - - it('returns empty array for non-existent agent', () => { - expect(registry.drainMessages('nope')).toEqual([]); - }); - }); - - describe('waitForMessages', () => { - it('resolves with queued input when a running agent is notified', async () => { - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - const waitPromise = registry.waitForMessages( - 'test-1', - new AbortController().signal, - ); - - registry.queueExternalInput('test-1', { - kind: 'notification', - text: 'event', - }); - - await expect(waitPromise).resolves.toEqual([ - { - kind: 'notification', - text: 'event', - }, - ]); - expect(registry.drainMessages('test-1')).toEqual([]); - }); - - it('resolves empty when the wait signal is aborted', async () => { - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - const waitAbort = new AbortController(); - const waitPromise = registry.waitForMessages('test-1', waitAbort.signal); - - waitAbort.abort(); - - await expect(waitPromise).resolves.toEqual([]); - }); - - it('resolves empty if the signal aborts immediately after listener registration', async () => { - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - let aborted = false; - const signal = { - get aborted() { - return aborted; - }, - addEventListener: vi.fn(() => { - aborted = true; - }), - removeEventListener: vi.fn(), - } as unknown as AbortSignal; - - await expect(registry.waitForMessages('test-1', signal)).resolves.toEqual( - [], - ); - expect(signal.removeEventListener).toHaveBeenCalled(); - }); - - it('wakes external input waiters without queueing input', async () => { - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - const waitPromise = registry.waitForMessages( - 'test-1', - new AbortController().signal, - ); - - registry.wakeExternalInputWaiters('test-1'); - - await expect(waitPromise).resolves.toEqual([]); - expect(registry.drainMessages('test-1')).toEqual([]); - }); - }); - - describe('session switch helpers', () => { - it('reset clears tracked entries without touching persisted sidecars', () => { - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - registry.register({ - agentId: 'test-2', - description: 'paused agent', - status: 'paused', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.reset(); - - expect(registry.getAll()).toEqual([]); - }); - }); - - describe('notification XML', () => { - it('includes output-file tag when outputFile is set', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - outputFile: '/tmp/agents/test-1.txt', - isBackgrounded: true, - }); - - registry.complete('test-1', 'done'); - - const [, modelText] = callback.mock.calls[0]; - expect(modelText).toContain( - '/tmp/agents/test-1.txt', - ); - }); - - it('omits output-file tag when outputFile is empty', () => { - // outputFile is mandatory on the contract but a caller may pass an - // empty string (e.g. an agent kind that explicitly opts out of disk - // persistence). In that case the notification XML should omit the - // `` tag — model-side parsers shouldn't see a path to - // a file that doesn't exist. - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'test-1', - description: 'test agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '', - }); - - registry.complete('test-1', 'done'); - - const [, modelText] = callback.mock.calls[0]; - expect(modelText).not.toContain(''); - }); - }); - - describe('foreground flavor', () => { - it('does not emit a task-notification on complete', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'fg-1', - description: 'sync agent', - isBackgrounded: false, - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - outputFile: '/tmp/test.jsonl', - }); - - registry.complete('fg-1', 'result text'); - - // Foreground entries deliver their result through the parent's normal - // tool-result channel; emitting the XML envelope on top would feed - // the parent model the same payload twice. - expect(callback).not.toHaveBeenCalled(); - // The status mutation still happens — internal invariants intact. - expect(registry.get('fg-1')!.status).toBe('completed'); - expect(registry.get('fg-1')!.notified).toBe(true); - }); - - it('does not emit a task-notification on fail', () => { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'fg-2', - description: 'sync agent', - isBackgrounded: false, - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - outputFile: '/tmp/test.jsonl', - }); - - registry.fail('fg-2', 'oops'); - - expect(callback).not.toHaveBeenCalled(); - }); - - it('is excluded from hasUnfinalizedTasks()', () => { - registry.register({ - agentId: 'fg-3', - description: 'sync agent', - isBackgrounded: false, - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - outputFile: '/tmp/test.jsonl', - }); - - // A still-running foreground entry must NOT keep the headless - // event loop alive — the parent's tool-call await already does that. - expect(registry.hasUnfinalizedTasks()).toBe(false); - }); - - it('cancel does not schedule the grace timer', () => { - // The grace-timer fallback only matters for background entries that - // might not see their natural completion handler fire. Foreground - // entries unregister themselves in agent.ts's finally path. - vi.useFakeTimers(); - try { - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'fg-4', - description: 'sync agent', - isBackgrounded: false, - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - outputFile: '/tmp/test.jsonl', - }); - - registry.cancel('fg-4'); - - // Advance well past the 5s grace window — no notification should fire. - vi.advanceTimersByTime(60_000); - expect(callback).not.toHaveBeenCalled(); - } finally { - vi.useRealTimers(); - } - }); - - it('unregisterForeground removes the entry and emits a status change', () => { - const onStatusChange = vi.fn(); - registry.setStatusChangeCallback(onStatusChange); - - registry.register({ - agentId: 'fg-5', - description: 'sync agent', - isBackgrounded: false, - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - outputFile: '/tmp/test.jsonl', - }); - onStatusChange.mockClear(); - - registry.unregisterForeground('fg-5'); - - expect(registry.get('fg-5')).toBeUndefined(); - expect(onStatusChange).toHaveBeenCalledTimes(1); - }); - - it('unregisterForeground throws if asked to remove a background entry', () => { - // Background entries must terminate via complete/fail/finalizeCancelled - // so the task-notification + headless holdback invariants stay intact. - // A silent no-op would mask caller bugs, so this throws. - registry.register({ - agentId: 'bg-1', - description: 'async agent', - isBackgrounded: true, - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - outputFile: '/tmp/test.jsonl', - }); - - expect(() => registry.unregisterForeground('bg-1')).toThrow( - /non-foreground entry bg-1/, - ); - expect(registry.get('bg-1')).toBeDefined(); - }); - - it('unregisterForeground is a no-op for unknown agent ids', () => { - // Idempotent for already-unregistered/never-registered ids — the - // foreground finally path runs unconditionally and shouldn't throw - // if a parallel cancel already cleared the entry. - expect(() => registry.unregisterForeground('missing')).not.toThrow(); - }); - - it('does not invoke the register callback for foreground entries', () => { - // Non-interactive bridges setRegisterCallback to a `task_started` - // SDK event. Foreground entries never produce a paired terminal - // task-notification (see emitNotification's flavor gate), so letting - // them fire `task_started` would leak orphaned in-flight tasks to - // SDK consumers. - const onRegister = vi.fn(); - registry.setRegisterCallback(onRegister); - - registry.register({ - agentId: 'fg-no-register-cb', - description: 'sync agent', - isBackgrounded: false, - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - outputFile: '/tmp/test.jsonl', - }); - - expect(onRegister).not.toHaveBeenCalled(); - - // Background entries still fire it. - registry.register({ - agentId: 'bg-fires-register-cb', - description: 'async agent', - isBackgrounded: true, - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - outputFile: '/tmp/test.jsonl', - }); - expect(onRegister).toHaveBeenCalledTimes(1); - expect(onRegister.mock.calls[0]![0].agentId).toBe('bg-fires-register-cb'); - }); - - it('unregisterForeground emits status change after removing the entry', () => { - // The entry is deleted from the Map before the status-change callback - // fires, so a callback that rebuilds its snapshot via getAll() no - // longer includes this entry. This ordering prevents the entry from - // lingering in React state with status='running' — the bug that - // caused "1 local agent" to stay visible after the foreground agent - // completed. - registry.register({ - agentId: 'fg-unregister-order', - description: 'sync agent', - isBackgrounded: false, - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - outputFile: '/tmp/test.jsonl', - }); - - let observedFromCallback: BackgroundTaskEntry | undefined; - let snapshotDuringCallback: BackgroundTaskEntry[] = []; - registry.setStatusChangeCallback((entry) => { - if (entry?.agentId === 'fg-unregister-order') { - observedFromCallback = registry.get(entry.agentId); - snapshotDuringCallback = registry.getAll(); - } - }); - - registry.unregisterForeground('fg-unregister-order'); - - // The entry has been deleted before the callback fires, so - // registry.get() returns undefined and getAll() omits it. - expect(observedFromCallback).toBeUndefined(); - expect(snapshotDuringCallback).toEqual([]); - expect(registry.get('fg-unregister-order')).toBeUndefined(); - }); - - it('background entries fire a task-notification on complete', () => { - // Counterpart to the foreground "does not emit" cases above — - // background entries deliver their result through the XML envelope, - // so the notification callback must fire on complete. - const callback = vi.fn(); - registry.setNotificationCallback(callback); - - registry.register({ - agentId: 'bg-notify-1', - description: 'async agent', - status: 'running', - startTime: Date.now(), - abortController: new AbortController(), - isBackgrounded: true, - outputFile: '/tmp/test.jsonl', - }); - - registry.complete('bg-notify-1', 'done'); - - expect(callback).toHaveBeenCalledOnce(); - }); - }); -}); diff --git a/packages/core/src/agents/background-tasks.ts b/packages/core/src/agents/background-tasks.ts index 5127301bbfb..d18838ec75b 100644 --- a/packages/core/src/agents/background-tasks.ts +++ b/packages/core/src/agents/background-tasks.ts @@ -5,897 +5,30 @@ */ /** - * @fileoverview BackgroundTaskRegistry — tracks background (async) sub-agents - * and, with `isBackgrounded: false`, the currently-running synchronous - * sub-agents whose UI is routed through the same pill+dialog while the - * parent turn waits on them. Both share the registry (and the dialog - * wiring) but differ in lifecycle: + * @fileoverview Re-export shim. The agent-task module + * (`tasks/agent-task.ts`) now owns the `AgentTask` type, lifecycle + * helpers, label builder, and the background-agent concurrency cap; + * this file re-exports them so external SDK consumers that imported + * from `'@qwen-code/qwen-code-core'` (which surfaces this module via + * `agents/index.ts`) keep their import paths working for one release. * - * - `isBackgrounded: true` entries persist across turns, emit a - * `` on terminal status (the parent's only return - * channel), and contribute to `hasUnfinalizedTasks()` so headless callers - * keep their loop alive. - * - `isBackgrounded: false` entries live for the duration of the parent's - * tool-call, are unregistered as soon as `execute()` returns, deliver - * their result through the normal tool-result channel (no XML envelope), - * and don't participate in the headless holdback. - */ - -import { createDebugLogger } from '../utils/debugLogger.js'; -import { escapeXml } from '../utils/xml.js'; -import { patchAgentMeta } from './agent-transcript.js'; -import type { AgentExternalInput } from './runtime/agent-types.js'; -import type { TaskBase, TaskRegistration, TaskStatus } from './tasks/types.js'; - -const debugLogger = createDebugLogger('BACKGROUND_TASKS'); - -const MAX_DESCRIPTION_LENGTH = 40; -const MAX_RECENT_ACTIVITIES = 5; -export const DEFAULT_MAX_CONCURRENT_BACKGROUND_AGENTS = 10; -export const BACKGROUND_AGENT_CONCURRENCY_ENV = - 'QWEN_CODE_MAX_BACKGROUND_AGENTS'; - -export function resolveMaxConcurrentBackgroundAgents( - env: Record = process.env, -): number { - const raw = env[BACKGROUND_AGENT_CONCURRENCY_ENV]; - if (raw === undefined || raw.trim() === '') { - return DEFAULT_MAX_CONCURRENT_BACKGROUND_AGENTS; - } - - const parsed = Number(raw); - if (!Number.isInteger(parsed) || parsed < 1) { - debugLogger.warn( - `Invalid ${BACKGROUND_AGENT_CONCURRENCY_ENV}=${JSON.stringify(raw)}, ` + - `using default (${DEFAULT_MAX_CONCURRENT_BACKGROUND_AGENTS})`, - ); - return DEFAULT_MAX_CONCURRENT_BACKGROUND_AGENTS; - } - - return parsed; -} - -export const MAX_CONCURRENT_BACKGROUND_AGENTS = - resolveMaxConcurrentBackgroundAgents(); - -/** - * Cap on how many fully-finalized terminal entries (those that have - * already emitted their terminal `task-notification`) the registry - * retains. Without this cap, every short-lived background subagent - * leaves a row in the Background tasks dialog and pill forever, - * crowding out the running entries the user actually opened the - * dialog to find. Mirrors the rationale + retention pattern in - * `MonitorRegistry.MAX_RETAINED_TERMINAL_MONITORS` and - * `BackgroundShellRegistry.MAX_RETAINED_TERMINAL_SHELLS`. - * - * Entries that are still `running`, `paused`, or `cancelled` but - * not yet notified are NEVER evicted — pruning a not-yet-notified - * cancelled entry would break the SDK contract that every - * `register` pairs with exactly one terminal `task-notification`. - */ -export const MAX_RETAINED_TERMINAL_AGENTS = 32; - -// Grace period after cancel() before emitting a fallback cancelled -// notification. The natural handler (bgBody) almost always settles and -// emits the terminal notification with the real partial result well -// within this window; the timeout only fires for pathological tools -// that ignore AbortSignal. Must be long enough that normal scheduler -// unwind wins the race, short enough that a stuck headless wait loop -// doesn't feel hung. -const CANCEL_GRACE_MS = 5000; - -/** - * Single source of truth for the human-facing label of a background - * entry. Shared by the notification payload (model-facing) and the TUI - * dialog (user-facing) so the two surfaces never drift. - * - * When `includePrefix` is true (default), returns `subagentType: desc`; - * when false, returns the bare truncated description — used where the - * subagent type is already rendered separately (e.g. the dialog header). - */ -export function buildBackgroundEntryLabel( - entry: { description: string; subagentType?: string }, - options: { includePrefix?: boolean } = {}, -): string { - const { includePrefix = true } = options; - let raw = entry.description; - if ( - entry.subagentType && - raw.toLowerCase().startsWith(entry.subagentType.toLowerCase() + ':') - ) { - raw = raw.slice(entry.subagentType.length + 1).trimStart(); - } - const truncated = - raw.length > MAX_DESCRIPTION_LENGTH - ? raw.slice(0, MAX_DESCRIPTION_LENGTH - 1) + '\u2026' - : raw; - return includePrefix && entry.subagentType - ? `${entry.subagentType}: ${truncated}` - : truncated; -} - -// Subagent-produced strings (description, result, error) can contain `<`, -// `>`, or literal `` — without escaping, a subagent -// summarizing HTML or another agent's notification could close the -// envelope early and forge sibling tags (e.g. a faked ) that the -// parent model would treat as trusted metadata. Use the shared helper. - -/** - * @deprecated Use `TaskStatus` from `./tasks/types.js`. Kept as a one-release - * alias so existing consumers (notably `nonInteractiveCli.ts`) compile - * unchanged; the underlying union is identical. - */ -export type BackgroundTaskStatus = TaskStatus; - -export interface AgentCompletionStats { - totalTokens: number; - toolUses: number; - durationMs: number; -} - -/** - * A compact record of a recent tool invocation — drives the Progress - * section of the detail dialog. The Agent tool maintains a rolling - * buffer of these on each background entry by subscribing to the - * subagent's event emitter. - */ -export interface BackgroundActivity { - /** Tool name (e.g. `Bash`, `Read`). */ - name: string; - /** Short one-line description — the tool's own render-friendly summary. */ - description: string; - /** Emission timestamp (ms). */ - at: number; -} - -/** - * Agent kind of `TaskState`. Tracks one running subagent — either a - * synchronous foreground run (`isBackgrounded: false`, awaited by the - * parent's tool-call) or an async background run (`isBackgrounded: true`, - * persists across turns and emits a terminal ``). - * - * Carries the shared `TaskBase` envelope plus agent-specific state: - * subagent config, prompt, stats, recent activity buffer, persisted - * sidecar metadata path, message queue, and resume hooks. - */ -export interface AgentTask extends TaskBase { - kind: 'agent'; - /** - * @deprecated Read `id` instead; kept as a synonym during the back-compat - * window. Always equals `id`. - */ - agentId: string; - subagentType?: string; - /** - * True if the task is running asynchronously (parent has moved on, the - * task persists across turns and emits a terminal XML notification). - * False if the parent's tool-call is synchronously awaiting it; the - * result is delivered through the normal tool-result channel and no - * XML envelope fires. Replaces the older `flavor: 'foreground' | - * 'background'` discriminator — same binary fact, named after the - * question every read site asks. - */ - isBackgrounded: boolean; - status: TaskStatus; - result?: string; - error?: string; - /** - * Present only when the task is intentionally kept paused but cannot be - * safely resumed under the current conditions. - */ - resumeBlockedReason?: string; - stats?: AgentCompletionStats; - toolUseId?: string; - /** - * The original user-supplied prompt for the background task. Surfaced - * verbatim in the detail dialog's Prompt section. Optional because - * resume-restored entries may not have it. - */ - prompt?: string; - /** - * Rolling buffer (newest last, capped at MAX_RECENT_ACTIVITIES) of - * recent tool invocations by this agent. Feeds the detail dialog's - * Progress section. Replaced as a new array each time an activity is - * appended so reference-based change detection works. Optional: - * callers may register without providing it, and `appendActivity` - * initializes the array lazily. - */ - recentActivities?: readonly BackgroundActivity[]; - /** Absolute path to the agent's sidecar metadata file. */ - metaPath?: string; - /** - * Inputs queued for delivery between tool rounds. - * Strings are parent `send_message` payloads; notification objects are - * owner-routed Monitor notifications. - */ - pendingMessages?: AgentExternalInput[]; - /** - * Persisted sidecar status to write when the current cancellation settles. - * Explicit user cancellation uses `cancelled`; shutdown interruption keeps - * `running` so `/resume` can recover the work later. - */ - persistedCancellationStatus?: Extract; -} - -/** - * @deprecated Renamed to `AgentTask`. Kept as a one-release type alias for - * external SDK consumers; will be removed in the release after PR 2 lands. - */ -export type BackgroundTaskEntry = AgentTask; - -/** - * Shape callers pass to {@link BackgroundTaskRegistry.register}; the - * registry derives the shared `TaskBase` envelope (`id`, `kind`, - * `outputOffset`, `notified`) from these and the surrounding context. - * `outputFile` is required here because every agent run reserves a JSONL - * transcript path at registration. - */ -export type AgentTaskRegistration = TaskRegistration; - -export interface NotificationMeta { - agentId: string; - status: TaskStatus; - stats?: AgentCompletionStats; - toolUseId?: string; -} - -export type BackgroundNotificationCallback = ( - displayText: string, - modelText: string, - meta: NotificationMeta, -) => void; - -export type BackgroundRegisterCallback = (entry: AgentTask) => void; - -interface BackgroundTaskCancelOptions { - notify?: boolean; - persistedStatus?: Extract; -} - -/** - * Fires on entry status transitions: `register`, `complete`, `fail`, - * `cancel`, `finalizeCancelled`, `finalizeCancellationIfPending`, - * `abandon`, `unregisterForeground`, and `reset`. Intentionally does - * NOT fire on `appendActivity` so consumers that only care about the - * roster don't re-render on every tool call a background agent makes. - * - * Ordering relative to the registry mutation falls into two camps: - * - **Keeps the entry around** (`register` / `complete` / `fail` / - * `cancel` / `finalizeCancelled` / - * `finalizeCancellationIfPending` / `abandon`): emit while the - * entry is still in the Map (the status field has been mutated - * in place to its terminal value), so a callback that re-reads - * `registry.get(entry.agentId)` sees the entry. Snapshot-style - * consumers calling `getAll()` see the new status too. - * - **Removes the entry** (`unregisterForeground`, `reset`): - * deletes from the Map BEFORE emitting so snapshot-style - * consumers drop the row. The `entry` arg carries the agent's - * last live state for log / display consumers; `registry.get` - * and `getAll` already reflect the deletion. - */ -export type BackgroundStatusChangeCallback = (entry?: AgentTask) => void; - -/** Fires on `appendActivity` — scoped to detail-view consumers. */ -export type BackgroundActivityChangeCallback = (entry: AgentTask) => void; - -type MessageWaiter = () => void; - -export interface BackgroundTaskRegistryOptions { - maxConcurrentBackgroundAgents?: number; -} - -export class BackgroundTaskRegistry { - private readonly agents = new Map(); - private readonly messageWaiters = new Map>(); - private readonly maxConcurrentBackgroundAgents: number; - private notificationCallback?: BackgroundNotificationCallback; - private registerCallback?: BackgroundRegisterCallback; - private statusChangeCallback?: BackgroundStatusChangeCallback; - private activityChangeCallback?: BackgroundActivityChangeCallback; - - constructor(options: BackgroundTaskRegistryOptions = {}) { - const configured = - options.maxConcurrentBackgroundAgents ?? MAX_CONCURRENT_BACKGROUND_AGENTS; - this.maxConcurrentBackgroundAgents = - Number.isInteger(configured) && configured >= 1 - ? configured - : MAX_CONCURRENT_BACKGROUND_AGENTS; - } - - assertCanStartBackgroundAgent(): void { - const running = this.getRunningBackgroundCount(); - if (running >= this.maxConcurrentBackgroundAgents) { - debugLogger.warn( - `Background agent concurrency cap reached: ` + - `${running}/${this.maxConcurrentBackgroundAgents}. ` + - `Refusing new background agent.`, - ); - throw new Error( - `Cannot start background agent: maximum concurrent background agents ` + - `(${this.maxConcurrentBackgroundAgents}) reached. Stop an existing ` + - `agent first.`, - ); - } - } - - register(registration: AgentTaskRegistration): AgentTask { - if (registration.isBackgrounded && registration.status === 'running') { - const existing = this.agents.get(registration.agentId); - const isReplacingRunning = - existing?.isBackgrounded === true && existing.status === 'running'; - if (!isReplacingRunning) { - this.assertCanStartBackgroundAgent(); - } - } - - // Mutate the registration in place to graduate it to an `AgentTask`. - // Returning the same reference lets callers (e.g. the resume service) - // continue using their local variable post-register and lets external - // consumers see updates the registry makes without an extra `get()`. - const entry = registration as AgentTask; - entry.id = registration.agentId; - entry.kind = 'agent'; - entry.outputOffset = 0; - entry.notified = false; - entry.pendingMessages = registration.pendingMessages ?? []; - this.agents.set(entry.agentId, entry); - debugLogger.info(`Registered background agent: ${entry.agentId}`); - - // Foreground entries are paired with a synchronous tool-call result on - // the parent's response and never emit a terminal `task_notification` - // (see emitNotification's isBackgrounded gate). Letting them fire the - // register callback would emit a `task_started` SDK event without a - // matching completion event, breaking the lifecycle contract for SDK - // consumers. - if (entry.isBackgrounded && this.registerCallback) { - try { - this.registerCallback(entry); - } catch (error) { - debugLogger.error('Failed to emit register callback:', error); - } - } - this.emitStatusChange(entry); - return entry; - } - - // Transition a still-running entry to 'completed' and emit the terminal - // notification. No-op if the entry is already terminal *and* has been - // notified — protects against duplicate emission when cancel aborts the - // signal and the natural handler also races to completion. - complete( - agentId: string, - result: string, - stats?: AgentCompletionStats, - ): void { - const entry = this.agents.get(agentId); - if (!entry) return; - // Allow running → completed (normal path) and cancelled → completed - // (cancel raced the natural handler: the reasoning loop finished with - // a real result before the abort landed, and we prefer to surface that - // real result over the bare cancel). - if (entry.status !== 'running' && entry.status !== 'cancelled') return; - if (entry.notified) return; - - entry.status = 'completed'; - entry.endTime = Date.now(); - entry.result = result; - entry.stats = stats; - debugLogger.info(`Background agent completed: ${agentId}`); - - this.emitNotification(entry); - this.emitStatusChange(entry); - } - - /** - * Remove a foreground entry from the registry without emitting any - * terminal notification. Called by the foreground tool-call's `finally` - * path, which has already delivered the result through the tool-result - * channel — the registry entry has served its UI-surfacing purpose. - * Background entries must go through complete/fail/finalizeCancelled - * instead, so this throws if asked to remove one. - */ - unregisterForeground(agentId: string): void { - const entry = this.agents.get(agentId); - if (!entry) return; - if (entry.isBackgrounded) { - throw new Error( - `unregisterForeground called on non-foreground entry ${agentId} ` + - `(isBackgrounded=true). ` + - `Background entries must terminate via complete/fail/finalizeCancelled.`, - ); - } - // Delete BEFORE emitting so snapshot-style consumers (those that - // re-pull `getAll()` from inside the callback) no longer include - // this entry. The reverse order (emit-then-delete) caused the - // foreground agent to linger as `status='running'` in the footer - // pill / dialog: the callback's `getAll()` still saw it, and no - // second status-change fired after the deletion. Diverges from - // complete/fail/cancel/finalize ordering on purpose — those - // keep the entry around (terminal state) so callbacks can inspect - // it on re-read; unregister removes it outright. - this.agents.delete(agentId); - this.emitStatusChange(entry); - debugLogger.info(`Unregistered foreground agent: ${agentId}`); - } - - // See complete() for the cancelled → terminal path rationale. - fail(agentId: string, error: string, stats?: AgentCompletionStats): void { - const entry = this.agents.get(agentId); - if (!entry) return; - if (entry.status !== 'running' && entry.status !== 'cancelled') return; - if (entry.notified) return; - - entry.status = 'failed'; - entry.endTime = Date.now(); - entry.error = error; - entry.stats = stats; - debugLogger.info(`Background agent failed: ${agentId}`); - - this.emitNotification(entry); - this.emitStatusChange(entry); - } - - // Cancellation aborts the signal and marks the entry as cancelled, but - // does *not* emit the terminal notification immediately. The natural - // completion path (bgBody) fires complete()/fail()/finalizeCancelled() - // with the real partial/final result, which carries far more information - // than a bare "cancelled" message. A deferred fallback handles the rare - // case where a tool ignores AbortSignal and bgBody never settles — the - // timeout lands on finalizeCancellationIfPending(), which is a no-op - // once the natural handler has already emitted. - // - // Foreground entries (`isBackgrounded === false`) take a partial path - // through this method: status flips to 'cancelled' and the meta sidecar - // is patched, but the Map entry is *not* removed. Removal is the caller's - // responsibility via `unregisterForeground()` in the tool-call's finally - // path — without that follow-up, the foreground entry leaks. Callers - // outside `agent.ts` that invoke `cancel()` on a foreground entry must - // pair it with `unregisterForeground()`. - cancel(agentId: string, options: BackgroundTaskCancelOptions = {}): void { - const entry = this.agents.get(agentId); - if (!entry || entry.status !== 'running') return; - const persistedStatus = options.persistedStatus ?? 'cancelled'; - - entry.abortController.abort(); - entry.status = 'cancelled'; - entry.endTime = Date.now(); - entry.persistedCancellationStatus = persistedStatus; - if (entry.metaPath) { - patchAgentMeta(entry.metaPath, { - status: persistedStatus, - lastUpdatedAt: new Date().toISOString(), - lastError: undefined, - }); - } - debugLogger.info(`Background agent cancelled: ${agentId}`); - this.emitStatusChange(entry); - - // Foreground entries don't emit XML notifications and unregister - // themselves in the tool-call's finally path, so the grace timer - // would only ever no-op for them. - if (!entry.isBackgrounded) return; - - if (options.notify === false) { - // Session reset paths intentionally suppress the old task's terminal - // notification so it cannot leak into a new conversation. - entry.notified = true; - return; - } - - const timer = setTimeout(() => { - this.finalizeCancellationIfPending(agentId); - }, CANCEL_GRACE_MS); - timer.unref?.(); - } - - /** - * Marks a paused interrupted task as intentionally discarded/cancelled - * without emitting a task-notification. Used when the user explicitly - * abandons a recovered task instead of resuming it. - */ - abandon(agentId: string): void { - const entry = this.agents.get(agentId); - if (!entry || entry.status !== 'paused') return; - - entry.status = 'cancelled'; - entry.endTime = Date.now(); - entry.notified = true; - debugLogger.info(`Abandoned paused background agent: ${agentId}`); - this.emitStatusChange(entry); - } - - // Emit the terminal cancelled notification once the agent's natural - // handler has confirmed that the reasoning loop ended because of the - // abort (terminateMode === CANCELLED). Attaches the partial result and - // stats so the parent model still sees whatever work the agent had - // captured before the abort landed, instead of a bare "cancelled" line. - finalizeCancelled( - agentId: string, - partialResult: string, - stats?: AgentCompletionStats, - ): void { - const entry = this.agents.get(agentId); - if (!entry) return; - if (entry.status !== 'running' && entry.status !== 'cancelled') return; - if (entry.notified) return; - - entry.status = 'cancelled'; - entry.endTime ??= Date.now(); - if (partialResult) entry.result = partialResult; - entry.stats = stats; - this.emitNotification(entry); - this.emitStatusChange(entry); - } - - // Emit the terminal cancelled notification for entries that were cancelled - // but for which no natural handler delivered a follow-up complete()/fail()/ - // finalizeCancelled(). Used by shutdown paths (abortAll) to guarantee the - // SDK contract (every registered agent produces exactly one - // task-notification). - finalizeCancellationIfPending(agentId: string): void { - const entry = this.agents.get(agentId); - if (!entry || entry.status !== 'cancelled' || entry.notified) return; - this.emitNotification(entry); - this.emitStatusChange(entry); - } - - /** - * Append a recent tool activity to a running entry's rolling buffer. - * No-op if the entry is not running — late events after a cancellation - * shouldn't leak into the Progress section. - */ - appendActivity(agentId: string, activity: BackgroundActivity): void { - const entry = this.agents.get(agentId); - if (!entry || entry.status !== 'running') return; - - const prior = entry.recentActivities ?? []; - const next = [...prior, activity]; - if (next.length > MAX_RECENT_ACTIVITIES) { - next.splice(0, next.length - MAX_RECENT_ACTIVITIES); - } - entry.recentActivities = next; - this.emitActivityChange(entry); - } - - get(agentId: string): AgentTask | undefined { - return this.agents.get(agentId); - } - - /** - * Snapshot of every entry regardless of status. Used by the TUI - * footer/dialog to render rows for still-running AND terminal-state - * tasks; the headless holdback loop keys off `hasUnfinalizedTasks` - * instead, so callers that only need the running slice can filter - * this snapshot at the call site. - */ - getAll(): AgentTask[] { - return Array.from(this.agents.values()); - } - - private getRunningBackgroundCount(): number { - return Array.from(this.agents.values()).filter( - (entry) => entry.isBackgrounded && entry.status === 'running', - ).length; - } - - /** - * True if any registered task has not yet emitted its terminal - * task-notification. Covers `running` (still executing) and - * `cancelled`-but-not-finalized (cancel requested, but the natural - * handler hasn't fired finalizeCancelled() yet). Headless callers - * must keep their event loop alive while this returns true, so every - * task_started is paired with a matching task_notification. - */ - hasUnfinalizedTasks(): boolean { - for (const entry of this.agents.values()) { - // Foreground entries block the parent tool-call synchronously, so the - // headless event loop is already pinned by the `await` on the caller's - // promise — counting them here would be redundant and would also keep - // the loop alive for entries that don't even emit a notification. - if (!entry.isBackgrounded) continue; - if (entry.status === 'running') return true; - if (entry.status === 'cancelled' && !entry.notified) return true; - } - return false; - } - - /** - * Drops every in-memory entry without touching sidecar state. - * - * Used only when switching to a different session after the caller has - * already established that no live work from the current session is still - * running. Paused/interrupted entries remain recoverable from disk because - * their sidecars keep the persisted status. - */ - reset(): void { - const firstEntry = this.agents.values().next().value as - | AgentTask - | undefined; - if (!firstEntry) return; - for (const agentId of this.agents.keys()) { - this.wakeMessageWaiters(agentId); - } - this.agents.clear(); - this.emitStatusChange(firstEntry); - } - - /** - * Enqueue a message for delivery to a running background agent. - * The agent drains this queue between tool rounds. - */ - queueMessage(agentId: string, message: string): boolean { - return this.queueExternalInput(agentId, message); - } - - /** - * Enqueue generalized external input for an agent. Use queueMessage for the - * parent send_message text path; this lower-level API also accepts - * structured inputs such as owner-routed Monitor notifications. - */ - queueExternalInput(agentId: string, input: AgentExternalInput): boolean { - const entry = this.agents.get(agentId); - if (!entry || entry.status !== 'running') return false; - const queue = entry.pendingMessages!; - queue.push(input); - debugLogger.info( - `Queued message for background agent ${agentId} (${queue.length} pending)`, - ); - this.wakeMessageWaiters(agentId); - return true; - } - - /** - * Drain all pending messages for an agent. Returns the messages - * and clears the queue. Called by the agent's reasoning loop. - */ - drainMessages(agentId: string): AgentExternalInput[] { - const entry = this.agents.get(agentId); - if (!entry || !entry.pendingMessages!.length) return []; - const messages = entry.pendingMessages!.splice(0); - debugLogger.info( - `Drained ${messages.length} message(s) for background agent ${agentId}`, - ); - return messages; - } - - async waitForMessages( - agentId: string, - signal: AbortSignal, - ): Promise { - const immediate = this.drainMessages(agentId); - if (immediate.length > 0) return immediate; - - const entry = this.agents.get(agentId); - if (!entry || entry.status !== 'running' || signal.aborted) return []; - - return new Promise((resolve) => { - const cleanup = () => { - signal.removeEventListener('abort', onAbort); - const waiters = this.messageWaiters.get(agentId); - if (!waiters) return; - waiters.delete(onWake); - if (waiters.size === 0) { - this.messageWaiters.delete(agentId); - } - }; - const resolveWithDrain = () => { - cleanup(); - resolve(this.drainMessages(agentId)); - }; - const onWake = () => resolveWithDrain(); - const onAbort = () => { - cleanup(); - resolve([]); - }; - - let waiters = this.messageWaiters.get(agentId); - if (!waiters) { - waiters = new Set(); - this.messageWaiters.set(agentId, waiters); - } - waiters.add(onWake); - signal.addEventListener('abort', onAbort, { once: true }); - if (signal.aborted) { - cleanup(); - resolve([]); - return; - } - }); - } - - wakeExternalInputWaiters(agentId: string): void { - this.wakeMessageWaiters(agentId); - } - - setNotificationCallback( - cb: BackgroundNotificationCallback | undefined, - ): void { - this.notificationCallback = cb; - } - - setRegisterCallback(cb: BackgroundRegisterCallback | undefined): void { - this.registerCallback = cb; - } - - setStatusChangeCallback( - cb: BackgroundStatusChangeCallback | undefined, - ): void { - this.statusChangeCallback = cb; - } - - setActivityChangeCallback( - cb: BackgroundActivityChangeCallback | undefined, - ): void { - this.activityChangeCallback = cb; - } - - abortAll(options: BackgroundTaskCancelOptions = {}): void { - const cancelOptions: BackgroundTaskCancelOptions = { - persistedStatus: 'running', - ...options, - }; - for (const entry of Array.from(this.agents.values())) { - if (entry.status === 'running') { - this.cancel(entry.agentId, cancelOptions); - } - - if (cancelOptions.notify === false) { - entry.notified = true; - continue; - } - - // Shutdown path: no natural handler will run, so emit the cancelled - // notification here to honour the one-notification-per-agent contract. - this.finalizeCancellationIfPending(entry.agentId); - } - debugLogger.info('Aborted all background agents'); - } - - private buildDisplayLabel(entry: AgentTask): string { - return buildBackgroundEntryLabel(entry); - } - - private emitNotification(entry: AgentTask): void { - // Mark notified *before* invoking the callback so that a re-entrant - // terminal call inside the callback chain (cancel → complete race) - // sees the flag and short-circuits, rather than firing twice. - if (entry.notified) return; - entry.notified = true; - - // Foreground entries return their result through the parent's normal - // tool-result channel (the `returnDisplay` field on the synchronous - // tool-call). Emitting the XML envelope on top would feed the parent - // model the same payload twice. - if (!entry.isBackgrounded) return; - - if (!this.notificationCallback) return; - - const statusText = - entry.status === 'completed' - ? 'completed' - : entry.status === 'failed' - ? 'failed' - : 'was cancelled'; - - const label = this.buildDisplayLabel(entry); - const displayLine = `Background agent "${label}" ${statusText}.`; - - const xmlParts: string[] = [ - '', - `${escapeXml(entry.agentId)}`, - ]; - if (entry.toolUseId) { - xmlParts.push(`${escapeXml(entry.toolUseId)}`); - } - xmlParts.push( - `${escapeXml(entry.status)}`, - `Agent "${escapeXml(entry.description)}" ${statusText}.`, - ); - if (entry.result) { - xmlParts.push(`${escapeXml(entry.result)}`); - } - if (entry.error) { - xmlParts.push(`Error: ${escapeXml(entry.error)}`); - } - if (entry.outputFile) { - xmlParts.push( - `${escapeXml(entry.outputFile)}`, - ); - } - if (entry.stats) { - xmlParts.push( - '', - `${entry.stats.totalTokens}`, - `${entry.stats.toolUses}`, - `${entry.stats.durationMs}`, - '', - ); - } - xmlParts.push(''); - - const meta: NotificationMeta = { - agentId: entry.agentId, - status: entry.status, - stats: entry.stats, - toolUseId: entry.toolUseId, - }; - - try { - this.notificationCallback(displayLine, xmlParts.join('\n'), meta); - } catch (error) { - debugLogger.error('Failed to emit background notification:', error); - } - } - - private emitStatusChange(entry?: AgentTask): void { - this.pruneTerminalEntries(); - if (!this.statusChangeCallback) return; - try { - this.statusChangeCallback(entry); - } catch (error) { - debugLogger.error('Failed to emit background status change:', error); - } - } - - /** - * Evict the oldest fully-finalized terminal entries (those with - * `notified === true`) once their count exceeds - * `MAX_RETAINED_TERMINAL_AGENTS`. Sorted by `endTime` (then - * `startTime` as a tiebreaker for entries that share an endTime). - * - * Running, paused, and cancelled-but-not-yet-notified entries are - * excluded from the eviction set: - * - running / paused: the user explicitly cares about live work, - * and pruning a paused entry would silently drop a recoverable - * task without giving the user a chance to resume / abandon it. - * - cancelled but not notified: the natural handler (or grace - * timer) is still going to fire `finalizeCancelled` / - * `finalizeCancellationIfPending`. Evicting now would break the - * SDK contract that every `register` pairs with exactly one - * terminal `task-notification`. - * - * The caller (typically `emitStatusChange`) is responsible for - * invoking this after every transition that mutates `notified` or - * `endTime`. Cap-exceeded eviction is a best-effort: a transition - * that sets `notified = true` outside the status-change path (the - * `cancel({ notify: false })` shortcut and `abortAll`'s loop body) - * may briefly carry a few extra entries until the next transition - * triggers another prune. Both of those paths are reset / shutdown - * adjacent — the registry is about to be cleared via `reset()` - * anyway, so the extra retention does not leak across sessions. - */ - private pruneTerminalEntries(): void { - const evictable = Array.from(this.agents.values()) - .filter((entry) => entry.notified === true) - .sort( - (a, b) => - (a.endTime ?? a.startTime) - (b.endTime ?? b.startTime) || - a.startTime - b.startTime, - ); - - while (evictable.length > MAX_RETAINED_TERMINAL_AGENTS) { - const oldest = evictable.shift(); - if (oldest) { - this.agents.delete(oldest.agentId); - } - } - } - - private wakeMessageWaiters(agentId: string): void { - const waiters = this.messageWaiters.get(agentId); - if (!waiters) return; - this.messageWaiters.delete(agentId); - for (const waiter of waiters) { - waiter(); - } - } - - private emitActivityChange(entry: AgentTask): void { - if (!this.activityChangeCallback) return; - try { - this.activityChangeCallback(entry); - } catch (error) { - debugLogger.error('Failed to emit background activity change:', error); - } - } -} + * Removal: scheduled for the release after PR 2 lands. New code should + * import from `'../tasks/agent-task.js'` directly. + */ + +export { + type AgentTask, + type AgentTaskRegistration, + type AgentCompletionStats, + type BackgroundActivity, + type BackgroundNotificationCallback, + type BackgroundRegisterCallback, + type BackgroundTaskEntry, + type BackgroundTaskStatus, + type NotificationMeta, + buildBackgroundEntryLabel, + BACKGROUND_AGENT_CONCURRENCY_ENV, + DEFAULT_MAX_CONCURRENT_BACKGROUND_AGENTS, + MAX_CONCURRENT_BACKGROUND_AGENTS, + resolveMaxConcurrentBackgroundAgents, +} from '../tasks/agent-task.js'; diff --git a/packages/core/src/agents/index.ts b/packages/core/src/agents/index.ts index 24e285ab0a8..ae50021eed2 100644 --- a/packages/core/src/agents/index.ts +++ b/packages/core/src/agents/index.ts @@ -16,6 +16,4 @@ export * from './backends/index.js'; export * from './arena/index.js'; export * from './runtime/index.js'; -export * from './background-tasks.js'; export * from './background-agent-resume.js'; -export * from './tasks/types.js'; diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 472db6144df..f31b05c5ef3 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -82,10 +82,13 @@ import { } from '../permissions/denialTracking.js'; import { SubagentManager } from '../subagents/subagent-manager.js'; import type { SubagentConfig } from '../subagents/types.js'; -import { BackgroundTaskRegistry } from '../agents/background-tasks.js'; -import { MonitorRegistry } from '../services/monitorRegistry.js'; +import { TaskRegistry } from '../tasks/registry.js'; +import { registerTaskKind } from '../tasks/dispatcher.js'; +import { AgentTaskKind, agentAbortAll } from '../tasks/agent-task.js'; +import { ShellTaskKind, shellAbortAll } from '../tasks/shell-task.js'; +import { MonitorTaskKind, monitorAbortAll } from '../tasks/monitor-task.js'; +import { DreamTaskKind } from '../tasks/dream-task.js'; import { BackgroundAgentResumeService } from '../agents/background-agent-resume.js'; -import { BackgroundShellRegistry } from '../services/backgroundShellRegistry.js'; import { FileReadCache } from '../services/fileReadCache.js'; import { resolveStopHookBlockingCap } from '../hooks/stopHookCap.js'; import { @@ -174,6 +177,17 @@ import { import { resolveModelId } from '../utils/modelId.js'; import type { ClaudeMarketplaceConfig } from '../extension/claude-converter.js'; +// Register every task kind with the dispatcher at module load. Mirrors +// claw-code's `tasks.ts` pattern: a single place where each kind's +// `Task` implementation is wired up. Placed after all imports so the +// executable side-effect doesn't visually interrupt the import block. +// Idempotent — re-importing this module (or constructing additional +// Config instances) does not re-register. +registerTaskKind(AgentTaskKind); +registerTaskKind(ShellTaskKind); +registerTaskKind(MonitorTaskKind); +registerTaskKind(DreamTaskKind); + // Re-export types export type { AnyToolInvocation, FileFilteringOptions, MCPOAuthConfig }; export { @@ -1026,10 +1040,8 @@ export class Config { private subagentManager!: SubagentManager; private memoryPressureConfig?: MemoryPressureConfig; private memoryPressureMonitor?: MemoryPressureMonitor; - private readonly backgroundTaskRegistry = new BackgroundTaskRegistry(); - private readonly monitorRegistry = new MonitorRegistry(); + private readonly taskRegistry = new TaskRegistry(); private backgroundAgentResumeService?: BackgroundAgentResumeService; - private readonly backgroundShellRegistry = new BackgroundShellRegistry(); // Field initializer runs once on the parent Config; child Configs // built via Object.create(parent) intentionally do NOT pick this up // — see getFileReadCache() for the per-instance lazy initialization @@ -2587,9 +2599,9 @@ export class Config { await this.toolRegistry.stop(); } - this.backgroundTaskRegistry.abortAll(); - this.monitorRegistry.abortAll({ notify: false }); - this.backgroundShellRegistry.abortAll(); + agentAbortAll(this.taskRegistry); + monitorAbortAll(this.taskRegistry, { notify: false }); + shellAbortAll(this.taskRegistry); await this.cleanupArenaRuntime(); } catch (error) { @@ -3829,12 +3841,15 @@ export class Config { return this.subagentManager; } - getBackgroundTaskRegistry(): BackgroundTaskRegistry { - return this.backgroundTaskRegistry; - } - - getMonitorRegistry(): MonitorRegistry { - return this.monitorRegistry; + /** + * Unified task registry covering agents, shells, and monitors. Dream + * consolidation tasks live in {@link MemoryManager} and are surfaced + * via the dream adapter (`tasks/dream-task.ts`); the dispatcher's + * `kill` table treats all four kinds uniformly even though the + * registry only holds three. + */ + getTaskRegistry(): TaskRegistry { + return this.taskRegistry; } getBackgroundAgentResumeService(): BackgroundAgentResumeService { @@ -3848,7 +3863,7 @@ export class Config { async loadPausedBackgroundAgents( sessionId: string = this.getSessionId(), - ): Promise> { + ): Promise> { return this.getBackgroundAgentResumeService().loadPausedBackgroundAgents( sessionId, ); @@ -3857,7 +3872,7 @@ export class Config { async resumeBackgroundAgent( agentId: string, initialMessage?: string, - ): Promise { + ): Promise { return this.getBackgroundAgentResumeService().resumeBackgroundAgent( agentId, initialMessage, @@ -3870,10 +3885,6 @@ export class Config { ); } - getBackgroundShellRegistry(): BackgroundShellRegistry { - return this.backgroundShellRegistry; - } - /** * Session-scoped cache that tracks Read / Edit / WriteFile operations * on files. The cache must be **per-Config-instance** so that each diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index d6e24a605f5..c3e2f1b6794 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -279,6 +279,7 @@ export * from './prompts/mcp-prompts.js'; export * from './skills/index.js'; export * from './subagents/index.js'; export * from './agents/index.js'; +export * from './tasks/index.js'; // ============================================================================ // Follow-up Suggestions diff --git a/packages/core/src/services/backgroundShellRegistry.test.ts b/packages/core/src/services/backgroundShellRegistry.test.ts deleted file mode 100644 index a9398aeb552..00000000000 --- a/packages/core/src/services/backgroundShellRegistry.test.ts +++ /dev/null @@ -1,752 +0,0 @@ -/** - * @license - * Copyright 2026 Qwen Team - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - constants as fsConstants, - mkdtempSync, - rmSync, - symlinkSync, - writeFileSync, -} from 'node:fs'; -import { tmpdir } from 'node:os'; -import { join } from 'node:path'; -import { afterEach, describe, expect, it, vi } from 'vitest'; -import { - BackgroundShellRegistry, - MAX_NOTIFICATION_OUTPUT_TAIL_BYTES, - MAX_RETAINED_TERMINAL_SHELLS, - type ShellTaskRegistration, -} from './backgroundShellRegistry.js'; - -let tmpDirs: string[] = []; - -afterEach(() => { - for (const dir of tmpDirs) { - rmSync(dir, { recursive: true, force: true }); - } - tmpDirs = []; -}); - -function makeOutputFile(content: string): string { - const dir = mkdtempSync(join(tmpdir(), 'qwen-shell-notification-')); - tmpDirs.push(dir); - const file = join(dir, 'shell.output'); - writeFileSync(file, content); - return file; -} - -function makeTempDir(): string { - const dir = mkdtempSync(join(tmpdir(), 'qwen-shell-notification-')); - tmpDirs.push(dir); - return dir; -} - -function makeEntry( - overrides: Partial = {}, -): ShellTaskRegistration { - return { - shellId: 's1', - command: 'sleep 60', - cwd: '/tmp', - status: 'running', - startTime: 1000, - outputPath: '/tmp/s1.output', - abortController: new AbortController(), - ...overrides, - }; -} - -describe('BackgroundShellRegistry', () => { - describe('register / get / getAll', () => { - it('round-trips a registered entry by id', () => { - const reg = new BackgroundShellRegistry(); - const e = makeEntry({ shellId: 'a' }); - reg.register(e); - expect(reg.get('a')).toBe(e); - }); - - it('returns undefined for unknown id', () => { - const reg = new BackgroundShellRegistry(); - expect(reg.get('missing')).toBeUndefined(); - }); - - it('lists all entries via getAll', () => { - const reg = new BackgroundShellRegistry(); - const a = makeEntry({ shellId: 'a' }); - const b = makeEntry({ shellId: 'b' }); - reg.register(a); - reg.register(b); - const all = reg.getAll(); - expect(all).toHaveLength(2); - expect(all).toContain(a); - expect(all).toContain(b); - }); - }); - - describe('complete', () => { - it('transitions running → completed with exitCode and endTime', () => { - const reg = new BackgroundShellRegistry(); - reg.register(makeEntry({ shellId: 'a' })); - reg.complete('a', 0, 2000); - const e = reg.get('a')!; - expect(e.status).toBe('completed'); - expect(e.exitCode).toBe(0); - expect(e.endTime).toBe(2000); - }); - - it('is a no-op when entry is not running', () => { - const reg = new BackgroundShellRegistry(); - reg.register(makeEntry({ shellId: 'a' })); - reg.cancel('a', 1500); - reg.complete('a', 0, 2000); - const e = reg.get('a')!; - expect(e.status).toBe('cancelled'); - expect(e.exitCode).toBeUndefined(); - }); - - it('is a no-op for unknown id', () => { - const reg = new BackgroundShellRegistry(); - expect(() => reg.complete('missing', 0, 0)).not.toThrow(); - }); - }); - - describe('fail', () => { - it('transitions running → failed with error and endTime', () => { - const reg = new BackgroundShellRegistry(); - reg.register(makeEntry({ shellId: 'a' })); - reg.fail('a', 'spawn error', 2000); - const e = reg.get('a')!; - expect(e.status).toBe('failed'); - expect(e.error).toBe('spawn error'); - expect(e.endTime).toBe(2000); - }); - - it('is a no-op when entry is not running', () => { - const reg = new BackgroundShellRegistry(); - reg.register(makeEntry({ shellId: 'a' })); - reg.complete('a', 0, 1500); - reg.fail('a', 'late error', 2000); - const e = reg.get('a')!; - expect(e.status).toBe('completed'); - expect(e.error).toBeUndefined(); - }); - }); - - describe('callbacks', () => { - it('fires register callback synchronously when an entry is added', () => { - const reg = new BackgroundShellRegistry(); - const seen: string[] = []; - reg.setRegisterCallback((entry) => seen.push(entry.shellId)); - - reg.register(makeEntry({ shellId: 'a' })); - reg.register(makeEntry({ shellId: 'b' })); - - expect(seen).toEqual(['a', 'b']); - }); - - it('fires statusChange callback on register too (mirrors BackgroundTaskRegistry)', () => { - const reg = new BackgroundShellRegistry(); - const seen: string[] = []; - reg.setStatusChangeCallback((entry) => { - if (entry) seen.push(entry.shellId); - }); - reg.register(makeEntry({ shellId: 'a' })); - reg.register(makeEntry({ shellId: 'b' })); - expect(seen).toEqual(['a', 'b']); - }); - - it('fires statusChange callback on complete / fail / cancel', () => { - const reg = new BackgroundShellRegistry(); - reg.register(makeEntry({ shellId: 'a' })); - reg.register(makeEntry({ shellId: 'b' })); - reg.register(makeEntry({ shellId: 'c' })); - const transitions: Array<{ id: string; status: string }> = []; - reg.setStatusChangeCallback((entry) => { - if (entry) { - transitions.push({ id: entry.shellId, status: entry.status }); - } - }); - - reg.complete('a', 0, 1000); - reg.fail('b', 'boom', 1100); - reg.cancel('c', 1200); - - expect(transitions).toEqual([ - { id: 'a', status: 'completed' }, - { id: 'b', status: 'failed' }, - { id: 'c', status: 'cancelled' }, - ]); - }); - - it('does not fire statusChange when a transition is a no-op', () => { - const reg = new BackgroundShellRegistry(); - const transitions: string[] = []; - reg.setStatusChangeCallback((entry) => { - if (entry) transitions.push(entry.shellId); - }); - reg.register(makeEntry({ shellId: 'a' })); - reg.complete('a', 0, 1000); - transitions.length = 0; - - reg.complete('a', 0, 2000); // already terminal - reg.fail('a', 'late', 2000); // already terminal - reg.cancel('a', 2000); // already terminal - reg.requestCancel('a'); // already terminal — also no fire - - expect(transitions).toEqual([]); - }); - - it('keeps the registry usable when a callback throws', () => { - const reg = new BackgroundShellRegistry(); - reg.setRegisterCallback(() => { - throw new Error('subscriber blew up'); - }); - - expect(() => reg.register(makeEntry({ shellId: 'a' }))).not.toThrow(); - expect(reg.get('a')!.status).toBe('running'); - }); - - it('clears subscriber when set to undefined', () => { - const reg = new BackgroundShellRegistry(); - const seen: string[] = []; - reg.setRegisterCallback((e) => seen.push(e.shellId)); - reg.register(makeEntry({ shellId: 'a' })); - reg.setRegisterCallback(undefined); - reg.register(makeEntry({ shellId: 'b' })); - expect(seen).toEqual(['a']); - }); - - it('setNotificationCallback(undefined) clears the callback', () => { - // useGeminiStream's cleanup relies on this contract to avoid - // leaked callbacks firing into torn-down React state on unmount. - // If a future refactor breaks the clearing path, stale callbacks - // would fire silently — no test would catch it without this guard. - const reg = new BackgroundShellRegistry(); - const callback = vi.fn(); - reg.setNotificationCallback(callback); - reg.register(makeEntry({ shellId: 'a' })); - reg.setNotificationCallback(undefined); - reg.complete('a', 0, 2000); - expect(callback).not.toHaveBeenCalled(); - }); - }); - - describe('notifications', () => { - it('emits one task-notification when a shell completes', () => { - const reg = new BackgroundShellRegistry(); - const callback = vi.fn(); - const outputPath = makeOutputFile('first line\nfinal result\n'); - reg.setNotificationCallback(callback); - reg.register( - makeEntry({ - shellId: 'a', - command: 'npm test', - cwd: '/repo', - outputPath, - pid: 1234, - }), - ); - - reg.complete('a', 0, 2000); - - expect(callback).toHaveBeenCalledTimes(1); - const [displayText, modelText, meta] = callback.mock.calls[0]; - expect(displayText).toBe('Background shell "npm test" completed.'); - expect(modelText).toContain(''); - expect(modelText).toContain('a'); - expect(modelText).toContain('shell'); - expect(modelText).toContain('completed'); - expect(modelText).toContain('npm test'); - expect(modelText).toContain('/repo'); - expect(modelText).toContain('1234'); - expect(modelText).toContain('0'); - expect(modelText).toContain( - 'first line\nfinal result', - ); - expect(modelText).toContain(`${outputPath}`); - expect(meta).toEqual({ - shellId: 'a', - status: 'completed', - exitCode: 0, - }); - }); - - it('truncates long commands for display, summary, and model XML', () => { - const reg = new BackgroundShellRegistry(); - const callback = vi.fn(); - const command = `node -e ${'a'.repeat(700)}`; - const displayCommand = command.slice(0, 77) + '...'; - const modelCommand = command.slice(0, 497) + '...'; - reg.setNotificationCallback(callback); - reg.register(makeEntry({ shellId: 'a', command })); - - reg.complete('a', 0, 2000); - - const [displayText, modelText] = callback.mock.calls[0]; - expect(displayText).toBe( - `Background shell "${displayCommand}" completed.`, - ); - expect(modelText).toContain( - `Shell command "${displayCommand}" completed.`, - ); - expect(modelText).toContain( - `${modelCommand}`, - ); - expect(modelText).not.toContain(command); - }); - - it('escapes XML and strips display control characters on failure', () => { - const reg = new BackgroundShellRegistry(); - const callback = vi.fn(); - reg.setNotificationCallback(callback); - reg.register( - makeEntry({ - shellId: 'a&b', - command: 'echo "'); - - const [, modelText] = callback.mock.calls[0] as [string, string]; - expect(modelText).toContain('<script>'); - expect(modelText).not.toContain('