mirror of
https://github.com/earendil-works/pi.git
synced 2026-06-18 15:54:04 +08:00
feat(coding-agent): add after_provider_response hook closes #3128
This commit is contained in:
@@ -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,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) {
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user