From 93b5ccf0bd4aed0b61f9563e3c5e44cab6babd1e Mon Sep 17 00:00:00 2001
From: Onur Solmaz <2453968+osolmaz@users.noreply.github.com>
Date: Fri, 26 Jun 2026 13:03:19 +0800
Subject: [PATCH 1/2] feat: split localpi token status rates

---
 docs/runtime-specification.md            |  2 +-
 src/pi/extension-sources/token-status.ts | 68 ++++++++++++++++++++----
 tests/extensions.test.ts                 |  5 ++
 3 files changed, 64 insertions(+), 11 deletions(-)

diff --git a/docs/runtime-specification.md b/docs/runtime-specification.md
index a22c6ba..4cf385e 100644
--- a/docs/runtime-specification.md
+++ b/docs/runtime-specification.md
@@ -150,7 +150,7 @@ session dir: ~/.local/state/localpi/sessions
 Localpi installs two default extensions:
 
 - tool approval gate: ask before each tool call, and tell the model clearly when a tool call was blocked
-- token status: show live output token estimate while streaming and final exact token stats when usage data is available
+- token status: show live generation speed while streaming, then final prefill and generation rates when usage data is available
 
 ## System Prompt
 
diff --git a/src/pi/extension-sources/token-status.ts b/src/pi/extension-sources/token-status.ts
index d123400..c0e2555 100644
--- a/src/pi/extension-sources/token-status.ts
+++ b/src/pi/extension-sources/token-status.ts
@@ -11,6 +11,7 @@ type Usage = {
 
 type TurnState = {
   startedAt: number;
+  firstOutputAt?: number;
   outputText: string;
   estimatedOutputTokens: number;
   lastStatusAt: number;
@@ -33,6 +34,7 @@ export default function localpiTokenStatus(pi: ExtensionAPI): void {
     if (!ctx.hasUI || state === undefined) {
       return;
     }
+    const now = Date.now();
     const update = textUpdateFromUnknown(event.assistantMessageEvent ?? event.message ?? event);
     if (update.kind === "delta") {
       state.outputText += update.text;
@@ -40,11 +42,14 @@ export default function localpiTokenStatus(pi: ExtensionAPI): void {
       state.outputText = update.text;
     }
     state.estimatedOutputTokens = Math.ceil(state.outputText.length / 4);
-    if (Date.now() - state.lastStatusAt < 250) {
+    if (state.firstOutputAt === undefined && state.outputText.length > 0) {
+      state.firstOutputAt = now;
+    }
+    if (now - state.lastStatusAt < 250) {
       return;
     }
-    state.lastStatusAt = Date.now();
-    ctx.ui.setStatus("localpi-perf", ctx.ui.theme.fg("dim", statusText(state)));
+    state.lastStatusAt = now;
+    ctx.ui.setStatus("localpi-perf", ctx.ui.theme.fg("dim", statusText(state, now)));
   });
 
   pi.on("turn_end", (event, ctx) => {
@@ -65,7 +70,10 @@ export default function localpiTokenStatus(pi: ExtensionAPI): void {
     const input = usage?.input ?? 0;
     const cacheRead = usage?.cacheRead ?? 0;
     const cacheWrite = usage?.cacheWrite ?? 0;
-    const elapsedSeconds = elapsed(state);
+    const now = Date.now();
+    const elapsedSeconds = elapsed(state, now);
+    const decodeSeconds = generationElapsed(state, now);
+    const prefillText = prefillStatusText(state, input, cacheRead, cacheWrite);
     const context = ctx.getContextUsage();
     const contextText =
       context && context.percent !== null
@@ -77,7 +85,8 @@ export default function localpiTokenStatus(pi: ExtensionAPI): void {
       ctx.ui.theme.fg(
         "dim",
         [
-          \`\${(output / elapsedSeconds).toFixed(1)} tok/s\`,
+          \`gen \${(output / decodeSeconds).toFixed(1)} tok/s\`,
+          prefillText,
           \`out \${output}\`,
           \`in \${input}\`,
           cacheRead > 0 ? \`cache \${cacheRead}\` : undefined,
@@ -98,13 +107,52 @@ export default function localpiTokenStatus(pi: ExtensionAPI): void {
   });
 }
 
-function statusText(state: TurnState): string {
-  const elapsedSeconds = elapsed(state);
-  return \`\${(state.estimatedOutputTokens / elapsedSeconds).toFixed(1)} tok/s | out ~\${state.estimatedOutputTokens} | \${elapsedSeconds.toFixed(1)}s\`;
+function statusText(state: TurnState, now: number): string {
+  const elapsedSeconds = elapsed(state, now);
+  if (state.firstOutputAt === undefined) {
+    return \`prefill \${elapsedSeconds.toFixed(1)}s | out ~\${state.estimatedOutputTokens}\`;
+  }
+  const decodeSeconds = generationElapsed(state, now);
+  const prefillSeconds = secondsBetween(state.startedAt, state.firstOutputAt);
+  return [
+    \`gen \${(state.estimatedOutputTokens / decodeSeconds).toFixed(1)} tok/s\`,
+    \`out ~\${state.estimatedOutputTokens}\`,
+    \`prefill \${prefillSeconds.toFixed(1)}s\`,
+    \`total \${elapsedSeconds.toFixed(1)}s\`
+  ].join(" | ");
+}
+
+function prefillStatusText(
+  state: TurnState,
+  input: number,
+  cacheRead: number,
+  cacheWrite: number
+): string | undefined {
+  if (state.firstOutputAt === undefined || input <= 0) {
+    return undefined;
+  }
+  const tokens = prefillTokenCount(input, cacheRead, cacheWrite);
+  const seconds = secondsBetween(state.startedAt, state.firstOutputAt);
+  return \`prefill \${(tokens / seconds).toFixed(1)} tok/s\`;
+}
+
+function prefillTokenCount(input: number, cacheRead: number, cacheWrite: number): number {
+  if (cacheWrite > 0) {
+    return cacheWrite;
+  }
+  return Math.max(input - cacheRead, 0);
+}
+
+function generationElapsed(state: TurnState, now: number): number {
+  return secondsBetween(state.firstOutputAt ?? state.startedAt, now);
+}
+
+function elapsed(state: TurnState, now: number): number {
+  return secondsBetween(state.startedAt, now);
 }
 
-function elapsed(state: TurnState): number {
-  return Math.max((Date.now() - state.startedAt) / 1000, 0.001);
+function secondsBetween(start: number, end: number): number {
+  return Math.max((end - start) / 1000, 0.001);
 }
 
 type TextUpdate = {
diff --git a/tests/extensions.test.ts b/tests/extensions.test.ts
index 64a14b3..531e921 100644
--- a/tests/extensions.test.ts
+++ b/tests/extensions.test.ts
@@ -27,6 +27,11 @@ describe("Pi extensions", () => {
       expect(status).toContain("tok/s");
       expect(status).toContain("message_update");
       expect(status).toContain("currentTurn");
+      expect(status).toContain("firstOutputAt");
+      expect(status).toContain("generationElapsed");
+      expect(status).toContain("prefillTokenCount");
+      expect(status).toContain("prefill ");
+      expect(status).toContain("gen ");
       expect(status).toContain("outputText += update.text");
       expect(status).toContain('kind: "delta"');
       expect(status).not.toContain("turns.get(event.turnIndex)");

From f0617d5e2566b183dcb1d497a2cf575e61f38071 Mon Sep 17 00:00:00 2001
From: Onur Solmaz <2453968+osolmaz@users.noreply.github.com>
Date: Fri, 26 Jun 2026 13:08:18 +0800
Subject: [PATCH 2/2] fix: count uncached prefill token buckets

---
 src/pi/extension-sources/token-status.ts | 14 +++++---------
 tests/extensions.test.ts                 |  1 +
 2 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/src/pi/extension-sources/token-status.ts b/src/pi/extension-sources/token-status.ts
index c0e2555..f6d493f 100644
--- a/src/pi/extension-sources/token-status.ts
+++ b/src/pi/extension-sources/token-status.ts
@@ -73,7 +73,7 @@ export default function localpiTokenStatus(pi: ExtensionAPI): void {
     const now = Date.now();
     const elapsedSeconds = elapsed(state, now);
     const decodeSeconds = generationElapsed(state, now);
-    const prefillText = prefillStatusText(state, input, cacheRead, cacheWrite);
+    const prefillText = prefillStatusText(state, input, cacheWrite);
     const context = ctx.getContextUsage();
     const contextText =
       context && context.percent !== null
@@ -125,22 +125,18 @@ function statusText(state: TurnState, now: number): string {
 function prefillStatusText(
   state: TurnState,
   input: number,
-  cacheRead: number,
   cacheWrite: number
 ): string | undefined {
-  if (state.firstOutputAt === undefined || input <= 0) {
+  const tokens = prefillTokenCount(input, cacheWrite);
+  if (state.firstOutputAt === undefined || tokens <= 0) {
     return undefined;
   }
-  const tokens = prefillTokenCount(input, cacheRead, cacheWrite);
   const seconds = secondsBetween(state.startedAt, state.firstOutputAt);
   return \`prefill \${(tokens / seconds).toFixed(1)} tok/s\`;
 }
 
-function prefillTokenCount(input: number, cacheRead: number, cacheWrite: number): number {
-  if (cacheWrite > 0) {
-    return cacheWrite;
-  }
-  return Math.max(input - cacheRead, 0);
+function prefillTokenCount(input: number, cacheWrite: number): number {
+  return Math.max(input + cacheWrite, 0);
 }
 
 function generationElapsed(state: TurnState, now: number): number {
diff --git a/tests/extensions.test.ts b/tests/extensions.test.ts
index 531e921..7b679bc 100644
--- a/tests/extensions.test.ts
+++ b/tests/extensions.test.ts
@@ -30,6 +30,7 @@ describe("Pi extensions", () => {
       expect(status).toContain("firstOutputAt");
       expect(status).toContain("generationElapsed");
       expect(status).toContain("prefillTokenCount");
+      expect(status).toContain("input + cacheWrite");
       expect(status).toContain("prefill ");
       expect(status).toContain("gen ");
       expect(status).toContain("outputText += update.text");