Skip to content

Commit 3b56639

Browse files
easonliang28easonliang28
committed
feat(task): pass abort signal to condense metadata
1 parent 60a0dcf commit 3b56639

2 files changed

Lines changed: 311 additions & 5 deletions

File tree

‎src/core/task/Task.ts‎

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1555,6 +1555,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
15551555
const metadata: ApiHandlerCreateMessageMetadata = {
15561556
mode,
15571557
taskId: this.taskId,
1558+
...(this.currentRequestAbortController?.signal
1559+
? {
1560+
abortSignal: this.currentRequestAbortController.signal,
1561+
}
1562+
: {}),
15581563
...(allTools.length > 0
15591564
? {
15601565
tools: allTools,
@@ -3763,6 +3768,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
37633768
const metadata: ApiHandlerCreateMessageMetadata = {
37643769
mode,
37653770
taskId: this.taskId,
3771+
...(this.currentRequestAbortController?.signal
3772+
? {
3773+
abortSignal: this.currentRequestAbortController.signal,
3774+
}
3775+
: {}),
37663776
...(allTools.length > 0
37673777
? {
37683778
tools: allTools,
@@ -3979,6 +3989,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
39793989
const contextMgmtMetadata: ApiHandlerCreateMessageMetadata = {
39803990
mode,
39813991
taskId: this.taskId,
3992+
...(this.currentRequestAbortController?.signal
3993+
? {
3994+
abortSignal: this.currentRequestAbortController.signal,
3995+
}
3996+
: {}),
39823997
...(contextMgmtTools.length > 0
39833998
? {
39843999
tools: contextMgmtTools,
@@ -4141,10 +4156,15 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
41414156

41424157
const shouldIncludeTools = allTools.length > 0
41434158

4159+
// Create an AbortController to allow cancelling the request mid-stream
4160+
this.currentRequestAbortController = new AbortController()
4161+
const abortSignal = this.currentRequestAbortController.signal
4162+
41444163
const metadata: ApiHandlerCreateMessageMetadata = {
41454164
mode: mode,
41464165
taskId: this.taskId,
41474166
suppressPreviousResponseId: this.skipPrevResponseIdOnce,
4167+
abortSignal,
41484168
// Include tools whenever they are present.
41494169
...(shouldIncludeTools
41504170
? {
@@ -4157,11 +4177,6 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
41574177
}
41584178
: {}),
41594179
}
4160-
4161-
// Create an AbortController to allow cancelling the request mid-stream
4162-
this.currentRequestAbortController = new AbortController()
4163-
const abortSignal = this.currentRequestAbortController.signal
4164-
metadata.abortSignal = abortSignal
41654180
// Reset the flag after using it
41664181
this.skipPrevResponseIdOnce = false
41674182

‎src/core/task/__tests__/Task.spec.ts‎

Lines changed: 291 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { TelemetryService } from "@roo-code/telemetry"
1111

1212
import { Task } from "../Task"
1313
import { createRateLimitClock } from "../RateLimitClock"
14+
import { summarizeConversation } from "../../condense"
1415
import { ClineProvider } from "../../webview/ClineProvider"
1516
import { ApiStreamChunk } from "../../../api/transform/stream"
1617
import { ContextProxy } from "../../config/ContextProxy"
@@ -1796,6 +1797,24 @@ describe("Cline", () => {
17961797
expect(cancelSpy).toHaveBeenCalled()
17971798
})
17981799
describe("abortSignal", () => {
1800+
it("should pass AbortController signal to condenseContext metadata when a current request exists", async () => {
1801+
const task = new Task({
1802+
provider: mockProvider,
1803+
apiConfiguration: mockApiConfig,
1804+
task: "test task",
1805+
startTask: false,
1806+
})
1807+
1808+
task.currentRequestAbortController = new AbortController()
1809+
vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt")
1810+
1811+
await task.condenseContext()
1812+
1813+
expect(summarizeConversation).toHaveBeenCalled()
1814+
const [options] = vi.mocked(summarizeConversation).mock.calls.at(-1)!
1815+
expect(options.metadata?.abortSignal).toBeInstanceOf(AbortSignal)
1816+
})
1817+
17991818
it("should pass AbortController signal to createMessage metadata", async () => {
18001819
const task = new Task({
18011820
provider: mockProvider,
@@ -1865,6 +1884,109 @@ describe("Cline", () => {
18651884
expect(metadata!.abortSignal).toBeInstanceOf(AbortSignal)
18661885
})
18671886

1887+
it("should invoke abort on currentRequestAbortController during first-chunk wait", async () => {
1888+
const task = new Task({
1889+
provider: mockProvider,
1890+
apiConfiguration: mockApiConfig,
1891+
task: "test task",
1892+
startTask: false,
1893+
})
1894+
1895+
const abortSpy = vi.fn()
1896+
task.currentRequestAbortController = {
1897+
abort: abortSpy,
1898+
signal: new AbortController().signal,
1899+
} as AbortController
1900+
1901+
task.cancelCurrentRequest()
1902+
1903+
expect(abortSpy).toHaveBeenCalledTimes(1)
1904+
expect(task.currentRequestAbortController).toBeUndefined()
1905+
})
1906+
1907+
it("should reject streaming consumption when aborted between chunks", async () => {
1908+
const task = new Task({
1909+
provider: mockProvider,
1910+
apiConfiguration: mockApiConfig,
1911+
task: "test task",
1912+
startTask: false,
1913+
})
1914+
1915+
vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt")
1916+
vi.spyOn(task.api, "getModel").mockReturnValue({
1917+
id: mockApiConfig.apiModelId!,
1918+
info: {
1919+
supportsImages: false,
1920+
supportsPromptCache: true,
1921+
contextWindow: 200000,
1922+
maxTokens: 4096,
1923+
inputPrice: 0.3,
1924+
outputPrice: 1.5,
1925+
} as ModelInfo,
1926+
})
1927+
1928+
const providerState = await mockProvider.getState()
1929+
vi.spyOn(mockProvider, "getState").mockResolvedValue({
1930+
...providerState,
1931+
apiConfiguration: mockApiConfig,
1932+
autoApprovalEnabled: true,
1933+
requestDelaySeconds: 0,
1934+
})
1935+
1936+
const createMessageSpy = vi.fn((_systemPrompt, _messages, metadata) => {
1937+
let callCount = 0
1938+
return {
1939+
[Symbol.asyncIterator]() {
1940+
return this
1941+
},
1942+
next: () => {
1943+
callCount++
1944+
if (callCount === 1) {
1945+
return Promise.resolve({
1946+
done: false,
1947+
value: { type: "text", text: "first chunk" },
1948+
})
1949+
}
1950+
return new Promise<IteratorResult<ApiStreamChunk>>((resolve, reject) => {
1951+
if (metadata?.abortSignal?.aborted) {
1952+
return reject(new Error("Request cancelled by user"))
1953+
}
1954+
metadata?.abortSignal?.addEventListener("abort", () => {
1955+
reject(new Error("Request cancelled by user"))
1956+
})
1957+
})
1958+
},
1959+
async return() {
1960+
return { done: true, value: undefined }
1961+
},
1962+
async throw(e: any) {
1963+
throw e
1964+
},
1965+
[Symbol.asyncDispose]: async () => {},
1966+
} as AsyncGenerator<ApiStreamChunk>
1967+
})
1968+
vi.spyOn(task.api, "createMessage").mockImplementation(createMessageSpy)
1969+
1970+
task.apiConversationHistory = [
1971+
{
1972+
role: "user" as const,
1973+
content: [{ type: "text" as const, text: "test message" }],
1974+
ts: Date.now(),
1975+
},
1976+
] as any
1977+
1978+
const streamIterator = task.attemptApiRequest(0)
1979+
await expect(streamIterator.next()).resolves.toMatchObject({
1980+
done: false,
1981+
value: { type: "text", text: "first chunk" },
1982+
})
1983+
1984+
task.cancelCurrentRequest()
1985+
1986+
await expect(streamIterator.next()).rejects.toThrow("Request cancelled by user")
1987+
expect(createMessageSpy).toHaveBeenCalledTimes(1)
1988+
})
1989+
18681990
it("should use the same AbortController signal as currentRequestAbortController", async () => {
18691991
const task = new Task({
18701992
provider: mockProvider,
@@ -1933,6 +2055,175 @@ describe("Cline", () => {
19332055
// The signal in metadata should be the same as the one from currentRequestAbortController
19342056
expect(metadataSignal).toBe(task.currentRequestAbortController!.signal)
19352057
})
2058+
2059+
it("should omit createMessage abortSignal metadata when no current request exists before condense metadata checks", async () => {
2060+
const task = new Task({
2061+
provider: mockProvider,
2062+
apiConfiguration: mockApiConfig,
2063+
task: "test task",
2064+
startTask: false,
2065+
})
2066+
2067+
vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt")
2068+
vi.spyOn(task.api, "getModel").mockReturnValue({
2069+
id: mockApiConfig.apiModelId!,
2070+
info: {
2071+
supportsImages: false,
2072+
supportsPromptCache: true,
2073+
contextWindow: 200000,
2074+
maxTokens: 4096,
2075+
inputPrice: 0.3,
2076+
outputPrice: 1.5,
2077+
} as ModelInfo,
2078+
})
2079+
2080+
const providerState = await mockProvider.getState()
2081+
vi.spyOn(mockProvider, "getState").mockResolvedValue({
2082+
...providerState,
2083+
apiConfiguration: mockApiConfig,
2084+
autoApprovalEnabled: true,
2085+
requestDelaySeconds: 0,
2086+
})
2087+
2088+
const mockStream = {
2089+
async *[Symbol.asyncIterator]() {
2090+
yield { type: "text", text: "response" }
2091+
},
2092+
async next() {
2093+
return { done: true, value: { type: "text", text: "response" } }
2094+
},
2095+
async return() {
2096+
return { done: true, value: undefined }
2097+
},
2098+
async throw(e: any) {
2099+
throw e
2100+
},
2101+
[Symbol.asyncDispose]: async () => {},
2102+
} as AsyncGenerator<ApiStreamChunk>
2103+
2104+
const createMessageSpy = vi.spyOn(task.api, "createMessage").mockReturnValue(mockStream)
2105+
task.apiConversationHistory = [
2106+
{
2107+
role: "user" as const,
2108+
content: [{ type: "text" as const, text: "test message" }],
2109+
ts: Date.now(),
2110+
},
2111+
] as any
2112+
2113+
expect(task.currentRequestAbortController).toBeUndefined()
2114+
2115+
const iterator = task.attemptApiRequest(0)
2116+
await iterator.next()
2117+
2118+
const [, , metadata] = createMessageSpy.mock.calls[0]!
2119+
expect(metadata).toBeDefined()
2120+
expect("abortSignal" in metadata!).toBe(true)
2121+
expect(metadata!.abortSignal).toBeInstanceOf(AbortSignal)
2122+
})
2123+
})
2124+
2125+
it("should propagate AbortController signal through attemptApiRequest context-window retry path", async () => {
2126+
const task = new Task({
2127+
provider: mockProvider,
2128+
apiConfiguration: mockApiConfig,
2129+
task: "test task",
2130+
startTask: false,
2131+
})
2132+
2133+
vi.spyOn(task as any, "getSystemPrompt").mockResolvedValue("mock system prompt")
2134+
vi.spyOn(task, "getTokenUsage").mockReturnValue({
2135+
totalCost: 0,
2136+
totalTokensIn: 0,
2137+
totalTokensOut: 0,
2138+
contextTokens: 120000,
2139+
})
2140+
vi.spyOn(task.api, "getModel").mockReturnValue({
2141+
id: mockApiConfig.apiModelId!,
2142+
info: {
2143+
supportsImages: false,
2144+
supportsPromptCache: true,
2145+
contextWindow: 1000,
2146+
maxTokens: 4096,
2147+
inputPrice: 0.3,
2148+
outputPrice: 1.5,
2149+
} as ModelInfo,
2150+
})
2151+
const providerState = await mockProvider.getState()
2152+
vi.spyOn(mockProvider, "getState").mockResolvedValue({
2153+
...providerState,
2154+
apiConfiguration: mockApiConfig,
2155+
mode: "code",
2156+
autoCondenseContext: true,
2157+
autoCondenseContextPercent: 80,
2158+
requestDelaySeconds: 0,
2159+
customModes: [],
2160+
experiments: {},
2161+
disabledTools: [],
2162+
customSupportPrompts: {},
2163+
autoApprovalEnabled: true,
2164+
profileThresholds: {},
2165+
currentApiConfigName: "default",
2166+
})
2167+
2168+
task.apiConversationHistory = [
2169+
{
2170+
role: "user" as const,
2171+
content: [{ type: "text" as const, text: "test message" }],
2172+
ts: Date.now(),
2173+
},
2174+
] as any
2175+
2176+
let firstCall = true
2177+
const retryStream = {
2178+
async *[Symbol.asyncIterator]() {
2179+
yield { type: "text", text: "retried response" }
2180+
},
2181+
async next() {
2182+
return { done: false, value: { type: "text", text: "retried response" } }
2183+
},
2184+
async return() {
2185+
return { done: true, value: undefined }
2186+
},
2187+
async throw(e: any) {
2188+
throw e
2189+
},
2190+
[Symbol.asyncDispose]: async () => {},
2191+
} as AsyncGenerator<ApiStreamChunk>
2192+
2193+
const contextWindowErrorStream = {
2194+
[Symbol.asyncIterator]() {
2195+
return this
2196+
},
2197+
async next() {
2198+
throw { status: 400, message: "context length exceeded" }
2199+
},
2200+
async return() {
2201+
return { done: true, value: undefined }
2202+
},
2203+
async throw(e: any) {
2204+
throw e
2205+
},
2206+
[Symbol.asyncDispose]: async () => {},
2207+
} as AsyncGenerator<ApiStreamChunk>
2208+
2209+
vi.spyOn(task.api, "createMessage").mockImplementation(() => {
2210+
if (firstCall) {
2211+
firstCall = false
2212+
return contextWindowErrorStream
2213+
}
2214+
return retryStream
2215+
})
2216+
2217+
const iterator = task.attemptApiRequest(0)
2218+
await expect(iterator.next()).resolves.toMatchObject({
2219+
done: false,
2220+
value: { type: "text", text: "retried response" },
2221+
})
2222+
2223+
expect(summarizeConversation).toHaveBeenCalled()
2224+
const [options] = vi.mocked(summarizeConversation).mock.calls.at(-1)!
2225+
expect(options.metadata?.taskId).toBe(task.taskId)
2226+
expect(options.metadata?.abortSignal).toBeUndefined()
19362227
})
19372228
})
19382229
})

0 commit comments

Comments
 (0)