Skip to content

Commit b7156e8

Browse files
proyectoauraorgproyectoauraorg
committed
fix(api): propagate AbortSignal to OpenAI-compatible providers for mid-stream cancellation
Both BaseOpenAiCompatibleProvider and OpenAICompatibleHandler did not create an AbortController or pass a signal to their underlying SDK calls. When a user hit Stop, the Task-level Promise.race would break out of the chunk loop, but the HTTP request continued running server-side until it finished naturally. This fix adds an AbortController to both providers: - BaseOpenAiCompatibleProvider: signal passed via requestOptions to chat.completions.create() and checked in the stream iteration loop - OpenAICompatibleHandler: abortSignal passed to streamText() and generateText() from the AI SDK - Both providers clean up the controller in finally blocks - completePrompt() also gets abort support in both providers Includes 5 regression tests covering signal propagation, mid-stream abort, and controller cleanup. Closes: #404
1 parent 3df406e commit b7156e8

3 files changed

Lines changed: 284 additions & 106 deletions

File tree

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

Lines changed: 131 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -354,7 +354,7 @@ describe("BaseOpenAiCompatibleProvider", () => {
354354
stream: true,
355355
stream_options: { include_usage: true },
356356
}),
357-
undefined,
357+
expect.objectContaining({ signal: expect.any(AbortSignal) }),
358358
)
359359
})
360360

@@ -545,4 +545,134 @@ describe("BaseOpenAiCompatibleProvider", () => {
545545
expect(endChunks).toHaveLength(0)
546546
})
547547
})
548+
549+
describe("Abort/cancellation support (fixes #404)", () => {
550+
it("should pass abort signal to chat.completions.create in createMessage", async () => {
551+
mockCreate.mockImplementationOnce(() => {
552+
return {
553+
[Symbol.asyncIterator]: () => ({
554+
next: vi
555+
.fn()
556+
.mockResolvedValueOnce({
557+
done: false,
558+
value: { choices: [{ delta: { content: "Hello" } }] },
559+
})
560+
.mockResolvedValueOnce({ done: true }),
561+
}),
562+
}
563+
})
564+
565+
const stream = handler.createMessage("system prompt", [])
566+
const chunks = []
567+
for await (const chunk of stream) {
568+
chunks.push(chunk)
569+
}
570+
571+
// Verify signal was passed to the SDK
572+
expect(mockCreate).toHaveBeenCalledWith(
573+
expect.objectContaining({ model: "test-model" }),
574+
expect.objectContaining({ signal: expect.any(AbortSignal) }),
575+
)
576+
})
577+
578+
it("should stop yielding chunks when abort is signaled mid-stream", async () => {
579+
let resolveSecondChunk: () => void
580+
const secondChunkPromise = new Promise<void>((resolve) => {
581+
resolveSecondChunk = resolve
582+
})
583+
584+
mockCreate.mockImplementationOnce(() => {
585+
let callCount = 0
586+
return {
587+
[Symbol.asyncIterator]: () => ({
588+
next: vi.fn().mockImplementation(async () => {
589+
callCount++
590+
if (callCount === 1) {
591+
return {
592+
done: false,
593+
value: { choices: [{ delta: { content: "Before abort" } }] },
594+
}
595+
}
596+
// Wait until the test signals abort
597+
await secondChunkPromise
598+
return {
599+
done: false,
600+
value: { choices: [{ delta: { content: "After abort" } }] },
601+
}
602+
}),
603+
}),
604+
}
605+
})
606+
607+
const stream = handler.createMessage("system prompt", [])
608+
609+
// Collect chunks with abort after first
610+
const chunks: any[] = []
611+
const iterator = stream[Symbol.asyncIterator]()
612+
613+
// Get first chunk
614+
const first = await iterator.next()
615+
chunks.push(first.value)
616+
617+
// Access the private abortController to signal abort
618+
// The abortController is created in createMessage, so it exists now
619+
const controller = (handler as any).abortController as AbortController
620+
expect(controller).toBeDefined()
621+
controller.abort()
622+
resolveSecondChunk!()
623+
624+
// The stream should stop after abort
625+
const second = await iterator.next()
626+
// After abort, the for-await loop breaks, so we get done: true
627+
expect(second.done).toBe(true)
628+
})
629+
630+
it("should pass abort signal to chat.completions.create in completePrompt", async () => {
631+
mockCreate.mockResolvedValueOnce({
632+
choices: [{ message: { content: "response" } }],
633+
})
634+
635+
await handler.completePrompt("test prompt")
636+
637+
expect(mockCreate).toHaveBeenCalledWith(
638+
expect.objectContaining({ model: "test-model" }),
639+
expect.objectContaining({ signal: expect.any(AbortSignal) }),
640+
)
641+
})
642+
643+
it("should clean up abortController after createMessage completes", async () => {
644+
mockCreate.mockImplementationOnce(() => {
645+
return {
646+
[Symbol.asyncIterator]: () => ({
647+
next: vi
648+
.fn()
649+
.mockResolvedValueOnce({
650+
done: false,
651+
value: { choices: [{ delta: { content: "done" } }] },
652+
})
653+
.mockResolvedValueOnce({ done: true }),
654+
}),
655+
}
656+
})
657+
658+
const stream = handler.createMessage("system prompt", [])
659+
for await (const _chunk of stream) {
660+
// consume stream
661+
}
662+
663+
// abortController should be cleaned up
664+
expect((handler as any).abortController).toBeUndefined()
665+
})
666+
667+
it("should clean up abortController after completePrompt completes", async () => {
668+
mockCreate.mockResolvedValueOnce({
669+
choices: [{ message: { content: "response" } }],
670+
})
671+
672+
await handler.completePrompt("test prompt")
673+
674+
// abortController should be cleaned up
675+
expect((handler as any).abortController).toBeUndefined()
676+
})
677+
})
548678
})

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

