feat(api): add abort signal support for API providers and task (#434)

- Add AbortController support to openai-compatible provider
- Pass abort signal through Task execution chain
- Add comprehensive tests for abort signal behavior across all providers

Author: easonliang28

src/api/index.ts

@@ -90,6 +90,12 @@ export interface ApiHandlerCreateMessageMetadata {
 	 * Only applies to providers that support function calling restrictions (e.g., Gemini).
 	 */
 	allowedFunctionNames?: string[]
+	/**
+	 * Abort signal for cancelling the HTTP request mid-stream.
+	 * Passed through to AI SDK's streamText() so the underlying HTTP request is aborted
+	 * when the user clicks stop, preventing wasted API tokens/compute on the provider side.
+	 */
+	abortSignal?: AbortSignal
 }
 
 export interface ApiHandler {

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

@@ -1540,4 +1540,60 @@ describe("VertexHandler", () => {
 			})
 		})
 	})
+
+	describe("abortSignal support", () => {
+		it("should pass abortSignal to messages.create when provided in metadata", async () => {
+			const handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockCreate = handler["client"].messages.create as any
+			mockCreate.mockClear()
+
+			const controller = new AbortController()
+			const mockAbortSignal = controller.signal
+
+			const systemPrompt = "You are a helpful assistant."
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
+			]
+
+			for await (const _chunk of handler.createMessage(systemPrompt, messages, {
+				taskId: "test",
+				abortSignal: mockAbortSignal,
+			})) {
+				break
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][1]
+			expect(callArgs?.signal).toBe(mockAbortSignal)
+		})
+
+		it("should pass undefined signal when abortSignal is not provided", async () => {
+			const handler = new AnthropicVertexHandler({
+				apiModelId: "claude-3-5-sonnet-v2@20241022",
+				vertexProjectId: "test-project",
+				vertexRegion: "us-central1",
+			})
+
+			const mockCreate = handler["client"].messages.create as any
+			mockCreate.mockClear()
+
+			const systemPrompt = "You are a helpful assistant."
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
+			]
+
+			for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
+				break
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][1]
+			expect(callArgs?.signal).toBeUndefined()
+		})
+	})
 })

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

@@ -1055,4 +1055,57 @@ describe("AnthropicHandler", () => {
 			})
 		})
 	})
+
+	describe("abortSignal support", () => {
+		it("should pass abortSignal to messages.create when provided in metadata", async () => {
+			const handler = new AnthropicHandler({
+				apiKey: "test-api-key",
+				apiModelId: "claude-3-5-sonnet-20241022",
+			})
+
+			const controller = new AbortController()
+			const mockAbortSignal = controller.signal
+
+			mockCreate.mockResolvedValueOnce({
+				[Symbol.asyncIterator]: async function* () {
+					yield { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } }
+				},
+			})
+
+			for await (const _chunk of handler.createMessage(
+				"system",
+				[{ role: "user", content: [{ type: "text", text: "Hello!" }] }],
+				{ taskId: "test", abortSignal: mockAbortSignal },
+			)) {
+				break
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][1]
+			expect(callArgs?.signal).toBe(mockAbortSignal)
+		})
+
+		it("should pass undefined signal when abortSignal is not provided", async () => {
+			const handler = new AnthropicHandler({
+				apiKey: "test-api-key",
+				apiModelId: "claude-3-5-sonnet-20241022",
+			})
+
+			mockCreate.mockResolvedValueOnce({
+				[Symbol.asyncIterator]: async function* () {
+					yield { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } }
+				},
+			})
+
+			for await (const _chunk of handler.createMessage("system", [
+				{ role: "user", content: [{ type: "text", text: "Hello!" }] },
+			])) {
+				break
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][1]
+			expect(callArgs?.signal).toBeUndefined()
+		})
+	})
 })

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

@@ -358,7 +358,7 @@ describe("BaseOpenAiCompatibleProvider", () => {
 					stream: true,
 					stream_options: { include_usage: true },
 				}),
-				undefined,
+				expect.any(Object),
 			)
 		})
 

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

