Skip to content

Commit de408ee

Browse files
committed
feat(api): integrate external abort signal with internal controllers
Previously, these three providers used their own AbortControllers but did not listen to the external abort signal from the Task execution chain. This meant clicking stop would not immediately cancel requests to these providers. Now all providers consistently respect the user's stop action. Also add optional metadata parameter with abortSignal to completePrompt methods across 22 provider implementations: - Anthropic, Anthropic Vertex, Base OpenAI Compatible - Gemini/Vertex, LiteLLM, LM Studio, MiniMax, Mistral - Native Ollama, OpenAI, OpenAI Compatible, OpenRouter - Qwen Code, Requesty, Unbound, Vercel AI Gateway - Opencode Go, xAI, Zoo Gateway, VSCode LM, Poe, Bedrock Add abortSignal tests for all modified providers. - Add comment explaining Mistral completePrompt does not support non-streaming abort (Mistral SDK limitation)
1 parent 05219ef commit de408ee

40 files changed

Lines changed: 609 additions & 97 deletions

src/api/providers/__tests__/anthropic-vertex.spec.ts

Lines changed: 54 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -763,18 +763,21 @@ describe("VertexHandler", () => {
763763

764764
const result = await handler.completePrompt("Test prompt")
765765
expect(result).toBe("Test response")
766-
expect(handler["client"].messages.create).toHaveBeenCalledWith({
767-
model: "claude-3-5-sonnet-v2@20241022",
768-
max_tokens: 8192,
769-
temperature: 0,
770-
messages: [
771-
{
772-
role: "user",
773-
content: [{ type: "text", text: "Test prompt", cache_control: { type: "ephemeral" } }],
774-
},
775-
],
776-
stream: false,
777-
})
766+
expect(handler["client"].messages.create).toHaveBeenCalledWith(
767+
{
768+
model: "claude-3-5-sonnet-v2@20241022",
769+
max_tokens: 8192,
770+
temperature: 0,
771+
messages: [
772+
{
773+
role: "user",
774+
content: [{ type: "text", text: "Test prompt", cache_control: { type: "ephemeral" } }],
775+
},
776+
],
777+
stream: false,
778+
},
779+
undefined,
780+
)
778781
})
779782

780783
it("should handle API errors for Claude", async () => {
@@ -824,6 +827,45 @@ describe("VertexHandler", () => {
824827
const result = await handler.completePrompt("Test prompt")
825828
expect(result).toBe("")
826829
})
830+
831+
it("should pass abortSignal to messages.create when provided in metadata", async () => {
832+
handler = new AnthropicVertexHandler({
833+
apiModelId: "claude-3-5-sonnet-v2@20241022",
834+
vertexProjectId: "test-project",
835+
vertexRegion: "us-central1",
836+
})
837+
838+
const controller = new AbortController()
839+
const mockAbortSignal = controller.signal
840+
841+
const mockCreate = vitest.fn().mockResolvedValue({
842+
content: [{ type: "text", text: "Test response" }],
843+
})
844+
;(handler["client"].messages as any).create = mockCreate
845+
846+
await handler.completePrompt("Test prompt", { taskId: "test", abortSignal: mockAbortSignal })
847+
848+
const callArgs = mockCreate.mock.calls[0][1]
849+
expect(callArgs?.signal).toBe(mockAbortSignal)
850+
})
851+
852+
it("should pass undefined signal when abortSignal is not provided", async () => {
853+
handler = new AnthropicVertexHandler({
854+
apiModelId: "claude-3-5-sonnet-v2@20241022",
855+
vertexProjectId: "test-project",
856+
vertexRegion: "us-central1",
857+
})
858+
859+
const mockCreate = vitest.fn().mockResolvedValue({
860+
content: [{ type: "text", text: "Test response" }],
861+
})
862+
;(handler["client"].messages as any).create = mockCreate
863+
864+
await handler.completePrompt("Test prompt")
865+
866+
const callArgs = mockCreate.mock.calls[0][1]
867+
expect(callArgs?.signal).toBeUndefined()
868+
})
827869
})
828870

