Skip to content

Commit 28c54a7

Browse files
feat: add abort signal core plumbing (#674)
* feat: add abort singal core plumbing * feat(task): pass abort signal to condense metadata --------- Co-authored-by: Eason Liang <easonliang28@gmail.com>
1 parent 518bae4 commit 28c54a7

3 files changed

Lines changed: 541 additions & 5 deletions

File tree

src/api/index.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,12 @@ export interface ApiHandlerCreateMessageMetadata {
9090
* Only applies to providers that support function calling restrictions (e.g., Gemini).
9191
*/
9292
allowedFunctionNames?: string[]
93+
/**
94+
* Abort signal for cancelling the HTTP request mid-stream.
95+
* Passed through to AI SDK's streamText() so the underlying HTTP request is aborted
96+
* when the user clicks stop, preventing wasted API tokens/compute on the provider side.
97+
*/
98+
abortSignal?: AbortSignal
9399
}
94100

95101
export interface ApiHandler {

src/core/task/Task.ts

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1557,6 +1557,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
15571557
const metadata: ApiHandlerCreateMessageMetadata = {
15581558
mode,
15591559
taskId: this.taskId,
1560+
...(this.currentRequestAbortController?.signal
1561+
? {
1562+
abortSignal: this.currentRequestAbortController.signal,
1563+
}
1564+
: {}),
15601565
...(allTools.length > 0
15611566
? {
15621567
tools: allTools,
@@ -3765,6 +3770,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
37653770
const metadata: ApiHandlerCreateMessageMetadata = {
37663771
mode,
37673772
taskId: this.taskId,
3773+
...(this.currentRequestAbortController?.signal
3774+
? {
3775+
abortSignal: this.currentRequestAbortController.signal,
3776+
}
3777+
: {}),
37683778
...(allTools.length > 0
37693779
? {
37703780
tools: allTools,
@@ -3981,6 +3991,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
39813991
const contextMgmtMetadata: ApiHandlerCreateMessageMetadata = {
39823992
mode,
39833993
taskId: this.taskId,
3994+
...(this.currentRequestAbortController?.signal
3995+
? {
3996+
abortSignal: this.currentRequestAbortController.signal,
3997+
}
3998+
: {}),
39843999
...(contextMgmtTools.length > 0
39854000
? {
39864001
tools: contextMgmtTools,
@@ -4143,10 +4158,15 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
41434158

41444159
const shouldIncludeTools = allTools.length > 0
41454160

4161+
// Create an AbortController to allow cancelling the request mid-stream
4162+
this.currentRequestAbortController = new AbortController()
4163+
const abortSignal = this.currentRequestAbortController.signal
4164+
41464165
const metadata: ApiHandlerCreateMessageMetadata = {
41474166
mode: mode,
41484167
taskId: this.taskId,
41494168
suppressPreviousResponseId: this.skipPrevResponseIdOnce,
4169+
abortSignal,
41504170
// Include tools whenever they are present.
41514171
...(shouldIncludeTools
41524172
? {
@@ -4159,10 +4179,6 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
41594179
}
41604180
: {}),
41614181
}
4162-
4163-
// Create an AbortController to allow cancelling the request mid-stream
4164-
this.currentRequestAbortController = new AbortController()
4165-
const abortSignal = this.currentRequestAbortController.signal
41664182
// Reset the flag after using it
41674183
this.skipPrevResponseIdOnce = false
41684184

@@ -4201,9 +4217,12 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
42014217
this.isWaitingForFirstChunk = false
42024218
} catch (error) {
42034219
this.isWaitingForFirstChunk = false
4204-
this.currentRequestAbortController = undefined
42054220
const isContextWindowExceededError = checkContextWindowExceededError(error)
42064221

4222+
if (!isContextWindowExceededError) {
4223+
this.currentRequestAbortController = undefined
4224+
}
4225+
42074226
// If it's a context window error and we haven't exceeded max retries for this error type
42084227
if (isContextWindowExceededError && retryAttempt < MAX_CONTEXT_WINDOW_RETRIES) {
42094228
console.warn(

0 commit comments

Comments
 (0)