@@ -637,6 +637,48 @@ describe("DeepSeekHandler", () => {
 			const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial")
 			expect(toolCallChunks.length).toBeGreaterThan(0)
 			expect(toolCallChunks[0].name).toBe("get_weather")
+
+			describe("abortSignal support", () => {
+				it("should pass abortSignal to chat.completions.create when provided in metadata", async () => {
+					const handler = new DeepSeekHandler({ ...mockOptions, apiKey: "test-key" })
+					const systemPrompt = "You are a helpful assistant."
+					const messages: Anthropic.Messages.MessageParam[] = [
+						{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
+					]
+
+					const controller = new AbortController()
+					const mockAbortSignal = controller.signal
+
+					await handler.createMessage(systemPrompt, messages, {
+						taskId: "test",
+						abortSignal: mockAbortSignal,
+					})
+					for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
+						break
+					}
+
+					expect(mockCreate).toHaveBeenCalled()
+					const callArgs = mockCreate.mock.calls[0][0]
+					expect(callArgs.signal).toBe(mockAbortSignal)
+				})
+
+				it("should not include signal when abortSignal is not provided", async () => {
+					const handler = new DeepSeekHandler({ ...mockOptions, apiKey: "test-key" })
+					const systemPrompt = "You are a helpful assistant."
+					const messages: Anthropic.Messages.MessageParam[] = [
+						{ role: "user", content: [{ type: "text" as const, text: "Hello!" }] },
+					]
+
+					await handler.createMessage(systemPrompt, messages)
+					for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
+						break
+					}
+
+					expect(mockCreate).toHaveBeenCalled()
+					const callArgs = mockCreate.mock.calls[0][0]
+					expect(callArgs.signal).toBeUndefined()
+				})
+			})
 		})
 	})
 })

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

@@ -95,25 +95,46 @@ describe("FireworksHandler", () => {
 	})
 
 	it.each([
-		{ modelId: "accounts/fireworks/models/glm-5p1" as const, contextWindow: 202752, inputPrice: 1.4, outputPrice: 4.4, cacheReadsPrice: 0.26 },
-		{ modelId: "accounts/fireworks/models/kimi-k2p6" as const, contextWindow: 262144, inputPrice: 0.95, outputPrice: 4.0, cacheReadsPrice: 0.16 },
-		{ modelId: "accounts/fireworks/models/deepseek-v4-pro" as const, contextWindow: 1048576, inputPrice: 1.74, outputPrice: 3.48, cacheReadsPrice: 0.14 },
-	])("should expose newly added model $modelId", ({ modelId, contextWindow, inputPrice, outputPrice, cacheReadsPrice }) => {
-		expect(fireworksModels[modelId]).toBeDefined()
-		const info = fireworksModels[modelId]
-		expect(info.maxTokens).toBeGreaterThan(0)
-		expect(info.contextWindow).toBe(contextWindow)
-		expect(info.inputPrice).toBe(inputPrice)
-		expect(info.outputPrice).toBe(outputPrice)
-		expect(info.cacheReadsPrice).toBe(cacheReadsPrice)
-		expect(info.description).toBeTruthy()
-
-		const handlerWithModel = new FireworksHandler({
-			apiModelId: modelId,
-			fireworksApiKey: "test-fireworks-api-key",
-		})
-		expect(handlerWithModel.getModel().id).toBe(modelId)
-	})
+		{
+			modelId: "accounts/fireworks/models/glm-5p1" as const,
+			contextWindow: 202752,
+			inputPrice: 1.4,
+			outputPrice: 4.4,
+			cacheReadsPrice: 0.26,
+		},
+		{
+			modelId: "accounts/fireworks/models/kimi-k2p6" as const,
+			contextWindow: 262144,
+			inputPrice: 0.95,
+			outputPrice: 4.0,
+			cacheReadsPrice: 0.16,
+		},
+		{
+			modelId: "accounts/fireworks/models/deepseek-v4-pro" as const,
+			contextWindow: 1048576,
+			inputPrice: 1.74,
+			outputPrice: 3.48,
+			cacheReadsPrice: 0.14,
+		},
+	])(
+		"should expose newly added model $modelId",
+		({ modelId, contextWindow, inputPrice, outputPrice, cacheReadsPrice }) => {
+			expect(fireworksModels[modelId]).toBeDefined()
+			const info = fireworksModels[modelId]
+			expect(info.maxTokens).toBeGreaterThan(0)
+			expect(info.contextWindow).toBe(contextWindow)
+			expect(info.inputPrice).toBe(inputPrice)
+			expect(info.outputPrice).toBe(outputPrice)
+			expect(info.cacheReadsPrice).toBe(cacheReadsPrice)
+			expect(info.description).toBeTruthy()
+
+			const handlerWithModel = new FireworksHandler({
+				apiModelId: modelId,
+				fireworksApiKey: "test-fireworks-api-key",
+			})
+			expect(handlerWithModel.getModel().id).toBe(modelId)
+		},
+	)
 
 	it("should return Kimi K2 Instruct model with correct configuration", () => {
 		const testModelId: FireworksModelId = "accounts/fireworks/models/kimi-k2-instruct"
@@ -465,7 +486,7 @@ describe("FireworksHandler", () => {
 				stream: true,
 				stream_options: { include_usage: true },
 			}),
-			undefined,
+			expect.any(Object),
 		)
 	})
 
