From 4a065e3d8948d6e8111203b93ee87354df72a327 Mon Sep 17 00:00:00 2001 From: oldschoola Date: Sun, 19 Jul 2026 13:03:38 -0700 Subject: [PATCH] fix(status-line): aggregate vibe worker tok/s into the status-line badge In /vibe mode the director is often idle while workers stream, so the status-line tok/s badge showed a stale/zero rate even while parallel workers were actively generating tokens. The badge now aggregates the main session's live tok/s with every live vibe worker's tok/s, and falls back to the main session's own cached rate when no workers are streaming. - Extract calculateTokensPerSecond to utils/token-rate.ts (neutral location) so vibe/runtime.ts can depend on it without the render layer depending on the heavy vibe/task graph. - Add aggregateVibeWorkerTokensPerSecond to vibe/runtime.ts: sums each live worker's rate via the shared calculator, returns null when no workers are streaming so the main rate shines through. - StatusLineComponent takes the aggregator via an injected setVibeWorkerTokenRateProvider callback (wired in interactive-mode) keeping the render layer off the vibe/task dependency graph. - #getTokensPerSecond splits into #getMainSessionTokensPerSecond (preserves the sticky per-assistant-message cache) plus the worker-aggregation path, so the director-idle case no longer short-circuits to null. --- packages/coding-agent/CHANGELOG.md | 4 + .../modes/components/status-line/component.ts | 56 ++++++++++++- .../src/modes/interactive-mode.ts | 9 +- .../status-line => utils}/token-rate.ts | 6 ++ .../src/vibe/__tests__/token-rate.test.ts | 83 +++++++++++++++++++ packages/coding-agent/src/vibe/runtime.ts | 50 +++++++++++ .../test/status-line-token-rate.test.ts | 2 +- 7 files changed, 207 insertions(+), 3 deletions(-) rename packages/coding-agent/src/{modes/components/status-line => utils}/token-rate.ts (85%) create mode 100644 packages/coding-agent/src/vibe/__tests__/token-rate.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 6f32dc885..8c6b6aa93 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed the status-line `tok/s` badge ignoring vibe worker sessions: in `/vibe` mode the director is often idle while workers stream, so the badge showed a stale/zero rate while parallel work was actively generating tokens. The rate now aggregates the main session's live tok/s with every live vibe worker's tok/s, and falls back to the main session's own cached rate when no workers are streaming. + ## [17.0.5] - 2026-07-18 ### Added diff --git a/packages/coding-agent/src/modes/components/status-line/component.ts b/packages/coding-agent/src/modes/components/status-line/component.ts index 96608efa5..6adc87835 100644 --- a/packages/coding-agent/src/modes/components/status-line/component.ts +++ b/packages/coding-agent/src/modes/components/status-line/component.ts @@ -11,13 +11,13 @@ import { limitMatchesActiveAccount } from "../../../slash-commands/helpers/activ import { type ActiveRepoContext, resolveActiveRepoContextSync } from "../../../utils/active-repo-context"; import * as git from "../../../utils/git"; import { getSessionAccentAnsi, getSessionAccentHex } from "../../../utils/session-color"; +import { calculateTokensPerSecond } from "../../../utils/token-rate"; import { sanitizeStatusText } from "../../shared"; import { theme } from "../../theme/theme"; import { canReuseCachedPr, createPrCacheContext, isSamePrCacheContext, type PrCacheContext } from "./git-utils"; import { getPreset } from "./presets"; import { renderSegment, type SegmentContext } from "./segments"; import { getSeparator } from "./separators"; -import { calculateTokensPerSecond } from "./token-rate"; import type { CollabStatus, EffectiveStatusLineSettings, @@ -276,6 +276,13 @@ export class StatusLineComponent implements Component { #loopModeStatus: SegmentContext["loopMode"] = null; #goalModeStatus: { enabled: boolean; paused: boolean } | null = null; #vibeModeStatus: { enabled: boolean } | null = null; + /** + * Injected aggregator that returns the aggregate tok/s of this session's + * live vibe worker sessions, or null when no workers are streaming. Kept as + * a callback so the render layer doesn't import the heavy vibe/task + * dependency graph; interactive-mode wires it to VibeSessionRegistry. + */ + #vibeWorkerTokenRate: (() => number | null) | null = null; #collabStatus: CollabStatus | null = null; #focusedAgentId: string | undefined; #activeRepoCache: ActiveRepoCache | undefined; @@ -508,6 +515,17 @@ export class StatusLineComponent implements Component { this.#vibeModeStatus = status ?? null; } + /** + * Inject the aggregator that returns the aggregate tok/s of this session's + * live vibe worker sessions (null when no workers are streaming). Wired by + * interactive-mode, which owns the VibeSessionRegistry coupling, so the + * render layer stays off the heavy vibe/task dependency graph. Pass + * `undefined` to clear. + */ + setVibeWorkerTokenRateProvider(provider: (() => number | null) | undefined): void { + this.#vibeWorkerTokenRate = provider ?? null; + } + setCollabStatus(status: CollabStatus | null): void { this.#collabStatus = status; } @@ -747,6 +765,31 @@ export class StatusLineComponent implements Component { } #getTokensPerSecond(): number | null { + // Aggregate tok/s across the main session AND every live vibe worker. + // In vibe mode the director is often idle while workers stream, so the + // main session's own rate alone would show a stale/zero value while + // parallel work is actively generating tokens. + const workerRate = this.#getVibeWorkerTokensPerSecond(); + if (workerRate !== null) { + // At least one worker is streaming — add the director's live rate + // only when it is itself streaming (a finalized last-turn rate would + // double-count and overstate throughput). + const mainRate = this.session.isStreaming ? calculateTokensPerSecond(this.session.state.messages, true) : 0; + return (mainRate ?? 0) + workerRate; + } + + // No workers streaming — fall back to the main session's own rate with + // its sticky per-assistant-message cache so the badge doesn't flicker + // off in the brief gap between stream end and the finalized message. + return this.#getMainSessionTokensPerSecond(); + } + + /** + * Main session's tok/s with sticky caching keyed on the last assistant + * message timestamp. Preserves the pre-aggregation behavior when no vibe + * workers are active. + */ + #getMainSessionTokensPerSecond(): number | null { let lastAssistantTimestamp: number | null = null; for (let i = this.session.state.messages.length - 1; i >= 0; i--) { const message = this.session.state.messages[i]; @@ -776,6 +819,17 @@ export class StatusLineComponent implements Component { return null; } + /** + * Aggregate tok/s across every live vibe worker session owned by this + * session. Returns null when no workers are streaming (so the main + * session's own rate shines through unchanged). The aggregation itself is + * injected via {@link setVibeWorkerTokenRateProvider} to keep this render + * layer off the heavy vibe/task dependency graph. + */ + #getVibeWorkerTokensPerSecond(): number | null { + return this.#vibeWorkerTokenRate?.() ?? null; + } + #getUsageContextKey(session: AgentSession): string { const activeProvider = session.state.model?.provider ?? session.model?.provider ?? ""; if (!activeProvider) return ""; diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index 9d3a2a59a..d0e5f1279 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -138,7 +138,7 @@ import { getEditorCommand, openInEditor } from "../utils/external-editor"; import { getSessionAccentAnsi, getSessionAccentHex } from "../utils/session-color"; import { messageHasDisplayableThinking } from "../utils/thinking-display"; import { popTerminalTitle, pushTerminalTitle, setSessionTerminalTitle } from "../utils/title-generator"; -import { VibeSessionRegistry } from "../vibe/runtime"; +import { aggregateVibeWorkerTokensPerSecond, VibeSessionRegistry } from "../vibe/runtime"; import type { AssistantMessageComponent } from "./components/assistant-message"; import type { BashExecutionComponent } from "./components/bash-execution"; import { ChatBlock, type ChatBlockHost } from "./components/chat-block"; @@ -730,6 +730,13 @@ export class InteractiveMode implements InteractiveModeContext { this.editorContainer.addChild(this.editor); this.statusLine = new StatusLineComponent(session); this.statusLine.setAutoCompactEnabled(session.autoCompactionEnabled); + // Vibe worker tok/s aggregator — keeps the status-line render layer off + // the heavy vibe/task dependency graph. The director is often idle while + // workers stream, so without this the tok/s badge would show a stale + // value while parallel work is actively generating tokens. + this.statusLine.setVibeWorkerTokenRateProvider(() => + aggregateVibeWorkerTokensPerSecond(this.session.getAgentId() ?? MAIN_AGENT_ID), + ); // Lazy provider — the top border rebuild coalesces to at most one // invocation per painted frame instead of firing on every session event // (#4145). The TUI throttles renders at ~30fps, so a long-running eval diff --git a/packages/coding-agent/src/modes/components/status-line/token-rate.ts b/packages/coding-agent/src/utils/token-rate.ts similarity index 85% rename from packages/coding-agent/src/modes/components/status-line/token-rate.ts rename to packages/coding-agent/src/utils/token-rate.ts index 5b35edfa8..a0e5fab49 100644 --- a/packages/coding-agent/src/modes/components/status-line/token-rate.ts +++ b/packages/coding-agent/src/utils/token-rate.ts @@ -1,3 +1,9 @@ +/** + * Token-throughput calculator shared by the status line (main session tok/s + * badge) and the vibe worker aggregation ({@link aggregateVibeWorkerTokensPerSecond}). + * Lives in `utils/` so neither the render layer nor the vibe runtime has to + * depend on the other for a pure arithmetic helper. + */ const MIN_DURATION_MS = 100; type AssistantUsage = { diff --git a/packages/coding-agent/src/vibe/__tests__/token-rate.test.ts b/packages/coding-agent/src/vibe/__tests__/token-rate.test.ts new file mode 100644 index 000000000..b47822ef6 --- /dev/null +++ b/packages/coding-agent/src/vibe/__tests__/token-rate.test.ts @@ -0,0 +1,83 @@ +/** + * Contracts: vibe worker tok/s aggregation. + * + * 1. Returns null when the owner has no vibe worker sessions. + * 2. Sums the tok/s of every live worker the owner has registered, reading + * each worker's last assistant message through the same calculator the + * main status line uses. + * 3. Returns null when workers exist but none are streaming (no live rate to + * aggregate) — so the caller's own rate shines through unchanged. + * 4. Skips workers whose AgentRegistry session is detached (parked/reviving), + * so a stale roster entry can't contribute a phantom zero. + */ +import { afterEach, describe, expect, it } from "bun:test"; +import { AgentRegistry, MAIN_AGENT_ID } from "../../registry/agent-registry"; +import type { AgentSession } from "../../session/agent-session"; +import { aggregateVibeWorkerTokensPerSecond, VibeSessionRegistry } from "../runtime"; + +const OWNER = "test-owner"; + +/** Minimal fake AgentSession: just the messages + isStreaming the aggregator reads. */ +function fakeSession(messages: unknown[], isStreaming: boolean): AgentSession { + return { state: { messages }, isStreaming } as unknown as AgentSession; +} + +/** A finalized assistant message with a known duration → deterministic tok/s. */ +function assistantMessage(output: number, durationMs: number, timestamp = 1000) { + return { role: "assistant", timestamp, duration: durationMs, usage: { output } }; +} + +function registerWorker(id: string, session: AgentSession | null, ownerId = OWNER) { + VibeSessionRegistry.global().registerRecordForTests({ id, ownerId }); + if (session) { + AgentRegistry.global().register({ + id, + displayName: id, + kind: "sub", + session, + status: "running", + }); + } +} + +describe("aggregateVibeWorkerTokensPerSecond", () => { + afterEach(() => { + AgentRegistry.resetGlobalForTests(); + VibeSessionRegistry.resetGlobalForTests(); + }); + + it("returns null when the owner has no worker sessions", () => { + expect(aggregateVibeWorkerTokensPerSecond(OWNER)).toBeNull(); + }); + + it("sums tok/s across every live streaming worker", () => { + // 100 tokens in 1000ms → 100 tok/s. + registerWorker("w1", fakeSession([assistantMessage(100, 1000)], true)); + // 50 tokens in 500ms → 100 tok/s. + registerWorker("w2", fakeSession([assistantMessage(50, 500)], true)); + expect(aggregateVibeWorkerTokensPerSecond(OWNER)).toBe(200); + }); + + it("returns null when workers exist but none have a live rate", () => { + // Not streaming, no duration → calculateTokensPerSecond returns null. + registerWorker("w1", fakeSession([assistantMessage(100, 0)], false)); + expect(aggregateVibeWorkerTokensPerSecond(OWNER)).toBeNull(); + }); + + it("ignores workers whose AgentRegistry session is detached", () => { + registerWorker("w1", fakeSession([assistantMessage(100, 1000)], true)); + // w2 is in the vibe roster but has no live AgentRegistry session. + registerWorker("w2", null); + expect(aggregateVibeWorkerTokensPerSecond(OWNER)).toBe(100); + }); + + it("scopes to the requesting owner — other owners' workers don't count", () => { + registerWorker("w1", fakeSession([assistantMessage(100, 1000)], true), OWNER); + registerWorker("w2", fakeSession([assistantMessage(50, 500)], true), "other-owner"); + expect(aggregateVibeWorkerTokensPerSecond(OWNER)).toBe(100); + }); + + it("returns null for the main-agent owner id when no workers are registered", () => { + expect(aggregateVibeWorkerTokensPerSecond(MAIN_AGENT_ID)).toBeNull(); + }); +}); diff --git a/packages/coding-agent/src/vibe/runtime.ts b/packages/coding-agent/src/vibe/runtime.ts index 08610420c..f28d46844 100644 --- a/packages/coding-agent/src/vibe/runtime.ts +++ b/packages/coding-agent/src/vibe/runtime.ts @@ -33,6 +33,7 @@ import { type AgentDefinition, type AgentProgress, oneLineLabel, type SingleResu import type { ToolSession } from "../tools"; import { formatDuration } from "../tools/render-utils"; import { ToolError } from "../tools/tool-errors"; +import { calculateTokensPerSecond } from "../utils/token-rate"; /** The two worker CLI flavors the director drives. */ export type VibeCli = "fast" | "good"; @@ -207,6 +208,26 @@ export class VibeSessionRegistry { VibeSessionRegistry.#global = undefined; } + /** + * Insert a bare worker record without the spawn/job machinery. Test-only — + * lets {@link aggregateVibeWorkerTokensPerSecond} be exercised against a + * fake roster + AgentRegistry session without driving a real turn. + */ + registerRecordForTests(record: { id: string; cli?: VibeCli; ownerId: string; state?: VibeSessionState }): void { + this.#records.set(record.id, { + id: record.id, + cli: record.cli ?? "fast", + ownerId: record.ownerId, + agent: getBundledAgent("sonic")!, + state: record.state ?? "running", + createdAt: Date.now(), + lastActivityAt: Date.now(), + queue: [], + turnCount: 0, + killed: false, + }); + } + readonly #records = new Map(); #manager(session: ToolSession): AsyncJobManager { @@ -708,3 +729,32 @@ export class VibeSessionRegistry { return text; } } + +/** + * Aggregate tok/s across every live vibe worker session owned by `ownerId`. + * Returns null when no workers are streaming (so callers can fall back to + * their own rate unchanged). The director is often idle while workers stream, + * so without this aggregation the status-line tok/s badge would show a stale + * value while parallel work is actively generating tokens. + * + * Reads each worker's last assistant message via {@link calculateTokensPerSecond} + * — the same leaf calculator the main status line uses — so worker rates are + * computed identically to the main session's rate. + */ +export function aggregateVibeWorkerTokensPerSecond(ownerId: string): number | null { + const ids = VibeSessionRegistry.global().listIds(ownerId); + if (ids.length === 0) return null; + let total = 0; + let any = false; + const registry = AgentRegistry.global(); + for (const id of ids) { + const workerSession = registry.get(id)?.session; + if (!workerSession) continue; + const rate = calculateTokensPerSecond(workerSession.state.messages, workerSession.isStreaming); + if (rate !== null) { + total += rate; + any = true; + } + } + return any ? total : null; +} diff --git a/packages/coding-agent/test/status-line-token-rate.test.ts b/packages/coding-agent/test/status-line-token-rate.test.ts index 8de667d3b..7d891806a 100644 --- a/packages/coding-agent/test/status-line-token-rate.test.ts +++ b/packages/coding-agent/test/status-line-token-rate.test.ts @@ -2,9 +2,9 @@ import { beforeAll, describe, expect, it } from "bun:test"; import { stripVTControlCharacters } from "node:util"; import type { AssistantMessage } from "@oh-my-pi/pi-ai"; import { renderSegment } from "@oh-my-pi/pi-coding-agent/modes/components/status-line/segments"; -import { calculateTokensPerSecond } from "@oh-my-pi/pi-coding-agent/modes/components/status-line/token-rate"; import type { SegmentContext } from "@oh-my-pi/pi-coding-agent/modes/components/status-line/types"; import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; +import { calculateTokensPerSecond } from "@oh-my-pi/pi-coding-agent/utils/token-rate"; beforeAll(async () => { await initTheme();