From 202c417347461e1644348834e841e629121902ad Mon Sep 17 00:00:00 2001 From: Slava Goltser Date: Wed, 8 Jul 2026 17:00:17 -0400 Subject: [PATCH] fix(pi): add get_session_stats fallback from get_state for old OMP binaries (#1886) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(pi): add get_session_stats fallback from get_state for old OMP binaries Older Oh My Pi binaries don't support the get_session_stats RPC command, leaving session usage stats blank in Paseo. This adds a compat layer in cli-runtime.ts where getSessionStats() tries get_session_stats first, then falls back to extracting context window usage from get_state. The fallback provides contextUsage.tokens and contextWindow even when the binary lacks full token/cost reporting — enough to keep the context meter working. Token counts and cost still require the newer RPC. Added three tests: - Fallback triggers when get_session_stats is unsupported (throws) - No fallback when get_session_stats returns valid data - Returns empty object when both commands fail * fix(pi): use nullish check instead of truthy for cost field in getSessionStats --- .../agent/providers/pi/cli-runtime.test.ts | 131 ++++++++++++++++-- .../server/agent/providers/pi/cli-runtime.ts | 30 +++- 2 files changed, 150 insertions(+), 11 deletions(-) diff --git a/packages/server/src/server/agent/providers/pi/cli-runtime.test.ts b/packages/server/src/server/agent/providers/pi/cli-runtime.test.ts index 56f1994ee..88ac3905c 100644 --- a/packages/server/src/server/agent/providers/pi/cli-runtime.test.ts +++ b/packages/server/src/server/agent/providers/pi/cli-runtime.test.ts @@ -61,16 +61,29 @@ function replyToCommands( const line = buffer.slice(0, newlineIndex); buffer = buffer.slice(newlineIndex + 1); const command = JSON.parse(line) as Record; - const result = handler(command); - child.stdout.write( - `${JSON.stringify({ - id: command.id, - type: "response", - command: command.type, - success: true, - data: result, - })}\n`, - ); + try { + const result = handler(command); + child.stdout.write( + `${JSON.stringify({ + id: command.id, + type: "response", + command: command.type, + success: true, + data: result, + })}\n`, + ); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + child.stdout.write( + `${JSON.stringify({ + id: command.id, + type: "response", + command: command.type, + success: false, + error: message, + })}\n`, + ); + } } }); } @@ -257,4 +270,102 @@ describe("PiCliRuntime", () => { expect(child.killedSignals).toContain("SIGTERM"); }); + + test("falls back to get_state when get_session_stats is unsupported", async () => { + const child = createPiChild(); + let commandSequence: string[] = []; + replyToCommands(child, (command) => { + commandSequence.push(String(command.type)); + if (command.type === "get_session_stats") { + // Simulate older OMP binary that doesn't support this RPC + throw new Error(`Unknown command: ${command.type}`); + } + // get_state returns contextUsage for fallback + return { + sessionId: "pi-session-1", + thinkingLevel: "medium" as const, + isStreaming: false, + isCompacting: false, + steeringMode: "one-at-a-time" as const, + followUpMode: "one-at-a-time" as const, + interruptMode: "immediate" as const, + messageCount: 0, + queuedMessageCount: 0, + todoPhases: [], + contextUsage: { tokens: 1_100, contextWindow: 200_000 }, + }; + }); + const session = await createRuntime(child).startSession({ cwd: "/workspace/project" }); + + const stats = await session.getSessionStats(); + + expect(stats.contextUsage).toEqual({ + tokens: 1_100, + contextWindow: 200_000, + }); + expect(commandSequence).toEqual(["get_session_stats", "get_state"]); + }); + + test("returns full stats from get_session_stats without falling back", async () => { + const child = createPiChild(); + let fallbackCalled = false; + replyToCommands(child, (command) => { + if (command.type === "get_state") { + fallbackCalled = true; + } + return { + tokens: { input: 500, output: 300, cacheRead: 100 }, + cost: 0.02, + contextUsage: { tokens: 800, contextWindow: 200_000 }, + }; + }); + const session = await createRuntime(child).startSession({ cwd: "/workspace/project" }); + + const stats = await session.getSessionStats(); + + expect(stats).toMatchObject({ + tokens: { input: 500, output: 300, cacheRead: 100 }, + cost: 0.02, + contextUsage: { tokens: 800, contextWindow: 200_000 }, + }); + // Should NOT have called get_state as a fallback + expect(fallbackCalled).toBe(false); + }); + + test("does not fall back when get_session_stats returns cost:0", async () => { + const child = createPiChild(); + let fallbackCalled = false; + replyToCommands(child, (command) => { + if (command.type === "get_state") { + fallbackCalled = true; + } + return { + tokens: { input: 200, output: 100 }, + cost: 0, + contextUsage: { tokens: 500, contextWindow: 200_000 }, + }; + }); + const session = await createRuntime(child).startSession({ cwd: "/workspace/project" }); + + const stats = await session.getSessionStats(); + + expect(stats.tokens).toEqual({ input: 200, output: 100 }); + expect(stats.cost).toBe(0); + expect(stats.contextUsage).toEqual({ tokens: 500, contextWindow: 200_000 }); + // Should NOT have called get_state as a fallback + expect(fallbackCalled).toBe(false); + }); + + test("returns empty object when both get_session_stats and get_state fail", async () => { + const child = createPiChild(); + replyToCommands(child, (command) => { + throw new Error(`Unknown command: ${command.type}`); + }); + const session = await createRuntime(child).startSession({ cwd: "/workspace/project" }); + + const stats = await session.getSessionStats(); + + // Neither RPC returned usable data — should resolve with empty object + expect(stats).toEqual({}); + }); }); diff --git a/packages/server/src/server/agent/providers/pi/cli-runtime.ts b/packages/server/src/server/agent/providers/pi/cli-runtime.ts index dd74fe87c..da702a25a 100644 --- a/packages/server/src/server/agent/providers/pi/cli-runtime.ts +++ b/packages/server/src/server/agent/providers/pi/cli-runtime.ts @@ -180,7 +180,35 @@ class PiCliRuntimeSession implements PiRuntimeSession { } async getSessionStats(): Promise { - return (await this.request({ type: "get_session_stats" })) as PiSessionStats; + // COMPAT(piGetStateFallback): added in v0.1.X — older Oh My Pi binaries + // lack the `get_session_stats` RPC command; fall back to extracting + // context window usage from `get_state`. Drop this gate when the floor + // daemon supports get_session_stats (added ~v0.1.97). + let stats: PiSessionStats | undefined; + try { + stats = (await this.request({ type: "get_session_stats" })) as PiSessionStats; + } catch { + // get_session_stats not supported by this binary — will try get_state below + } + if (stats?.tokens == null && stats?.cost == null && stats?.contextUsage == null) { + try { + const state = (await this.request({ type: "get_state" })) as Record; + const ctx = state.contextUsage as + | { tokens?: number | null; contextWindow?: number | null } + | undefined; + if (ctx) { + return { + contextUsage: { + tokens: typeof ctx.tokens === "number" ? ctx.tokens : undefined, + contextWindow: typeof ctx.contextWindow === "number" ? ctx.contextWindow : undefined, + }, + }; + } + } catch { + // get_state also failed — nothing we can do + } + } + return stats ?? {}; } async getCommands(): Promise {