feat(coding-agent): add after_provider_response hook closes #3128

This commit is contained in:
Mario Zechner
2026-04-16 20:27:04 +02:00
Unverified
parent aa78fa91fe
commit d131fcd4ba
19 changed files with 120 additions and 25 deletions
+4
View File
@@ -98,6 +98,7 @@ export interface AgentOptions {
streamFn?: StreamFn;
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
onPayload?: SimpleStreamOptions["onPayload"];
onResponse?: SimpleStreamOptions["onResponse"];
beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
steeringMode?: QueueMode;
@@ -165,6 +166,7 @@ export class Agent {
public streamFn: StreamFn;
public getApiKey?: (provider: string) => Promise<string | undefined> | 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,
+2 -6
View File
@@ -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
@@ -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<string, string> = {};
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) {
+5 -1
View File
@@ -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 };
@@ -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);
+1
View File
@@ -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"),
@@ -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();
@@ -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<string, string> {
const out: Record<string, string> = {};
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;
@@ -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;
@@ -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, {
@@ -10,6 +10,7 @@ export function buildBaseOptions(model: Model<Api>, options?: SimpleStreamOption
sessionId: options?.sessionId,
headers: options?.headers,
onPayload: options?.onPayload,
onResponse: options?.onResponse,
maxRetryDelayMs: options?.maxRetryDelayMs,
metadata: options?.metadata,
};
+10
View File
@@ -57,6 +57,11 @@ export type CacheRetention = "none" | "short" | "long";
export type Transport = "sse" | "websocket" | "auto";
export interface ProviderResponse {
status: number;
headers: Record<string, string>;
}
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<Api>) => unknown | undefined | Promise<unknown | undefined>;
/**
* Optional callback invoked after an HTTP response is received and before
* its body stream is consumed.
*/
onResponse?: (response: ProviderResponse, model: Model<Api>) => void | Promise<void>;
/**
* Optional custom HTTP headers to include in API requests.
* Merged with provider defaults; can override default headers.
+7
View File
@@ -0,0 +1,7 @@
export function headersToRecord(headers: Headers): Record<string, string> {
const result: Record<string, string> = {};
for (const [key, value] of headers.entries()) {
result[key] = value;
}
return result;
}
+18 -1
View File
@@ -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` |
@@ -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");
});
}
@@ -20,6 +20,7 @@ export type {
} from "./runner.js";
export { ExtensionRunner } from "./runner.js";
export type {
AfterProviderResponseEvent,
AgentEndEvent,
AgentStartEvent,
// Re-exports
@@ -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<string, string>;
}
/** 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<BeforeProviderRequestEvent, BeforeProviderRequestEventResult>,
): void;
on(event: "after_provider_response", handler: ExtensionHandler<AfterProviderResponseEvent>): void;
on(event: "before_agent_start", handler: ExtensionHandler<BeforeAgentStartEvent, BeforeAgentStartEventResult>): void;
on(event: "agent_start", handler: ExtensionHandler<AgentStartEvent>): void;
on(event: "agent_end", handler: ExtensionHandler<AgentEndEvent>): void;
+11
View File
@@ -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;
@@ -142,6 +142,17 @@ export async function createHarness(options: HarnessOptions = {}): Promise<Harne
}
return runner.emitBeforeProviderRequest(payload);
},
onResponse: async (response) => {
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;