feat(api): abort signal for vscode-lm, gemini, mistral (SDK-specific quirks)#732
Draft
easonLiangWorldedtech wants to merge 16 commits into
Draft
Conversation
…nfiguration Body: Implement generic request configuration builder with chainable methods (addAbortSignal, addHeaders, setOption), static factory methods (fromMetadata, mergeAbortSignals), and 40 unit tests.
… and multi-SDK examples
…ls early-abort - Fix README TOC: change #how-mergesignals-works to #how-mergeabortsignals-works to match the actual heading anchor - Simplify mergeAbortSignals: return primarySignal directly when it's already aborted instead of creating a new AbortController
…onfigBuilder (Zoo-Code-Org#615) - Add default empty object parameter to addHeaders() so calling with undefined no longer throws TypeError from Object.keys(undefined) - Reorder mergeAbortSignals to check primarySignal.aborted before allocating AbortController, preventing unnecessary controller creation
…roviders - anthropic.spec.ts: timeout trigger test + signal instance verification - native-ollama.spec.ts: explicit assertion no second arg passed (no signal forwarding) - openai-codex-native-tool-calls.spec.ts: removed weak toHaveProperty(signal) assertion - vercel-ai-gateway.spec.ts: temperature test uses correct undefined second arg fix: add timeoutMs forwarding to completePrompt methods - vercel-ai-gateway.ts: use Object.keys(createOptions).length > 0 check instead of truthy check - poe.ts: merge signal and timeoutMs properly with combined abort logic - config-builder/README.md: update documentation for mergeAbortSignals behavior test: add timeoutMs coverage for poe, moonshot, minimax, mistral, xai providers - poe.spec.ts: signal+timeoutMs merge, timeoutMs only, timeoutMs=0 cases - moonshot.spec.ts: same timeoutMs tests for openai-compatible pattern - minimax.spec.ts: signal+timeoutMs, timeoutMs only, truthy check behavior - mistral.spec.ts: same timeoutMs coverage - xai.spec.ts: signal+timeoutMs, timeoutMs only, truthy check behavior test: add timeoutMs coverage for anthropic-vertex, base-openai-compatible, bedrock, openai-native - anthropic-vertex.spec.ts: signal passing test (no timeoutMs support) - base-openai-compatible-provider-timeout.spec.ts: completePrompt with signal+timeoutMs, timeoutMs only, truthy check behavior - bedrock.spec.ts: timeoutMs coverage for adaptive thinking path - openai-native.spec.ts: signal and timeoutMs merging tests test: add signal+timeoutMs merge tests for fireworks, lite-llm, lmstudio - fireworks.spec.ts: added merge signal and timeoutMs together test - lite-llm.spec.ts: added same combined signal+timeoutMs test - lmstudio.spec.ts: aligned with same abort signal pattern fix: replace tautological assertion in poe.spec.ts - poe.spec.ts: completePrompt should prefer signal over timeoutMs test now asserts abortSignal is a distinct AbortSignal from controller.signal instead of always-true instanceof check test: add missing error catch and timeout cleanup tests - vscode-lm.spec.ts: added 'should handle errors in completePrompt' test - poe.spec.ts: added 'completePrompt should clear timeout when user signal aborts' test - opencode-go.spec.ts: added OpenAI path completePrompt tests (signal, timeoutMs, merged)
…tial-request test
…-code, lmstudio-native
easonLiangWorldedtech
requested review from
JamesRobert20,
edelauna,
hannesrudolph,
navedmerchant and
taltas
as code owners
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds ChangesCompletion option plumbing
Model cache export
Sequence Diagram(s)Estimated review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related issues
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 3❌ Failed checks (3 warnings)
✅ Passed checks (2 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 20
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/api/providers/openai-native.ts (1)
1487-1563: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
timeoutMsis accepted here but never applied.This method now takes
CompletePromptOptions, but the new logic only threadsabortSignalinto the request. Any caller passing{ timeoutMs }gets no per-call deadline on OpenAI Native, which breaks the sharedcompletePromptcontract.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/api/providers/openai-native.ts` around lines 1487 - 1563, `completePrompt` is merging only `abortSignal` and ignoring the `timeoutMs` from `CompletePromptOptions`, so per-call deadlines are not enforced. Update `completePrompt` in `openai-native.ts` to apply `timeoutMs` by creating/merging a timeout-driven abort signal with the existing class-level controller and the caller’s signal, then pass that merged signal into `responses.create` and keep the cleanup behavior consistent with the current `RequestConfigBuilder.mergeAbortSignals` flow.
🧹 Nitpick comments (2)
src/utils/__tests__/enhance-prompt.spec.ts (1)
40-68: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd one test for explicit options forwarding.
These assertions only cover the
undefinedpath. The newsingleCompletionHandler(..., options)branch should have a package-local unit test that verifiesabortSignalandtimeoutMsare passed through unchanged tobuildApiHandler(...).completePrompt(...).Proposed test
+ it("forwards completion options to the handler", async () => { + const controller = new AbortController() + const options = { abortSignal: controller.signal, timeoutMs: 5000 } + + await singleCompletionHandler(mockApiConfig, "Test prompt", options) + + const handler = buildApiHandler(mockApiConfig) + expect((handler as any).completePrompt).toHaveBeenCalledWith("Test prompt", options) + })As per coding guidelines, "Use package-local unit tests for pure logic, parsing, state transitions, validation, serialization, request construction, retry decisions, and error handling."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/utils/__tests__/enhance-prompt.spec.ts` around lines 40 - 68, Add a package-local unit test in enhance-prompt.spec.ts for the new singleCompletionHandler(..., options) path. Verify that when options are provided, singleCompletionHandler forwards abortSignal and timeoutMs unchanged into buildApiHandler(...).completePrompt(...), instead of only covering the undefined-options case. Use the existing singleCompletionHandler and buildApiHandler mocks to assert the call shape clearly.Source: Coding guidelines
src/api/providers/__tests__/vscode-lm.spec.ts (1)
415-420: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThis never exercises the post-start abort path.
Line 416 aborts the controller before
createMessage()is invoked, so this spec hits the already-aborted branch again instead of verifying the listener bridge described in the test name. Abort after the request has started if you want this to cover the cancellation state transition. As per coding guidelines,**/*.{test,spec}.{ts,tsx,js}: Use package-local unit tests for pure logic, parsing, state transitions, validation, serialization, request construction, retry decisions, and error handling.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/api/providers/__tests__/vscode-lm.spec.ts` around lines 415 - 420, The abort test is using an already-aborted signal before `createMessage()` starts, so it does not cover the post-start cancellation path named by the spec. Update the `vscode-lm.spec.ts` test around `handler.createMessage()` to start the request first and then abort the `AbortController` so it exercises the listener bridge/state transition instead of the pre-aborted branch. Keep the change within this package-local unit test and use the existing `controller`, `abortSignal`, and `createMessage` setup to verify the intended cancellation behavior.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/api/index.ts`:
- Around line 47-53: The shared completion contract is exposing
CompletePromptOptions.timeoutMs through SingleCompletionHandler before all
providers actually support it. Update the provider implementations referenced by
CompletePromptOptions and SingleCompletionHandler, especially anthropic-vertex
and bedrock, to honor the timeout override when completing prompts, or otherwise
narrow/remove timeoutMs from the exported API until every provider consistently
forwards it alongside abortSignal.
In `@src/api/providers/__tests__/anthropic.spec.ts`:
- Around line 522-551: The Anthropic timeout test does not verify that the
timeout path actually fired because timeoutTriggered is never asserted. Update
the test in AnthropicHandler/completePrompt to explicitly await or assert the
rejection path, and add a check like expect(timeoutTriggered).toBe(true) after
the wait so the test fails if the request only stays pending. Make sure the
assertion still uses the existing mockCreateWithTimeout, controller.signal, and
timeoutMs setup.
In `@src/api/providers/__tests__/bedrock.spec.ts`:
- Around line 1629-1651: The test in AwsBedrockHandler.completePrompt currently
only asserts that send() was called, so it does not verify that timeoutMs is
actually propagated. Update the spec to inspect the arguments passed to the
mocked client send() call, or assert the derived AbortSignal/abort behavior, so
the test proves the timeout option is used by completePrompt rather than just
reaching the send path.
In `@src/api/providers/__tests__/request-config-builder.spec.ts`:
- Around line 65-77: The RequestConfigBuilder replacement test is seeding the
wrong property, so it never proves that addAbortSignal actually overrides an
existing signal. Update the test setup to initialize the builder with the same
signal field that RequestConfigBuilder.build() uses, then verify addAbortSignal
replaces that existing signal with controller2.signal in the
RequestConfigBuilder/addAbortSignal flow.
In `@src/api/providers/__tests__/vscode-lm.spec.ts`:
- Around line 473-478: The test around handler["currentRequestCancellation"] is
asserting the old leaked behavior by expecting the CancellationTokenSource to
remain defined and dispose() not to be called. Update the vscode-lm.spec.ts
expectations to match the intended cleanup in the request handler/finally path:
verify that the cancellation source is disposed after completion and no longer
retained, using the relevant handler/currentRequestCancellation and dispose
assertions.
In `@src/api/providers/config-builder/README.md`:
- Around line 51-55: The markdown links for RequestConfigBuilder currently point
to the wrong relative path and resolve one directory too high. Update the source
links in the table to reference RequestConfigBuilder from this README’s location
so they correctly target request-config-builder.ts, including both the Generic
Methods and Static Methods entries.
In `@src/api/providers/fetchers/__tests__/modelCache.spec.ts`:
- Around line 474-485: The writeModels test only checks that the promise
resolves, so it does not verify the request construction performed by
writeModels(). Update the test in writeModels to assert the mocked safeWriteJson
call was made with the expected cache path and serialized payload derived from
mockModels and the provider name, using the existing safeWriteJson mock and the
writeModels helper as the main symbols to locate the code.
In `@src/api/providers/gemini.ts`:
- Around line 595-607: The GenerateContentConfig assembly in gemini.ts is
putting the cancellation token into httpOptions.signal, but `@google/genai`
expects it on config.abortSignal. Update the prompt/config construction so the
abort signal from options.abortSignal is assigned to
GenerateContentConfig.abortSignal directly, and keep httpOptions only for
transport settings like timeout and baseUrl; use the existing promptConfig
building area to make this change.
In `@src/api/providers/lite-llm.ts`:
- Around line 337-344: The request options builder in lite-llm currently uses a
truthy check for timeoutMs, which skips the intentional sentinel value 0. Update
the createOptions logic in the request setup so it checks options.timeoutMs
explicitly against undefined before assigning createOptions.timeout, while
keeping the existing abortSignal handling unchanged. Use the createOptions block
in the same function as the reference point.
In `@src/api/providers/minimax.ts`:
- Around line 300-302: The timeout handling in minimax request setup is treating
0 as missing, so an explicit zero-ms override is ignored. Update the timeout
assignment logic in the code path that builds requestOptions so it checks for
presence of options.timeoutMs rather than truthiness, and preserve zero-valued
timeouts when passed through this path.
In `@src/api/providers/openai-codex.ts`:
- Around line 1156-1160: The completePrompt path in OpenAICodexProvider is still
ignoring the timeoutMs override from CompletePromptOptions, so callers can end
up with an unbounded fetch. Update completePrompt (and the matching later call
site in the same provider) to derive a timeout-backed AbortSignal from
options.timeoutMs and merge it with the existing abort signals via
RequestConfigBuilder.mergeAbortSignals, keeping the shared completion-options
contract intact.
- Around line 1156-1166: Keep completePrompt() from mutating the shared
this.abortController, since createMessage() relies on it for live stream
handling. Make the abort controller in completePrompt() request-local, wire
mergedSignal to that local controller, and pass it through the fetch path
without reassigning the instance field or touching the createMessage()
controller lifecycle in finally. Also apply the same local-controller cleanup to
the related completePrompt abort handling near the later referenced section.
In `@src/api/providers/openai-compatible.ts`:
- Around line 212-233: The merged abort handling in openai-compatible.ts ignores
an explicit timeoutMs of 0 because both timeout branches require timeoutMs > 0.
Update the abort-signal setup around the merged AbortController logic so
timeoutMs=0 is treated as a valid immediate deadline instead of falling through
to “no timeout,” while preserving the existing behavior for abortSignal-only and
positive timeout cases. Reference the options.abortSignal / timeoutMs checks and
the generateOptions.abortSignal assignment when making the change.
In `@src/api/providers/openrouter.ts`:
- Around line 599-605: The requestOptions builder in openrouter.ts is dropping
an explicit timeout override because it uses a truthy check on
options.timeoutMs, so completePrompt(..., { timeoutMs: 0 }) falls back to the
default timeout. Update the OpenAI.RequestOptions assembly in the requestOptions
block to treat timeoutMs as present when it is 0, matching the existing contract
used by the shared OpenAI-compatible path, while keeping the abortSignal and
headers merge behavior unchanged.
In `@src/api/providers/requesty.ts`:
- Around line 223-225: The Requesty timeout handling in the provider should
match the other handlers by treating non-positive values as unset. Update the
check around createOptions.timeout in the requesty.ts logic so it only assigns
options.timeoutMs when it is greater than 0, and leave it unchanged for 0 or
negative values to keep CompletePromptOptions behavior consistent.
In `@src/api/providers/unbound.ts`:
- Around line 212-214: The timeout override in the unbound provider drops an
explicit 0 value because the `options.timeoutMs` check is truthy, so update the
timeout assignment in `unbound.ts` to preserve `timeoutMs=0` by checking for
`undefined` instead of truthiness. Make the change in the same
`createOptions.timeout` setup used by the provider’s request path so the
caller-provided timeout is forwarded exactly, matching the sibling
OpenAI-compatible handling.
In `@src/api/providers/vercel-ai-gateway.ts`:
- Around line 141-142: The timeout override in vercel-ai-gateway’s createOptions
handling is using a truthy check, so completePrompt and similar paths ignore an
explicit timeoutMs: 0. Update the condition in vercel-ai-gateway.ts to match the
OpenAI-compatible path by checking for undefined instead of truthiness, keeping
zero as a valid override value. Use the existing createOptions logic and any
shared timeout handling around options.timeoutMs as the place to fix it.
In `@src/api/providers/vscode-lm.ts`:
- Around line 387-398: The per-request CancellationTokenSource created in the
vscode-lm request flow is only cleaned up on error, so successful requests leave
currentRequestCancellation pointing at a completed token and leak the source.
Update the request handling around the tokenSource setup in vscode-lm to use a
try/finally around the request body, and in the finally block clear
currentRequestCancellation and dispose the tokenSource on every exit path,
including successful streams.
In `@src/api/providers/xai.ts`:
- Around line 154-155: The timeout override in xai completePrompt is using a
truthy check, so timeoutMs: 0 is ignored and the default timeout remains in
effect. Update the requestOptions timeout assignment in completePrompt to check
for the field being defined rather than truthy, matching the shared OpenAI-style
timeout handling so an explicit zero value is forwarded correctly.
In `@src/api/providers/zoo-gateway.ts`:
- Around line 302-303: The timeout forwarding branch in zoo-gateway’s request
setup is only checking for truthy values, so negative timeoutMs values can still
be passed through as if they were valid. Update the timeout handling around the
createOptions assignment to ignore any non-positive timeoutMs and only forward
values greater than 0, preserving the existing “no timeout” behavior for invalid
inputs.
---
Outside diff comments:
In `@src/api/providers/openai-native.ts`:
- Around line 1487-1563: `completePrompt` is merging only `abortSignal` and
ignoring the `timeoutMs` from `CompletePromptOptions`, so per-call deadlines are
not enforced. Update `completePrompt` in `openai-native.ts` to apply `timeoutMs`
by creating/merging a timeout-driven abort signal with the existing class-level
controller and the caller’s signal, then pass that merged signal into
`responses.create` and keep the cleanup behavior consistent with the current
`RequestConfigBuilder.mergeAbortSignals` flow.
---
Nitpick comments:
In `@src/api/providers/__tests__/vscode-lm.spec.ts`:
- Around line 415-420: The abort test is using an already-aborted signal before
`createMessage()` starts, so it does not cover the post-start cancellation path
named by the spec. Update the `vscode-lm.spec.ts` test around
`handler.createMessage()` to start the request first and then abort the
`AbortController` so it exercises the listener bridge/state transition instead
of the pre-aborted branch. Keep the change within this package-local unit test
and use the existing `controller`, `abortSignal`, and `createMessage` setup to
verify the intended cancellation behavior.
In `@src/utils/__tests__/enhance-prompt.spec.ts`:
- Around line 40-68: Add a package-local unit test in enhance-prompt.spec.ts for
the new singleCompletionHandler(..., options) path. Verify that when options are
provided, singleCompletionHandler forwards abortSignal and timeoutMs unchanged
into buildApiHandler(...).completePrompt(...), instead of only covering the
undefined-options case. Use the existing singleCompletionHandler and
buildApiHandler mocks to assert the call shape clearly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 3c5bd7b7-6c3c-450e-94af-5628ae31d081
📒 Files selected for processing (70)
src/api/index.tssrc/api/providers/__tests__/anthropic-vertex.spec.tssrc/api/providers/__tests__/anthropic.spec.tssrc/api/providers/__tests__/base-openai-compatible-provider-timeout.spec.tssrc/api/providers/__tests__/bedrock.spec.tssrc/api/providers/__tests__/complete-prompt-options.spec.tssrc/api/providers/__tests__/deepseek.spec.tssrc/api/providers/__tests__/fireworks.spec.tssrc/api/providers/__tests__/gemini-abort-signal.spec.tssrc/api/providers/__tests__/gemini-handler.spec.tssrc/api/providers/__tests__/gemini.spec.tssrc/api/providers/__tests__/lite-llm.spec.tssrc/api/providers/__tests__/lmstudio-native-tools.spec.tssrc/api/providers/__tests__/lmstudio.spec.tssrc/api/providers/__tests__/mimo.spec.tssrc/api/providers/__tests__/minimax.spec.tssrc/api/providers/__tests__/mistral-abort-signal.spec.tssrc/api/providers/__tests__/mistral.spec.tssrc/api/providers/__tests__/moonshot.spec.tssrc/api/providers/__tests__/native-ollama.spec.tssrc/api/providers/__tests__/openai-codex-native-tool-calls.spec.tssrc/api/providers/__tests__/openai-native.spec.tssrc/api/providers/__tests__/openai.spec.tssrc/api/providers/__tests__/opencode-go.spec.tssrc/api/providers/__tests__/openrouter.spec.tssrc/api/providers/__tests__/poe.spec.tssrc/api/providers/__tests__/qwen-code-native-tools.spec.tssrc/api/providers/__tests__/request-config-builder.spec.tssrc/api/providers/__tests__/requesty.spec.tssrc/api/providers/__tests__/sambanova.spec.tssrc/api/providers/__tests__/unbound.spec.tssrc/api/providers/__tests__/vercel-ai-gateway.spec.tssrc/api/providers/__tests__/vertex.spec.tssrc/api/providers/__tests__/vscode-lm.spec.tssrc/api/providers/__tests__/xai.spec.tssrc/api/providers/__tests__/zai.spec.tssrc/api/providers/__tests__/zoo-gateway.spec.tssrc/api/providers/anthropic-vertex.tssrc/api/providers/anthropic.tssrc/api/providers/base-openai-compatible-provider.tssrc/api/providers/bedrock.tssrc/api/providers/config-builder/README.mdsrc/api/providers/config-builder/request-config-builder.tssrc/api/providers/fake-ai.tssrc/api/providers/fetchers/__tests__/modelCache.spec.tssrc/api/providers/fetchers/modelCache.tssrc/api/providers/gemini.tssrc/api/providers/index.tssrc/api/providers/lite-llm.tssrc/api/providers/lm-studio.tssrc/api/providers/minimax.tssrc/api/providers/mistral.tssrc/api/providers/native-ollama.tssrc/api/providers/openai-codex.tssrc/api/providers/openai-compatible.tssrc/api/providers/openai-native.tssrc/api/providers/openai.tssrc/api/providers/opencode-go.tssrc/api/providers/openrouter.tssrc/api/providers/poe.tssrc/api/providers/qwen-code.tssrc/api/providers/requesty.tssrc/api/providers/unbound.tssrc/api/providers/vercel-ai-gateway.tssrc/api/providers/vscode-lm.tssrc/api/providers/xai.tssrc/api/providers/zoo-gateway.tssrc/core/task/__tests__/Task.spec.tssrc/utils/__tests__/enhance-prompt.spec.tssrc/utils/single-completion-handler.ts
Comment on lines
+47
to
+53
| export interface CompletePromptOptions extends Pick<ApiHandlerCreateMessageMetadata, "abortSignal"> { | ||
| /** Optional timeout override (ms) — falls back to provider default if omitted */ | ||
| timeoutMs?: number | ||
| } | ||
|
|
||
| export interface SingleCompletionHandler { | ||
| completePrompt(prompt: string): Promise<string> | ||
| completePrompt(prompt: string, metadata?: CompletePromptOptions): Promise<string> |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Don't export timeoutMs as a shared completion contract until every provider honors it.
CompletePromptOptions.timeoutMs is now part of the public SingleCompletionHandler API, but src/api/providers/anthropic-vertex.ts (Line 299) and src/api/providers/bedrock.ts (Line 841) still ignore it and only forward abort state. That makes timeout behavior provider-dependent behind one shared interface. Either wire the override everywhere or narrow/document the contract before exporting it.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/api/index.ts` around lines 47 - 53, The shared completion contract is
exposing CompletePromptOptions.timeoutMs through SingleCompletionHandler before
all providers actually support it. Update the provider implementations
referenced by CompletePromptOptions and SingleCompletionHandler, especially
anthropic-vertex and bedrock, to honor the timeout override when completing
prompts, or otherwise narrow/remove timeoutMs from the exported API until every
provider consistently forwards it alongside abortSignal.
Comment on lines
+522
to
+551
| it("should trigger timeout when timeoutMs elapses before request completes", async () => { | ||
| const mockCreateWithTimeout = vitest | ||
| .fn() | ||
| .mockImplementation( | ||
| async () => | ||
| new Promise((resolve) => | ||
| setTimeout(() => resolve({ content: [{ type: "text", text: "response" }] }), 500), | ||
| ), | ||
| ) | ||
|
|
||
| const handlerTimeout = new AnthropicHandler(mockOptions) | ||
| // Replace the mock on the existing handler's client | ||
| handlerTimeout["client"].messages.create = mockCreateWithTimeout | ||
|
|
||
| const controller = new AbortController() | ||
| let timeoutTriggered = false | ||
| handlerTimeout | ||
| .completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 50 }) | ||
| .catch(() => { | ||
| timeoutTriggered = true | ||
| }) | ||
|
|
||
| // Wait for timeout to trigger (50ms timeout + buffer) | ||
| await new Promise((resolve) => setTimeout(resolve, 150)) | ||
|
|
||
| // Verify the API was called with timeout options | ||
| expect(mockCreateWithTimeout).toHaveBeenCalled() | ||
| // User signal should not be aborted (timeout mechanism aborts its own internal signal) | ||
| expect(controller.signal.aborted).toBe(false) | ||
| }) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
This timeout test never proves the timeout path ran.
timeoutTriggered is set in the rejection handler, but the test never asserts it. As written, this still passes if completePrompt() just stays pending past 150ms and never times out. Please assert the rejection (or at least expect(timeoutTriggered).toBe(true)) so the test actually covers timeout/error handling. As per coding guidelines, **/*.{test,spec}.{ts,tsx,js}: Use package-local unit tests for pure logic, parsing, state transitions, validation, serialization, request construction, retry decisions, and error handling.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/api/providers/__tests__/anthropic.spec.ts` around lines 522 - 551, The
Anthropic timeout test does not verify that the timeout path actually fired
because timeoutTriggered is never asserted. Update the test in
AnthropicHandler/completePrompt to explicitly await or assert the rejection
path, and add a check like expect(timeoutTriggered).toBe(true) after the wait so
the test fails if the request only stays pending. Make sure the assertion still
uses the existing mockCreateWithTimeout, controller.signal, and timeoutMs setup.
Source: Coding guidelines
Comment on lines
+1629
to
+1651
| it("completePrompt should pass timeoutMs through to client", async () => { | ||
| const mockSend = vi.fn() | ||
|
|
||
| const handler = new AwsBedrockHandler({ | ||
| apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0", | ||
| awsAccessKey: "test-access-key", | ||
| awsSecretKey: "test-secret-key", | ||
| awsRegion: "us-east-1", | ||
| }) | ||
|
|
||
| const clientInstance = (handler as any).client | ||
| expect(clientInstance).toBeDefined() | ||
| clientInstance.send = mockSend | ||
|
|
||
| mockSend.mockResolvedValueOnce({ | ||
| output: { message: { content: [{ type: "text", text: "response" }] }, stopReason: null }, | ||
| }) | ||
|
|
||
| await handler.completePrompt("test prompt", { timeoutMs: 5000 }) | ||
|
|
||
| // bedrock.ts uses truthy check for timeoutMs, so it creates AbortSignal.timeout | ||
| expect(mockSend).toHaveBeenCalled() | ||
| }) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
This assertion doesn't validate timeoutMs propagation.
Line 1650 only checks that send() ran, so the test still passes if timeoutMs is ignored completely. Please inspect the second argument or drive the derived abort path so this actually verifies the timeout behavior its name claims to cover. As per coding guidelines, **/*.{test,spec}.{ts,tsx,js}: Use package-local unit tests for pure logic, parsing, state transitions, validation, serialization, request construction, retry decisions, and error handling.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/api/providers/__tests__/bedrock.spec.ts` around lines 1629 - 1651, The
test in AwsBedrockHandler.completePrompt currently only asserts that send() was
called, so it does not verify that timeoutMs is actually propagated. Update the
spec to inspect the arguments passed to the mocked client send() call, or assert
the derived AbortSignal/abort behavior, so the test proves the timeout option is
used by completePrompt rather than just reaching the send path.
Source: Coding guidelines
Comment on lines
+65
to
+77
| test("should replace existing signal if metadata contains abortSignal", () => { | ||
| const controller1 = new AbortController() | ||
| const controller2 = new AbortController() | ||
|
|
||
| const builder = new RequestConfigBuilder({ abortSignal: controller1.signal }) | ||
| builder.addAbortSignal({ | ||
| taskId: "test-task", | ||
| abortSignal: controller2.signal, | ||
| } as ApiHandlerCreateMessageMetadata) | ||
|
|
||
| const config = builder.build() as { signal?: AbortSignal } | ||
| expect(config?.signal).toBe(controller2.signal) | ||
| }) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
This test never exercises signal replacement.
addAbortSignal() writes to signal, but the fixture seeds abortSignal, so the assertion still passes even if replacement is broken and both keys coexist.
Suggested fix
- const builder = new RequestConfigBuilder({ abortSignal: controller1.signal })
+ const builder = new RequestConfigBuilder({ signal: controller1.signal })
builder.addAbortSignal({
taskId: "test-task",
abortSignal: controller2.signal,
} as ApiHandlerCreateMessageMetadata)
const config = builder.build() as { signal?: AbortSignal }
expect(config?.signal).toBe(controller2.signal)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| test("should replace existing signal if metadata contains abortSignal", () => { | |
| const controller1 = new AbortController() | |
| const controller2 = new AbortController() | |
| const builder = new RequestConfigBuilder({ abortSignal: controller1.signal }) | |
| builder.addAbortSignal({ | |
| taskId: "test-task", | |
| abortSignal: controller2.signal, | |
| } as ApiHandlerCreateMessageMetadata) | |
| const config = builder.build() as { signal?: AbortSignal } | |
| expect(config?.signal).toBe(controller2.signal) | |
| }) | |
| test("should replace existing signal if metadata contains abortSignal", () => { | |
| const controller1 = new AbortController() | |
| const controller2 = new AbortController() | |
| const builder = new RequestConfigBuilder({ signal: controller1.signal }) | |
| builder.addAbortSignal({ | |
| taskId: "test-task", | |
| abortSignal: controller2.signal, | |
| } as ApiHandlerCreateMessageMetadata) | |
| const config = builder.build() as { signal?: AbortSignal } | |
| expect(config?.signal).toBe(controller2.signal) | |
| }) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/api/providers/__tests__/request-config-builder.spec.ts` around lines 65 -
77, The RequestConfigBuilder replacement test is seeding the wrong property, so
it never proves that addAbortSignal actually overrides an existing signal.
Update the test setup to initialize the builder with the same signal field that
RequestConfigBuilder.build() uses, then verify addAbortSignal replaces that
existing signal with controller2.signal in the
RequestConfigBuilder/addAbortSignal flow.
Comment on lines
+473
to
+478
| // After completion, the token source is still referenced but will be cleaned up on next request | ||
| const cancellationAfter = handler["currentRequestCancellation"] as any | ||
| expect(cancellationAfter).toBeDefined() | ||
| // dispose happens when ensureCleanState is called (on next request or error) | ||
| // For now, just verify the token source exists and dispose method was not yet called | ||
| expect(cancellationAfter.dispose).not.toHaveBeenCalled() |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
This assertion locks in the leak the PR is meant to remove.
Lines 473-478 expect the CancellationTokenSource to remain live and dispose() not to run, but the stated goal for vscode-lm.ts is to dispose the source in finally. If the implementation is correct, this test will fail; if the test passes, it masks the cleanup bug.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/api/providers/__tests__/vscode-lm.spec.ts` around lines 473 - 478, The
test around handler["currentRequestCancellation"] is asserting the old leaked
behavior by expecting the CancellationTokenSource to remain defined and
dispose() not to be called. Update the vscode-lm.spec.ts expectations to match
the intended cleanup in the request handler/finally path: verify that the
cancellation source is disposed after completion and no longer retained, using
the relevant handler/currentRequestCancellation and dispose assertions.
Comment on lines
+212
to
+214
| if (options?.timeoutMs) { | ||
| createOptions.timeout = options.timeoutMs | ||
| } |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Preserve explicit timeoutMs=0.
Line 212 uses a truthy check, so { timeoutMs: 0 } falls back to the client default timeout instead of forwarding the caller's override. The sibling OpenAI-compatible path already uses !== undefined for this case.
Suggested fix
- if (options?.timeoutMs) {
+ if (options?.timeoutMs !== undefined) {
createOptions.timeout = options.timeoutMs
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if (options?.timeoutMs) { | |
| createOptions.timeout = options.timeoutMs | |
| } | |
| if (options?.timeoutMs !== undefined) { | |
| createOptions.timeout = options.timeoutMs | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/api/providers/unbound.ts` around lines 212 - 214, The timeout override in
the unbound provider drops an explicit 0 value because the `options.timeoutMs`
check is truthy, so update the timeout assignment in `unbound.ts` to preserve
`timeoutMs=0` by checking for `undefined` instead of truthiness. Make the change
in the same `createOptions.timeout` setup used by the provider’s request path so
the caller-provided timeout is forwarded exactly, matching the sibling
OpenAI-compatible handling.
Comment on lines
+141
to
+142
| if (options?.timeoutMs) { | ||
| createOptions.timeout = options.timeoutMs |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Preserve an explicit timeoutMs: 0.
Line 141 uses a truthy check, so completePrompt(..., { timeoutMs: 0 }) is treated the same as omitting the override. The shared OpenAI-compatible path already uses !== undefined, which keeps explicit zero values intact.
Suggested fix
- if (options?.timeoutMs) {
+ if (options?.timeoutMs !== undefined) {
createOptions.timeout = options.timeoutMs
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if (options?.timeoutMs) { | |
| createOptions.timeout = options.timeoutMs | |
| if (options?.timeoutMs !== undefined) { | |
| createOptions.timeout = options.timeoutMs |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/api/providers/vercel-ai-gateway.ts` around lines 141 - 142, The timeout
override in vercel-ai-gateway’s createOptions handling is using a truthy check,
so completePrompt and similar paths ignore an explicit timeoutMs: 0. Update the
condition in vercel-ai-gateway.ts to match the OpenAI-compatible path by
checking for undefined instead of truthiness, keeping zero as a valid override
value. Use the existing createOptions logic and any shared timeout handling
around options.timeoutMs as the place to fix it.
Comment on lines
+387
to
+398
| const tokenSource = new vscode.CancellationTokenSource() | ||
| this.currentRequestCancellation = tokenSource | ||
|
|
||
| // Bridge external AbortSignal to VSCode CancellationToken. | ||
| // metadata.abortSignal is set by the task layer when the user clicks "Stop". | ||
| if (metadata?.abortSignal) { | ||
| if (metadata.abortSignal.aborted) { | ||
| tokenSource.cancel() | ||
| } else { | ||
| metadata.abortSignal.addEventListener("abort", () => tokenSource.cancel(), { once: true }) | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Dispose the bridged CancellationTokenSource on the success path too.
This new per-request tokenSource is only cleaned up when an error hits ensureCleanState(), on the next request, or on handler disposal. Successful streams now leave currentRequestCancellation pointing at a completed request, which is the leak the linked objective was trying to remove. Wrap the request body in try/finally, clear currentRequestCancellation, and dispose the source there.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/api/providers/vscode-lm.ts` around lines 387 - 398, The per-request
CancellationTokenSource created in the vscode-lm request flow is only cleaned up
on error, so successful requests leave currentRequestCancellation pointing at a
completed token and leak the source. Update the request handling around the
tokenSource setup in vscode-lm to use a try/finally around the request body, and
in the finally block clear currentRequestCancellation and dispose the
tokenSource on every exit path, including successful streams.
Comment on lines
+154
to
+155
| if (options?.timeoutMs) { | ||
| requestOptions.timeout = options.timeoutMs |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Don't collapse timeoutMs: 0 into “no override.”
Line 154 uses a truthy check, so completePrompt(..., { timeoutMs: 0 }) is treated the same as omitting the field entirely. In this handler that leaves the constructor-level default timeout in place, which diverges from the shared OpenAI-style path that forwards any defined timeout value.
Suggested fix
- if (options?.timeoutMs) {
+ if (options?.timeoutMs !== undefined) {
requestOptions.timeout = options.timeoutMs
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if (options?.timeoutMs) { | |
| requestOptions.timeout = options.timeoutMs | |
| if (options?.timeoutMs !== undefined) { | |
| requestOptions.timeout = options.timeoutMs |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/api/providers/xai.ts` around lines 154 - 155, The timeout override in xai
completePrompt is using a truthy check, so timeoutMs: 0 is ignored and the
default timeout remains in effect. Update the requestOptions timeout assignment
in completePrompt to check for the field being defined rather than truthy,
matching the shared OpenAI-style timeout handling so an explicit zero value is
forwarded correctly.
Comment on lines
+302
to
+303
| if (options?.timeoutMs) { | ||
| createOptions.timeout = options.timeoutMs |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Ignore non-positive timeoutMs before forwarding it.
This branch accepts negative values because the check is only truthy. The shared completion path treats timeout as valid only when it is > 0, so this can turn an invalid caller input into an SDK-level request error instead of the existing “no timeout” behavior.
Suggested fix
- if (options?.timeoutMs) {
+ if ((options?.timeoutMs ?? 0) > 0) {
createOptions.timeout = options.timeoutMs
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if (options?.timeoutMs) { | |
| createOptions.timeout = options.timeoutMs | |
| if ((options?.timeoutMs ?? 0) > 0) { | |
| createOptions.timeout = options.timeoutMs |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/api/providers/zoo-gateway.ts` around lines 302 - 303, The timeout
forwarding branch in zoo-gateway’s request setup is only checking for truthy
values, so negative timeoutMs values can still be passed through as if they were
valid. Update the timeout handling around the createOptions assignment to ignore
any non-positive timeoutMs and only forward values greater than 0, preserving
the existing “no timeout” behavior for invalid inputs.
… signal propagation
easonLiangWorldedtech
force-pushed
the
feat/abort-signal-core/sdk-providers
branch
from
dc3acb3 to
8d73895
Compare
easonLiangWorldedtech
force-pushed
the
feat/abort-signal-core/sdk-providers
branch
from
8d73895 to
a5ab2d4
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Wires
AbortSignalto 3 SDK-backed providers with individual constraints — CancellationToken bridging, config placement, and type-safe casting.Closes #618 (part of #404, depends on #615)
Providers
vscode-lm.ts
CancellationToken↔AbortSignal: createsCancellationTokenSource, wiresmetadata.abortSignalto callsource.cancel()source.dispose()infinallyblock (prevents memory leak per vscode#205966)removeEventListenerinside{ once: true }callbackgemini.ts
configobject per@google/genaiSDK specGenerateContentParameterstype — noanycast, uses typed assertion insteadmistral.ts
completePromptverified against SDK types and documentedChanges at a Glance
Closes #618
Summary by CodeRabbit