Browse Source

fix: rearm continuation without message ids

Alvin Unreal 3 weeks ago
parent
commit
66b78bff9e

+ 29 - 12
src/hooks/task-session-manager/continuation-attempt-gate.ts

@@ -10,13 +10,17 @@ type AttemptState =
   | { status: 'reserved'; owner: symbol }
   | { status: 'consumed' };
 
+type RearmIdentity = string | symbol;
+
 type ContinuationAttemptStore = {
   attempts: Map<string, AttemptState>;
   /**
-   * Last external user message ID that rearmed each session. Process-global so
-   * two hook instances observing the same chat.message open only one epoch.
+   * Last external user-message identity that rearmed each session.
+   * string = chat.message ID; symbol = same-process object identity fallback.
    */
-  lastRearmMessageID: Map<string, string>;
+  lastRearmIdentity: Map<string, RearmIdentity>;
+  /** Stable symbols for ID-less output.message object identity (same process). */
+  messageObjectIdentity: WeakMap<object, symbol>;
 };
 
 const STORE_KEY = Symbol.for('oh-my-opencode-slim.continuation-attempt-gate');
@@ -27,11 +31,22 @@ function getStore(): ContinuationAttemptStore {
   };
   globalWithStore[STORE_KEY] ??= {
     attempts: new Map(),
-    lastRearmMessageID: new Map(),
+    lastRearmIdentity: new Map(),
+    messageObjectIdentity: new WeakMap(),
   };
   return globalWithStore[STORE_KEY];
 }
 
