feat: split localpi token status rates

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
 

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,18 +34,22 @@ 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;
     } else if (update.text.length > state.outputText.length) {
       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, 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,48 @@ 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,
+  cacheWrite: number
+): string | undefined {
+  const tokens = prefillTokenCount(input, cacheWrite);
+  if (state.firstOutputAt === undefined || tokens <= 0) {
+    return undefined;
+  }
+  const seconds = secondsBetween(state.startedAt, state.firstOutputAt);
+  return \`prefill \${(tokens / seconds).toFixed(1)} tok/s\`;
+}
+
+function prefillTokenCount(input: number, cacheWrite: number): number {
+  return Math.max(input + cacheWrite, 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 = {

tests/extensions.test.ts

@@ -27,6 +27,12 @@ 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("input + cacheWrite");
+      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)");