Add context event for non-destructive message modification before LLM calls

- Add contextTransform option to Agent (runs before messageTransformer)
- Deep copy messages before passing to contextTransform (modifications are ephemeral)
- Add ContextEvent and ContextEventResult types
- Add emitContext() to HookRunner (chains multiple handlers)
- Wire up in sdk.ts when creating Agent with hooks

Enables dynamic context pruning: hooks can modify messages sent to LLM
without changing session data. See discussion #330.
This commit is contained in:
Mario Zechner 2025-12-27 19:33:41 +01:00
parent 9e165d1d81
commit 77fe3f1a13
6 changed files with 89 additions and 2 deletions

View file

@ -227,7 +227,7 @@ interface HookUIContext {
See also: `CustomEntry<T>` for storing hook state that does NOT participate in context.
**New: `context` event (TODO)**
**New: `context` event **
Fires before messages are sent to the LLM, allowing hooks to modify context non-destructively.

View file

@ -13,6 +13,8 @@ export type {
AgentEndEvent,
AgentStartEvent,
BashToolResultEvent,
ContextEvent,
ContextEventResult,
CustomToolResultEvent,
EditToolResultEvent,
ExecOptions,

View file

@ -2,10 +2,13 @@
* Hook runner - executes hooks and manages their lifecycle.
*/
import type { AppMessage } from "@mariozechner/pi-agent-core";
import type { ModelRegistry } from "../model-registry.js";
import type { SessionManager } from "../session-manager.js";
import type { AppendEntryHandler, LoadedHook, SendMessageHandler } from "./loader.js";
import type {
ContextEvent,
ContextEventResult,
HookError,
HookEvent,
HookEventContext,
@ -304,4 +307,43 @@ export class HookRunner {
return result;
}
/**
* Emit a context event to all hooks.
* Handlers are chained - each gets the previous handler's output (if any).
* Returns the final modified messages, or undefined if no modifications.
*/
async emitContext(messages: AppMessage[]): Promise<AppMessage[] | undefined> {
const ctx = this.createContext();
let currentMessages = messages;
let modified = false;
for (const hook of this.hooks) {
const handlers = hook.handlers.get("context");
if (!handlers || handlers.length === 0) continue;
for (const handler of handlers) {
try {
const event: ContextEvent = { type: "context", messages: currentMessages };
const timeout = createTimeout(this.timeout);
const handlerResult = await Promise.race([handler(event, ctx), timeout.promise]);
timeout.clear();
if (handlerResult && (handlerResult as ContextEventResult).messages) {
currentMessages = (handlerResult as ContextEventResult).messages!;
modified = true;
}
} catch (err) {
const message = err instanceof Error ? err.message : String(err);
this.emitError({
hookPath: hook.path,
event: "context",
error: message,
});
}
}
}
return modified ? currentMessages : undefined;
}
}

View file

@ -146,6 +146,17 @@ export type SessionEvent =
fromHook: boolean;
});
/**
* Event data for context event.
* Fired before messages are sent to the LLM, allowing hooks to modify context non-destructively.
* Original session messages are NOT modified - only the messages sent to the LLM are affected.
*/
export interface ContextEvent {
type: "context";
/** Messages about to be sent to the LLM */
messages: AppMessage[];
}
/**
* Event data for agent_start event.
* Fired when an agent loop starts (once per user prompt).
@ -301,6 +312,7 @@ export function isLsToolResult(e: ToolResultEvent): e is LsToolResultEvent {
*/
export type HookEvent =
| SessionEvent
| ContextEvent
| AgentStartEvent
| AgentEndEvent
| TurnStartEvent
@ -312,6 +324,15 @@ export type HookEvent =
// Event Results
// ============================================================================
/**
* Return type for context event handlers.
* Allows hooks to modify messages before they're sent to the LLM.
*/
export interface ContextEventResult {
/** Modified messages to send instead of the original */
messages?: AppMessage[];
}
/**
* Return type for tool_call event handlers.
* Allows hooks to block tool execution.
@ -417,6 +438,8 @@ export interface RegisteredCommand {
export interface HookAPI {
// biome-ignore lint/suspicious/noConfusingVoidType: void allows handlers to not return anything
on(event: "session", handler: HookHandler<SessionEvent, SessionEventResult | void>): void;
// biome-ignore lint/suspicious/noConfusingVoidType: void allows handlers to not return anything
on(event: "context", handler: HookHandler<ContextEvent, ContextEventResult | void>): void;
on(event: "agent_start", handler: HookHandler<AgentStartEvent>): void;
on(event: "agent_end", handler: HookHandler<AgentEndEvent>): void;
on(event: "turn_start", handler: HookHandler<TurnStartEvent>): void;

View file

@ -589,6 +589,11 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
tools: allToolsArray,
},
messageTransformer,
contextTransform: hookRunner
? async (messages) => {
return hookRunner.emitContext(messages);
}
: undefined,
queueMode: settingsManager.getQueueMode(),
transport: new ProviderTransport({
getApiKey: async () => {