diff --git a/src/main/agent-runtime/agent-hooks.ts b/src/main/agent-runtime/agent-hooks.ts index 690f3d49..8a499b7d 100644 --- a/src/main/agent-runtime/agent-hooks.ts +++ b/src/main/agent-runtime/agent-hooks.ts @@ -12,10 +12,9 @@ * `buildAgentHooks` collapses the wiring into a single object whose * shape mirrors the subset of `AgentOptions` that we populate. The * factory spreads the result into `new Agent({...})`, so the agent - * contract is still the pi-defined one — this is a pure organizational - * refactor. Behavior (RPC shapes, failure modes, tool-name semantics) - * is preserved verbatim by routing through the same `WorkerRpcClient` - * the factory previously held. + * contract is still the pi-defined one. Tool policy stays routed + * through `WorkerRpcClient`; context compaction runs inside the worker + * so the full transcript never crosses the worker-main boundary. * * The signatures below are pinned to pi via `pi-hook-contract`: any * drift in `AgentOptions.beforeToolCall` / `afterToolCall` / @@ -34,21 +33,28 @@ import type { AfterToolCallRpcResult, BeforeToolCallRpcPayload, BeforeToolCallRpcResult, - TransformContextRpcPayload, - TransformContextRpcResult, } from "./hook-rpc"; import type { WorkerRpcClient } from "../worker/entry/rpc"; +import { + compactMessagesIfNeeded, + type CompactionTelemetry, +} from "./context-compactor"; + +export interface ContextCompactionTelemetry extends CompactionTelemetry { + runId: string; + emittedAt: string; +} /** * Inputs the aggregator needs to produce a hook bundle. `rpc` is the - * worker-side RPC client that ships hook invocations back to main - * where the actual policy (approval gate, artifact projection, context - * compaction) lives. `runId` is the routing key main uses to look up - * the per-run `ActiveRunContext` in `RuntimeBridge`. + * worker-side RPC client that ships policy hook invocations back to + * main. `runId` is the routing key main uses to look up the per-run + * `ActiveRunContext` in `RuntimeBridge`. */ export interface AgentHooksDeps { runId: string; rpc: WorkerRpcClient; + onContextCompaction?: (telemetry: ContextCompactionTelemetry) => void; } /** @@ -74,11 +80,11 @@ export interface AgentHookBundle { * must never silently bypass main's approval gate. * - `afterToolCall`: fail-open RPC. Projection is observational; if * main cannot be reached the tool result stands unchanged. - * - `transformContext`: fail-open RPC. The pi contract forbids throws - * here, so on RPC failure we return the original transcript. + * - `transformContext`: worker-local compaction. The transcript stays + * inside the worker; main receives telemetry only. */ export function buildAgentHooks(deps: AgentHooksDeps): AgentHookBundle { - const { runId, rpc } = deps; + const { runId, rpc, onContextCompaction } = deps; const beforeToolCall: AgentHookBundle["beforeToolCall"] = async ( ctx, @@ -129,12 +135,24 @@ export function buildAgentHooks(deps: AgentHooksDeps): AgentHookBundle { const transformContext: AgentHookBundle["transformContext"] = async ( messages, ): Promise => { - const payload: TransformContextRpcPayload = { messages }; - const response = await rpc.invoke(runId, "transformContext", payload); - if (!response.ok) return messages; - const result = response.result as TransformContextRpcResult; - if (result.changed && result.messages) return result.messages; - return messages; + try { + const outcome = compactMessagesIfNeeded(messages); + if (outcome.changed && outcome.messages) { + if (outcome.telemetry) { + onContextCompaction?.({ + ...outcome.telemetry, + runId, + emittedAt: new Date().toISOString(), + }); + } + return outcome.messages; + } + return messages; + } catch (error) { + throw new Error( + `transformContext compaction failed: ${error instanceof Error ? error.message : String(error)}`, + ); + } }; return { beforeToolCall, afterToolCall, transformContext }; diff --git a/src/main/agent-runtime/context-compactor.ts b/src/main/agent-runtime/context-compactor.ts index d754e6b4..2e9b3487 100644 --- a/src/main/agent-runtime/context-compactor.ts +++ b/src/main/agent-runtime/context-compactor.ts @@ -51,6 +51,15 @@ interface MessageLike { export interface CompactionOutcome { changed: boolean; messages: AgentMessage[] | null; + telemetry: CompactionTelemetry | null; +} + +export interface CompactionTelemetry { + estimatedTokens: number; + originalMessageCount: number; + compactedMessageCount: number; + compactedOlderMessageCount: number; + retainedRecentMessageCount: number; } /** @@ -69,16 +78,16 @@ export interface CompactionOutcome { export function compactMessagesIfNeeded(messages: AgentMessage[]): CompactionOutcome { const estimatedTokens = estimateTokens(messages); if (estimatedTokens < COMPACTION_TRIGGER_TOKENS) { - return { changed: false, messages: null }; + return { changed: false, messages: null, telemetry: null }; } if (messages.length <= RECENT_MESSAGE_TAIL + 1) { - return { changed: false, messages: null }; + return { changed: false, messages: null, telemetry: null }; } const desiredCutoff = messages.length - RECENT_MESSAGE_TAIL; const cutoff = findSafeCutoff(messages, desiredCutoff); if (cutoff === null) { - return { changed: false, messages: null }; + return { changed: false, messages: null, telemetry: null }; } const older = messages.slice(0, cutoff); @@ -91,7 +100,18 @@ export function compactMessagesIfNeeded(messages: AgentMessage[]): CompactionOut timestamp: Date.now(), } as AgentMessage; - return { changed: true, messages: [summary, ...recent] }; + const compactedMessages = [summary, ...recent]; + return { + changed: true, + messages: compactedMessages, + telemetry: { + estimatedTokens, + originalMessageCount: messages.length, + compactedMessageCount: compactedMessages.length, + compactedOlderMessageCount: older.length, + retainedRecentMessageCount: recent.length, + }, + }; } /** diff --git a/src/main/agent-runtime/hook-handlers.ts b/src/main/agent-runtime/hook-handlers.ts index 9809f1f6..e1c0c36b 100644 --- a/src/main/agent-runtime/hook-handlers.ts +++ b/src/main/agent-runtime/hook-handlers.ts @@ -1,15 +1,13 @@ /** * Main-side implementations of pi's tool-execution hooks. Invoked by - * `RuntimeBridge.handleRpc` when the worker proxies `beforeToolCall`, - * `afterToolCall`, or `transformContext`. + * `RuntimeBridge.handleRpc` when the worker proxies `beforeToolCall` + * or `afterToolCall`. * * Each handler is deliberately isolated to one concern: * - `beforeToolCall`: operational-approval gate (blocks destructive * tools only while an approval that gates EXECUTION is pending) * - `afterToolCall`: filesystem artifact projection (bash/write/edit * outputs become `work_products`) - * - `transformContext`: per-LLM-call auto-compaction via a pure - * reducer * * Policy lives here; `RuntimeBridge` stays a pure IPC router. */ @@ -20,14 +18,11 @@ import type { AfterToolCallRpcResult, BeforeToolCallRpcPayload, BeforeToolCallRpcResult, - TransformContextRpcPayload, - TransformContextRpcResult, } from "./hook-rpc"; import { projectToolExecutionToWorkProduct, type ArtifactProjectorRunContext, } from "./artifact-projector"; -import { compactMessagesIfNeeded } from "./context-compactor"; import { isCofounderMandatoryApproval, isOperationalApproval, @@ -194,18 +189,3 @@ export function handleAfterToolCall( } return NO_OP_AFTER; } - -/** - * Context-transform hook. Delegates to `compactMessagesIfNeeded`, - * which is a pure, provider-free deterministic compactor. Runs before - * every LLM call; short transcripts short-circuit on the token-budget - * check so the steady-state cost is negligible. - */ -export function handleTransformContext( - _deps: HookHandlerDeps, - _ctx: ActiveRunContext, - payload: TransformContextRpcPayload, -): TransformContextRpcResult { - const outcome = compactMessagesIfNeeded(payload.messages); - return { changed: outcome.changed, messages: outcome.messages }; -} diff --git a/src/main/agent-runtime/hook-rpc.ts b/src/main/agent-runtime/hook-rpc.ts index fd731adb..d5e09ccd 100644 --- a/src/main/agent-runtime/hook-rpc.ts +++ b/src/main/agent-runtime/hook-rpc.ts @@ -1,21 +1,19 @@ /** * Hook RPC contracts — the parent↔worker half of pi-agent-core's - * `beforeToolCall` / `afterToolCall` / `transformContext` hooks. + * `beforeToolCall` / `afterToolCall` hooks. * * The Agent instance lives in the utility process. Policy (approval, - * compaction, filesystem projection) lives in main. The hooks are - * therefore proxied: the worker invokes an RPC that main answers from - * the same `RuntimeBridge` that already owns agent-event and - * `persistMessage` handling. + * filesystem projection) lives in main. The hooks are therefore proxied: + * the worker invokes an RPC that main answers from the same + * `RuntimeBridge` that already owns agent-event and `persistMessage` + * handling. Context compaction runs in the worker process and reports + * telemetry through a worker message. * * Payloads intentionally exclude pi's `AgentContext` (systemPrompt, * messages, tools) because main policy never needs it. Stripping it * keeps each hook IPC well under the structured-clone budget even for * large transcripts. */ - -import type { AgentMessage } from "@mariozechner/pi-agent-core"; - // ─── beforeToolCall ─────────────────────────────────────────────────────── export interface BeforeToolCallRpcPayload { @@ -51,19 +49,3 @@ export interface AfterToolCallRpcResult { contentOverride: { type: "text"; text: string }[] | null; isErrorOverride: boolean | null; } - -// ─── transformContext ───────────────────────────────────────────────────── - -export interface TransformContextRpcPayload { - messages: AgentMessage[]; -} - -export interface TransformContextRpcResult { - /** - * Transformed transcript. When `changed` is false the worker keeps its - * original array instance instead of replacing it — saves structured - * clone on the common "no compaction yet" path. - */ - changed: boolean; - messages: AgentMessage[] | null; -} diff --git a/src/main/agent-runtime/runtime-bridge.ts b/src/main/agent-runtime/runtime-bridge.ts index b1ae41bb..164442a2 100644 --- a/src/main/agent-runtime/runtime-bridge.ts +++ b/src/main/agent-runtime/runtime-bridge.ts @@ -20,14 +20,12 @@ import { captureBaselineIfRequired, reconcileWorkspaceArtifacts } from "./artifa import { handleBeforeToolCall, handleAfterToolCall, - handleTransformContext, type ActiveRunContext, type ActiveRunLifecycleState, } from "./hook-handlers"; import type { AfterToolCallRpcPayload, BeforeToolCallRpcPayload, - TransformContextRpcPayload, } from "./hook-rpc"; import { routeInternalToolCall } from "./internal-tools/router"; import type { InvokeInternalToolRpcPayload } from "@main/worker/runtime/internal-tools/internal-tool-rpc"; @@ -85,9 +83,9 @@ type RpcMethod = WorkerRpcMethod; * - Forwards `agent-event` from the worker to the renderer (push channel). * - Persists each `message_end` message via `SqliteSessionStore`. * - Calls `CostWriter` for assistant messages with usage. - * - Proxies pi's tool-execution hooks (`beforeToolCall` / `afterToolCall` - * / `transformContext`) into policy handlers: approval-gate checks, - * filesystem artifact projection, and context compaction. + * - Proxies pi's tool-execution hooks (`beforeToolCall` / `afterToolCall`) + * into policy handlers: approval-gate checks and filesystem artifact + * projection. * * Policy itself lives in `hook-handlers.ts`; this class is a pure IPC * router plus an `ActiveRunContext` cache. @@ -563,8 +561,6 @@ export class RuntimeBridge { return this.handleBeforeToolCallRpc(request.runId, request.payload as BeforeToolCallRpcPayload); case "afterToolCall": return this.handleAfterToolCallRpc(request.runId, request.payload as AfterToolCallRpcPayload); - case "transformContext": - return this.handleTransformContextRpc(request.runId, request.payload as TransformContextRpcPayload); case "invokeInternalTool": return this.handleInvokeInternalToolRpc(request.runId, request.payload as InvokeInternalToolRpcPayload); default: { @@ -733,15 +729,4 @@ export class RuntimeBridge { return { ok: true, result }; } - private handleTransformContextRpc( - runId: string, - payload: TransformContextRpcPayload, - ): { ok: true; result: unknown } | { ok: false; error: string } { - const ctx = this.activeRuns.get(runId); - if (!ctx) { - return { ok: false, error: `transformContext: no active run ${runId}` }; - } - const result = handleTransformContext(this.deps, ctx, payload); - return { ok: true, result }; - } } diff --git a/src/main/worker/entry/message-handler.ts b/src/main/worker/entry/message-handler.ts index a8c4a2d9..bd225fea 100644 --- a/src/main/worker/entry/message-handler.ts +++ b/src/main/worker/entry/message-handler.ts @@ -3,6 +3,7 @@ import type { workerMessageSchema } from "./types"; import { handleRunLog } from "../handlers/run-log-handler"; import { handleRunStatus } from "../handlers/run-status-handler"; import { handleRunFinished } from "../handlers/run-finished-handler"; +import { handleContextCompaction } from "../handlers/context-compaction-handler"; import type { HandlerContext, HandlerDeps } from "../handlers/handler-context"; export async function handleWorkerMessage( @@ -67,6 +68,9 @@ export async function handleWorkerMessage( } return; } + case "context-compaction": + handleContextCompaction(event, deps, context); + return; case "rpc-request": { const bridge = deps.getRuntimeBridge?.(); const response = bridge diff --git a/src/main/worker/entry/types.ts b/src/main/worker/entry/types.ts index 8bda95d4..70237e63 100644 --- a/src/main/worker/entry/types.ts +++ b/src/main/worker/entry/types.ts @@ -24,7 +24,6 @@ export const workerRpcMethodSchema = z.enum([ "prepareRun", "beforeToolCall", "afterToolCall", - "transformContext", "invokeInternalTool", ]); @@ -65,6 +64,16 @@ export const workerMessageSchema = z.union([ emittedAt: z.string(), event: z.record(z.string(), z.unknown()).and(z.object({ type: z.string() })), }), + z.object({ + type: z.literal("context-compaction"), + runId: z.string().uuid(), + emittedAt: z.string(), + estimatedTokens: z.number(), + originalMessageCount: z.number().int(), + compactedMessageCount: z.number().int(), + compactedOlderMessageCount: z.number().int(), + retainedRecentMessageCount: z.number().int(), + }), z.object({ type: z.literal("rpc-request"), rpcId: z.string(), diff --git a/src/main/worker/handlers/context-compaction-handler.ts b/src/main/worker/handlers/context-compaction-handler.ts new file mode 100644 index 00000000..2074431a --- /dev/null +++ b/src/main/worker/handlers/context-compaction-handler.ts @@ -0,0 +1,30 @@ +import { persistRunEvent } from "../observability/run-event-recorder"; +import type { HandlerContext, HandlerDeps, WorkerMessageEvent } from "./handler-context"; + +export function handleContextCompaction( + event: Extract, + deps: HandlerDeps, + _context: HandlerContext, +): void { + const runState = deps.db.runs.getRunAuthState(event.runId); + if (!runState) { + deps.logger.warn(`[context-compaction] run ${event.runId} disappeared before telemetry could be recorded.`); + return; + } + persistRunEvent(deps.db, deps.logger, { + companyId: runState.companyId, + runId: event.runId, + agentId: runState.agentId, + eventType: "context_compaction", + level: "info", + message: `Compacted ${event.compactedOlderMessageCount} older messages into a summary.`, + payload: { + emittedAt: event.emittedAt, + estimatedTokens: event.estimatedTokens, + originalMessageCount: event.originalMessageCount, + compactedMessageCount: event.compactedMessageCount, + compactedOlderMessageCount: event.compactedOlderMessageCount, + retainedRecentMessageCount: event.retainedRecentMessageCount, + }, + }); +} diff --git a/src/main/worker/runtime/helpers.ts b/src/main/worker/runtime/helpers.ts index 73b02223..cfd641bf 100644 --- a/src/main/worker/runtime/helpers.ts +++ b/src/main/worker/runtime/helpers.ts @@ -77,6 +77,16 @@ export type ChildMessage = emittedAt: string; event: { type: string } & Record; } + | { + type: "context-compaction"; + runId: string; + emittedAt: string; + estimatedTokens: number; + originalMessageCount: number; + compactedMessageCount: number; + compactedOlderMessageCount: number; + retainedRecentMessageCount: number; + } | { type: "rpc-request"; rpcId: string; runId: string; method: string; payload: unknown } | { type: "run-agent-state"; runId: string; state: RunAgentLifecycleState } | { type: "steer-ack"; runId: string; correlationId: string; accepted: boolean }; diff --git a/src/main/worker/runtime/index.ts b/src/main/worker/runtime/index.ts index c6ce3a23..044c9c74 100644 --- a/src/main/worker/runtime/index.ts +++ b/src/main/worker/runtime/index.ts @@ -178,6 +178,18 @@ export class WorkerRunner { apiKey: prepared.apiKey, extraHeaders, rpc, + onContextCompaction: (telemetry) => { + post({ + type: "context-compaction", + runId: telemetry.runId, + emittedAt: telemetry.emittedAt, + estimatedTokens: telemetry.estimatedTokens, + originalMessageCount: telemetry.originalMessageCount, + compactedMessageCount: telemetry.compactedMessageCount, + compactedOlderMessageCount: telemetry.compactedOlderMessageCount, + retainedRecentMessageCount: telemetry.retainedRecentMessageCount, + }); + }, }); this.activeAgents.set(job.runId, agent); post({ type: "run-agent-state", runId: job.runId, state: "agent_active" }); diff --git a/src/main/worker/runtime/internal-tools/internal-tool-rpc.ts b/src/main/worker/runtime/internal-tools/internal-tool-rpc.ts index e4186d76..53a099d1 100644 --- a/src/main/worker/runtime/internal-tools/internal-tool-rpc.ts +++ b/src/main/worker/runtime/internal-tools/internal-tool-rpc.ts @@ -9,7 +9,7 @@ * none of which the worker should replicate. * * This is structurally the same pattern as `beforeToolCall`/ - * `afterToolCall`/`transformContext` in `hook-rpc.ts`; what's different + * `afterToolCall` in `hook-rpc.ts`; what's different * is that those are side-effect hooks around pi's own tools, whereas * this is a primary-effect tool whose business logic lives entirely in * main. Keeping the payload shape minimal (name + JSON args) mirrors diff --git a/src/main/worker/runtime/pi-agent-factory.ts b/src/main/worker/runtime/pi-agent-factory.ts index 2a11dbe1..2a6fb521 100644 --- a/src/main/worker/runtime/pi-agent-factory.ts +++ b/src/main/worker/runtime/pi-agent-factory.ts @@ -9,8 +9,8 @@ * thrown hooks — a thrown `beforeToolCall` produces an error tool * result, a thrown `afterToolCall` likewise — so RPC failures degrade * to a deterministic tool-error instead of crashing the run. The - * `transformContext` contract explicitly forbids throwing, so on RPC - * failure we return the original messages unchanged. + * `transformContext` runs the deterministic compactor in the worker + * process and lets compaction failures fail the run explicitly. */ import { @@ -20,6 +20,7 @@ import { } from "@mariozechner/pi-agent-core"; import { streamSimple, type Api, type Model } from "@mariozechner/pi-ai"; import { buildAgentHooks } from "@main/agent-runtime/agent-hooks"; +import type { ContextCompactionTelemetry } from "@main/agent-runtime/agent-hooks"; import type { WorkerRpcClient } from "../entry/rpc"; // Side-effect import: drags the hook signature contract into the typecheck // graph so any drift in pi-agent-core hook shapes breaks the build here. @@ -40,6 +41,7 @@ export interface PiAgentFactoryInput { apiKey: string; extraHeaders: Record | undefined; rpc: WorkerRpcClient; + onContextCompaction?: (telemetry: ContextCompactionTelemetry) => void; } /** Map the UI-facing `thinkingEffort` string onto pi's `ThinkingLevel`. */ @@ -61,7 +63,11 @@ export function resolveThinkingLevel(effort: string | null | undefined): Thinkin } export function createPiAgent(input: PiAgentFactoryInput): Agent { - const hooks = buildAgentHooks({ runId: input.runId, rpc: input.rpc }); + const hooks = buildAgentHooks({ + runId: input.runId, + rpc: input.rpc, + onContextCompaction: input.onContextCompaction, + }); return new Agent({ initialState: { systemPrompt: input.systemPrompt, diff --git a/tests/unit/context-compaction-worker.test.ts b/tests/unit/context-compaction-worker.test.ts new file mode 100644 index 00000000..77f85a2b --- /dev/null +++ b/tests/unit/context-compaction-worker.test.ts @@ -0,0 +1,66 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; +import { describe, expect, it, vi } from "vitest"; +import { buildAgentHooks } from "../../src/main/agent-runtime/agent-hooks"; +import type { WorkerRpcClient } from "../../src/main/worker/entry/rpc"; + +function createRpc() { + return { + invoke: vi.fn(async () => ({ ok: true, result: undefined })), + } as unknown as WorkerRpcClient & { invoke: ReturnType }; +} + +function largeTranscript(): AgentMessage[] { + const messages: AgentMessage[] = []; + for (let i = 0; i < 46; i += 1) { + const role = i % 2 === 0 ? "user" : "assistant"; + messages.push({ + role, + content: [{ + type: "text", + text: `${role} turn ${i}\n${"context ".repeat(6_000)}`, + }], + timestamp: i, + } as AgentMessage); + } + return messages; +} + +describe("worker-side context compaction", () => { + it("compacts inside the worker without transformContext RPC", async () => { + const rpc = createRpc(); + const onContextCompaction = vi.fn(); + const hooks = buildAgentHooks({ + runId: "8b3f07b2-841a-4e16-9d1c-9d2a27c3c2ad", + rpc, + onContextCompaction, + }); + + const messages = largeTranscript(); + const transformed = await hooks.transformContext(messages); + + expect(rpc.invoke).not.toHaveBeenCalledWith( + expect.any(String), + "transformContext", + expect.anything(), + ); + expect(transformed.length).toBeLessThan(messages.length); + expect(transformed[0]).toMatchObject({ role: "user" }); + expect(onContextCompaction).toHaveBeenCalledWith(expect.objectContaining({ + runId: "8b3f07b2-841a-4e16-9d1c-9d2a27c3c2ad", + originalMessageCount: messages.length, + compactedMessageCount: transformed.length, + })); + }); + + it("propagates compaction failures instead of returning the original transcript", async () => { + const rpc = createRpc(); + const hooks = buildAgentHooks({ + runId: "b9f6aa93-0db7-4f39-8f28-2d87742ef414", + rpc, + }); + + await expect(hooks.transformContext(null as unknown as AgentMessage[])).rejects.toThrow( + "transformContext compaction failed", + ); + }); +});