diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index b61f8d14..42a7048c 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -8,6 +8,7 @@ import { type ImageContent, type Message, type Model, + type SimpleStreamOptions, streamSimple, type TextContent, type ThinkingBudgets, @@ -74,6 +75,11 @@ export interface AgentOptions { */ getApiKey?: (provider: string) => Promise | string | undefined; + /** + * Inspect or replace provider payloads before they are sent. + */ + onPayload?: SimpleStreamOptions["onPayload"]; + /** * Custom token budgets for thinking levels (token-based providers only). */ @@ -117,6 +123,7 @@ export class Agent { public streamFn: StreamFn; private _sessionId?: string; public getApiKey?: (provider: string) => Promise | string | undefined; + private _onPayload?: SimpleStreamOptions["onPayload"]; private runningPrompt?: Promise; private resolveRunningPrompt?: () => void; private _thinkingBudgets?: ThinkingBudgets; @@ -132,6 +139,7 @@ export class Agent { this.streamFn = opts.streamFn || streamSimple; this._sessionId = opts.sessionId; this.getApiKey = opts.getApiKey; + this._onPayload = opts.onPayload; this._thinkingBudgets = opts.thinkingBudgets; this._transport = opts.transport ?? "sse"; this._maxRetryDelayMs = opts.maxRetryDelayMs; @@ -429,6 +437,7 @@ export class Agent { model, reasoning, sessionId: this._sessionId, + onPayload: this._onPayload, transport: this._transport, thinkingBudgets: this._thinkingBudgets, maxRetryDelayMs: this._maxRetryDelayMs, diff --git a/packages/ai/src/providers/amazon-bedrock.ts b/packages/ai/src/providers/amazon-bedrock.ts index 4e4dfffb..52b42b4d 100644 --- a/packages/ai/src/providers/amazon-bedrock.ts +++ b/packages/ai/src/providers/amazon-bedrock.ts @@ -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 }); diff --git a/packages/ai/src/providers/anthropic.ts b/packages/ai/src/providers/anthropic.ts index ca2825df..ba59f147 100644 --- a/packages/ai/src/providers/anthropic.ts +++ b/packages/ai/src/providers/anthropic.ts @@ -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 }); diff --git a/packages/ai/src/providers/azure-openai-responses.ts b/packages/ai/src/providers/azure-openai-responses.ts index 1b82c4e4..4f474e1a 100644 --- a/packages/ai/src/providers/azure-openai-responses.ts +++ b/packages/ai/src/providers/azure-openai-responses.ts @@ -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, diff --git a/packages/ai/src/providers/google-gemini-cli.ts b/packages/ai/src/providers/google-gemini-cli.ts index cc4bf6c3..67eccdbc 100644 --- a/packages/ai/src/providers/google-gemini-cli.ts +++ b/packages/ai/src/providers/google-gemini-cli.ts @@ -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 = { diff --git a/packages/ai/src/providers/google-vertex.ts b/packages/ai/src/providers/google-vertex.ts index 97a3f869..7bdaf5eb 100644 --- a/packages/ai/src/providers/google-vertex.ts +++ b/packages/ai/src/providers/google-vertex.ts @@ -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 }); diff --git a/packages/ai/src/providers/google.ts b/packages/ai/src/providers/google.ts index 59c50245..991d5c90 100644 --- a/packages/ai/src/providers/google.ts +++ b/packages/ai/src/providers/google.ts @@ -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 }); diff --git a/packages/ai/src/providers/mistral.ts b/packages/ai/src/providers/mistral.ts index 29e141f1..95d3f839 100644 --- a/packages/ai/src/providers/mistral.ts +++ b/packages/ai/src/providers/mistral.ts @@ -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); diff --git a/packages/ai/src/providers/openai-codex-responses.ts b/packages/ai/src/providers/openai-codex-responses.ts index 90ec3ecd..8c9b8aae 100644 --- a/packages/ai/src/providers/openai-codex-responses.ts +++ b/packages/ai/src/providers/openai-codex-responses.ts @@ -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"; diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index 58ac9b36..2d28bf4e 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -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 }); diff --git a/packages/ai/src/providers/openai-responses.ts b/packages/ai/src/providers/openai-responses.ts index b83bfe87..a2e56404 100644 --- a/packages/ai/src/providers/openai-responses.ts +++ b/packages/ai/src/providers/openai-responses.ts @@ -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, diff --git a/packages/ai/src/providers/register-builtins.ts b/packages/ai/src/providers/register-builtins.ts index e9b7483a..e75f84de 100644 --- a/packages/ai/src/providers/register-builtins.ts +++ b/packages/ai/src/providers/register-builtins.ts @@ -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; 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(); diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 2447efd4..cc4a4309 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -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) => unknown | undefined | Promise; /** * Optional custom HTTP headers to include in API requests. * Merged with provider defaults; can override default headers. diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 66a555f2..66282de2 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -243,6 +243,7 @@ user sends prompt ──────────────────── │ │ │ │ │ ├─► turn_start │ │ │ ├─► context (can modify messages) │ │ + │ ├─► before_provider_request (can inspect or replace payload) │ │ │ │ │ │ LLM responds, may call tools: │ │ │ │ ├─► tool_call (can block) │ │ @@ -489,6 +490,21 @@ pi.on("context", async (event, ctx) => { }); ``` +#### before_provider_request + +Fired after the provider-specific payload is built, right before the request is sent. Handlers run in extension load order. Returning `undefined` keeps the payload unchanged. Returning any other value replaces the payload for later handlers and for the actual request. + +```typescript +pi.on("before_provider_request", (event, ctx) => { + console.log(JSON.stringify(event.payload, null, 2)); + + // Optional: replace payload + // return { ...event.payload, temperature: 0 }; +}); +``` + +This is mainly useful for debugging provider serialization and cache behavior. + ### Model Events #### model_select @@ -1934,6 +1950,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")` | | `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 new file mode 100644 index 00000000..55d8cf59 --- /dev/null +++ b/packages/coding-agent/examples/extensions/provider-payload.ts @@ -0,0 +1,14 @@ +import { appendFileSync } from "node:fs"; +import { join } from "node:path"; +import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; + +export default function (pi: ExtensionAPI) { + const logFile = join(process.cwd(), ".pi", "provider-payload.log"); + + pi.on("before_provider_request", (event) => { + appendFileSync(logFile, `${JSON.stringify(event.payload, null, 2)}\n\n`, "utf8"); + + // Optional: replace the payload instead of only logging it. + // return { ...event.payload, temperature: 0 }; + }); +} diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index 7dada9bc..e24b839e 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -32,6 +32,8 @@ export type { BashToolResultEvent, BeforeAgentStartEvent, BeforeAgentStartEventResult, + BeforeProviderRequestEvent, + BeforeProviderRequestEventResult, // Context CompactOptions, // Events - Agent diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index ede4c04b..05f2d786 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -13,6 +13,7 @@ import type { SessionManager } from "../session-manager.js"; import type { BeforeAgentStartEvent, BeforeAgentStartEventResult, + BeforeProviderRequestEvent, CompactOptions, ContextEvent, ContextEventResult, @@ -105,6 +106,7 @@ type RunnerEmitEvent = Exclude< | ToolResultEvent | UserBashEvent | ContextEvent + | BeforeProviderRequestEvent | BeforeAgentStartEvent | ResourcesDiscoverEvent | InputEvent @@ -710,6 +712,40 @@ export class ExtensionRunner { return currentMessages; } + async emitBeforeProviderRequest(payload: unknown): Promise { + const ctx = this.createContext(); + let currentPayload = payload; + + for (const ext of this.extensions) { + const handlers = ext.handlers.get("before_provider_request"); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + try { + const event: BeforeProviderRequestEvent = { + type: "before_provider_request", + payload: currentPayload, + }; + const handlerResult = await handler(event, ctx); + if (handlerResult !== undefined) { + currentPayload = handlerResult; + } + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + const stack = err instanceof Error ? err.stack : undefined; + this.emitError({ + extensionPath: ext.path, + event: "before_provider_request", + error: message, + stack, + }); + } + } + } + + return currentPayload; + } + async emitBeforeAgentStart( prompt: string, images: ImageContent[] | undefined, diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 5499dcff..b50f5332 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -493,6 +493,12 @@ export interface ContextEvent { messages: AgentMessage[]; } +/** Fired before a provider request is sent. Can replace the payload. */ +export interface BeforeProviderRequestEvent { + type: "before_provider_request"; + payload: unknown; +} + /** Fired after user submits prompt but before agent loop. */ export interface BeforeAgentStartEvent { type: "before_agent_start"; @@ -807,6 +813,7 @@ export type ExtensionEvent = | ResourcesDiscoverEvent | SessionEvent | ContextEvent + | BeforeProviderRequestEvent | BeforeAgentStartEvent | AgentStartEvent | AgentEndEvent @@ -832,6 +839,8 @@ export interface ContextEventResult { messages?: AgentMessage[]; } +export type BeforeProviderRequestEventResult = unknown; + export interface ToolCallEventResult { block?: boolean; reason?: string; @@ -944,6 +953,10 @@ export interface ExtensionAPI { on(event: "session_before_tree", handler: ExtensionHandler): void; on(event: "session_tree", handler: ExtensionHandler): void; on(event: "context", handler: ExtensionHandler): void; + on( + event: "before_provider_request", + 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 aa074c30..8c572370 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -292,6 +292,13 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} tools: [], }, convertToLlm: convertToLlmWithBlockImages, + onPayload: async (payload, _model) => { + const runner = extensionRunnerRef.current; + if (!runner?.hasHandlers("before_provider_request")) { + return payload; + } + return runner.emitBeforeProviderRequest(payload); + }, sessionId: sessionManager.getSessionId(), transformContext: async (messages) => { const runner = extensionRunnerRef.current; diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 7a15cc82..b83c17fe 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -56,6 +56,8 @@ export type { AppAction, BashToolCallEvent, BeforeAgentStartEvent, + BeforeProviderRequestEvent, + BeforeProviderRequestEventResult, CompactOptions, ContextEvent, ContextUsage,