Skip to content

Commit 37cff20

Browse files
easonliang28easonliang28
committed
feat: add abort signal pass-through to providers
1 parent ccf07eb commit 37cff20

41 files changed

Lines changed: 2170 additions & 48 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎src/api/index.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,12 @@ export interface ApiHandlerCreateMessageMetadata {
9090
* Only applies to providers that support function calling restrictions (e.g., Gemini).
9191
*/
9292
allowedFunctionNames?: string[]
93+
/**
94+
* Abort signal for cancelling the HTTP request mid-stream.
95+
* Passed through to AI SDK's streamText() so the underlying HTTP request is aborted
96+
* when the user clicks stop, preventing wasted API tokens/compute on the provider side.
97+
*/
98+
abortSignal?: AbortSignal
9399
}
94100

95101
export interface ApiHandler {

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

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1611,4 +1611,157 @@ describe("VertexHandler", () => {
16111611
})
16121612
})
16131613
})
1614+
1615+
describe("abort signal", () => {
1616+
it("should handle abort signal triggered during request", async () => {
1617+
const controller = new AbortController()
1618+
const handler = new AnthropicVertexHandler({
1619+
apiModelId: "claude-3-sonnet",
1620+
vertexProjectId: "test-project",
1621+
vertexRegion: "us-central1",
1622+
})
1623+
1624+
const mockStream = async function* () {
1625+
await new Promise((resolve) => setTimeout(resolve, 10))
1626+
if (controller.signal.aborted) {
1627+
throw new Error("AbortError: The operation was aborted")
1628+
}
1629+
yield {
1630+
type: "message_start",
1631+
message: { usage: { input_tokens: 10, output_tokens: 0 } },
1632+
}
1633+
}
1634+
1635+
;(handler["client"].messages as any).create = vitest.fn().mockResolvedValue(mockStream())
1636+
1637+
const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }], {
1638+
taskId: "test",
1639+
tools: [],
1640+
abortSignal: controller.signal,
1641+
})
1642+
1643+
const chunks: any[] = []
1644+
for await (const chunk of stream) {
1645+
chunks.push(chunk)
1646+
}
1647+
1648+
expect(chunks.length).toBeGreaterThan(0)
1649+
})
1650+
1651+
it("should not pass signal when abortSignal is undefined", async () => {
1652+
const handler = new AnthropicVertexHandler({
1653+
apiModelId: "claude-3-sonnet",
1654+
vertexProjectId: "test-project",
1655+
vertexRegion: "us-central1",
1656+
})
1657+
1658+
const mockStream = async function* () {
1659+
yield {
1660+
type: "message_start",
1661+
message: { usage: { input_tokens: 10, output_tokens: 5 } },
1662+
}
1663+
yield {
1664+
type: "content_block_start",
1665+
content_block: { type: "text", text: "" },
1666+
}
1667+
yield {
1668+
type: "content_block_delta",
1669+
delta: { type: "text_delta", text: "response" },
1670+
}
1671+
}
1672+
1673+
;(handler["client"].messages as any).create = vitest.fn().mockResolvedValue(mockStream())
1674+
1675+
const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }])
1676+
1677+
const chunks: any[] = []
1678+
for await (const chunk of stream) {
1679+
chunks.push(chunk)
1680+
}
1681+
1682+
expect(chunks.length).toBeGreaterThan(0)
1683+
})
1684+
1685+
it("should abort immediately if signal is already aborted", async () => {
1686+
const controller = new AbortController()
1687+
controller.abort()
1688+
1689+
const testHandler = new AnthropicVertexHandler({
1690+
apiModelId: "claude-3-sonnet",
1691+
vertexProjectId: "test-project",
1692+
vertexRegion: "us-central1",
1693+
})
1694+
1695+
testHandler["client"].messages.create = vitest.fn().mockImplementation(async (options, requestOptions) => {
1696+
// Verify that the signal was passed and is already aborted
1697+
expect(requestOptions).toHaveProperty("signal", controller.signal)
1698+
expect(controller.signal.aborted).toBe(true)
1699+
1700+
return {
1701+
[Symbol.asyncIterator]: async function* () {
1702+
if (controller.signal.aborted) {
1703+
throw new Error("AbortError: The operation was aborted")
1704+
}
1705+
yield {
1706+
type: "message_start",
1707+
message: { usage: { input_tokens: 10, output_tokens: 5 } },
1708+
}
1709+
},
1710+
}
1711+
})
1712+
1713+
const stream = testHandler.createMessage("system", [{ role: "user", content: "Hello" }], {
1714+
taskId: "test",
1715+
tools: [],
1716+
abortSignal: controller.signal,
1717+
})
1718+
1719+
await expect(async () => {
1720+
for await (const _chunk of stream) {
1721+
// consume stream
1722+
}
1723+
}).rejects.toThrow(/abort/i)
1724+
})
1725+
1726+
it("should pass signal when provided", async () => {
1727+
const controller = new AbortController()
1728+
let capturedRequestOptions: any
1729+
1730+
const testHandler = new AnthropicVertexHandler({
1731+
apiModelId: "claude-3-sonnet",
1732+
vertexProjectId: "test-project",
1733+
vertexRegion: "us-central1",
1734+
})
1735+
1736+
testHandler["client"].messages.create = vitest.fn().mockImplementation(async (options, requestOptions) => {
1737+
capturedRequestOptions = requestOptions
1738+
return {
1739+
[Symbol.asyncIterator]: async function* () {
1740+
yield {
1741+
type: "message_start",
1742+
message: { usage: { input_tokens: 10, output_tokens: 5 } },
1743+
}
1744+
yield {
1745+
type: "content_block_delta",
1746+
delta: { type: "text_delta", text: "response" },
1747+
}
1748+
},
1749+
}
1750+
})
1751+
1752+
const stream = testHandler.createMessage("system", [{ role: "user", content: "Hello" }], {
1753+
taskId: "test",
1754+
tools: [],
1755+
abortSignal: controller.signal,
1756+
})
1757+
1758+
const chunks: any[] = []
1759+
for await (const chunk of stream) {
1760+
chunks.push(chunk)
1761+
}
1762+
1763+
expect(chunks.length).toBeGreaterThan(0)
1764+
expect(capturedRequestOptions).toHaveProperty("signal", controller.signal)
1765+
})
1766+
})
16141767
})

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

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1057,4 +1057,78 @@ describe("AnthropicHandler", () => {
10571057
})
10581058
})
10591059
})
1060+
1061+
describe("abort signal", () => {
1062+
it("should pass abortSignal to the SDK options", async () => {
1063+
const controller = new AbortController()
1064+
1065+
mockCreate.mockImplementation(async (options, requestOptions) => {
1066+
// Verify that the signal was passed
1067+
expect(requestOptions).toHaveProperty("signal", controller.signal)
1068+
return {
1069+
async *[Symbol.asyncIterator]() {
1070+
yield {
1071+
type: "message_start",
1072+
message: { usage: { input_tokens: 10, output_tokens: 5 } },
1073+
}
1074+
yield {
1075+
type: "content_block_delta",
1076+
delta: { type: "text_delta", text: "response" },
1077+
}
1078+
},
1079+
}
1080+
})
1081+
1082+
const handler = new AnthropicHandler(mockOptions)
1083+
const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }], {
1084+
taskId: "test",
1085+
tools: [],
1086+
abortSignal: controller.signal,
1087+
})
1088+
1089+
const chunks: any[] = []
1090+
for await (const chunk of stream) {
1091+
chunks.push(chunk)
1092+
}
1093+
1094+
expect(chunks.length).toBeGreaterThan(0)
1095+
})
1096+
1097+
it("should work normally without abortSignal", async () => {
1098+
const handler = new AnthropicHandler(mockOptions)
1099+
const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }])
1100+
1101+
const chunks: any[] = []
1102+
for await (const chunk of stream) {
1103+
chunks.push(chunk)
1104+
}
1105+
1106+
expect(chunks.length).toBeGreaterThan(0)
1107+
})
1108+
1109+
it("should not pass signal when abortSignal is undefined", async () => {
1110+
mockCreate.mockImplementation(async (options, requestOptions) => {
1111+
// When no abortSignal is provided, requestOptions should be undefined or not have signal
1112+
expect(requestOptions).toBeUndefined()
1113+
return {
1114+
async *[Symbol.asyncIterator]() {
1115+
yield {
1116+
type: "message_start",
1117+
message: { usage: { input_tokens: 10, output_tokens: 5 } },
1118+
}
1119+
},
1120+
}
1121+
})
1122+
1123+
const handler = new AnthropicHandler(mockOptions)
1124+
const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }])
1125+
1126+
const chunks: any[] = []
1127+
for await (const chunk of stream) {
1128+
chunks.push(chunk)
1129+
}
1130+
1131+
expect(chunks.length).toBeGreaterThan(0)
1132+
})
1133+
})
10601134
})

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

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -551,4 +551,92 @@ describe("BaseOpenAiCompatibleProvider", () => {
551551
expect(endChunks).toHaveLength(0)
552552
})
553553
})
554+
555+
describe("abort signal", () => {
556+
it("should handle abort signal triggered during request", async () => {
557+
const controller = new AbortController()
558+
handler = new TestOpenAiCompatibleProvider("test-api-key")
559+
560+
mockCreate.mockImplementation(async (options: unknown) => {
561+
return {
562+
async *[Symbol.asyncIterator]() {
563+
while (!controller.signal.aborted) {
564+
await new Promise((resolve) => setTimeout(resolve, 10))
565+
if (controller.signal.aborted) {
566+
throw new Error("AbortError: The operation was aborted")
567+
}
568+
yield {
569+
choices: [{ delta: { content: "response" } }],
570+
usage: null,
571+
}
572+
}
573+
},
574+
}
575+
})
576+
577+
const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }] as any, {
578+
taskId: "test",
579+
tools: [],
580+
abortSignal: controller.signal,
581+
})
582+
583+
setTimeout(() => controller.abort(), 50)
584+
585+
await expect(async () => {
586+
for await (const _chunk of stream) {
587+
// consume stream
588+
}
589+
}).rejects.toThrow(/abort/i)
590+
})
591+
592+
it("should work normally without abortSignal", async () => {
593+
handler = new TestOpenAiCompatibleProvider("test-api-key")
594+
595+
mockCreate.mockResolvedValue({
596+
async *[Symbol.asyncIterator]() {
597+
yield { choices: [{ delta: { content: "Hello" } }], usage: null }
598+
yield { choices: [{ delta: {} }], usage: { prompt_tokens: 10, completion_tokens: 5 } }
599+
},
600+
})
601+
602+
const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }] as any)
603+
604+
const chunks: any[] = []
605+
for await (const chunk of stream) {
606+
chunks.push(chunk)
607+
}
608+
609+
expect(chunks.length).toBeGreaterThan(0)
610+
})
611+
612+
it("should abort immediately if signal is already aborted", async () => {
613+
const controller = new AbortController()
614+
controller.abort()
615+
616+
handler = new TestOpenAiCompatibleProvider("test-api-key")
617+
618+
mockCreate.mockImplementation(async (options: unknown) => {
619+
return {
620+
async *[Symbol.asyncIterator]() {
621+
if (controller.signal.aborted) {
622+
throw new Error("AbortError: The operation was aborted")
623+
}
624+
yield { choices: [{ delta: { content: "response" } }], usage: null }
625+
},
626+
}
627+
})
628+
629+
const stream = handler.createMessage("system", [{ role: "user", content: "Hello" }] as any, {
630+
taskId: "test",
631+
tools: [],
632+
abortSignal: controller.signal,
633+
})
634+
635+
await expect(async () => {
636+
for await (const _chunk of stream) {
637+
// consume stream
638+
}
639+
}).rejects.toThrow(/abort/i)
640+
})
641+
})
554642
})

0 commit comments

Comments
 (0)