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
parent aa78fa91fe
commit d131fcd4ba
19 changed files with 120 additions and 25 deletions

View File

@@ -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) {

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 };

View File

@@ -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);

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"),

View File

@@ -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();

View File

@@ -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;

View File

@@ -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;

View File

@@ -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, {

View File

@@ -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,
};

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.

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;
}