|
| 1 | +import type { StreamFn } from '@earendil-works/pi-agent-core'; |
| 2 | +import { |
| 3 | + createAssistantMessageEventStream, |
| 4 | + streamSimple, |
| 5 | + type AssistantMessage, |
| 6 | + type AssistantMessageEvent, |
| 7 | + type AssistantMessageEventStream, |
| 8 | + type Model, |
| 9 | + type Api, |
| 10 | +} from '@earendil-works/pi-ai'; |
| 11 | +import { |
| 12 | + LLM_FIRST_EVENT_TIMEOUT_MS, |
| 13 | + LLM_FIRST_EVENT_TIMEOUT_ENV, |
| 14 | + LLM_STREAM_IDLE_TIMEOUT_ENV, |
| 15 | + LLM_STREAM_IDLE_TIMEOUT_MS, |
| 16 | +} from './defaults.js'; |
| 17 | + |
| 18 | +/** |
| 19 | + * Per-attempt watchdog for provider streams. |
| 20 | + * |
| 21 | + * `firstEventTimeoutMs` bounds the wait for the first stream event. The |
| 22 | + * built-in providers emit `start` only after the HTTP response headers |
| 23 | + * arrive, so this is effectively a first-byte timeout: a request the provider |
| 24 | + * accepted but never answers fails here and is retried instead of waiting for |
| 25 | + * the 20 min outer request bound (or a custom fetch's own transport default). |
| 26 | + * |
| 27 | + * `idleTimeoutMs` bounds the gap between consecutive events after the first |
| 28 | + * one, so a stream that stalls mid-response is also failed. |
| 29 | + * |
| 30 | + * `0` disables the corresponding timer. |
| 31 | + */ |
| 32 | +export interface LLMStreamTimeouts { |
| 33 | + readonly firstEventTimeoutMs: number; |
| 34 | + readonly idleTimeoutMs: number; |
| 35 | +} |
| 36 | + |
| 37 | +export interface LLMStreamTimeoutConfig { |
| 38 | + readonly firstEventTimeoutMs?: number; |
| 39 | + readonly streamIdleTimeoutMs?: number; |
| 40 | +} |
| 41 | + |
| 42 | +/** Marker included in every synthesized timeout message; matches the shared timeout classifier. */ |
| 43 | +export const LLM_STREAM_TIMEOUT_MESSAGE_PREFIX = 'LLM request timed out'; |
| 44 | + |
| 45 | +/** |
| 46 | + * Resolve effective timeouts. Precedence: explicit host config, then the |
| 47 | + * environment override, then the built-in default. Invalid values (negative, |
| 48 | + * non-integer, non-numeric) fall through to the next source. |
| 49 | + */ |
| 50 | +export function resolveLLMStreamTimeouts( |
| 51 | + config: LLMStreamTimeoutConfig = {}, |
| 52 | + env: Readonly<Record<string, string | undefined>> = readProcessEnv(), |
| 53 | +): LLMStreamTimeouts { |
| 54 | + return { |
| 55 | + firstEventTimeoutMs: |
| 56 | + validTimeout(config.firstEventTimeoutMs) ?? |
| 57 | + parseTimeout(env[LLM_FIRST_EVENT_TIMEOUT_ENV]) ?? |
| 58 | + LLM_FIRST_EVENT_TIMEOUT_MS, |
| 59 | + idleTimeoutMs: |
| 60 | + validTimeout(config.streamIdleTimeoutMs) ?? |
| 61 | + parseTimeout(env[LLM_STREAM_IDLE_TIMEOUT_ENV]) ?? |
| 62 | + LLM_STREAM_IDLE_TIMEOUT_MS, |
| 63 | + }; |
| 64 | +} |
| 65 | + |
| 66 | +/** |
| 67 | + * Wrap a StreamFn so each invocation (each physical request, including every |
| 68 | + * retry attempt made by `withLLMRetry`) is guarded by the watchdog. |
| 69 | + * |
| 70 | + * On timeout the inner request is aborted through a linked signal and the |
| 71 | + * returned stream ends with a terminal `error` event whose stop reason is |
| 72 | + * `error` (not `aborted`), so the retry layer treats it as a retryable |
| 73 | + * transport timeout rather than a user cancellation. A caller abort is passed |
| 74 | + * through unchanged and never reported as a timeout. |
| 75 | + */ |
| 76 | +export function withLLMStreamTimeouts( |
| 77 | + inner: StreamFn | undefined, |
| 78 | + timeouts: LLMStreamTimeouts, |
| 79 | +): StreamFn { |
| 80 | + const base = inner ?? streamSimple; |
| 81 | + if (timeouts.firstEventTimeoutMs <= 0 && timeouts.idleTimeoutMs <= 0) return base; |
| 82 | + return ((model, context, options) => { |
| 83 | + const callerSignal = options?.signal; |
| 84 | + const controller = new AbortController(); |
| 85 | + const out = createAssistantMessageEventStream(); |
| 86 | + let finished = false; |
| 87 | + let timer: ReturnType<typeof setTimeout> | undefined; |
| 88 | + let lastPartial: AssistantMessage | undefined; |
| 89 | + // Resolves when the wrapper is finished for any reason, so the read loop |
| 90 | + // below stops even if the inner stream ignores the abort and never yields |
| 91 | + // or settles again. |
| 92 | + let signalStopped!: () => void; |
| 93 | + const stopped = new Promise<typeof STOPPED>((resolve) => { |
| 94 | + signalStopped = () => resolve(STOPPED); |
| 95 | + }); |
| 96 | + |
| 97 | + const clearTimer = () => { |
| 98 | + if (timer !== undefined) clearTimeout(timer); |
| 99 | + timer = undefined; |
| 100 | + }; |
| 101 | + const onCallerAbort = () => { |
| 102 | + clearTimer(); |
| 103 | + controller.abort(callerSignal?.reason); |
| 104 | + }; |
| 105 | + const finish = () => { |
| 106 | + finished = true; |
| 107 | + clearTimer(); |
| 108 | + callerSignal?.removeEventListener('abort', onCallerAbort); |
| 109 | + signalStopped(); |
| 110 | + }; |
| 111 | + const arm = (ms: number, phase: 'first' | 'idle') => { |
| 112 | + clearTimer(); |
| 113 | + if (ms <= 0 || finished) return; |
| 114 | + timer = setTimeout(() => { |
| 115 | + if (finished || callerSignal?.aborted) return; |
| 116 | + finish(); |
| 117 | + const message = timeoutMessage(phase, ms); |
| 118 | + controller.abort(new LLMStreamTimeoutError(message)); |
| 119 | + const error: AssistantMessage = { |
| 120 | + ...snapshot(lastPartial, model), |
| 121 | + stopReason: 'error', |
| 122 | + errorMessage: message, |
| 123 | + }; |
| 124 | + out.push({ type: 'error', reason: 'error', error }); |
| 125 | + out.end(); |
| 126 | + }, ms); |
| 127 | + // Never keep the process alive only for the watchdog. |
| 128 | + (timer as { unref?: () => void }).unref?.(); |
| 129 | + }; |
| 130 | + |
| 131 | + if (callerSignal?.aborted) { |
| 132 | + controller.abort(callerSignal.reason); |
| 133 | + } else { |
| 134 | + callerSignal?.addEventListener('abort', onCallerAbort, { once: true }); |
| 135 | + arm(timeouts.firstEventTimeoutMs, 'first'); |
| 136 | + } |
| 137 | + |
| 138 | + void (async () => { |
| 139 | + try { |
| 140 | + const opened = await Promise.race([ |
| 141 | + Promise.resolve(base(model, context, { ...(options ?? {}), signal: controller.signal })), |
| 142 | + stopped, |
| 143 | + ]); |
| 144 | + if (opened === STOPPED) return; |
| 145 | + const stream: AssistantMessageEventStream = opened; |
| 146 | + const iterator = stream[Symbol.asyncIterator](); |
| 147 | + for (;;) { |
| 148 | + const next = await Promise.race([iterator.next(), stopped]); |
| 149 | + if (next === STOPPED || finished) { |
| 150 | + releaseIterator(iterator); |
| 151 | + return; |
| 152 | + } |
| 153 | + if (next.done) break; |
| 154 | + const event = next.value; |
| 155 | + const partial = partialOf(event); |
| 156 | + if (partial) lastPartial = partial; |
| 157 | + if (event.type === 'done' || event.type === 'error') { |
| 158 | + finish(); |
| 159 | + out.push(event); |
| 160 | + out.end(); |
| 161 | + return; |
| 162 | + } |
| 163 | + arm(timeouts.idleTimeoutMs, 'idle'); |
| 164 | + out.push(event); |
| 165 | + } |
| 166 | + if (!finished) { |
| 167 | + // Inner stream ended without a terminal event; mirror its result. |
| 168 | + finish(); |
| 169 | + const final = await stream.result(); |
| 170 | + out.push( |
| 171 | + final.stopReason === 'error' || final.stopReason === 'aborted' |
| 172 | + ? { type: 'error', reason: final.stopReason, error: final } |
| 173 | + : { type: 'done', reason: final.stopReason, message: final }, |
| 174 | + ); |
| 175 | + out.end(); |
| 176 | + } |
| 177 | + } catch (error) { |
| 178 | + if (finished) return; |
| 179 | + finish(); |
| 180 | + const aborted = callerSignal?.aborted === true; |
| 181 | + out.push({ |
| 182 | + type: 'error', |
| 183 | + reason: aborted ? 'aborted' : 'error', |
| 184 | + error: { |
| 185 | + ...snapshot(lastPartial, model), |
| 186 | + stopReason: aborted ? 'aborted' : 'error', |
| 187 | + errorMessage: error instanceof Error ? error.message : String(error), |
| 188 | + }, |
| 189 | + }); |
| 190 | + out.end(); |
| 191 | + } |
| 192 | + })(); |
| 193 | + |
| 194 | + return out; |
| 195 | + }) as StreamFn; |
| 196 | +} |
| 197 | + |
| 198 | +const STOPPED: unique symbol = Symbol('llm-stream-timeout-stopped'); |
| 199 | + |
| 200 | +/** Ask the inner stream to stop without waiting on one that may never settle. */ |
| 201 | +function releaseIterator(iterator: AsyncIterator<AssistantMessageEvent>): void { |
| 202 | + void Promise.resolve() |
| 203 | + .then(() => iterator.return?.()) |
| 204 | + .catch(() => undefined); |
| 205 | +} |
| 206 | + |
| 207 | +export class LLMStreamTimeoutError extends Error { |
| 208 | + override readonly name = 'TimeoutError'; |
| 209 | +} |
| 210 | + |
| 211 | +function timeoutMessage(phase: 'first' | 'idle', ms: number): string { |
| 212 | + return phase === 'first' |
| 213 | + ? `${LLM_STREAM_TIMEOUT_MESSAGE_PREFIX}: no response from the provider within ${ms}ms. ` + |
| 214 | + `Raise it with ${LLM_FIRST_EVENT_TIMEOUT_ENV} (milliseconds, 0 disables) or the host's per-model firstEventTimeoutMs.` |
| 215 | + : `${LLM_STREAM_TIMEOUT_MESSAGE_PREFIX}: the provider stream was idle for ${ms}ms. ` + |
| 216 | + `Raise it with ${LLM_STREAM_IDLE_TIMEOUT_ENV} (milliseconds, 0 disables) or the host's per-model streamIdleTimeoutMs.`; |
| 217 | +} |
| 218 | + |
| 219 | +function partialOf(event: AssistantMessageEvent): AssistantMessage | undefined { |
| 220 | + if (event.type === 'done') return event.message; |
| 221 | + if (event.type === 'error') return event.error; |
| 222 | + return event.partial; |
| 223 | +} |
| 224 | + |
| 225 | +/** Detach from the provider's mutable partial so a late abort cannot rewrite the reported message. */ |
| 226 | +function snapshot(partial: AssistantMessage | undefined, model: Model<Api>): AssistantMessage { |
| 227 | + return partial ? { ...partial, content: [...partial.content] } : blankAssistant(model); |
| 228 | +} |
| 229 | + |
| 230 | +function blankAssistant(model: Model<Api>): AssistantMessage { |
| 231 | + return { |
| 232 | + role: 'assistant', |
| 233 | + content: [], |
| 234 | + api: model.api, |
| 235 | + provider: model.provider, |
| 236 | + model: model.id, |
| 237 | + usage: { |
| 238 | + input: 0, |
| 239 | + output: 0, |
| 240 | + cacheRead: 0, |
| 241 | + cacheWrite: 0, |
| 242 | + totalTokens: 0, |
| 243 | + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, |
| 244 | + }, |
| 245 | + stopReason: 'error', |
| 246 | + timestamp: Date.now(), |
| 247 | + }; |
| 248 | +} |
| 249 | + |
| 250 | +function validTimeout(value: number | undefined): number | undefined { |
| 251 | + return typeof value === 'number' && Number.isSafeInteger(value) && value >= 0 ? value : undefined; |
| 252 | +} |
| 253 | + |
| 254 | +function parseTimeout(raw: string | undefined): number | undefined { |
| 255 | + if (raw === undefined || raw.trim() === '' || !/^\d+$/u.test(raw.trim())) return undefined; |
| 256 | + return validTimeout(Number(raw.trim())); |
| 257 | +} |
| 258 | + |
| 259 | +function readProcessEnv(): Readonly<Record<string, string | undefined>> { |
| 260 | + return (globalThis as { process?: { env?: Record<string, string | undefined> } }).process?.env ?? {}; |
| 261 | +} |
0 commit comments