Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/api/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
29 changes: 24 additions & 5 deletions src/core/task/Task.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1555,6 +1555,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
const metadata: ApiHandlerCreateMessageMetadata = {
mode,
taskId: this.taskId,
...(this.currentRequestAbortController?.signal
? {
abortSignal: this.currentRequestAbortController.signal,
}
: {}),
...(allTools.length > 0
? {
tools: allTools,
Expand Down Expand Up @@ -3763,6 +3768,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
const metadata: ApiHandlerCreateMessageMetadata = {
mode,
taskId: this.taskId,
...(this.currentRequestAbortController?.signal
? {
abortSignal: this.currentRequestAbortController.signal,
}
: {}),
Comment thread
coderabbitai[bot] marked this conversation as resolved.
...(allTools.length > 0
? {
tools: allTools,
Expand Down Expand Up @@ -3979,6 +3989,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
const contextMgmtMetadata: ApiHandlerCreateMessageMetadata = {
mode,
taskId: this.taskId,
...(this.currentRequestAbortController?.signal
? {
abortSignal: this.currentRequestAbortController.signal,
}
: {}),
...(contextMgmtTools.length > 0
? {
tools: contextMgmtTools,
Expand Down Expand Up @@ -4141,10 +4156,15 @@ export class Task extends EventEmitter<TaskEvents> 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
? {
Expand All @@ -4157,10 +4177,6 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
}
: {}),
}

// Create an AbortController to allow cancelling the request mid-stream
this.currentRequestAbortController = new AbortController()
const abortSignal = this.currentRequestAbortController.signal
// Reset the flag after using it
this.skipPrevResponseIdOnce = false

Expand Down Expand Up @@ -4199,9 +4215,12 @@ export class Task extends EventEmitter<TaskEvents> 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(
Expand Down
Loading
Loading