Skip to content

Commit 74fac19

Browse files
halfaipghalfaipg
committed
fix(status): honest context-fill bar and live tok/s
Two visible status-bar bugs the user surfaced on Codebase Auto: The fill bar pegged at 0% because it was reading state.usage (cumulative across turns, only updated when message_end carries a usage payload). The OAuth-fronted proxy was stripping usage, so the bar never moved. Switch to turnUsage when present and a char-based estimate when not, plus include in-flight streaming so the bar fills mid-turn instead of jumping at message_end. The tok/s readout showed 3 tok/s because the rate was averaged from the moment the streaming flag flipped — which includes long thinking-wait periods before any text shows up. Replace with a 4-second sliding window over recent char growth, so a model that thought for 60s and then streams quickly reports the real recent rate, not the time-since-streaming-began average.
1 parent 8aff684 commit 74fac19

3 files changed

Lines changed: 219 additions & 24 deletions

File tree

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "codebase-cli",
3-
"version": "2.0.0-pre.36",
3+
"version": "2.0.0-pre.37",
44
"description": "Codebase CLI — a TypeScript coding agent on the pi-mono runtime. OAuth-aware, any LLM provider, single install.",
55
"keywords": [
66
"ai",

‎src/ui/Status.test.ts‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
import type { AgentMessage } from "@earendil-works/pi-agent-core";
2+
import { describe, expect, it } from "vitest";
3+
import { type ChatState, EMPTY_USAGE } from "../types.js";
4+
import { estimateContextTokens } from "./Status.js";
5+
6+
const MODEL = { provider: "test", id: "test-model", name: "Test" };
7+
8+
function baseState(): ChatState {
9+
return {
10+
messages: [],
11+
tools: new Map(),
12+
status: "idle",
13+
usage: EMPTY_USAGE,
14+
model: MODEL,
15+
};
16+
}
17+
18+
describe("estimateContextTokens", () => {
19+
it("returns 0 for an empty conversation", () => {
20+
expect(estimateContextTokens(baseState())).toBe(0);
21+
});
22+
23+
it("prefers turnUsage.input + cacheRead when the provider reports them", () => {
24+
const state: ChatState = {
25+
...baseState(),
26+
turnUsage: {
27+
input: 1000,
28+
output: 0,
29+
cacheRead: 500,
30+
cacheWrite: 0,
31+
totalTokens: 1500,
32+
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
33+
},
34+
};
35+
expect(estimateContextTokens(state)).toBe(1500);
36+
});
37+
38+
it("falls back to char-based estimation when turnUsage is missing", () => {
39+
const state: ChatState = {
40+
...baseState(),
41+
messages: [
42+
{ role: "user", content: "hello world", timestamp: 1 } as AgentMessage,
43+
{
44+
role: "assistant",
45+
content: [{ type: "text", text: "x".repeat(400) }],
46+
api: "chat",
47+
provider: "p",
48+
model: "m",
49+
usage: EMPTY_USAGE,
50+
stopReason: "stop",
51+
timestamp: 2,
52+
} as AgentMessage,
53+
],
54+
};
55+
// 11 ("hello world") + 400 (x's) = 411 chars; 411 / 4 = ~103 tokens
56+
expect(estimateContextTokens(state)).toBeGreaterThan(95);
57+
expect(estimateContextTokens(state)).toBeLessThan(115);
58+
});
59+
60+
it("falls back when turnUsage is present but reports zero (proxy stripped usage)", () => {
61+
const state: ChatState = {
62+
...baseState(),
63+
turnUsage: {
64+
input: 0,
65+
output: 0,
66+
cacheRead: 0,
67+
cacheWrite: 0,
68+
totalTokens: 0,
69+
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
70+
},
71+
messages: [{ role: "user", content: "x".repeat(400), timestamp: 1 } as AgentMessage],
72+
};
73+
expect(estimateContextTokens(state)).toBeGreaterThan(95);
74+
});
75+
76+
it("includes streaming content so the bar fills mid-turn, not just at message_end", () => {
77+
const baseline: ChatState = {
78+
...baseState(),
79+
turnUsage: {
80+
input: 1000,
81+
output: 0,
82+
cacheRead: 0,
83+
cacheWrite: 0,
84+
totalTokens: 1000,
85+
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
86+
},
87+
};
88+
const withStream: ChatState = {
89+
...baseline,
90+
streaming: {
91+
role: "assistant",
92+
content: [{ type: "text", text: "y".repeat(800) }],
93+
} as unknown as AgentMessage,
94+
};
95+
// 1000 reported + 800/4 = 200 estimated streaming tokens = 1200
96+
expect(estimateContextTokens(withStream)).toBe(1200);
97+
expect(estimateContextTokens(baseline)).toBe(1000);
98+
});
99+
100+
it("counts tool-call args toward the fallback estimate", () => {
101+
const args = { path: "src/lib/some/very/long/path.ts", limit: 200 };
102+
const state: ChatState = {
103+
...baseState(),
104+
messages: [
105+
{
106+
role: "assistant",
107+
content: [{ type: "toolCall", id: "t1", name: "read_file", arguments: args }],
108+
api: "chat",
109+
provider: "p",
110+
model: "m",
111+
usage: EMPTY_USAGE,
112+
stopReason: "stop",
113+
timestamp: 1,
114+
} as unknown as AgentMessage,
115+
],
116+
};
117+
expect(estimateContextTokens(state)).toBeGreaterThan(0);
118+
});
119+
120+
it("counts thinking blocks toward the fallback estimate", () => {
121+
const state: ChatState = {
122+
...baseState(),
123+
messages: [
124+
{
125+
role: "assistant",
126+
content: [{ type: "thinking", thinking: "z".repeat(800) }],
127+
api: "chat",
128+
provider: "p",
129+
model: "m",
130+
usage: EMPTY_USAGE,
131+
stopReason: "stop",
132+
timestamp: 1,
133+
} as unknown as AgentMessage,
134+
],
135+
};
136+
// 800 / 4 = 200 tokens
137+
expect(estimateContextTokens(state)).toBe(200);
138+
});
139+
});

‎src/ui/Status.tsx‎

Lines changed: 79 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,14 @@
11
import { basename } from "node:path";
2+
import type { AgentMessage } from "@earendil-works/pi-agent-core";
23
import { Box, Text } from "ink";
34
import { useEffect, useRef, useState } from "react";
45
import type { ChatState } from "../types.js";
56
import { Throbber } from "./Throbber.js";
67

8+
/** Average chars-per-token across the major model families. Used only as a
9+
* fallback when the provider doesn't return usage info on message_end. */
10+
const CHARS_PER_TOKEN = 4;
11+
712
interface StatusProps {
813
state: ChatState;
914
cwd?: string;
@@ -67,7 +72,7 @@ export function Status({ state, cwd, contextWindow = 200_000 }: StatusProps) {
6772
const tokRate = useTokenRate(state);
6873
const elapsedSec = useBusyElapsed(busy);
6974
const u = state.usage;
70-
const usedTokens = u.input + u.cacheRead;
75+
const usedTokens = estimateContextTokens(state);
7176
const ctxPct = contextWindow > 0 ? Math.min(100, Math.round((usedTokens / contextWindow) * 100)) : 0;
7277
const cwdLabel = cwd ? basename(cwd) || "/" : "";
7378
const modelLabel = state.model.name || state.model.id;
@@ -145,38 +150,89 @@ function findRunningTool(state: ChatState): string | undefined {
145150
}
146151

147152
/**
148-
* Estimate the live token-output rate during streaming. Pi-ai only
149-
* surfaces accurate `usage` at message_end, so for the live counter
150-
* we approximate from the streaming message's character length using
151-
* the common ~4-chars-per-token rule. Cheap, no extra deps, and
152-
* accurate enough for a status-bar readout.
153+
* Estimate the live token-output rate during streaming using a 4-second
154+
* sliding window over recent character growth. This avoids dragging the
155+
* rate down with the pre-text wait period — a thinking-heavy model that
156+
* spent 60s reasoning before emitting its first token should show "120
157+
* tok/s" once it starts streaming, not "5 tok/s averaged with the wait."
153158
*
154-
* Returns undefined when not streaming, or when too few chars have
155-
* accumulated for the rate to be meaningful (so the bar doesn't
156-
* flicker a noisy "9999 tok/s" in the first 100ms).
159+
* Returns undefined when not streaming, or when the window doesn't yet
160+
* have enough samples / delta for the rate to be meaningful (so the bar
161+
* doesn't flicker noisy values in the first half-second).
157162
*/
158163
function useTokenRate(state: ChatState): number | undefined {
159-
const startRef = useRef<number | undefined>(undefined);
160-
const [tick, setTick] = useState(0);
164+
const samplesRef = useRef<Array<{ t: number; c: number }>>([]);
165+
const charsRef = useRef(0);
166+
const [, setTick] = useState(0);
161167
const streaming = state.status === "streaming";
168+
169+
// Keep the latest char count in a ref so the interval callback always
170+
// reads the live value rather than the closure-captured one.
171+
charsRef.current = streaming ? streamingChars(state) : 0;
172+
162173
useEffect(() => {
163174
if (!streaming) {
164-
startRef.current = undefined;
175+
samplesRef.current = [];
165176
return;
166177
}
167-
if (startRef.current === undefined) startRef.current = Date.now();
168-
const id = setInterval(() => setTick((t) => t + 1), 500);
178+
const sample = () => {
179+
const now = Date.now();
180+
samplesRef.current.push({ t: now, c: charsRef.current });
181+
const cutoff = now - 4000;
182+
while (samplesRef.current.length > 0 && samplesRef.current[0].t < cutoff) {
183+
samplesRef.current.shift();
184+
}
185+
setTick((n) => n + 1);
186+
};
187+
sample(); // seed immediately
188+
const id = setInterval(sample, 500);
169189
return () => clearInterval(id);
170190
}, [streaming]);
171-
if (!streaming || !startRef.current) return undefined;
172-
const elapsedSec = (Date.now() - startRef.current) / 1000;
173-
if (elapsedSec < 0.5) return undefined;
174-
void tick; // force re-eval on each interval
175-
const chars = streamingChars(state);
176-
if (chars < 40) return undefined;
177-
const tokens = chars / 4;
178-
const rate = tokens / elapsedSec;
179-
return Math.round(rate);
191+
192+
if (!streaming || samplesRef.current.length < 2) return undefined;
193+
const oldest = samplesRef.current[0];
194+
const newest = samplesRef.current[samplesRef.current.length - 1];
195+
const dt = (newest.t - oldest.t) / 1000;
196+
if (dt < 0.5) return undefined;
197+
const dc = newest.c - oldest.c;
198+
if (dc < 10) return undefined;
199+
return Math.round(dc / CHARS_PER_TOKEN / dt);
200+
}
201+
202+
/**
203+
* Tokens currently in the model's context, for the status-bar fill meter.
204+
* Prefers the last-turn's reported `input + cacheRead` from pi-ai, since
205+
* that's literally what the model saw. Falls back to char-based estimation
206+
* when the provider strips usage (e.g. some OAuth-fronted proxies) so the
207+
* bar still grows as the conversation grows. Streaming content is added
208+
* on top of the prior-turn baseline so the bar visibly fills during a turn
209+
* instead of jumping at message_end.
210+
*/
211+
export function estimateContextTokens(state: ChatState): number {
212+
if (state.turnUsage && state.turnUsage.input + state.turnUsage.cacheRead > 0) {
213+
const reported = state.turnUsage.input + state.turnUsage.cacheRead;
214+
const streamingExtra = Math.round(streamingChars(state) / CHARS_PER_TOKEN);
215+
return reported + streamingExtra;
216+
}
217+
let chars = 0;
218+
for (const msg of state.messages) chars += messageChars(msg);
219+
if (state.streaming) chars += messageChars(state.streaming);
220+
return Math.round(chars / CHARS_PER_TOKEN);
221+
}
222+
223+
function messageChars(message: AgentMessage): number {
224+
if (typeof message.content === "string") return message.content.length;
225+
if (!Array.isArray(message.content)) return 0;
226+
let total = 0;
227+
for (const block of message.content) {
228+
if (block.type === "text") total += block.text.length;
229+
else if (block.type === "thinking") total += block.thinking.length;
230+
else if (block.type === "toolCall") {
231+
total += block.name.length;
232+
total += JSON.stringify(block.arguments ?? {}).length;
233+
}
234+
}
235+
return total;
180236
}
181237

182238
/** Sum the visible text length of all text/thinking blocks in the live streaming message. */

0 commit comments

Comments
 (0)