Skip to content

Commit e72f4d8

Browse files
easonliang28easonliang28
committed
fix(api): standardize timeoutMs handling across all providers and fix signal propagation
1 parent 899e717 commit e72f4d8

25 files changed

Lines changed: 180 additions & 89 deletions

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -946,11 +946,11 @@ describe("VertexHandler", () => {
946946
await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 5000 })
947947
expect(mockCreate).toHaveBeenCalledWith(
948948
expect.objectContaining({ model: expect.any(String) }),
949-
{ signal: controller.signal }, // only signal is passed, not timeoutMs
949+
expect.objectContaining({ signal: controller.signal, timeout: 5000 }),
950950
)
951951
})
952952

953-
it("completePrompt should not pass timeoutMs when no signal provided", async () => {
953+
it("completePrompt should pass timeoutMs when provided", async () => {
954954
const mockCreate = vitest.fn().mockResolvedValue({
955955
content: [{ type: "text", text: "response" }],
956956
})
@@ -959,7 +959,7 @@ describe("VertexHandler", () => {
959959
await handler.completePrompt("test prompt", { timeoutMs: 3000 })
960960
expect(mockCreate).toHaveBeenCalledWith(
961961
expect.objectContaining({ model: expect.any(String) }),
962-
undefined, // anthropic-vertex only passes signal, not timeoutMs
962+
expect.objectContaining({ timeout: 3000 }),
963963
)
964964
})
965965
})

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

Lines changed: 7 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -519,35 +519,14 @@ describe("AnthropicHandler", () => {
519519
)
520520
})
521521

522-
it("should trigger timeout when timeoutMs elapses before request completes", async () => {
523-
const mockCreateWithTimeout = vitest
524-
.fn()
525-
.mockImplementation(
526-
async () =>
527-
new Promise((resolve) =>
528-
setTimeout(() => resolve({ content: [{ type: "text", text: "response" }] }), 500),
529-
),
530-
)
531-
532-
const handlerTimeout = new AnthropicHandler(mockOptions)
533-
// Replace the mock on the existing handler's client
534-
handlerTimeout["client"].messages.create = mockCreateWithTimeout
535-
522+
it("should pass timeoutMs through to client alongside abortSignal", async () => {
536523
const controller = new AbortController()
537-
let timeoutTriggered = false
538-
handlerTimeout
539-
.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 50 })
540-
.catch(() => {
541-
timeoutTriggered = true
542-
})
543-
544-
// Wait for timeout to trigger (50ms timeout + buffer)
545-
await new Promise((resolve) => setTimeout(resolve, 150))
546-
547-
// Verify the API was called with timeout options
548-
expect(mockCreateWithTimeout).toHaveBeenCalled()
549-
// User signal should not be aborted (timeout mechanism aborts its own internal signal)
550-
expect(controller.signal.aborted).toBe(false)
524+
mockCreate.mockResolvedValueOnce({ content: [{ type: "text", text: "response" }] })
525+
await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 5000 })
526+
expect(mockCreate).toHaveBeenCalledWith(
527+
expect.objectContaining({ model: mockOptions.apiModelId }),
528+
expect.objectContaining({ signal: controller.signal, timeout: 5000 }),
529+
)
551530
})
552531

553532
it("should pass the same signal instance", async () => {

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

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1646,8 +1646,36 @@ describe("AwsBedrockHandler", () => {
16461646

16471647
await handler.completePrompt("test prompt", { timeoutMs: 5000 })
16481648

1649-
// bedrock.ts uses truthy check for timeoutMs, so it creates AbortSignal.timeout
16501649
expect(mockSend).toHaveBeenCalled()
1650+
// Verify the second argument (sendOptions) contains an abortSignal derived from timeoutMs
1651+
const sendOptions = mockSend.mock.calls[0][1]
1652+
expect(sendOptions).toBeDefined()
1653+
expect(sendOptions?.abortSignal).toBeDefined()
1654+
})
1655+
1656+
it("completePrompt should merge abortSignal and timeoutMs", async () => {
1657+
const mockSend = vi.fn()
1658+
1659+
const handler = new AwsBedrockHandler({
1660+
apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0",
1661+
awsAccessKey: "test-access-key",
1662+
awsSecretKey: "test-secret-key",
1663+
awsRegion: "us-east-1",
1664+
})
1665+
1666+
const clientInstance = (handler as any).client
1667+
clientInstance.send = mockSend
1668+
1669+
mockSend.mockResolvedValueOnce({
1670+
output: { message: { content: [{ type: "text", text: "response" }] }, stopReason: null },
1671+
})
1672+
1673+
const controller = new AbortController()
1674+
await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 5000 })
1675+
1676+
expect(mockSend).toHaveBeenCalled()
1677+
const sendOptions = mockSend.mock.calls[0][1]
1678+
expect(sendOptions?.abortSignal).toBeDefined()
16511679
})
16521680
})
16531681
})