Lines changed: 88 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
3636
protected readonly options: ApiHandlerOptions
3737

3838
protected client: OpenAI
39+
// Abort controller for cancelling ongoing requests (fixes #404)
40+
private abortController?: AbortController
3941

4042
constructor({
4143
providerName,
@@ -115,87 +117,103 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
115117
messages: Anthropic.Messages.MessageParam[],
116118
metadata?: ApiHandlerCreateMessageMetadata,
117119
): ApiStream {
118-
const stream = await this.createStream(systemPrompt, messages, metadata)
120+
// Create AbortController for cancellation (fixes #404)
121+
this.abortController = new AbortController()
119122

120-
const matcher = new TagMatcher(
121-
"think",
122-
(chunk) =>
123-
({
124-
type: chunk.matched ? "reasoning" : "text",
125-
text: chunk.data,
126-
}) as const,
127-
)
128-
129-
let lastUsage: OpenAI.CompletionUsage | undefined
130-
const activeToolCallIds = new Set<string>()
123+
try {
124+
const stream = await this.createStream(systemPrompt, messages, metadata, {
125+
signal: this.abortController.signal,
126+
})
127+
128+
const matcher = new TagMatcher(
129+
"think",
130+
(chunk) =>
131+
({
132+
type: chunk.matched ? "reasoning" : "text",
133+
text: chunk.data,
134+
}) as const,
135+
)
136+
137+
let lastUsage: OpenAI.CompletionUsage | undefined
138+
const activeToolCallIds = new Set<string>()
139+
140+
for await (const chunk of stream) {
141+
// Check if request was aborted (fixes #404)
142+
if (this.abortController?.signal.aborted) {
143+
break
144+
}
131145

132-
for await (const chunk of stream) {
133-
// Check for provider-specific error responses (e.g., MiniMax base_resp)
134-
const chunkAny = chunk as any
135-
if (chunkAny.base_resp?.status_code && chunkAny.base_resp.status_code !== 0) {
136-
throw new Error(
137-
`${this.providerName} API Error (${chunkAny.base_resp.status_code}): ${chunkAny.base_resp.status_msg || "Unknown error"}`,
138-
)
139-
}
146+
// Check for provider-specific error responses (e.g., MiniMax base_resp)
147+
const chunkAny = chunk as any
148+
if (chunkAny.base_resp?.status_code && chunkAny.base_resp.status_code !== 0) {
149+
throw new Error(
150+
`${this.providerName} API Error (${chunkAny.base_resp.status_code}): ${chunkAny.base_resp.status_msg || "Unknown error"}`,
151+
)
152+
}
140153

141-
const delta = chunk.choices?.[0]?.delta
142-
const finishReason = chunk.choices?.[0]?.finish_reason
154+
const delta = chunk.choices?.[0]?.delta
155+
const finishReason = chunk.choices?.[0]?.finish_reason
143156

144-
if (delta?.content) {
145-
for (const processedChunk of matcher.update(delta.content)) {
146-
yield processedChunk
157+
if (delta?.content) {
158+
for (const processedChunk of matcher.update(delta.content)) {
159+
yield processedChunk
160+
}
147161
}
148-
}
149162

150-
if (delta) {
151-
for (const key of ["reasoning_content", "reasoning"] as const) {
152-
if (key in delta) {
153-
const reasoning_content = ((delta as any)[key] as string | undefined) || ""
154-
if (reasoning_content?.trim()) {
155-
yield { type: "reasoning", text: reasoning_content }
163+
if (delta) {
164+
for (const key of ["reasoning_content", "reasoning"] as const) {
165+
if (key in delta) {
166+
const reasoning_content = ((delta as any)[key] as string | undefined) || ""
167+
if (reasoning_content?.trim()) {
168+
yield { type: "reasoning", text: reasoning_content }
169+
}
170+
break
156171
}
157-
break
158172
}
159173
}
160-
}
161174

162-
// Emit raw tool call chunks - NativeToolCallParser handles state management
163-
if (delta?.tool_calls) {
164-
for (const toolCall of delta.tool_calls) {
165-
if (toolCall.id) {
166-
activeToolCallIds.add(toolCall.id)
175+
// Emit raw tool call chunks - NativeToolCallParser handles state management
176+
if (delta?.tool_calls) {
177+
for (const toolCall of delta.tool_calls) {
178+
if (toolCall.id) {
179+
activeToolCallIds.add(toolCall.id)
180+
}
181+
yield {
182+
type: "tool_call_partial",
183+
index: toolCall.index,
184+
id: toolCall.id,
185+
name: toolCall.function?.name,
186+
arguments: toolCall.function?.arguments,
187+
}
167188
}
168-
yield {
169-
type: "tool_call_partial",
170-
index: toolCall.index,
171-
id: toolCall.id,
172-
name: toolCall.function?.name,
173-
arguments: toolCall.function?.arguments,
189+
}
190+
191+
// Emit tool_call_end events when finish_reason is "tool_calls"
192+
// This ensures tool calls are finalized even if the stream doesn't properly close
193+
if (finishReason === "tool_calls" && activeToolCallIds.size > 0) {
194+
for (const id of activeToolCallIds) {
195+
yield { type: "tool_call_end", id }
174196
}
197+
activeToolCallIds.clear()
175198
}
176-
}
177199

178-
// Emit tool_call_end events when finish_reason is "tool_calls"
179-
// This ensures tool calls are finalized even if the stream doesn't properly close
180-
if (finishReason === "tool_calls" && activeToolCallIds.size > 0) {
181-
for (const id of activeToolCallIds) {
182-
yield { type: "tool_call_end", id }
200+
if (chunk.usage) {
201+
lastUsage = chunk.usage
183202
}
184-
activeToolCallIds.clear()
185203
}
186204

187-
if (chunk.usage) {
188-
lastUsage = chunk.usage
205+
if (lastUsage) {
206+
yield this.processUsageMetrics(lastUsage, this.getModel().info)
189207
}
190-
}
191208

192-
if (lastUsage) {
193-
yield this.processUsageMetrics(lastUsage, this.getModel().info)
194-
}
195-
196-
// Process any remaining content
197-
for (const processedChunk of matcher.final()) {
198-
yield processedChunk
209+
// Process any remaining content
210+
for (const processedChunk of matcher.final()) {
211+
yield processedChunk
212+
}
213+
} catch (error) {
214+
throw handleOpenAIError(error, this.providerName)
215+
} finally {
216+
this.abortController = undefined
199217
}
200218
}
201219

@@ -222,6 +240,9 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
222240
async completePrompt(prompt: string): Promise<string> {
223241
const { id: modelId, info: modelInfo } = this.getModel()
224242

243+
// Create AbortController for cancellation (fixes #404)
244+
this.abortController = new AbortController()
245+
225246
const params: OpenAI.Chat.Completions.ChatCompletionCreateParams = {
226247
model: modelId,
227248
messages: [{ role: "user", content: prompt }],
@@ -233,7 +254,9 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
233254
}
234255

235256
try {
236-
const response = await this.client.chat.completions.create(params)
257+
const response = await this.client.chat.completions.create(params, {
258+
signal: this.abortController.signal,
259+
})
237260

238261
// Check for provider-specific error responses (e.g., MiniMax base_resp)
239262
const responseAny = response as any
@@ -246,6 +269,8 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
246269
return response.choices?.[0]?.message.content || ""
247270
} catch (error) {
248271
throw handleOpenAIError(error, this.providerName)
272+
} finally {
273+
this.abortController = undefined
249274
}
250275
}
251276

0 commit comments

Comments
 (0)