Skip to content

Commit 159c42e

Browse files
easonLiangWorldedtecheasonliang28easonLiangWorldedtechdsh-agentedelauna
authored
feat(api): abort signal support for gemini, mistral, lite-llm (completePrompt + createMessage) (#1303)
* feat(api): abort signal support for gemini, mistral, lite-llm (completePrompt + createMessage) * fix(api): unify timeoutMs:0 handling across gemini/mistral/lite-llm + fix test title * test(api): close changed-line coverage gaps in gemini, mistral, lite-llm * test(api): use provider identifiers in gemini-handler spec Replace raw gemini apiProvider literals in the two abort-signal spec cases with providerIdentifiers.gemini, matching the rest of the file and the zoo/no-raw-provider-identifiers rule that CI lint enforces. * fix(api): harden gemini/mistral abort, timeout and base-url handling Address CodeRabbit findings on the abort-signal series: - gemini: reject non-HTTPS (non-loopback) googleGeminiBaseUrl before requests so API keys are never sent over cleartext (CWE-319) - mistral: route completePrompt timeout through mergeAbortSignalAndTimeout so the timeout actually cancels the request, instead of a dead timeoutMs field - gemini/mistral/lite-llm specs: request-local signal identity assertions, readiness barrier instead of fixed delay, makeApiHandlerOptions over casts * chore: retrigger CodeRabbit review (no-op) The incremental review for the previous head was stuck in a phantom "review finished" state on the CodeRabbit side (the review object never materialized), so this no-op commit moves the head to a fresh sha and forces a new incremental review. No code changes. * fix(api): allow IPv6 loopback host and strengthen abort-signal spec kills (mutation gate) * test(api): kill remaining abort-signal mutation survivors in gemini and mistral specs (mutation gate) * fix(api): address review findings on abort-bridge listener assertions and 127. hostname anchoring - gemini/lite-llm/mistral specs: capture the bridge's registered abort listener and assert the exact reference on removeEventListener instead of expect.any(Function) (CodeRabbit finding). - gemini spec: add a foo127.bar case — a hostname new URL() accepts that contains the 127. substring but is not loopback — proving the anchored 127. check is observable; drop the Stryker disable Regex directive in gemini.ts whose unobservable-anchor rationale is wrong (hosts such as a127.0.0.1 are rejected by new URL() itself). * fix(api): restrict Gemini loopback HTTP exception to literal 127.0.0.0/8 A public hostname may legitimately start with a 127. label (e.g. 127.example.test), which the prefix-style /^127\./ test misclassified as loopback and allowed as a cleartext base URL with API-key auth. isLoopbackHostname now requires a literal IPv4 loopback address: four dot-separated parts, first part exactly 127, remaining parts decimal octets 0-255. localhost and [::1] handling is unchanged. Adds regression tests rejecting http://127.example.test and http://10.0.0.1 and accepting the inclusive boundary 127.255.255.255. * test(api): restructure Gemini loopback octet check for mutation observability Extract the octet predicate into isOctet() so the single-line 127.0.0.0/8 check has no every/slice method-chain mutants, and pin the remaining unobservable ConditionalExpression variants with true-rationale Stryker directives. Add reject tests for 127.0.0.a (non-numeric last octet), 127.0.a.b (non-numeric middle octet) and 127.0.0.1.a (five-part host). * test(api): collapse duplicated Gemini base-URL rejection tests into a table-driven case * refactor(api): leverage abort-signal utils and RequestConfigBuilder for gemini/mistral/lite-llm Address review feedback (edelauna, #1303): - catch blocks: isRequestAborted(error, signal) catches SDK-native abort errors that surface before the signal flag propagates (error-name and exact-message branches), and createAbortError normalizes the thrown error to the series abort contract - createMessage bridge: RequestConfigBuilder.addMergedSignal (AbortSignal.any) replaces the manual AbortController + addEventListener/removeEventListener plumbing; the finally cleanup blocks are gone - pre-abort fast-fail uses the throwIfAborted helper - specs: message assertions updated to the helper messages, listener-identity assertions inverted to assert no manual listener management, and new regression tests pin the SDK-native-abort branch per provider * fix(api): run LiteLLM pre-abort fast-fail before model discovery Address CodeRabbit finding on #1303 (review 5174654329): LiteLLMHandler.createMessage awaited fetchModel() before throwIfAborted, so an already-aborted request could reach provider model discovery (getModels/refreshModels) on a cold cache, and a discovery failure would escape as a model-fetch error instead of AbortError. completePrompt had the same shape. Both methods now fast-fail with the canonical AbortError before model discovery begins. Tests: the completePrompt aborted-signal case is rewritten as a true in-flight abort (still covering catch-block classification via signal.aborted, without leaking a once-mock), a new completePrompt pre-abort test added, and the createMessage pre-abort test now pins that fetchModel is never called for an already-aborted request. * fix(api): stop provider streams after abort instead of leaking buffered content Class h of the abort-signal contract (streaming yield granularity): a for-await loop keeps pulling and yielding buffered content after the request is aborted, and the SDK swallows the mid-stream AbortError, so an aborted request could finish as a normal stream completion. createMessage in gemini/mistral/lite-llm now: - breaks at the top of the loop once the request-local signal is aborted (for-await pulls the already-buffered element before the check runs, so the break prevents processing, not pulling); - re-checks before every yield site reachable after a suspension point (a guard before the first yield of an iteration is dead code - the top check runs without a suspension point before it - and is intentionally absent); - throws the canonical AbortError after the loop, so a break or a swallowed mid-stream abort surfaces as "The <Provider> request was aborted" instead of a normal completion. completePrompt in gemini/mistral additionally fast-fails before building the request (lite-llm received that in the previous commit). Tests: new "streaming loop abort defense" suites per provider (per-yield-site mid-chunk aborts plus a structural break/post-loop case with a pull counter, using the unguarded first-yield shape so only the break can prevent a leak), completePrompt pre-abort fast-fail tests, and the pre-aborted catch-path cases rewritten as true in-flight aborts (the mock stays pending until the external signal aborts - no once-mock leaks). The lite-llm in-flight mock attaches a no-op rejection handler to its pending element, which the top-of-loop break correctly never reads. Extra kill tests for lines swept into the diff hunks: responseId capture, non-Error throw rethrow/wrap branches (createMessage and completePrompt), supportsTemperature true/false temperature config, and a usage-only empty-choices chunk. Dead write-only hasContent/hasReasoning flags in gemini createMessage removed (they generated equivalent BooleanLiteral mutants inside the diff hunks). Mutation gate: three equivalent OptionalChaining mutants on the top-of-loop checks carry mutator-specific Stryker directives (requestSignal is always set by the addMergedSignal call two lines above - a request-local controller signal exists even without an external signal). Verified: tsc clean; 209/209 across the six delta suites; eslint --prune-suppressions with flat counts; zdt mutation preflight 204 valid / 204 killed / 0 survived (exit 0); zdt align check exit 0 (0 ERROR; 3 WARN = statically unguarded first-yield sites subsumed by the loop-top check; 2 DIVERGENCE = recorded per-method mechanism divergence). * test(api): assert the mapped usage values in the LiteLLM usage-only case The usage-only empty-choices test checked only the chunk type. Assert the mapped values (inputTokens: 5, outputTokens: 7) so the prompt_tokens / completion_tokens mapping in this branch is pinned, not just its presence. Addresses the CodeRabbit finding on the structural kill-test suite. * fix(api): settle LiteLLM requests with AbortError during model discovery Model discovery (RouterProvider.fetchModel) is a shared single-flight fetch, so a per-request signal must not be threaded into it: one caller's abort would reject the in-flight discovery of every other waiter. Instead, race each request's discovery against its own signal via rejectOnAbort (new utility in utils/abort-signal) and normalize aborted discovery failures with createAbortError("LiteLLM"): - createMessage: race discovery against metadata.abortSignal - completePrompt: race discovery against the merged signal (mergeAbortSignalAndTimeout: external abort + per-request timeout, timeoutMs <= 0 disabling the timeout), so a stalled cold-cache fetch cannot outlive the requested timeout - non-abort discovery failures propagate unchanged (exact identity) - the completePrompt "abort while in flight" test now lands the abort after discovery (readiness barrier) so it still exercises the SDK catch's signal.aborted classification; the bridging test now asserts the transient discovery-race listener is the only manual listener on the external signal and that it is detached at settle - kill tests: mid-discovery abort, preserved non-abort failure, concurrent failure+abort normalization, single-flight isolation (a concurrent non-aborted request sharing the fetch is unaffected), and the timeoutMs boundary; unit tests for rejectOnAbort itself Addresses the CodeRabbit out-of-diff finding on the current head ("propagate cancellation through cold-cache model discovery"). * ci: retrigger e2e-mock after flaky markdown-lists timeout The e2e-mock run at 4f746b8 failed only on the 'Markdown List Rendering' suite (3 tests, 30s waitUntilCompleted timeout); every other required check was green. The identical failure signature (same suite, same 30s 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 to retrigger the workflow. * ci: retrigger e2e-mock (second flake: restart-persistence 30s timeout) Run 36404841531 failed on the 'Run mocked restart-persistence E2E test' step only (the main mocked suite, incl. Markdown List Rendering, passed). The verify phase rehydrated the persisted task, hit TaskMessagesReadError (readFileWithMissingRetry's single 1-10ms ENOENT retry lost the race under CI load), and the 30s history-sequence poll timed out. The identical tree (4f746b8) passed that step in run 36396120353, and the e2e-mock workflow shows chronic cross-fork flakiness over recent weeks (repeated re-rolls on other branches). No content change in this PR touches e2e or task persistence. * fix(api): bound LiteLLM completePrompt by a single deadline and classify discovery timeouts distinctly Address the CodeRabbit findings on completePrompt: - Compute one deadline before model discovery and pass only the remaining budget to the SDK via createOptions.timeout, so a request's total lifetime is bounded by the configured timeoutMs instead of discovery plus a fresh full timeout. - A per-request timeout elapsing during discovery now surfaces as a timeout-specific error (name TimeoutError) rather than the abort contract error: a timeout is not a user cancellation. Caller aborts still settle with the standard AbortError, and fetcher-originated AbortError-shaped failures keep their existing normalization. - The shared-discovery sibling test now drains the task queue before asserting the sibling has not settled, making the assertion meaningful instead of timing-trivially-true. - A new spec asserts the SDK receives the remaining (not full) timeout budget after discovery. * chore: re-trigger CI (flaky api coverage lane) * test: pin the clock for the remaining-timeout assertion (flaky under coverage instrumentation) * test(abort-r1): pin the clock in the LiteLLM remaining-budget test * chore: trigger CodeRabbit re-review (findings already addressed) * test(abort-r1): pin the clock in the lite-llm timeout passthrough test * fix(api): enforce the absolute request deadline through LiteLLM SDK retries The completePrompt timeout was 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. Enforce one absolute deadline (start + timeoutMs) through the completion call as well: a deadline signal merged with the caller's signal aborts in-flight attempts and makes the SDK short-circuit its next attempt at the pre-flight signal check, and a rejectOnAbort race settles the caller at the deadline even while the SDK is mid-retry-sleep. Retries stay enabled within the remaining budget. Timeout expiry stays distinct from caller cancellation: with the caller still running, deadline expiry surfaces as a TimeoutError (previously only used for discovery) instead of masquerading as the cancellation AbortError; a caller abort still yields the standard AbortError. Reply to the maintainer review on fork PR 1303 (lite-llm deadline). * refactor(api): move the Gemini HTTPS-only endpoint policy out of this rollout The cleartext-endpoint guard (assertSecureGeminiBaseUrl and its loopback exception) changes which custom endpoints are accepted independently of cancellation. Per the maintainer review on fork PR 1303, the policy and its CWE-319 security tests now live in their own PR (branch feat/gemini-https-endpoint-policy), keeping this rollout focused on abort-signal and timeout plumbing. No behavior change in this branch beyond dropping the endpoint validation, which this branch never shipped to main. * test(api): strengthen stream-failure assertions and timer cleanup; simplify deadline signal merge Address the re-review at head b94c13c: - gemini.spec: the two non-abort stream-failure tests only asserted "not Cannot read properties", which also passes when the wrapper is removed (raw rethrow) or when the catch misclassifies the failure as an abort. Pin the wrapping path by its observable shape instead: a NEW Error, a non-raw message, and a non-AbortError name. (The exact localized template string is not assertable in unit tests - i18next is not initialized with resources here, as the sibling telemetry test notes; the proposed t() equality compares undefined against "".) Negative control: raw-rethrow and abort-misclassification mutations each turn exactly these two tests red. - lite-llm.spec: restore real timers in a finally block so a failed assertion cannot leak fake timers into the deadline and discovery-timeout tests. - lite-llm: collapse the deadline/caller signal branch into one expression. mergeAbortSignals returns the primary signal directly when no secondary exists, so the guarded branch was an equivalent ConditionalExpression mutant; the ternary's mutants are killed by the existing caller-only and deadline-only tests. No behavior change in production code. * chore(api): remove Stryker suppression comments from touched files Per repo policy (AGENTS.md: Stryker suppression comments are banned; the mutation gate rejects them in touched files). Equivalence arguments for the resulting mutants are recorded in the PR body. * fix(api): reject non-finite and cap overflowing mergeAbortSignalAndTimeout delays Port of the #1300 timeout guard finding (CodeRabbit): Node's AbortSignal.timeout throws ERR_OUT_OF_RANGE for non-integer delays (Infinity) and clamps delays above 2^31-1 ms to 1 ms, so a non-finite timeoutMs threw a RangeError instead of running the request, and an overflowing one aborted almost immediately. The helper now treats non-finite values as "no timeout" and caps finite delays at 2^31-1 ms. Regression tests added; negative controls (reverting the guard) fail exactly the two new tests. * fix(api): normalize timeoutMs before constructing deadline signals CodeRabbit threads on a1cb932: - mergeAbortSignalAndTimeout rounds fractional delays up before AbortSignal.timeout's uint32 validation (1500.5/0.5 threw ERR_OUT_OF_RANGE for Mistral completePrompt, LiteLLM discovery, RequestConfigBuilder.addMergedSignal and Bedrock); regression test for timeoutMs 1.5. - getRequestTimeoutMs now rejects non-finite values and rounds up, and the LiteLLM completion flow clamps remainingMs to [1, 2^31-1] whole milliseconds: fractional values no longer throw from the deadline signal, Infinity disables the discovery timeout instead of throwing outside the completion try, and an overflowing value no longer clamps to a 1 ms deadline. The normalized value is what reaches createOptions.timeout. Regression test covers all three values; reverting the normalization fails exactly that test. * test(api): assert the rounded delay reaching AbortSignal.timeout CodeRabbit: the fractional-timeout test only proved the call did not throw and eventually aborted - a round-down mutation (1.5 -> 1) would still pass. The test now spies AbortSignal.timeout and pins the received delay: 2 for 1.5 and 1 for 0.5 (never 0). Negative control: mutating Math.ceil to Math.floor fails exactly this test. * fix(lite-llm): settle caller cancellation during a non-signal-aware SDK retry sleep The completion settle race watched only the deadline signal, so a caller cancelling while the OpenAI SDK slept between retries waited for the sleep to end (or for the deadline). The race now watches the merged caller+deadline signal; the catch still classifies deadline expiry as a TimeoutError first, so timeout and cancellation stay distinguishable. * chore(lite-llm): derive the deadline signal from the shared timeout helper The lint rule forbids AbortSignal.timeout on a variable delay (it throws on Infinity/fractions and clamps large values to 1 ms); mergeAbortSignalAndTimeout applies the series' capping and finiteness rules instead. Behaviour is unchanged for the remainingMs values this branch produces. * fix(api): report a per-request timeout as TimeoutError, not as cancellation completePrompt mapped every abort-shaped SDK failure to AbortError, so an HTTP timeout expiry looked like a user stop. Distinguish them: a live caller signal plus a configured timeout means the deadline elapsed. Moved the shared createTimeoutError into utils/request-timeout so the LiteLLM unit in this PR and Gemini use one definition, and strengthened the two non-abort completePrompt assertions to pin the wrapping path (instance, non-AbortError name, non-raw message). * fix(providers): use the shared timeout helpers in gemini and lite-llm The abort-signal hygiene lint gate rejects a raw timeoutMs in a timeout property and a hand-written 2^31-1 cap. gemini now forwards hasUsableTimeout/toSdkTimeoutMs (rounded up, capped) and lite-llm uses toSdkTimeoutMs for the remaining deadline budget. Lint passes on both files; provider suites 204 passed. * test(mistral): cover iterator failures raised during the streaming loop Review follow-up: the changed createMessage catch wraps the whole for-await loop, so iterator failures now map to the provider error. Adds focused tests for an Error and a non-Error value thrown by the returned async iterator mid-stream: both assert the Mistral completion error message and that telemetry captures the original message exactly once. mistral.spec.ts 46 passed. * fix(mistral): classify timeouts and merged-signal aborts Review follow-ups: completePrompt distinguishes a per-request timeout (createTimeoutError) from a caller abort and pins it with a fake-timer RequestTimeoutError test; createMessage checks the merged request signal so a plain error after cancellation surfaces as AbortError; the wider stream try is documented. mistral.spec 48 passed. * test(mistral): build typed createMessage metadata in the abort test * fix(mistral): classify the first abort cause and type the timeout mock Review follow-ups: use the merged signal reason so a timeout-first abort stays TimeoutError; add ordering tests; drop the timeout mock cast. mistral.spec 50 passed. * test(mistral): reject the request only after both abort signals are set Review follow-up: capture the reject callback, advance the timeout and abort the caller before rejecting, so the timeout-first classification assertion is meaningful. mistral.spec 50 passed. * chore(lint): ban hand-built AbortError and new @ts-ignore in api specs --------- Co-authored-by: Eason Liang <[email protected]> Co-authored-by: easonLiangWorldedtech <[email protected]> Co-authored-by: dsh-agent <dsh@local> Co-authored-by: Elliott de Launay <[email protected]>
1 parent 16fced1 commit 159c42e

14 files changed

Lines changed: 3017 additions & 133 deletions

‎packages/config-eslint/api-provider-hygiene.js‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,12 +42,24 @@ export const apiProviderHygieneSelectors = [
4242
},
4343
]
4444

45+
// Hand-built AbortError values drift from the shared contract (name, message). Use createAbortError.
46+
export const abortErrorSelectors = [
47+
{
48+
selector: "AssignmentExpression[left.property.name='name'][right.value='AbortError']",
49+
message: "Use createAbortError from providers/utils/abort-signal instead of building an AbortError by hand.",
50+
},
51+
{
52+
selector: "NewExpression[callee.name='DOMException'][arguments.1.value='AbortError']",
53+
message: "Use createAbortError from providers/utils/abort-signal instead of building an AbortError by hand.",
54+
},
55+
]
56+
4557
export const apiProviderHygieneConfig = {
4658
files: ["api/**/*.ts"],
4759
// The helper module owns the shared names. Tests build their own fixtures.
4860
ignores: ["api/providers/utils/abort-signal.ts", "**/__tests__/**", "**/*.test.ts", "**/*.spec.ts"],
4961
rules: {
50-
"no-restricted-syntax": ["error", ...apiProviderHygieneSelectors],
62+
"no-restricted-syntax": ["error", ...apiProviderHygieneSelectors, ...abortErrorSelectors],
5163
},
5264
}
5365

@@ -59,3 +71,13 @@ export const apiProviderTestSelectors = [
5971
"Use it.each for repeated cases. A loop of expects in one test stops at the first failing case and does not name it.",
6072
},
6173
]
74+
75+
// A new @ts-ignore in a provider spec hides a private-member access. Use bracket notation or a typed double.
76+
// Ratchet: files that still hold one are listed here. Remove a file from the list when it is clean.
77+
export const apiSpecTsIgnoreConfig = {
78+
files: ["api/**/__tests__/**/*.spec.ts"],
79+
ignores: ["api/providers/__tests__/gemini-handler.spec.ts", "api/providers/__tests__/vscode-lm.spec.ts"],
80+
rules: {
81+
"@typescript-eslint/ban-ts-comment": ["error", { "ts-ignore": true }],
82+
},
83+
}

‎packages/config-eslint/api-provider-hygiene.test.js‎

Lines changed: 72 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { describe, it } from "node:test"
44
import { Linter } from "eslint"
55
import typescriptEslint from "typescript-eslint"
66

7-
import { apiProviderHygieneConfig, apiProviderTestSelectors } from "./api-provider-hygiene.js"
7+
import { apiProviderHygieneConfig, apiProviderTestSelectors, apiSpecTsIgnoreConfig } from "./api-provider-hygiene.js"
88

99
const linter = new Linter({ configType: "flat" })
1010
const config = [{ files: ["**/*.ts"], languageOptions: { parser: typescriptEslint.parser } }, apiProviderHygieneConfig]
@@ -174,3 +174,74 @@ describe("api provider hygiene: spec loops", () => {
174174
})
175175
}
176176
})
177+
178+
describe("api provider hygiene: hand-built AbortError", () => {
179+
for (const code of [
180+
'const e = new Error("x"); e.name = "AbortError"',
181+
'throw new DOMException("aborted", "AbortError")',
182+
'Object.assign(new Error("x"), {}).name = "AbortError"',
183+
]) {
184+
it(`flags: ${code}`, () => {
185+
assert.equal(lint(code).length, 1)
186+
})
187+
}
188+
189+
for (const code of [
190+
'import { createAbortError } from "./utils/abort-signal"\nthrow createAbortError("X")',
191+
'if (error.name === "AbortError") {}',
192+
'const e = new Error("x"); e.name = "TimeoutError"',
193+
'new DOMException("timed out", "TimeoutError")',
194+
]) {
195+
it(`allows: ${code}`, () => {
196+
assert.deepEqual(lint(code), [])
197+
})
198+
}
199+
200+
it("does not apply to the helper module", () => {
201+
assert.deepEqual(
202+
lint('const e = new Error("x"); e.name = "AbortError"', "api/providers/utils/abort-signal.ts"),
203+
[],
204+
)
205+
})
206+
207+
it("does not apply to tests", () => {
208+
assert.deepEqual(lint('new DOMException("x", "AbortError")', "api/providers/__tests__/a.spec.ts"), [])
209+
})
210+
})
211+
212+
describe("api provider hygiene: @ts-ignore in specs", () => {
213+
const specLinter = new Linter({ configType: "flat" })
214+
const specConfig = [
215+
{
216+
files: ["**/*.ts"],
217+
languageOptions: { parser: typescriptEslint.parser },
218+
plugins: { "@typescript-eslint": typescriptEslint.plugin },
219+
},
220+
apiSpecTsIgnoreConfig,
221+
]
222+
const lintSpecFile = (code, filename = "api/providers/__tests__/example.spec.ts") =>
223+
specLinter.verify(code, specConfig, { filename }).map((m) => m.message)
224+
225+
it("flags @ts-ignore in a spec", () => {
226+
assert.equal(lintSpecFile("// @ts-ignore access private client\nhandler.client").length, 1)
227+
})
228+
229+
it("allows @ts-expect-error with a description", () => {
230+
assert.deepEqual(lintSpecFile("// @ts-expect-error private client\nhandler.client"), [])
231+
})
232+
233+
it("allows bracket notation", () => {
234+
assert.deepEqual(lintSpecFile('handler["client"]'), [])
235+
})
236+
237+
it("skips files in the ratchet list", () => {
238+
assert.deepEqual(
239+
lintSpecFile("// @ts-ignore x\nhandler.client", "api/providers/__tests__/gemini-handler.spec.ts"),
240+
[],
241+
)
242+
})
243+
244+
it("does not apply outside src/api", () => {
245+
assert.deepEqual(lintSpecFile("// @ts-ignore x\nhandler.client", "services/__tests__/example.spec.ts"), [])
246+
})
247+
})

