@pontia/pi-client-plugin 0.5.2 → 0.5.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pontia/pi-client-plugin",
-  "version": "0.5.2",
+  "version": "0.5.3",
   "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", () => {

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"

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