diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 92af2c791..ee8e6e591 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -119,6 +119,7 @@ export interface AgentOptions { getApiKey?: (provider: string) => Promise | string | undefined; onPayload?: SimpleStreamOptions["onPayload"]; onResponse?: SimpleStreamOptions["onResponse"]; + onProviderStreamEvent?: SimpleStreamOptions["onProviderStreamEvent"]; beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise; afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise; finishTurn?: FinishTurn; @@ -196,6 +197,7 @@ export class Agent { public getApiKey?: (provider: string) => Promise | 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, diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index bed5dd427..2a339944e 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -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({ diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 3fd01039a..8ff552c5f 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added `onProviderStreamEvent` to observe parsed provider stream events before normalization, including provider-specific fields not retained in assistant messages ([#9784](https://github.com/earendil-works/pi/issues/9784)). + ## [0.87.1] - 2026-09-22 ### Added diff --git a/packages/ai/README.md b/packages/ai/README.md index 7395672ff..e5f19c4d0 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -38,6 +38,7 @@ Unified LLM API with provider collections, automatic auth resolution, token and - [Aborting Requests](#aborting-requests) - [Continuing After Abort](#continuing-after-abort) - [Debugging Provider Payloads](#debugging-provider-payloads) + - [Observing Provider Stream Events](#observing-provider-stream-events) - [Custom Providers](#custom-providers) - [createProvider()](#createprovider) - [Calling API Implementations Directly](#calling-api-implementations-directly) @@ -1030,6 +1031,25 @@ const response = await models.complete(model, context, { The callback is supported by `stream`, `complete`, `streamSimple`, and `completeSimple`. +### Observing Provider Stream Events + +Use `onProviderStreamEvent` to inspect provider-specific fields that Pi does not include in `AssistantMessage`. The callback receives the parsed event available to the adapter before Pi normalizes it. Treat the event as read-only because mutations can affect normalization. This is not guaranteed to be the original HTTP bytes or SSE frame. + +```typescript +const openRouterModel = models.getModel('openrouter', 'openrouter/auto')!; +const response = await models.complete(openRouterModel, context, { + headers: { "X-OpenRouter-Metadata": "enabled" }, + onProviderStreamEvent: (data) => { + const chunk = data as Record; + if (chunk.openrouter_metadata) { + console.log(chunk.openrouter_metadata); + } + }, +}); +``` + +Callbacks are awaited in stream order, so slow callbacks delay stream consumption and thrown errors fail the request. SDK-backed adapters can expose only fields retained by their SDK. + ## Custom Providers ### createProvider() diff --git a/packages/ai/src/api/anthropic-messages.ts b/packages/ai/src/api/anthropic-messages.ts index f25b2b11d..a935e01c5 100644 --- a/packages/ai/src/api/anthropic-messages.ts +++ b/packages/ai/src/api/anthropic-messages.ts @@ -599,6 +599,7 @@ export const stream: StreamFunction<"anthropic-messages", AnthropicOptions> = ( const blocks = output.content as Block[]; for await (const event of iterateAnthropicEvents(response, options?.signal)) { + await options?.onProviderStreamEvent?.(event, model); if (event.type === "message_start") { output.responseId = event.message.id; const transformations = event.message.input_transformations; diff --git a/packages/ai/src/api/azure-openai-responses.ts b/packages/ai/src/api/azure-openai-responses.ts index b275d9f75..558f683ce 100644 --- a/packages/ai/src/api/azure-openai-responses.ts +++ b/packages/ai/src/api/azure-openai-responses.ts @@ -130,7 +130,10 @@ export const stream: StreamFunction<"azure-openai-responses", AzureOpenAIRespons await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model); stream.push({ type: "start", partial: output }); - await processResponsesStream(openaiStream, output, stream, model, { grammarToolInputProperties }); + await processResponsesStream(openaiStream, output, stream, model, { + onProviderStreamEvent: options?.onProviderStreamEvent, + grammarToolInputProperties, + }); if (options?.signal?.aborted) { throw new Error("Request was aborted"); diff --git a/packages/ai/src/api/bedrock-converse-stream.ts b/packages/ai/src/api/bedrock-converse-stream.ts index ee66b4599..74f32f5ee 100644 --- a/packages/ai/src/api/bedrock-converse-stream.ts +++ b/packages/ai/src/api/bedrock-converse-stream.ts @@ -294,6 +294,7 @@ export const stream: StreamFunction<"bedrock-converse-stream", BedrockOptions> = } for await (const item of response.stream!) { + await options.onProviderStreamEvent?.(item, model); if (item.messageStart) { if (item.messageStart.role !== ConversationRole.ASSISTANT) { throw new Error("Unexpected assistant message start but got user message start instead"); diff --git a/packages/ai/src/api/google-generative-ai.ts b/packages/ai/src/api/google-generative-ai.ts index 177e92af2..24fe9134c 100644 --- a/packages/ai/src/api/google-generative-ai.ts +++ b/packages/ai/src/api/google-generative-ai.ts @@ -104,6 +104,7 @@ export const stream: StreamFunction<"google-generative-ai", GoogleOptions> = ( const blocks = output.content; const blockIndex = () => blocks.length - 1; for await (const chunk of googleStream) { + await options?.onProviderStreamEvent?.(chunk, model); // @google/genai documents GenerateContentResponse.responseId as an output-only field // used to identify each response. Keep the first non-empty one from the stream. output.responseId ||= chunk.responseId; diff --git a/packages/ai/src/api/google-vertex.ts b/packages/ai/src/api/google-vertex.ts index cfbe20d37..4e667df18 100644 --- a/packages/ai/src/api/google-vertex.ts +++ b/packages/ai/src/api/google-vertex.ts @@ -113,6 +113,7 @@ export const stream: StreamFunction<"google-vertex", GoogleVertexOptions> = ( const blocks = output.content; const blockIndex = () => blocks.length - 1; for await (const chunk of googleStream) { + await options?.onProviderStreamEvent?.(chunk, model); // Vertex uses the same @google/genai GenerateContentResponse type as Gemini. // responseId is documented there as an output-only identifier for each response. output.responseId ||= chunk.responseId; diff --git a/packages/ai/src/api/mistral-conversations.ts b/packages/ai/src/api/mistral-conversations.ts index 837d1ec11..899ac3cbf 100644 --- a/packages/ai/src/api/mistral-conversations.ts +++ b/packages/ai/src/api/mistral-conversations.ts @@ -150,7 +150,7 @@ export const stream: StreamFunction<"mistral-conversations", MistralOptions> = ( } const mistralStream = await requestMistralStream(model, payload, apiKey, options); stream.push({ type: "start", partial: output }); - await consumeChatStream(model, output, stream, mistralStream); + await consumeChatStream(model, output, stream, mistralStream, options?.onProviderStreamEvent); if (options?.signal?.aborted) { throw new Error("Request was aborted"); @@ -559,6 +559,7 @@ async function consumeChatStream( output: AssistantMessage, stream: AssistantMessageEventStream, mistralStream: AsyncIterable, + onProviderStreamEvent: StreamOptions["onProviderStreamEvent"], ): Promise { let currentBlock: TextContent | ThinkingContent | null = null; const blocks = output.content; @@ -588,6 +589,7 @@ async function consumeChatStream( for await (const event of mistralStream) { const chunk = event.data; + await onProviderStreamEvent?.(chunk, model); // Mistral's streamed CompletionChunk carries an id field. Keep the first non-empty one, // mirroring how OpenAI-style streaming exposes a stable response identifier per stream. output.responseId ||= chunk.id; diff --git a/packages/ai/src/api/openai-codex-responses.ts b/packages/ai/src/api/openai-codex-responses.ts index 8c0e54466..9c1831a02 100644 --- a/packages/ai/src/api/openai-codex-responses.ts +++ b/packages/ai/src/api/openai-codex-responses.ts @@ -665,12 +665,18 @@ async function processStream( grammarToolInputProperties: ReadonlyMap, options?: OpenAICodexResponsesOptions, ): Promise { - await processResponsesStream(mapCodexEvents(parseSSE(response, options?.signal), output), output, stream, model, { - serviceTier: options?.serviceTier, - grammarToolInputProperties, - resolveServiceTier: resolveCodexServiceTier, - applyServiceTierPricing: (usage, serviceTier) => applyServiceTierPricing(usage, serviceTier, model), - }); + await processResponsesStream( + mapCodexEvents(parseSSE(response, options?.signal), output, model, options?.onProviderStreamEvent), + output, + stream, + model, + { + serviceTier: options?.serviceTier, + grammarToolInputProperties, + resolveServiceTier: resolveCodexServiceTier, + applyServiceTierPricing: (usage, serviceTier) => applyServiceTierPricing(usage, serviceTier, model), + }, + ); } class CodexApiError extends Error { @@ -697,8 +703,20 @@ class CodexProtocolError extends Error { } } +class ProviderStreamEventCallbackError extends Error { + constructor(cause: unknown) { + super(formatThrownValue(cause)); + this.name = "ProviderStreamEventCallbackError"; + this.cause = cause; + } +} + function isCodexNonTransportError(error: unknown): boolean { - return error instanceof CodexApiError || error instanceof CodexProtocolError; + return ( + error instanceof CodexApiError || + error instanceof CodexProtocolError || + error instanceof ProviderStreamEventCallbackError + ); } function isWebSocketConnectionLimitReachedError(error: unknown): boolean { @@ -725,8 +743,16 @@ function extractCodexEventError(event: Record): { code?: string async function* mapCodexEvents( events: AsyncIterable>, output: AssistantMessage, + model: Model<"openai-codex-responses">, + onProviderStreamEvent: StreamOptions["onProviderStreamEvent"] | undefined, ): AsyncGenerator { for await (const event of events) { + try { + await onProviderStreamEvent?.(event, model); + } catch (error) { + // Keep callback failures out of Codex's WebSocket retry and SSE fallback path. + throw new ProviderStreamEventCallbackError(error); + } const type = typeof event.type === "string" ? event.type : undefined; if (!type) continue; @@ -1516,7 +1542,12 @@ async function processWebSocketStream( socket.send(JSON.stringify({ type: "response.create", ...requestBody })); await processResponsesStream( startWebSocketOutputOnFirstEvent( - mapCodexEvents(parseWebSocket(socket, options?.signal, idleTimeoutMs), output), + mapCodexEvents( + parseWebSocket(socket, options?.signal, idleTimeoutMs), + output, + model, + options?.onProviderStreamEvent, + ), onStart, ), output, diff --git a/packages/ai/src/api/openai-completions.ts b/packages/ai/src/api/openai-completions.ts index f5a1bc668..4ea3695b9 100644 --- a/packages/ai/src/api/openai-completions.ts +++ b/packages/ai/src/api/openai-completions.ts @@ -551,6 +551,7 @@ export const stream: StreamFunction<"openai-completions", OpenAICompletionsOptio }; for await (const chunk of openaiStream) { + await options?.onProviderStreamEvent?.(chunk, model); if (!chunk || typeof chunk !== "object") continue; // OpenAI documents ChatCompletionChunk.id as the unique chat completion identifier, diff --git a/packages/ai/src/api/openai-responses-shared.ts b/packages/ai/src/api/openai-responses-shared.ts index 80bf45370..9fc599c2e 100644 --- a/packages/ai/src/api/openai-responses-shared.ts +++ b/packages/ai/src/api/openai-responses-shared.ts @@ -20,6 +20,7 @@ import type { ImageContent, Model, StopReason, + StreamOptions, SystemMessage, TextContent, TextSignatureV1, @@ -107,6 +108,7 @@ function convertToolResultOutput( } export interface OpenAIResponsesStreamOptions { + onProviderStreamEvent?: StreamOptions["onProviderStreamEvent"]; serviceTier?: ResponseCreateParamsStreaming["service_tier"]; grammarToolInputProperties?: ReadonlyMap; resolveServiceTier?: ( @@ -596,6 +598,7 @@ export async function processResponsesStream( }; for await (const event of openaiStream) { + await options?.onProviderStreamEvent?.(event, model); if (event.type === "response.created") { output.responseId = event.response.id; } else if (event.type === "response.output_item.added") { diff --git a/packages/ai/src/api/openai-responses.ts b/packages/ai/src/api/openai-responses.ts index f7a0c6088..811a0d4b9 100644 --- a/packages/ai/src/api/openai-responses.ts +++ b/packages/ai/src/api/openai-responses.ts @@ -178,6 +178,7 @@ export const stream: StreamFunction<"openai-responses", OpenAIResponsesOptions> stream.push({ type: "start", partial: output }); await processResponsesStream(openaiStream, output, stream, model, { + onProviderStreamEvent: options?.onProviderStreamEvent, serviceTier: options?.serviceTier, grammarToolInputProperties, applyServiceTierPricing: (usage, serviceTier) => applyServiceTierPricing(usage, serviceTier, model), diff --git a/packages/ai/src/api/pi-messages.ts b/packages/ai/src/api/pi-messages.ts index 4ed6ceadd..2137b551f 100644 --- a/packages/ai/src/api/pi-messages.ts +++ b/packages/ai/src/api/pi-messages.ts @@ -412,6 +412,7 @@ export const stream: StreamFunction<"pi-messages", PiMessagesOptions> = ( } for await (const piEvent of readPiMessagesEvents(response.body)) { + await options?.onProviderStreamEvent?.(piEvent, model); const event = convertEvent(piEvent); eventStream.push(event); if (event.type === "done" || event.type === "error") { diff --git a/packages/ai/src/api/simple-options.ts b/packages/ai/src/api/simple-options.ts index 63e34f3c7..59f7196b4 100644 --- a/packages/ai/src/api/simple-options.ts +++ b/packages/ai/src/api/simple-options.ts @@ -42,6 +42,7 @@ export function buildBaseOptions( headers: options?.headers, onPayload: options?.onPayload, onResponse: options?.onResponse, + onProviderStreamEvent: options?.onProviderStreamEvent, timeoutMs: options?.timeoutMs, websocketConnectTimeoutMs: options?.websocketConnectTimeoutMs, maxRetries: options?.maxRetries, diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 7646a9c05..07c38cfe8 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -189,6 +189,12 @@ export interface StreamOptions extends ProviderRequestOptions> { * its body stream is consumed. */ onResponse?: (response: ProviderResponse, model: Model) => void | Promise; + /** + * Optional observer for each parsed provider stream event before Pi normalization. + * Event data is adapter-owned and must be treated as read-only. + * Adapter support is explicit; unsupported adapters do not invoke it. + */ + onProviderStreamEvent?: (data: unknown, model: Model) => void | Promise; temperature?: number; /** * Arbitrary sampling parameters merged into the request body as-is, after the named request diff --git a/packages/ai/test/anthropic-sse-parsing.test.ts b/packages/ai/test/anthropic-sse-parsing.test.ts index 36de46d26..b39780bff 100644 --- a/packages/ai/test/anthropic-sse-parsing.test.ts +++ b/packages/ai/test/anthropic-sse-parsing.test.ts @@ -4,7 +4,7 @@ import { describe, expect, it } from "vitest"; import { stream as streamAnthropic } from "../src/api/anthropic-messages.ts"; import { transformMessages } from "../src/api/transform-messages.ts"; import { getModel, normalizeContext } from "../src/compat.ts"; -import type { Model, ToolCall } from "../src/types.ts"; +import type { Api, Model, ToolCall } from "../src/types.ts"; function createSseResponse(events: Array<{ event: string; data: string }>): Response { const body = events.map(({ event, data }) => `event: ${event}\ndata: ${data}\n`).join("\n"); @@ -110,6 +110,35 @@ function createResponseModelSseResponse(model: string, contentBlock: ResponseCon } describe("Anthropic raw SSE parsing", () => { + it("forwards parsed provider stream events in order", async () => { + const model = getModel("anthropic", "claude-haiku-4-5"); + const providerEvents: unknown[] = []; + const eventModels: Model[] = []; + const result = await streamAnthropic( + model, + normalizeContext({ messages: [{ role: "user", content: "Hello", timestamp: 1 }] }), + { + client: createFakeAnthropicClient(createSseResponse(minimalAnthropicEvents)), + onProviderStreamEvent: async (event, eventModel) => { + await Promise.resolve(); + providerEvents.push(event); + eventModels.push(eventModel); + }, + }, + ).result(); + + expect(result.stopReason).toBe("stop"); + expect(providerEvents.map((event) => (event as { type: string }).type)).toEqual([ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ]); + expect(eventModels).toEqual([model, model, model, model, model, model]); + }); + it("keeps signed thinking replayable when a proxy relabels the model", async () => { // Regression test for earendil-works/pi#9188. const model = getModel("anthropic", "claude-opus-5"); diff --git a/packages/ai/test/azure-openai-base-url.test.ts b/packages/ai/test/azure-openai-base-url.test.ts index 0e37cb6ee..9ac8ce4be 100644 --- a/packages/ai/test/azure-openai-base-url.test.ts +++ b/packages/ai/test/azure-openai-base-url.test.ts @@ -1,9 +1,10 @@ import { arch, platform, release } from "node:os"; +import type { ResponseStreamEvent } from "openai/resources/responses/responses.js"; import { Type } from "typebox"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { stream as streamAzureOpenAIResponses } from "../src/api/azure-openai-responses.ts"; import { getModel, normalizeContext } from "../src/compat.ts"; -import type { Context, Model } from "../src/types.ts"; +import type { Api, Context, Model } from "../src/types.ts"; interface CapturedAzureClientOptions { apiKey: string; @@ -22,6 +23,7 @@ interface CapturedAzureResponsesPayload { const azureMock = vi.hoisted(() => ({ constructorCalls: [] as CapturedAzureClientOptions[], lastParams: undefined as CapturedAzureResponsesPayload | undefined, + streamEvents: undefined as ResponseStreamEvent[] | undefined, })); vi.mock("openai", () => { @@ -29,7 +31,16 @@ vi.mock("openai", () => { responses = { create: (params: CapturedAzureResponsesPayload) => { azureMock.lastParams = params; - throw new Error("mock create"); + const events = azureMock.streamEvents; + if (!events) throw new Error("mock create"); + return { + withResponse: async () => ({ + data: (async function* () { + for (const event of events) yield event; + })(), + response: { status: 200, headers: new Headers() }, + }), + }; }, }; @@ -55,6 +66,7 @@ const originalAzureOpenAIApiKey = process.env.AZURE_OPENAI_API_KEY; beforeEach(() => { azureMock.constructorCalls.length = 0; azureMock.lastParams = undefined; + azureMock.streamEvents = undefined; delete process.env.AZURE_OPENAI_BASE_URL; delete process.env.AZURE_OPENAI_RESOURCE_NAME; delete process.env.AZURE_OPENAI_API_VERSION; @@ -218,6 +230,38 @@ describe("azure-openai-responses base URL normalization", () => { }); }); +describe("azure-openai-responses provider stream events", () => { + it("forwards parsed events in order before normalizing the response", async () => { + azureMock.streamEvents = [ + { type: "response.created", sequence_number: 0, response: { id: "resp_azure" } } as ResponseStreamEvent, + { + type: "response.completed", + sequence_number: 1, + response: { id: "resp_azure", status: "completed" }, + } as ResponseStreamEvent, + ]; + const model = getModel("azure-openai-responses", "gpt-4o-mini"); + const received: unknown[] = []; + const eventModels: Model[] = []; + const result = await streamAzureOpenAIResponses(model, normalizeContext(context), { + apiKey: "test-api-key", + azureBaseUrl: "https://my-resource.openai.azure.com", + onProviderStreamEvent: async (event, eventModel) => { + await Promise.resolve(); + received.push(event); + eventModels.push(eventModel); + }, + }).result(); + + expect(received).toEqual(azureMock.streamEvents); + expect(received[0]).toBe(azureMock.streamEvents[0]); + expect(received[1]).toBe(azureMock.streamEvents[1]); + expect(eventModels).toEqual([model, model]); + expect(result.stopReason).toBe("stop"); + expect(result.responseId).toBe("resp_azure"); + }); +}); + describe("azure-openai-responses user agent", () => { it("uses pi's User-Agent by default", async () => { expect((await captureClientHeaders())["User-Agent"]).toBe(PI_USER_AGENT); diff --git a/packages/ai/test/bedrock-raw-stop-reason.test.ts b/packages/ai/test/bedrock-raw-stop-reason.test.ts index 60dc7517f..c9dcac753 100644 --- a/packages/ai/test/bedrock-raw-stop-reason.test.ts +++ b/packages/ai/test/bedrock-raw-stop-reason.test.ts @@ -1,7 +1,8 @@ -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; const bedrockMock = vi.hoisted(() => ({ stopReason: "end_turn" as string, + streamEvents: undefined as unknown[] | undefined, })); vi.mock("@aws-sdk/client-bedrock-runtime", () => { @@ -16,6 +17,10 @@ vi.mock("@aws-sdk/client-bedrock-runtime", () => { return { $metadata: { httpStatusCode: 200, requestId: "request-id" }, stream: (async function* () { + if (bedrockMock.streamEvents) { + yield* bedrockMock.streamEvents; + return; + } yield { messageStart: { role: "assistant" } }; yield { messageStop: { stopReason: bedrockMock.stopReason } }; })(), @@ -52,12 +57,60 @@ vi.mock("@aws-sdk/client-bedrock-runtime", () => { import { stream as streamBedrock } from "../src/api/bedrock-converse-stream.ts"; import { getModel, normalizeContext } from "../src/compat.ts"; +import type { Api, Model } from "../src/types.ts"; + +beforeEach(() => { + bedrockMock.streamEvents = undefined; +}); const model = getModel("amazon-bedrock", "us.anthropic.claude-opus-4-8"); const context = normalizeContext({ messages: [{ role: "user", content: "hello", timestamp: Date.now() }], }); +describe("Bedrock provider stream events", () => { + it("forwards SDK stream items in order before normalizing them", async () => { + bedrockMock.streamEvents = [ + { messageStart: { role: "assistant" } }, + { contentBlockDelta: { contentBlockIndex: 0, delta: { text: "hello" } } }, + { messageStop: { stopReason: "end_turn", additionalModelResponseFields: { source: "test" } } }, + { metadata: { usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 } } }, + ]; + const received: unknown[] = []; + const eventModels: Model[] = []; + const result = await streamBedrock(model, context, { + cacheRetention: "none", + onProviderStreamEvent: async (item, eventModel) => { + await Promise.resolve(); + received.push(item); + eventModels.push(eventModel); + }, + }).result(); + + expect(received).toEqual(bedrockMock.streamEvents); + for (const [index, item] of received.entries()) expect(item).toBe(bedrockMock.streamEvents[index]); + expect(eventModels).toEqual([model, model, model, model]); + expect(result.stopReason).toBe("stop"); + expect(result.content).toEqual([{ type: "text", text: "hello" }]); + }); + + it("forwards SDK error items before reporting them", async () => { + const exception = new Error("bedrock stream failed"); + bedrockMock.streamEvents = [{ messageStart: { role: "assistant" } }, { internalServerException: exception }]; + const received: unknown[] = []; + const result = await streamBedrock(model, context, { + cacheRetention: "none", + onProviderStreamEvent: (item) => { + received.push(item); + }, + }).result(); + + expect(received).toEqual(bedrockMock.streamEvents); + expect(result.stopReason).toBe("error"); + expect(result.errorMessage).toBe("bedrock stream failed"); + }); +}); + describe("Bedrock raw stop reasons", () => { it("preserves raw Bedrock stop reasons for successful stops", async () => { bedrockMock.stopReason = "end_turn"; diff --git a/packages/ai/test/google-raw-stop-reason.test.ts b/packages/ai/test/google-raw-stop-reason.test.ts index c9c5cc099..37b79a66e 100644 --- a/packages/ai/test/google-raw-stop-reason.test.ts +++ b/packages/ai/test/google-raw-stop-reason.test.ts @@ -1,10 +1,11 @@ import { arch, platform, release } from "node:os"; -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; const googleGenAiMock = vi.hoisted(() => ({ constructorCalls: [] as Array>, finishReason: "MALFORMED_FUNCTION_CALL", includeFunctionCall: false, + streamChunks: undefined as Array> | undefined, })); vi.mock("@google/genai", () => { @@ -15,6 +16,10 @@ vi.mock("@google/genai", () => { models = { generateContentStream: async function* () { + if (googleGenAiMock.streamChunks) { + yield* googleGenAiMock.streamChunks; + return; + } yield { responseId: "google-response-id", candidates: [ @@ -88,6 +93,11 @@ vi.mock("@google/genai", () => { import { stream as streamGoogleGenerativeAi } from "../src/api/google-generative-ai.ts"; import { stream as streamGoogleVertex } from "../src/api/google-vertex.ts"; import { getModel, normalizeContext } from "../src/compat.ts"; +import type { Api, Model, StreamOptions } from "../src/types.ts"; + +beforeEach(() => { + googleGenAiMock.streamChunks = undefined; +}); const PI_USER_AGENT = `pi (${platform()} ${release()}; ${arch()})`; @@ -182,6 +192,63 @@ describe("Google raw stop reasons", () => { }); }); +describe("Google provider stream events", () => { + const googleModel = getModel("google", "gemini-2.5-flash"); + const vertexModel = getModel("google-vertex", "gemini-3-flash-preview"); + const adapters = [ + { + name: "Google Generative AI", + model: googleModel, + createStream: (onProviderStreamEvent: StreamOptions["onProviderStreamEvent"]) => + streamGoogleGenerativeAi(googleModel, context, { + apiKey: "test-api-key", + onProviderStreamEvent, + }), + }, + { + name: "Google Vertex", + model: vertexModel, + createStream: (onProviderStreamEvent: StreamOptions["onProviderStreamEvent"]) => + streamGoogleVertex(vertexModel, context, { + project: "test-project", + location: "us-central1", + onProviderStreamEvent, + }), + }, + ]; + + it.each(adapters)( + "forwards each SDK chunk in order before normalizing it for $name", + async ({ model, createStream }) => { + googleGenAiMock.streamChunks = [ + { + responseId: "resp_google", + candidates: [{ content: { parts: [{ text: "hello" }] } }], + }, + { + candidates: [{ finishReason: "STOP" }], + usageMetadata: { promptTokenCount: 2, candidatesTokenCount: 1, totalTokenCount: 3 }, + }, + ]; + const received: unknown[] = []; + const eventModels: Model[] = []; + const result = await createStream(async (chunk, eventModel) => { + await Promise.resolve(); + received.push(chunk); + eventModels.push(eventModel); + }).result(); + + expect(received).toEqual(googleGenAiMock.streamChunks); + expect(received[0]).toBe(googleGenAiMock.streamChunks[0]); + expect(received[1]).toBe(googleGenAiMock.streamChunks[1]); + expect(eventModels).toEqual([model, model]); + expect(result.stopReason).toBe("stop"); + expect(result.responseId).toBe("resp_google"); + expect(result.content).toEqual([{ type: "text", text: "hello" }]); + }, + ); +}); + describe("Google Generative AI user agent", () => { it("uses pi's User-Agent by default", async () => { expect((await captureGoogleHeaders())["User-Agent"]).toBe(PI_USER_AGENT); diff --git a/packages/ai/test/mistral-http-transport.test.ts b/packages/ai/test/mistral-http-transport.test.ts index 9180837f6..83c91c8fc 100644 --- a/packages/ai/test/mistral-http-transport.test.ts +++ b/packages/ai/test/mistral-http-transport.test.ts @@ -3,7 +3,7 @@ import { Type } from "typebox"; import { describe, expect, it } from "vitest"; import { stream as streamMistral } from "../src/api/mistral-conversations.ts"; import { getModel, normalizeContext } from "../src/compat.ts"; -import type { FetchFunction, ProviderResponse } from "../src/types.ts"; +import type { Api, FetchFunction, Model, ProviderResponse } from "../src/types.ts"; const PI_USER_AGENT = `pi (${platform()} ${release()}; ${arch()})`; @@ -322,6 +322,41 @@ describe("Mistral HTTP transport", () => { expect(message.usage).toMatchObject({ input: 7, output: 4, cacheRead: 3, cacheWrite: 0, totalTokens: 14 }); }); + it("forwards each parsed SSE payload before normalizing it", async () => { + const model = getModel("mistral", "mistral-large-latest"); + const context = normalizeContext({ + messages: [{ role: "user", content: "hello", timestamp: 1 }], + }); + const events = [ + { + id: "response-1", + provider_metadata: { request: "test" }, + choices: [{ index: 0, finish_reason: null, delta: { content: "hello" } }], + }, + { + id: "response-1", + choices: [{ index: 0, finish_reason: "stop", delta: {} }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + ]; + const received: unknown[] = []; + const eventModels: Model[] = []; + const result = await streamMistral(model, context, { + apiKey: "test", + fetch: async () => createSseResponse(events), + onProviderStreamEvent: async (event, eventModel) => { + await Promise.resolve(); + received.push(event); + eventModels.push(eventModel); + }, + }).result(); + + expect(received).toEqual(events); + expect(eventModels).toEqual([model, model]); + expect(result.stopReason).toBe("stop"); + expect(result.content).toEqual([{ type: "text", text: "hello" }]); + }); + it("parses SSE and UTF-8 sequences split across transport chunks", async () => { const model = getModel("mistral", "mistral-large-latest"); const context = normalizeContext({ diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index b756a1479..9aabc09c3 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -11,7 +11,7 @@ import { stream as streamOpenAICodexResponses, streamSimple as streamSimpleOpenAICodexResponses, } from "../src/api/openai-codex-responses.ts"; -import type { Context, Model } from "../src/types.ts"; +import type { Api, Context, Model } from "../src/types.ts"; import { normalizeContext } from "../src/utils/transcript.ts"; const originalAgentDir = process.env.PI_CODING_AGENT_DIR; @@ -98,7 +98,7 @@ function buildSSEPayload({ } describe("openai-codex streaming", () => { - it("streams SSE responses into AssistantMessageEventStream", async () => { + it("streams SSE responses and forwards raw provider events", async () => { const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-")); process.env.PI_CODING_AGENT_DIR = tempDir; @@ -192,9 +192,15 @@ describe("openai-codex streaming", () => { messages: [{ role: "user", content: "Say hello", timestamp: Date.now() }], }); + const providerEvents: unknown[] = []; + const providerEventModels: Model[] = []; const streamResult = streamOpenAICodexResponses(model, context, { apiKey: token, transport: "sse", + onProviderStreamEvent: (event, eventModel) => { + providerEvents.push(event); + providerEventModels.push(eventModel); + }, }); let sawTextDelta = false; let sawDone = false; @@ -211,6 +217,14 @@ describe("openai-codex streaming", () => { expect(sawTextDelta).toBe(true); expect(sawDone).toBe(true); + expect(providerEvents.map((event) => (event as { type?: string }).type)).toEqual([ + "response.output_item.added", + "response.content_part.added", + "response.output_text.delta", + "response.output_item.done", + "response.completed", + ]); + expect(providerEventModels).toEqual([model, model, model, model, model]); }); // Regression test for https://github.com/earendil-works/pi/issues/9047 @@ -1266,9 +1280,11 @@ describe("openai-codex streaming", () => { }); await streamResult.result(); }); - it("forwards auto transport from streamSimple options and uses cached websocket context", async () => { + it("forwards auto transport and raw provider events from streamSimple", async () => { const token = mockToken(); const sentBodies: unknown[] = []; + const providerEvents: unknown[] = []; + const providerEventModels: Model[] = []; let capturedWebSocketHeaders: Record | undefined; const fetchMock = vi.fn(async () => new Response("unexpected fetch", { status: 500 })); @@ -1317,7 +1333,7 @@ describe("openai-codex streaming", () => { }, }, { - type: "response.completed", + type: "response.done", response: { status: "completed", end_turn: false, @@ -1369,10 +1385,22 @@ describe("openai-codex streaming", () => { apiKey: token, sessionId: "session-auto", transport: "auto", + onProviderStreamEvent: (event, eventModel) => { + providerEvents.push(event); + providerEventModels.push(eventModel); + }, }).result(); expect(result.endTurn).toBe(false); expect(sentBodies).toHaveLength(1); + expect(providerEvents.map((event) => (event as { type?: string }).type)).toEqual([ + "response.output_item.added", + "response.content_part.added", + "response.output_text.delta", + "response.output_item.done", + "response.done", + ]); + expect(providerEventModels).toEqual([model, model, model, model, model]); expect(capturedWebSocketHeaders?.["session-id"]).toBe("session-auto"); expect(capturedWebSocketHeaders?.session_id).toBeUndefined(); expect(capturedWebSocketHeaders?.["x-client-request-id"]).toBe("session-auto"); diff --git a/packages/ai/test/openai-completions-provider-stream-event.test.ts b/packages/ai/test/openai-completions-provider-stream-event.test.ts new file mode 100644 index 000000000..b025599d3 --- /dev/null +++ b/packages/ai/test/openai-completions-provider-stream-event.test.ts @@ -0,0 +1,94 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { completeSimple } from "../src/compat.ts"; +import type { Model } from "../src/types.ts"; + +const mockState = vi.hoisted(() => ({ + chunks: [] as unknown[], +})); + +vi.mock("openai", () => { + class FakeOpenAI { + chat = { + completions: { + create: () => { + const stream = { + async *[Symbol.asyncIterator]() { + for (const chunk of mockState.chunks) yield chunk; + }, + }; + const promise = Promise.resolve(stream) as Promise & { + withResponse: () => Promise<{ + data: typeof stream; + response: { status: number; headers: Headers }; + }>; + }; + promise.withResponse = async () => ({ + data: stream, + response: { status: 200, headers: new Headers({ "x-request-id": "req-1" }) }, + }); + return promise; + }, + }, + }; + } + return { default: FakeOpenAI }; +}); + +function openRouterModel(): Model<"openai-completions"> { + return { + id: "openrouter/auto", + name: "OpenRouter Auto", + api: "openai-completions", + provider: "openrouter", + baseUrl: "https://openrouter.ai/api/v1", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 200_000, + maxTokens: 8192, + }; +} + +describe("openai-completions provider stream events", () => { + beforeEach(() => { + mockState.chunks = []; + }); + + // Regression test for #9784. + it("exposes provider chunks including OpenRouter metadata", async () => { + const firstChunk = { + id: "chatcmpl-1", + model: "anthropic/claude-sonnet-4.6", + choices: [{ index: 0, delta: { content: "hello" } }], + }; + const finalChunk = { + id: "chatcmpl-1", + model: "anthropic/claude-sonnet-4.6", + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + usage: { + prompt_tokens: 10, + completion_tokens: 2, + total_tokens: 12, + cost: 0.0012, + is_byok: false, + }, + openrouter_metadata: { strategy: "direct", region: "iad" }, + }; + mockState.chunks = [firstChunk, finalChunk]; + const events: unknown[] = []; + + const message = await completeSimple( + openRouterModel(), + { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }, + { + apiKey: "test", + onProviderStreamEvent: (data) => { + events.push(data); + }, + }, + ); + + expect(message.content).toEqual([{ type: "text", text: "hello" }]); + expect(events).toEqual([firstChunk, finalChunk]); + }); +}); diff --git a/packages/ai/test/openai-responses-terminal-event.test.ts b/packages/ai/test/openai-responses-terminal-event.test.ts index 8b0553c25..89f54f8b4 100644 --- a/packages/ai/test/openai-responses-terminal-event.test.ts +++ b/packages/ai/test/openai-responses-terminal-event.test.ts @@ -2,7 +2,7 @@ import type { ResponseStreamEvent } from "openai/resources/responses/responses.j import { describe, expect, it, vi } from "vitest"; import { stream as streamOpenAIResponses } from "../src/api/openai-responses.ts"; import { processResponsesStream } from "../src/api/openai-responses-shared.ts"; -import type { AssistantMessage, AssistantMessageEvent, Model } from "../src/types.ts"; +import type { Api, AssistantMessage, AssistantMessageEvent, Model } from "../src/types.ts"; import { AssistantMessageEventStream } from "../src/utils/event-stream.ts"; import { normalizeContext } from "../src/utils/transcript.ts"; @@ -215,6 +215,35 @@ describe("OpenAI Responses terminal event handling", () => { ); }); + it("forwards parsed provider stream events in order", async () => { + const model = createModel(); + const context = normalizeContext({ + systemPrompt: "", + messages: [{ role: "user", content: [{ type: "text", text: "hi" }], timestamp: 0 }], + tools: [], + }); + const providerEvents: unknown[] = []; + const eventModels: Model[] = []; + const stream = streamOpenAIResponses(model, context, { + apiKey: "test", + onProviderStreamEvent: async (event, eventModel) => { + await Promise.resolve(); + providerEvents.push(event); + eventModels.push(eventModel); + }, + }); + + await stream.result(); + + expect(providerEvents).toHaveLength(3); + expect(providerEvents.map((event) => (event as ResponseStreamEvent).type)).toEqual([ + "response.created", + "response.output_item.added", + "response.reasoning_text.delta", + ]); + expect(eventModels).toEqual([model, model, model]); + }); + it("emits an error final result when the wrapper stream ends before a terminal response event", async () => { const model = createModel(); const context = normalizeContext({ diff --git a/packages/ai/test/pi-messages.test.ts b/packages/ai/test/pi-messages.test.ts index 809ef7516..211c24cb0 100644 --- a/packages/ai/test/pi-messages.test.ts +++ b/packages/ai/test/pi-messages.test.ts @@ -165,6 +165,34 @@ describe("pi-messages", () => { }); }); + it("forwards parsed wire events in order before converting them", async () => { + const wireEvents = [ + { type: "start" }, + { type: "text_start", contentIndex: 0 }, + { type: "text_delta", contentIndex: 0, delta: "Hello", gatewayField: "upstream-value" }, + { type: "text_end", contentIndex: 0, content: "Hello" }, + { type: "done", reason: "stop", usage, responseId: "resp_1" }, + ]; + const { baseUrl } = await startServer({ events: wireEvents }); + const model = createModel(baseUrl); + const received: unknown[] = []; + const eventModels: Model[] = []; + const message = await streamSimple(model, normalizeContext(context), { + apiKey: "test-key", + onProviderStreamEvent: async (event, eventModel) => { + await Promise.resolve(); + received.push(event); + eventModels.push(eventModel); + }, + }).result(); + + expect(received).toEqual(wireEvents); + expect(eventModels).toEqual(wireEvents.map(() => model)); + expect(message.stopReason).toBe("stop"); + expect(message.responseId).toBe("resp_1"); + expect(message.content).toEqual([{ type: "text", text: "Hello", textSignature: undefined }]); + }); + it("appends debug=1 and reports response headers via onResponse", async () => { const { baseUrl, requests } = await startServer({ headers: { "x-pi-gateway-upstream-provider": "anthropic" }, diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 60333a329..8e72902d9 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added the `provider_stream_event` extension event for observing parsed provider events before normalization, with an opt-in `/debug-provider` example viewer ([#9784](https://github.com/earendil-works/pi/issues/9784)). + ### Fixed - Prevented managed git packages from automatically installing Pi peer dependencies and added warnings for extension packages that list host-provided modules in `dependencies` ([#9863](https://github.com/earendil-works/pi/issues/9863)). diff --git a/packages/coding-agent/docs/custom-provider.md b/packages/coding-agent/docs/custom-provider.md index ae17b48c0..472be8354 100644 --- a/packages/coding-agent/docs/custom-provider.md +++ b/packages/coding-agent/docs/custom-provider.md @@ -102,9 +102,10 @@ The stream must also honor request instrumentation supplied through `SimpleStrea - Call `options.onPayload` before sending the provider request and use any replacement payload it returns. - Call `options.onResponse` after receiving the response but before consuming its body. +- Await `options.onProviderStreamEvent?.(providerEvent, model)` for each parsed provider event before normalizing it. - Pass through the abort signal and provider-scoped environment. -These hooks power extension request inspection and response-header events. Omitting them makes the provider behave differently from Pi’s built-in providers. +These hooks power extension request inspection, response-header events, and provider-stream observation. Omitting them makes the provider behave differently from Pi’s built-in providers. ## Report failures and usage diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 9a3d06454..a0a833e05 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -102,6 +102,12 @@ Events cover resource discovery, sessions, agent and message lifecycle, provider `message_end` can replace a finalized message while preserving its role. `tool_call` can mutate input or block execution. `tool_result` handlers compose, with each handler seeing prior changes. + + +`provider_stream_event` fires for each parsed provider stream event before Pi normalizes it. The event identifies the provider, API, and model; `event.data` is the earliest structured value available to Pi, not necessarily the original HTTP bytes or SSE frame. Treat it as read-only because mutation can affect normalization. The event is notification-only and is not persisted. + +Handlers are awaited in stream order, so slow handlers delay stream consumption. Handler errors are reported without changing the provider response. See [`debug-provider.ts`](../examples/extensions/debug-provider.ts) for an opt-in viewer that groups raw events by assistant message. + `context` transforms conversation messages without prompt and tool system messages; Pi restores that state afterward. Use `context_with_system` only when a request-local transformation must own the complete transcript, and keep a system message at index zero. diff --git a/packages/coding-agent/examples/extensions/README.md b/packages/coding-agent/examples/extensions/README.md index 30904dba5..90e7ead29 100644 --- a/packages/coding-agent/examples/extensions/README.md +++ b/packages/coding-agent/examples/extensions/README.md @@ -113,6 +113,7 @@ cp permission-gate.ts ~/.pi/agent/extensions/ |-----------|-------------| | `message-renderer.ts` | Custom message rendering with colors and expandable details via `registerMessageRenderer` | | `entry-renderer.ts` | TUI-only session entry rendering via `appendEntry` and `registerEntryRenderer` | +| `debug-provider.ts` | Toggle raw provider stream capture with `/debug-provider` and inspect each assistant message in a TUI-only session entry | | `event-bus.ts` | Inter-extension communication via `pi.events` | ### Session Metadata diff --git a/packages/coding-agent/examples/extensions/debug-provider.ts b/packages/coding-agent/examples/extensions/debug-provider.ts new file mode 100644 index 000000000..14aae7027 --- /dev/null +++ b/packages/coding-agent/examples/extensions/debug-provider.ts @@ -0,0 +1,90 @@ +/** + * Raw provider event viewer. + * + * Usage: /debug-provider [on|off] + * With no argument, the command toggles capture. Captured events are persisted + * as expandable custom entries. + */ + +import { type ExtensionAPI, keyHint } from "@earendil-works/pi-coding-agent"; +import { Box, Text } from "@earendil-works/pi-tui"; + +const ENTRY_TYPE = "debug-provider-events"; +const STATUS_KEY = "debug-provider"; + +interface ProviderDebugEntry { + provider: string; + api: string; + model: string; + events: unknown[]; +} + +export default function (pi: ExtensionAPI) { + let enabled = false; + let activeEvents: unknown[] | undefined; + let completedEntry: ProviderDebugEntry | undefined; + + pi.registerEntryRenderer(ENTRY_TYPE, (entry, { expanded }, theme) => { + const data = entry.data; + if (!data) return new Text(theme.fg("warning", "[provider debug] Missing event data"), 0, 0); + + const box = new Box(1, 1, (text) => theme.bg("customMessageBg", text)); + const count = `${data.events.length} event${data.events.length === 1 ? "" : "s"}`; + const expandHint = expanded ? "" : ` (${keyHint("app.tools.expand", "to view events")})`; + box.addChild( + new Text( + `${theme.fg("accent", "[provider debug]")} ${data.provider}/${data.model} (${data.api}) · ${count}${expandHint}`, + 0, + 0, + ), + ); + if (expanded) { + box.addChild(new Text(JSON.stringify(data.events, null, 2), 0, 0)); + } + return box; + }); + + pi.registerCommand("debug-provider", { + description: "Toggle capture of raw provider stream events", + handler: async (args, ctx) => { + const requestedState = args.trim().toLowerCase(); + if (requestedState !== "" && requestedState !== "on" && requestedState !== "off") { + ctx.ui.notify("Usage: /debug-provider [on|off]", "warning"); + return; + } + + enabled = requestedState === "" ? !enabled : requestedState === "on"; + if (!enabled) { + activeEvents = undefined; + completedEntry = undefined; + } + ctx.ui.setStatus(STATUS_KEY, enabled ? "provider debug" : undefined); + ctx.ui.notify(`Provider event capture ${enabled ? "enabled" : "disabled"}`, "info"); + }, + }); + + pi.on("turn_start", () => { + activeEvents = enabled ? [] : undefined; + }); + + pi.on("provider_stream_event", (event) => { + activeEvents?.push(structuredClone(event.data)); + }); + + pi.on("message_end", (event) => { + if (event.message.role !== "assistant" || !activeEvents) return; + completedEntry = { + provider: event.message.provider, + api: event.message.api, + model: event.message.model, + events: activeEvents, + }; + activeEvents = undefined; + }); + + pi.on("turn_end", () => { + if (!completedEntry) return; + pi.appendEntry(ENTRY_TYPE, completedEntry); + completedEntry = undefined; + }); +} diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index ad0287629..f9c989b35 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -130,6 +130,7 @@ export type { // Provider Registration ProviderConfig, ProviderModelConfig, + ProviderStreamEvent, ReadToolCallEvent, ReadToolResultEvent, // Commands diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 48b12ccc9..1580b5873 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -27,6 +27,7 @@ import type { OAuthLoginCallbacks, Provider, ProviderHeaders, + ProviderId, RefreshModelsContext, SimpleStreamOptions, TextContent, @@ -733,6 +734,15 @@ export interface AfterProviderResponseEvent { headers: Record; } +/** Fired for a parsed provider stream event before Pi normalizes it. */ +export interface ProviderStreamEvent { + type: "provider_stream_event"; + provider: ProviderId; + api: Api; + model: string; + data: unknown; +} + /** Fired after user submits prompt but before agent loop. */ export interface BeforeAgentStartEvent { type: "before_agent_start"; @@ -1177,6 +1187,7 @@ export type ExtensionEvent = | BeforeProviderRequestEvent | BeforeProviderHeadersEvent | AfterProviderResponseEvent + | ProviderStreamEvent | BeforeAgentStartEvent | AgentStartEvent | AgentEndEvent @@ -1390,6 +1401,7 @@ export interface ExtensionAPI { ): () => void; on(event: "before_provider_headers", handler: ExtensionHandler): () => void; on(event: "after_provider_response", handler: ExtensionHandler): () => void; + on(event: "provider_stream_event", handler: ExtensionHandler): () => void; on( event: "before_agent_start", handler: ExtensionHandler, @@ -1643,7 +1655,9 @@ export interface ProviderConfig { * (`getCurrentSystemPrompt(context.messages)`, `getCurrentTools(context.messages)`). * Implementations must invoke `options.onPayload` before sending the provider request and use any * returned replacement payload. They must invoke `options.onResponse` after receiving the response - * and before consuming its body, matching built-in providers. + * and before consuming its body, matching built-in providers. Implementations may invoke + * `options.onProviderStreamEvent(data, model)` with parsed stream events before normalization. + * Event data is adapter-owned and must be treated as read-only. */ streamSimple?: ( model: Model, diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 2ab852514..53e4afd2c 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -362,6 +362,20 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} headers: response.headers, }); }; + const handleProviderStreamEvent: NonNullable = async ( + data, + model, + ) => { + const runner = extensionRunnerRef.current; + if (!runner?.hasHandlers("provider_stream_event")) return; + await runner.emit({ + data, + type: "provider_stream_event", + provider: model.provider, + api: model.api, + model: model.id, + }); + }; const agent = new Agent({ initialState: { @@ -386,6 +400,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} }, onPayload: transformProviderPayload, onResponse: handleProviderResponse, + onProviderStreamEvent: handleProviderStreamEvent, sessionId: sessionManager.getSessionId(), transformContext: async (messages) => { const runner = extensionRunnerRef.current; diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 8746a894b..a8b3ff8cd 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -135,6 +135,7 @@ export type { ProjectTrustHandler, ProviderConfig, ProviderModelConfig, + ProviderStreamEvent, ReadToolCallEvent, RegisteredCommand, RegisteredTool, diff --git a/packages/coding-agent/test/sdk-stream-options.test.ts b/packages/coding-agent/test/sdk-stream-options.test.ts index 09c5c718c..143b0ad5c 100644 --- a/packages/coding-agent/test/sdk-stream-options.test.ts +++ b/packages/coding-agent/test/sdk-stream-options.test.ts @@ -1,4 +1,4 @@ -import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { mkdirSync, mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { @@ -11,6 +11,8 @@ import { } from "@earendil-works/pi-ai"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; +import type { ExtensionFactory } from "../src/core/extensions/types.ts"; +import { DefaultResourceLoader } from "../src/core/resource-loader.ts"; import { createAgentSession } from "../src/core/sdk.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { type Settings, SettingsManager } from "../src/core/settings-manager.ts"; @@ -81,15 +83,18 @@ describe("createAgentSession stream options", () => { api: Api, settings: Partial, requestOptions: SimpleStreamOptions = {}, - extensionSource?: string, + extensionFactory?: ExtensionFactory, + providerEvent?: unknown, ): Promise { const model = createModel(api); const settingsManager = SettingsManager.inMemory(settings); - if (extensionSource) { - const extensionsDir = join(agentDir, "extensions"); - mkdirSync(extensionsDir, { recursive: true }); - writeFileSync(join(extensionsDir, "headers.ts"), extensionSource); - } + const resourceLoader = new DefaultResourceLoader({ + cwd, + agentDir, + settingsManager, + extensionFactories: extensionFactory ? [extensionFactory] : [], + }); + await resourceLoader.reload(); const authStorage = AuthStorage.create(join(agentDir, "auth.json")); await authStorage.modify(model.provider, async () => ({ type: "api_key", key: "test-api-key" })); @@ -99,9 +104,16 @@ describe("createAgentSession stream options", () => { modelRegistry.registerProvider(model.provider, { api, headers: { "x-provider": "provider" }, - streamSimple: (_model, _context, providerOptions) => { + streamSimple: (requestModel, _context, providerOptions) => { capturedOptions = providerOptions; - return createDoneStream(api); + if (providerEvent === undefined) return createDoneStream(api); + + const stream = createAssistantMessageEventStream(); + void (async () => { + await providerOptions?.onProviderStreamEvent?.(providerEvent, requestModel); + stream.end(createDoneMessage(api)); + })(); + return stream; }, }); @@ -114,11 +126,20 @@ describe("createAgentSession stream options", () => { modelRuntime, settingsManager, sessionManager, + resourceLoader, }); try { - const stream = await session.agent.streamFunction(model, normalizeContext({ messages: [] }), requestOptions); - await stream.result(); + if (providerEvent === undefined) { + const stream = await session.agent.streamFunction( + model, + normalizeContext({ messages: [] }), + requestOptions, + ); + await stream.result(); + } else { + await session.prompt("test"); + } return capturedOptions; } finally { session.dispose(); @@ -247,12 +268,41 @@ describe("createAgentSession stream options", () => { expect(options?.maxRetryDelayMs).toBe(3000); }); + // Regression test for #9784. + it("forwards provider stream events to extensions", async () => { + const providerEvent = { openrouter_metadata: { strategy: "direct" } }; + const extensionEvents: unknown[] = []; + + const options = await captureStreamOptions( + "openai-completions", + {}, + {}, + (pi) => { + pi.on("provider_stream_event", (event) => { + extensionEvents.push(event); + }); + }, + providerEvent, + ); + + expect(options?.onProviderStreamEvent).toEqual(expect.any(Function)); + expect(extensionEvents).toEqual([ + { + data: providerEvent, + type: "provider_stream_event", + provider: "capture-provider", + api: "openai-completions", + model: "capture-model", + }, + ]); + }); + it("runs before_provider_headers on assembled headers without forwarding the transform", async () => { const options = await captureStreamOptions( "openai-completions", {}, { headers: { "x-explicit": "explicit" } }, - `export default function (pi) { + (pi) => { pi.on("before_provider_headers", (event) => { event.headers["x-hook"] = [ event.headers["x-provider"], @@ -260,7 +310,7 @@ describe("createAgentSession stream options", () => { event.headers["x-explicit"], ].join(":"); }); - }`, + }, ); expect(options?.headers).toMatchObject({