Browse Source

refactor: migrate hooks to SessionLifecycle coordinator

Michael Henke 1 month ago
parent
commit
c834286aa4

+ 23 - 16
src/hooks/foreground-fallback/index.test.ts

@@ -1,4 +1,5 @@
 import { beforeEach, describe, expect, mock, test } from 'bun:test';
+import { SessionLifecycle } from '../session-lifecycle';
 import { ForegroundFallbackManager, isRateLimitError } from './index';
 
 type ForegroundFallbackClient = ConstructorParameters<
@@ -735,9 +736,16 @@ describe('ForegroundFallbackManager subagent.session.created', () => {
 // ---------------------------------------------------------------------------
 
 describe('ForegroundFallbackManager session.deleted', () => {
-  test('cleans up session state on session.deleted preventing memory leaks', async () => {
+  test('cleans up session state on session.deleted via coordinator', async () => {
+    const coordinator = new SessionLifecycle(() => {});
     const { client, mocks } = createMockClient();
-    const mgr = new ForegroundFallbackManager(client, makeChains(), true);
+    const mgr = new ForegroundFallbackManager(
+      client,
+      makeChains(),
+      true,
+      3,
+      coordinator,
+    );
 
     // Populate all maps for this session
     await mgr.handleEvent({
@@ -752,11 +760,8 @@ describe('ForegroundFallbackManager session.deleted', () => {
       },
     });
 
-    // Delete the session
-    await mgr.handleEvent({
-      type: 'session.deleted',
-      properties: { sessionID: 'sess-del' },
-    });
+    // Cleanup via coordinator
+    coordinator.dispatchSessionDeleted('sess-del');
 
     // After deletion, a new rate-limit on the same ID should behave as a fresh
     // session (no prior model known → uses chain from start, dedup cleared)
@@ -789,11 +794,16 @@ describe('ForegroundFallbackManager session.deleted', () => {
     ).resolves.toBeUndefined();
   });
 
-  test('cleans up state using info.id shape (top-level session deletion)', async () => {
-    // OpenCode emits { properties: { info: { id } } } for top-level sessions
-    // and { properties: { sessionID } } for subagent sessions. Both must clean up.
+  test('cleans up state using info.id shape via coordinator', async () => {
+    const coordinator = new SessionLifecycle(() => {});
     const { client, mocks } = createMockClient();
-    const mgr = new ForegroundFallbackManager(client, makeChains(), true);
+    const mgr = new ForegroundFallbackManager(
+      client,
+      makeChains(),
+      true,
+      3,
+      coordinator,
+    );
 
     // Seed state for the session
     await mgr.handleEvent({
@@ -808,11 +818,8 @@ describe('ForegroundFallbackManager session.deleted', () => {
       },
     });
 
-    // Delete via the info.id shape
-    await mgr.handleEvent({
-      type: 'session.deleted',
-      properties: { info: { id: 'sess-info-del' } },
-    });
+    // Cleanup via coordinator
+    coordinator.dispatchSessionDeleted('sess-info-del');
 
     // State is cleared: a new rate-limit on same ID should behave as fresh session
     await mgr.handleEvent({

+ 18 - 14
src/hooks/foreground-fallback/index.ts

@@ -22,6 +22,7 @@ import {
   abortSessionWithTimeout,
   parseModelReference,
 } from '../../utils/session';
+import type { SessionLifecycle } from '../session-lifecycle';
 import { isUserMessageWithParts } from '../types';
 
 type OpencodeClient = PluginInput['client'];
@@ -121,7 +122,20 @@ export class ForegroundFallbackManager {
     private readonly enabled: boolean,
     /** Consecutive 429s tolerated on the same model before swap/abort. */
     private readonly maxRetries: number = 3,
-  ) {}
+    coordinator?: SessionLifecycle,
+  ) {
+    if (coordinator) {
+      coordinator.onSessionDeleted((id) => {
+        this.sessionModel.delete(id);
+        this.sessionAgent.delete(id);
+        this.sessionTried.delete(id);
+        this.inProgress.delete(id);
+        this.lastTrigger.delete(id);
+        this.lastTriggerModel.delete(id);
+        this.sessionRetries.delete(id);
+      });
+    }
+  }
 
   /**
    * Process an OpenCode plugin event.
@@ -225,24 +239,14 @@ export class ForegroundFallbackManager {
       }
 
       case 'session.deleted': {
-        // Clean up all per-session state to prevent unbounded memory growth
-        // in long-running instances with many subagent sessions.
-        // OpenCode emits two shapes depending on context:
-        //   { properties: { sessionID } }   - subagent / task sessions
-        //   { properties: { info: { id } } } - top-level session deletion
-        // Mirror the same dual-shape lookup used elsewhere in the plugin.
         const props = event.properties as
           | { sessionID?: string; info?: { id?: string } }
           | undefined;
         const id = extractSessionId(props?.info, props?.sessionID);
         if (id) {
-          this.sessionModel.delete(id);
-          this.sessionAgent.delete(id);
-          this.sessionTried.delete(id);
-          this.inProgress.delete(id);
-          this.lastTrigger.delete(id);
-          this.lastTriggerModel.delete(id);
-          this.sessionRetries.delete(id);
+          log('[foreground-fallback] session.deleted observed', {
+            sessionID: id,
+          });
         }
         break;
       }

+ 3 - 3
src/hooks/phase-reminder/index.ts

@@ -7,7 +7,7 @@
  */
 import { PHASE_REMINDER } from '../../config/constants';
 import { SLIM_INTERNAL_INITIATOR_MARKER } from '../../utils';
-import { hasPendingSession } from '../post-file-tool-nudge';
+import type { SessionLifecycle } from '../session-lifecycle';
 import { isUserMessageWithParts } from '../types';
 
 export { PHASE_REMINDER };
@@ -17,7 +17,7 @@ export { PHASE_REMINDER };
  * This hook runs right before sending to API, so it doesn't affect UI display.
  * Only injects for the orchestrator agent.
  */
-export function createPhaseReminderHook() {
+export function createPhaseReminderHook(coordinator?: SessionLifecycle) {
   return {
     'experimental.chat.messages.transform': async (
       _input: Record<string, never>,
@@ -55,7 +55,7 @@ export function createPhaseReminderHook() {
       // injection via system prompt — skip message-level injection.
       const sessionId = (lastUserMessage as { info?: { sessionID?: string } })
         ?.info?.sessionID;
-      if (sessionId && hasPendingSession(sessionId)) {
+      if (sessionId && coordinator?.hasPendingSession(sessionId)) {
         return;
       }
 

+ 31 - 20
src/hooks/post-file-tool-nudge/index.test.ts

@@ -1,11 +1,13 @@
 import { describe, expect, test } from 'bun:test';
 
 import { PHASE_REMINDER } from '../../config/constants';
+import { SessionLifecycle } from '../session-lifecycle';
 import { createPostFileToolNudgeHook } from './index';
 
 describe('post-file-tool-nudge hook', () => {
   test('records pending session on Read tool', async () => {
-    const hook = createPostFileToolNudgeHook();
+    const coordinator = new SessionLifecycle(() => {});
+    const hook = createPostFileToolNudgeHook({ coordinator });
     const output = { system: [] };
 
     await hook['tool.execute.after']({ tool: 'Read', sessionID: 's1' }, {});
@@ -18,7 +20,8 @@ describe('post-file-tool-nudge hook', () => {
   });
 
   test('records pending session on Write tool', async () => {
-    const hook = createPostFileToolNudgeHook();
+    const coordinator = new SessionLifecycle(() => {});
+    const hook = createPostFileToolNudgeHook({ coordinator });
     const output = { system: [] };
 
     await hook['tool.execute.after']({ tool: 'Write', sessionID: 's1' }, {});
@@ -31,7 +34,8 @@ describe('post-file-tool-nudge hook', () => {
   });
 
   test('does not mutate tool output', async () => {
-    const hook = createPostFileToolNudgeHook();
+    const coordinator = new SessionLifecycle(() => {});
+    const hook = createPostFileToolNudgeHook({ coordinator });
     const toolOutput = { output: 'real content' };
 
     await hook['tool.execute.after'](
@@ -43,7 +47,8 @@ describe('post-file-tool-nudge hook', () => {
   });
 
   test('deduplicates multiple Read/Write calls in same session', async () => {
-    const hook = createPostFileToolNudgeHook();
+    const coordinator = new SessionLifecycle(() => {});
+    const hook = createPostFileToolNudgeHook({ coordinator });
 
     await hook['tool.execute.after']({ tool: 'read', sessionID: 's1' }, {});
     await hook['tool.execute.after']({ tool: 'write', sessionID: 's1' }, {});
@@ -59,7 +64,8 @@ describe('post-file-tool-nudge hook', () => {
   });
 
   test('consumes pending marker after injection', async () => {
-    const hook = createPostFileToolNudgeHook();
+    const coordinator = new SessionLifecycle(() => {});
+    const hook = createPostFileToolNudgeHook({ coordinator });
 
     await hook['tool.execute.after']({ tool: 'Read', sessionID: 's1' }, {});
     await hook['experimental.chat.system.transform'](
@@ -78,7 +84,8 @@ describe('post-file-tool-nudge hook', () => {
   });
 
   test('ignores non-file tools', async () => {
-    const hook = createPostFileToolNudgeHook();
+    const coordinator = new SessionLifecycle(() => {});
+    const hook = createPostFileToolNudgeHook({ coordinator });
     const output = { system: [] };
 
     await hook['tool.execute.after']({ tool: 'bash', sessionID: 's1' }, {});
@@ -91,7 +98,11 @@ describe('post-file-tool-nudge hook', () => {
   });
 
   test('skips injection when shouldInject returns false', async () => {
-    const hook = createPostFileToolNudgeHook({ shouldInject: () => false });
+    const coordinator = new SessionLifecycle(() => {});
+    const hook = createPostFileToolNudgeHook({
+      shouldInject: () => false,
+      coordinator,
+    });
     const output = { system: [] };
 
     await hook['tool.execute.after']({ tool: 'Read', sessionID: 's1' }, {});
@@ -104,7 +115,8 @@ describe('post-file-tool-nudge hook', () => {
   });
 
   test('ignores Read/Write without sessionID', async () => {
-    const hook = createPostFileToolNudgeHook();
+    const coordinator = new SessionLifecycle(() => {});
+    const hook = createPostFileToolNudgeHook({ coordinator });
     const output = { system: [] };
 
     await hook['tool.execute.after']({ tool: 'read' }, {});
@@ -116,13 +128,12 @@ describe('post-file-tool-nudge hook', () => {
     expect(output.system).toHaveLength(0);
   });
 
-  test('cleans up pending marker on session.deleted', async () => {
-    const hook = createPostFileToolNudgeHook();
+  test('cleans up pending marker on session.deleted via coordinator', async () => {
+    const coordinator = new SessionLifecycle(() => {});
+    const hook = createPostFileToolNudgeHook({ coordinator });
 
     await hook['tool.execute.after']({ tool: 'Read', sessionID: 's1' }, {});
-    await hook.event({
-      event: { type: 'session.deleted', properties: { sessionID: 's1' } },
-    });
+    coordinator.dispatchSessionDeleted('s1');
 
     const output = { system: [] };
     await hook['experimental.chat.system.transform'](
@@ -133,13 +144,12 @@ describe('post-file-tool-nudge hook', () => {
     expect(output.system).toHaveLength(0);
   });
 
-  test('cleans up on session.deleted with info.id shape', async () => {
-    const hook = createPostFileToolNudgeHook();
+  test('cleans up pending marker via coordinator with info.id shape', async () => {
+    const coordinator = new SessionLifecycle(() => {});
+    const hook = createPostFileToolNudgeHook({ coordinator });
 
     await hook['tool.execute.after']({ tool: 'Read', sessionID: 's1' }, {});
-    await hook.event({
-      event: { type: 'session.deleted', properties: { info: { id: 's1' } } },
-    });
+    coordinator.dispatchSessionDeleted('s1');
 
     const output = { system: [] };
     await hook['experimental.chat.system.transform'](
@@ -152,8 +162,9 @@ describe('post-file-tool-nudge hook', () => {
 
   test('composed: phase-reminder skips when post-file-tool-nudge handles system', async () => {
     const { createPhaseReminderHook } = await import('../phase-reminder/index');
-    const nudgeHook = createPostFileToolNudgeHook();
-    const phaseHook = createPhaseReminderHook();
+    const coordinator = new SessionLifecycle(() => {});
+    const nudgeHook = createPostFileToolNudgeHook({ coordinator });
+    const phaseHook = createPhaseReminderHook(coordinator);
 
     // Simulate Read tool call
     await nudgeHook['tool.execute.after'](

+ 13 - 51
src/hooks/post-file-tool-nudge/index.ts

@@ -7,81 +7,43 @@
  */
 
 import { PHASE_REMINDER } from '../../config/constants';
-import { extractSessionId } from '../../utils';
+import type { SessionLifecycle } from '../session-lifecycle';
 
-interface ToolExecuteAfterInput {
-  tool: string;
-  sessionID?: string;
-  callID?: string;
-}
+const FILE_TOOLS = new Set(['Read', 'read', 'Write', 'write']);
 
 interface PostFileToolNudgeOptions {
   shouldInject?: (sessionID: string) => boolean;
-}
-
-const FILE_TOOLS = new Set(['Read', 'read', 'Write', 'write']);
-
-// Module-scoped for coordination with phase-reminder hook.
-const pendingSessionIds = new Set<string>();
-const everPendingSessionIds = new Set<string>();
-
-/** Check if a session was marked pending by a file tool AND has not yet been
- *  consumed by system.transform. Allows phase-reminder to skip injection
- *  when post-file-tool-nudge already handles it. */
-export function hasPendingSession(sessionId: string): boolean {
-  return (
-    everPendingSessionIds.has(sessionId) && !pendingSessionIds.has(sessionId)
-  );
+  coordinator?: SessionLifecycle;
 }
 
 export function createPostFileToolNudgeHook(
   options: PostFileToolNudgeOptions = {},
 ) {
+  const { coordinator } = options;
+
+  if (coordinator) {
+    coordinator.onSessionDeleted((sid) => coordinator.clearSession(sid));
+  }
+
   return {
     'tool.execute.after': async (
-      input: ToolExecuteAfterInput,
+      input: { tool: string; sessionID?: string; callID?: string },
       _output: unknown,
     ): Promise<void> => {
-      if (!FILE_TOOLS.has(input.tool) || !input.sessionID) {
-        return;
-      }
-
-      pendingSessionIds.add(input.sessionID);
-      everPendingSessionIds.add(input.sessionID);
+      if (!FILE_TOOLS.has(input.tool) || !input.sessionID) return;
+      coordinator?.markPending(input.sessionID);
     },
     'experimental.chat.system.transform': async (
       input: { sessionID?: string },
       output: { system: string[] },
     ): Promise<void> => {
-      if (!input.sessionID || !pendingSessionIds.delete(input.sessionID)) {
+      if (!input.sessionID || !coordinator?.consumePending(input.sessionID)) {
         return;
       }
-
-      // Track consumption so phase-reminder can check without consuming.
-      // (already tracked via everPendingSessionIds — delete from pending is
-      // sufficient signal)
-
       if (options.shouldInject && !options.shouldInject(input.sessionID)) {
         return;
       }
-
       output.system.push(PHASE_REMINDER);
     },
-    event: async (input: {
-      event: {
-        type: string;
-        properties?: { info?: { id?: string }; sessionID?: string };
-      };
-    }): Promise<void> => {
-      if (input.event.type !== 'session.deleted') return;
-      const sid = extractSessionId(
-        input.event.properties?.info,
-        input.event.properties?.sessionID,
-      );
-      if (sid) {
-        pendingSessionIds.delete(sid);
-        everPendingSessionIds.delete(sid);
-      }
-    },
   };
 }

+ 12 - 18
src/hooks/task-session-manager/index.test.ts

@@ -1,4 +1,5 @@
 import { describe, expect, mock, test } from 'bun:test';
+import { SessionLifecycle } from '../../hooks/session-lifecycle';
 import { BackgroundJobBoard } from '../../utils';
 import { createTaskSessionManagerHook } from './index';
 
@@ -9,6 +10,7 @@ function createHook(options?: {
   backgroundJobBoard?: BackgroundJobBoard;
   sessionStatus?: unknown;
   isFallbackInProgress?: (sessionID: string) => boolean;
+  coordinator?: SessionLifecycle;
 }) {
   const hook = createTaskSessionManagerHook(
     {
@@ -27,6 +29,7 @@ function createHook(options?: {
       backgroundJobBoard: options?.backgroundJobBoard,
       shouldManageSession: options?.shouldManageSession ?? (() => true),
       isFallbackInProgress: options?.isFallbackInProgress,
+      coordinator: options?.coordinator,
     },
   );
 
@@ -1743,7 +1746,8 @@ describe('task-session-manager hook', () => {
   });
 
   test('cleans up background jobs when parent or child is deleted', async () => {
-    const { hook } = createHook();
+    const coordinator = new SessionLifecycle(() => {});
+    const { hook } = createHook({ coordinator });
 
     await hook['tool.execute.before'](
       {
@@ -1770,12 +1774,7 @@ describe('task-session-manager hook', () => {
       },
     );
 
-    await hook.event({
-      event: {
-        type: 'session.deleted',
-        properties: { sessionID: 'child-1' },
-      },
-    });
+    coordinator.dispatchSessionDeleted('child-1');
 
     const messages = createMessages('parent-1', 'do something');
     await hook['experimental.chat.messages.transform']({}, messages);
@@ -1784,7 +1783,8 @@ describe('task-session-manager hook', () => {
   });
 
   test('cleans pending calls when parent session is deleted', async () => {
-    const { hook } = createHook();
+    const coordinator = new SessionLifecycle(() => {});
+    const { hook } = createHook({ coordinator });
 
     await hook['tool.execute.before'](
       {
@@ -1800,12 +1800,7 @@ describe('task-session-manager hook', () => {
       },
     );
 
-    await hook.event({
-      event: {
-        type: 'session.deleted',
-        properties: { sessionID: 'parent-1' },
-      },
-    });
+    coordinator.dispatchSessionDeleted('parent-1');
 
     await hook['tool.execute.after'](
       {
@@ -2028,8 +2023,9 @@ describe('task-session-manager hook', () => {
   });
 
   test('parent deletion clears jobs and pending calls', async () => {
+    const coordinator = new SessionLifecycle(() => {});
     const board = new BackgroundJobBoard();
-    const { hook } = createHook({ backgroundJobBoard: board });
+    const { hook } = createHook({ backgroundJobBoard: board, coordinator });
     await hook['tool.execute.before'](
       { tool: 'task', sessionID: 'parent-1', callID: 'call-1' },
       { args: { subagent_type: 'oracle', description: 'architecture review' } },
@@ -2041,9 +2037,7 @@ describe('task-session-manager hook', () => {
       description: 'architecture review',
     });
 
-    await hook.event({
-      event: { type: 'session.deleted', properties: { sessionID: 'parent-1' } },
-    });
+    coordinator.dispatchSessionDeleted('parent-1');
     await hook['tool.execute.after'](
       { tool: 'task', sessionID: 'parent-1', callID: 'call-1' },
       { output: ['task_id: child-2', 'state: running'].join('\n') },

+ 17 - 26
src/hooks/task-session-manager/index.ts

@@ -12,6 +12,7 @@ import {
 import { isRecord as isObjectRecord } from '../../utils/guards';
 import { log } from '../../utils/logger';
 import { isRateLimitError } from '../foreground-fallback/index';
+import type { SessionLifecycle } from '../session-lifecycle';
 import {
   isUserMessageWithParts,
   type MessagePart,
@@ -87,10 +88,11 @@ export function createTaskSessionManagerHook(
     shouldManageSession: (sessionID: string) => boolean;
     /** Optional guard: when provided, idle events for a session that is
      *  currently undergoing a foreground-fallback abort/re-prompt cycle
-     *  will NOT trigger idle reconciliation. Prevents marking a still-
+     *  will NOT trigger idle reconciliation. prevents marking a still-
      *  active child job as completed when the session was aborted for
      *  model fallback rather than natural completion. */
     isFallbackInProgress?: (sessionID: string) => boolean;
+    coordinator?: SessionLifecycle;
   },
 ) {
   const backgroundJobBoard =
@@ -108,6 +110,17 @@ export function createTaskSessionManagerHook(
   const processedInjectedCompletionOrder: string[] = [];
   const terminalJobsInjectedByParent = new Map<string, Set<string>>();
 
+  if (options.coordinator) {
+    options.coordinator.onSessionDeleted((sessionId) => {
+      backgroundJobBoard.drop(sessionId);
+      backgroundJobBoard.clearParent(sessionId);
+      terminalJobsInjectedByParent.delete(sessionId);
+      taskContextTracker.clearSession(sessionId);
+      taskContextTracker.prune(backgroundJobBoard);
+      pendingCallTracker.clearSession(sessionId);
+    });
+  }
+
   function updateBackgroundJobFromOutput(
     output: unknown,
   ): BackgroundJobRecord | undefined {
@@ -702,31 +715,9 @@ export function createTaskSessionManagerHook(
       );
       if (!sessionId) return;
 
-      log(
-        '[task-session-manager] session.deleted observed; clearing job state',
-        {
-          sessionID: sessionId,
-          deletedJob: (() => {
-            const record = backgroundJobBoard.get(sessionId);
-            return record
-              ? {
-                  state: record.state,
-                  parentSessionID: record.parentSessionID,
-                  alias: record.alias,
-                }
-              : undefined;
-          })(),
-          childJobCount: backgroundJobBoard.list(sessionId).length,
-          managesSession: options.shouldManageSession(sessionId),
-        },
-      );
-
-      backgroundJobBoard.drop(sessionId);
-      backgroundJobBoard.clearParent(sessionId);
-      terminalJobsInjectedByParent.delete(sessionId);
-      taskContextTracker.clearSession(sessionId);
-      taskContextTracker.prune(backgroundJobBoard);
-      pendingCallTracker.clearSession(sessionId);
+      log('[task-session-manager] session.deleted observed', {
+        sessionID: sessionId,
+      });
     },
   };