@ian-pascoe/pi-codemode 0.1.0 → 0.3.0

package/src/pi-codemode-extension.ts

@@ -39,7 +39,7 @@ import {
 } from "./pi-tool-bridge.js";
 
 const CODEMODE_EXECUTE_DESCRIPTION =
-  "Execute a TypeScript Cell in a persistent isolated Deno CodeMode Session. Top-level declarations become Notebook Bindings. Use the read-only tools object for registered Pi tools.";
+  "Execute a TypeScript Cell in a persistent isolated Deno CodeMode Session. Reuse a Session ID to retain Notebook Bindings; a new Session reclaims the least-recently-used idle Session at capacity. Use the read-only tools object for registered Pi tools. Cells may call console.log, console.info, console.warn, console.error, and console.debug; captured output arrives only with terminal results.";
 
 type PiCodeModeGeneration = {
   readonly captured: CapturedPiAgentSession;
@@ -126,11 +126,11 @@ class PiCodeModeLifecycleController {
     this.pi.on("session_start", async (_event, context) => this.startSession(context));
     this.pi.on("before_agent_start", () => this.synchronizeCurrentGeneration());
     this.pi.on("tool_execution_end", () => this.synchronizeCurrentGeneration());
-    this.pi.on("session_shutdown", async (event) => this.shutdownSession(event.reason));
+    this.pi.on("session_shutdown", async () => this.shutdownSession());
   }
 
   private async startSession(context: ExtensionContext): Promise<void> {
-    await this.shutdownSession("replacement");
+    await this.shutdownSession();
     const capturedResult = capturePiAgentSession(this.pi);
     if (!capturedResult.ok) {
       this.notifyWarning(context, capturedResult.warning);
@@ -207,6 +207,10 @@ class PiCodeModeLifecycleController {
       },
       result: async (input) => coordinator.result(input.sessionId),
       cancel: async (input) => coordinator.cancel(input.sessionId),
+      sessions: async () => ({
+        result: "success",
+        sessions: [...coordinator.listSessions()],
+      }),
     };
     generation = {
       captured,
@@ -243,7 +247,7 @@ class PiCodeModeLifecycleController {
       } catch {
         // Observer cleanup is presentation-only; execution resources still require release.
       }
-      await coordinator.shutdown("startup failure");
+      await coordinator.shutdown();
       await sessionFiles.close();
       if (this.generation === generation) this.generation = undefined;
       this.notifyWarning(
@@ -253,13 +257,16 @@ class PiCodeModeLifecycleController {
       return;
     }
 
-    for (const definition of createRenderedCodeModeToolDefinitions(
-      operations,
-      generation.executeDescription,
-      (sessionId) => coordinator.formatSessionPrefix(sessionId),
-    )) {
-      this.pi.registerTool(definition);
-    }
+    const [executeTool, resultTool, cancelTool, sessionsTool] =
+      createRenderedCodeModeToolDefinitions(
+        operations,
+        generation.executeDescription,
+        (sessionId) => coordinator.formatSessionPrefix(sessionId),
+      );
+    this.pi.registerTool(executeTool);
+    this.pi.registerTool(resultTool);
+    this.pi.registerTool(cancelTool);
+    this.pi.registerTool(sessionsTool);
     generation.toolsRegistered = true;
     this.synchronizeGeneration(generation);
   }
@@ -382,12 +389,8 @@ class PiCodeModeLifecycleController {
       now: CODEMODE_SYSTEM_RUNTIME.now,
       signal: AbortSignal.any([batch.signal, terminationController.signal]),
       onTerminate: () => terminationController.abort(),
-    };
-    if (outerAssistantMessage !== undefined) {
-      Object.assign(bridgeOptions, { outerAssistantMessage });
-    }
-    if (batch.onUpdate !== undefined) {
-      Object.assign(bridgeOptions, {
+      ...(outerAssistantMessage !== undefined && { outerAssistantMessage }),
+      ...(batch.onUpdate !== undefined && {
         onUpdate: (_callId: string, update: AgentToolResult<unknown>) => {
           const outerUpdate: AgentToolResult<CodeModeResultDetails> = {
             content: update.content,
@@ -395,8 +398,8 @@ class PiCodeModeLifecycleController {
           };
           batch.onUpdate?.(outerUpdate);
         },
-      });
-    }
+      }),
+    };
     const bridged = await executePiToolBridgeBatch(bridgeCaptured, bridgeOptions);
     const bridgedResults = new Map<string, CodeModeNestedToolResult>(
       bridged.calls.map((outcome) => [
@@ -424,16 +427,16 @@ class PiCodeModeLifecycleController {
           "Pi CodeMode nested tool returned no result",
         ),
     );
-    const batchResult = { results, presentation: bridged.presentation };
-    if (bridged.usage !== undefined) Object.assign(batchResult, { usage: bridged.usage });
-    if (bridged.addedToolNames.length > 0) {
-      Object.assign(batchResult, { addedToolNames: bridged.addedToolNames });
-    }
-    if (bridged.terminate) Object.assign(batchResult, { terminate: true });
-    return batchResult;
+    return {
+      results,
+      presentation: bridged.presentation,
+      ...(bridged.usage !== undefined && { usage: bridged.usage }),
+      ...(bridged.addedToolNames.length > 0 && { addedToolNames: bridged.addedToolNames }),
+      ...(bridged.terminate && { terminate: true }),
+    };
   }
 
-  private async shutdownSession(reason: string): Promise<void> {
+  private async shutdownSession(): Promise<void> {
     const generation = this.generation;
     if (generation === undefined || !generation.active) return;
     generation.active = false;
@@ -443,7 +446,7 @@ class PiCodeModeLifecycleController {
       // Observer cleanup is presentation-only; execution resources still require release.
     }
     try {
-      await generation.coordinator.shutdown(reason);
+      await generation.coordinator.shutdown();
     } finally {
       try {
         await generation.sessionFiles.close();

package/src/pi-tool-bridge.ts

@@ -12,6 +12,7 @@ import {
   parseCodeModeJsonValue,
   type CodeModeJsonValue,
 } from "./codemode-tool-contract.js";
+import { addCodeModeUsage, boundedCodeModeElapsedMs } from "./codemode-session-coordinator.js";
 import type { CapturedPiAgentSession } from "./pi-agent-session-capture.js";
 
 const PI_TOOL_BRIDGE_RESULT_LIMIT_BYTES = 8 * 1024 * 1024;
@@ -505,36 +506,6 @@ function translatePiToolResult(
   };
 }
 
-function addUsage(left: Usage | undefined, right: Usage | undefined): Usage | undefined {
-  if (right === undefined) return left;
-  if (left === undefined) return right;
-  const reasoning =
-    left.reasoning === undefined && right.reasoning === undefined
-      ? undefined
-      : (left.reasoning ?? 0) + (right.reasoning ?? 0);
-  const cacheWrite1h =
-    left.cacheWrite1h === undefined && right.cacheWrite1h === undefined
-      ? undefined
-      : (left.cacheWrite1h ?? 0) + (right.cacheWrite1h ?? 0);
-  const combined: Usage = {
-    input: left.input + right.input,
-    output: left.output + right.output,
-    cacheRead: left.cacheRead + right.cacheRead,
-    cacheWrite: left.cacheWrite + right.cacheWrite,
-    totalTokens: left.totalTokens + right.totalTokens,
-    cost: {
-      input: left.cost.input + right.cost.input,
-      output: left.cost.output + right.cost.output,
-      cacheRead: left.cost.cacheRead + right.cost.cacheRead,
-      cacheWrite: left.cost.cacheWrite + right.cost.cacheWrite,
-      total: left.cost.total + right.cost.total,
-    },
-  };
-  if (cacheWrite1h !== undefined) combined.cacheWrite1h = cacheWrite1h;
-  if (reasoning !== undefined) combined.reasoning = reasoning;
-  return combined;
-}
-
 function terminatedPiToolCall(callId: string): FinalizedPiToolCall {
   return createBridgeFailure(
     callId,
@@ -553,12 +524,6 @@ function presentationOutcome(
     : "failed";
 }
 
-function elapsedMilliseconds(startedAt: number, finishedAt: number): number {
-  const elapsed = Math.round(finishedAt - startedAt);
-  if (!Number.isFinite(elapsed)) return 0;
-  return Math.min(Number.MAX_SAFE_INTEGER, Math.max(0, elapsed));
-}
-
 function timedFinalizedPiToolCall(
   call: PiToolBridgeCall,
   finalized: FinalizedPiToolCall,
@@ -569,7 +534,7 @@ function timedFinalizedPiToolCall(
     finalized,
     presentation: {
       callId: call.callId,
-      elapsedMs: elapsedMilliseconds(startedAt, now()),
+      elapsedMs: boundedCodeModeElapsedMs(startedAt, now()),
       name: call.name,
       outcome: presentationOutcome(finalized.outcome),
     },
@@ -590,7 +555,9 @@ function collectPiToolBridgeBatch(
   const addedToolNames: string[] = [];
   const seenToolNames = new Set<string>();
   for (const { finalized } of timedCalls) {
-    usage = addUsage(usage, finalized.usage);
+    if (finalized.usage !== undefined) {
+      usage = addCodeModeUsage(usage, finalized.usage);
+    }
     for (const toolName of finalized.addedToolNames) {
       if (seenToolNames.has(toolName)) continue;
       seenToolNames.add(toolName);