‎src/api/providers/__tests__/gemini-handler.spec.ts‎

Lines changed: 50 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,7 @@ vi.mock("@roo-code/telemetry", () => ({
1212
}))
1313

1414
import { GeminiHandler } from "../gemini"
15-
import type { ApiHandlerOptions } from "../../../shared/api"
16-
import { providerIdentifiers } from "@roo-code/types/provider-identifiers"
15+
import { makeApiHandlerOptions } from "../../../test-utils/api"
1716

1817
describe("GeminiHandler backend support", () => {
1918
beforeEach(() => {
@@ -24,11 +23,7 @@ describe("GeminiHandler backend support", () => {
2423
// URL context and grounding are mutually exclusive with function declarations
2524
// in Gemini API, so createMessage only uses function declarations.
2625
// URL context/grounding are only added in completePrompt.
27-
const options = {
28-
apiProvider: providerIdentifiers.gemini,
29-
enableUrlContext: true,
30-
enableGrounding: true,
31-
} as ApiHandlerOptions
26+
const options = makeApiHandlerOptions()
3227
const handler = new GeminiHandler(options)
3328
const stub = vi.fn().mockReturnValue((async function* () {})())
3429
// @ts-ignore access private client
@@ -41,11 +36,7 @@ describe("GeminiHandler backend support", () => {
4136
})
4237

4338
it("completePrompt passes config overrides without tools when URL context and grounding disabled", async () => {
44-
const options = {
45-
apiProvider: providerIdentifiers.gemini,
46-
enableUrlContext: false,
47-
enableGrounding: false,
48-
} as ApiHandlerOptions
39+
const options = makeApiHandlerOptions()
4940
const handler = new GeminiHandler(options)
5041
const stub = vi.fn().mockResolvedValue({ text: "ok" })
5142
// @ts-ignore access private client
@@ -56,12 +47,39 @@ describe("GeminiHandler backend support", () => {
5647
expect(promptConfig.tools).toBeUndefined()
5748
})
5849

50+
it("completePrompt should pass abort signal through to client via config.abortSignal", async () => {
51+
const options = makeApiHandlerOptions()
52+
const handler = new GeminiHandler(options)
53+
54+
const controller = new AbortController()
55+
const stub = vi.fn().mockResolvedValue({ text: "response" })
56+
handler["client"].models.generateContent = stub
57+
58+
await handler.completePrompt("test prompt", { abortSignal: controller.signal })
59+
60+
expect(stub).toHaveBeenCalledWith(
61+
expect.objectContaining({
62+
config: expect.objectContaining({
63+
abortSignal: controller.signal,
64+
}),
65+
}),
66+
)
67+
})
68+
69+
it("completePrompt should work without options (backward compatible)", async () => {
70+
const options = makeApiHandlerOptions()
71+
const handler = new GeminiHandler(options)
72+
73+
const stub = vi.fn().mockResolvedValue({ text: "response" })
74+
handler["client"].models.generateContent = stub
75+
76+
const result = await handler.completePrompt("test prompt")
77+
expect(result).toBe("response")
78+
})
79+
5980
describe("error scenarios", () => {
6081
it("should handle grounding metadata extraction failure gracefully", async () => {
61-
const options = {
62-
apiProvider: providerIdentifiers.gemini,
63-
enableGrounding: true,
64-
} as ApiHandlerOptions
82+
const options = makeApiHandlerOptions()
6583
const handler = new GeminiHandler(options)
6684

6785
const mockStream = async function* () {
@@ -93,10 +111,7 @@ describe("GeminiHandler backend support", () => {
93111
})
94112

95113
it("should handle malformed grounding metadata", async () => {
96-
const options = {
97-
apiProvider: providerIdentifiers.gemini,
98-
enableGrounding: true,
99-
} as ApiHandlerOptions
114+
const options = makeApiHandlerOptions()
100115
const handler = new GeminiHandler(options)
101116

102117
const mockStream = async function* () {
@@ -144,11 +159,7 @@ describe("GeminiHandler backend support", () => {
144159
})
145160

146161
it("should handle API errors when tools are enabled", async () => {
147-
const options = {
148-
apiProvider: providerIdentifiers.gemini,
149-
enableUrlContext: true,
150-
enableGrounding: true,
151-
} as ApiHandlerOptions
162+
const options = makeApiHandlerOptions()
152163
const handler = new GeminiHandler(options)
153164

154165
const mockError = new Error("API rate limit exceeded")
@@ -192,9 +203,7 @@ describe("GeminiHandler backend support", () => {
192203
]
193204

194205
it("should ignore allowedFunctionNames because Gemini rejects larger restriction lists", async () => {
195-
const options = {
196-
apiProvider: providerIdentifiers.gemini,
197-
} as ApiHandlerOptions
206+
const options = makeApiHandlerOptions()
198207
const handler = new GeminiHandler(options)
199208
const stub = vi.fn().mockReturnValue((async function* () {})())
200209
// @ts-ignore access private client
@@ -213,9 +222,7 @@ describe("GeminiHandler backend support", () => {
213222
})
214223

215224
it("should include all tools when allowedFunctionNames is provided", async () => {
216-
const options = {
217-
apiProvider: providerIdentifiers.gemini,
218-
} as ApiHandlerOptions
225+
const options = makeApiHandlerOptions()
219226
const handler = new GeminiHandler(options)
220227
const stub = vi.fn().mockReturnValue((async function* () {})())
221228
// @ts-ignore access private client
@@ -236,9 +243,7 @@ describe("GeminiHandler backend support", () => {
236243
})
237244

238245
it("should not pass large allowedFunctionNames lists to Gemini", async () => {
239-
const options = {
240-
apiProvider: providerIdentifiers.gemini,
241-
} as ApiHandlerOptions
246+
const options = makeApiHandlerOptions()
242247
const handler = new GeminiHandler(options)
243248
const stub = vi.fn().mockReturnValue((async function* () {})())
244249
// @ts-ignore access private client
@@ -267,9 +272,7 @@ describe("GeminiHandler backend support", () => {
267272
})
268273

269274
it("should not pass allowedFunctionNames even when history includes tool calls", async () => {
270-
const options = {
271-
apiProvider: providerIdentifiers.gemini,
272-
} as ApiHandlerOptions
275+
const options = makeApiHandlerOptions()
273276
const handler = new GeminiHandler(options)
274277
const stub = vi.fn().mockReturnValue((async function* () {})())
275278
// @ts-ignore access private client
@@ -304,9 +307,7 @@ describe("GeminiHandler backend support", () => {
304307
})
305308

306309
it("should fall back to tool_choice when allowedFunctionNames is provided", async () => {
307-
const options = {
308-
apiProvider: providerIdentifiers.gemini,
309-
} as ApiHandlerOptions
310+
const options = makeApiHandlerOptions()
310311
const handler = new GeminiHandler(options)
311312
const stub = vi.fn().mockReturnValue((async function* () {})())
312313
// @ts-ignore access private client
@@ -327,9 +328,7 @@ describe("GeminiHandler backend support", () => {
327328
})
328329

329330
it("should fall back to tool_choice when allowedFunctionNames is empty", async () => {
330-
const options = {
331-
apiProvider: providerIdentifiers.gemini,
332-
} as ApiHandlerOptions
331+
const options = makeApiHandlerOptions()
333332
const handler = new GeminiHandler(options)
334333
const stub = vi.fn().mockReturnValue((async function* () {})())
335334
// @ts-ignore access private client
@@ -351,9 +350,7 @@ describe("GeminiHandler backend support", () => {
351350
})
352351

353352
it("should not set toolConfig when allowedFunctionNames is undefined and no tool_choice", async () => {
354-
const options = {
355-
apiProvider: providerIdentifiers.gemini,
356-
} as ApiHandlerOptions
353+
const options = makeApiHandlerOptions()
357354
const handler = new GeminiHandler(options)
358355
const stub = vi.fn().mockReturnValue((async function* () {})())
359356
// @ts-ignore access private client
@@ -374,9 +371,7 @@ describe("GeminiHandler backend support", () => {
374371

375372
describe("Gemini schema compatibility", () => {
376373
it("should strip broad JSON Schema metadata from function declarations", async () => {
377-
const options = {
378-
apiProvider: providerIdentifiers.gemini,
379-
} as ApiHandlerOptions
374+
const options = makeApiHandlerOptions()
380375
const handler = new GeminiHandler(options)
381376
const stub = vi.fn().mockReturnValue((async function* () {})())
382377
// @ts-ignore access private client
@@ -435,9 +430,7 @@ describe("GeminiHandler backend support", () => {
435430
})
436431

437432
it("should collapse composition and type arrays in function declaration schemas", async () => {
438-
const options = {
439-
apiProvider: providerIdentifiers.gemini,
440-
} as ApiHandlerOptions
433+
const options = makeApiHandlerOptions()
441434
const handler = new GeminiHandler(options)
442435
const stub = vi.fn().mockReturnValue((async function* () {})())
443436
// @ts-ignore access private client
@@ -496,7 +489,7 @@ describe("GeminiHandler backend support", () => {
496489
})
497490

498491
it("should deep-merge allOf fragments instead of overwriting earlier properties", async () => {
499-
const options = { apiProvider: providerIdentifiers.gemini } as ApiHandlerOptions
492+
const options = makeApiHandlerOptions()
500493
const handler = new GeminiHandler(options)
501494
const stub = vi.fn().mockReturnValue((async function* () {})())
502495
// @ts-ignore access private client
@@ -541,7 +534,7 @@ describe("GeminiHandler backend support", () => {
541534
})
542535

543536
it("should resolve $ref entries before dropping $defs", async () => {
544-
const options = { apiProvider: providerIdentifiers.gemini } as ApiHandlerOptions
537+
const options = makeApiHandlerOptions()
545538
const handler = new GeminiHandler(options)
546539
const stub = vi.fn().mockReturnValue((async function* () {})())
547540
// @ts-ignore access private client
@@ -590,7 +583,7 @@ describe("GeminiHandler backend support", () => {
590583
})
591584

592585
it("should preserve top-level properties and required entries when allOf is also present", async () => {
593-
const options = { apiProvider: providerIdentifiers.gemini } as ApiHandlerOptions
586+
const options = makeApiHandlerOptions()
594587
const handler = new GeminiHandler(options)
595588
const stub = vi.fn().mockReturnValue((async function* () {})())
596589
// @ts-ignore access private client
@@ -632,7 +625,7 @@ describe("GeminiHandler backend support", () => {
632625
})
633626

634627
it("should stop recursive $ref expansion before the sanitized schema becomes cyclic", async () => {
635-
const options = { apiProvider: providerIdentifiers.gemini } as ApiHandlerOptions
628+
const options = makeApiHandlerOptions()
636629
const handler = new GeminiHandler(options)
637630
const stub = vi.fn().mockReturnValue((async function* () {})())
638631
// @ts-ignore access private client
@@ -684,7 +677,7 @@ describe("GeminiHandler backend support", () => {
684677
})
685678

686679
it("should preserve parameter names that collide with stripped schema keywords", async () => {
687-
const options = { apiProvider: providerIdentifiers.gemini } as ApiHandlerOptions
680+
const options = makeApiHandlerOptions()
688681
const handler = new GeminiHandler(options)
689682
const stub = vi.fn().mockReturnValue((async function* () {})())
690683
// @ts-ignore access private client

0 commit comments

Comments
 (0)