From a2515cf43f5345a5af51c80c610b788c6a77c1c2 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Sat, 27 Dec 2025 19:57:14 +0100 Subject: [PATCH] Wire context event to preprocessor for per-LLM-call execution - Change from contextTransform (runs once at agent start) to preprocessor - preprocessor runs before EACH LLM call inside the agent loop - ContextEvent now uses Message[] (pi-ai format) instead of AppMessage[] - Deep copy handled by pi-ai preprocessor, not Agent This enables: - Pruning rules applied on every turn (not just agent start) - /prune during long agent loop takes effect immediately - Compaction can use same transforms (future work) --- packages/agent/src/agent.ts | 22 +++++-------------- .../agent/src/transports/ProviderTransport.ts | 1 + packages/agent/src/transports/types.ts | 2 ++ .../coding-agent/src/core/hooks/runner.ts | 12 +++++----- packages/coding-agent/src/core/hooks/types.ts | 11 +++++----- packages/coding-agent/src/core/sdk.ts | 2 +- 6 files changed, 22 insertions(+), 28 deletions(-) diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index c857bfec..39b4c943 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -55,8 +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; + // Called before each LLM call inside the agent loop - can modify messages (e.g., for pruning) + preprocessor?: (messages: Message[]) => 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"; } @@ -77,7 +77,7 @@ export class Agent { private abortController?: AbortController; private transport: AgentTransport; private messageTransformer: (messages: AppMessage[]) => Message[] | Promise; - private contextTransform?: (messages: AppMessage[]) => Promise; + private preprocessor?: (messages: Message[]) => Promise; private messageQueue: Array> = []; private queueMode: "all" | "one-at-a-time"; private runningPrompt?: Promise; @@ -87,7 +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.preprocessor = opts.preprocessor; this.queueMode = opts.queueMode || "one-at-a-time"; } @@ -286,6 +286,7 @@ export class Agent { tools: this._state.tools, model, reasoning, + preprocessor: this.preprocessor, getQueuedMessages: async () => { if (this.queueMode === "one-at-a-time") { if (this.messageQueue.length > 0) { @@ -302,18 +303,7 @@ export class Agent { }, }; - // 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); + const llmMessages = await this.messageTransformer(this._state.messages); return { llmMessages, cfg, model }; } diff --git a/packages/agent/src/transports/ProviderTransport.ts b/packages/agent/src/transports/ProviderTransport.ts index 024db0e4..ee494553 100644 --- a/packages/agent/src/transports/ProviderTransport.ts +++ b/packages/agent/src/transports/ProviderTransport.ts @@ -60,6 +60,7 @@ export class ProviderTransport implements AgentTransport { // Resolve API key per assistant response (important for expiring OAuth tokens) getApiKey: this.options.getApiKey, getQueuedMessages: cfg.getQueuedMessages, + preprocessor: cfg.preprocessor, }; } diff --git a/packages/agent/src/transports/types.ts b/packages/agent/src/transports/types.ts index 736ba0c3..f74dcac5 100644 --- a/packages/agent/src/transports/types.ts +++ b/packages/agent/src/transports/types.ts @@ -9,6 +9,8 @@ export interface AgentRunConfig { model: Model; reasoning?: ReasoningEffort; getQueuedMessages?: () => Promise[]>; + /** Called before each LLM call - can modify messages (e.g., for pruning) */ + preprocessor?: (messages: Message[]) => Promise; } /** diff --git a/packages/coding-agent/src/core/hooks/runner.ts b/packages/coding-agent/src/core/hooks/runner.ts index 42623afc..df1bd087 100644 --- a/packages/coding-agent/src/core/hooks/runner.ts +++ b/packages/coding-agent/src/core/hooks/runner.ts @@ -2,7 +2,7 @@ * Hook runner - executes hooks and manages their lifecycle. */ -import type { AppMessage } from "@mariozechner/pi-agent-core"; +import type { Message } from "@mariozechner/pi-ai"; import type { ModelRegistry } from "../model-registry.js"; import type { SessionManager } from "../session-manager.js"; import type { AppendEntryHandler, LoadedHook, SendMessageHandler } from "./loader.js"; @@ -311,12 +311,13 @@ export class HookRunner { /** * 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. + * Returns the final modified messages, or the original if no modifications. + * + * Note: Messages are already deep-copied by the caller (pi-ai preprocessor). */ - async emitContext(messages: AppMessage[]): Promise { + async emitContext(messages: Message[]): Promise { const ctx = this.createContext(); let currentMessages = messages; - let modified = false; for (const hook of this.hooks) { const handlers = hook.handlers.get("context"); @@ -331,7 +332,6 @@ export class HookRunner { if (handlerResult && (handlerResult as ContextEventResult).messages) { currentMessages = (handlerResult as ContextEventResult).messages!; - modified = true; } } catch (err) { const message = err instanceof Error ? err.message : String(err); @@ -344,6 +344,6 @@ export class HookRunner { } } - return modified ? currentMessages : undefined; + return currentMessages; } } diff --git a/packages/coding-agent/src/core/hooks/types.ts b/packages/coding-agent/src/core/hooks/types.ts index a0089300..480696b7 100644 --- a/packages/coding-agent/src/core/hooks/types.ts +++ b/packages/coding-agent/src/core/hooks/types.ts @@ -6,7 +6,7 @@ */ import type { AppMessage } from "@mariozechner/pi-agent-core"; -import type { ImageContent, Model, TextContent, ToolResultMessage } from "@mariozechner/pi-ai"; +import type { ImageContent, Message, Model, TextContent, ToolResultMessage } from "@mariozechner/pi-ai"; import type { Component } from "@mariozechner/pi-tui"; import type { Theme } from "../../modes/interactive/theme/theme.js"; import type { CompactionPreparation, CompactionResult } from "../compaction.js"; @@ -148,13 +148,14 @@ export type SessionEvent = /** * Event data for context event. - * Fired before messages are sent to the LLM, allowing hooks to modify context non-destructively. + * Fired before each LLM call, allowing hooks to modify context non-destructively. * Original session messages are NOT modified - only the messages sent to the LLM are affected. + * Messages are already in LLM format (Message[], not AppMessage[]). */ export interface ContextEvent { type: "context"; - /** Messages about to be sent to the LLM */ - messages: AppMessage[]; + /** Messages about to be sent to the LLM (deep copy, safe to modify) */ + messages: Message[]; } /** @@ -330,7 +331,7 @@ export type HookEvent = */ export interface ContextEventResult { /** Modified messages to send instead of the original */ - messages?: AppMessage[]; + messages?: Message[]; } /** diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index d563a099..30ac53a1 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -589,7 +589,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} tools: allToolsArray, }, messageTransformer, - contextTransform: hookRunner + preprocessor: hookRunner ? async (messages) => { return hookRunner.emitContext(messages); }