diff --git a/src/screens/REPL.tsx b/src/screens/REPL.tsx index 00a7e3d6b2..1d7a8a8c8d 100644 --- a/src/screens/REPL.tsx +++ b/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) diff --git a/src/utils/QueryGuard.test.ts b/src/utils/QueryGuard.test.ts index e3efe7875c..9283685765 100644 --- a/src/utils/QueryGuard.test.ts +++ b/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() }) }) diff --git a/src/utils/QueryGuard.ts b/src/utils/QueryGuard.ts index a953ccdcf6..dba2cf40fe 100644 --- a/src/utils/QueryGuard.ts +++ b/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 | 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) }