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

93
packages/coding-agent/src/core/model-runtime.ts
View file

@ -710,14 +710,19 @@ export class ModelRuntime implements Models {
options?: ModelsApiStreamOptions<TApi>,
): 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<TApi>);
});
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<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 () => {
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<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 () => {
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(

33
packages/coding-agent/src/core/provider-composer.ts
View file

@ -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,

53
packages/coding-agent/test/suite/regressions/8409-aborted-tool-auth-setup.test.ts Normal file
View file

@ -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");
});
});