From aa472c4b3c0b96f68ec216dc92191663781f55ad Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Sun, 21 Jun 2026 03:41:10 +0800 Subject: [PATCH 1/2] feat: add abort singal core plumbing --- src/api/index.ts | 6 ++ src/core/task/Task.ts | 1 + src/core/task/__tests__/Task.spec.ts | 139 +++++++++++++++++++++++++++ 3 files changed, 146 insertions(+) diff --git a/src/api/index.ts b/src/api/index.ts index e52b41200b..0c901f8e23 100644 --- a/src/api/index.ts +++ b/src/api/index.ts @@ -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 { diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 2f1b370b48..53b5768e92 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -4161,6 +4161,7 @@ export class Task extends EventEmitter implements TaskLike { // Create an AbortController to allow cancelling the request mid-stream this.currentRequestAbortController = new AbortController() const abortSignal = this.currentRequestAbortController.signal + metadata.abortSignal = abortSignal // Reset the flag after using it this.skipPrevResponseIdOnce = false diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index dd63135802..e84b023de2 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -1795,6 +1795,145 @@ describe("Cline", () => { // Verify cancelCurrentRequest was called expect(cancelSpy).toHaveBeenCalled() }) + describe("abortSignal", () => { + it("should pass AbortController signal to createMessage metadata", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + + // Mock required methods for attemptApiRequest to work without hanging + vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt") + + vi.spyOn(task.api, "getModel").mockReturnValue({ + id: mockApiConfig.apiModelId!, + info: { + supportsImages: false, + supportsPromptCache: true, + contextWindow: 200000, + maxTokens: 4096, + inputPrice: 0.3, + outputPrice: 1.5, + } as ModelInfo, + }) + + const providerState = await mockProvider.getState() + vi.spyOn(mockProvider, "getState").mockResolvedValue({ + ...providerState, + apiConfiguration: mockApiConfig, + autoApprovalEnabled: true, + requestDelaySeconds: 0, + }) + + // Mock the API stream response + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { type: "text", text: "response" } + }, + async next() { + return { done: true, value: { type: "text", text: "response" } } + }, + async return() { + return { done: true, value: undefined } + }, + async throw(e: any) { + throw e + }, + [Symbol.asyncDispose]: async () => {}, + } as AsyncGenerator + + const createMessageSpy = vi.spyOn(task.api, "createMessage").mockReturnValue(mockStream) + + task.apiConversationHistory = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: "test message" }], + ts: Date.now(), + }, + ] as any + + const iterator = task.attemptApiRequest(0) + await iterator.next() + + // Verify createMessage was called with metadata containing abortSignal + expect(createMessageSpy).toHaveBeenCalled() + const [, , metadata] = createMessageSpy.mock.calls[0]! + + expect(metadata).toBeDefined() + expect(metadata!.abortSignal).toBeInstanceOf(AbortSignal) + }) + + it("should use the same AbortController signal as currentRequestAbortController", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + + // Mock required methods for attemptApiRequest to work without hanging + vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt") + + vi.spyOn(task.api, "getModel").mockReturnValue({ + id: mockApiConfig.apiModelId!, + info: { + supportsImages: false, + supportsPromptCache: true, + contextWindow: 200000, + maxTokens: 4096, + inputPrice: 0.3, + outputPrice: 1.5, + } as ModelInfo, + }) + + const providerState = await mockProvider.getState() + vi.spyOn(mockProvider, "getState").mockResolvedValue({ + ...providerState, + apiConfiguration: mockApiConfig, + autoApprovalEnabled: true, + requestDelaySeconds: 0, + }) + + // Mock the API stream response + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { type: "text", text: "response" } + }, + async next() { + return { done: true, value: { type: "text", text: "response" } } + }, + async return() { + return { done: true, value: undefined } + }, + async throw(e: any) { + throw e + }, + [Symbol.asyncDispose]: async () => {}, + } as AsyncGenerator + + const createMessageSpy = vi.spyOn(task.api, "createMessage").mockReturnValue(mockStream) + + task.apiConversationHistory = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: "test message" }], + ts: Date.now(), + }, + ] as any + + const iterator = task.attemptApiRequest(0) + await iterator.next() + + // Get the signal from metadata + const [, , metadata] = createMessageSpy.mock.calls[0]! + const metadataSignal = metadata!.abortSignal + + // The signal in metadata should be the same as the one from currentRequestAbortController + expect(metadataSignal).toBe(task.currentRequestAbortController!.signal) + }) + }) }) }) From b26ec5d0385d6dc0ade07d9d872b61aefc3a1ca2 Mon Sep 17 00:00:00 2001 From: Eason Liang Date: Tue, 23 Jun 2026 00:20:07 +0800 Subject: [PATCH 2/2] feat(task): pass abort signal to condense metadata --- src/core/task/Task.ts | 30 ++- src/core/task/__tests__/Task.spec.ts | 372 +++++++++++++++++++++++++++ 2 files changed, 396 insertions(+), 6 deletions(-) diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 53b5768e92..10b22ffaa3 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -1555,6 +1555,11 @@ export class Task extends EventEmitter implements TaskLike { const metadata: ApiHandlerCreateMessageMetadata = { mode, taskId: this.taskId, + ...(this.currentRequestAbortController?.signal + ? { + abortSignal: this.currentRequestAbortController.signal, + } + : {}), ...(allTools.length > 0 ? { tools: allTools, @@ -3763,6 +3768,11 @@ export class Task extends EventEmitter implements TaskLike { const metadata: ApiHandlerCreateMessageMetadata = { mode, taskId: this.taskId, + ...(this.currentRequestAbortController?.signal + ? { + abortSignal: this.currentRequestAbortController.signal, + } + : {}), ...(allTools.length > 0 ? { tools: allTools, @@ -3979,6 +3989,11 @@ export class Task extends EventEmitter implements TaskLike { const contextMgmtMetadata: ApiHandlerCreateMessageMetadata = { mode, taskId: this.taskId, + ...(this.currentRequestAbortController?.signal + ? { + abortSignal: this.currentRequestAbortController.signal, + } + : {}), ...(contextMgmtTools.length > 0 ? { tools: contextMgmtTools, @@ -4141,10 +4156,15 @@ export class Task extends EventEmitter implements TaskLike { const shouldIncludeTools = allTools.length > 0 + // Create an AbortController to allow cancelling the request mid-stream + this.currentRequestAbortController = new AbortController() + const abortSignal = this.currentRequestAbortController.signal + const metadata: ApiHandlerCreateMessageMetadata = { mode: mode, taskId: this.taskId, suppressPreviousResponseId: this.skipPrevResponseIdOnce, + abortSignal, // Include tools whenever they are present. ...(shouldIncludeTools ? { @@ -4157,11 +4177,6 @@ export class Task extends EventEmitter implements TaskLike { } : {}), } - - // Create an AbortController to allow cancelling the request mid-stream - this.currentRequestAbortController = new AbortController() - const abortSignal = this.currentRequestAbortController.signal - metadata.abortSignal = abortSignal // Reset the flag after using it this.skipPrevResponseIdOnce = false @@ -4200,9 +4215,12 @@ export class Task extends EventEmitter implements TaskLike { this.isWaitingForFirstChunk = false } catch (error) { this.isWaitingForFirstChunk = false - this.currentRequestAbortController = undefined const isContextWindowExceededError = checkContextWindowExceededError(error) + if (!isContextWindowExceededError) { + this.currentRequestAbortController = undefined + } + // If it's a context window error and we haven't exceeded max retries for this error type if (isContextWindowExceededError && retryAttempt < MAX_CONTEXT_WINDOW_RETRIES) { console.warn( diff --git a/src/core/task/__tests__/Task.spec.ts b/src/core/task/__tests__/Task.spec.ts index e84b023de2..5c10771686 100644 --- a/src/core/task/__tests__/Task.spec.ts +++ b/src/core/task/__tests__/Task.spec.ts @@ -11,6 +11,7 @@ import { TelemetryService } from "@roo-code/telemetry" import { Task } from "../Task" import { createRateLimitClock } from "../RateLimitClock" +import { summarizeConversation } from "../../condense" import { ClineProvider } from "../../webview/ClineProvider" import { ApiStreamChunk } from "../../../api/transform/stream" import { ContextProxy } from "../../config/ContextProxy" @@ -1796,6 +1797,42 @@ describe("Cline", () => { expect(cancelSpy).toHaveBeenCalled() }) describe("abortSignal", () => { + it("should pass AbortController signal to condenseContext metadata when a current request exists", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + + task.currentRequestAbortController = new AbortController() + vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt") + + await task.condenseContext() + + expect(summarizeConversation).toHaveBeenCalled() + const [options] = vi.mocked(summarizeConversation).mock.calls.at(-1)! + expect(options.metadata?.abortSignal).toBeInstanceOf(AbortSignal) + }) + + it("should omit abortSignal from condenseContext metadata when no current request exists", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + + vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt") + + await task.condenseContext() + + expect(summarizeConversation).toHaveBeenCalled() + const [options] = vi.mocked(summarizeConversation).mock.calls.at(-1)! + expect(options.metadata).toBeDefined() + expect("abortSignal" in (options.metadata ?? {})).toBe(false) + }) + it("should pass AbortController signal to createMessage metadata", async () => { const task = new Task({ provider: mockProvider, @@ -1865,6 +1902,109 @@ describe("Cline", () => { expect(metadata!.abortSignal).toBeInstanceOf(AbortSignal) }) + it("should invoke abort on currentRequestAbortController during first-chunk wait", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + + const abortSpy = vi.fn() + task.currentRequestAbortController = { + abort: abortSpy, + signal: new AbortController().signal, + } as AbortController + + task.cancelCurrentRequest() + + expect(abortSpy).toHaveBeenCalledTimes(1) + expect(task.currentRequestAbortController).toBeUndefined() + }) + + it("should reject streaming consumption when aborted between chunks", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + + vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt") + vi.spyOn(task.api, "getModel").mockReturnValue({ + id: mockApiConfig.apiModelId!, + info: { + supportsImages: false, + supportsPromptCache: true, + contextWindow: 200000, + maxTokens: 4096, + inputPrice: 0.3, + outputPrice: 1.5, + } as ModelInfo, + }) + + const providerState = await mockProvider.getState() + vi.spyOn(mockProvider, "getState").mockResolvedValue({ + ...providerState, + apiConfiguration: mockApiConfig, + autoApprovalEnabled: true, + requestDelaySeconds: 0, + }) + + const createMessageSpy = vi.fn((_systemPrompt, _messages, metadata) => { + let callCount = 0 + return { + [Symbol.asyncIterator]() { + return this + }, + next: () => { + callCount++ + if (callCount === 1) { + return Promise.resolve({ + done: false, + value: { type: "text", text: "first chunk" }, + }) + } + return new Promise>((resolve, reject) => { + if (metadata?.abortSignal?.aborted) { + return reject(new Error("Request cancelled by user")) + } + metadata?.abortSignal?.addEventListener("abort", () => { + reject(new Error("Request cancelled by user")) + }) + }) + }, + async return() { + return { done: true, value: undefined } + }, + async throw(e: any) { + throw e + }, + [Symbol.asyncDispose]: async () => {}, + } as AsyncGenerator + }) + vi.spyOn(task.api, "createMessage").mockImplementation(createMessageSpy) + + task.apiConversationHistory = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: "test message" }], + ts: Date.now(), + }, + ] as any + + const streamIterator = task.attemptApiRequest(0) + await expect(streamIterator.next()).resolves.toMatchObject({ + done: false, + value: { type: "text", text: "first chunk" }, + }) + + task.cancelCurrentRequest() + + await expect(streamIterator.next()).rejects.toThrow("Request cancelled by user") + expect(createMessageSpy).toHaveBeenCalledTimes(1) + }) + it("should use the same AbortController signal as currentRequestAbortController", async () => { const task = new Task({ provider: mockProvider, @@ -1933,6 +2073,238 @@ describe("Cline", () => { // The signal in metadata should be the same as the one from currentRequestAbortController expect(metadataSignal).toBe(task.currentRequestAbortController!.signal) }) + + it("should omit createMessage abortSignal metadata when no current request exists before condense metadata checks", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + + vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt") + vi.spyOn(task.api, "getModel").mockReturnValue({ + id: mockApiConfig.apiModelId!, + info: { + supportsImages: false, + supportsPromptCache: true, + contextWindow: 200000, + maxTokens: 4096, + inputPrice: 0.3, + outputPrice: 1.5, + } as ModelInfo, + }) + + const providerState = await mockProvider.getState() + vi.spyOn(mockProvider, "getState").mockResolvedValue({ + ...providerState, + apiConfiguration: mockApiConfig, + autoApprovalEnabled: true, + requestDelaySeconds: 0, + }) + + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { type: "text", text: "response" } + }, + async next() { + return { done: true, value: { type: "text", text: "response" } } + }, + async return() { + return { done: true, value: undefined } + }, + async throw(e: any) { + throw e + }, + [Symbol.asyncDispose]: async () => {}, + } as AsyncGenerator + + const createMessageSpy = vi.spyOn(task.api, "createMessage").mockReturnValue(mockStream) + task.apiConversationHistory = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: "test message" }], + ts: Date.now(), + }, + ] as any + + expect(task.currentRequestAbortController).toBeUndefined() + + const iterator = task.attemptApiRequest(0) + await iterator.next() + + const [, , metadata] = createMessageSpy.mock.calls[0]! + expect(metadata).toBeDefined() + expect("abortSignal" in metadata!).toBe(true) + expect(metadata!.abortSignal).toBeInstanceOf(AbortSignal) + }) + + it("should keep createMessage abortSignal metadata unaborted before cancellation", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + + vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt") + vi.spyOn(task.api, "getModel").mockReturnValue({ + id: mockApiConfig.apiModelId!, + info: { + supportsImages: false, + supportsPromptCache: true, + contextWindow: 200000, + maxTokens: 4096, + inputPrice: 0.3, + outputPrice: 1.5, + } as ModelInfo, + }) + + const providerState = await mockProvider.getState() + vi.spyOn(mockProvider, "getState").mockResolvedValue({ + ...providerState, + apiConfiguration: mockApiConfig, + autoApprovalEnabled: true, + requestDelaySeconds: 0, + }) + + const mockStream = { + async *[Symbol.asyncIterator]() { + yield { type: "text", text: "response" } + }, + async next() { + return { done: false, value: { type: "text", text: "response" } } + }, + async return() { + return { done: true, value: undefined } + }, + async throw(e: any) { + throw e + }, + [Symbol.asyncDispose]: async () => {}, + } as AsyncGenerator + + const createMessageSpy = vi.spyOn(task.api, "createMessage").mockReturnValue(mockStream) + task.apiConversationHistory = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: "test message" }], + ts: Date.now(), + }, + ] as any + + const iterator = task.attemptApiRequest(0) + await iterator.next() + + const [, , metadata] = createMessageSpy.mock.calls[0]! + expect(metadata?.abortSignal).toBeInstanceOf(AbortSignal) + expect(metadata?.abortSignal?.aborted).toBe(false) + }) + }) + + it("should propagate AbortController signal through attemptApiRequest context-window retry path", async () => { + const task = new Task({ + provider: mockProvider, + apiConfiguration: mockApiConfig, + task: "test task", + startTask: false, + }) + + vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt") + vi.spyOn(task, "getTokenUsage").mockReturnValue({ + totalCost: 0, + totalTokensIn: 0, + totalTokensOut: 0, + contextTokens: 120000, + }) + vi.spyOn(task.api, "getModel").mockReturnValue({ + id: mockApiConfig.apiModelId!, + info: { + supportsImages: false, + supportsPromptCache: true, + contextWindow: 1000, + maxTokens: 4096, + inputPrice: 0.3, + outputPrice: 1.5, + } as ModelInfo, + }) + const providerState = await mockProvider.getState() + vi.spyOn(mockProvider, "getState").mockResolvedValue({ + ...providerState, + apiConfiguration: mockApiConfig, + mode: "code", + autoCondenseContext: true, + autoCondenseContextPercent: 80, + requestDelaySeconds: 0, + customModes: [], + experiments: {}, + disabledTools: [], + customSupportPrompts: {}, + autoApprovalEnabled: true, + profileThresholds: {}, + currentApiConfigName: "default", + }) + + task.apiConversationHistory = [ + { + role: "user" as const, + content: [{ type: "text" as const, text: "test message" }], + ts: Date.now(), + }, + ] as any + + let firstCall = true + const retryStream = { + async *[Symbol.asyncIterator]() { + yield { type: "text", text: "retried response" } + }, + async next() { + return { done: false, value: { type: "text", text: "retried response" } } + }, + async return() { + return { done: true, value: undefined } + }, + async throw(e: any) { + throw e + }, + [Symbol.asyncDispose]: async () => {}, + } as AsyncGenerator + + const contextWindowErrorStream = { + [Symbol.asyncIterator]() { + return this + }, + async next() { + throw { status: 400, message: "context length exceeded" } + }, + async return() { + return { done: true, value: undefined } + }, + async throw(e: any) { + throw e + }, + [Symbol.asyncDispose]: async () => {}, + } as AsyncGenerator + + vi.spyOn(task.api, "createMessage").mockImplementation(() => { + if (firstCall) { + firstCall = false + return contextWindowErrorStream + } + return retryStream + }) + + const iterator = task.attemptApiRequest(0) + await expect(iterator.next()).resolves.toMatchObject({ + done: false, + value: { type: "text", text: "retried response" }, + }) + + expect(summarizeConversation).toHaveBeenCalled() + const [options] = vi.mocked(summarizeConversation).mock.calls.at(-1)! + expect(options.metadata?.taskId).toBe(task.taskId) + expect(options.metadata?.abortSignal).toBeInstanceOf(AbortSignal) + expect(options.metadata?.abortSignal?.aborted).toBe(false) }) }) })