[codex] fix(query): abort active work on QueryGuard timeout

src/screens/REPL.tsx

@@ -1726,6 +1726,28 @@ export function REPL({
   const mrOnBeforeQuery = useCallback(async (_input: string, _allMessages: MessageType[], _newMessageCount: number) => true, []);
   const mrOnTurnComplete = useCallback(async (_allMessages: MessageType[], _aborted: boolean) => { }, []);
   const mrRender = useCallback(() => null, []);
+  const abortTimedOutQuery = useCallback(() => {
+    const activeAbortController = abortControllerRef.current;
+    if (activeAbortController && !activeAbortController.signal.aborted) {
+      activeAbortController.abort('query-timeout');
+    }
+    if (feature('TOKEN_BUDGET')) {
+      snapshotOutputTokensForTurn(null);
+    }
+
+    // QueryGuard calls this before forceEnd(); defer UI cleanup until after
+    // the guard has released so the normal stale-generation finally path skips.
+    queueMicrotask(() => {
+      resetLoadingState();
+      setAbortController(null);
+      void mrOnTurnComplete(messagesRef.current, true);
+    });
+  }, [mrOnTurnComplete, resetLoadingState]);
+
+  useEffect(() => {
+    return queryGuard.setTimeoutHandler(abortTimedOutQuery);
+  }, [abortTimedOutQuery, queryGuard]);
+
   const showSpinner = (!toolJSX || toolJSX.showSpinner === true) && toolUseConfirmQueue.length === 0 && promptQueue.length === 0 && (
     // Show spinner during input processing, API call, while teammates are running,
     // or while pending task notifications are queued (prevents spinner bounce between consecutive notifications)

src/utils/QueryGuard.test.ts

@@ -1,7 +1,12 @@
-import { describe, test, expect, vi } from 'vitest'
+import { afterEach, describe, test, expect, vi } from 'vitest'
 import { QueryGuard } from './QueryGuard.js'
 
 describe('QueryGuard', () => {
+  afterEach(() => {
+    vi.useRealTimers()
+    vi.restoreAllMocks()
+  })
+
   test('starts idle', () => {
     const guard = new QueryGuard()
     expect(guard.isActive).toBe(false)
@@ -52,8 +57,50 @@ describe('QueryGuard', () => {
     // At timeout
     vi.advanceTimersByTime(1)
     expect(guard.isActive).toBe(false)
+  })
 
-    vi.useRealTimers()
+  test('timeout notifies owner with the timed-out generation', () => {
+    vi.useFakeTimers()
+    const guard = new QueryGuard()
+    const onTimeout = vi.fn()
+    guard.setTimeoutHandler(onTimeout)
+
+    const gen = guard.tryStart()!
+    vi.advanceTimersByTime(5 * 60 * 1000)
+
+    expect(onTimeout).toHaveBeenCalledTimes(1)
+    expect(onTimeout).toHaveBeenCalledWith(gen)
+    expect(guard.isActive).toBe(false)
+  })
+
+  test('timeout handler cleanup prevents stale notification', () => {
+    vi.useFakeTimers()
+    const guard = new QueryGuard()
+    const onTimeout = vi.fn()
+    const cleanup = guard.setTimeoutHandler(onTimeout)
+    cleanup()
+
+    guard.tryStart()
+    vi.advanceTimersByTime(5 * 60 * 1000)
+
+    expect(onTimeout).not.toHaveBeenCalled()
+    expect(guard.isActive).toBe(false)
+  })
+
+  test('timeout handler errors do not escape the watchdog callback', () => {
+    vi.useFakeTimers()
+    const guard = new QueryGuard()
+    const handlerError = new Error('handler failed')
+    const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {})
+    guard.setTimeoutHandler(() => {
+      throw handlerError
+    })
+
+    guard.tryStart()
+
+    expect(() => vi.advanceTimersByTime(5 * 60 * 1000)).not.toThrow()
+    expect(guard.isActive).toBe(false)
+    expect(consoleError).toHaveBeenCalledWith('[QueryGuard] Timeout handler failed', handlerError)
   })
 
   test('timeout is cleared when end() is called normally', () => {
@@ -72,6 +119,5 @@ describe('QueryGuard', () => {
     expect(guard.isActive).toBe(true)
 
     guard.forceEnd()
-    vi.useRealTimers()
   })
 })

src/utils/QueryGuard.ts

@@ -31,12 +31,14 @@
 import { createSignal } from './signal.js'
 
 const QUERY_TIMEOUT_MS = 5 * 60 * 1000 // 5 minutes
+type QueryTimeoutHandler = (generation: number) => void
 
 export class QueryGuard {
   private _status: 'idle' | 'dispatching' | 'running' = 'idle'
   private _generation = 0
   private _changed = createSignal()
   private _timeoutId: ReturnType<typeof setTimeout> | null = null
+  private _timeoutHandler: QueryTimeoutHandler | null = null
 
   /**
    * Reserve the guard for queue processing. Transitions idle → dispatching.
@@ -114,6 +116,20 @@ export class QueryGuard {
     return this._generation
   }
 
+  /**
+   * Register a single owner callback for watchdog timeouts. The callback runs
+   * before forceEnd(), so callers can abort in-flight work while the timed-out
+   * generation is still current.
+   */
+  setTimeoutHandler(handler: QueryTimeoutHandler | null): () => void {
+    this._timeoutHandler = handler
+    return () => {
+      if (this._timeoutHandler === handler) {
+        this._timeoutHandler = null
+      }
+    }
+  }
+
   // --
   // useSyncExternalStore interface
 
@@ -138,7 +154,13 @@ export class QueryGuard {
     this._timeoutId = setTimeout(() => {
       if (this._status === 'running') {
         console.error(`[QueryGuard] Query timeout after ${QUERY_TIMEOUT_MS}ms — force-ending to prevent infinite spinner`)
-        this.forceEnd()
+        try {
+          this._timeoutHandler?.(this._generation)
+        } catch (error) {
+          console.error('[QueryGuard] Timeout handler failed', error)
+        } finally {
+          this.forceEnd()
+        }
       }
     }, QUERY_TIMEOUT_MS)
   }