fix(core): remediate subagent memory leaks using AbortSignal in MessageBus (#25048)

packages/core/src/confirmation-bus/message-bus.test.ts

@@ -348,4 +348,66 @@ describe('MessageBus', () => {
       );
     });
   });
+
+  describe('subscribe with AbortSignal', () => {
+    it('should remove listener when signal is aborted', async () => {
+      const handler = vi.fn();
+      const controller = new AbortController();
+
+      messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler, {
+        signal: controller.signal,
+      });
+
+      const message: ToolExecutionSuccess<string> = {
+        type: MessageBusType.TOOL_EXECUTION_SUCCESS as const,
+        toolCall: { name: 'test' },
+        result: 'test',
+      };
+
+      controller.abort();
+
+      await messageBus.publish(message);
+
+      expect(handler).not.toHaveBeenCalled();
+    });
+
+    it('should not add listener if signal is already aborted', async () => {
+      const handler = vi.fn();
+      const controller = new AbortController();
+      controller.abort();
+
+      messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler, {
+        signal: controller.signal,
+      });
+
+      const message: ToolExecutionSuccess<string> = {
+        type: MessageBusType.TOOL_EXECUTION_SUCCESS as const,
+        toolCall: { name: 'test' },
+        result: 'test',
+      };
+
+      await messageBus.publish(message);
+
+      expect(handler).not.toHaveBeenCalled();
+    });
+
+    it('should remove abort listener when unsubscribe is called', async () => {
+      const handler = vi.fn();
+      const controller = new AbortController();
+      const signal = controller.signal;
+
+      const removeEventListenerSpy = vi.spyOn(signal, 'removeEventListener');
+
+      messageBus.subscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler, {
+        signal,
+      });
+
+      messageBus.unsubscribe(MessageBusType.TOOL_EXECUTION_SUCCESS, handler);
+
+      expect(removeEventListenerSpy).toHaveBeenCalledWith(
+        'abort',
+        expect.any(Function),
+      );
+    });
+  });
 });

packages/core/src/confirmation-bus/message-bus.ts

@@ -13,6 +13,11 @@ import { safeJsonStringify } from '../utils/safeJsonStringify.js';
 import { debugLogger } from '../utils/debugLogger.js';
 
 export class MessageBus extends EventEmitter {
+  private listenerToAbortCleanup = new WeakMap<
+    object,
+    Map<string, () => void>
+  >();
+
   constructor(
     private readonly policyEngine: PolicyEngine,
     private readonly debug = false,
@@ -145,7 +150,36 @@ export class MessageBus extends EventEmitter {
   subscribe<T extends Message>(
     type: T['type'],
     listener: (message: T) => void,
+    options?: { signal?: AbortSignal },
   ): void {
+    if (options?.signal) {
+      const signal = options.signal;
+      if (signal.aborted) return;
+
+      if (this.listenerToAbortCleanup.get(listener)?.has(type)) return;
+
+      const abortHandler = () => {
+        this.off(type, listener);
+        const typeToCleanup = this.listenerToAbortCleanup.get(listener);
+        if (typeToCleanup) {
+          typeToCleanup.delete(type);
+          if (typeToCleanup.size === 0) {
+            this.listenerToAbortCleanup.delete(listener);
+          }
+        }
+      };
+      signal.addEventListener('abort', abortHandler, { once: true });
+
+      let typeToCleanup = this.listenerToAbortCleanup.get(listener);
+      if (!typeToCleanup) {
+        typeToCleanup = new Map<string, () => void>();
+        this.listenerToAbortCleanup.set(listener, typeToCleanup);
+      }
+      typeToCleanup.set(type, () => {
+        signal.removeEventListener('abort', abortHandler);
+      });
+    }
+
     this.on(type, listener);
   }
 
@@ -154,6 +188,17 @@ export class MessageBus extends EventEmitter {
     listener: (message: T) => void,
   ): void {
     this.off(type, listener);
+    const typeToCleanup = this.listenerToAbortCleanup.get(listener);
+    if (typeToCleanup) {
+      const cleanup = typeToCleanup.get(type);
+      if (cleanup) {
+        cleanup();
+        typeToCleanup.delete(type);
+      }
+      if (typeToCleanup.size === 0) {
+        this.listenerToAbortCleanup.delete(listener);
+      }
+    }
   }
 
   /**