From 855272710f37e1be95c422c99a2738ea3f0bcb14 Mon Sep 17 00:00:00 2001 From: mrsibe Date: Sun, 27 Sep 2026 14:08:31 +0800 Subject: [PATCH] =?UTF-8?q?feat(chat):=20guard=20the=20stream=20=E2=80=94?= =?UTF-8?q?=20idle=20deadline,=20pre-token=20retry,=20configurable=20ceili?= =?UTF-8?q?ng?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #150, the first half of #143. Three things the turn lifecycle could not do, all of them about a turn that goes wrong rather than about what a turn *is*. **A stream that goes quiet is cancelled.** Nothing noticed a provider that stopped producing while holding the connection open: the turn sat in `streaming` forever, the composer stayed locked, and a hung provider looked exactly like a slow one. There is a deadline now, reset by every event — reasoning deltas included, so a model that thinks for a long time is not cut off mid-thought. It *cancels* rather than settling, because whether a silent provider is worth another attempt is the retry policy's call, so `ChatExecution` learned to cancel without deciding anything. **A transient failure before the first token is retried.** Retrying is the wrong answer for a deterministic failure — same wait, same error, plus tokens — so only recognisably transient failures qualify (rate limit, 429, 5xx, overloaded, timeout, the transport's own errors), and anything unrecognised is reported as it is. Two more conditions: nothing the reader would see may have been sent, and the attempts are bounded with backoff. Structural events (`start`, `start-step`) do not count as content, which is what makes the retry reachable at all: verified against the pinned SDK, a second attempt's parts append to an empty message, so the reader sees one answer. **The output ceiling stops being a hard-coded constant.** `DEFAULT_MAX_TOKENS = 2048` was sent on every call and was not configurable anywhere. On a reasoning model the thinking tokens share that budget with the answer, so the reader saw `truncated` early even with a perfectly healthy stream. The connection now carries `maxOutputTokens`, and absent means no ceiling is sent at all — the model's own default applies. It is a per-connection field for the same reason the model id is: the right value belongs to the model, not to KnowNote. Settings shows Automatic or a custom value, chat connections only. Each attempt gets its own signal (`ChatExecution.beginAttempt`), so a retry after a timeout or a cancel does not inherit an aborted one, and the outcome a cancelled attempt reports is translated to `failed(timeout)` however the cancellation surfaced — a rejected read, or the SDK turning it into an `abort` chunk. Verified: `npm test` 336 pass (6 new in `test/retryPolicy.test.ts`; five in `test/chatStreamManager.test.ts` covering a retry that succeeds, a transient failure after content that is *not* retried, a deterministic failure that is not retried, bounded attempts, and a silent stream that is cancelled and fails as a timeout; plus the ceiling's two cases in `test/streamChat.test.ts`), `npm run typecheck`, `npm run check:design`, `npm run build`, `npm run build:unpack` + `smoke:packaged` (23 checks). Based on #142's branch: #149 is not merged yet, and this work is independent of it. --- src/main/ipc/validation.ts | 4 +- src/main/models/ModelClient.ts | 8 +- src/main/services/chat/ChatExecution.ts | 36 +++- src/main/services/chat/ChatStreamManager.ts | 182 +++++++++++++++--- src/main/services/chat/retryPolicy.ts | 52 +++++ .../settings/ModelConnectionForm.tsx | 42 ++++ src/renderer/src/locales/en-US/settings.json | 4 + src/renderer/src/locales/zh-CN/settings.json | 4 + src/shared/types/chat.ts | 6 +- src/shared/types/connection.ts | 8 + src/shared/utils/chatExecution.ts | 2 +- src/shared/utils/uiMessage.ts | 20 ++ test/chatStreamManager.test.ts | 180 +++++++++++++++++ test/retryPolicy.test.ts | 99 ++++++++++ test/streamChat.test.ts | 39 ++++ 15 files changed, 651 insertions(+), 35 deletions(-) create mode 100644 src/main/services/chat/retryPolicy.ts create mode 100644 test/retryPolicy.test.ts diff --git a/src/main/ipc/validation.ts b/src/main/ipc/validation.ts index 8c1ea84..52ae433 100644 --- a/src/main/ipc/validation.ts +++ b/src/main/ipc/validation.ts @@ -93,7 +93,9 @@ export const ConnectionSchemas = { protocol: z.enum(API_PROTOCOLS), baseUrl: z.string().min(1, 'Base URL 不能为空'), apiKey: z.string(), - modelId: z.string().min(1, 'Model ID 不能为空') + modelId: z.string().min(1, 'Model ID 不能为空'), + // Absent means "no ceiling": the provider's own default applies (#150). + maxOutputTokens: z.number().int().positive().optional() }) }), diff --git a/src/main/models/ModelClient.ts b/src/main/models/ModelClient.ts index f37cca0..b876ba3 100644 --- a/src/main/models/ModelClient.ts +++ b/src/main/models/ModelClient.ts @@ -15,7 +15,6 @@ import type { EmbeddingConfig, EmbeddingResult } from './types' import Logger from '../../shared/utils/logger' const DEFAULT_TEMPERATURE = 0.7 -const DEFAULT_MAX_TOKENS = 2048 /** * One provider call, as a stream. @@ -93,11 +92,16 @@ export class ModelClient { streamChat(messages: APIMessage[], options: { signal?: AbortSignal } = {}): ChatStream { Logger.debug('ModelClient', `Streaming with model: ${this.connection.modelId}`) + const maxOutputTokens = this.connection.maxOutputTokens + const result = streamText({ model: this.getAIModel(), messages: convertToCoreMessages(messages), temperature: DEFAULT_TEMPERATURE, - maxOutputTokens: DEFAULT_MAX_TOKENS, + // Only when the connection asks for a ceiling. Absent means the model's own + // default, which is what a reasoning model needs: its thinking shares this + // budget with the answer, so 2048 of them was a ceiling reached early (#150). + ...(maxOutputTokens !== undefined && { maxOutputTokens }), // Owned by the caller, so that stopping a turn actually cancels the provider // request rather than only recording that it was stopped. abortSignal: options.signal diff --git a/src/main/services/chat/ChatExecution.ts b/src/main/services/chat/ChatExecution.ts index e1848ea..cfa6c9d 100644 --- a/src/main/services/chat/ChatExecution.ts +++ b/src/main/services/chat/ChatExecution.ts @@ -36,7 +36,13 @@ export class ChatExecution { private currentStatus: ChatExecutionStatus = 'pending' private settledOutcome: ChatExecutionOutcome | undefined - private readonly controller = new AbortController() + /** + * The attempt in flight. + * + * Replaced by `beginAttempt`, because a retry must not inherit a signal that a + * timeout or a cancel has already aborted (#150). + */ + private attemptController = new AbortController() constructor(init: { id: string; messageId: string; sessionId: string; notebookId?: string }) { this.id = init.id @@ -60,9 +66,31 @@ export class ChatExecution { return this.settledOutcome !== undefined } - /** Handed to the provider call, so a stop cancels the request itself. */ + /** The signal of the attempt in flight. */ get signal(): AbortSignal { - return this.controller.signal + return this.attemptController.signal + } + + /** + * Open a fresh attempt, and hand back the signal that cancels it. + * + * Called before every provider call, including a retry: the previous attempt's + * signal is spent. + */ + beginAttempt(): AbortSignal { + this.attemptController = new AbortController() + return this.attemptController.signal + } + + /** + * Cancel the attempt in flight without deciding anything about the turn. + * + * The idle deadline uses this: stopping a request that has gone quiet is not a + * statement about what the turn was, and the attempt that was cancelled reports + * back as a failure the retry policy can look at (#150). + */ + cancel(): void { + this.attemptController.abort() } /** The first chunk, which is what turns "waiting" into "streaming". */ @@ -95,7 +123,7 @@ export class ChatExecution { const settled = this.settle(abortedOutcome(reason)) if (!settled) return null - this.controller.abort() + this.attemptController.abort() return settled } } diff --git a/src/main/services/chat/ChatStreamManager.ts b/src/main/services/chat/ChatStreamManager.ts index 7df6e41..57215af 100644 --- a/src/main/services/chat/ChatStreamManager.ts +++ b/src/main/services/chat/ChatStreamManager.ts @@ -12,7 +12,13 @@ import type { import type { Citation, CitationContext } from '../../../shared/types/citation' import { classifyTerminal, failedOutcome } from '../../../shared/utils/chatExecution' import { resolveCitations } from '../../../shared/utils/citationResolution' -import { isTerminalChunk, messageReasoning, messageText } from '../../../shared/utils/uiMessage' +import { + carriesContent, + isTerminalChunk, + messageReasoning, + messageText +} from '../../../shared/utils/uiMessage' +import { retryDelayMs, shouldRetry } from './retryPolicy' import { estimateTokens } from '../../../shared/utils/tokenEstimate' import Logger from '../../../shared/utils/logger' import { ChatExecution } from './ChatExecution' @@ -58,6 +64,18 @@ import type { SessionAutoSwitchService } from '../SessionAutoSwitchService' * was stopped. */ +/** + * How long a turn may produce nothing before its request is cancelled. + * + * Generous on purpose: this is for a stream that has died, not for a model that is + * slow to start. Every event resets it, so a reasoning model is never cut off + * mid-thought (#150). + */ +const IDLE_TIMEOUT_MS = 60_000 + +/** Waits before the next attempt. */ +const wait = (ms: number): Promise => new Promise((resolve) => setTimeout(resolve, ms)) + /** What the turn's persistence needs, as a port: see `turnStore.ts`. */ export interface ChatTurnStore { createTurn(sessionId: string): { id: string } @@ -100,6 +118,19 @@ interface TurnEntry { usage?: ChatTokenUsage /** The message of an `error` chunk, when the provider reported one mid-stream. */ streamError?: string + /** + * Whether anything the reader would see has been forwarded. + * + * Sticky across attempts on purpose: once content has been sent, no later attempt + * may be started, whatever the failure (#150). + */ + contentSent?: boolean + /** Set when the idle deadline cancelled the attempt, so its failure reads as one. */ + timedOut?: boolean + /** The attempt in flight, 1-based; logged, and what the retry bound counts. */ + attempts?: number + /** The watchdog for the attempt in flight. */ + idle?: { reset: () => void; stop: () => void } /** Monotonic per execution, so a consumer can detect a gap. */ seq: number } @@ -112,6 +143,8 @@ export class ChatStreamManager { private readonly deps: { store: ChatTurnStore sessionAutoSwitchService: Pick + /** Overridable so a test does not have to wait a minute to see the deadline. */ + idleTimeoutMs?: number } ) {} @@ -199,40 +232,134 @@ export class ChatStreamManager { return executionId === undefined ? undefined : this.turns.get(executionId) } + /** + * Run the turn, attempting it again while the retry policy says so (#150). + * + * Each attempt is a fresh provider call with its own signal. A retry is only + * reachable while nothing has reached the reader, so the answer on screen is never + * duplicated or silently replaced by a second attempt. + */ private async run(entry: TurnEntry): Promise { + for (let attempt = 1; ; attempt += 1) { + entry.attempts = attempt + const outcome = await this.attempt(entry) + + // A stop the reader asked for has already settled the turn. + if (entry.execution.isSettled) return + + if (shouldRetry({ outcome, attempt, contentSent: entry.contentSent === true })) { + const delay = retryDelayMs(attempt + 1) + Logger.warn( + 'ChatStreamManager', + `Turn ${entry.execution.id} attempt ${attempt} failed (${outcome.status}${ + outcome.status === 'failed' ? `: ${outcome.error.message}` : '' + }); retrying in ${delay}ms` + ) + await wait(delay) + if (entry.execution.isSettled) return + continue + } + + this.settleTurn(entry, outcome) + return + } + } + + /** + * One provider call, from opening the attempt to its outcome. + * + * Nothing is settled here: a failure is a value the retry policy looks at, and the + * turn ends only once that policy has had its say. + */ + private async attempt(entry: TurnEntry): Promise { const { execution, request } = entry const client = request.client - if (!client) return + if (!client) { + return failedOutcome('Chat model not configured, please configure in settings', 'error') + } + // Per-attempt facts, cleared so a retry cannot inherit the previous attempt's. + entry.message = undefined + entry.finishReason = undefined + entry.usage = undefined + entry.streamError = undefined + entry.timedOut = false + + const signal = execution.beginAttempt() + entry.idle = this.watchIdle(entry) + + let outcome: ChatExecutionOutcome try { - const { events } = client.streamChat(request.messages, { signal: execution.signal }) + const { events } = client.streamChat(request.messages, { signal }) const [toRenderer, toAssemble] = events.tee() - - await Promise.all([ - this.forward(entry, toRenderer), - // The SDK assembles, and nothing here re-derives the message from deltas. - // - // Kept on the entry as each snapshot arrives rather than assigned once at the - // end: a turn can be stopped or fail while this branch is still running, and - // the part of the answer that had arrived is what gets persisted (#138 - // invariant 5). - (async () => { - for await (const snapshot of readUIMessageStream({ stream: toAssemble })) { - entry.message = snapshot - } - })() - ]) + await Promise.all([this.forward(entry, toRenderer), this.assemble(entry, toAssemble)]) + // The stream ended without being stopped. What that means is + // `classifyTerminal`'s decision: a `finish` chunk's reason, or — when no + // terminal chunk arrived at all — an EOF, which is not success (#138 + // invariant 1). + outcome = this.outcomeFor(entry) } catch (error) { - // A transport error arrives as a rejection rather than as an event. The - // partial answer is kept, and the turn ends as a failure. - this.settleTurn(entry, failedOutcome((error as Error).message, 'error')) - return + // A transport error arrives as a rejection: the partial answer is kept, and the + // failure is reported as the connection's. + outcome = failedOutcome((error as Error).message, 'error') + } finally { + entry.idle?.stop() + entry.idle = undefined } - // The stream ended without being stopped. What that means is `classifyTerminal`'s - // decision: a `finish` chunk's reason, or — when no terminal chunk arrived at all - // — an EOF, which is not success (#138 invariant 1). - this.settleTurn(entry, this.outcomeFor(entry)) + // However the cancellation surfaced — a rejected read, or the SDK turning it into + // an `abort` chunk — the reason this attempt failed is the deadline that cancelled + // it (#150). + return entry.timedOut ? failedOutcome('The model stopped responding.', 'timeout') : outcome + } + + /** + * The SDK assembles; nothing here re-derives the message from deltas. + * + * Kept on the entry as each snapshot arrives rather than assigned once at the end: + * a turn can be stopped or fail while this is still running, and the part of the + * answer that had arrived is what gets persisted (#138 invariant 5). + */ + private async assemble(entry: TurnEntry, stream: ReadableStream): Promise { + for await (const snapshot of readUIMessageStream({ stream })) { + entry.message = snapshot + } + } + + /** + * The deadline for a stream that has gone quiet. + * + * Reset by every event — reasoning deltas included, so a model that thinks for a + * long time is not killed mid-thought. It *cancels* rather than settling: whether a + * silent provider is worth trying again is the retry policy's decision, not the + * watchdog's. + */ + private watchIdle(entry: TurnEntry): { reset: () => void; stop: () => void } { + let timer: NodeJS.Timeout + + const onIdle = (): void => { + if (entry.execution.isSettled) return + Logger.warn( + 'ChatStreamManager', + `Turn ${entry.execution.id} produced nothing for ${this.deps.idleTimeoutMs ?? IDLE_TIMEOUT_MS}ms; cancelling the request` + ) + entry.timedOut = true + entry.execution.cancel() + } + + const arm = (): void => { + timer = setTimeout(onIdle, this.deps.idleTimeoutMs ?? IDLE_TIMEOUT_MS) + } + + arm() + + return { + reset: () => { + clearTimeout(timer) + arm() + }, + stop: () => clearTimeout(timer) + } } /** @@ -258,6 +385,9 @@ export class ChatStreamManager { } entry.execution.begin() + entry.idle?.reset() + if (carriesContent(value)) entry.contentSent = true + entry.seq += 1 entry.request.emit({ type: 'chunk', diff --git a/src/main/services/chat/retryPolicy.ts b/src/main/services/chat/retryPolicy.ts new file mode 100644 index 0000000..9830797 --- /dev/null +++ b/src/main/services/chat/retryPolicy.ts @@ -0,0 +1,52 @@ +import type { ChatExecutionOutcome } from '../../../shared/types/chat' + +/** + * Whether a failed turn is worth attempting again (#150). + * + * The rule is deliberately narrow. A failure that is retried when it will fail the + * same way every time costs the reader time, tokens and a second wait for the same + * answer — so only failures that are recognisably transient are retried, and + * anything unrecognised is reported as it is. + * + * Two conditions beyond the kind of failure, both of which the caller supplies: + * + * - **nothing has been shown yet.** Once content has reached the reader, a retry + * would either duplicate it or silently replace it. + * - **the attempts are bounded.** Backoff on top, so a provider that is down is not + * hammered. + */ + +/** How many attempts one turn may make, the first included. */ +export const MAX_ATTEMPTS = 3 + +/** How long to wait before attempt `next` (1-based). */ +export const retryDelayMs = (next: number): number => 500 * 2 ** Math.max(0, next - 2) + +/** + * Failures that mean "try again", by the text the provider or transport gave us. + * + * A short list on purpose. Providers do not agree on error shapes and the UI event + * carries only a message, so this is a best-effort reading of what is transient — + * and the default is *not* to retry, which is the safe direction to be wrong in. + */ +const TRANSIENT = + /(rate.?limit|too many requests|\b429\b|\b50[0-9]\b|overloaded|temporarily unavailable|timeout|timed out|ECONNRESET|ECONNREFUSED|ECONNABORTED|ETIMEDOUT|EAI_AGAIN|EPIPE|socket hang up|premature close|fetch failed|network)/i + +export const isRetryableFailure = (outcome: ChatExecutionOutcome): boolean => { + if (outcome.status !== 'failed') return false + + // A stream that went quiet and was cancelled, and a stream that ended without + // saying anything, are both the connection's fault rather than the request's. + if (outcome.reason === 'timeout' || outcome.reason === 'unexpected_eof') return true + + return TRANSIENT.test(outcome.error.message) +} + +export const shouldRetry = (options: { + outcome: ChatExecutionOutcome + /** The attempt that just failed, 1-based. */ + attempt: number + /** Whether any content has reached the reader. */ + contentSent: boolean +}): boolean => + options.attempt < MAX_ATTEMPTS && !options.contentSent && isRetryableFailure(options.outcome) diff --git a/src/renderer/src/components/settings/ModelConnectionForm.tsx b/src/renderer/src/components/settings/ModelConnectionForm.tsx index c26db89..179fe76 100644 --- a/src/renderer/src/components/settings/ModelConnectionForm.tsx +++ b/src/renderer/src/components/settings/ModelConnectionForm.tsx @@ -60,6 +60,12 @@ const EMPTY_CONNECTION: ModelConnection = { modelId: '' } +/** A ceiling below 1 is not a ceiling; clearing the input means "automatic" (#150). */ +const toCeiling = (raw: string): number | undefined => { + const parsed = Number.parseInt(raw, 10) + return Number.isFinite(parsed) && parsed > 0 ? parsed : undefined +} + export default function ModelConnectionForm({ capability, connection, @@ -286,6 +292,42 @@ export default function ModelConnectionForm({ {fetchError &&

{fetchError}

} + {/* Output ceiling. Chat only: an embedding connection has no output (#150). */} + {capability === 'chat' && ( +
+ + + {value.maxOutputTokens !== undefined && ( + update({ maxOutputTokens: toCeiling(e.target.value) })} + /> + )} +

{t('maxOutputTokensHint')}

+
+ )} + {/* Test connection */}