feat(coding-agent): add provider payload hook

This commit is contained in:
Mario Zechner
2026-03-07 14:23:25 +01:00
parent e3adaf1bd9
commit a3f05423d9
20 changed files with 157 additions and 32 deletions

View File

@@ -146,7 +146,7 @@ export const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOpt
const client = new BedrockRuntimeClient(config);
const cacheRetention = resolveCacheRetention(options.cacheRetention);
const commandInput = {
let commandInput = {
modelId: model.id,
messages: convertMessages(context, model, cacheRetention),
system: buildSystemPrompt(context.systemPrompt, model, cacheRetention),
@@ -154,7 +154,10 @@ export const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOpt
toolConfig: convertToolConfig(context.tools, options.toolChoice),
additionalModelRequestFields: buildAdditionalModelRequestFields(model, options),
};
options?.onPayload?.(commandInput);
const nextCommandInput = await options?.onPayload?.(commandInput, model);
if (nextCommandInput !== undefined) {
commandInput = nextCommandInput as typeof commandInput;
}
const command = new ConverseStreamCommand(commandInput);
const response = await client.send(command, { abortSignal: options.signal });

View File

@@ -235,8 +235,11 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti
options?.headers,
copilotDynamicHeaders,
);
const params = buildParams(model, context, isOAuthToken, options);
options?.onPayload?.(params);
let params = buildParams(model, context, isOAuthToken, options);
const nextParams = await options?.onPayload?.(params, model);
if (nextParams !== undefined) {
params = nextParams as MessageCreateParamsStreaming;
}
const anthropicStream = client.messages.stream({ ...params, stream: true }, { signal: options?.signal });
stream.push({ type: "start", partial: output });

View File

