mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-07-09 17:28:45 +00:00
feat(coding-agent): add before_provider_headers extension hook (#6350)
Some checks are pending
CI / build-check-test (push) Waiting to run
Some checks are pending
CI / build-check-test (push) Waiting to run
* 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:
parent
cfaa52e1c2
commit
244f1deaf1
7 changed files with 136 additions and 8 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ export type {
|
|||
BashToolResultEvent,
|
||||
BeforeAgentStartEvent,
|
||||
BeforeAgentStartEventResult,
|
||||
BeforeProviderHeadersEvent,
|
||||
BeforeProviderRequestEvent,
|
||||
BeforeProviderRequestEventResult,
|
||||
BuildSystemPromptOptions,
|
||||
|
|
|
|||
|
|
@ -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<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(
|
||||
prompt: string,
|
||||
images: ImageContent[] | undefined,
|
||||
|
|
|
|||
|
|
@ -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<BeforeProviderRequestEvent, BeforeProviderRequestEventResult>,
|
||||
): void;
|
||||
on(event: "before_provider_headers", handler: ExtensionHandler<BeforeProviderHeadersEvent>): void;
|
||||
on(event: "after_provider_response", handler: ExtensionHandler<AfterProviderResponseEvent>): void;
|
||||
on(event: "before_agent_start", handler: ExtensionHandler<BeforeAgentStartEvent, BeforeAgentStartEventResult>): 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 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) => {
|
||||
|
|
|
|||
|
|
@ -69,6 +69,7 @@ export type {
|
|||
BashToolCallEvent,
|
||||
BeforeAgentStartEvent,
|
||||
BeforeAgentStartEventResult,
|
||||
BeforeProviderHeadersEvent,
|
||||
BeforeProviderRequestEvent,
|
||||
BeforeProviderRequestEventResult,
|
||||
BuildSystemPromptOptions,
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue