Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/api/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
176 changes: 166 additions & 10 deletions src/api/providers/__tests__/anthropic-vertex.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -263,10 +263,13 @@ describe("VertexHandler", () => {
],
stream: true,
// Tools are now always present (minimum 6 from ALWAYS_AVAILABLE_TOOLS)
tools: expect.any(Array),
tool_choice: expect.any(Object),
tools: [],
tool_choice: {
disable_parallel_tool_use: false,
type: "auto",
},
}),
undefined,
{},
)
})

Expand Down Expand Up @@ -481,7 +484,7 @@ describe("VertexHandler", () => {
}),
],
}),
undefined,
{},
)
})

Expand Down Expand Up @@ -1156,7 +1159,7 @@ describe("VertexHandler", () => {
}

// Verify the API was called without the beta header
expect(mockCreate).toHaveBeenCalledWith(expect.anything(), undefined)
expect(mockCreate).toHaveBeenCalledWith(expect.anything(), {})
})
})

Expand Down Expand Up @@ -1246,7 +1249,7 @@ describe("VertexHandler", () => {
thinking: { type: "enabled", budget_tokens: 4096 },
temperature: 1.0, // Thinking requires temperature 1.0
}),
undefined,
{},
)
})

Expand All @@ -1273,7 +1276,7 @@ describe("VertexHandler", () => {
expect.objectContaining({
thinking: { type: "adaptive" },
}),
undefined,
{},
)

const request = mockCreate.mock.calls[0][0]
Expand Down Expand Up @@ -1302,7 +1305,7 @@ describe("VertexHandler", () => {
expect.objectContaining({
thinking: { type: "adaptive" },
}),
undefined,
{},
)

const request = mockCreate.mock.calls[0][0]
Expand Down Expand Up @@ -1393,7 +1396,7 @@ describe("VertexHandler", () => {
]),
tool_choice: { type: "auto", disable_parallel_tool_use: false },
}),
undefined,
{},
)
})

Expand Down Expand Up @@ -1446,7 +1449,7 @@ describe("VertexHandler", () => {
}),
]),
}),
undefined,
{},
)
})

Expand Down Expand Up @@ -1611,4 +1614,157 @@ describe("VertexHandler", () => {
})
})
})

describe("abort signal", () => {
it("should handle abort signal triggered during request", async () => {
const controller = new AbortController()
const handler = new AnthropicVertexHandler({
apiModelId: "claude-3-sonnet",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})

const mockStream = async function* () {
await new Promise((resolve) => setTimeout(resolve, 10))
if (controller.signal.aborted) {
throw new Error("AbortError: The operation was aborted")
}
yield {
type: "message_start",
message: { usage: { input_tokens: 10, output_tokens: 0 } },
}
}

;(handler["client"].messages as any).create = vitest.fn().mockResolvedValue(mockStream())

const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }], {
taskId: "test",
tools: [],
abortSignal: controller.signal,
})

const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}

expect(chunks.length).toBeGreaterThan(0)
})
Comment on lines +1619 to +1652

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Abort-path test is not actually aborting.

Line 1619 names a mid-request abort scenario, but controller.abort() is never triggered. This currently validates normal streaming only, not cancellation behavior.

Suggested fix
 it("should handle abort signal triggered during request", async () => {
   const controller = new AbortController()
+  let capturedRequestOptions: any

-  ;(handler["client"].messages as any).create = vitest.fn().mockResolvedValue(mockStream())
+  ;(handler["client"].messages as any).create = vitest.fn().mockImplementation(async (_params, requestOptions) => {
+    capturedRequestOptions = requestOptions
+    return mockStream()
+  })

   const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }], {
     taskId: "test",
     tools: [],
     abortSignal: controller.signal,
   })
+
+  setTimeout(() => controller.abort(), 20)

-  const chunks: any[] = []
-  for await (const chunk of stream) {
-    chunks.push(chunk)
-  }
-
-  expect(chunks.length).toBeGreaterThan(0)
+  await expect(async () => {
+    for await (const _chunk of stream) {}
+  }).rejects.toThrow(/abort/i)
+  expect(capturedRequestOptions).toHaveProperty("signal", controller.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__/anthropic-vertex.spec.ts` around lines 1619 -
1652, The test "should handle abort signal triggered during request" at the
AnthropicVertexHandler test file is not actually testing the abort scenario
because controller.abort() is never called during the test execution. To fix
this, add a call to controller.abort() inside the for-await loop that iterates
through the stream chunks to actually trigger the abort signal during streaming.
This will ensure the mockStream function's abort check is evaluated when the
signal is truly aborted, validating the cancellation behavior that the test is
supposed to cover.


it("should not pass signal when abortSignal is undefined", async () => {
const handler = new AnthropicVertexHandler({
apiModelId: "claude-3-sonnet",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})

const mockStream = async function* () {
yield {
type: "message_start",
message: { usage: { input_tokens: 10, output_tokens: 5 } },
}
yield {
type: "content_block_start",
content_block: { type: "text", text: "" },
}
yield {
type: "content_block_delta",
delta: { type: "text_delta", text: "response" },
}
}

;(handler["client"].messages as any).create = vitest.fn().mockResolvedValue(mockStream())

const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }])

const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}

expect(chunks.length).toBeGreaterThan(0)
})
Comment on lines +1654 to +1686

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

“No abortSignal” test should assert request options, not only chunks.

Line 1654 currently verifies output only; it does not prove signal omission. Capture/assert the second SDK arg (for this provider, {} without signal) to prevent false positives.

🤖 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-vertex.spec.ts` around lines 1654 -
1686, The test "should not pass signal when abortSignal is undefined" currently
only verifies that chunks are received but does not actually assert that the
signal was omitted from the request. After the mocked
handler["client"].messages.create is called within the loop consuming the
stream, capture the arguments passed to this mock and assert that the second
argument (the options/config object passed to the SDK) is either an empty object
or explicitly does not contain a signal property. This will ensure the test
actually validates that abortSignal is not passed when undefined, rather than
only checking the output.


it("should abort immediately if signal is already aborted", async () => {
const controller = new AbortController()
controller.abort()

const testHandler = new AnthropicVertexHandler({
apiModelId: "claude-3-sonnet",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})

testHandler["client"].messages.create = vitest.fn().mockImplementation(async (options, requestOptions) => {
// Verify that the signal was passed and is already aborted
expect(requestOptions).toHaveProperty("signal", controller.signal)
expect(controller.signal.aborted).toBe(true)

return {
[Symbol.asyncIterator]: async function* () {
if (controller.signal.aborted) {
throw new Error("AbortError: The operation was aborted")
}
yield {
type: "message_start",
message: { usage: { input_tokens: 10, output_tokens: 5 } },
}
},
}
})

const stream = testHandler.createMessage("system", [{ role: "user", content: "Hello" }], {
taskId: "test",
tools: [],
abortSignal: controller.signal,
})

await expect(async () => {
for await (const _chunk of stream) {
// consume stream
}
}).rejects.toThrow(/abort/i)
})

it("should pass signal when provided", async () => {
const controller = new AbortController()
let capturedRequestOptions: any

const testHandler = new AnthropicVertexHandler({
apiModelId: "claude-3-sonnet",
vertexProjectId: "test-project",
vertexRegion: "us-central1",
})

testHandler["client"].messages.create = vitest.fn().mockImplementation(async (options, requestOptions) => {
capturedRequestOptions = requestOptions
return {
[Symbol.asyncIterator]: async function* () {
yield {
type: "message_start",
message: { usage: { input_tokens: 10, output_tokens: 5 } },
}
yield {
type: "content_block_delta",
delta: { type: "text_delta", text: "response" },
}
},
}
})

const stream = testHandler.createMessage("system", [{ role: "user", content: "Hello" }], {
taskId: "test",
tools: [],
abortSignal: controller.signal,
})

const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}

expect(chunks.length).toBeGreaterThan(0)
expect(capturedRequestOptions).toHaveProperty("signal", controller.signal)
})
})
})
74 changes: 74 additions & 0 deletions src/api/providers/__tests__/anthropic.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1057,4 +1057,78 @@ describe("AnthropicHandler", () => {
})
})
})

describe("abort signal", () => {
it("should pass abortSignal to the SDK options", async () => {
const controller = new AbortController()

mockCreate.mockImplementation(async (options, requestOptions) => {
// Verify that the signal was passed
expect(requestOptions).toHaveProperty("signal", controller.signal)
return {
async *[Symbol.asyncIterator]() {
yield {
type: "message_start",
message: { usage: { input_tokens: 10, output_tokens: 5 } },
}
yield {
type: "content_block_delta",
delta: { type: "text_delta", text: "response" },
}
},
}
})

const handler = new AnthropicHandler(mockOptions)
const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }], {
taskId: "test",
tools: [],
abortSignal: controller.signal,
})

const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}

expect(chunks.length).toBeGreaterThan(0)
})

it("should work normally without abortSignal", async () => {
const handler = new AnthropicHandler(mockOptions)
const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }])

const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}

expect(chunks.length).toBeGreaterThan(0)
})

it("should not pass signal when abortSignal is undefined", async () => {
mockCreate.mockImplementation(async (options, requestOptions) => {
// When no abortSignal is provided, requestOptions should be undefined or not have signal
expect(requestOptions).toBeUndefined()
return {
async *[Symbol.asyncIterator]() {
yield {
type: "message_start",
message: { usage: { input_tokens: 10, output_tokens: 5 } },
}
},
}
})

const handler = new AnthropicHandler(mockOptions)
const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }])

const chunks: any[] = []
for await (const chunk of stream) {
chunks.push(chunk)
}

expect(chunks.length).toBeGreaterThan(0)
})
})
})
Loading
Loading