feat: expose provider stream events to extensions (#9901)

Add onProviderStreamEvent through pi-ai and agent core, and surface
provider_stream_event to coding-agent extensions.

Closes #9784
This commit is contained in:
David Brailovsky 2026-09-23 16:03:30 +02:00 • committed by GitHub
parent 8d897edaa6
commit 002fc83852
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
36 changed files with 729 additions and 36 deletions

View file

@ -119,6 +119,7 @@ export interface AgentOptions {
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
onPayload?: SimpleStreamOptions["onPayload"];
onResponse?: SimpleStreamOptions["onResponse"];
onProviderStreamEvent?: SimpleStreamOptions["onProviderStreamEvent"];
beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
finishTurn?: FinishTurn;
@ -196,6 +197,7 @@ export class Agent {
public getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
public onPayload?: SimpleStreamOptions["onPayload"];
public onResponse?: SimpleStreamOptions["onResponse"];
public onProviderStreamEvent?: SimpleStreamOptions["onProviderStreamEvent"];
public beforeToolCall?: (
context: BeforeToolCallContext,
signal?: AbortSignal,
@ -235,6 +237,7 @@ export class Agent {
this.getApiKey = runtimeOptions.getApiKey;
this.onPayload = runtimeOptions.onPayload;
this.onResponse = runtimeOptions.onResponse;
this.onProviderStreamEvent = runtimeOptions.onProviderStreamEvent;
this.beforeToolCall = runtimeOptions.beforeToolCall;
this.afterToolCall = runtimeOptions.afterToolCall;
this.finishTurn = runtimeOptions.finishTurn;
@ -469,6 +472,7 @@ export class Agent {
sessionId: this.sessionId,
onPayload: this.onPayload,
onResponse: this.onResponse,
onProviderStreamEvent: this.onProviderStreamEvent,
transport: this.transport,
thinkingBudgets: this.thinkingBudgets,
maxRetryDelayMs: this.maxRetryDelayMs,

View file

@ -1166,6 +1166,28 @@ describe("Agent", () => {
expect(agent.peekQueuedMessages()).toEqual([followUp]);
});
it("forwards provider stream event observers through AgentOptions", async () => {
const providerEvents: unknown[] = [];
const agent = new Agent({
onProviderStreamEvent: (data) => {
providerEvents.push(data);
},
streamFn: (model, _context, options) => {
const stream = new MockAssistantStream();
queueMicrotask(async () => {
await options?.onProviderStreamEvent?.({ request_cost: 0.01 }, model);
const message = createAssistantMessage("ok");
stream.push({ type: "done", reason: "stop", message });
});
return stream;
},
});
await agent.prompt("hello");
expect(providerEvents).toEqual([{ request_cost: 0.01 }]);
});
it("forwards sessionId to streamFunction options", async () => {
let receivedSessionId: string | undefined;
const agent = new Agent({