829871
describe("getModel", () => {

src/api/providers/__tests__/anthropic.spec.ts

Lines changed: 28 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -432,14 +432,17 @@ describe("AnthropicHandler", () => {
432432
it("should complete prompt successfully", async () => {
433433
const result = await handler.completePrompt("Test prompt")
434434
expect(result).toBe("Test response")
435-
expect(mockCreate).toHaveBeenCalledWith({
436-
model: mockOptions.apiModelId,
437-
messages: [{ role: "user", content: "Test prompt" }],
438-
max_tokens: 8192,
439-
temperature: 0,
440-
thinking: undefined,
441-
stream: false,
442-
})
435+
expect(mockCreate).toHaveBeenCalledWith(
436+
{
437+
model: mockOptions.apiModelId,
438+
messages: [{ role: "user", content: "Test prompt" }],
439+
max_tokens: 8192,
440+
temperature: 0,
441+
thinking: undefined,
442+
stream: false,
443+
},
444+
undefined,
445+
)
443446
})
444447

445448
it("should handle API errors", async () => {
@@ -462,6 +465,23 @@ describe("AnthropicHandler", () => {
462465
const result = await handler.completePrompt("Test prompt")
463466
expect(result).toBe("")
464467
})
468+
469+
it("should pass abortSignal to messages.create when provided in metadata", async () => {
470+
const controller = new AbortController()
471+
const mockAbortSignal = controller.signal
472+
473+
await handler.completePrompt("Test prompt", { taskId: "test", abortSignal: mockAbortSignal })
474+
475+
const callArgs = mockCreate.mock.calls[0][1]
476+
expect(callArgs?.signal).toBe(mockAbortSignal)
477+
})
478+
479+
it("should pass undefined signal when abortSignal is not provided", async () => {
480+
await handler.completePrompt("Test prompt")
481+
482+
const callArgs = mockCreate.mock.calls[0][1]
483+
expect(callArgs?.signal).toBeUndefined()
484+
})
465485
})
466486

467487
describe("getModel", () => {

src/api/providers/__tests__/gemini.spec.ts

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,33 @@ describe("GeminiHandler", () => {
156156
const result = await handler.completePrompt("Test prompt")
157157
expect(result).toBe("")
158158
})
159+
160+
it("should pass abortSignal to generateContent when provided in metadata", async () => {
161+
const mockGenerateContent = vitest.fn().mockResolvedValue({
162+
text: "Test response",
163+
})
164+
;(handler["client"].models as any).generateContent = mockGenerateContent
165+
166+
const controller = new AbortController()
167+
const mockAbortSignal = controller.signal
168+
169+
await handler.completePrompt("Test prompt", { taskId: "test", abortSignal: mockAbortSignal })
170+
171+
const callArgs = mockGenerateContent.mock.calls[0][0]
172+
expect(callArgs.signal).toBe(mockAbortSignal)
173+
})
174+
175+
it("should pass undefined signal when abortSignal is not provided", async () => {
176+
const mockGenerateContent = vitest.fn().mockResolvedValue({
177+
text: "Test response",
178+
})
179+
;(handler["client"].models as any).generateContent = mockGenerateContent
180+
181+
await handler.completePrompt("Test prompt")
182+
183+
const callArgs = mockGenerateContent.mock.calls[0][0]
184+
expect(callArgs.signal).toBeUndefined()
185+
})
159186
})
160187

161188
describe("getModel", () => {

src/api/providers/__tests__/lite-llm.spec.ts

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -393,6 +393,31 @@ describe("LiteLLMHandler", () => {
393393
expect(createCall.max_tokens).toBeUndefined()
394394
expect(createCall.max_completion_tokens).toBeUndefined()
395395
})
396+
397+
it("should pass abortSignal to chat.completions.create when provided in metadata", async () => {
398+
const controller = new AbortController()
399+
const mockAbortSignal = controller.signal
400+
401+
mockCreate.mockResolvedValue({
402+
choices: [{ message: { content: "Test response" } }],
403+
})
404+
405+
await handler.completePrompt("Test prompt", { taskId: "test", abortSignal: mockAbortSignal })
406+
407+
const callArgs = mockCreate.mock.calls[0][1]
408+
expect(callArgs?.signal).toBe(mockAbortSignal)
409+
})
410+
411+
it("should pass undefined signal when abortSignal is not provided", async () => {
412+
mockCreate.mockResolvedValue({
413+
choices: [{ message: { content: "Test response" } }],
414+
})
415+
416+
await handler.completePrompt("Test prompt")
417+
418+
const callArgs = mockCreate.mock.calls[0][1]
419+
expect(callArgs?.signal).toBeUndefined()
420+
})
396421
})
397422

398423
describe("Gemini thought signature injection", () => {

src/api/providers/__tests__/lmstudio.spec.ts

Lines changed: 34 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -131,12 +131,15 @@ describe("LmStudioHandler", () => {
131131
it("should complete prompt successfully", async () => {
132132
const result = await handler.completePrompt("Test prompt")
133133
expect(result).toBe("Test response")
134-
expect(mockCreate).toHaveBeenCalledWith({
135-
model: mockOptions.lmStudioModelId,
136-
messages: [{ role: "user", content: "Test prompt" }],
137-
temperature: 0,
138-
stream: false,
139-
})
134+
expect(mockCreate).toHaveBeenCalledWith(
135+
{
136+
model: mockOptions.lmStudioModelId,
137+
messages: [{ role: "user", content: "Test prompt" }],
138+
temperature: 0,
139+
stream: false,
140+
},
141+
undefined,
142+
)
140143
})
141144

142145
it("should handle API errors", async () => {
@@ -153,6 +156,31 @@ describe("LmStudioHandler", () => {
153156
const result = await handler.completePrompt("Test prompt")
154157
expect(result).toBe("")
155158
})
159+
160+
it("should pass abortSignal to chat.completions.create when provided in metadata", async () => {
161+
const controller = new AbortController()
162+
const mockAbortSignal = controller.signal
163+
164+
mockCreate.mockResolvedValueOnce({
165+
choices: [{ message: { content: "Test response" } }],
166+
})
167+
168+
await handler.completePrompt("Test prompt", { taskId: "test", abortSignal: mockAbortSignal })
169+
170+
const callArgs = mockCreate.mock.calls[0][1]
171+
expect(callArgs?.signal).toBe(mockAbortSignal)
172+
})
173+
174+
it("should pass undefined signal when abortSignal is not provided", async () => {
175+
mockCreate.mockResolvedValueOnce({
176+
choices: [{ message: { content: "Test response" } }],
177+
})
178+
179+
await handler.completePrompt("Test prompt")
180+
181+
const callArgs = mockCreate.mock.calls[0][1]
182+
expect(callArgs?.signal).toBeUndefined()
183+
})
156184
})
157185

158186
describe("getModel", () => {

src/api/providers/__tests__/minimax.spec.ts

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,31 @@ describe("MiniMaxHandler", () => {
218218
await expect(handler.completePrompt("test prompt")).rejects.toThrow()
219219
})
220220

221+
it("should pass abortSignal to messages.create when provided in metadata", async () => {
222+
const controller = new AbortController()
223+
const mockAbortSignal = controller.signal
224+
225+
mockCreate.mockResolvedValueOnce({
226+
content: [{ type: "text", text: "Test response" }],
227+
})
228+
229+
await handler.completePrompt("test prompt", { taskId: "test", abortSignal: mockAbortSignal })
230+
231+
const callArgs = mockCreate.mock.calls[0][1]
232+
expect(callArgs?.signal).toBe(mockAbortSignal)
233+
})
234+
235+
it("should pass undefined signal when abortSignal is not provided", async () => {
236+
mockCreate.mockResolvedValueOnce({
237+
content: [{ type: "text", text: "Test response" }],
238+
})
239+
240+
await handler.completePrompt("test prompt")
241+
242+
const callArgs = mockCreate.mock.calls[0][1]
243+
expect(callArgs?.signal).toBeUndefined()
244+
})
245+
221246
it("createMessage should yield text content from stream", async () => {
222247
const testContent = "This is test content from MiniMax stream"
223248

src/api/providers/__tests__/native-ollama.spec.ts

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,31 @@ describe("NativeOllamaHandler", () => {
231231
}),
232232
)
233233
})
234+
235+
it("should pass abortSignal to chat when provided in metadata", async () => {
236+
const controller = new AbortController()
237+
const mockAbortSignal = controller.signal
238+
239+
mockedData.mockChat.mockResolvedValue({
240+
message: { content: "Test response" },
241+
})
242+
243+
await handler.completePrompt("Test prompt", { taskId: "test", abortSignal: mockAbortSignal })
244+
245+
const callArgs = mockedData.mockChat.mock.calls[0][0]
246+
expect(callArgs.signal).toBe(mockAbortSignal)
247+
})
248+
249+
it("should pass undefined signal when abortSignal is not provided", async () => {
250+
mockedData.mockChat.mockResolvedValue({
251+
message: { content: "Test response" },
252+
})
253+
254+
await handler.completePrompt("Test prompt")
255+
256+
const callArgs = mockedData.mockChat.mock.calls[0][0]
257+
expect(callArgs.signal).toBeUndefined()
258+
})
234259
})
235260

236261
describe("error handling", () => {

src/api/providers/__tests__/openai.spec.ts

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -625,6 +625,23 @@ describe("OpenAiHandler", () => {
625625
const result = await handler.completePrompt("Test prompt")
626626
expect(result).toBe("")
627627
})
628+
629+
it("should pass abortSignal to chat.completions.create when provided in metadata", async () => {
630+
const controller = new AbortController()
631+
const mockAbortSignal = controller.signal
632+
633+
await handler.completePrompt("Test prompt", { taskId: "test", abortSignal: mockAbortSignal })
634+
635+
const callArgs = mockCreate.mock.calls[0][1]
636+
expect(callArgs?.signal).toBe(mockAbortSignal)
637+
})
638+
639+
it("should pass undefined signal when abortSignal is not provided", async () => {
640+
await handler.completePrompt("Test prompt")
641+
642+
const callArgs = mockCreate.mock.calls[0][1]
643+
expect(callArgs?.signal).toBeUndefined()
644+
})
628645
})
629646

630647
describe("getModel", () => {

src/api/providers/__tests__/opencode-go.spec.ts

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,5 +174,28 @@ describe("OpencodeGoHandler", () => {
174174
const handler = new OpencodeGoHandler(mockOptions)
175175
await expect(handler.completePrompt("ping")).rejects.toThrow("Opencode Go completion error: boom")
176176
})
177+
178+
it("should pass abortSignal to chat.completions.create when provided in metadata", async () => {
179+
mockCreate.mockResolvedValue({ choices: [{ message: { content: "the answer" } }] })
180+
const handler = new OpencodeGoHandler(mockOptions)
181+
182+
const controller = new AbortController()
183+
const mockAbortSignal = controller.signal
184+
185+
await handler.completePrompt("ping", { taskId: "test", abortSignal: mockAbortSignal })
186+
187+
const callArgs = mockCreate.mock.calls[0][1]
188+
expect(callArgs?.signal).toBe(mockAbortSignal)
189+
})
190+
191+
it("should pass undefined signal when abortSignal is not provided", async () => {
192+
mockCreate.mockResolvedValue({ choices: [{ message: { content: "the answer" } }] })
193+
const handler = new OpencodeGoHandler(mockOptions)
194+
195+
await handler.completePrompt("ping")
196+
197+
const callArgs = mockCreate.mock.calls[0][1]
198+
expect(callArgs?.signal).toBeUndefined()
199+
})
177200
})
178201
})

0 commit comments

Comments
 (0)