feat(api): pass through abortSignal to 19 provider implementations

Author: easonliang28

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

@@ -538,6 +538,74 @@ describe("VertexHandler", () => {
 			expect(usageChunks[0]).toHaveProperty("cacheWriteTokens", 5)
 			expect(usageChunks[0]).toHaveProperty("cacheReadTokens", 3)
 		})
+
+		describe("abort signal", () => {
+			const systemPrompt = "You are a helpful assistant"
+			const mockMessages: Anthropic.Messages.MessageParam[] = [
+				{
+					role: "user",
+					content: "Hello",
+				},
+			]
+
+			it("should pass abort signal through to client in createMessage", async () => {
+				handler = new AnthropicVertexHandler({
+					apiModelId: "claude-3-5-sonnet-v2@20241022",
+					vertexProjectId: "test-project",
+					vertexRegion: "us-central1",
+				})
+
+				const controller = new AbortController()
+
+				const mockCreate = vitest.fn().mockImplementation(async () => ({
+					async *[Symbol.asyncIterator]() {
+						yield { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } }
+					},
+				}))
+				;(handler["client"].messages as any).create = mockCreate
+
+				const stream = handler.createMessage(systemPrompt, mockMessages, {
+					taskId: "test-task",
+					abortSignal: controller.signal as any,
+				})
+				for await (const _ of stream) {
+					// consume stream
+				}
+
+				expect(mockCreate).toHaveBeenCalledWith(
+					expect.any(Object),
+					expect.objectContaining({ signal: controller.signal }),
+				)
+			})
+
+			it("should pass the exact same signal reference (reference identity)", async () => {
+				handler = new AnthropicVertexHandler({
+					apiModelId: "claude-3-5-sonnet-v2@20241022",
+					vertexProjectId: "test-project",
+					vertexRegion: "us-central1",
+				})
+
+				const controller = new AbortController()
+
+				const mockCreate = vitest.fn().mockImplementation(async () => ({
+					async *[Symbol.asyncIterator]() {
+						yield { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } }
+					},
+				}))
+				;(handler["client"].messages as any).create = mockCreate
+
+				const stream = handler.createMessage(systemPrompt, mockMessages, {
+					taskId: "test-task",
+					abortSignal: controller.signal as any,
+				})
+				for await (const _ of stream) {
+					// consume stream
+				}
+
+				const callOptions = mockCreate.mock.calls[0][1]
+				expect(callOptions?.signal).toBe(controller.signal)
+			})
+		})
 	})
 
 	describe("thinking functionality", () => {

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

@@ -425,8 +425,42 @@ describe("AnthropicHandler", () => {
 			const requestOptions = mockCreate.mock.calls[mockCreate.mock.calls.length - 1]?.[1]
 			expect(requestBody?.thinking).toEqual({ type: "adaptive" })
 			expect(requestBody?.temperature).toBeUndefined()
-			expect(requestBody?.max_tokens).toBe(32768)
-			expect(requestOptions?.headers?.["anthropic-beta"]).toContain("prompt-caching-2024-07-31")
+		})
+
+		describe("abort signal", () => {
+			it("should pass abort signal through to client in createMessage", async () => {
+				const controller = new AbortController()
+				const testHandler = new AnthropicHandler(mockOptions)
+
+				const stream = testHandler.createMessage(systemPrompt, [], {
+					taskId: "test-task",
+					abortSignal: controller.signal as any,
+				})
+				for await (const _ of stream) {
+					// consume stream
+				}
+
+				expect(mockCreate).toHaveBeenCalledWith(
+					expect.any(Object),
+					expect.objectContaining({ signal: controller.signal }),
+				)
+			})
+
+			it("should pass the exact same signal reference (reference identity)", async () => {
+				const controller = new AbortController()
+				const testHandler = new AnthropicHandler(mockOptions)
+
+				const stream = testHandler.createMessage(systemPrompt, [], {
+					taskId: "test-task",
+					abortSignal: controller.signal as any,
+				})
+				for await (const _ of stream) {
+					// consume stream
+				}
+
+				const callOptions = mockCreate.mock.calls[0][1]
+				expect(callOptions?.signal).toBe(controller.signal)
+			})
 		})
 	})
 

