Browse Source

refactor(hooks): simplify HookRegistry+SessionLifecycle per review feedback

- Delete HookRegistry class — revert to explicit sequential hook calls
- Delete extractSessionId utility — inline at call sites
- Remove dead TTL logic from SessionLifecycle
- Fix hasPendingSession empty-string rejection (|| instead of ??)
- Use getSessionId() consistently in MultiplexerSessionManager
Michael Henke 1 month ago
parent
commit
56e9cdee72

+ 1 - 2
src/hooks/foreground-fallback/index.ts

@@ -16,7 +16,6 @@
 
 import type { PluginInput } from '@opencode-ai/plugin';
 import { ALL_AGENT_NAMES } from '../../config/constants';
-import { extractSessionId } from '../../utils';
 import { log } from '../../utils/logger';
 import {
   abortSessionWithTimeout,
@@ -242,7 +241,7 @@ export class ForegroundFallbackManager {
         const props = event.properties as
           | { sessionID?: string; info?: { id?: string } }
           | undefined;
-        const id = extractSessionId(props?.info, props?.sessionID);
+        const id = props?.info?.id || props?.sessionID;
         if (id) {
           log('[foreground-fallback] session.deleted observed', {
             sessionID: id,

+ 0 - 37
src/hooks/hook-registry.test.ts

@@ -1,37 +0,0 @@
-import { describe, expect, test } from 'bun:test';
-import { HookRegistry } from './hook-registry';
-
-describe('HookRegistry', () => {
-  test('dispatch runs handlers in registration order', async () => {
-    const r = new HookRegistry();
-    const order: number[] = [];
-    r.register('test', async () => {
-      order.push(1);
-    });
-    r.register('test', async () => {
-      order.push(2);
-    });
-    await r.dispatch('test', {}, {});
-    expect(order).toEqual([1, 2]);
-  });
-
-  test('unregistered hook point is no-op', async () => {
-    const r = new HookRegistry();
-    await r.dispatch('none', {}, {});
-  });
-
-  test('handlers returns empty for unregistered point', () => {
-    const r = new HookRegistry();
-    expect(r.handlers('x')).toEqual([]);
-  });
-
-  test('dispatch passes input and output to handlers', async () => {
-    const r = new HookRegistry();
-    const captured: unknown[] = [];
-    r.register('test', async (i, o) => {
-      captured.push(i, o);
-    });
-    await r.dispatch('test', { a: 1 }, { b: 2 });
-    expect(captured).toEqual([{ a: 1 }, { b: 2 }]);
-  });
-});

+ 0 - 43
src/hooks/hook-registry.ts

@@ -1,43 +0,0 @@
-export class HookRegistry {
-  #handlers = new Map<
-    string,
-    Array<(input: unknown, output: unknown) => Promise<void>>
-  >();
-  #firedHookPoints = new Set<string>();
-
-  register(
-    hookPoint: string,
-    handler: (input: unknown, output: unknown) => Promise<void>,
-  ): void {
-    if (this.#firedHookPoints.has(hookPoint)) {
-      console.warn(
-        `[hook-registry] "${hookPoint}" already dispatched; late registration may miss events`,
-      );
-    }
-    const group = this.#handlers.get(hookPoint);
-    if (group) {
-      group.push(handler);
-    } else {
-      this.#handlers.set(hookPoint, [handler]);
-    }
-  }
-
-  async dispatch(
-    hookPoint: string,
-    input: unknown,
-    output: unknown,
-  ): Promise<void> {
-    this.#firedHookPoints.add(hookPoint);
-    const group = this.#handlers.get(hookPoint);
-    if (!group) return;
-    for (const handler of group) {
-      await handler(input, output);
-    }
-  }
-
-  handlers(
-    hookPoint: string,
-  ): ReadonlyArray<(input: unknown, output: unknown) => Promise<void>> {
-    return this.#handlers.get(hookPoint) ?? [];
-  }
-}

+ 0 - 1
src/hooks/index.ts

@@ -9,7 +9,6 @@ export {
   ForegroundFallbackManager,
   isRateLimitError,
 } from './foreground-fallback';
-export { HookRegistry } from './hook-registry';
 export { processImageAttachments } from './image-hook';
 export { createJsonErrorRecoveryHook } from './json-error-recovery/hook';
 export { createLoopCommandHook } from './loop-command';

+ 0 - 15
src/hooks/session-lifecycle.test.ts

@@ -43,21 +43,6 @@ describe('SessionLifecycle', () => {
     expect(lc.hasPendingSession('s1')).toBe(false);
   });
 
-  test('hasPendingSession returns false after TTL expiry', () => {
-    const lc = new SessionLifecycle(noop);
-    const now = Date.now();
-    lc.markPending('s1');
-
-    // Simulate time passing beyond TTL (5 minutes)
-    const originalDateNow = Date.now;
-    Date.now = () => now + SessionLifecycle.PENDING_TTL_MS + 1;
-    try {
-      expect(lc.hasPendingSession('s1')).toBe(false);
-    } finally {
-      Date.now = originalDateNow;
-    }
-  });
-
   test('clearSession removes all state', () => {
     const lc = new SessionLifecycle(noop);
     lc.markPending('s1');

+ 0 - 12
src/hooks/session-lifecycle.ts

@@ -1,10 +1,7 @@
 export class SessionLifecycle {
-  static readonly PENDING_TTL_MS = 5 * 60 * 1000;
-
   #cleanupCallbacks: Array<(sessionId: string) => void> = [];
   #pendingSessionIds = new Set<string>();
   #everPendingSessionIds = new Set<string>();
-  #pendingTimestamps = new Map<string, number>();
   #log: (msg: string, meta?: Record<string, unknown>) => void;
 
   constructor(log: (msg: string, meta?: Record<string, unknown>) => void) {
@@ -31,24 +28,16 @@ export class SessionLifecycle {
   markPending(sessionId: string): void {
     this.#pendingSessionIds.add(sessionId);
     this.#everPendingSessionIds.add(sessionId);
-    this.#pendingTimestamps.set(sessionId, Date.now());
   }
 
   /** Atomic — only one caller gets true per markPending call. */
   consumePending(sessionId: string): boolean {
     const had = this.#pendingSessionIds.has(sessionId);
     this.#pendingSessionIds.delete(sessionId);
-    this.#pendingTimestamps.delete(sessionId);
     return had;
   }
 
   hasPendingSession(sessionId: string): boolean {
-    const ts = this.#pendingTimestamps.get(sessionId);
-    if (ts && Date.now() - ts > SessionLifecycle.PENDING_TTL_MS) {
-      this.#pendingTimestamps.delete(sessionId);
-      this.#pendingSessionIds.delete(sessionId);
-      return false;
-    }
     return (
       this.#everPendingSessionIds.has(sessionId) &&
       !this.#pendingSessionIds.has(sessionId)
@@ -58,6 +47,5 @@ export class SessionLifecycle {
   clearSession(sessionId: string): void {
     this.#pendingSessionIds.delete(sessionId);
     this.#everPendingSessionIds.delete(sessionId);
-    this.#pendingTimestamps.delete(sessionId);
   }
 }

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

@@ -3,7 +3,6 @@ import {
   BackgroundJobBoard,
   type BackgroundJobRecord,
   deriveTaskSessionLabel,
-  extractSessionId,
   parseTaskIdFromTaskOutput,
   parseTaskLaunchOutput,
   parseTaskStatusOutput,
@@ -592,10 +591,8 @@ export function createTaskSessionManagerHook(
           (input.event.properties as { status?: { type?: string } } | undefined)
             ?.status?.type === 'idle')
       ) {
-        const sessionId = extractSessionId(
-          input.event.properties?.info,
-          input.event.properties?.sessionID,
-        );
+        const sessionId =
+          input.event.properties?.info?.id || input.event.properties?.sessionID;
         const job = sessionId ? backgroundJobBoard.get(sessionId) : undefined;
         log('[task-session-manager] idle/status idle observed', {
           sessionID: sessionId,
@@ -647,10 +644,8 @@ export function createTaskSessionManagerHook(
       }
 
       if (input.event.type === 'session.error') {
-        const sessionId = extractSessionId(
-          input.event.properties?.info,
-          input.event.properties?.sessionID,
-        );
+        const sessionId =
+          input.event.properties?.info?.id || input.event.properties?.sessionID;
         if (sessionId && options.shouldManageSession(sessionId)) {
           // Only clear injected terminal jobs for fatal errors.
           // Rate-limit errors are recovered by ForegroundFallbackManager
@@ -673,10 +668,8 @@ export function createTaskSessionManagerHook(
         (input.event.properties as { status?: { type?: string } } | undefined)
           ?.status?.type === 'busy'
       ) {
-        const sessionId = extractSessionId(
-          input.event.properties?.info,
-          input.event.properties?.sessionID,
-        );
+        const sessionId =
+          input.event.properties?.info?.id || input.event.properties?.sessionID;
         const before = sessionId
           ? backgroundJobBoard.get(sessionId)
           : undefined;
@@ -709,10 +702,8 @@ export function createTaskSessionManagerHook(
       }
 
       if (input.event.type !== 'session.deleted') return;
-      const sessionId = extractSessionId(
-        input.event.properties?.info,
-        input.event.properties?.sessionID,
-      );
+      const sessionId =
+        input.event.properties?.info?.id || input.event.properties?.sessionID;
       if (!sessionId) return;
 
       log('[task-session-manager] session.deleted observed', {

+ 52 - 73
src/index.ts

@@ -31,7 +31,6 @@ import {
   createReflectCommandHook,
   createTaskSessionManagerHook,
   ForegroundFallbackManager,
-  HookRegistry,
   SessionLifecycle,
 } from './hooks';
 import { processImageAttachments } from './hooks/image-hook';
@@ -56,7 +55,6 @@ import { recordTuiAgentModel, recordTuiAgentModels } from './tui-state';
 import {
   BackgroundJobBoard,
   createDisplayNameMentionRewriter,
-  extractSessionId,
   resolveRuntimeAgentName,
 } from './utils';
 import { isPluginDisabledByEnv } from './utils/env';
@@ -141,13 +139,23 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
   let autoUpdateChecker: ReturnType<typeof createAutoUpdateCheckerHook>;
   let sessionAgentMap: Map<string, string>;
   let sessionLifecycle: SessionLifecycle;
-  let hookRegistry: HookRegistry;
+
   let chatHeadersHook: ReturnType<typeof createChatHeadersHook>;
   let foregroundFallback: ForegroundFallbackManager;
   let deepworkCommandHook: ReturnType<typeof createDeepworkCommandHook>;
   let reflectCommandHook: ReturnType<typeof createReflectCommandHook>;
   let loopCommandHook: ReturnType<typeof createLoopCommandHook>;
   let taskSessionManagerHook: ReturnType<typeof createTaskSessionManagerHook>;
+  let phaseReminder: ReturnType<typeof createPhaseReminderHook>;
+  let filterAvailableSkills: ReturnType<typeof createFilterAvailableSkillsHook>;
+  let postFileToolNudge: ReturnType<typeof createPostFileToolNudgeHook>;
+  let delegateTaskRetry: ReturnType<typeof createDelegateTaskRetryHook>;
+  let applyPatch: ReturnType<typeof createApplyPatchHook>;
+  let jsonErrorRecovery: ReturnType<typeof createJsonErrorRecoveryHook>;
+  let postFileToolNudgeAfter: (i: unknown, o: unknown) => Promise<void>;
+  let delegateTaskRetryAfter: (i: unknown, o: unknown) => Promise<void>;
+  let jsonErrorRecoveryAfter: (i: unknown, o: unknown) => Promise<void>;
+  let taskSessionManagerAfter: (i: unknown, o: unknown) => Promise<void>;
   let backgroundJobBoard: BackgroundJobBoard;
   let interviewManager: ReturnType<typeof createInterviewManager>;
   let presetManager: ReturnType<typeof createPresetManager>;
@@ -303,8 +311,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       coordinator: sessionLifecycle,
     });
 
-    // Initialize hooks and register them with the registry
-    hookRegistry = new HookRegistry();
+    // Initialize hooks and wrapPostToolHook helper for error isolation
 
     // Wrap tool.execute.after handlers with per-hook error isolation.
     // Preserves the old runPostToolHook behavior: one failing hook doesn't
@@ -333,75 +340,34 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       };
     };
 
-    const phaseReminder = createPhaseReminderHook(sessionLifecycle);
-    hookRegistry.register('experimental.chat.messages.transform', (i, o) =>
-      phaseReminder['experimental.chat.messages.transform'](
-        i as never,
-        o as never,
-      ),
-    );
+    phaseReminder = createPhaseReminderHook(sessionLifecycle);
 
-    const filterAvailableSkills = createFilterAvailableSkillsHook(ctx, config);
-    hookRegistry.register('experimental.chat.messages.transform', (i, o) =>
-      filterAvailableSkills['experimental.chat.messages.transform'](
-        i as never,
-        o as never,
-      ),
-    );
+    filterAvailableSkills = createFilterAvailableSkillsHook(ctx, config);
 
-    const postFileToolNudge = createPostFileToolNudgeHook({
+    postFileToolNudge = createPostFileToolNudgeHook({
       shouldInject: (sessionID) =>
         sessionAgentMap.get(sessionID) === 'orchestrator',
       coordinator: sessionLifecycle,
     });
-    hookRegistry.register('experimental.chat.system.transform', (i, o) =>
-      postFileToolNudge['experimental.chat.system.transform'](
-        i as never,
-        o as never,
-      ),
-    );
-    hookRegistry.register(
-      'tool.execute.after',
-      wrapPostToolHook('post-file-tool-nudge', (i, o) =>
-        postFileToolNudge['tool.execute.after'](i as never, o as never),
-      ),
-    );
 
-    const delegateTaskRetry = createDelegateTaskRetryHook(ctx);
-    hookRegistry.register(
-      'tool.execute.after',
-      wrapPostToolHook('delegate-task-retry', (i, o) =>
-        delegateTaskRetry['tool.execute.after'](i as never, o as never),
-      ),
-    );
+    delegateTaskRetry = createDelegateTaskRetryHook(ctx);
 
-    const applyPatch = createApplyPatchHook(ctx);
-    hookRegistry.register('tool.execute.before', (i, o) =>
-      applyPatch['tool.execute.before'](i as never, o as never),
-    );
+    applyPatch = createApplyPatchHook(ctx);
 
-    const jsonErrorRecovery = createJsonErrorRecoveryHook(ctx);
-    hookRegistry.register(
-      'tool.execute.after',
-      wrapPostToolHook('json-error-recovery', (i, o) =>
-        jsonErrorRecovery['tool.execute.after'](i as never, o as never),
-      ),
-    );
+    jsonErrorRecovery = createJsonErrorRecoveryHook(ctx);
 
-    hookRegistry.register('tool.execute.before', (i, o) =>
-      taskSessionManagerHook['tool.execute.before'](i as never, o as never),
+    // Pre-created wrapped handlers for tool.execute.after (error-isolated)
+    postFileToolNudgeAfter = wrapPostToolHook('post-file-tool-nudge', (i, o) =>
+      postFileToolNudge['tool.execute.after'](i as never, o as never),
+    );
+    delegateTaskRetryAfter = wrapPostToolHook('delegate-task-retry', (i, o) =>
+      delegateTaskRetry['tool.execute.after'](i as never, o as never),
     );
-    hookRegistry.register('experimental.chat.messages.transform', (i, o) =>
-      taskSessionManagerHook['experimental.chat.messages.transform'](
-        i as never,
-        o as never,
-      ),
+    jsonErrorRecoveryAfter = wrapPostToolHook('json-error-recovery', (i, o) =>
+      jsonErrorRecovery['tool.execute.after'](i as never, o as never),
     );
-    hookRegistry.register(
-      'tool.execute.after',
-      wrapPostToolHook('task-session-manager', (i, o) =>
-        taskSessionManagerHook['tool.execute.after'](i as never, o as never),
-      ),
+    taskSessionManagerAfter = wrapPostToolHook('task-session-manager', (i, o) =>
+      taskSessionManagerHook['tool.execute.after'](i as never, o as never),
     );
     interviewManager = createInterviewManager(ctx, config);
     presetManager = createPresetManager(ctx, config);
@@ -960,7 +926,7 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
         const props = input.event.properties as
           | { info?: { id?: string }; sessionID?: string }
           | undefined;
-        const sessionID = extractSessionId(props?.info, props?.sessionID);
+        const sessionID = props?.info?.id || props?.sessionID;
 
         if (sessionID) {
           sessionLifecycle.dispatchSessionDeleted(sessionID);
@@ -976,7 +942,11 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
     },
 
     'tool.execute.before': async (input, output) => {
-      await hookRegistry.dispatch('tool.execute.before', input, output);
+      await applyPatch['tool.execute.before'](input as never, output as never);
+      await taskSessionManagerHook['tool.execute.before'](
+        input as never,
+        output as never,
+      );
     },
 
     'command.execute.before': async (input, output) => {
@@ -1101,10 +1071,9 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       }
 
       // Inject ephemeral post-file-tool-nudge reminder
-      await hookRegistry.dispatch(
-        'experimental.chat.system.transform',
-        input,
-        output,
+      await postFileToolNudge['experimental.chat.system.transform'](
+        input as never,
+        output as never,
       );
 
       // Collapse to single system message for provider compatibility.
@@ -1147,15 +1116,25 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
         log,
       });
 
-      await hookRegistry.dispatch(
-        'experimental.chat.messages.transform',
-        input,
-        typedOutput,
+      await phaseReminder['experimental.chat.messages.transform'](
+        input as never,
+        typedOutput as never,
+      );
+      await filterAvailableSkills['experimental.chat.messages.transform'](
+        input as never,
+        typedOutput as never,
+      );
+      await taskSessionManagerHook['experimental.chat.messages.transform'](
+        input as never,
+        typedOutput as never,
       );
     },
 
     'tool.execute.after': async (input, output) => {
-      await hookRegistry.dispatch('tool.execute.after', input, output);
+      await postFileToolNudgeAfter(input, output);
+      await delegateTaskRetryAfter(input, output);
+      await jsonErrorRecoveryAfter(input, output);
+      await taskSessionManagerAfter(input, output);
     },
   };
 };

+ 1 - 5
src/multiplexer/session-manager.ts

@@ -6,7 +6,6 @@ import {
   isServerRunning,
   type Multiplexer,
 } from '../multiplexer';
-import { extractSessionId } from '../utils';
 import type {
   BackgroundJobBoard,
   BackgroundJobState,
@@ -608,10 +607,7 @@ export class MultiplexerSessionManager {
   }
 
   private getSessionId(event: SessionEvent): string | undefined {
-    return extractSessionId(
-      event.properties?.info,
-      event.properties?.sessionID,
-    );
+    return event.properties?.info?.id || event.properties?.sessionID;
   }
 
   private backgroundJobState(

+ 0 - 23
src/utils/extract-session-id.test.ts

@@ -1,23 +0,0 @@
-import { describe, expect, test } from 'bun:test';
-import { extractSessionId } from './extract-session-id';
-
-describe('extractSessionId', () => {
-  test('prefers info.id over sessionID', () => {
-    expect(extractSessionId({ id: 'i' }, 's')).toBe('i');
-  });
-
-  test('falls back to sessionID when info.id missing', () => {
-    expect(extractSessionId({}, 's')).toBe('s');
-    expect(extractSessionId({ id: undefined }, 's')).toBe('s');
-  });
-
-  test('returns undefined when both missing', () => {
-    expect(extractSessionId(undefined, undefined)).toBeUndefined();
-    expect(extractSessionId(null, null)).toBeUndefined();
-    expect(extractSessionId({}, undefined)).toBeUndefined();
-  });
-
-  test('handles null info', () => {
-    expect(extractSessionId(null, 's')).toBe('s');
-  });
-});

+ 0 - 7
src/utils/extract-session-id.ts

@@ -1,7 +0,0 @@
-export function extractSessionId(
-  info: { id?: string } | undefined | null,
-  sessionID: string | undefined | null,
-): string | undefined {
-  // ponytail: ?? undefined converts null to undefined for TS strict mode
-  return info?.id ?? sessionID ?? undefined;
-}

+ 0 - 1
src/utils/index.ts

@@ -1,6 +1,5 @@
 export * from './agent-variant';
 export * from './background-job-board';
-export * from './extract-session-id';
 export * from './internal-initiator';
 export { getLogDir, initLogger, log } from './logger';
 export * from './polling';