diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index ae969640367..78542fb4c9c 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -82,6 +82,7 @@ import { collectToolCallIdsFromHistory, normalizeModelToolCallIds, } from './toolCallIdUtils.js'; +import { createStreamIdleWatchdog, linkAbortSignal, type StreamIdleWatchdog, type InvalidStreamErrorType } from './streamIdleWatchdog.js'; const debugLogger = createDebugLogger('QWEN_CODE_CHAT'); @@ -993,9 +994,9 @@ function stripThoughtPartsFromContent(content: Content): Content | null { * which should trigger a retry. */ export class InvalidStreamError extends Error { - readonly type: 'NO_FINISH_REASON' | 'NO_RESPONSE_TEXT'; + readonly type: InvalidStreamErrorType; - constructor(message: string, type: 'NO_FINISH_REASON' | 'NO_RESPONSE_TEXT') { + constructor(message: string, type: InvalidStreamErrorType) { super(message); this.name = 'InvalidStreamError'; this.type = type; @@ -2605,12 +2606,15 @@ export class GeminiChat { params: SendMessageParameters, prompt_id: string, ): Promise> { + const streamAbortController = new AbortController(); + const cleanupAbortLink = linkAbortSignal(params.config?.abortSignal, streamAbortController); + const streamWatchdog = createStreamIdleWatchdog(model, streamAbortController); const apiCall = () => this.config.getContentGenerator().generateContentStream( { model, contents: requestContents, - config: { ...this.generationConfig, ...params.config }, + config: { ...this.generationConfig, ...params.config, abortSignal: streamAbortController.signal }, }, prompt_id, ); @@ -2660,7 +2664,13 @@ export class GeminiChat { }, }); - return this.processStreamResponse(model, streamResponse); + try { + return this.processStreamResponse(model, streamResponse, streamWatchdog, cleanupAbortLink); + } catch (error) { + streamWatchdog?.cleanup(); + cleanupAbortLink?.(); + throw error; + } } /** @@ -2964,6 +2974,8 @@ export class GeminiChat { private async *processStreamResponse( model: string, streamResponse: AsyncGenerator, + _streamWatchdog?: StreamIdleWatchdog, + _cleanupAbortLink?: () => void, ): AsyncGenerator { // Collect ALL parts from the model response (including thoughts for recording) const allModelParts: Part[] = []; diff --git a/packages/core/src/core/streamIdleWatchdog.ts b/packages/core/src/core/streamIdleWatchdog.ts new file mode 100644 index 00000000000..fec59607c91 --- /dev/null +++ b/packages/core/src/core/streamIdleWatchdog.ts @@ -0,0 +1,79 @@ +import { createDebugLogger } from '../utils/debugLogger.js'; +import { InvalidStreamError } from './geminiChat.js'; + +const debugLogger = createDebugLogger('QWEN_CODE_WATCHDOG'); + +export const DEFAULT_STREAM_IDLE_TIMEOUT_MS = 90_000; +const STREAM_IDLE_WARNING_FRACTION = 0.5; + +export type InvalidStreamErrorType = + | 'NO_FINISH_REASON' + | 'NO_RESPONSE_TEXT' + | 'STREAM_IDLE_TIMEOUT'; + +export interface StreamIdleWatchdog { + next(nextPromise: Promise>): Promise>; + cleanup(): void; +} + +export function isStreamWatchdogDisabled(): boolean { + const value = process.env['QWEN_CODE_DISABLE_STREAM_WATCHDOG']; + return value === '1' || value === 'true'; +} + +export function getStreamIdleTimeoutMs(): number | undefined { + if (isStreamWatchdogDisabled()) return undefined; + const value = process.env['QWEN_CODE_STREAM_IDLE_TIMEOUT_MS']; + if (!value) return DEFAULT_STREAM_IDLE_TIMEOUT_MS; + const parsed = Number(value); + if (!Number.isFinite(parsed) || parsed <= 0) { + debugLogger.warn('Ignoring invalid STREAM_IDLE_TIMEOUT_MS: ' + value); + return DEFAULT_STREAM_IDLE_TIMEOUT_MS; + } + return parsed; +} + +export function linkAbortSignal( + signal: AbortSignal | undefined, + controller: AbortController, +): () => void { + if (!signal) return () => {}; + if (signal.aborted) { controller.abort(signal.reason); return () => {}; } + const onAbort = () => controller.abort(signal.reason); + signal.addEventListener('abort', onAbort, { once: true }); + return () => signal.removeEventListener('abort', onAbort); +} + +export function createStreamIdleWatchdog( + model: string, + abortController: AbortController, +): StreamIdleWatchdog | undefined { + const timeoutMs = getStreamIdleTimeoutMs(); + if (timeoutMs === undefined) return undefined; + const warningMs = Math.max(1, Math.floor(timeoutMs * STREAM_IDLE_WARNING_FRACTION)); + let timeoutId: ReturnType | undefined; + let warningId: ReturnType | undefined; + const clearTimers = () => { + if (timeoutId !== undefined) { clearTimeout(timeoutId); timeoutId = undefined; } + if (warningId !== undefined) { clearTimeout(warningId); warningId = undefined; } + }; + return { + next(nextPromise: Promise>): Promise> { + const timeoutPromise = new Promise((_, reject) => { + warningId = setTimeout(() => { + debugLogger.warn('Stream idle for ' + warningMs + 'ms from ' + model); + }, warningMs); + timeoutId = setTimeout(() => { + clearTimers(); + abortController.abort(); + reject(new InvalidStreamError( + 'Stream idle timeout after ' + timeoutMs + 'ms', + 'STREAM_IDLE_TIMEOUT', + )); + }, timeoutMs); + }); + return Promise.race([nextPromise, timeoutPromise]).finally(clearTimers); + }, + cleanup() { clearTimers(); }, + }; +}