Skip to content

feat(api): abort signal support for gemini, mistral, lite-llm (completePrompt + createMessage) - #1303

Merged
edelauna merged 57 commits into
Zoo-Code-Org:mainfrom
easonLiangWorldedtech:feat/abort-r1-gemini-mistral-lite
Oct 30, 2026
Merged

edelauna merged 57 commits into
Zoo-Code-Org:mainfrom
easonLiangWorldedtech:feat/abort-r1-gemini-mistral-lite

Conversation

@easonLiangWorldedtech

@easonLiangWorldedtech easonLiangWorldedtech commented Aug 20, 2026 •

Copy link
Copy Markdown
Contributor

Add abort-signal support to the Gemini, Mistral, and LiteLLM providers.

Changes

  • gemini.ts
    • completePrompt: forwards CompletePromptOptions.abortSignal to GenerateContentConfig.abortSignal and timeoutMs to httpOptions.timeout (httpOptions is omitted entirely when nothing is set). On catch, a user-initiated abort re-throws as a standard DOMException with name = "AbortError".
    • createMessage: bridges metadata.abortSignal into a request-local AbortController passed to the SDK via config.abortSignal. A pre-aborted signal rejects immediately with AbortError; the abort listener is stored in a named const and removed in finally.
  • mistral.ts
    • completePrompt: forwards abortSignal via fetchOptions.signal and timeoutMs to the Mistral SDK RequestOptions (options arg omitted when empty, preserving the legacy 1-arg call shape). Abort normalization in catch as above.
    • createMessage: same request-local controller bridging as Gemini; the stream call only receives { fetchOptions: { signal } } when a signal is present.
  • lite-llm.ts
    • completePrompt: forwards abortSignal via OpenAI.RequestOptions.signal; timeoutMs is only forwarded when > 0 because the OpenAI SDK treats a 0 timeout as an immediate abort.
    • createMessage: same bridging pattern; the in-flight chat.completions.create(...).withResponse() call receives the request signal alongside the existing X-Zoo-Session-ID header.
  • vertex.ts: unchanged — VertexHandler inherits the new behavior from GeminiHandler.

Tests

  • Ported the reference completePrompt request-options coverage (gemini, vertex, gemini-handler, mistral, lite-llm specs).
  • New createMessage bridging regression tests per provider: pre-aborted signal rejects immediately with name = "AbortError" (no SDK call), and a mid-flight external abort propagates into the in-flight request and surfaces as AbortError on the stream.
  • Mistral spec additionally covers timeoutMs: 0 forwarding (valid for the Mistral SDK, which uses a truthy check) and no-signal call-shape preservation.
  • Vertex spec verifies the inherited bridging behavior.
  • All 5 specs green; tsc --noEmit and per-file ESLint (--max-warnings=0) clean.

Part of the abort-signal series (round 1). Builds on #674, #901, #1008. Addresses #404.

Review feedback addressed (2026-09-11)

Per @edelauna's review, this revision leverages the shared abort-signal work already on main:

  • Catch normalization: all six catch sites (gemini / mistral / lite-llm, createMessage + completePrompt) now use isRequestAborted(error, signal) — which also catches SDK-native abort errors that surface before the signal flag propagates — and throw via createAbortError(<provider>).
  • Signal bridging: the per-provider manual AbortController + addEventListener/removeEventListener + finally cleanup is replaced with RequestConfigBuilder.addMergedSignal (feat(api): introduce RequestConfigBuilder for SDK-agnostic abort signal support #1008) — AbortSignal.any internally, no manual listener management.
  • Tests: new per-provider regression tests pin the SDK-native-abort branch; the listener-identity assertions now assert that no manual listeners are attached to the external signal.

Series alignment check (zdt align check)

exit 0 — 0 ERROR. Findings recorded, not refactored:

  • 3 WARN — statically unguarded first-yield sites (gemini fallback-text / lite-llm reasoning /
    mistral string-content yields). A guard before the FIRST yield of a loop iteration is dead code:
    the top-of-loop break runs without a suspension point before it, so it can never throw. The
    loop-top check + post-loop AbortError cover those sites; the structural kill-tests pin them.
  • 2 DIVERGENCE — per-method mechanism signature across the series (bridging builder + signal
    discovery: gemini/mistral vs bare for lite-llm createMessage/completePrompt). Recorded as a
    design decision; aligning would widen this PR beyond its unit.

Review feedback addressed (2026-09-23)

Per the CodeRabbit out-of-diff finding on cold-cache model discovery (LiteLLMHandler.createMessage awaits fetchModel() before the abort check):

  • The fetchModel() awaits in both createMessage and completePrompt now race against the request signal via rejectOnAbort (the shared series helper in utils/abort-signal.ts), and an abort landing while discovery fails is normalized to the standard AbortError via isRequestAborted + createAbortError("LiteLLM").
  • fetchModel() itself is a shared single-flight lookup, so it carries no per-request signal; cancellation is expressed per-request at the caller, and the underlying discovery keeps running so concurrent requests still get the populated cache.

New head: f0f1fa5b9. Tests: 91 green across lite-llm.spec.ts (66, incl. the pending-discovery + abort regression) and abort-signal.spec.ts (25, incl. the rejectOnAbort contract); tsc --noEmit and per-file ESLint (--max-warnings=0) clean; local Stryker preflight on the unit delta (base 0ec250e3d): 226 valid mutants, 226 killed, PASS.


CI note (2026-09-28)

  • Merged upstream/main into the branch (head 4f746b855) to clear drift; no code change from this unit.
  • e2e-mock failed at 4f746b855 solely on the upstream Markdown List Rendering suite (3 tests, 30s waitUntilCompleted timeout); every other required check was green. The same failure signature (same suite, same timeout) occurred in the same CI window on two unrelated external forks, and the previous head of this branch passed e2e-mock with the identical test file present. This PR's diff touches no e2e or mock content.
  • Empty commit 0388aeccd re-triggers the workflow (no content change). That run failed only on the restart-persistence step (verify-phase TaskMessagesReadError — the single 1-10ms ENOENT retry in readFileWithMissingRetry lost the race under CI load — then a 30s history-sequence poll timeout); the identical tree passed that step in the previous run, and the e2e-mock workflow shows chronic cross-fork flakiness over recent weeks.
  • Second empty commit c299aa213 re-triggers the workflow.

Review feedback addressed (2026-10-13)

Responding to the maintainer's two inline comments on fork PR 1303. New head: b94c13c1f.

  • LiteLLM absolute deadline (lite-llm.ts, completePrompt): the timeout was previously forwarded to the OpenAI SDK only as the per-attempt timeout option while the SDK's default two retries stayed enabled — each retry adds a backoff sleep that is not signal-aware, so a 30ms budget could elapse ~1,530ms across three attempts, exceeding the promised discovery-plus-completion budget. One absolute deadline (start + timeoutMs) is now enforced THROUGH the completion call: a deadline signal for the remaining budget is merged with the caller's signal and passed as the SDK call's signal (in-flight attempts abort at the deadline; the SDK's next attempt short-circuits at its pre-flight signal.aborted check), and because the SDK's retry sleep is not signal-aware, the call is additionally raced against the deadline via rejectOnAbort so the caller settles at the deadline even mid-backoff. Retries stay enabled within the remaining budget. Timeout expiry stays distinct from caller cancellation: deadline expiry with the caller still running surfaces as TimeoutError ("The LiteLLM completion request timed out") and never masquerades as the cancellation AbortError (the SDK reports deadline expiry as an abort-shaped APIUserAbortError); a caller abort still yields the standard AbortError and wins when both fire. Three regression tests pin this (real-SDK integration test with a hanging fetch stub asserting one attempt / bounded elapsed / TimeoutError; deadline expiry reported as a timeout, not caller cancellation; caller cancellation still AbortError), and the merged signal+timeout test now pins the deadline-bounded signal the SDK receives. Negative control: reverting the fix turns exactly those tests red.
  • Gemini HTTPS-only endpoint policy carved out: the assertSecureGeminiBaseUrl policy (helpers plus both call-site guards) and its CWE-319 security test matrix now live in their own PR — upstream PR 2033 (branch feat/gemini-https-endpoint-policy, base main) — because the policy changes which custom endpoints are accepted independently of cancellation. This branch keeps cancellation content only: the remaining Gemini diff is the abort-signal bridging and httpOptions timeout forwarding.
  • Refreshed against upstream main: merged current main to clear the dirty state. The only conflict was the tail of lite-llm.spec.ts (both sides appended blocks; kept both — main's model-catalog scope test and this branch's completePrompt suite). Main's model-cache changes (fetchModels / getModelCacheScope on the ApiHandler interface) auto-merged cleanly alongside this branch's abort changes.

