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

@@ -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<Api>, error: unknown, timestamp: number): AssistantMessage {
+function createSetupErrorMessage(
+	model: Model<Api>,
+	error: unknown,
+	timestamp: number,
+	stopReason: "error" | "aborted",
+): AssistantMessage {
 	return {
 		role: "assistant",
 		content: [],
@@ -16,7 +21,7 @@ function createSetupErrorMessage(model: Model<Api>, 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<Api>,
 	setup: () => Promise<AsyncIterable<AssistantMessageEvent>>,
 	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 reason = signal?.aborted ? "aborted" : "error";
-			const event: AssistantMessageEvent = { type: "error", reason: "error", error: message };
+			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<ProviderStreams>, 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 () => {
+			lazyStream(
-				const implementation = await load();
+				model,
-				if (!implementation.fetchDeferred) throw new Error("API does not support deferred responses");
+				async () => {
-				return implementation.fetchDeferred(model, handle, options);
+					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) => {

packages/ai/src/models.ts

@@ -915,6 +915,7 @@ class ModelsImpl implements MutableModels {
 				return provider.stream(requestModel, transcript, requestOptions as ApiStreamOptions<TApi>);
 			},
 			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<TApi extends Api = Api>(input: CreateProviderOpti
 	const dispatch = (
 		model: Model<Api>,
 		run: (streams: ProviderStreams) => AssistantMessageEventStream,
+		signal?: AbortSignal,
 	): AssistantMessageEventStream => {
 		const streams = apiFor(model);
 		if (!streams) {
-			return lazyStream(model, async () => {
+			return lazyStream(
-				throw new ModelsError("stream", `Provider ${input.id} has no API implementation for "${model.api}"`);
+				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<TApi extends Api = Api>(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 () => {
+			lazyStream(
-				const implementation = apiFor(model);
+				model,
-				if (!implementation?.fetchDeferred) {
+				async () => {
-					throw new ModelsError(
+					const implementation = apiFor(model);
-						"provider",
+					if (!implementation?.fetchDeferred) {
-						`Provider ${input.id} does not support deferred responses for "${model.api}"`,
+						throw new ModelsError(
-					);
+							"provider",
-				}
+							`Provider ${input.id} does not support deferred responses for "${model.api}"`,
-				return implementation.fetchDeferred(model, handle, options);
+						);
-			});
+					}
+					return implementation.fetchDeferred(model, handle, options);
+				},
+				undefined,
+				options?.signal,
+			);
 	}
 	if (streams.some((entry) => entry.cancelDeferred !== undefined)) {
 		provider.cancelDeferred = async (model, handle, options) => {

packages/coding-agent/src/core/model-runtime.ts

@@ -710,14 +710,19 @@ export class ModelRuntime implements Models {
 		options?: ModelsApiStreamOptions<TApi>,
 	): AssistantMessageEventStream {
 		const transcript = normalizeContext(context);
-		return lazyStream(model, async () => {
+		return lazyStream(
-			assertChatModel(model);
+			model,
-			const prepared = await this.prepareRequest(
+			async () => {
-				model,
+				assertChatModel(model);
-				options as (StreamOptions & ModelsRequestTransforms) | undefined,
+				const prepared = await this.prepareRequest(
-			);
+					model,
-			return prepared.provider.stream(prepared.model, transcript, prepared.options as ApiStreamOptions<TApi>);
+					options as (StreamOptions & ModelsRequestTransforms) | undefined,
-		});
+				);
+				return prepared.provider.stream(prepared.model, transcript, prepared.options as ApiStreamOptions<TApi>);
+			},
+			undefined,
+			options?.signal,
+		);
 	}
 
 	complete<TApi extends Api>(
@@ -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 () => {
+			return lazyStream(
-				const route = await this.resolveModel(model, transcript.messages, {
+				model,
-					reason: "direct",
+				async () => {
-					thinkingLevel: options?.reasoning ?? "off",
+					const route = await this.resolveModel(model, transcript.messages, {
-					signal: options?.signal,
+						reason: "direct",
-				});
+						thinkingLevel: options?.reasoning ?? "off",
-				const { maxTokens: limit } = route.model;
+						signal: options?.signal,
-				const maxTokens = options?.maxTokens && limit > 0 ? Math.min(options.maxTokens, limit) : options?.maxTokens;
+					});
-				const reasoning = route.thinkingLevel === "off" ? undefined : route.thinkingLevel;
+					const { maxTokens: limit } = route.model;
-				// Caller credentials were resolved for the virtual model's provider. Another provider
+					const maxTokens =
-				// resolves its own, so they are not sent to the wrong vendor.
+						options?.maxTokens && limit > 0 ? Math.min(options.maxTokens, limit) : options?.maxTokens;
-				const { apiKey, headers, env, ...rest } = options ?? {};
+					const reasoning = route.thinkingLevel === "off" ? undefined : route.thinkingLevel;
-				const auth = route.model.provider === model.provider ? { apiKey, headers, env } : {};
+					// Caller credentials were resolved for the virtual model's provider. Another provider
-				return this.streamSimple(route.model, context, { ...rest, ...auth, maxTokens, reasoning });
+					// 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 () => {
+		return lazyStream(
-			assertChatModel(model);
+			model,
-			const prepared = await this.prepareRequest(model, options);
+			async () => {
-			return prepared.provider.streamSimple(prepared.model, transcript, prepared.options as SimpleStreamOptions);
+				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<Api>, context: Context, options?: ModelsSimpleStreamOptions): Promise<AssistantMessage> {
@@ -765,14 +781,19 @@ export class ModelRuntime implements Models {
 		handle: DeferredHandle,
 		options?: ModelsDeferredFetchOptions,
 	): AssistantMessageEventStream {
-		return lazyStream(model, async () => {
+		return lazyStream(
-			assertChatModel(model);
+			model,
-			const prepared = await this.prepareRequest(model, options);
+			async () => {
-			if (!prepared.provider.fetchDeferred) {
+				assertChatModel(model);
-				throw new ModelsError("provider", `Provider ${model.provider} does not support deferred responses`);
+				const prepared = await this.prepareRequest(model, options);
-			}
+				if (!prepared.provider.fetchDeferred) {
-			return prepared.provider.fetchDeferred(prepared.model, handle, prepared.options as DeferredFetchOptions);
+					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(

packages/coding-agent/src/core/provider-composer.ts

@@ -591,21 +591,26 @@ export function composeModelProvider(
 		options: StreamOptions | undefined,
 		simple: boolean,
 	): AssistantMessageEventStream =>
-		lazyStream(model, async () => {
+		lazyStream(
-			if (extension?.streamSimple && model.api === extension.api) {
+			model,
-				return extension.streamSimple(model, context, options as SimpleStreamOptions);
+			async () => {
-			}
+				if (extension?.streamSimple && model.api === extension.api) {
-			if (base && supportsBaseApi(model)) {
+					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)
+					? api.streamSimple(model, context, options as SimpleStreamOptions)
-					: base.stream(model, context, options);
+					: api.stream(model, context, options);
-			}
+			},
-			const api = getApiProvider(model.api);
+			undefined,
-			if (!api) throw new Error(`No API provider registered for api: ${model.api}`);
+			options?.signal,
-			return simple
+		);
-				? api.streamSimple(model, context, options as SimpleStreamOptions)
-				: api.stream(model, context, options);
-		});
 
 	const provider: Provider = {
 		id: providerId,

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<void>((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<void>((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");
+	});
+});