+function resolveRearmIdentity(identity: string | object): RearmIdentity {
+  if (typeof identity === 'string') return identity;
+  const store = getStore();
+  const existing = store.messageObjectIdentity.get(identity);
+  if (existing) return existing;
+  const token = Symbol('continuation-rearm-message');
+  store.messageObjectIdentity.set(identity, token);
+  return token;
+}
+
 /**
  * Atomically reserve a continuation attempt.
  * Returns an owner token on success, or null if already reserved/consumed.
@@ -80,19 +95,20 @@ export function releaseContinuationAttempt(
 
 /**
  * Open a new continuation epoch for a real external user message.
- * Idempotent per (sessionID, messageID): a second observe of the same message
- * (e.g. another hook instance) is a no-op and does not rearm again.
- * Returns true when this call cleared attempt state.
+ * Idempotent per (sessionID, identity): string message IDs or same-process
+ * object identity (WeakMap→symbol). A second observe of the same identity
+ * does not rearm again. Returns true when this call cleared attempt state.
  */
 export function rearmContinuationForUserMessage(
   sessionID: string,
-  messageID: string,
+  identity: string | object,
 ): boolean {
   const store = getStore();
-  if (store.lastRearmMessageID.get(sessionID) === messageID) {
+  const resolved = resolveRearmIdentity(identity);
+  if (store.lastRearmIdentity.get(sessionID) === resolved) {
     return false;
   }
-  store.lastRearmMessageID.set(sessionID, messageID);
+  store.lastRearmIdentity.set(sessionID, resolved);
   store.attempts.delete(sessionID);
   return true;
 }
@@ -104,7 +120,7 @@ export function rearmContinuationForUserMessage(
 export function clearContinuationAttempt(sessionID: string): void {
   const store = getStore();
   store.attempts.delete(sessionID);
-  store.lastRearmMessageID.delete(sessionID);
+  store.lastRearmIdentity.delete(sessionID);
 }
 
 export function hasConsumedContinuationAttempt(sessionID: string): boolean {
@@ -115,5 +131,6 @@ export function hasConsumedContinuationAttempt(sessionID: string): boolean {
 export function resetContinuationAttemptGateForTests(): void {
   const store = getStore();
   store.attempts.clear();
-  store.lastRearmMessageID.clear();
+  store.lastRearmIdentity.clear();
+  // WeakMap entries are not enumerable; leave for GC. Tests use fresh objects.
 }

+ 9 - 5
src/hooks/task-session-manager/continuation-token-manager.ts

@@ -80,12 +80,16 @@ export function createContinuationTokenManager(options?: {
 
   /**
    * Real external user message: process-global attempt clear is idempotent per
-   * message ID (only the first observe opens a new epoch). Always invalidate
-   * this instance's local timers/tokens/reservations so a pre-message idle
-   * timer on a second hook cannot fire SDK reads after the shared observe.
+   * message identity (string ID or same-process message object). Always
+   * invalidate this instance's local timers/tokens/reservations so a
+   * pre-message idle timer on a second hook cannot fire SDK reads after the
+   * shared observe.
    */
-  function rearmForUserMessage(sessionID: string, messageID: string): void {
-    rearmContinuationForUserMessage(sessionID, messageID);
+  function rearmForUserMessage(
+    sessionID: string,
+    messageIdentity: string | object,
+  ): void {
+    rearmContinuationForUserMessage(sessionID, messageIdentity);
     invalidateContinuation(sessionID);
   }
 

+ 207 - 0
src/hooks/task-session-manager/index.test.ts

@@ -4438,6 +4438,213 @@ describe('task-session-manager hook', () => {
     expect(promptAsync).toHaveBeenCalledTimes(2);
   });
 
+  test('output.message.id rearms when input.messageID is missing', async () => {
+    const promptAsync = mock(async () => ({}));
+    const { hook } = createHook({
+      idleReconcileDelayMs: 0,
+      sessionClient: {
+        todo: mock(async () => ({ data: [{ status: 'pending' }] })),
+        children: mock(async () => ({ data: [] })),
+        status: mock(async () => ({ data: {} })),
+        promptAsync,
+      },
+    });
+
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(1);
+
+    hook.observeChatMessage(
+      { sessionID: 'parent-1' },
+      {
+        message: {
+          id: 'msg-output-id-only',
+          role: 'user',
+          sessionID: 'parent-1',
+        },
+        parts: [{ type: 'text', text: 'continue' }],
+      },
+    );
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(2);
+  });
+
+  test('ID-less output.message object identity rearms once', async () => {
+    const promptAsync = mock(async () => ({}));
+    const { hook } = createHook({
+      idleReconcileDelayMs: 0,
+      sessionClient: {
+        todo: mock(async () => ({ data: [{ status: 'pending' }] })),
+        children: mock(async () => ({ data: [] })),
+        status: mock(async () => ({ data: {} })),
+        promptAsync,
+      },
+    });
+    const message = { role: 'user', sessionID: 'parent-1' };
+
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(1);
+
+    hook.observeChatMessage(
+      { sessionID: 'parent-1' },
+      { message, parts: [{ type: 'text', text: 'continue' }] },
+    );
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(2);
+
+    // Same object again must not open another epoch.
+    hook.observeChatMessage(
+      { sessionID: 'parent-1' },
+      { message, parts: [{ type: 'text', text: 'continue' }] },
+    );
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(2);
+  });
+
+  test('two hooks share ID-less output.message object identity', async () => {
+    const promptAsync = mock(async () => ({}));
+    const sessionClient = {
+      todo: mock(async () => ({ data: [{ status: 'pending' }] })),
+      children: mock(async () => ({ data: [] })),
+      status: mock(async () => ({ data: {} })),
+      promptAsync,
+    };
+    const makeHook = () =>
+      createTaskSessionManagerHook(
+        {
+          client: { session: sessionClient },
+          directory: '/tmp',
+          worktree: '/tmp',
+        } as never,
+        {
+          maxSessionsPerAgent: 2,
+          continueOnIdle: true,
+          idleReconcileDelayMs: 0,
+          shouldManageSession: () => true,
+        },
+      );
+    const hookA = makeHook();
+    const hookB = makeHook();
+    const message = { role: 'user', sessionID: 'parent-1' };
+    const output = {
+      message,
+      parts: [{ type: 'text', text: 'continue' }],
+    };
+
+    await hookA.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(1);
+
+    hookA.observeChatMessage({ sessionID: 'parent-1' }, output);
+    await hookA.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(2);
+
+    hookB.observeChatMessage({ sessionID: 'parent-1' }, output);
+    await hookB.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(2);
+  });
+
+  test('distinct ID-less message objects each open a new epoch', async () => {
+    const promptAsync = mock(async () => ({}));
+    const { hook } = createHook({
+      idleReconcileDelayMs: 0,
+      sessionClient: {
+        todo: mock(async () => ({ data: [{ status: 'pending' }] })),
+        children: mock(async () => ({ data: [] })),
+        status: mock(async () => ({ data: {} })),
+        promptAsync,
+      },
+    });
+
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(1);
+
+    const text = 'identical text must not dedupe distinct objects';
+    hook.observeChatMessage(
+      { sessionID: 'parent-1' },
+      {
+        message: { role: 'user', sessionID: 'parent-1' },
+        parts: [{ type: 'text', text }],
+      },
+    );
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(2);
+
+    hook.observeChatMessage(
+      { sessionID: 'parent-1' },
+      {
+        message: { role: 'user', sessionID: 'parent-1' },
+        parts: [{ type: 'text', text }],
+      },
+    );
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(3);
+  });
+
+  test('missing id and output.message fails closed without rearm', async () => {
+    const promptAsync = mock(async () => ({}));
+    const { hook } = createHook({
+      idleReconcileDelayMs: 0,
+      sessionClient: {
+        todo: mock(async () => ({ data: [{ status: 'pending' }] })),
+        children: mock(async () => ({ data: [] })),
+        status: mock(async () => ({ data: {} })),
+        promptAsync,
+      },
+    });
+
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(1);
+
+    // sessionID only on input; no messageID and no output.message object.
+    hook.observeChatMessage(
+      {
+        sessionID: 'parent-1',
+        parts: [{ type: 'text', text: 'continue' }],
+      },
+      { parts: [{ type: 'text', text: 'continue' }] },
+    );
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(promptAsync).toHaveBeenCalledTimes(1);
+  });
+
   test('same user message observed by two hooks rearms only one new epoch', async () => {
     const promptAsync = mock(async () => ({}));
     const sessionClient = {

+ 6 - 6
src/hooks/task-session-manager/index.ts

@@ -207,18 +207,18 @@ export function createTaskSessionManagerHook(
       const parts = Array.isArray(outputRecord?.parts)
         ? outputRecord.parts
         : inputMessage?.parts;
-      // Stable identity from chat.message (input.messageID or output.message.id).
-      // Required for process-global idempotent rearm across hook instances.
-      const messageID =
+      // Safe identity order (Oracle): input.messageID → output.message.id →
+      // same-process output.message object → fail closed.
+      const messageIdentity: string | object | undefined =
         typeof inputMessage?.messageID === 'string' &&
         inputMessage.messageID.length > 0
           ? inputMessage.messageID
           : typeof outputMessage?.id === 'string' && outputMessage.id.length > 0
             ? outputMessage.id
-            : undefined;
+            : outputMessage;
       if (
         !sessionID ||
-        !messageID ||
+        messageIdentity === undefined ||
         (typeof outputMessage?.role === 'string' &&
           outputMessage.role !== 'user') ||
         !options.shouldManageSession(sessionID) ||
@@ -235,7 +235,7 @@ export function createTaskSessionManagerHook(
       ) {
         return;
       }
-      continuationTokens.rearmForUserMessage(sessionID, messageID);
+      continuationTokens.rearmForUserMessage(sessionID, messageIdentity);
     },
 
     'tool.execute.before': (

+ 8 - 1
src/index.ts

@@ -1088,9 +1088,16 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
     // Track which agent each session uses (needed for serve-mode prompt
     // injection)
     'chat.message': async (
-      input: { sessionID: string; agent?: string; parts?: unknown[] },
+      input: {
+        sessionID: string;
+        agent?: string;
+        parts?: unknown[];
+        /** OpenCode chat.message message identity when present. */
+        messageID?: string;
+      },
       output?: {
         message?: {
+          id?: string;
           agent?: string;
           role?: string;
           sessionID?: string;