@pontia/pi-client-plugin 0.5.2 → 0.6.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pontia/pi-client-plugin",
-  "version": "0.5.2",
+  "version": "0.6.0",
   "description": "Pi extension that connects coding-agent sessions to Pontia.",
   "keywords": [
     "pi-package",

package/src/compaction.ts

@@ -0,0 +1,225 @@
+import { randomUUID } from "node:crypto";
+import type {
+  SessionBeforeCompactEvent,
+  SessionCompactEvent,
+  SessionCompactFailedEvent,
+  CompactionResult,
+} from "@earendil-works/pi-coding-agent";
+import type { InternalEvent } from "./events.js";
+import type { SessionContext } from "./session.js";
+import { RpcError } from "./rpc-error.js";
+
+type Trigger = "manual" | "threshold" | "overflow" | "unknown";
+interface Operation {
+  id: string;
+  trigger: Trigger;
+  started?: InternalEvent;
+  terminal?: InternalEvent;
+  acknowledged: Set<string>;
+  reporting: Promise<void>;
+  awaitingNativeResult: boolean;
+  failureKey?: string;
+}
+interface State {
+  report: (event: InternalEvent) => Promise<boolean>;
+  active?: Operation;
+  operations: Map<string, Operation>;
+  signals: WeakMap<AbortSignal, Operation>;
+  failures: WeakSet<object>;
+  completed: Map<string, Operation>;
+}
+
+// Pi reloads extensions in the same process. Keep only control metadata across
+// reload, so a replaced hook or reconnect cannot allocate another operation.
+const stateKey = Symbol.for("pontia.pi.compaction");
+const processState = globalThis as typeof globalThis & { [stateKey]?: Map<string, State> };
+const states = (processState[stateKey] ??= new Map());
+
+export class PiCompaction {
+  private readonly state: State;
+
+  constructor(
+    home: string,
+    private readonly context: SessionContext,
+    report: (event: InternalEvent) => Promise<boolean>,
+  ) {
+    const key = JSON.stringify([
+      home,
+      context.sessionId,
+      context.runtimeId,
+      context.clientSessionKey,
+    ]);
+    let state = states.get(key);
+    if (!state) {
+      state = {
+        report,
+        operations: new Map(),
+        signals: new WeakMap(),
+        failures: new WeakSet(),
+        completed: new Map(),
+      };
+      states.set(key, state);
+    }
+    state.report = report;
+    this.state = state;
+  }
+
+  get running(): boolean {
+    return !!this.state.active && !this.state.active.terminal;
+  }
+
+  async reconnect(): Promise<void> {
+    for (const operation of this.state.operations.values()) {
+      if (operation.started) await this.publish(operation, !operation.awaitingNativeResult);
+    }
+  }
+
+  request(operationId: string, compact: () => void): void {
+    if (this.state.operations.has(operationId)) return;
+    if (this.running) throw new RpcError(-32006, "Pi compaction is already running");
+    const operation = this.create("manual", operationId);
+    operation.awaitingNativeResult = true;
+    try {
+      compact();
+    } catch (error) {
+      this.state.operations.delete(operationId);
+      this.state.active = undefined;
+      throw error;
+    }
+  }
+
+  async before(event: SessionBeforeCompactEvent): Promise<void> {
+    let operation = this.state.signals.get(event.signal);
+    if (!operation) {
+      operation = this.running ? this.state.active! : this.create(trigger(event.reason));
+      this.state.signals.set(event.signal, operation);
+    }
+    await this.publish(operation);
+  }
+
+  async completed(event: SessionCompactEvent): Promise<void> {
+    const entry = event.compactionEntry;
+    const operation = this.state.completed.get(entry.id) ?? this.current(event.reason);
+    this.state.completed.set(entry.id, operation);
+    operation.terminal ??= this.fact("session.compaction_completed", operation, {
+      native_compaction_id: entry.id ?? null,
+      tokens_before: tokens(entry.tokensBefore),
+      // Pi's extension event does not expose a post-compaction token count.
+      tokens_after: null,
+    });
+    // The manual context callback supplies Pi's estimatedTokensAfter result.
+    // Automatic compaction exposes only the extension event, so it stays null.
+    await this.publish(operation, !operation.awaitingNativeResult);
+  }
+
+  async failed(event: SessionCompactFailedEvent, leaf: string | null): Promise<void> {
+    if (this.state.failures.has(event)) return;
+    this.state.failures.add(event);
+    const failureKey = JSON.stringify([
+      event.reason,
+      event.aborted,
+      event.errorMessage ?? null,
+      leaf,
+    ]);
+    const previous = this.state.active;
+    const operation =
+      this.running || previous?.failureKey === failureKey
+        ? previous!
+        : this.create(trigger(event.reason));
+    operation.failureKey = failureKey;
+    operation.terminal ??= this.fact("session.compaction_failed", operation, {
+      aborted: event.aborted,
+      error_message: event.errorMessage ?? null,
+    });
+    operation.awaitingNativeResult = false;
+    await this.publish(operation);
+  }
+
+  // The callback covers native failures occurring before session_before_compact
+  // (for example, a session too small to compact). Hooks remain authoritative
+  // for cancellation and supply richer diagnostics when present.
+  async callbackError(operationId: string, error: Error): Promise<void> {
+    const operation = this.state.operations.get(operationId);
+    if (!operation) return;
+    operation.awaitingNativeResult = false;
+    operation.terminal ??= this.fact("session.compaction_failed", operation, {
+      aborted: false,
+      error_message: error.message,
+    });
+    await this.publish(operation);
+  }
+
+  async callbackComplete(operationId: string, result: CompactionResult): Promise<void> {
+    const operation = this.state.operations.get(operationId);
+    if (!operation) return;
+    operation.awaitingNativeResult = false;
+    operation.terminal ??= this.fact("session.compaction_completed", operation, {
+      native_compaction_id: null,
+      tokens_before: tokens(result.tokensBefore),
+      tokens_after: null,
+    });
+    if (
+      !operation.acknowledged.has(operation.terminal.type) &&
+      operation.terminal.type === "session.compaction_completed"
+    ) {
+      operation.terminal.data.tokens_after = tokens(result.estimatedTokensAfter);
+    }
+    await this.publish(operation);
+  }
+
+  private current(reason: unknown): Operation {
+    return this.running ? this.state.active! : this.create(trigger(reason));
+  }
+
+  private create(reason: Trigger, id = `compact_${randomUUID()}`): Operation {
+    const operation: Operation = {
+      id,
+      trigger: reason,
+      acknowledged: new Set(),
+      reporting: Promise.resolve(),
+      awaitingNativeResult: false,
+    };
+    this.state.operations.set(id, operation);
+    this.state.active = operation;
+    return operation;
+  }
+
+  private fact(
+    type: InternalEvent["type"],
+    operation: Operation,
+    metadata: Record<string, unknown> = {},
+  ): InternalEvent {
+    return {
+      session_id: this.context.sessionId,
+      type,
+      data: { operation_id: operation.id, trigger: operation.trigger, ...metadata },
+    };
+  }
+
+  private publish(operation: Operation, includeTerminal = true): Promise<void> {
+    operation.started ??= this.fact("session.compaction_started", operation);
+    operation.reporting = operation.reporting
+      .catch(() => {})
+      .then(async () => {
+        for (const event of [
+          operation.started!,
+          includeTerminal ? operation.terminal : undefined,
+        ]) {
+          if (!event || operation.acknowledged.has(event.type)) continue;
+          if (!(await this.state.report(event))) return;
+          operation.acknowledged.add(event.type);
+        }
+      });
+    return operation.reporting;
+  }
+}
+
+function trigger(reason: unknown): Trigger {
+  return reason === "manual" || reason === "threshold" || reason === "overflow"
+    ? reason
+    : "unknown";
+}
+
+function tokens(value: unknown): number | null {
+  return typeof value === "number" && Number.isSafeInteger(value) && value >= 0 ? value : null;
+}

package/src/control-socket.ts

@@ -4,7 +4,7 @@ import { readHistory, type HistoryContext } from "./history.js";
 import { createConnection, type Socket } from "node:net";
 import { isAbsolute, join } from "node:path";
 
-export const CONTROL_VERSION = 7;
+export const CONTROL_VERSION = 8;
 export const MAX_CONTROL_FRAME_BYTES = 64 * 1024;
 // Match the former HTTP event body limit, including the RPC envelope.
 export const MAX_RPC_FRAME_BYTES = 2 * 1024 * 1024 + 1024;
@@ -33,6 +33,9 @@ export interface LifecycleControl {
   interrupt(): void;
   shutdown(): void;
 }
+export interface CompactionControl {
+  compact(operationId: string): void;
+}
 export interface PiConnection {
   request(method: string, params: object): Promise<unknown>;
   registered(identity: ControlIdentity): void;
@@ -65,6 +68,7 @@ class RpcSocket {
   private models?: ModelControl;
   private lifecycle?: LifecycleControl;
   private history?: () => HistoryContext;
+  private compaction?: CompactionControl;
   private onReplay?: (inboxMessageId: string) => void;
 
   constructor(
@@ -74,6 +78,7 @@ class RpcSocket {
     onReplay?: (inboxMessageId: string) => void,
     lifecycle?: LifecycleControl,
     history?: () => HistoryContext,
+    compaction?: CompactionControl,
   ) {
     this.socket = socket;
     this.onSubmit = onSubmit;
@@ -81,6 +86,7 @@ class RpcSocket {
     this.onReplay = onReplay;
     this.lifecycle = lifecycle;
     this.history = history;
+    this.compaction = compaction;
     socket.on("data", (chunk: Buffer) => this.receive(chunk));
     socket.on("error", (error) => this.fail(error));
     socket.on("close", () =>
@@ -195,6 +201,26 @@ class RpcSocket {
         respond({ result: { pong: true } });
         return;
       }
+      if (message.method === "compact" && this.compaction) {
+        if (
+          Object.keys(params).some((key) => key !== "operation_id") ||
+          typeof params.operation_id !== "string" ||
+          !params.operation_id.trim()
+        ) {
+          error(-32602, "compact requires an operation_id");
+          return;
+        }
+        try {
+          this.compaction.compact(params.operation_id);
+          respond({ result: { accepted: true } });
+        } catch (failure) {
+          error(
+            failure instanceof RpcError ? failure.code : -32006,
+            failure instanceof Error ? failure.message : String(failure),
+          );
+        }
+        return;
+      }
       if (this.lifecycle && (message.method === "interrupt" || message.method === "shutdown")) {
         try {
           this.lifecycle[message.method as "interrupt" | "shutdown"]();
@@ -310,6 +336,7 @@ export async function connectPi(
   onReplay?: (inboxMessageId: string) => void,
   lifecycle?: LifecycleControl,
   history?: () => HistoryContext,
+  compaction?: CompactionControl,
 ): Promise<PiConnection> {
   const path = piSocketPath(pontiaHome);
   let identity: ControlIdentity | undefined;
@@ -381,6 +408,7 @@ export async function connectPi(
       reportingOnly ? undefined : onReplay,
       reportingOnly ? undefined : lifecycle,
       reportingOnly ? undefined : history,
+      reportingOnly ? undefined : compaction,
     );
     peers.add(peer);
     socket.once("close", () => {
@@ -434,16 +462,6 @@ export async function connectPi(
   current = await open();
   return {
     async request(method, params) {
-      if (method === "turn.startFailure") {
-        if (!identity?.clientSessionKey || stopped)
-          throw new Error("Pi reporting identity is unavailable");
-        await ready();
-        if (!current || stopped) throw new Error("Pi connection is unavailable");
-        return current.request(method, {
-          ...params,
-          client_session_key: identity.clientSessionKey,
-        });
-      }
       await ready();
       if (!current || stopped) return Promise.reject(new Error("Pi connection is unavailable"));
       return current.request(method, params);

package/src/events.ts

@@ -9,6 +9,9 @@ export type InternalEventType =
   | "session.exited"
   | "session.model_updated"
   | "session.context_usage_updated"
+  | "session.compaction_started"
+  | "session.compaction_completed"
+  | "session.compaction_failed"
   | "turn.started"
   | "turn.output"
   | "turn.completed"
@@ -58,6 +61,7 @@ export interface PiTopologyContext {
 export interface InternalEvent {
   session_id: string;
   turn_id?: string;
+  report_key?: string;
   type: InternalEventType;
   data: Record<string, unknown>;
   [legacyField: string]: any;
@@ -311,6 +315,7 @@ export function contextUsageFromPiHook(
 
 export function buildTurnStartedEvent(
   context: TurnContext,
+  reportKey: string,
   previousLeafId: string | null = null,
   topologyContext?: PiTopologyContext,
 ): InternalEvent {
@@ -324,7 +329,12 @@ export function buildTurnStartedEvent(
   };
   if (context.inboxMessageId) data.inbox_message_id = context.inboxMessageId;
   if (topologyContext) data.topology_context = topologyContext;
-  return { session_id: context.sessionId, type: "turn.started", data };
+  return {
+    session_id: context.sessionId,
+    report_key: reportKey,
+    type: "turn.started",
+    data,
+  };
 }
 
 export function buildTurnOutputEvent(context: ActiveTurnContext, output: string): InternalEvent {

package/src/index.ts

@@ -1,6 +1,7 @@
 import { AsyncLocalStorage } from "node:async_hooks";
 import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
 import type { HistoryContext } from "./history.js";
+import { PiCompaction } from "./compaction.js";
 import { connectPi, RpcError, type PiConnection, type ControlInput } from "./control-socket.js";
 import { defaultHookLogFile, type EnvLike, type TurnContext } from "./context.js";
 import { appendDiagnostic, type DiagnosticEntry } from "./diagnostics.js";
@@ -187,6 +188,13 @@ export function createPontiaPiExtension(
   let controlGeneration = 0;
   let piContext: ExtensionContext | undefined;
   let changingModel = false;
+  const compaction = () => {
+    const context = boundSessionContext;
+    if (!context || reportingDisabled || !readyReported) return undefined;
+    return new PiCompaction(currentPontiaHome(), context, async (event) =>
+      reportAccepted(await makeReporter(currentHookLogFile()).report(context, event)),
+    );
+  };
   const directInput = new AsyncLocalStorage<{
     submission: ControlInput;
     sessionContext: SessionContext;
@@ -317,6 +325,7 @@ export function createPontiaPiExtension(
           if (generation !== controlGeneration)
             throw new Error("Pi session changed during reconnect");
           await reportModel();
+          await compaction()?.reconnect();
         },
       },
       (inboxMessageId) => {
@@ -345,6 +354,26 @@ export function createPontiaPiExtension(
           throw new Error("Pi history native Session identity changed");
         return { generation, sessionManager: manager } satisfies HistoryContext;
       },
+      {
+        compact(operationId) {
+          const ctx = lifecycleContext();
+          if (ctx.sessionManager.getSessionId() !== boundSessionContext?.clientSessionKey)
+            throw new RpcError(-32006, "Pi compaction native Session identity changed");
+          const observation = compaction()!;
+          if (!ctx.isIdle() || ctx.hasPendingMessages() || queuedControlInput || changingModel)
+            throw new RpcError(-32006, "Pi is busy; compaction was not requested");
+          observation.request(operationId, () =>
+            ctx.compact({
+              onComplete: (result) => {
+                void observation.callbackComplete(operationId, result).catch(controlError);
+              },
+              onError: (error) => {
+                void observation.callbackError(operationId, error).catch(controlError);
+              },
+            }),
+          );
+        },
+      },
     );
     if (generation !== controlGeneration) {
       await socket.close();
@@ -519,6 +548,18 @@ export function createPontiaPiExtension(
     await reportModel().catch(controlError);
   });
 
+  pi.on("session_before_compact", async (event) => {
+    await compaction()?.before(event).catch(controlError);
+  });
+  pi.on("session_compact", async (event) => {
+    await compaction()?.completed(event).catch(controlError);
+  });
+  pi.on("session_compact_failed", async (event, ctx) => {
+    await compaction()
+      ?.failed(event, leafIdFromHookContext(ctx ?? piContext))
+      .catch(controlError);
+  });
+
   pi.on("session_start", async (event, ctx) => {
     if (reportingDisabled) return;
     const reason = (event as unknown as Record<string, unknown> | undefined)?.reason;
@@ -658,6 +699,7 @@ export function createPontiaPiExtension(
         await makeReporter(logFile).report(context, buildSessionReadyEvent(context)),
       );
       await reportModel().catch(controlError);
+      await compaction()?.reconnect().catch(controlError);
     } catch (error) {
       await closeControlSocket();
       const logFile = currentHookLogFile();
@@ -698,7 +740,7 @@ export function createPontiaPiExtension(
 
   pi.on("agent_start", async (_event, ctx) => {
     const topologyContext = topologyContextFromHookContext(ctx);
-    const previousLeafId = leafIdFromHookContext(ctx);
+    const reportKey = leafIdFromHookContext(ctx);
     if (reportingDisabled) return;
     try {
       const dispatch = directInput.getStore();
@@ -708,6 +750,16 @@ export function createPontiaPiExtension(
           : undefined;
       if (dispatch) dispatch.consumed = true;
       const logFile = currentHookLogFile();
+      if (!reportKey) {
+        activeTurn = undefined;
+        pendingPrompt = undefined;
+        await logDiagnostic(logFile, {
+          level: "error",
+          code: "turn_start_entry_id_missing",
+          message: "pi agent_start has no native session entry id; turn.started was not reported",
+        });
+        return;
+      }
       if (!boundSessionContext) {
         const hookSessionDetails = piSessionDetailsFromHookContext(ctx);
         const sessionDetails = hookSessionDetails.clientSessionKey
@@ -774,7 +826,7 @@ export function createPontiaPiExtension(
       lastContextUsageJson = undefined;
       const started = await reporter.report(
         turnContext,
-        buildTurnStartedEvent(turnContext, previousLeafId, topologyContext),
+        buildTurnStartedEvent(turnContext, reportKey, reportKey, topologyContext),
       );
       const canonicalTurnId = typeof started === "boolean" ? turnContext.turnId : started.turnId;
       if (!reportAccepted(started) || !canonicalTurnId) {

package/src/reporter.ts

@@ -23,72 +23,52 @@ export class EventReporter {
   }
 
   async report(context: { runtimeId: string }, event: InternalEvent): Promise<EventReportResult> {
-    try {
-      const body = await this.connection.request("event.report", {
-        runtime_id: context.runtimeId,
-        event,
-      });
-      const record =
-        body && typeof body === "object" ? (body as Record<string, unknown>) : undefined;
-      if (
-        event.type === "turn.started" &&
-        (record?.accepted !== true || typeof record?.turn_id !== "string" || !record.turn_id)
-      ) {
-        await this.reportStartFailure(event, "missing_turn_id");
-        return { accepted: false };
-      }
-      if (record?.accepted !== true)
-        throw new Error("Pi event RPC returned an invalid acknowledgement");
-      return {
-        accepted: true,
-        eventId: typeof record.event_id === "string" ? record.event_id : undefined,
-        turnId: typeof record.turn_id === "string" ? record.turn_id : undefined,
-      };
-    } catch (error) {
-      await appendDiagnostic(this.logFile, {
-        level: "error",
-        code: "pi_event_report_failed",
-        message: `Pi event RPC failed for ${event.type}`,
-        details: {
-          event_type: event.type,
-          error: error instanceof Error ? error.message : String(error),
-        },
-      });
-      await this.reportStartFailure(
-        event,
-        error instanceof RpcError ? "event_rejected" : "transport_failed",
-      );
-      return { accepted: false };
-    }
-  }
+    const attempts = event.type === "turn.started" ? 3 : 1;
+    let lastError: unknown;
 
-  private async reportStartFailure(
-    event: InternalEvent,
-    reason: "event_rejected" | "transport_failed" | "missing_turn_id",
-  ): Promise<void> {
-    if (event.type !== "turn.started" || typeof event.data.runtime_id !== "string") return;
-    // A lost turn.started response is ambiguous: retrying it could create a
-    // second Turn. Only this fenced, idempotent failure notification is retried.
-    for (let attempt = 0; attempt < 3; attempt++) {
+    for (let attempt = 0; attempt < attempts; attempt++) {
       if (attempt > 0) await new Promise((resolve) => setTimeout(resolve, 100 * attempt));
       try {
-        const result = (await this.connection.request("turn.startFailure", {
-          session_id: event.session_id,
-          runtime_id: event.data.runtime_id,
-          reason,
-        })) as { accepted?: unknown } | null;
-        if (result?.accepted === true) return;
+        const body = await this.connection.request("event.report", {
+          runtime_id: context.runtimeId,
+          event,
+        });
+        const record =
+          body && typeof body === "object" ? (body as Record<string, unknown>) : undefined;
+        if (record?.accepted !== true) {
+          if (record?.accepted === false) {
+            lastError = new Error("Pi event report was rejected");
+            break;
+          }
+          throw new Error("Pi event RPC returned an invalid acknowledgement");
+        }
+        if (
+          event.type === "turn.started" &&
+          (typeof record.turn_id !== "string" || !record.turn_id)
+        ) {
+          throw new Error("Pi turn.started acknowledgement has no canonical turn_id");
+        }
+        return {
+          accepted: true,
+          eventId: typeof record.event_id === "string" ? record.event_id : undefined,
+          turnId: typeof record.turn_id === "string" ? record.turn_id : undefined,
+        };
       } catch (error) {
+        lastError = error;
         if (error instanceof RpcError && error.code !== -32603) break;
-        // Exhaustion is visible locally; an unreachable server cannot persist
-        // a state transition. Never claim that this notification succeeded.
       }
     }
+
     await appendDiagnostic(this.logFile, {
       level: "error",
-      code: "turn_start_failure_report_failed",
-      message: "could not persist turn.started reporting failure; workflow state is unconfirmed",
-      details: { session_id: event.session_id, reason },
+      code: "pi_event_report_failed",
+      message: `Pi event RPC failed for ${event.type}`,
+      details: {
+        event_type: event.type,
+        report_key: event.report_key,
+        error: lastError instanceof Error ? lastError.message : String(lastError),
+      },
     });
+    return { accepted: false };
   }
 }