src/api/providers/__tests__/base-openai-compatible-provider.spec.ts

@@ -454,6 +454,56 @@ describe("BaseOpenAiCompatibleProvider", () => {
 		})
 	})
 
+	describe("abort signal", () => {
+		it("should pass abort signal through to client in createMessage", async () => {
+			const controller = new AbortController()
+
+			mockCreate.mockImplementationOnce(() => ({
+				[Symbol.asyncIterator]: () => ({
+					async next() {
+						return { done: true }
+					},
+				}),
+			}))
+
+			const stream = handler.createMessage("system prompt", [], {
+				taskId: "test-task",
+				abortSignal: controller.signal as any,
+			})
+			for await (const _ of stream) {
+				// consume stream
+			}
+
+			expect(mockCreate).toHaveBeenCalledWith(
+				expect.any(Object),
+				expect.objectContaining({ signal: controller.signal }),
+			)
+		})
+
+		it("should pass the exact same signal reference (reference identity)", async () => {
+			const controller = new AbortController()
+
+			mockCreate.mockImplementationOnce(() => ({
+				[Symbol.asyncIterator]: () => ({
+					async next() {
+						return { done: true }
+					},
+				}),
+			}))
+
+			const stream = handler.createMessage("system prompt", [], {
+				taskId: "test-task",
+				abortSignal: controller.signal as any,
+			})
+			for await (const _ of stream) {
+				// consume stream
+			}
+
+			const callOptions = mockCreate.mock.calls[0][1]
+			expect(callOptions?.signal).toBe(controller.signal)
+		})
+	})
+
 	describe("Tool call handling", () => {
 		it("should yield tool_call_end events when finish_reason is tool_calls", async () => {
 			mockCreate.mockImplementationOnce(() => {

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

@@ -455,6 +455,42 @@ describe("DeepSeekHandler", () => {
 			const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning")
 			expect(reasoningChunks).toEqual([{ type: "reasoning", text: "primary thought" }])
 		})
+
+		describe("abort signal", () => {
+			it("should pass abort signal through to client in createMessage", async () => {
+				const controller = new AbortController()
+				const testHandler = new DeepSeekHandler(mockOptions)
+
+				const stream = testHandler.createMessage(systemPrompt, messages, {
+					taskId: "test-task",
+					abortSignal: controller.signal as any,
+				})
+				for await (const _ of stream) {
+					// consume stream
+				}
+
+				expect(mockCreate).toHaveBeenCalledWith(
+					expect.any(Object),
+					expect.objectContaining({ signal: controller.signal }),
+				)
+			})
+
+			it("should pass the exact same signal reference (reference identity)", async () => {
+				const controller = new AbortController()
+				const testHandler = new DeepSeekHandler(mockOptions)
+
+				const stream = testHandler.createMessage(systemPrompt, messages, {
+					taskId: "test-task",
+					abortSignal: controller.signal as any,
+				})
+				for await (const _ of stream) {
+					// consume stream
+				}
+
+				const callOptions = mockCreate.mock.calls[0][1]
+				expect(callOptions?.signal).toBe(controller.signal)
+			})
+		})
 	})
 
 	describe("processUsageMetrics", () => {
@@ -563,7 +599,7 @@ describe("DeepSeekHandler", () => {
 				expect.objectContaining({
 					thinking: { type: "enabled" },
 				}),
-				{}, // Empty path options for non-Azure URLs
+				undefined,
 			)
 			const callArgs = mockCreate.mock.calls[0][0]
 			expect(callArgs.reasoning_effort).toBeUndefined()
@@ -586,7 +622,7 @@ describe("DeepSeekHandler", () => {
 					reasoning_effort: "high",
 					max_completion_tokens: 200_000,
 				}),
-				{},
+				undefined,
 			)
 		})
 
@@ -623,7 +659,7 @@ describe("DeepSeekHandler", () => {
 					thinking: { type: "enabled" },
 					reasoning_effort: "max",
 				}),
-				{},
+				undefined,
 			)
 		})
 

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

@@ -922,6 +922,62 @@ describe("LiteLLMHandler", () => {
 		})
 	})
 
