diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index df6c55e5e..7691ed120 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -294,6 +294,7 @@ user sends prompt ──────────────────── │ │ │ │ │ ├─► turn_start │ │ │ ├─► context (can modify messages) │ │ + │ ├─► before_provider_headers (can mutate headers) | │ ├─► before_provider_request (can inspect or replace payload) │ ├─► after_provider_response (status + headers, before stream consume) │ │ │ │ @@ -643,6 +644,24 @@ pi.on("context", async (event, ctx) => { }); ``` +#### before_provider_headers + +Fired after the outgoing HTTP headers are assembled. Use it to add, override, or remove request headers. + +Handlers mutate `event.headers` in place. Set a key to a string to add or override it, or to `null` to delete it. + +```typescript +pi.on("before_provider_headers", (event, ctx) => { + // Add or override — e.g. a session id for gateway tracing/attribution + event.headers["x-session-id"] = ctx.sessionManager.getSessionId(); + + // Drop a tracking header pi adds for this call + event.headers["X-OpenRouter-Title"] = null; +}); +``` + +Runs once per provider request; retries reuse the same headers rather than re-firing the hook. + #### 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. diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index 56174d58e..cbe772f16 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -35,6 +35,7 @@ export type { BashToolResultEvent, BeforeAgentStartEvent, BeforeAgentStartEventResult, + BeforeProviderHeadersEvent, BeforeProviderRequestEvent, BeforeProviderRequestEventResult, BuildSystemPromptOptions, diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index a1817f7b6..c05d61bc7 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -3,7 +3,7 @@ */ import type { AgentMessage } from "@earendil-works/pi-agent-core"; -import type { ImageContent, Model } from "@earendil-works/pi-ai"; +import type { ImageContent, Model, ProviderHeaders } from "@earendil-works/pi-ai"; import type { KeyId } from "@earendil-works/pi-tui"; import { type Theme, theme } from "../../modes/interactive/theme/theme.ts"; import type { ResourceDiagnostic } from "../diagnostics.ts"; @@ -14,6 +14,7 @@ import type { BuildSystemPromptOptions } from "../system-prompt.ts"; import type { BeforeAgentStartEvent, BeforeAgentStartEventResult, + BeforeProviderHeadersEvent, BeforeProviderRequestEvent, CompactOptions, ContextEvent, @@ -145,6 +146,7 @@ type RunnerEmitEvent = Exclude< | UserBashEvent | ContextEvent | BeforeProviderRequestEvent + | BeforeProviderHeadersEvent | BeforeAgentStartEvent | MessageEndEvent | ResourcesDiscoverEvent @@ -1013,6 +1015,37 @@ export class ExtensionRunner { return currentPayload; } + async emitBeforeProviderHeaders(headers: ProviderHeaders): Promise { + const ctx = this.createContext(); + + for (const ext of this.extensions) { + const handlers = ext.handlers.get("before_provider_headers"); + if (!handlers || handlers.length === 0) continue; + + for (const handler of handlers) { + try { + // Handlers mutate `headers` in place; the return value is ignored. + const event: BeforeProviderHeadersEvent = { + type: "before_provider_headers", + headers, + }; + await handler(event, ctx); + } 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_headers", + error: message, + stack, + }); + } + } + } + + return headers; + } + 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 1e0f73127..91406c00a 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -24,6 +24,7 @@ import type { Model, OAuthCredentials, OAuthLoginCallbacks, + ProviderHeaders, SimpleStreamOptions, TextContent, ToolResultMessage, @@ -663,6 +664,16 @@ export interface BeforeProviderRequestEvent { payload: unknown; } +/** + * Fired after request headers are assembled, before the provider HTTP call. + * Handlers mutate `headers` in place (e.g. to inject tracing/session headers); + * the return value is ignored. A `null` value deletes that header. + */ +export interface BeforeProviderHeadersEvent { + type: "before_provider_headers"; + headers: ProviderHeaders; +} + /** Fired after a provider response is received and before the response stream is consumed. */ export interface AfterProviderResponseEvent { type: "after_provider_response"; @@ -1005,6 +1016,7 @@ export type ExtensionEvent = | SessionEvent | ContextEvent | BeforeProviderRequestEvent + | BeforeProviderHeadersEvent | AfterProviderResponseEvent | BeforeAgentStartEvent | AgentStartEvent @@ -1171,6 +1183,7 @@ export interface ExtensionAPI { event: "before_provider_request", handler: ExtensionHandler, ): void; + on(event: "before_provider_headers", handler: ExtensionHandler): void; on(event: "after_provider_response", handler: ExtensionHandler): void; on(event: "before_agent_start", handler: ExtensionHandler): void; on(event: "agent_start", handler: ExtensionHandler): void; diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 75b2ef732..4fc986708 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -313,6 +313,19 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const timeoutMs = options?.timeoutMs ?? providerRetrySettings.timeoutMs ?? effectiveTimeoutMs; const websocketConnectTimeoutMs = options?.websocketConnectTimeoutMs ?? settingsManager.getWebSocketConnectTimeoutMs(); + let headers = mergeProviderAttributionHeaders( + model, + settingsManager, + options?.sessionId, + auth.headers, + options?.headers, + ); + // Let extensions inject/adjust per-request headers (e.g. tracing, session correlation) + // after static assembly, before the provider HTTP call. + const headerRunner = extensionRunnerRef.current; + if (headerRunner?.hasHandlers("before_provider_headers")) { + headers = await headerRunner.emitBeforeProviderHeaders(headers ?? {}); + } return streamSimple(model, context, { ...options, apiKey: auth.apiKey, @@ -321,13 +334,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} websocketConnectTimeoutMs, maxRetries: options?.maxRetries ?? providerRetrySettings.maxRetries, maxRetryDelayMs: options?.maxRetryDelayMs ?? providerRetrySettings.maxRetryDelayMs, - headers: mergeProviderAttributionHeaders( - model, - settingsManager, - options?.sessionId, - auth.headers, - options?.headers, - ), + headers, }); }, onPayload: async (payload, _model) => { diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index df4c76123..29c821757 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -69,6 +69,7 @@ export type { BashToolCallEvent, BeforeAgentStartEvent, BeforeAgentStartEventResult, + BeforeProviderHeadersEvent, BeforeProviderRequestEvent, BeforeProviderRequestEventResult, BuildSystemPromptOptions, diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 61be09686..ace4bdc50 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -906,4 +906,58 @@ describe("ExtensionRunner", () => { expect(runner.hasHandlers("agent_end")).toBe(false); }); }); + + describe("before_provider_headers", () => { + it("lets a handler mutate headers in place and preserves existing headers", async () => { + const extCode = ` + export default function(pi) { + pi.on("before_provider_headers", (event) => { + event.headers["X-Turn-Index"] = "3"; + }); + } + `; + fs.writeFileSync(path.join(extensionsDir, "headers.ts"), extCode); + + const result = await discoverAndLoadExtensions([], tempDir, tempDir); + const runner = new ExtensionRunner(result.extensions, result.runtime, tempDir, sessionManager, modelRegistry); + + expect(runner.hasHandlers("before_provider_headers")).toBe(true); + + const headers = await runner.emitBeforeProviderHeaders({ "User-Agent": "kimchi/1.0" }); + expect(headers["X-Turn-Index"]).toBe("3"); + expect(headers["User-Agent"]).toBe("kimchi/1.0"); + }); + + it("isolates a throwing handler and still applies the others", async () => { + const throwing = ` + export default function(pi) { + pi.on("before_provider_headers", () => { + throw new Error("header handler boom"); + }); + } + `; + const good = ` + export default function(pi) { + pi.on("before_provider_headers", (event) => { + event.headers["X-Good"] = "yes"; + }); + } + `; + fs.writeFileSync(path.join(extensionsDir, "a-throwing.ts"), throwing); + fs.writeFileSync(path.join(extensionsDir, "b-good.ts"), good); + + const result = await discoverAndLoadExtensions([], tempDir, tempDir); + const runner = new ExtensionRunner(result.extensions, result.runtime, tempDir, sessionManager, modelRegistry); + const errors: Array<{ event: string; error: string }> = []; + runner.onError((err) => errors.push(err)); + + const headers = await runner.emitBeforeProviderHeaders({ "User-Agent": "x" }); + + expect(headers["X-Good"]).toBe("yes"); + expect(headers["User-Agent"]).toBe("x"); + expect(errors).toHaveLength(1); + expect(errors[0].event).toBe("before_provider_headers"); + expect(errors[0].error).toContain("header handler boom"); + }); + }); });