@@ -491,7 +512,7 @@ describe("FireworksHandler", () => {
 			expect.objectContaining({
 				temperature: 0.5,
 			}),
-			undefined,
+			expect.any(Object),
 		)
 	})
 
@@ -518,7 +539,7 @@ describe("FireworksHandler", () => {
 			expect.objectContaining({
 				temperature: 1.0,
 			}),
-			undefined,
+			expect.any(Object),
 		)
 	})
 
@@ -546,7 +567,7 @@ describe("FireworksHandler", () => {
 			expect.objectContaining({
 				temperature: 0.7,
 			}),
-			undefined,
+			expect.any(Object),
 		)
 	})
 

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

@@ -366,4 +366,48 @@ describe("GeminiHandler", () => {
 			expect(mockCaptureException).toHaveBeenCalled()
 		})
 	})
+
+	describe("abortSignal support", () => {
+		it("should pass abortSignal to generateContentStream when provided in metadata", async () => {
+			const mockGenerateContentStream = vitest.fn().mockResolvedValue({
+				[Symbol.asyncIterator]: async function* () {
+					yield { text: "Hello" }
+				},
+			})
+
+			handler["client"].models.generateContentStream = mockGenerateContentStream
+
+			const controller = new AbortController()
+			const mockAbortSignal = controller.signal
+
+			for await (const _chunk of handler.createMessage("system", [{ role: "user", content: "Hello!" }], {
+				taskId: "test",
+				abortSignal: mockAbortSignal,
+			})) {
+				break
+			}
+
+			expect(mockGenerateContentStream).toHaveBeenCalled()
+			const callArgs = mockGenerateContentStream.mock.calls[0][0]
+			expect(callArgs.signal).toBe(mockAbortSignal)
+		})
+
+		it("should pass undefined signal when abortSignal is not provided", async () => {
+			const mockGenerateContentStream = vitest.fn().mockResolvedValue({
+				[Symbol.asyncIterator]: async function* () {
+					yield { text: "Hello" }
+				},
+			})
+
+			handler["client"].models.generateContentStream = mockGenerateContentStream
+
+			for await (const _chunk of handler.createMessage("system", [{ role: "user", content: "Hello!" }])) {
+				break
+			}
+
+			expect(mockGenerateContentStream).toHaveBeenCalled()
+			const callArgs = mockGenerateContentStream.mock.calls[0][0]
+			expect(callArgs.signal).toBeUndefined()
+		})
+	})
 })

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

@@ -1115,4 +1115,58 @@ describe("LiteLLMHandler", () => {
 			expect(id1).not.toBe(id2)
 		})
 	})
+
+	describe("abortSignal support", () => {
+		const mockStream = {
+			async *[Symbol.asyncIterator]() {
+				yield {
+					choices: [{ delta: { content: "test response" } }],
+				}
+			},
+		}
+
+		beforeEach(() => {
+			mockCreate.mockReturnValue({
+				withResponse: vi.fn().mockResolvedValue({ data: mockStream }),
+			})
+		})
+
+		it("should pass abortSignal to chat.completions.create when provided in metadata", async () => {
+			const handler = new LiteLLMHandler(mockOptions)
+			const systemPrompt = "You are a helpful assistant."
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{ role: "user", content: [{ type: "text", text: "Hello!" }] },
+			]
+
+			const controller = new AbortController()
+			const mockAbortSignal = controller.signal
+
+			for await (const _chunk of handler.createMessage(systemPrompt, messages, {
+				taskId: "test",
+				abortSignal: mockAbortSignal,
+			})) {
+				break
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][1]
+			expect(callArgs?.signal).toBe(mockAbortSignal)
+		})
+
+		it("should pass undefined signal when abortSignal is not provided", async () => {
+			const handler = new LiteLLMHandler(mockOptions)
+			const systemPrompt = "You are a helpful assistant."
+			const messages: Anthropic.Messages.MessageParam[] = [
+				{ role: "user", content: [{ type: "text", text: "Hello!" }] },
+			]
+
+			for await (const _chunk of handler.createMessage(systemPrompt, messages)) {
+				break
+			}
+
+			expect(mockCreate).toHaveBeenCalled()
+			const callArgs = mockCreate.mock.calls[0][1]
+			expect(callArgs?.signal).toBeUndefined()
+		})
+	})
 })