+	describe("abort signal", () => {
+		it("should pass abort signal through to client in createMessage", async () => {
+			const mockStream = {
+				async *[Symbol.asyncIterator]() {
+					yield {
+						choices: [{ delta: { content: "response" } }],
+						usage: null,
+					}
+				},
+			}
+
+			mockCreate.mockReturnValue({
+				withResponse: vi.fn().mockResolvedValue({ data: mockStream }),
+			})
+
+			const controller = new AbortController()
+			const stream = handler.createMessage("system", [{ role: "user", content: "Test" }], {
+				taskId: "test-task",
+				abortSignal: controller.signal as any,
+			})
+			for await (const _ of stream) {
+			}
+
+			expect(mockCreate).toHaveBeenCalledWith(
+				expect.any(Object),
+				expect.objectContaining({ signal: controller.signal }),
+			)
+		})
+
+		it("should pass the exact same signal reference (reference identity)", async () => {
+			const mockStream = {
+				async *[Symbol.asyncIterator]() {
+					yield {
+						choices: [{ delta: { content: "response" } }],
+						usage: null,
+					}
+				},
+			}
+
+			mockCreate.mockReturnValue({
+				withResponse: vi.fn().mockResolvedValue({ data: mockStream }),
+			})
+
+			const controller = new AbortController()
+			const stream = handler.createMessage("system", [{ role: "user", content: "Test" }], {
+				taskId: "test-task",
+				abortSignal: controller.signal as any,
+			})
+			for await (const _ of stream) {
+			}
+
+			const callOptions = mockCreate.mock.calls[0][1]
+			expect(callOptions?.signal).toBe(controller.signal)
+		})
+	})
+
 	describe("tool ID normalization", () => {
 		it("should truncate tool IDs longer than 64 characters", async () => {
 			const optionsWithBedrock: ApiHandlerOptions = {

src/api/providers/__tests__/lmstudio-native-tools.spec.ts

@@ -83,6 +83,7 @@ describe("LmStudioHandler Native Tools", () => {
 						}),
 					]),
 				}),
+				undefined,
 			)
 			// parallel_tool_calls should be true by default when not explicitly set
 			const callArgs = mockCreate.mock.calls[0][0]
@@ -109,6 +110,7 @@ describe("LmStudioHandler Native Tools", () => {
 				expect.objectContaining({
 					tool_choice: "auto",
 				}),
+				undefined,
 			)
 		})
 
@@ -221,6 +223,7 @@ describe("LmStudioHandler Native Tools", () => {
 				expect.objectContaining({
 					parallel_tool_calls: true,
 				}),
+				undefined,
 			)
 		})
 

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

@@ -127,6 +127,42 @@ describe("LmStudioHandler", () => {
 				}
 			}).rejects.toThrow("Please check the LM Studio developer logs to debug what went wrong")
 		})
+
+		describe("abort signal", () => {
+			it("should pass abort signal through to client in createMessage", async () => {
+				const controller = new AbortController()
+				const testHandler = new LmStudioHandler(mockOptions)
+
+				const stream = testHandler.createMessage(systemPrompt, messages, {
+					taskId: "test-task",
+					abortSignal: controller.signal as any,
+				})
+				for await (const _ of stream) {
+					// consume stream
+				}
+
+				expect(mockCreate).toHaveBeenCalledWith(
+					expect.any(Object),
+					expect.objectContaining({ signal: controller.signal }),
+				)
+			})
+
+			it("should pass the exact same signal reference (reference identity)", async () => {
+				const controller = new AbortController()
+				const testHandler = new LmStudioHandler(mockOptions)
+
+				const stream = testHandler.createMessage(systemPrompt, messages, {
+					taskId: "test-task",
+					abortSignal: controller.signal as any,
+				})
+				for await (const _ of stream) {
+					// consume stream
+				}
+
+				const callOptions = mockCreate.mock.calls[0][1]
+				expect(callOptions?.signal).toBe(controller.signal)
+			})
+		})
 	})
 
 	describe("completePrompt", () => {