feat(coding-agent): add before_provider_headers extension hook (#6350)
* feat(coding-agent): add before_provider_headers extension hook Extensions can already rewrite the request payload through before_provider_request, but there is no way to adjust the outgoing HTTP headers of a provider call. This hook fills that gap for cases like request tracing, session correlation, or tenant routing. Handlers mutate the headers map in place - a null value deletes a header - and the return value is ignored, so a handler cannot accidentally drop auth or attribution headers by forgetting to spread. * docs(coding-agent): document before_provider_headers extension hook Add the before_provider_headers section and lifecycle-diagram entry to extensions.md; drop the now-superseded proposal file.
This commit is contained in:
@@ -294,6 +294,7 @@ user sends prompt ────────────────────
|
|||||||
│ │ │ │
|
│ │ │ │
|
||||||
│ ├─► turn_start │ │
|
│ ├─► turn_start │ │
|
||||||
│ ├─► context (can modify messages) │ │
|
│ ├─► context (can modify messages) │ │
|
||||||
|
│ ├─► before_provider_headers (can mutate headers) |
|
||||||
│ ├─► before_provider_request (can inspect or replace payload)
|
│ ├─► before_provider_request (can inspect or replace payload)
|
||||||
│ ├─► after_provider_response (status + headers, before stream consume)
|
│ ├─► 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
|
#### 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.
|
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.
|
||||||
|
|||||||
@@ -35,6 +35,7 @@ export type {
|
|||||||
BashToolResultEvent,
|
BashToolResultEvent,
|
||||||
BeforeAgentStartEvent,
|
BeforeAgentStartEvent,
|
||||||
BeforeAgentStartEventResult,
|
BeforeAgentStartEventResult,
|
||||||
|
BeforeProviderHeadersEvent,
|
||||||
BeforeProviderRequestEvent,
|
BeforeProviderRequestEvent,
|
||||||
BeforeProviderRequestEventResult,
|
BeforeProviderRequestEventResult,
|
||||||
BuildSystemPromptOptions,
|
BuildSystemPromptOptions,
|
||||||
|
|||||||
@@ -3,7 +3,7 @@
|
|||||||
*/
|
*/
|
||||||
|
|
||||||
import type { AgentMessage } from "@earendil-works/pi-agent-core";
|
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 { KeyId } from "@earendil-works/pi-tui";
|
||||||
import { type Theme, theme } from "../../modes/interactive/theme/theme.ts";
|
import { type Theme, theme } from "../../modes/interactive/theme/theme.ts";
|
||||||
import type { ResourceDiagnostic } from "../diagnostics.ts";
|
import type { ResourceDiagnostic } from "../diagnostics.ts";
|
||||||
@@ -14,6 +14,7 @@ import type { BuildSystemPromptOptions } from "../system-prompt.ts";
|
|||||||
import type {
|
import type {
|
||||||
BeforeAgentStartEvent,
|
BeforeAgentStartEvent,
|
||||||
BeforeAgentStartEventResult,
|
BeforeAgentStartEventResult,
|
||||||
|
BeforeProviderHeadersEvent,
|
||||||
BeforeProviderRequestEvent,
|
BeforeProviderRequestEvent,
|
||||||
CompactOptions,
|
CompactOptions,
|
||||||
ContextEvent,
|
ContextEvent,
|
||||||
@@ -145,6 +146,7 @@ type RunnerEmitEvent = Exclude<
|
|||||||
| UserBashEvent
|
| UserBashEvent
|
||||||
| ContextEvent
|
| ContextEvent
|
||||||
| BeforeProviderRequestEvent
|
| BeforeProviderRequestEvent
|
||||||
|
| BeforeProviderHeadersEvent
|
||||||
| BeforeAgentStartEvent
|
| BeforeAgentStartEvent
|
||||||
| MessageEndEvent
|
| MessageEndEvent
|
||||||
| ResourcesDiscoverEvent
|
| ResourcesDiscoverEvent
|
||||||
@@ -1013,6 +1015,37 @@ export class ExtensionRunner {
|
|||||||
return currentPayload;
|
return currentPayload;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async emitBeforeProviderHeaders(headers: ProviderHeaders): Promise<ProviderHeaders> {
|
||||||
|
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(
|
async emitBeforeAgentStart(
|
||||||
prompt: string,
|
prompt: string,
|
||||||
images: ImageContent[] | undefined,
|
images: ImageContent[] | undefined,
|
||||||
|
|||||||
@@ -24,6 +24,7 @@ import type {
|
|||||||
Model,
|
Model,
|
||||||
OAuthCredentials,
|
OAuthCredentials,
|
||||||
OAuthLoginCallbacks,
|
OAuthLoginCallbacks,
|
||||||
|
ProviderHeaders,
|
||||||
SimpleStreamOptions,
|
SimpleStreamOptions,
|
||||||
TextContent,
|
TextContent,
|
||||||
ToolResultMessage,
|
ToolResultMessage,
|
||||||
@@ -663,6 +664,16 @@ export interface BeforeProviderRequestEvent {
|
|||||||
payload: unknown;
|
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. */
|
/** Fired after a provider response is received and before the response stream is consumed. */
|
||||||
export interface AfterProviderResponseEvent {
|
export interface AfterProviderResponseEvent {
|
||||||
type: "after_provider_response";
|
type: "after_provider_response";
|
||||||
@@ -1005,6 +1016,7 @@ export type ExtensionEvent =
|
|||||||
| SessionEvent
|
| SessionEvent
|
||||||
| ContextEvent
|
| ContextEvent
|
||||||
| BeforeProviderRequestEvent
|
| BeforeProviderRequestEvent
|
||||||
|
| BeforeProviderHeadersEvent
|
||||||
| AfterProviderResponseEvent
|
| AfterProviderResponseEvent
|
||||||
| BeforeAgentStartEvent
|
| BeforeAgentStartEvent
|
||||||
| AgentStartEvent
|
| AgentStartEvent
|
||||||
@@ -1171,6 +1183,7 @@ export interface ExtensionAPI {
|
|||||||
event: "before_provider_request",
|
event: "before_provider_request",
|
||||||
handler: ExtensionHandler<BeforeProviderRequestEvent, BeforeProviderRequestEventResult>,
|
handler: ExtensionHandler<BeforeProviderRequestEvent, BeforeProviderRequestEventResult>,
|
||||||
): void;
|
): void;
|
||||||
|
on(event: "before_provider_headers", handler: ExtensionHandler<BeforeProviderHeadersEvent>): void;
|
||||||
on(event: "after_provider_response", handler: ExtensionHandler<AfterProviderResponseEvent>): void;
|
on(event: "after_provider_response", handler: ExtensionHandler<AfterProviderResponseEvent>): void;
|
||||||
on(event: "before_agent_start", handler: ExtensionHandler<BeforeAgentStartEvent, BeforeAgentStartEventResult>): void;
|
on(event: "before_agent_start", handler: ExtensionHandler<BeforeAgentStartEvent, BeforeAgentStartEventResult>): void;
|
||||||
on(event: "agent_start", handler: ExtensionHandler<AgentStartEvent>): void;
|
on(event: "agent_start", handler: ExtensionHandler<AgentStartEvent>): void;
|
||||||
|
|||||||
@@ -313,6 +313,19 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
|
|||||||
const timeoutMs = options?.timeoutMs ?? providerRetrySettings.timeoutMs ?? effectiveTimeoutMs;
|
const timeoutMs = options?.timeoutMs ?? providerRetrySettings.timeoutMs ?? effectiveTimeoutMs;
|
||||||
const websocketConnectTimeoutMs =
|
const websocketConnectTimeoutMs =
|
||||||
options?.websocketConnectTimeoutMs ?? settingsManager.getWebSocketConnectTimeoutMs();
|
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, {
|
return streamSimple(model, context, {
|
||||||
...options,
|
...options,
|
||||||
apiKey: auth.apiKey,
|
apiKey: auth.apiKey,
|
||||||
@@ -321,13 +334,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
|
|||||||
websocketConnectTimeoutMs,
|
websocketConnectTimeoutMs,
|
||||||
maxRetries: options?.maxRetries ?? providerRetrySettings.maxRetries,
|
maxRetries: options?.maxRetries ?? providerRetrySettings.maxRetries,
|
||||||
maxRetryDelayMs: options?.maxRetryDelayMs ?? providerRetrySettings.maxRetryDelayMs,
|
maxRetryDelayMs: options?.maxRetryDelayMs ?? providerRetrySettings.maxRetryDelayMs,
|
||||||
headers: mergeProviderAttributionHeaders(
|
headers,
|
||||||
model,
|
|
||||||
settingsManager,
|
|
||||||
options?.sessionId,
|
|
||||||
auth.headers,
|
|
||||||
options?.headers,
|
|
||||||
),
|
|
||||||
});
|
});
|
||||||
},
|
},
|
||||||
onPayload: async (payload, _model) => {
|
onPayload: async (payload, _model) => {
|
||||||
|
|||||||
@@ -69,6 +69,7 @@ export type {
|
|||||||
BashToolCallEvent,
|
BashToolCallEvent,
|
||||||
BeforeAgentStartEvent,
|
BeforeAgentStartEvent,
|
||||||
BeforeAgentStartEventResult,
|
BeforeAgentStartEventResult,
|
||||||
|
BeforeProviderHeadersEvent,
|
||||||
BeforeProviderRequestEvent,
|
BeforeProviderRequestEvent,
|
||||||
BeforeProviderRequestEventResult,
|
BeforeProviderRequestEventResult,
|
||||||
BuildSystemPromptOptions,
|
BuildSystemPromptOptions,
|
||||||
|
|||||||
@@ -906,4 +906,58 @@ describe("ExtensionRunner", () => {
|
|||||||
expect(runner.hasHandlers("agent_end")).toBe(false);
|
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");
|
||||||
|
});
|
||||||
|
});
|
||||||
});
|
});
|
||||||
|
|||||||
Reference in New Issue
Block a user