diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index ab4620d8..c857bfec 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -55,6 +55,8 @@ export interface AgentOptions { transport: AgentTransport; // Transform app messages to LLM-compatible messages before sending to transport messageTransformer?: (messages: AppMessage[]) => Message[] | Promise; + // Called before messageTransformer - can modify messages before they're sent to LLM (non-destructive) + contextTransform?: (messages: AppMessage[]) => Promise; // Queue mode: "all" = send all queued messages at once, "one-at-a-time" = send one queued message per turn queueMode?: "all" | "one-at-a-time"; } @@ -75,6 +77,7 @@ export class Agent { private abortController?: AbortController; private transport: AgentTransport; private messageTransformer: (messages: AppMessage[]) => Message[] | Promise; + private contextTransform?: (messages: AppMessage[]) => Promise; private messageQueue: Array> = []; private queueMode: "all" | "one-at-a-time"; private runningPrompt?: Promise; @@ -84,6 +87,7 @@ export class Agent { this._state = { ...this._state, ...opts.initialState }; this.transport = opts.transport; this.messageTransformer = opts.messageTransformer || defaultMessageTransformer; + this.contextTransform = opts.contextTransform; this.queueMode = opts.queueMode || "one-at-a-time"; } @@ -298,7 +302,18 @@ export class Agent { }, }; - const llmMessages = await this.messageTransformer(this._state.messages); + // Apply context transform (hooks can modify messages non-destructively) + // Deep copy so modifications don't affect the original state + let messagesToSend = this._state.messages; + if (this.contextTransform) { + const messagesCopy = JSON.parse(JSON.stringify(messagesToSend)) as AppMessage[]; + const transformed = await this.contextTransform(messagesCopy); + if (transformed) { + messagesToSend = transformed; + } + } + + const llmMessages = await this.messageTransformer(messagesToSend); return { llmMessages, cfg, model }; } diff --git a/packages/coding-agent/docs/session-tree-plan.md b/packages/coding-agent/docs/session-tree-plan.md index 65c79bed..319cf632 100644 --- a/packages/coding-agent/docs/session-tree-plan.md +++ b/packages/coding-agent/docs/session-tree-plan.md @@ -227,7 +227,7 @@ interface HookUIContext { See also: `CustomEntry` 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. diff --git a/packages/coding-agent/src/core/hooks/index.ts b/packages/coding-agent/src/core/hooks/index.ts index 5df5f5fa..4f98b704 100644 --- a/packages/coding-agent/src/core/hooks/index.ts +++ b/packages/coding-agent/src/core/hooks/index.ts @@ -13,6 +13,8 @@ export type { AgentEndEvent, AgentStartEvent, BashToolResultEvent, + ContextEvent, + ContextEventResult, CustomToolResultEvent, EditToolResultEvent, ExecOptions, diff --git a/packages/coding-agent/src/core/hooks/runner.ts b/packages/coding-agent/src/core/hooks/runner.ts index 6af0b77c..42623afc 100644 --- a/packages/coding-agent/src/core/hooks/runner.ts +++ b/packages/coding-agent/src/core/hooks/runner.ts @@ -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 { + 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; + } } diff --git a/packages/coding-agent/src/core/hooks/types.ts b/packages/coding-agent/src/core/hooks/types.ts index 07f7b3f7..a0089300 100644 --- a/packages/coding-agent/src/core/hooks/types.ts +++ b/packages/coding-agent/src/core/hooks/types.ts @@ -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): void; + // biome-ignore lint/suspicious/noConfusingVoidType: void allows handlers to not return anything + on(event: "context", handler: HookHandler): void; on(event: "agent_start", handler: HookHandler): void; on(event: "agent_end", handler: HookHandler): void; on(event: "turn_start", handler: HookHandler): void; diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 072cb009..d563a099 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -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 () => {