Verification at this head: 209 tests green across all six specs this PR touches (gemini, gemini-handler, mistral, vertex, lite-llm, request-timeout); tsc --noEmit → 0 errors; per-file ESLint (--max-warnings=0) clean with no suppression-count changes; prettier --check clean on all touched files.


Review feedback addressed (re-review round, 2026-10-13)

The re-review at head b94c13c1f opened three threads; all fixed at head d382f33c1.

  • Weak assertions in the gemini stream-failure tests (assertion identity): both non-abort stream-failure tests now pin the wrapping path by its observable shape — a NEW Error instance (not the raw SDK error), a non-raw message, and a non-AbortError name — so neither a removed wrapper (raw rethrow) nor an abort-path misclassification can pass. Design-contract note: the proposed exact t("common:errors.gemini.generate_stream", ...) equality is not assertable in this environment — i18next is not initialized with resources in unit tests (documented in this spec's sibling telemetry test), so t() returns undefined and the provider's wrapped message is ""; the proposed toBe(t(...)) compares undefined against "" and fails on green code. Negative control: raw-rethrow and abort-misclassification mutations each turn exactly these two tests red; byte-exact restore returns them green.
  • Fake-timer leak risk (cleanup terminal paths): the merged-signal test body is wrapped in try/finally with vi.useRealTimers() in the finally block, matching the neighbouring tests; a failed assertion can no longer leak fake timers into the deadline or discovery-timeout tests.
  • Equivalent ConditionalExpression mutant (lite-llm deadline signal branch): collapsed to callerAbortSignal ? mergeAbortSignals(callerAbortSignal, deadlineSignal) : deadlineSignal — mergeAbortSignals returns the primary signal directly when no secondary exists, so all four caller/deadline combinations behave exactly as before and the guarded branch's equivalent mutant is gone; the ternary's own mutants are killed by the existing caller-only and deadline-only tests. No behavior change in production code.

Verification at this head: 129/129 green across gemini.spec.ts + lite-llm.spec.ts; local mutation preflight on the unit delta green (240 valid / 240 killed / 0 survived / 0 noCoverage); tsc --noEmit → 0 errors; per-file ESLint (--max-warnings=0) clean with no suppression-count changes; prettier --check clean on all touched files.

Stryker directive removal (equivalence notes)

Six directives removed (gemini.ts 1, lite-llm.ts 4, mistral.ts 1):

  • OptionalChaining on requestSignal?.aborted (x3): requestSignal is the result of the merged-signal helper called unconditionally above the check in the same method scope, so the optional chain can never observe nullish - the mutant is equivalent by a local invariant.
  • StringLiteral on the race's providerName "LiteLLM" (x3): every error escaping rejectOnAbort/the deadline race is reclassified to the contract abort error before reaching callers, so the provider label embedded in the intermediate error is unobservable through any public path.

Timeout guard hardening (helper port)

Port of the #1300 CodeRabbit finding to the shared mergeAbortSignalAndTimeout helper: non-finite timeoutMs (Infinity) now means "no timeout" instead of throwing ERR_OUT_OF_RANGE from AbortSignal.timeout, and finite delays are capped at 2^31-1 ms (Node clamps larger delays to 1ms, which aborted requests almost immediately). Regression tests added; negative controls fail exactly the two new tests. Identical helper fix lands in every unit of the abort series (merge-order aware); canonical copy in #1293.

Fractional-timeout rounding wave

Follow-up CR round: Node validates AbortSignal.timeout delays as uint32, so finite fractions (1500.5, 0.5) threw ERR_OUT_OF_RANGE like Infinity did. mergeAbortSignalAndTimeout now rounds delays up (a positive fraction never becomes 0 ms) before capping at 2^31-1 ms; regression test mergeAbortSignalAndTimeout(undefined, 1.5). Identical helper patch in every unit of the series (merge-order aware, canonical in #1293).

Also (lite-llm): getRequestTimeoutMs rejects non-finite values and rounds up; remainingMs clamped to [1, 2^31-1] whole ms before the deadline signal and before forwarding as createOptions.timeout. Regression test covers 1500.5 / Infinity / 2^31.

Assertion strengthening (rounding wave follow-up)

The fractional-timeout test now spies AbortSignal.timeout and pins the received delay (2 for 1.5, 1 for 0.5, never 0) instead of only asserting no-throw plus eventual abort; a round-down mutation (Math.ceil -> Math.floor) fails exactly this test. Identical patch in every unit of the series (canonical in #1293).

Loading
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants