Merge pull request #437 from alvinunreal/fix/async-lifecycle-hardening

src/hooks/foreground-fallback/index.test.ts

@@ -1,34 +1,45 @@
 import { beforeEach, describe, expect, mock, test } from 'bun:test';
 import { ForegroundFallbackManager, isRateLimitError } from './index';
 
+type ForegroundFallbackClient = ConstructorParameters<
+  typeof ForegroundFallbackManager
+>[0];
+
 // ---------------------------------------------------------------------------
 // Helpers
 // ---------------------------------------------------------------------------
 
 function createMockClient(overrides?: {
   promptAsyncImpl?: (args: unknown) => Promise<unknown>;
+  abortImpl?: () => Promise<unknown>;
+  includePromptAsync?: boolean;
   messagesData?: Array<{ info: { role: string }; parts: unknown[] }>;
 }) {
   const promptAsync = mock(async (args: unknown) => {
     if (overrides?.promptAsyncImpl) return overrides.promptAsyncImpl(args);
     return {};
   });
-  const abort = mock(async () => ({}));
+  const abort = mock(async () => {
+    if (overrides?.abortImpl) return overrides.abortImpl();
+    return {};
+  });
   const messages = mock(async () => ({
     data: overrides?.messagesData ?? [
       { info: { role: 'user' }, parts: [{ type: 'text', text: 'hello' }] },
     ],
   }));
+  const session: Record<string, unknown> = {
+    abort,
+    messages,
+  };
+  if (overrides?.includePromptAsync !== false) {
+    session.promptAsync = promptAsync;
+  }
 
   return {
     client: {
-      session: {
-        abort,
-        messages,
-        // promptAsync is cast at runtime — expose via the session object
-        promptAsync,
-      },
-    } as unknown as Parameters<typeof ForegroundFallbackManager>[0],
+      session,
+    } as unknown as ForegroundFallbackClient,
     mocks: { promptAsync, abort, messages },
   };
 }
@@ -183,6 +194,42 @@ describe('ForegroundFallbackManager session.error', () => {
 
     expect(mocks.promptAsync).not.toHaveBeenCalled();
   });
+
+  test('does not abort when promptAsync is unavailable', async () => {
+    const { client, mocks } = createMockClient({ includePromptAsync: false });
+    const mgr = new ForegroundFallbackManager(client, makeChains(), true);
+
+    await mgr.handleEvent({
+      type: 'session.error',
+      properties: {
+        sessionID: 'sess-no-prompt-async',
+        error: { message: 'Rate limit exceeded' },
+      },
+    });
+
+    expect(mocks.abort).not.toHaveBeenCalled();
+    expect(mocks.promptAsync).not.toHaveBeenCalled();
+  });
+
+  test('continues fallback when abort rejects', async () => {
+    const { client, mocks } = createMockClient({
+      abortImpl: async () => {
+        throw new Error('abort failed');
+      },
+    });
+    const mgr = new ForegroundFallbackManager(client, makeChains(), true);
+
+    await mgr.handleEvent({
+      type: 'session.error',
+      properties: {
+        sessionID: 'sess-abort-rejects',
+        error: { message: 'Rate limit exceeded' },
+      },
+    });
+
+    expect(mocks.abort).toHaveBeenCalledTimes(1);
+    expect(mocks.promptAsync).toHaveBeenCalledTimes(1);
+  });
 });
 
 // ---------------------------------------------------------------------------

src/hooks/foreground-fallback/index.ts

@@ -15,6 +15,7 @@
  */
 
 import type { PluginInput } from '@opencode-ai/plugin';
+import { abortSessionWithTimeout } from '../../utils/session';
 import { log } from '../../utils/logger';
 
 type OpencodeClient = PluginInput['client'];
