fix(runtime): compact context inside worker

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<AgentMessage[]> => {
-    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 };

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,
+    },
+  };
 }
 
 /**

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 };
-}

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;
-}

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 };
-  }
 }

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

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(),

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<WorkerMessageEvent, { type: "context-compaction" }>,
+  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,
+    },
+  });
+}

src/main/worker/runtime/helpers.ts

@@ -77,6 +77,16 @@ export type ChildMessage =
       emittedAt: string;
       event: { type: string } & Record<string, unknown>;
     }
+  | {
+      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 };

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" });

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

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<string, string> | 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,

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<typeof vi.fn> };
+}
+
+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",
+    );
+  });
+});