refactor(core): collapse three task registries into one TaskRegistry

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<typeof vi.fn>;
     tryCompressChat: ReturnType<typeof vi.fn>;
   };
-  let mockBackgroundTaskRegistry: {
-    setNotificationCallback: ReturnType<typeof vi.fn>;
-  };
-  let mockMonitorRegistry: {
-    setNotificationCallback: ReturnType<typeof vi.fn>;
-  };
-  let mockBackgroundShellRegistry: {
-    setNotificationCallback: ReturnType<typeof vi.fn>;
-  };
   let mockToolRegistry: {
     getTool: ReturnType<typeof vi.fn>;
     ensureTool: ReturnType<typeof vi.fn>;
@@ -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,19 +4271,18 @@ 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);
       expect(internals.notificationQueue).toHaveLength(0);
       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);
     });
 

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 {