pi-vision-handoff 0.2.1 → 0.3.0

package/README.md

@@ -41,6 +41,7 @@ No `settings.json` touched. No per-provider glue. Pick a describer once and ever
 - **🗂️ Explicit overrides** — force handoff for specific models (e.g. a weak vision model) with `/vision-handoff add`.
 - **⚡ Cache-warmed** — `before_agent_start` describes attached images the moment you submit, so the request is rarely delayed.
 - **🛡️ Graceful degradation** — no API key? Describer unreachable? Aborted? The image is replaced with a clean `[Image: description unavailable]` placeholder instead of crashing your turn.
+- **📊 Usage reporting** — every real describer call reports model + tokens (and Neuralwatt energy/cost when the vision model is a Neuralwatt model), via `pi.appendEntry` + a `vision-handoff:usage` event for live consumers.
 - **🔧 Tunable** — cap description length (`maxDescriptionLines`; unbounded by default, so `read`'s native `ctrl+o` collapse handles compactness) and cache size, in the config file.
 
 ## Usage
@@ -176,6 +177,31 @@ terminal still renders the image inline via kitty.
 
 The describer call itself goes through pi's normal model machinery (`complete()`), **not** the agent event loop — so it never re-triggers `before_provider_request` (no recursion).
 
+### Usage reporting
+
+Every **real** describer call (cache misses only — cache hits emit nothing) reports its model + tokens so pi and other extensions can account for the handoff cost. When the vision model is a **Neuralwatt** model, the response's `: energy` / `: cost` / `: mcr-session` SSE comments are also captured (the OpenAI SDK discards comment lines, so the response body is teed and parsed — the same technique `pi-neuralwatt-provider` uses). For non-Neuralwatt vision models the energy fields are **omitted** (not zeroed), so consumers can distinguish "no energy" from "zero energy".
+
+Each record is published two ways, mirroring `pi-neuralwatt-provider`'s `neuralwatt:turn-energy` pattern:
+
+- **`pi.appendEntry("vision-handoff-usage", record)`** — persisted to the session log, so it replays on `/resume`, fork, and branch navigation.
+- **`pi.events.emit("vision-handoff:usage", record)`** — live event-bus channel a consumer (e.g. a `pi-tps`-style extension) can filter on to see tokens **and** energy in one payload.
+
+Record shape:
+
+```ts
+{
+  imageHash: string,            // sha256(mime + base64), first 32 hex chars
+  model: string, provider: string,
+  responseModel?: string, responseId?: string,
+  usage: Usage,                 // { input, output, cacheRead, cacheWrite, totalTokens, cost }
+  // Present ONLY when Neuralwatt SSE energy comments were captured:
+  energyJoules?: number, costUsd?: number,
+  energyRaw?: object, mcrSessionRaw?: object, costRaw?: object,
+}
+```
+
+Because `before_agent_start` fires several `describeImage()` calls fire-and-forget, the fetch interception is **refcounted** (installed only while ≥1 describe is in flight) and uses `AsyncLocalStorage` to route each teed response body to the describe call that issued it — so concurrent describes each attribute their own energy correctly without clobbering `globalThis.fetch`. When the vision model is a Neuralwatt model, `pi-neuralwatt-provider`'s own `streamNeuralwatt` tee nests on top and restores back to this interceptor; both tees read the same comment lines independently (the accepted duplication for easy filtering).
+
 ## Comparison with Alternatives
 
 | Approach | Pros | Cons |
@@ -190,7 +216,7 @@ The describer call itself goes through pi's normal model machinery (`complete()`
 
 ```bash
 pnpm install
-pnpm test          # Vitest unit tests (50 passing)
+pnpm test          # Vitest unit tests (77 passing)
 pnpm typecheck     # TypeScript validation
 pnpm lint:dead     # Dead code detection (knip)
 ```
@@ -201,10 +227,12 @@ pnpm lint:dead     # Dead code detection (knip)
 .
 ├── vision-handoff.ts            # Main extension: hooks, command, describer
 ├── src/
-│   ├── index.ts                  # Config schema, read/write, image-block helpers
+│   ├── index.ts                  # Config schema, read/write, image-block helpers (barrel)
+│   ├── usage.ts                  # Describer usage + Neuralwatt energy capture, fetch interceptor
 │   └── vision-model-selector.ts  # Interactive picker TUI component
 ├── __tests__/unit/
 │   ├── config-dir.test.ts        # Ensures getAgentDir() usage
+│   ├── usage.test.ts             # Energy parsing, usage records, concurrency-safe fetch routing
 │   └── vision-handoff.test.ts    # Config, refs, image-block extraction, insertion, truncation, round-trip
 ├── package.json
 ├── tsconfig.json

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-vision-handoff",
-  "version": "0.2.1",
+  "version": "0.3.0",
   "description": "Give text-only pi models vision — describe images with a vision model you pick via an interactive picker, then hand off the text description to non-vision models",
   "type": "module",
   "author": "Tom X Nguyen",

package/src/index.ts

@@ -10,6 +10,8 @@ import { getAgentDir } from "@earendil-works/pi-coding-agent";
 import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { join } from "node:path";
 
+export * from "./usage.js";
+
 /** Subdirectory under the pi agent dir where picker extensions store config. */
 const CONFIG_SUBDIR = "extensions";
 

package/src/usage.ts

@@ -0,0 +1,245 @@
+/**
+ * Describer usage + energy capture for pi-vision-handoff.
+ *
+ * One record is produced per REAL describer provider call (cache hits emit
+ * nothing):
+ *   - model + tokens: from complete()'s AssistantMessage.usage
+ *   - energy + cost + raw MCR/energy/cost payloads: from Neuralwatt SSE comment
+ *     lines parsed out of the teed response body (readEnergyFromTee). Present
+ *     ONLY when the vision model is a Neuralwatt model — non-Neuralwatt models
+ *     produce no comment lines, so the energy fields are OMITTED (not zeroed)
+ *     for easy downstream filtering ("no energy" vs "zero energy").
+ *
+ * The caller (vision-handoff.ts) persists the record via pi.appendEntry (replays
+ * on session resume/branch) AND emits it on pi.events so a live consumer can
+ * filter on the one channel for tokens AND energy.
+ *
+ * Split out of vision-handoff.ts so the pure pieces (readEnergyFromTee,
+ * buildUsageRecord) and the concurrency-safe fetch interceptor are unit-testable
+ * through the normal src/ import path — vision-handoff.ts runs readConfig() at
+ * module load and pulls in the TUI selector, so it is not unit-test-friendly.
+ */
+
+import { AsyncLocalStorage } from "node:async_hooks";
+import type { Api, AssistantMessage, Model, Usage } from "@earendil-works/pi-ai";
+
+/** Custom session-entry type persisted via pi.appendEntry. */
+export const USAGE_ENTRY_TYPE = "vision-handoff-usage";
+
+/** Event-bus channel emitted via pi.events. */
+export const USAGE_EVENT_CHANNEL = "vision-handoff:usage";
+
+/** Parsed Neuralwatt SSE-comment energy/cost/MCR data for one describer call. */
+export interface VisionHandoffEnergyCapture {
+  energyJoules: number;
+  costUsd: number;
+  energyRaw: Record<string, unknown> | null;
+  mcrSessionRaw: Record<string, unknown> | null;
+  costRaw: Record<string, unknown> | null;
+}
+
+/** Shared empty capture. Safe to share: readEnergyFromTee returns a fresh object
+ *  and never mutates this; buildUsageRecord only reads from its argument. */
+export const EMPTY_ENERGY_CAPTURE: VisionHandoffEnergyCapture = {
+  energyJoules: 0,
+  costUsd: 0,
+  energyRaw: null,
+  mcrSessionRaw: null,
+  costRaw: null,
+};
+
+/** A single describer-call usage record. Energy fields are present only when
+ *  Neuralwatt SSE energy comments were captured (omitted, not zeroed, otherwise). */
+export interface VisionHandoffUsageRecord {
+  imageHash: string;
+  model: string;
+  provider: string;
+  responseModel?: string;
+  responseId?: string;
+  usage: Usage;
+  /** Neuralwatt energy — present only when SSE energy comments were captured. */
+  energyJoules?: number;
+  costUsd?: number;
+  energyRaw?: Record<string, unknown> | null;
+  mcrSessionRaw?: Record<string, unknown> | null;
+  costRaw?: Record<string, unknown> | null;
+}
+
+/**
+ * Parse Neuralwatt SSE comment lines (`: energy`, `: cost`, `: mcr-session`)
+ * from a teed response body into a fresh capture object. Mirrors
+ * pi-neuralwatt-provider's readEnergyFromTee but returns a result instead of
+ * mutating module state — vision-handoff describes images concurrently (the
+ * before_agent_start warm-up fires several in parallel), so each call needs its
+ * own capture routed via {@link describeAls}. For non-Neuralwatt vision models
+ * no comment lines are present and the returned capture stays empty.
+ */
+export async function readEnergyFromTee(
+  body: ReadableStream<Uint8Array>,
+): Promise<VisionHandoffEnergyCapture> {
+  const result: VisionHandoffEnergyCapture = { ...EMPTY_ENERGY_CAPTURE };
+  const reader = body.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+
+  function processLine(line: string): void {
+    const trimmed = line.trim();
+    if (trimmed.startsWith(": energy ")) {
+      try {
+        const energy = JSON.parse(trimmed.slice(9));
+        result.energyJoules += energy.energy_joules || 0;
+        result.energyRaw = energy;
+      } catch {
+        // malformed energy comment — ignore
+      }
+    } else if (trimmed.startsWith(": mcr-session ")) {
+      try {
+        result.mcrSessionRaw = JSON.parse(trimmed.slice(14));
+      } catch {
+        // malformed mcr-session comment — ignore
+      }
+    } else if (trimmed.startsWith(": cost ")) {
+      try {
+        const cost = JSON.parse(trimmed.slice(7));
+        result.costUsd += cost.request_cost_usd || 0;
+        result.costRaw = cost;
+      } catch {
+        // malformed cost comment — ignore
+      }
+    }
+  }
+
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      buffer += decoder.decode(value, { stream: true });
+      const lines = buffer.split("\n");
+      buffer = lines.pop() || "";
+      for (const line of lines) processLine(line);
+    }
+  } catch {
+    // tee stream may error if the main stream is aborted — that's fine
+  }
+
+  const final = decoder.decode(new Uint8Array(0), { stream: false });
+  const remaining = (buffer + final).trim();
+  if (remaining) processLine(remaining);
+
+  try {
+    reader.releaseLock();
+  } catch {
+    // ignore
+  }
+  return result;
+}
+
+/**
+ * Build a usage record from a describer response + energy capture. Returns
+ * null when there is nothing meaningful to report (e.g. a provider-level
+ * failure with zero tokens and no energy) so the caller can skip emitting.
+ *
+ * Energy fields are OMITTED entirely (not zeroed) when no Neuralwatt SSE energy
+ * comments were captured, so consumers can distinguish "no energy" from
+ * "zero energy".
+ */
+export function buildUsageRecord(
+  response: AssistantMessage,
+  capture: VisionHandoffEnergyCapture,
+  visionModel: Model<Api>,
+  imageHash: string,
+): VisionHandoffUsageRecord | null {
+  const hasEnergy = !!(
+    capture.energyRaw ||
+    capture.costRaw ||
+    capture.mcrSessionRaw ||
+    capture.energyJoules > 0 ||
+    capture.costUsd > 0
+  );
+  const hasUsage =
+    !!response.usage &&
+    (response.usage.totalTokens > 0 || response.usage.input > 0 || response.usage.output > 0);
+  if (!hasUsage && !hasEnergy) return null;
+
+  const record: VisionHandoffUsageRecord = {
+    imageHash,
+    model: response.model || visionModel.id,
+    provider: response.provider || visionModel.provider,
+    responseModel: response.responseModel,
+    responseId: response.responseId,
+    usage: response.usage,
+  };
+  if (hasEnergy) {
+    record.energyJoules = capture.energyJoules;
+    record.costUsd = capture.costUsd;
+    record.energyRaw = capture.energyRaw;
+    record.mcrSessionRaw = capture.mcrSessionRaw;
+    record.costRaw = capture.costRaw;
+  }
+  return record;
+}
+
+// ── Concurrency-safe fetch interceptor ─────────────────────────────────────
+//
+// The only body-interception point for complete() is globalThis.fetch (pi-ai's
+// StreamOptions.onResponse exposes headers only, not the body where the
+// `: energy` SSE comments live). before_agent_start fires several describeImage()
+// calls fire-and-forget, so a naïve save/patch/restore of globalThis.fetch would
+// clobber under concurrency. The fix is a refcounted shared interceptor
+// (installed only while ≥1 describe is in flight, so non-describe fetches pass
+// through unmodified) + AsyncLocalStorage to route each teed response body to
+// the describe call that issued it. Nested patches from other extensions (e.g.
+// pi-neuralwatt-provider's streamNeuralwatt, which also tees for its own energy
+// display) chain on top and restore back to this interceptor, so both tees read
+// the same comment lines independently.
+
+/** Per-describer-call ALS slot carrying the energy tee reader. */
+export interface DescribeContext {
+  energyReader: Promise<VisionHandoffEnergyCapture> | undefined;
+}
+
+/** Routes each teed response body to the describe call that issued it. */
+export const describeAls = new AsyncLocalStorage<DescribeContext>();
+
+let fetchInterceptRefCount = 0;
+let savedRealFetch: typeof globalThis.fetch | null = null;
+
+/** Install the globalThis.fetch interceptor. Refcounted: the first caller
+ *  patches fetch; later callers just bump the count. Idempotent per install. */
+export function installFetchInterceptor(): void {
+  if (fetchInterceptRefCount === 0) {
+    savedRealFetch = globalThis.fetch;
+    globalThis.fetch = interceptedFetch;
+  }
+  fetchInterceptRefCount++;
+}
+
+/** Remove the interceptor. Refcounted: only the last caller restores fetch. */
+export function uninstallFetchInterceptor(): void {
+  if (fetchInterceptRefCount > 0) fetchInterceptRefCount--;
+  if (fetchInterceptRefCount === 0 && savedRealFetch) {
+    globalThis.fetch = savedRealFetch;
+    savedRealFetch = null;
+  }
+}
+
+/** Current refcount — 0 means the interceptor is not installed. Test hook. */
+export function fetchInterceptorRefcount(): number {
+  return fetchInterceptRefCount;
+}
+
+async function interceptedFetch(input: RequestInfo | URL, init?: RequestInit): Promise<Response> {
+  const real = savedRealFetch ?? globalThis.fetch;
+  const response = await real(input, init);
+  const store = describeAls.getStore();
+  // Outside a describer call (no ALS store) or for bodiless responses: pass
+  // through untouched.
+  if (!store || !response.body) return response;
+  const [bodyForSdk, bodyForEnergy] = response.body.tee();
+  store.energyReader = readEnergyFromTee(bodyForEnergy);
+  return new Response(bodyForSdk, {
+    headers: response.headers,
+    status: response.status,
+    statusText: response.statusText,
+  });
+}

package/vision-handoff.ts

@@ -37,21 +37,35 @@ import {
   HANDOFF_COMMAND_DESCRIPTION,
   IMAGE_PLACEHOLDER_PREFIX,
   IMAGE_PLACEHOLDER_SUFFIX,
+  USAGE_ENTRY_TYPE,
+  USAGE_EVENT_CHANNEL,
+  EMPTY_ENERGY_CAPTURE,
+  buildUsageRecord,
+  describeAls,
   extractImageFromBlock,
   formatModelRef,
   insertImageDescriptions,
+  installFetchInterceptor,
   isVisionModel,
   makeReplacementText,
   parseModelRef,
   readConfig,
   truncateDescription,
+  uninstallFetchInterceptor,
   writeConfig,
+  type DescribeContext,
   type VisionHandoffConfig,
+  type VisionHandoffEnergyCapture,
+  type VisionHandoffUsageRecord,
 } from "./src/index.js";
 import { VisionModelSelectorComponent, type VisionModelSelectorResult } from "./src/vision-model-selector.js";
 
 const UNAVAILABLE = `${IMAGE_PLACEHOLDER_PREFIX}description unavailable${IMAGE_PLACEHOLDER_SUFFIX}`;
 
+// Usage reporter; wired to pi.appendEntry + pi.events.emit in the default
+// export. No-op until then so describeImage is safe to call before wiring.
+let reportUsage: (record: VisionHandoffUsageRecord) => void = () => {};
+
 let config: VisionHandoffConfig = readConfig();
 
 /** User prompt for the current agent turn, captured from before_agent_start. */