@@ -85,8 +85,11 @@ export const streamAzureOpenAIResponses: StreamFunction<"azure-openai-responses"
// Create Azure OpenAI client
const apiKey = options?.apiKey || getEnvApiKey(model.provider) || "";
const client = createClient(model, apiKey, options);
const params = buildParams(model, context, options, deploymentName);
options?.onPayload?.(params);
let params = buildParams(model, context, options, deploymentName);
const nextParams = await options?.onPayload?.(params, model);
if (nextParams !== undefined) {
params = nextParams as ResponseCreateParamsStreaming;
}
const openaiStream = await client.responses.create(
params,
options?.signal ? { signal: options.signal } : undefined,

View File

@@ -369,8 +369,11 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli", GoogleGe
const baseUrl = model.baseUrl?.trim();
const endpoints = baseUrl ? [baseUrl] : isAntigravity ? ANTIGRAVITY_ENDPOINT_FALLBACKS : [DEFAULT_ENDPOINT];
const requestBody = buildRequest(model, context, projectId, options, isAntigravity);
options?.onPayload?.(requestBody);
let requestBody = buildRequest(model, context, projectId, options, isAntigravity);
const nextRequestBody = await options?.onPayload?.(requestBody, model);
if (nextRequestBody !== undefined) {
requestBody = nextRequestBody as CloudCodeAssistRequest;
}
const headers = isAntigravity ? getAntigravityHeaders() : GEMINI_CLI_HEADERS;
const requestHeaders = {

View File

@@ -87,8 +87,11 @@ export const streamGoogleVertex: StreamFunction<"google-vertex", GoogleVertexOpt
const project = resolveProject(options);
const location = resolveLocation(options);
const client = createClient(model, project, location, options?.headers);
const params = buildParams(model, context, options);
options?.onPayload?.(params);
let params = buildParams(model, context, options);
const nextParams = await options?.onPayload?.(params, model);
if (nextParams !== undefined) {
params = nextParams as GenerateContentParameters;
}
const googleStream = await client.models.generateContentStream(params);
stream.push({ type: "start", partial: output });

View File

@@ -74,8 +74,11 @@ export const streamGoogle: StreamFunction<"google-generative-ai", GoogleOptions>
try {
const apiKey = options?.apiKey || getEnvApiKey(model.provider) || "";
const client = createClient(model, apiKey, options?.headers);
const params = buildParams(model, context, options);
options?.onPayload?.(params);
let params = buildParams(model, context, options);
const nextParams = await options?.onPayload?.(params, model);
if (nextParams !== undefined) {
params = nextParams as GenerateContentParameters;
}
const googleStream = await client.models.generateContentStream(params);
stream.push({ type: "start", partial: output });

View File

@@ -69,8 +69,11 @@ export const streamMistral: StreamFunction<"mistral-conversations", MistralOptio
const normalizeMistralToolCallId = createMistralToolCallIdNormalizer();
const transformedMessages = transformMessages(context.messages, model, (id) => normalizeMistralToolCallId(id));
const payload = buildChatPayload(model, context, transformedMessages, options);
options?.onPayload?.(payload);
let payload = buildChatPayload(model, context, transformedMessages, options);
const nextPayload = await options?.onPayload?.(payload, model);
if (nextPayload !== undefined) {
payload = nextPayload as ChatCompletionStreamRequest;
}
const mistralStream = await mistral.chat.stream(payload, buildRequestOptions(model, options));
stream.push({ type: "start", partial: output });
await consumeChatStream(model, output, stream, mistralStream);

View File

@@ -140,8 +140,11 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses"
}
const accountId = extractAccountId(apiKey);
const body = buildRequestBody(model, context, options);
options?.onPayload?.(body);
let body = buildRequestBody(model, context, options);
const nextBody = await options?.onPayload?.(body, model);
if (nextBody !== undefined) {
body = nextBody as RequestBody;
}
const headers = buildHeaders(model.headers, options?.headers, accountId, apiKey, options?.sessionId);
const bodyJson = JSON.stringify(body);
const transport = options?.transport || "sse";

View File

@@ -86,8 +86,11 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA
try {
const apiKey = options?.apiKey || getEnvApiKey(model.provider) || "";
const client = createClient(model, context, apiKey, options?.headers);
const params = buildParams(model, context, options);
options?.onPayload?.(params);
let params = buildParams(model, context, options);
const nextParams = await options?.onPayload?.(params, model);
if (nextParams !== undefined) {
params = nextParams as OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming;
}
const openaiStream = await client.chat.completions.create(params, { signal: options?.signal });
stream.push({ type: "start", partial: output });

View File

@@ -89,8 +89,11 @@ export const streamOpenAIResponses: StreamFunction<"openai-responses", OpenAIRes
// Create OpenAI client
const apiKey = options?.apiKey || getEnvApiKey(model.provider) || "";
const client = createClient(model, context, apiKey, options?.headers);
const params = buildParams(model, context, options);
options?.onPayload?.(params);
let params = buildParams(model, context, options);
const nextParams = await options?.onPayload?.(params, model);
if (nextParams !== undefined) {
params = nextParams as ResponseCreateParamsStreaming;
}
const openaiStream = await client.responses.create(
params,
options?.signal ? { signal: options.signal } : undefined,

View File

@@ -1,13 +1,7 @@
import { clearApiProviders, registerApiProvider } from "../api-registry.js";
import type {
AssistantMessage,
AssistantMessageEvent,
Context,
Model,
SimpleStreamOptions,
StreamOptions,
} from "../types.js";
import type { AssistantMessage, AssistantMessageEvent, Context, Model, SimpleStreamOptions } from "../types.js";
import { AssistantMessageEventStream } from "../utils/event-stream.js";
import type { BedrockOptions } from "./amazon-bedrock.js";
import { streamAnthropic, streamSimpleAnthropic } from "./anthropic.js";
import { streamAzureOpenAIResponses, streamSimpleAzureOpenAIResponses } from "./azure-openai-responses.js";
import { streamGoogle, streamSimpleGoogle } from "./google.js";
@@ -22,7 +16,7 @@ interface BedrockProviderModule {
streamBedrock: (
model: Model<"bedrock-converse-stream">,
context: Context,
options?: StreamOptions,
options?: BedrockOptions,
) => AsyncIterable<AssistantMessageEvent>;
streamSimpleBedrock: (
model: Model<"bedrock-converse-stream">,
@@ -83,7 +77,7 @@ function createLazyLoadErrorMessage(model: Model<"bedrock-converse-stream">, err
function streamBedrockLazy(
model: Model<"bedrock-converse-stream">,
context: Context,
options?: StreamOptions,
options?: BedrockOptions,
): AssistantMessageEventStream {
const outer = new AssistantMessageEventStream();

View File

@@ -79,9 +79,10 @@ export interface StreamOptions {
*/
sessionId?: string;
/**
* Optional callback for inspecting provider payloads before sending.
* Optional callback for inspecting or replacing provider payloads before sending.
* Return undefined to keep the payload unchanged.
*/
onPayload?: (payload: unknown) => void;
onPayload?: (payload: unknown, model: Model<Api>) => unknown | undefined | Promise<unknown | undefined>;
/**
* Optional custom HTTP headers to include in API requests.
* Merged with provider defaults; can override default headers.