Skip to content

Commit 9b97337

Browse files
committed
feat(api): abort signal core plumbing (#615)
- Add abortSignal?: AbortSignal to ApiHandlerCreateMessageMetadata interface - Update SingleCompletionHandler.completePrompt to accept metadata parameter - Move AbortController creation before metadata construction in Task.ts - Include abortSignal directly in metadata object literal (not post-mutation) - Forward metadata through single-completion-handler.ts - Add tests for signal identity and fresh controller per request
1 parent d3ba52d commit 9b97337

5 files changed

Lines changed: 58 additions & 11 deletions

File tree

src/api/index.ts

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ import {
4141
import { NativeOllamaHandler } from "./providers/native-ollama"
4242

4343
export interface SingleCompletionHandler {
44-
completePrompt(prompt: string): Promise<string>
44+
completePrompt(prompt: string, metadata?: ApiHandlerCreateMessageMetadata): Promise<string>
4545
}
4646

4747
export interface ApiHandlerCreateMessageMetadata {
@@ -90,6 +90,11 @@ export interface ApiHandlerCreateMessageMetadata {
9090
* Only applies to providers that support function calling restrictions (e.g., Gemini).
9191
*/
9292
allowedFunctionNames?: string[]
93+
/**
94+
* Abort signal from the Task's AbortController, used by providers to cancel
95+
* in-flight HTTP requests when the user presses Stop.
96+
*/
97+
abortSignal?: AbortSignal
9398
}
9499

95100
export interface ApiHandler {
Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
import { describe, it, expect } from "vitest"
2+
3+
import type { ApiHandlerCreateMessageMetadata } from "../../api"
4+
5+
describe("abort signal passing", () => {
6+
it("should pass the same AbortController signal instance to metadata.abortSignal", async () => {
7+
// Arrange: create an AbortController
8+
const controller = new AbortController()
9+
10+
// Act: simulate what Task.ts does - construct metadata with abortSignal
11+
const metadata: ApiHandlerCreateMessageMetadata = {
12+
taskId: "test-task-id",
13+
abortSignal: controller.signal,
14+
}
15+
16+
// Assert: signal identity (toBe, not just toBeInstanceOf)
17+
expect(metadata.abortSignal).toBe(controller.signal)
18+
})
19+
20+
it("should create a fresh AbortController for each request", async () => {
21+
// Arrange: simulate two sequential requests
22+
const controller1 = new AbortController()
23+
const metadata1: ApiHandlerCreateMessageMetadata = {
24+
taskId: "task-1",
25+
abortSignal: controller1.signal,
26+
}
27+
28+
const controller2 = new AbortController()
29+
const metadata2: ApiHandlerCreateMessageMetadata = {
30+
taskId: "task-2",
31+
abortSignal: controller2.signal,
32+
}
33+
34+
// Assert: different instances
35+
expect(metadata1.abortSignal).not.toBe(metadata2.abortSignal)
36+
expect(controller1.signal).not.toBe(controller2.signal)
37+
})
38+
})

src/core/task/Task.ts

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4144,10 +4144,15 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
41444144

41454145
const shouldIncludeTools = allTools.length > 0
41464146

4147+
// Create an AbortController FIRST so we can include its signal in metadata
4148+
this.currentRequestAbortController = new AbortController()
4149+
const abortSignal = this.currentRequestAbortController.signal
4150+
41474151
const metadata: ApiHandlerCreateMessageMetadata = {
41484152
mode: mode,
41494153
taskId: this.taskId,
41504154
suppressPreviousResponseId: this.skipPrevResponseIdOnce,
4155+
abortSignal: abortSignal,
41514156
// Include tools whenever they are present.
41524157
...(shouldIncludeTools
41534158
? {
@@ -4160,11 +4165,6 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
41604165
}
41614166
: {}),
41624167
}
4163-
4164-
// Create an AbortController to allow cancelling the request mid-stream
4165-
this.currentRequestAbortController = new AbortController()
4166-
const abortSignal = this.currentRequestAbortController.signal
4167-
// Reset the flag after using it
41684168
this.skipPrevResponseIdOnce = false
41694169

41704170
// The provider accepts reasoning items alongside standard messages; cast to the expected parameter type.

src/utils/__tests__/enhance-prompt.spec.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ describe("enhancePrompt", () => {
4242

4343
expect(result).toBe("Enhanced prompt")
4444
const handler = buildApiHandler(mockApiConfig)
45-
expect((handler as any).completePrompt).toHaveBeenCalledWith(`Test prompt`)
45+
expect((handler as any).completePrompt).toHaveBeenCalledWith(`Test prompt`, undefined)
4646
})
4747

4848
it("enhances prompt using custom enhancement prompt when provided", async () => {
@@ -64,7 +64,7 @@ describe("enhancePrompt", () => {
6464

6565
expect(result).toBe("Enhanced prompt")
6666
const handler = buildApiHandler(mockApiConfig)
67-
expect((handler as any).completePrompt).toHaveBeenCalledWith(`${customEnhancePrompt}\n\nTest prompt`)
67+
expect((handler as any).completePrompt).toHaveBeenCalledWith(`${customEnhancePrompt}\n\nTest prompt`, undefined)
6868
})
6969

7070
it("throws error for empty prompt input", async () => {

src/utils/single-completion-handler.ts

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,16 @@
11
import type { ProviderSettings } from "@roo-code/types"
22

3-
import { buildApiHandler, SingleCompletionHandler } from "../api"
3+
import { buildApiHandler, SingleCompletionHandler, type ApiHandlerCreateMessageMetadata } from "../api"
44

55
/**
66
* Enhances a prompt using the configured API without creating a full Cline instance or task history.
77
* This is a lightweight alternative that only uses the API's completion functionality.
88
*/
9-
export async function singleCompletionHandler(apiConfiguration: ProviderSettings, promptText: string): Promise<string> {
9+
export async function singleCompletionHandler(
10+
apiConfiguration: ProviderSettings,
11+
promptText: string,
12+
metadata?: ApiHandlerCreateMessageMetadata,
13+
): Promise<string> {
1014
if (!promptText) {
1115
throw new Error("No prompt text provided")
1216
}
@@ -21,5 +25,5 @@ export async function singleCompletionHandler(apiConfiguration: ProviderSettings
2125
throw new Error("The selected API provider does not support prompt enhancement")
2226
}
2327

24-
return (handler as SingleCompletionHandler).completePrompt(promptText)
28+
return (handler as SingleCompletionHandler).completePrompt(promptText, metadata)
2529
}

0 commit comments

Comments
 (0)