@@ -68,6 +69,7 @@ function parseModel(
 
 /** Prevent re-triggering within this window for the same session. */
 const DEDUP_WINDOW_MS = 5_000;
+const REPROMPT_DELAY_MS = 500;
 
 // ---------------------------------------------------------------------------
 // Manager
@@ -277,23 +279,13 @@ export class ForegroundFallbackManager {
         return;
       }
 
-      // Abort the currently rate-limited prompt so the session becomes idle.
-      try {
-        await this.client.session.abort({ path: { id: sessionID } });
-      } catch {
-        // Session may already be idle; safe to ignore.
-      }
-
-      // Give the server a moment to finalise the abort before re-prompting.
-      await new Promise((r) => setTimeout(r, 500));
-
       // promptAsync queues the prompt and returns immediately — this avoids
       // blocking the event handler while waiting for a full LLM response.
       // Cast required: promptAsync is not in the plugin TypeScript types for
       // oh-my-opencode-slim but IS present on the real OpenCode client at
       // runtime (verified by opencode-rate-limit-fallback reference impl).
       const sessionClient = this.client.session as unknown as {
-        promptAsync: (args: {
+        promptAsync?: (args: {
           path: { id: string };
           body: {
             parts: unknown[];
@@ -301,6 +293,25 @@ export class ForegroundFallbackManager {
           };
         }) => Promise<unknown>;
       };
+      if (typeof sessionClient.promptAsync !== 'function') {
+        log('[foreground-fallback] promptAsync unavailable', { sessionID });
+        return;
+      }
+
+      // Abort the currently rate-limited prompt so the session becomes idle.
+      try {
+        await abortSessionWithTimeout(this.client, sessionID);
+      } catch (error) {
+        // Session may already be idle or abort may be slow; keep fallback best-effort.
+        log('[foreground-fallback] abort did not complete cleanly', {
+          sessionID,
+          error: error instanceof Error ? error.message : String(error),
+        });
+      }
+
+      // Give the server a moment to finalise the abort before re-prompting.
+      await new Promise((r) => setTimeout(r, REPROMPT_DELAY_MS));
+
       await sessionClient.promptAsync({
         path: { id: sessionID },
         body: { parts: lastUser.parts, model: ref },

src/hooks/todo-continuation/index.test.ts

@@ -187,6 +187,23 @@ describe('createTodoContinuationHook', () => {
       );
     });
 
+    test('skips hygiene reminder when todo state lookup times out', async () => {
+      const ctx = createMockContext();
+      ctx.client.session.todo = mock(() => new Promise(() => {}));
+      const hook = createTodoContinuationHook(ctx);
+      const output = userMessages('primera request', 'main1', 'orchestrator');
+
+      await hook.handleMessagesTransform(output);
+      await hook.handleToolExecuteAfter({
+        tool: 'todowrite',
+        sessionID: 'main1',
+      });
+      await hook.handleToolExecuteAfter({ tool: 'read', sessionID: 'main1' });
+      await hook.handleMessagesTransform(output);
+
+      expect(allMessageText(output)).not.toContain(TODO_HYGIENE_REMINDER);
+    });
+
     test('compaction-like transform does not consume pending reminder', async () => {
       const ctx = createMockContext({
         todoResult: {

src/hooks/todo-continuation/index.ts

@@ -4,11 +4,13 @@ import {
   createInternalAgentTextPart,
   log,
   SLIM_INTERNAL_INITIATOR_MARKER,
+  withTimeout,
 } from '../../utils';
 import { createTodoHygiene } from './todo-hygiene';
 
 const HOOK_NAME = 'todo-continuation';
 const COMMAND_NAME = 'auto-continue';
+const TODO_STATE_TIMEOUT_MS = 500;
 
 const CONTINUATION_PROMPT =
   '[Auto-continue: enabled - there are incomplete todos remaining. Continue with the next uncompleted item. Press Esc to cancel. If you need user input or review for the next item, ask instead of proceeding.]';
@@ -212,12 +214,20 @@ export function createTodoContinuationHook(
     notificationBusyUntilBySession: new Map<string, number>(),
   };
 
+  async function fetchTodos(sessionID: string): Promise<TodoItem[]> {
+    const result = await withTimeout(
+      ctx.client.session.todo({
+        path: { id: sessionID },
+      }),
+      TODO_STATE_TIMEOUT_MS,
+      `Todo state lookup timed out after ${TODO_STATE_TIMEOUT_MS}ms`,
+    );
+    return result.data as TodoItem[];
+  }
+
   const hygiene = createTodoHygiene({
     getTodoState: async (sessionID) => {
-      const result = await ctx.client.session.todo({
-        path: { id: sessionID },
-      });
-      const todos = result.data as TodoItem[];
+      const todos = await fetchTodos(sessionID);
       const openTodos = todos.filter(
         (todo) => !TERMINAL_TODO_STATUSES.includes(todo.status),
       );
@@ -500,10 +510,7 @@ export function createTodoContinuationHook(
       // todos exist, automatically enable auto-continue.
       if (autoEnable && !state.enabled) {
         try {
-          const todosResult = await ctx.client.session.todo({
-            path: { id: sessionID },
-          });
-          const todos = todosResult.data as TodoItem[];
+          const todos = await fetchTodos(sessionID);
           const incompleteCount = todos.filter(
             (t) => !TERMINAL_TODO_STATUSES.includes(t.status),
           ).length;
@@ -544,10 +551,7 @@ export function createTodoContinuationHook(
       let hasIncompleteTodos = false;
       let incompleteCount = 0;
       try {
-        const todosResult = await ctx.client.session.todo({
-          path: { id: sessionID },
-        });
-        const todos = todosResult.data as TodoItem[];
+        const todos = await fetchTodos(sessionID);
         incompleteCount = todos.filter(
           (t) => !TERMINAL_TODO_STATUSES.includes(t.status),
         ).length;
@@ -838,10 +842,7 @@ export function createTodoContinuationHook(
     // Check for incomplete todos to decide on immediate continuation
     let hasIncompleteTodos = false;
     try {
-      const todosResult = await ctx.client.session.todo({
-        path: { id: input.sessionID },
-      });
-      const todos = todosResult.data as TodoItem[];
+      const todos = await fetchTodos(input.sessionID);
       hasIncompleteTodos = todos.some(
         (t) => !TERMINAL_TODO_STATUSES.includes(t.status),
       );

src/index.ts

@@ -1070,52 +1070,84 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
     // Post-tool hooks: retry guidance for delegation errors + file-tool
     // nudge
     'tool.execute.after': async (input, output) => {
-      await delegateTaskRetryHook['tool.execute.after'](
-        input as { tool: string },
-        output as { output: unknown },
+      const meta = input as {
+        tool?: string;
+        sessionID?: string;
+        callID?: string;
+      };
+      const runPostToolHook = async (
+        name: string,
+        fn: () => Promise<void>,
+      ): Promise<void> => {
+        try {
+          await fn();
+        } catch (error) {
+          log('[plugin] post-tool hook failed open', {
+            hook: name,
+            tool: meta.tool,
+            sessionID: meta.sessionID,
+            callID: meta.callID,
+            error: error instanceof Error ? error.message : String(error),
+          });
+        }
+      };
+
+      await runPostToolHook('delegate-task-retry', () =>
+        delegateTaskRetryHook['tool.execute.after'](
+          input as { tool: string },
+          output as { output: unknown },
+        ),
       );
 
-      await jsonErrorRecoveryHook['tool.execute.after'](
-        input as {
-          tool: string;
-          sessionID: string;
-          callID: string;
-        },
-        output as {
-          title: string;
-          output: unknown;
-          metadata: unknown;
-        },
+      await runPostToolHook('json-error-recovery', () =>
+        jsonErrorRecoveryHook['tool.execute.after'](
+          input as {
+            tool: string;
+            sessionID: string;
+            callID: string;
+          },
+          output as {
+            title: string;
+            output: unknown;
+            metadata: unknown;
+          },
+        ),
       );
 
-      await todoContinuationHook.handleToolExecuteAfter(
-        input as {
-          tool: string;
-          sessionID?: string;
-        },
-        output as { output?: unknown },
+      await runPostToolHook('todo-continuation', () =>
+        todoContinuationHook.handleToolExecuteAfter(
+          input as {
+            tool: string;
+            sessionID?: string;
+          },
+          output as { output?: unknown },
+        ),
       );
 
-      await postFileToolNudgeHook['tool.execute.after'](
-        input as {
-          tool: string;
-          sessionID?: string;
-          callID?: string;
-        },
-        output as {
-          title: string;
-          output: string;
-          metadata: Record<string, unknown>;
-        },
+      await runPostToolHook('post-file-tool-nudge', () =>
+        postFileToolNudgeHook['tool.execute.after'](
+          input as {
+            tool: string;
+            sessionID?: string;
+            callID?: string;
+          },
+          output as {
+            title: string;
+            output: string;
+            metadata: Record<string, unknown>;
+          },
+        ),
       );
 
-      await taskSessionManagerHook['tool.execute.after'](
-        input as {
-          tool: string;
-          sessionID?: string;
-          callID?: string;
-        },
-        output as { output: unknown },
+      await runPostToolHook('task-session-manager', () =>
+        taskSessionManagerHook['tool.execute.after'](
+          input as {
+            tool: string;
+            sessionID?: string;
+            callID?: string;
+          },
+          output as { output: unknown },
+        ),
       );
 
       if (input.tool.toLowerCase() === 'task') {

src/utils/session.test.ts

@@ -0,0 +1,79 @@
+import { describe, expect, mock, test } from 'bun:test';
+import {
+  abortSessionWithTimeout,
+  OperationTimeoutError,
+  promptWithTimeout,
+  withTimeout,
+} from './session';
+
+function never<T>(): Promise<T> {
+  return new Promise<T>(() => {});
+}
+
+describe('session utilities', () => {
+  test('withTimeout resolves without waiting for the timeout', async () => {
+    const result = await withTimeout(Promise.resolve('ok'), 50, 'too slow');
+
+    expect(result).toBe('ok');
+  });
+
+  test('withTimeout rejects with OperationTimeoutError when operation hangs', async () => {
+    await expect(withTimeout(never(), 5, 'too slow')).rejects.toThrow(
+      OperationTimeoutError,
+    );
+  });
+
+  test('promptWithTimeout aborts a timed-out prompt before rejecting', async () => {
+    const abort = mock(async () => ({}));
+    const prompt = mock(() => never());
+    const client = {
+      session: {
+        abort,
+        prompt,
+      },
+    } as any;
+
+    await expect(
+      promptWithTimeout(
+        client,
+        { path: { id: 's1' }, body: { parts: [] } },
+        5,
+      ),
+    ).rejects.toThrow('Prompt timed out after 5ms');
+
+    expect(abort).toHaveBeenCalledWith({ path: { id: 's1' } });
+  });
+
+  test('promptWithTimeout preserves timeout error when abort fails', async () => {
+    const abort = mock(async () => {
+      throw new Error('abort failed');
+    });
+    const prompt = mock(() => never());
+    const client = {
+      session: {
+        abort,
+        prompt,
+      },
+    } as any;
+
+    await expect(
+      promptWithTimeout(
+        client,
+        { path: { id: 's1' }, body: { parts: [] } },
+        5,
+      ),
+    ).rejects.toThrow('Prompt timed out after 5ms');
+  });
+
+  test('abortSessionWithTimeout rejects if abort hangs', async () => {
+    const client = {
+      session: {
+        abort: mock(() => never()),
+      },
+    } as any;
+
+    await expect(abortSessionWithTimeout(client, 's1', 5)).rejects.toThrow(
+      'Session abort timed out after 5ms',
+    );
+  });
+});

src/utils/session.ts

@@ -6,6 +6,49 @@ import type { PluginInput } from '@opencode-ai/plugin';
 
 type OpencodeClient = PluginInput['client'];
 
+export const SESSION_ABORT_TIMEOUT_MS = 1_000;
+
+export class OperationTimeoutError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = "OperationTimeoutError";
+  }
+}
+
+export async function withTimeout<T>(
+  operation: Promise<T>,
+  timeoutMs: number,
+  message: string
+): Promise<T> {
+  if (timeoutMs <= 0) return operation;
+
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  try {
+    return await Promise.race([
+      operation,
+      new Promise<never>((_, reject) => {
+        timer = setTimeout(() => {
+          reject(new OperationTimeoutError(message));
+        }, timeoutMs);
+      }),
+    ]);
+  } finally {
+    clearTimeout(timer);
+  }
+}
+
+export async function abortSessionWithTimeout(
+  client: OpencodeClient,
+  sessionId: string,
+  timeoutMs = SESSION_ABORT_TIMEOUT_MS
+): Promise<void> {
+  await withTimeout(
+    client.session.abort({ path: { id: sessionId } }),
+    timeoutMs,
+    `Session abort timed out after ${timeoutMs}ms`
+  );
+}
+
 /**
  * Extract the short model label from a "provider/model" string.
  * E.g. "openai/gpt-5.4-mini" → "gpt-5.4-mini"
@@ -72,11 +115,23 @@ export async function promptWithTimeout(
       promptPromise,
       new Promise<never>((_, reject) => {
         timer = setTimeout(() => {
-          client.session.abort({ path: { id: sessionId } }).catch(() => {});
-          reject(new Error(`Prompt timed out after ${timeoutMs}ms`));
+          reject(
+            new OperationTimeoutError(
+              `Prompt timed out after ${timeoutMs}ms`,
+            ),
+          );
         }, timeoutMs);
       }),
     ]);
+  } catch (error) {
+    if (error instanceof OperationTimeoutError) {
+      try {
+        await abortSessionWithTimeout(client, sessionId);
+      } catch {
+        // Best-effort cleanup: preserve the original prompt timeout error.
+      }
+    }
+    throw error;
   } finally {
     clearTimeout(timer);
   }