@pi-unipi/footer 2.1.0 → 2.1.2

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@pi-unipi/footer",
-  "version": "2.1.0",
-  "description": "Persistent status bar for Unipi \u2014 subscribes to UNIPI_EVENTS and renders key stats from all unipi packages",
+  "version": "2.1.2",
+  "description": "Persistent status bar for Unipi — subscribes to UNIPI_EVENTS and renders key stats from all unipi packages",
   "type": "module",
   "main": "src/index.ts",
   "license": "MIT",
@@ -32,7 +32,7 @@
     "access": "public"
   },
   "dependencies": {
-    "@pi-unipi/core": "2.1.0"
+    "@pi-unipi/core": "2.1.2"
   },
   "peerDependencies": {
     "@earendil-works/pi-coding-agent": "^0.80.0",

package/src/index.ts

@@ -96,6 +96,14 @@ export default function footerExtension(pi: ExtensionAPI): void {
     state.registry.registerGroup(group);
   }
 
+  // ─── TPS streaming-event hooks (registered once) ────────────────────────
+  // pi.on() has no unsubscribe, so we register these exactly once at factory
+  // time (not per session_start) to avoid duplicate handlers accumulating
+  // across session restarts. The streamingIndex counter is reset on each
+  // session_shutdown. These hooks feed the TPS tracker in real time; the
+  // 1s branch-scan in the refresh timer only reconciles persisted messages.
+  wireTpsStreamingEvents(pi);
+
   // ─── Session lifecycle ──────────────────────────────────────────────────
 
   pi.on("session_start", async (_event, ctx) => {
@@ -127,6 +135,7 @@ export default function footerExtension(pi: ExtensionAPI): void {
     }
     state.tuiRef = null;
     tpsTracker.reset();
+    resetTpsStreamingIndex();
   });
 
   // ─── Register commands ──────────────────────────────────────────────────
@@ -152,10 +161,14 @@ function setupFooterUI(pi: ExtensionAPI, ctx: ExtensionContext, state: FooterSta
   ctx.ui.setFooter((tui, _theme, footerData) => {
     state.tuiRef = tui;
 
-    // Start periodic refresh for time-sensitive segments (e.g. clock)
+    // Start periodic refresh for time-sensitive segments (e.g. clock, TPS)
     if (!state.refreshTimer) {
       state.refreshTimer = setInterval(() => {
-        // Feed TPS tracker with per-message data
+        // Re-seed TPS tracker from the session branch on each tick.
+        // Streaming events (message_start/update/end) handle live updates
+        // in real time; this scan reconciles the tracker with persisted
+        // messages after compactions, branch switches, or session reloads
+        // where in-flight streaming state may have been lost.
         try {
           const piCtx = state.piContext as Record<string, unknown> | undefined;
           if (piCtx?.sessionManager) {
@@ -168,9 +181,10 @@ function setupFooterUI(pi: ExtensionAPI, ctx: ExtensionContext, state: FooterSta
               const m = e.message;
               if (!m || m.role !== "assistant") continue;
               if (m.stopReason === "error" || m.stopReason === "aborted") continue;
-              const output = m.usage?.output ?? 0;
               const hasStop = !!m.stopReason;
-              tpsTracker.onMessageUpdate(msgIndex, output, hasStop);
+              // Pass the whole message: TPS tracker counts tokens from content,
+              // not from usage.output (which is 0 during streaming).
+              tpsTracker.onMessageUpdate(msgIndex, m, hasStop);
               msgIndex++;
             }
           }
@@ -251,3 +265,56 @@ function setupFooterUI(pi: ExtensionAPI, ctx: ExtensionContext, state: FooterSta
     };
   }, { placement: "belowEditor" });
 }
+
+// ─── TPS streaming-event hooks ──────────────────────────────────────────────
+
+/**
+ * Sequential index of the currently-streaming assistant message. Tracked
+ * locally because pi does not expose a stable message index on streaming
+ * events, and the TPS tracker keys records off this index.
+ */
+let tpsStreamingIndex = -1;
+
+/** Reset the streaming index (called on session_shutdown). */
+function resetTpsStreamingIndex(): void {
+  tpsStreamingIndex = -1;
+}
+
+/**
+ * Subscribe to pi's message streaming events and feed the TPS tracker in real
+ * time. This complements the 1s branch-scan in the refresh timer, which only
+ * sees persisted (completed) messages. Without these hooks the tracker would
+ * never observe an in-flight assistant message, so live TPS would stay frozen
+ * at the last completed message's value.
+ *
+ * Registered once at extension-factory time (pi.on has no unsubscribe, so we
+ * must not re-register per session_start or handlers would accumulate).
+ */
+function wireTpsStreamingEvents(pi: ExtensionAPI): void {
+  const safe = (fn: () => void) => {
+    try { fn(); } catch { /* TPS is best-effort */ }
+  };
+
+  pi.on("message_start", ((event: { message: unknown }) => safe(() => {
+    const m = event.message as Record<string, unknown> | undefined;
+    if (!m || m.role !== "assistant") return;
+    if (m.stopReason === "error" || m.stopReason === "aborted") return;
+    tpsStreamingIndex++;
+    tpsTracker.onMessageUpdate(tpsStreamingIndex, m, false);
+  })) as (event: unknown) => void);
+
+  pi.on("message_update", ((event: { message: unknown }) => safe(() => {
+    if (tpsStreamingIndex < 0) return;
+    const m = event.message as Record<string, unknown> | undefined;
+    if (!m || m.role !== "assistant") return;
+    tpsTracker.onMessageUpdate(tpsStreamingIndex, m, false);
+  })) as (event: unknown) => void);
+
+  pi.on("message_end", ((event: { message: unknown }) => safe(() => {
+    if (tpsStreamingIndex < 0) return;
+    const m = event.message as Record<string, unknown> | undefined;
+    if (!m || m.role !== "assistant") return;
+    if (m.stopReason === "error" || m.stopReason === "aborted") return;
+    tpsTracker.onMessageUpdate(tpsStreamingIndex, m, true);
+  })) as (event: unknown) => void);
+}

package/src/tps-tracker.ts

@@ -4,37 +4,112 @@
  * Per-message TPS calculation for live generation rate display.
  * Tracks individual assistant messages with start/stop timestamps
  * to measure generation rate excluding idle/tool-execution time.
+ *
+ * ## Token counting
+ *
+ * Output tokens are counted from the assistant message *content* (text +
+ * thinking blocks) via a lightweight word/char heuristic, NOT from
+ * `usage.output`. Rationale:
+ *   - `usage.output` is only populated by the provider at stream end (or in
+ *     a final `message_delta`), so it is 0 throughout generation and useless
+ *     for a live rate.
+ *   - Counting the streamed text ourselves gives a real, monotonically
+ *     increasing numerator that reflects how much the model has actually
+ *     produced so far.
+ *   - This matches what users perceive as "tokens per second" (visible
+ *     output rate) and avoids the previous bug where the displayed TPS was
+ *     always inflated to 90+ t/s (caused by dividing a post-hoc `usage.output`
+ *     total by a too-short, incorrectly measured elapsed window).
+ *
+ * The heuristic intentionally over-counts slightly (whitespace + punctuation
+ * are treated as tokens). That is acceptable for a live speed indicator and
+ * keeps the estimate in the right order of magnitude (within ~10-20% of the
+ * provider's reported output tokens for typical English/code output).
  */
 
 /** Per-message TPS record */
 interface MessageTpsRecord {
-  /** Message index in the session */
+  /** Message index in the session (sequential, assistant-only) */
   messageIndex: number;
-  /** Output tokens for this message */
+  /** Estimated output tokens produced so far for this message */
   outputTokens: number;
-  /** When generation started (Date.now()) */
+  /** When generation started (Date.now(), ms) */
   startedAt: number;
-  /** When generation completed (Date.now()), 0 if still generating */
+  /** When generation completed (Date.now(), ms), 0 if still generating */
   completedAt: number;
-  /** Computed TPS for this message */
+  /** Computed TPS for this message (final, once completed) */
   tps: number;
 }
 
+/**
+ * Estimate the number of tokens in a string.
+ *
+ * Uses a cheap word + punctuation split. Good enough for a live TPS display;
+ * we deliberately avoid pulling in a full BPE tokenizer for performance and
+ * footprint reasons.
+ */
+function estimateTokens(text: string): number {
+  if (!text) return 0;
+  // Count word-ish tokens and standalone punctuation.
+  // \p{L} / \p{N} keep this Unicode-aware for non-ASCII content.
+  const wordish = text.match(/[\p{L}\p{N}]+/gu);
+  const punct = text.match(/[^\s\p{L}\p{N}]+/gu);
+  const wordCount = wordish ? wordish.length : 0;
+  const punctCount = punct ? punct.length : 0;
+  return wordCount + punctCount;
+}
+
+/**
+ * Estimate output tokens from an assistant message by summing text and
+ * thinking content lengths. Tool-call argument JSON is also counted since
+ * it represents model output.
+ *
+ * Accepts the raw message shape from pi's session entries (loosely typed).
+ */
+function estimateOutputTokens(message: unknown): number {
+  if (!message || typeof message !== "object") return 0;
+  const m = message as Record<string, unknown>;
+  if (m.role !== "assistant") return 0;
+  const content = m.content;
+  if (!Array.isArray(content)) return 0;
+
+  let total = 0;
+  for (const block of content) {
+    if (!block || typeof block !== "object") continue;
+    const b = block as Record<string, unknown>;
+    const type = b.type;
+    if (type === "text" && typeof b.text === "string") {
+      total += estimateTokens(b.text);
+    } else if (type === "thinking" && typeof b.thinking === "string") {
+      total += estimateTokens(b.thinking);
+    } else if (type === "tool_use" || type === "toolcall") {
+      // Count tool name + serialized arguments as model output.
+      const name = typeof b.name === "string" ? b.name : "";
+      const input = b.input ?? b.arguments;
+      let inputStr: string;
+      try {
+        inputStr = typeof input === "string" ? input : JSON.stringify(input ?? "");
+      } catch {
+        inputStr = "";
+      }
+      total += estimateTokens(name) + estimateTokens(inputStr);
+    }
+  }
+  return total;
+}
+
 /**
  * Tracks per-message TPS and computes live/session metrics.
  *
- * Usage: Call `onMessageUpdate()` whenever output tokens change.
- * The tracker records generation start/stop per message and computes
- * live TPS from the current message and session averages excluding idle time.
+ * Usage: Call `onMessageUpdate()` on every tick (e.g. 1s interval) with the
+ * current assistant message state. The tracker records generation start/stop
+ * per message and computes live TPS from elapsed wall-clock time.
  */
 export class TpsTracker {
-  /** Per-message records */
+  /** Per-message records (one per assistant message, in order) */
   private records: MessageTpsRecord[] = [];
 
-  /** Highest message index seen so far (for dedup) */
-  private lastSeenMessageCount = 0;
-
-  /** Total output tokens across all completed messages */
+  /** Total estimated output tokens across all messages */
   private totalOutput = 0;
 
   /**
@@ -42,15 +117,16 @@ export class TpsTracker {
    * Call this on every tick (e.g. 1s interval) with the current state.
    *
    * @param messageIndex - Index of the assistant message (0-based, sequential)
-   * @param outputTokens - Output tokens for this message
+   * @param message - The assistant message object (content used to count tokens)
    * @param hasStopReason - Whether this message has completed (has stopReason)
    */
-  onMessageUpdate(messageIndex: number, outputTokens: number, hasStopReason: boolean): void {
+  onMessageUpdate(messageIndex: number, message: unknown, hasStopReason: boolean): void {
     const now = Date.now();
+    const outputTokens = estimateOutputTokens(message);
 
     // New message — create a record
     if (messageIndex >= this.records.length) {
-      // Fill gaps if indices jump
+      // Fill gaps if indices jump (shouldn't normally happen)
       while (this.records.length < messageIndex) {
         this.records.push({
           messageIndex: this.records.length,
@@ -62,11 +138,12 @@ export class TpsTracker {
       }
 
       if (hasStopReason && outputTokens > 0) {
-        // Fast message: already completed on first sighting
-        // Estimate duration: floor of 1 second, or outputTokens/100, whichever is smaller
-        const estimatedDuration = Math.max(1, outputTokens / 100);
+        // Fast message: already completed on first sighting.
+        // We never observed it streaming, so we can't measure real duration.
+        // Fall back to a conservative estimate instead of inventing a 1s
+        // window (which previously produced inflated TPS values).
+        const estimatedDuration = Math.max(0.5, outputTokens / 60);
         const tps = outputTokens / estimatedDuration;
-
         this.records.push({
           messageIndex,
           outputTokens,
@@ -85,7 +162,6 @@ export class TpsTracker {
           tps: 0,
         });
       }
-      this.lastSeenMessageCount = messageIndex + 1;
       return;
     }
 
@@ -93,23 +169,24 @@ export class TpsTracker {
     const record = this.records[messageIndex];
     if (!record) return;
 
-    record.outputTokens = outputTokens;
+    if (record.completedAt === 0) {
+      // Still generating — update token count (live TPS computed on demand)
+      record.outputTokens = outputTokens;
 
-    if (record.completedAt === 0 && hasStopReason) {
-      // Message just completed
-      record.completedAt = now;
-      const durationSec = (record.completedAt - record.startedAt) / 1000;
-      record.tps = durationSec > 0 ? outputTokens / durationSec : outputTokens;
-      this.totalOutput += outputTokens;
-    } else if (record.completedAt === 0) {
-      // Still generating — update output tokens (live TPS computed on demand)
+      if (hasStopReason) {
+        // Message just completed
+        record.completedAt = now;
+        const durationSec = (record.completedAt - record.startedAt) / 1000;
+        record.tps = durationSec > 0 ? record.outputTokens / durationSec : 0;
+        this.totalOutput += record.outputTokens;
+      }
     }
   }
 
   /**
    * Get the live TPS from the currently generating message.
    * Returns the instantaneous rate based on tokens generated so far
-   * in the current message divided by elapsed time.
+   * in the current message divided by elapsed wall-clock time.
    */
   getLiveTps(): number {
     // Find the last record that's still generating
@@ -177,7 +254,7 @@ export class TpsTracker {
   }
 
   /**
-   * Get total output tokens for the session.
+   * Get total output tokens for the session (estimated).
    */
   getTotalOutput(): number {
     // Include tokens from incomplete messages too
@@ -195,7 +272,6 @@ export class TpsTracker {
    */
   reset(): void {
     this.records = [];
-    this.lastSeenMessageCount = 0;
     this.totalOutput = 0;
   }
 }