diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index b7d69546f..1017239b9 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -98,6 +98,7 @@ export interface AgentOptions { streamFn?: StreamFn; getApiKey?: (provider: string) => Promise | string | undefined; onPayload?: SimpleStreamOptions["onPayload"]; + onResponse?: SimpleStreamOptions["onResponse"]; beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise; afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise; steeringMode?: QueueMode; @@ -165,6 +166,7 @@ export class Agent { public streamFn: StreamFn; public getApiKey?: (provider: string) => Promise | string | undefined; public onPayload?: SimpleStreamOptions["onPayload"]; + public onResponse?: SimpleStreamOptions["onResponse"]; public beforeToolCall?: ( context: BeforeToolCallContext, signal?: AbortSignal, @@ -192,6 +194,7 @@ export class Agent { this.streamFn = options.streamFn ?? streamSimple; this.getApiKey = options.getApiKey; this.onPayload = options.onPayload; + this.onResponse = options.onResponse; this.beforeToolCall = options.beforeToolCall; this.afterToolCall = options.afterToolCall; this.steeringQueue = new PendingMessageQueue(options.steeringMode ?? "one-at-a-time"); @@ -411,6 +414,7 @@ export class Agent { reasoning: this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel, sessionId: this.sessionId, onPayload: this.onPayload, + onResponse: this.onResponse, transport: this.transport, thinkingBudgets: this.thinkingBudgets, maxRetryDelayMs: this.maxRetryDelayMs, diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index e7c7ad1b6..39cbfe117 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,13 +2,9 @@ ## [Unreleased] -## [0.67.5] - 2026-04-16 +### Added -### Fixed - -- Fixed Opus 4.7 adaptive thinking configuration across Anthropic and Bedrock providers by recognizing Opus 4.7 adaptive-thinking support and mapping `xhigh` reasoning to provider-supported effort values ([#3286](https://github.com/badlogic/pi-mono/pull/3286) by [@markusylisiurunen](https://github.com/markusylisiurunen)) - -## [0.67.4] - 2026-04-16 +- Added `onResponse` to `StreamOptions` so callers can inspect provider HTTP status and headers after each response arrives and before the response stream is consumed ([#3128](https://github.com/badlogic/pi-mono/issues/3128)) ### Changed diff --git a/packages/ai/src/providers/amazon-bedrock.ts b/packages/ai/src/providers/amazon-bedrock.ts index 00333f0ce..f72356b59 100644 --- a/packages/ai/src/providers/amazon-bedrock.ts +++ b/packages/ai/src/providers/amazon-bedrock.ts @@ -168,6 +168,13 @@ export const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOpt const command = new ConverseStreamCommand(commandInput); const response = await client.send(command, { abortSignal: options.signal }); + if (response.$metadata.httpStatusCode !== undefined) { + const responseHeaders: Record = {}; + if (response.$metadata.requestId) { + responseHeaders["x-amzn-requestid"] = response.$metadata.requestId; + } + await options?.onResponse?.({ status: response.$metadata.httpStatusCode, headers: responseHeaders }, model); + } for await (const item of response.stream!) { if (item.messageStart) { diff --git a/packages/ai/src/providers/anthropic.ts b/packages/ai/src/providers/anthropic.ts index 1741b4001..1d6ee9bc2 100644 --- a/packages/ai/src/providers/anthropic.ts +++ b/packages/ai/src/providers/anthropic.ts @@ -26,6 +26,7 @@ import type { ToolResultMessage, } from "../types.js"; import { AssistantMessageEventStream } from "../utils/event-stream.js"; +import { headersToRecord } from "../utils/headers.js"; import { parseStreamingJson } from "../utils/json-parse.js"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.js"; @@ -258,7 +259,10 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti if (nextParams !== undefined) { params = nextParams as MessageCreateParamsStreaming; } - const anthropicStream = client.messages.stream({ ...params, stream: true }, { signal: options?.signal }); + const { data: anthropicStream, response } = await client.messages + .stream({ ...params, stream: true }, { signal: options?.signal }) + .withResponse(); + await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model); stream.push({ type: "start", partial: output }); type Block = (ThinkingContent | TextContent | (ToolCall & { partialJson: string })) & { index: number }; diff --git a/packages/ai/src/providers/azure-openai-responses.ts b/packages/ai/src/providers/azure-openai-responses.ts index 9dc5261ea..b99819350 100644 --- a/packages/ai/src/providers/azure-openai-responses.ts +++ b/packages/ai/src/providers/azure-openai-responses.ts @@ -12,6 +12,7 @@ import type { StreamOptions, } from "../types.js"; import { AssistantMessageEventStream } from "../utils/event-stream.js"; +import { headersToRecord } from "../utils/headers.js"; import { convertResponsesMessages, convertResponsesTools, processResponsesStream } from "./openai-responses-shared.js"; import { buildBaseOptions, clampReasoning } from "./simple-options.js"; @@ -90,10 +91,10 @@ export const streamAzureOpenAIResponses: StreamFunction<"azure-openai-responses" if (nextParams !== undefined) { params = nextParams as ResponseCreateParamsStreaming; } - const openaiStream = await client.responses.create( - params, - options?.signal ? { signal: options.signal } : undefined, - ); + const { data: openaiStream, response } = await client.responses + .create(params, options?.signal ? { signal: options.signal } : undefined) + .withResponse(); + await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model); stream.push({ type: "start", partial: output }); await processResponsesStream(openaiStream, output, stream, model); diff --git a/packages/ai/src/providers/faux.ts b/packages/ai/src/providers/faux.ts index 328fd87f7..7bba1b72e 100644 --- a/packages/ai/src/providers/faux.ts +++ b/packages/ai/src/providers/faux.ts @@ -435,6 +435,7 @@ export function registerFauxProvider(options: RegisterFauxProviderOptions = {}): queueMicrotask(async () => { try { + await streamOptions?.onResponse?.({ status: 200, headers: {} }, requestModel); if (!step) { let message = createErrorMessage( new Error("No more faux responses queued"), diff --git a/packages/ai/src/providers/google-gemini-cli.ts b/packages/ai/src/providers/google-gemini-cli.ts index 259ca0699..fa6f98f12 100644 --- a/packages/ai/src/providers/google-gemini-cli.ts +++ b/packages/ai/src/providers/google-gemini-cli.ts @@ -21,6 +21,7 @@ import type { ToolCall, } from "../types.js"; import { AssistantMessageEventStream } from "../utils/event-stream.js"; +import { headersToRecord } from "../utils/headers.js"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.js"; import { convertMessages, @@ -408,6 +409,10 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli", GoogleGe body: requestBodyJson, signal: options?.signal, }); + await options?.onResponse?.( + { status: response.status, headers: headersToRecord(response.headers) }, + model, + ); if (response.ok) { break; // Success, exit retry loop @@ -757,6 +762,10 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli", GoogleGe body: requestBodyJson, signal: options?.signal, }); + await options?.onResponse?.( + { status: currentResponse.status, headers: headersToRecord(currentResponse.headers) }, + model, + ); if (!currentResponse.ok) { const retryErrorText = await currentResponse.text(); diff --git a/packages/ai/src/providers/openai-codex-responses.ts b/packages/ai/src/providers/openai-codex-responses.ts index 4cca39b7f..248443ec4 100644 --- a/packages/ai/src/providers/openai-codex-responses.ts +++ b/packages/ai/src/providers/openai-codex-responses.ts @@ -33,6 +33,7 @@ import type { Usage, } from "../types.js"; import { AssistantMessageEventStream } from "../utils/event-stream.js"; +import { headersToRecord } from "../utils/headers.js"; import { convertResponsesMessages, convertResponsesTools, processResponsesStream } from "./openai-responses-shared.js"; import { buildBaseOptions, clampReasoning } from "./simple-options.js"; @@ -214,6 +215,10 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses" body: bodyJson, signal: options?.signal, }); + await options?.onResponse?.( + { status: response.status, headers: headersToRecord(response.headers) }, + model, + ); if (response.ok) { break; @@ -529,14 +534,6 @@ function getWebSocketConstructor(): WebSocketConstructor | null { return ctor as unknown as WebSocketConstructor; } -function headersToRecord(headers: Headers): Record { - const out: Record = {}; - for (const [key, value] of headers.entries()) { - out[key] = value; - } - return out; -} - function getWebSocketReadyState(socket: WebSocketLike): number | undefined { const readyState = (socket as { readyState?: unknown }).readyState; return typeof readyState === "number" ? readyState : undefined; diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index 6656540f3..e2ab39b18 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -27,6 +27,7 @@ import type { ToolResultMessage, } from "../types.js"; import { AssistantMessageEventStream } from "../utils/event-stream.js"; +import { headersToRecord } from "../utils/headers.js"; import { parseStreamingJson } from "../utils/json-parse.js"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.js"; import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.js"; @@ -91,7 +92,10 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA if (nextParams !== undefined) { params = nextParams as OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming; } - const openaiStream = await client.chat.completions.create(params, { signal: options?.signal }); + const { data: openaiStream, response } = await client.chat.completions + .create(params, { signal: options?.signal }) + .withResponse(); + await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model); stream.push({ type: "start", partial: output }); let currentBlock: TextContent | ThinkingContent | (ToolCall & { partialArgs?: string }) | null = null; diff --git a/packages/ai/src/providers/openai-responses.ts b/packages/ai/src/providers/openai-responses.ts index 35bfeffe6..cc99bbb0e 100644 --- a/packages/ai/src/providers/openai-responses.ts +++ b/packages/ai/src/providers/openai-responses.ts @@ -14,6 +14,7 @@ import type { Usage, } from "../types.js"; import { AssistantMessageEventStream } from "../utils/event-stream.js"; +import { headersToRecord } from "../utils/headers.js"; import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.js"; import { convertResponsesMessages, convertResponsesTools, processResponsesStream } from "./openai-responses-shared.js"; import { buildBaseOptions, clampReasoning } from "./simple-options.js"; @@ -96,10 +97,10 @@ export const streamOpenAIResponses: StreamFunction<"openai-responses", OpenAIRes if (nextParams !== undefined) { params = nextParams as ResponseCreateParamsStreaming; } - const openaiStream = await client.responses.create( - params, - options?.signal ? { signal: options.signal } : undefined, - ); + const { data: openaiStream, response } = await client.responses + .create(params, options?.signal ? { signal: options.signal } : undefined) + .withResponse(); + await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model); stream.push({ type: "start", partial: output }); await processResponsesStream(openaiStream, output, stream, model, { diff --git a/packages/ai/src/providers/simple-options.ts b/packages/ai/src/providers/simple-options.ts index 71c15847b..6e478e64c 100644 --- a/packages/ai/src/providers/simple-options.ts +++ b/packages/ai/src/providers/simple-options.ts @@ -10,6 +10,7 @@ export function buildBaseOptions(model: Model, options?: SimpleStreamOption sessionId: options?.sessionId, headers: options?.headers, onPayload: options?.onPayload, + onResponse: options?.onResponse, maxRetryDelayMs: options?.maxRetryDelayMs, metadata: options?.metadata, }; diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 3c2a0d1ab..79281c6ce 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -57,6 +57,11 @@ export type CacheRetention = "none" | "short" | "long"; export type Transport = "sse" | "websocket" | "auto"; +export interface ProviderResponse { + status: number; + headers: Record; +} + export interface StreamOptions { temperature?: number; maxTokens?: number; @@ -83,6 +88,11 @@ export interface StreamOptions { * Return undefined to keep the payload unchanged. */ onPayload?: (payload: unknown, model: Model) => unknown | undefined | Promise; + /** + * Optional callback invoked after an HTTP response is received and before + * its body stream is consumed. + */ + onResponse?: (response: ProviderResponse, model: Model) => void | Promise; /** * Optional custom HTTP headers to include in API requests. * Merged with provider defaults; can override default headers. diff --git a/packages/ai/src/utils/headers.ts b/packages/ai/src/utils/headers.ts new file mode 100644 index 000000000..fae8a0368 --- /dev/null +++ b/packages/ai/src/utils/headers.ts @@ -0,0 +1,7 @@ +export function headersToRecord(headers: Headers): Record { + const result: Record = {}; + for (const [key, value] of headers.entries()) { + result[key] = value; + } + return result; +} diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index af1a25134..7175bab48 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -246,6 +246,7 @@ user sends prompt ──────────────────── │ ├─► turn_start │ │ │ ├─► context (can modify messages) │ │ │ ├─► before_provider_request (can inspect or replace payload) + │ ├─► after_provider_response (status + headers, before stream consume) │ │ │ │ │ │ LLM responds, may call tools: │ │ │ │ ├─► tool_execution_start │ │ @@ -534,6 +535,22 @@ pi.on("before_provider_request", (event, ctx) => { This is mainly useful for debugging provider serialization and cache behavior. +#### after_provider_response + +Fired after an HTTP response is received and before its stream body is consumed. Handlers run in extension load order. + +```typescript +pi.on("after_provider_response", (event, ctx) => { + // event.status - HTTP status code + // event.headers - normalized response headers + if (event.status === 429) { + console.log("rate limited", event.headers["retry-after"]); + } +}); +``` + +Header availability depends on provider and transport. Providers that abstract HTTP responses may not expose headers. + ### Model Events #### model_select @@ -2231,7 +2248,7 @@ All examples in [examples/extensions/](../examples/extensions/). | `dirty-repo-guard.ts` | Warn on dirty git repo | `on("session_before_*")`, `exec` | | `input-transform.ts` | Transform user input | `on("input")` | | `model-status.ts` | React to model changes | `on("model_select")`, `setStatus` | -| `provider-payload.ts` | Inspect or patch provider payloads | `on("before_provider_request")` | +| `provider-payload.ts` | Inspect payloads and provider response headers | `on("before_provider_request")`, `on("after_provider_response")` | | `system-prompt-header.ts` | Display system prompt info | `on("agent_start")`, `getSystemPrompt` | | `claude-rules.ts` | Load rules from files | `on("session_start")`, `on("before_agent_start")` | | `file-trigger.ts` | File watcher triggers messages | `sendMessage` | diff --git a/packages/coding-agent/examples/extensions/provider-payload.ts b/packages/coding-agent/examples/extensions/provider-payload.ts index 55d8cf59f..738a44125 100644 --- a/packages/coding-agent/examples/extensions/provider-payload.ts +++ b/packages/coding-agent/examples/extensions/provider-payload.ts @@ -11,4 +11,8 @@ export default function (pi: ExtensionAPI) { // Optional: replace the payload instead of only logging it. // return { ...event.payload, temperature: 0 }; }); + + pi.on("after_provider_response", (event) => { + appendFileSync(logFile, `[${event.status}] ${JSON.stringify(event.headers)}\n\n`, "utf8"); + }); } diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index cca8c9142..686bd15eb 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -20,6 +20,7 @@ export type { } from "./runner.js"; export { ExtensionRunner } from "./runner.js"; export type { + AfterProviderResponseEvent, AgentEndEvent, AgentStartEvent, // Re-exports diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index f986db850..28af67033 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -543,6 +543,13 @@ export interface BeforeProviderRequestEvent { payload: unknown; } +/** Fired after a provider response is received and before the response stream is consumed. */ +export interface AfterProviderResponseEvent { + type: "after_provider_response"; + status: number; + headers: Record; +} + /** Fired after user submits prompt but before agent loop. */ export interface BeforeAgentStartEvent { type: "before_agent_start"; @@ -863,6 +870,7 @@ export type ExtensionEvent = | SessionEvent | ContextEvent | BeforeProviderRequestEvent + | AfterProviderResponseEvent | BeforeAgentStartEvent | AgentStartEvent | AgentEndEvent @@ -1010,6 +1018,7 @@ export interface ExtensionAPI { event: "before_provider_request", handler: ExtensionHandler, ): void; + on(event: "after_provider_response", handler: ExtensionHandler): void; on(event: "before_agent_start", handler: ExtensionHandler): void; on(event: "agent_start", handler: ExtensionHandler): void; on(event: "agent_end", handler: ExtensionHandler): void; diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 7e4512cb9..5831e1ff3 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -314,6 +314,17 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} } return runner.emitBeforeProviderRequest(payload); }, + onResponse: async (response, _model) => { + const runner = extensionRunnerRef.current; + if (!runner?.hasHandlers("after_provider_response")) { + return; + } + await runner.emit({ + type: "after_provider_response", + status: response.status, + headers: response.headers, + }); + }, sessionId: sessionManager.getSessionId(), transformContext: async (messages) => { const runner = extensionRunnerRef.current; diff --git a/packages/coding-agent/test/suite/harness.ts b/packages/coding-agent/test/suite/harness.ts index c1cdb81b5..865b8780e 100644 --- a/packages/coding-agent/test/suite/harness.ts +++ b/packages/coding-agent/test/suite/harness.ts @@ -142,6 +142,17 @@ export async function createHarness(options: HarnessOptions = {}): Promise { + const runner = extensionRunnerRef.current; + if (!runner?.hasHandlers("after_provider_response")) { + return; + } + await runner.emit({ + type: "after_provider_response", + status: response.status, + headers: response.headers, + }); + }, transformContext: async (messages: AgentMessage[]) => { const runner = extensionRunnerRef.current; if (!runner) return messages;