@@ -125,17 +139,38 @@ async function describeImage(
 
     const controller = new AbortController();
     const timer = setTimeout(() => controller.abort(), DESCRIBE_TIMEOUT_MS);
+    // Energy/token capture for this describer call. The fetch interceptor is
+    // refcount-installed around the complete() window and routes the teed
+    // response body to this describe's AsyncLocalStorage slot — so concurrent
+    // describes (before_agent_start warm-up fires several in parallel) each
+    // get their own capture without clobbering globalThis.fetch. For non-
+    // Neuralwatt vision models no SSE energy comments are present and the
+    // capture stays empty (energy fields omitted from the record).
+    const describeCtx: DescribeContext = { energyReader: undefined };
+    installFetchInterceptor();
     try {
-      const response = await complete(
-        visionModel,
-        { systemPrompt, messages: [userMessage] },
-        {
-          apiKey: auth.apiKey,
-          headers: auth.headers,
-          signal: controller.signal,
-          maxTokens: cfg.maxTokens,
-        },
+      const response = await describeAls.run(describeCtx, async () =>
+        complete(
+          visionModel,
+          { systemPrompt, messages: [userMessage] },
+          {
+            apiKey: auth.apiKey,
+            headers: auth.headers,
+            signal: controller.signal,
+            maxTokens: cfg.maxTokens,
+          },
+        ),
       );
+      let capture: VisionHandoffEnergyCapture = EMPTY_ENERGY_CAPTURE;
+      if (describeCtx.energyReader) {
+        try {
+          capture = await describeCtx.energyReader;
+        } catch {
+          // tee aborted with the main stream — keep the empty capture
+        }
+      }
+      const record = buildUsageRecord(response, capture, visionModel, key);
+      if (record) reportUsage(record);
       if (response.stopReason === "aborted" || response.stopReason === "error") {
         return UNAVAILABLE;
       }
@@ -159,6 +194,8 @@ async function describeImage(
     } catch {
       return UNAVAILABLE;
     } finally {
+      if (describeCtx.energyReader) describeCtx.energyReader.catch(() => {});
+      uninstallFetchInterceptor();
       clearTimeout(timer);
     }
   })();
@@ -213,6 +250,26 @@ function notifyUnresolvedVisionModel(ctx: ExtensionContext, ref: string): void {
 export default function (pi: ExtensionAPI) {
   config = readConfig();
 
+  // Wire the usage reporter to pi's persistence + event bus. appendEntry
+  // persists the record so it replays on session resume/branch, and the event
+  // lets live consumers filter on one channel for tokens AND energy. Each call
+  // is independently guarded so a persistence/emit failure never breaks a
+  // describer turn. Re-assigned every factory invocation (pi re-runs the
+  // factory on /new, /resume, fork, /reload) so the closure always references
+  // the live pi.
+  reportUsage = (record: VisionHandoffUsageRecord) => {
+    try {
+      pi.appendEntry(USAGE_ENTRY_TYPE, record);
+    } catch {
+      // never break the describer on persistence failure
+    }
+    try {
+      pi.events?.emit(USAGE_EVENT_CHANNEL, record);
+    } catch {
+      // never break the describer on emit failure
+    }
+  };
+
   pi.on("session_start", async () => {
     // Reload in case the user edited the config on disk from another session.
     config = readConfig();