sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

fix(ai): preserve aborted stop reason during lazy setup (#8635)

Browse source 
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.
This commit is contained in:
Sai Mallikarjun 2026-11-04 14:53:35 +05:30 • committed by GitHub
parent 7a53c310a3
commit f3a64af07b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 187 additions and 76 deletions
Download patch file Download diff file Expand all files Collapse all files

39
packages/ai/src/api/lazy.ts
View file

@ -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 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<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 () => {
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) => {

45
packages/ai/src/models.ts
View file

@ -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 () => {
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<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 () => {
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) => {