‎src/api/providers/__tests__/gemini-handler.spec.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ describe("GeminiHandler backend support", () => {
7272
expect(stub).toHaveBeenCalledWith(
7373
expect.objectContaining({
7474
config: expect.objectContaining({
75-
httpOptions: { signal: controller.signal },
75+
abortSignal: controller.signal,
7676
}),
7777
}),
7878
)

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

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -157,15 +157,16 @@ describe("GeminiHandler", () => {
157157
expect(result).toBe("")
158158
})
159159

160-
it("should pass abort signal through to client via httpOptions", async () => {
160+
it("should pass abort signal through to client via config.abortSignal", async () => {
161161
const controller = new AbortController()
162162
;(handler["client"].models.generateContent as any).mockResolvedValue({ text: "response" })
163163
await handler.completePrompt("test prompt", { abortSignal: controller.signal })
164164
expect(handler["client"].models.generateContent).toHaveBeenCalledWith({
165165
model: GEMINI_MODEL_NAME,
166166
contents: [{ role: "user", parts: [{ text: "test prompt" }] }],
167167
config: {
168-
httpOptions: { signal: controller.signal },
168+
abortSignal: controller.signal,
169+
httpOptions: undefined,
169170
temperature: 1,
170171
},
171172
})
@@ -185,15 +186,16 @@ describe("GeminiHandler", () => {
185186
})
186187
})
187188

188-
it("should pass timeoutMs through to client via httpOptions", async () => {
189+
it("should pass timeoutMs through to client via httpOptions with abortSignal on config", async () => {
189190
const controller = new AbortController()
190191
;(handler["client"].models.generateContent as any).mockResolvedValue({ text: "response" })
191192
await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 10000 })
192193
expect(handler["client"].models.generateContent).toHaveBeenCalledWith({
193194
model: GEMINI_MODEL_NAME,
194195
contents: [{ role: "user", parts: [{ text: "test prompt" }] }],
195196
config: {
196-
httpOptions: { signal: controller.signal, timeout: 10000 },
197+
abortSignal: controller.signal,
198+
httpOptions: { timeout: 10000 },
197199
temperature: 1,
198200
},
199201
})

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -266,14 +266,14 @@ describe("MiniMaxHandler", () => {
266266
})
267267
})
268268

269-
it("should not set timeout when timeoutMs=0 (truthy check)", async () => {
269+
it("should pass timeout when timeoutMs=0 (defined check)", async () => {
270270
mockCreate.mockResolvedValueOnce({
271271
content: [{ type: "text", text: "response" }],
272272
})
273273
await handler.completePrompt("test prompt", { timeoutMs: 0 })
274274
expect(mockCreate).toHaveBeenCalledWith(
275275
expect.objectContaining({ model: expect.any(String) }),
276-
undefined, // truthy check means 0 is falsy
276+
{ timeout: 0 }, // !== undefined check means 0 is passed through
277277
)
278278
})
279279

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ describe("MoonshotHandler", () => {
285285
)
286286
})
287287

288-
it("should not set abortSignal when timeoutMs is 0", async () => {
288+
it("should set an immediately-aborted abortSignal when timeoutMs is 0", async () => {
289289
mockGenerateText.mockResolvedValueOnce({ text: "response" })
290290

291291
await handler.completePrompt("test prompt", { timeoutMs: 0 })
@@ -295,7 +295,9 @@ describe("MoonshotHandler", () => {
295295
}),
296296
)
297297
const callArgs = mockGenerateText.mock.calls[0][0]
298-
expect(callArgs.abortSignal).toBeUndefined()
298+
// With !== undefined check, timeoutMs=0 creates an immediately-aborted signal
299+
expect(callArgs.abortSignal).toBeDefined()
300+
expect(callArgs.abortSignal.aborted).toBe(true)
299301
})
300302

301303
it("should propagate errors from generateText", async () => {

‎src/api/providers/__tests__/request-config-builder.spec.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ describe("RequestConfigBuilder", () => {
6666
const controller1 = new AbortController()
6767
const controller2 = new AbortController()
6868

69-
const builder = new RequestConfigBuilder({ abortSignal: controller1.signal })
69+
const builder = new RequestConfigBuilder({ signal: controller1.signal })
7070
builder.addAbortSignal({
7171
taskId: "test-task",
7272
abortSignal: controller2.signal,

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ describe("VertexHandler", () => {
138138
expect(result).toBe("")
139139
})
140140

141-
it("should pass abort signal through to client via httpOptions", async () => {
141+
it("should pass abort signal through to client via config.abortSignal", async () => {
142142
const controller = new AbortController()
143143
;(handler["client"].models.generateContent as any).mockResolvedValue({
144144
text: "response",
@@ -150,7 +150,8 @@ describe("VertexHandler", () => {
150150
model: expect.any(String),
151151
contents: [{ role: "user", parts: [{ text: "test prompt" }] }],
152152
config: expect.objectContaining({
153-
httpOptions: { signal: controller.signal },
153+
abortSignal: controller.signal,
154+
httpOptions: undefined,
154155
temperature: 1,
155156
}),
156157
}),

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -296,13 +296,13 @@ describe("XAIHandler", () => {
296296
})
297297
})
298298

299-
it("completePrompt should not set timeout when timeoutMs=0 (truthy check)", async () => {
299+
it("completePrompt should pass timeout when timeoutMs=0 (defined check)", async () => {
300300
mockResponsesCreate.mockResolvedValueOnce({ output_text: "response" })
301301

302302
await handler.completePrompt("test prompt", { timeoutMs: 0 })
303303
expect(mockResponsesCreate).toHaveBeenCalledWith(
304304
expect.objectContaining({ model: expect.any(String) }),
305-
undefined, // truthy check means 0 is falsy
305+
{ timeout: 0 }, // !== undefined check means 0 is passed through
306306
)
307307
})
308308

0 commit comments

Comments
 (0)