From f3a64af07b6f682d09f20dec5eaaab7a70c4e60a Mon Sep 17 00:00:00 2001 From: Sai Mallikarjun <106309863+Mallikarjun-0@users.noreply.github.com> Date: Wed, 4 Nov 2026 14:53:35 +0530 Subject: [PATCH] fix(ai): preserve aborted stop reason during lazy setup (#8635) Fixes #8409 ## Summary - Pass the request abort signal through lazy stream setup wrappers. - Report setup failures as aborted when the request signal is already aborted. - Added a regression test for aborting during a tool execution before next auth setup. --- packages/ai/src/api/lazy.ts | 39 +++++--- packages/ai/src/models.ts | 45 ++++++--- .../coding-agent/src/core/model-runtime.ts | 93 ++++++++++++------- .../src/core/provider-composer.ts | 33 ++++--- .../8409-aborted-tool-auth-setup.test.ts | 53 +++++++++++ 5 files changed, 187 insertions(+), 76 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/8409-aborted-tool-auth-setup.test.ts diff --git a/packages/ai/src/api/lazy.ts b/packages/ai/src/api/lazy.ts index d8db7440e..fad21164e 100644 --- a/packages/ai/src/api/lazy.ts +++ b/packages/ai/src/api/lazy.ts @@ -1,7 +1,12 @@ import type { Api, AssistantMessage, AssistantMessageEvent, Model, ProviderStreams } from "../types.ts"; import { AssistantMessageEventStream } from "../utils/event-stream.ts"; -function createSetupErrorMessage(model: Model, error: unknown, timestamp: number): AssistantMessage { +function createSetupErrorMessage( + model: Model, + error: unknown, + timestamp: number, + stopReason: "error" | "aborted", +): AssistantMessage { return { role: "assistant", content: [], @@ -16,7 +21,7 @@ function createSetupErrorMessage(model: Model, error: unknown, timestamp: n totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, - stopReason: "error", + stopReason, errorMessage: error instanceof Error ? error.message : String(error), timestamp, }; @@ -61,6 +66,7 @@ export function lazyStream( model: Model, setup: () => Promise>, observer?: LazyStreamObserver, + signal?: AbortSignal, ): AssistantMessageEventStream { const startedAt = Date.now(); const outer = new AssistantMessageEventStream(); @@ -74,8 +80,9 @@ export function lazyStream( setup() .then((inner) => forwardStream(outer, inner, observer, observeResult)) .catch((error) => { - const message = createSetupErrorMessage(model, error, startedAt); - const event: AssistantMessageEvent = { type: "error", reason: "error", error: message }; + const reason = signal?.aborted ? "aborted" : "error"; + const message = createSetupErrorMessage(model, error, startedAt, reason); + const event: AssistantMessageEvent = { type: "error", reason, error: message }; outer.push(event); observer?.onEvent?.(event); outer.end(message); @@ -98,18 +105,28 @@ export interface LazyApiCapabilities { export function lazyApi(load: () => Promise, capabilities?: LazyApiCapabilities): ProviderStreams { const api: ProviderStreams = { stream: (model, context, options) => - lazyStream(model, async () => (await load()).stream(model, context, options)), + lazyStream(model, async () => (await load()).stream(model, context, options), undefined, options?.signal), streamSimple: (model, context, options) => - lazyStream(model, async () => (await load()).streamSimple(model, context, options)), + lazyStream( + model, + async () => (await load()).streamSimple(model, context, options), + undefined, + options?.signal, + ), }; if (capabilities?.fetchDeferred) { api.fetchDeferred = (model, handle, options) => - lazyStream(model, async () => { - const implementation = await load(); - if (!implementation.fetchDeferred) throw new Error("API does not support deferred responses"); - return implementation.fetchDeferred(model, handle, options); - }); + lazyStream( + model, + async () => { + const implementation = await load(); + if (!implementation.fetchDeferred) throw new Error("API does not support deferred responses"); + return implementation.fetchDeferred(model, handle, options); + }, + undefined, + options?.signal, + ); } if (capabilities?.cancelDeferred) { api.cancelDeferred = async (model, handle, options) => { diff --git a/packages/ai/src/models.ts b/packages/ai/src/models.ts index 3d26dcb5f..5842affce 100644 --- a/packages/ai/src/models.ts +++ b/packages/ai/src/models.ts @@ -915,6 +915,7 @@ class ModelsImpl implements MutableModels { return provider.stream(requestModel, transcript, requestOptions as ApiStreamOptions); }, observeModelRequest(model, options, false), + options?.signal, ); } @@ -936,6 +937,7 @@ class ModelsImpl implements MutableModels { return provider.streamSimple(requestModel, transcript, requestOptions as SimpleStreamOptions); }, observeModelRequest(model, options, false), + options?.signal, ); } @@ -963,6 +965,7 @@ class ModelsImpl implements MutableModels { return provider.fetchDeferred(requestModel, handle, requestOptions as DeferredFetchOptions); }, observeModelRequest(model, options, true), + options?.signal, ); } @@ -1108,12 +1111,18 @@ export function createProvider(input: CreateProviderOpti const dispatch = ( model: Model, run: (streams: ProviderStreams) => AssistantMessageEventStream, + signal?: AbortSignal, ): AssistantMessageEventStream => { const streams = apiFor(model); if (!streams) { - return lazyStream(model, async () => { - throw new ModelsError("stream", `Provider ${input.id} has no API implementation for "${model.api}"`); - }); + return lazyStream( + model, + async () => { + throw new ModelsError("stream", `Provider ${input.id} has no API implementation for "${model.api}"`); + }, + undefined, + signal, + ); } return run(streams); }; @@ -1156,23 +1165,29 @@ export function createProvider(input: CreateProviderOpti : undefined, filterModels: input.filterModels, filterAllModels: input.filterAllModels, - stream: (model, context, options) => dispatch(model, (streams) => streams.stream(model, context, options)), + stream: (model, context, options) => + dispatch(model, (streams) => streams.stream(model, context, options), options?.signal), streamSimple: (model, context, options) => - dispatch(model, (streams) => streams.streamSimple(model, context, options)), + dispatch(model, (streams) => streams.streamSimple(model, context, options), options?.signal), }; if (streams.some((entry) => entry.fetchDeferred !== undefined)) { provider.fetchDeferred = (model, handle, options) => - lazyStream(model, async () => { - const implementation = apiFor(model); - if (!implementation?.fetchDeferred) { - throw new ModelsError( - "provider", - `Provider ${input.id} does not support deferred responses for "${model.api}"`, - ); - } - return implementation.fetchDeferred(model, handle, options); - }); + lazyStream( + model, + async () => { + const implementation = apiFor(model); + if (!implementation?.fetchDeferred) { + throw new ModelsError( + "provider", + `Provider ${input.id} does not support deferred responses for "${model.api}"`, + ); + } + return implementation.fetchDeferred(model, handle, options); + }, + undefined, + options?.signal, + ); } if (streams.some((entry) => entry.cancelDeferred !== undefined)) { provider.cancelDeferred = async (model, handle, options) => { diff --git a/packages/coding-agent/src/core/model-runtime.ts b/packages/coding-agent/src/core/model-runtime.ts index 8e03f4d1f..6a5520d68 100644 --- a/packages/coding-agent/src/core/model-runtime.ts +++ b/packages/coding-agent/src/core/model-runtime.ts @@ -710,14 +710,19 @@ export class ModelRuntime implements Models { options?: ModelsApiStreamOptions, ): AssistantMessageEventStream { const transcript = normalizeContext(context); - return lazyStream(model, async () => { - assertChatModel(model); - const prepared = await this.prepareRequest( - model, - options as (StreamOptions & ModelsRequestTransforms) | undefined, - ); - return prepared.provider.stream(prepared.model, transcript, prepared.options as ApiStreamOptions); - }); + return lazyStream( + model, + async () => { + assertChatModel(model); + const prepared = await this.prepareRequest( + model, + options as (StreamOptions & ModelsRequestTransforms) | undefined, + ); + return prepared.provider.stream(prepared.model, transcript, prepared.options as ApiStreamOptions); + }, + undefined, + options?.signal, + ); } complete( @@ -733,27 +738,38 @@ export class ModelRuntime implements Models { if (isVirtualModel(model)) { // Requests outside the agent loop are routed here. Callers sized them before routing, so // cap the output budget to the routed model. - return lazyStream(model, async () => { - const route = await this.resolveModel(model, transcript.messages, { - reason: "direct", - thinkingLevel: options?.reasoning ?? "off", - signal: options?.signal, - }); - const { maxTokens: limit } = route.model; - const maxTokens = options?.maxTokens && limit > 0 ? Math.min(options.maxTokens, limit) : options?.maxTokens; - const reasoning = route.thinkingLevel === "off" ? undefined : route.thinkingLevel; - // Caller credentials were resolved for the virtual model's provider. Another provider - // resolves its own, so they are not sent to the wrong vendor. - const { apiKey, headers, env, ...rest } = options ?? {}; - const auth = route.model.provider === model.provider ? { apiKey, headers, env } : {}; - return this.streamSimple(route.model, context, { ...rest, ...auth, maxTokens, reasoning }); - }); + return lazyStream( + model, + async () => { + const route = await this.resolveModel(model, transcript.messages, { + reason: "direct", + thinkingLevel: options?.reasoning ?? "off", + signal: options?.signal, + }); + const { maxTokens: limit } = route.model; + const maxTokens = + options?.maxTokens && limit > 0 ? Math.min(options.maxTokens, limit) : options?.maxTokens; + const reasoning = route.thinkingLevel === "off" ? undefined : route.thinkingLevel; + // Caller credentials were resolved for the virtual model's provider. Another provider + // resolves its own, so they are not sent to the wrong vendor. + const { apiKey, headers, env, ...rest } = options ?? {}; + const auth = route.model.provider === model.provider ? { apiKey, headers, env } : {}; + return this.streamSimple(route.model, context, { ...rest, ...auth, maxTokens, reasoning }); + }, + undefined, + options?.signal, + ); } - return lazyStream(model, async () => { - assertChatModel(model); - const prepared = await this.prepareRequest(model, options); - return prepared.provider.streamSimple(prepared.model, transcript, prepared.options as SimpleStreamOptions); - }); + return lazyStream( + model, + async () => { + assertChatModel(model); + const prepared = await this.prepareRequest(model, options); + return prepared.provider.streamSimple(prepared.model, transcript, prepared.options as SimpleStreamOptions); + }, + undefined, + options?.signal, + ); } completeSimple(model: Model, context: Context, options?: ModelsSimpleStreamOptions): Promise { @@ -765,14 +781,19 @@ export class ModelRuntime implements Models { handle: DeferredHandle, options?: ModelsDeferredFetchOptions, ): AssistantMessageEventStream { - return lazyStream(model, async () => { - assertChatModel(model); - const prepared = await this.prepareRequest(model, options); - if (!prepared.provider.fetchDeferred) { - throw new ModelsError("provider", `Provider ${model.provider} does not support deferred responses`); - } - return prepared.provider.fetchDeferred(prepared.model, handle, prepared.options as DeferredFetchOptions); - }); + return lazyStream( + model, + async () => { + assertChatModel(model); + const prepared = await this.prepareRequest(model, options); + if (!prepared.provider.fetchDeferred) { + throw new ModelsError("provider", `Provider ${model.provider} does not support deferred responses`); + } + return prepared.provider.fetchDeferred(prepared.model, handle, prepared.options as DeferredFetchOptions); + }, + undefined, + options?.signal, + ); } async fetchDeferred( diff --git a/packages/coding-agent/src/core/provider-composer.ts b/packages/coding-agent/src/core/provider-composer.ts index f7f5c4c59..50940ec6a 100644 --- a/packages/coding-agent/src/core/provider-composer.ts +++ b/packages/coding-agent/src/core/provider-composer.ts @@ -591,21 +591,26 @@ export function composeModelProvider( options: StreamOptions | undefined, simple: boolean, ): AssistantMessageEventStream => - lazyStream(model, async () => { - if (extension?.streamSimple && model.api === extension.api) { - return extension.streamSimple(model, context, options as SimpleStreamOptions); - } - if (base && supportsBaseApi(model)) { + lazyStream( + model, + async () => { + if (extension?.streamSimple && model.api === extension.api) { + return extension.streamSimple(model, context, options as SimpleStreamOptions); + } + if (base && supportsBaseApi(model)) { + return simple + ? base.streamSimple(model, context, options as SimpleStreamOptions) + : base.stream(model, context, options); + } + const api = getApiProvider(model.api); + if (!api) throw new Error(`No API provider registered for api: ${model.api}`); return simple - ? base.streamSimple(model, context, options as SimpleStreamOptions) - : base.stream(model, context, options); - } - const api = getApiProvider(model.api); - if (!api) throw new Error(`No API provider registered for api: ${model.api}`); - return simple - ? api.streamSimple(model, context, options as SimpleStreamOptions) - : api.stream(model, context, options); - }); + ? api.streamSimple(model, context, options as SimpleStreamOptions) + : api.stream(model, context, options); + }, + undefined, + options?.signal, + ); const provider: Provider = { id: providerId, diff --git a/packages/coding-agent/test/suite/regressions/8409-aborted-tool-auth-setup.test.ts b/packages/coding-agent/test/suite/regressions/8409-aborted-tool-auth-setup.test.ts new file mode 100644 index 000000000..b930bb7e7 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/8409-aborted-tool-auth-setup.test.ts @@ -0,0 +1,53 @@ +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; +import { Type } from "typebox"; +import { afterEach, describe, expect, it } from "vitest"; +import { createHarness, type Harness } from "../harness.ts"; + +describe("issue #8409 aborted tool auth setup", () => { + let harness: Harness | undefined; + + afterEach(() => { + harness?.cleanup(); + harness = undefined; + }); + + it("ends the turn as aborted when auth setup sees an already-aborted signal", async () => { + let markToolStarted = () => {}; + const toolStarted = new Promise((resolve) => { + markToolStarted = resolve; + }); + const blockingTool: AgentTool = { + name: "block", + label: "Block", + description: "Blocks until aborted", + parameters: Type.Object({}), + execute: async (_toolCallId, _params, signal) => { + markToolStarted(); + await new Promise((resolve) => { + if (signal?.aborted) resolve(); + else signal?.addEventListener("abort", () => resolve(), { once: true }); + }); + return { content: [{ type: "text", text: "tool aborted" }], details: {} }; + }, + }; + harness = await createHarness({ tools: [blockingTool] }); + const modelRuntime = harness.session.modelRuntime; + harness.session.agent.streamFunction = (model, context, options) => + modelRuntime.streamSimple(model, context, options); + harness.setResponses([fauxAssistantMessage(fauxToolCall("block", {}), { stopReason: "toolUse" })]); + + const prompt = harness.session.prompt("run the blocking tool"); + await toolStarted; + await harness.session.abort(); + await prompt; + + const finalAssistant = harness.session.messages + .slice() + .reverse() + .find((message) => message.role === "assistant"); + expect(finalAssistant?.role).toBe("assistant"); + if (finalAssistant?.role !== "assistant") throw new Error("Expected final assistant message"); + expect(finalAssistant.stopReason).toBe("aborted"); + }); +});