From 735ccda1682044d0d6825a1d17065105cb6676d5 Mon Sep 17 00:00:00 2001 From: KorenKrita Date: Fri, 24 Jul 2026 10:56:03 +0800 Subject: [PATCH] fix(pi): report authoritative context usage (#1106) --- cli/src/pi/loop.test.ts | 169 +++++++++++++++++++++++++++- cli/src/pi/loop.ts | 60 +++++++++- cli/src/pi/piEventConverter.test.ts | 68 +++++++++-- cli/src/pi/piEventConverter.ts | 54 ++++++--- cli/src/pi/schemas.ts | 36 ++++++ cli/src/pi/types.ts | 8 +- 6 files changed, 355 insertions(+), 40 deletions(-) diff --git a/cli/src/pi/loop.test.ts b/cli/src/pi/loop.test.ts index e611a313..c7dbd7a7 100644 --- a/cli/src/pi/loop.test.ts +++ b/cli/src/pi/loop.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; -import { parsePiModels, parsePiCommands, sendPiRpcAndWait, wireTransportEvents } from './loop'; +import { parsePiModels, parsePiCommands, parsePiContextUsage, sendPiRpcAndWait, wireTransportEvents } from './loop'; import type { PiResponseEvent } from './types'; import { PiSession } from './session'; import { PiTransport } from './piTransport'; @@ -19,9 +19,13 @@ vi.mock('@/agent/messageConverter', () => ({ convertAgentMessage: vi.fn((msg) => msg), })); -vi.mock('./PiEventConverter', () => ({ - convertPiEvent: vi.fn(() => []), -})); +vi.mock('./piEventConverter', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + convertPiEvent: vi.fn(() => []), + }; +}); vi.mock('./piMessageAccumulator', () => { return { @@ -191,6 +195,26 @@ describe('parsePiCommands', () => { expect(parsePiCommands(data)).toEqual([{ name: 'cmd', source: 'skill' }]); }); }); +// --- parsePiContextUsage --- + +describe('parsePiContextUsage', () => { + it('parses Pi authoritative context usage', () => { + expect(parsePiContextUsage({ + contextUsage: { tokens: 101_035, contextWindow: 200_000, percent: 50.5 }, + })).toEqual({ tokens: 101_035, contextWindow: 200_000 }); + }); + + it('preserves Pi explicit unknown context after compaction', () => { + expect(parsePiContextUsage({ + contextUsage: { tokens: null, contextWindow: 200_000 }, + })).toBeNull(); + }); + + it('returns unavailable for missing or malformed tokens', () => { + expect(parsePiContextUsage({})).toBeUndefined(); + expect(parsePiContextUsage({ contextUsage: { tokens: '101035' } })).toBeUndefined(); + }); +}); // --- wireTransportEvents (integration) --- @@ -217,6 +241,10 @@ describe('wireTransportEvents', () => { handler!(event); } + function getSentCommand(transport: PiTransport, index = 0): Record { + return (transport.send as ReturnType).mock.calls[index][0] as Record; + } + it('handles get_state response — updates model, provider, thinkingLevel', () => { const transport = createMockTransport(); const pendingLocalIds: string[] = []; @@ -331,14 +359,143 @@ describe('wireTransportEvents', () => { expect(session.client.emitMessagesConsumed).toHaveBeenCalledWith(['prompt-1'], undefined); }); - it('handles turn_end — stops streaming', () => { + it('publishes authoritative context usage after turn_end stats resolve', async () => { const transport = createMockTransport(); wireTransportEvents(transport, session, []); session.piIsStreaming = true; - emitEvent({ type: 'turn_end' }); + emitEvent({ + type: 'turn_end', + message: { + usage: { input: 100, output: 200, cacheRead: 10, cacheWrite: 5, totalTokens: 315 }, + stopReason: 'stop', + }, + }); expect(session.piIsStreaming).toBe(false); + expect(session.client.sendAgentMessage).not.toHaveBeenCalled(); + const command = getSentCommand(transport); + expect(command).toMatchObject({ type: 'get_session_stats' }); + + emitEvent({ + type: 'response', + id: command.id, + command: 'get_session_stats', + success: true, + data: { contextUsage: { tokens: 342, contextWindow: 200_000 } }, + }); + + await vi.waitFor(() => { + expect(session.client.sendAgentMessage).toHaveBeenCalledWith({ + type: 'usage', + inputTokens: 100, + outputTokens: 200, + totalTokens: 315, + cacheReadTokens: 10, + contextTokens: 342, + contextWindow: 200_000, + }); + }); + }); + + it('silently falls back to turn totalTokens when stats are unsupported', async () => { + const transport = createMockTransport(); + wireTransportEvents(transport, session, []); + + emitEvent({ + type: 'turn_end', + message: { + usage: { input: 100, output: 200, cacheRead: 10, cacheWrite: 5, totalTokens: 315 }, + }, + }); + const command = getSentCommand(transport); + + emitEvent({ + type: 'response', + id: command.id, + command: 'get_session_stats', + success: false, + error: 'Unknown command', + }); + + await vi.waitFor(() => { + expect(session.client.sendAgentMessage).toHaveBeenCalledWith(expect.objectContaining({ + type: 'usage', + contextTokens: 315, + })); + }); + expect(session.client.sendSessionEvent).not.toHaveBeenCalled(); + }); + + it('falls back to turn totalTokens when stats time out', async () => { + vi.useFakeTimers(); + try { + const transport = createMockTransport(); + wireTransportEvents(transport, session, []); + + emitEvent({ + type: 'turn_end', + message: { + usage: { input: 100, output: 200, cacheRead: 10, cacheWrite: 5, totalTokens: 315 }, + }, + }); + + expect(session.client.sendAgentMessage).not.toHaveBeenCalled(); + await vi.advanceTimersByTimeAsync(1_000); + + expect(session.client.sendAgentMessage).toHaveBeenCalledWith(expect.objectContaining({ + type: 'usage', + contextTokens: 315, + })); + } finally { + vi.useRealTimers(); + } + }); + + it('discards a stats response from an older completed turn', async () => { + const transport = createMockTransport(); + wireTransportEvents(transport, session, []); + + emitEvent({ + type: 'turn_end', + message: { + usage: { input: 10, output: 20, cacheRead: 0, cacheWrite: 0, totalTokens: 30 }, + }, + }); + emitEvent({ + type: 'turn_end', + message: { + usage: { input: 40, output: 50, cacheRead: 0, cacheWrite: 0, totalTokens: 90 }, + }, + }); + + const olderCommand = getSentCommand(transport, 0); + const latestCommand = getSentCommand(transport, 1); + emitEvent({ + type: 'response', + id: latestCommand.id, + command: 'get_session_stats', + success: true, + data: { contextUsage: { tokens: 120, contextWindow: 200_000 } }, + }); + + await vi.waitFor(() => { + expect(session.client.sendAgentMessage).toHaveBeenCalledTimes(1); + }); + expect(session.client.sendAgentMessage).toHaveBeenLastCalledWith(expect.objectContaining({ + contextTokens: 120, + })); + + emitEvent({ + type: 'response', + id: olderCommand.id, + command: 'get_session_stats', + success: true, + data: { contextUsage: { tokens: 45, contextWindow: 200_000 } }, + }); + await Promise.resolve(); + + expect(session.client.sendAgentMessage).toHaveBeenCalledTimes(1); }); it('handles agent_end — stops streaming', () => { diff --git a/cli/src/pi/loop.ts b/cli/src/pi/loop.ts index 58841aea..9fb2fda0 100644 --- a/cli/src/pi/loop.ts +++ b/cli/src/pi/loop.ts @@ -1,14 +1,14 @@ import { logger } from '@/ui/logger'; import { convertAgentMessage } from '@/agent/messageConverter'; import { PiTransport } from './piTransport'; -import { convertPiEvent } from './piEventConverter'; +import { convertPiEvent, convertPiTurnUsage } from './piEventConverter'; import { PiMessageAccumulator } from './piMessageAccumulator'; -import { parsePiModels, parsePiCommands, PiResponseEventSchema, PiStateDataSchema, PiSetModelDataSchema } from './schemas'; -import type { PiResponseEvent, PiRpcCommand, PiThinkingLevel } from './types'; +import { parsePiModels, parsePiCommands, parsePiContextUsage, PiResponseEventSchema, PiStateDataSchema, PiSetModelDataSchema } from './schemas'; +import type { PiContextUsage, PiResponseEvent, PiRpcCommand, PiThinkingLevel, PiTurnEndEvent } from './types'; import type { PiSession } from './session'; // --- Response parsers: re-exported from schemas.ts --- -export { parsePiModels, parsePiCommands } from './schemas'; +export { parsePiModels, parsePiCommands, parsePiContextUsage } from './schemas'; // --- Pending RPC resolver --- // Instance-scoped: created once by wireTransportEvents, stored on PiSession. @@ -143,7 +143,12 @@ function handleResponse( const error = response.error ?? 'Unknown Pi error'; logger.debug(`[pi] RPC error for ${command}: ${error}`); resolvePendingRpc(resolver, response); - session.sendSessionEvent({ type: 'message', message: error }); + // get_session_stats is a best-effort compatibility probe. Older Pi + // versions may reject it, so fall back silently instead of surfacing an + // error event to the user on every completed turn. + if (command !== 'get_session_stats') { + session.sendSessionEvent({ type: 'message', message: error }); + } if (command === 'prompt' && pendingLocalIds.length > 0) { const oldestLocalId = pendingLocalIds.shift()!; session.emitMessagesConsumed([oldestLocalId], { clearQueuedThinkingGrace: true }); @@ -257,6 +262,43 @@ function handleResponse( } } +const PI_CONTEXT_USAGE_RPC_TIMEOUT_MS = 1_000; + +async function publishPiTurnUsage( + event: PiTurnEndEvent, + transport: PiTransport, + session: PiSession, + isLatestRequest: () => boolean, +): Promise { + let contextUsage: PiContextUsage | null | undefined; + try { + const stats = await sendPiRpcAndWait( + session, + transport, + { type: 'get_session_stats' }, + PI_CONTEXT_USAGE_RPC_TIMEOUT_MS, + ); + contextUsage = parsePiContextUsage(stats); + } catch (error) { + // Unsupported/failed stats RPC: convertPiTurnUsage falls back to the + // positive per-turn totalTokens value. The fallback is intentionally + // local to Pi so providers with different usage semantics are untouched. + logger.debug(`[pi] get_session_stats unavailable, using turn usage fallback: ${error instanceof Error ? error.message : String(error)}`); + contextUsage = undefined; + } + + // RPC responses can arrive after a newer turn has already completed. + // Publishing only the newest request prevents stale context values from + // overwriting a later turn's usage state. + if (!isLatestRequest()) return; + + const usageMessage = convertPiTurnUsage(event, contextUsage); + if (!usageMessage) return; + + const converted = convertAgentMessage(usageMessage); + if (converted) session.sendAgentMessage(converted); +} + // --- Wire transport events to session --- export function wireTransportEvents( @@ -266,6 +308,7 @@ export function wireTransportEvents( ): void { session.rpcResolver = new PiRpcResolver(); const assistantMessageAccumulator = new PiMessageAccumulator(); + let latestContextUsageRequest = 0; transport.onEvent((event) => { // Debug: log all event types to diagnose missing Pi output @@ -312,6 +355,13 @@ export function wireTransportEvents( } } else if (event.type === 'turn_end') { session.updateThinkingState(false); + const requestVersion = ++latestContextUsageRequest; + void publishPiTurnUsage( + event as PiTurnEndEvent, + transport, + session, + () => requestVersion === latestContextUsageRequest, + ); } else if (event.type === 'agent_end') { session.piIsStreaming = false; } diff --git a/cli/src/pi/piEventConverter.test.ts b/cli/src/pi/piEventConverter.test.ts index aaa0a5d0..75cf6200 100644 --- a/cli/src/pi/piEventConverter.test.ts +++ b/cli/src/pi/piEventConverter.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect } from 'vitest'; -import { convertPiEvent } from './piEventConverter'; +import { convertPiEvent, convertPiTurnUsage } from './piEventConverter'; import type { PiAgentEvent } from './types'; describe('convertPiEvent', () => { @@ -132,7 +132,7 @@ describe('convertPiEvent', () => { expect((result[0] as any).id).toBeUndefined(); }); - it('should convert turn_end to usage + turn_complete (2 messages)', () => { + it('should defer turn usage and convert only turn completion', () => { const result = convertPiEvent({ type: 'turn_end', message: { @@ -148,20 +148,68 @@ describe('convertPiEvent', () => { toolResults: [] }); - expect(result).toHaveLength(2); - expect(result[0]).toEqual({ + expect(result).toEqual([{ + type: 'turn_complete', + stopReason: 'stop' + }]); + }); + + it('should build usage from Pi authoritative context stats', () => { + const result = convertPiTurnUsage({ + type: 'turn_end', + message: { + usage: { input: 100, output: 200, cacheRead: 10, cacheWrite: 5, totalTokens: 315 } + } + }, { tokens: 342, contextWindow: 200_000 }); + + expect(result).toEqual({ type: 'usage', inputTokens: 100, outputTokens: 200, totalTokens: 315, - cacheReadTokens: 10 + cacheReadTokens: 10, + contextTokens: 342, + contextWindow: 200_000 }); - expect(result[1]).toEqual({ - type: 'turn_complete', - stopReason: 'stop' + }); + + it('should fall back to positive totalTokens when stats are unavailable', () => { + const result = convertPiTurnUsage({ + type: 'turn_end', + message: { + usage: { input: 100, output: 200, cacheRead: 10, cacheWrite: 5, totalTokens: 315 } + } + }, undefined); + + expect(result).toMatchObject({ + type: 'usage', + totalTokens: 315, + contextTokens: 315 }); }); + it('should preserve prior usage when Pi explicitly reports unknown context', () => { + const result = convertPiTurnUsage({ + type: 'turn_end', + message: { + usage: { input: 100, output: 200, cacheRead: 10, cacheWrite: 5, totalTokens: 315 } + } + }, null); + + expect(result).toBeNull(); + }); + + it('should skip all-zero error or aborted usage', () => { + const result = convertPiTurnUsage({ + type: 'turn_end', + message: { + usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0 } + } + }, { tokens: 342, contextWindow: 200_000 }); + + expect(result).toBeNull(); + }); + it('should convert turn_end with toolUse stopReason', () => { const result = convertPiEvent({ type: 'turn_end', @@ -172,8 +220,8 @@ describe('convertPiEvent', () => { toolResults: [] }); - expect(result).toHaveLength(2); - expect(result[1]).toEqual({ + expect(result).toHaveLength(1); + expect(result[0]).toEqual({ type: 'turn_complete', stopReason: 'toolUse' }); diff --git a/cli/src/pi/piEventConverter.ts b/cli/src/pi/piEventConverter.ts index 360b83a9..75ac0170 100644 --- a/cli/src/pi/piEventConverter.ts +++ b/cli/src/pi/piEventConverter.ts @@ -4,9 +4,42 @@ import type { PiAgentEvent, PiToolExecutionStartEvent, PiToolExecutionEndEvent, - PiTurnEndEvent + PiTurnEndEvent, + PiContextUsage, + PiUsage } from './types'; +function hasMeaningfulUsage(usage: PiUsage | undefined): usage is PiUsage { + return usage !== undefined + && Number.isFinite(usage.totalTokens) + && usage.totalTokens > 0; +} + +/** + * Builds the turn usage update after Pi's session stats request settles. + * + * undefined stats fall back to the turn's totalTokens for older Pi versions. + * null means Pi explicitly reported an unknown context size, so the previous + * valid HAPI usage state is preserved by not publishing an update. + */ +export function convertPiTurnUsage( + event: PiTurnEndEvent, + contextUsage: PiContextUsage | null | undefined, +): AgentMessage | null { + const usage = event.message?.usage; + if (!hasMeaningfulUsage(usage) || contextUsage === null) return null; + + return { + type: 'usage', + inputTokens: usage.input ?? 0, + outputTokens: usage.output ?? 0, + totalTokens: usage.totalTokens, + cacheReadTokens: usage.cacheRead, + contextTokens: contextUsage?.tokens ?? usage.totalTokens, + contextWindow: contextUsage?.contextWindow, + }; +} + /** * Converts Pi AgentEvent to HAPI AgentMessage array. * @@ -40,25 +73,10 @@ export function convertPiEvent(event: PiAgentEvent): AgentMessage[] { case 'turn_end': { const e = event as PiTurnEndEvent; - const messages: AgentMessage[] = []; - const usage = e.message?.usage; - - if (usage) { - messages.push({ - type: 'usage', - inputTokens: usage.input ?? 0, - outputTokens: usage.output ?? 0, - totalTokens: usage.totalTokens, - cacheReadTokens: usage.cacheRead - }); - } - - messages.push({ + return [{ type: 'turn_complete', stopReason: e.message?.stopReason ?? 'stop' - }); - - return messages; + }]; } // Lifecycle and other events — not converted to AgentMessage. diff --git a/cli/src/pi/schemas.ts b/cli/src/pi/schemas.ts index 358b4413..e588100b 100644 --- a/cli/src/pi/schemas.ts +++ b/cli/src/pi/schemas.ts @@ -12,6 +12,7 @@ import { z } from 'zod'; import { PI_THINKING_LEVELS } from '@hapi/protocol'; import type { PiModelSummary } from '@hapi/protocol/apiTypes'; +import type { PiContextUsage } from './types'; // ============================================================================ // 字段级容错 schema @@ -23,6 +24,17 @@ const asOptStr = z.unknown().optional().transform(v => typeof v === 'string' ? v /** 提取 number 值，非 number 或缺失返回 undefined */ const asOptNum = z.unknown().optional().transform(v => typeof v === 'number' ? v : undefined); +/** Extract a finite positive number, otherwise return undefined. */ +const asOptPositiveNum = z.unknown().optional().transform(v => + typeof v === 'number' && Number.isFinite(v) && v > 0 ? v : undefined, +); + +/** Context usage tokens may be null immediately after compaction. */ +const asContextTokens = z.unknown().optional().transform((v): number | null | undefined => { + if (v === null) return null; + return typeof v === 'number' && Number.isFinite(v) && v >= 0 ? v : undefined; +}); + /** 提取 boolean 值，非 boolean 或缺失返回 undefined */ const asOptBool = z.unknown().optional().transform(v => typeof v === 'boolean' ? v : undefined); @@ -135,6 +147,13 @@ const PiModelsResponseDataSchema = z.object({ .map(r => r.data), ); +const PiSessionStatsDataSchema = z.object({ + contextUsage: z.object({ + tokens: asContextTokens, + contextWindow: asOptPositiveNum, + }).passthrough().optional(), +}).passthrough(); + // ============================================================================ // Pi State (get_state response data) // ============================================================================ @@ -203,3 +222,20 @@ export function parsePiModels(data: unknown) { const result = PiModelsResponseDataSchema.safeParse(data) return result.success ? result.data : [] } + +/** + * Parse Pi's authoritative current context-window estimate. + * + * undefined: stats unavailable/malformed; callers may fall back to turn usage. + * null: Pi explicitly reports unknown (for example, immediately after compaction). + */ +export function parsePiContextUsage(data: unknown): PiContextUsage | null | undefined { + const result = PiSessionStatsDataSchema.safeParse(data); + if (!result.success || !result.data.contextUsage) return undefined; + + const { tokens, contextWindow } = result.data.contextUsage; + if (tokens === null) return null; + if (tokens === undefined) return undefined; + + return contextWindow === undefined ? { tokens } : { tokens, contextWindow }; +} diff --git a/cli/src/pi/types.ts b/cli/src/pi/types.ts index ded9607b..51107b40 100644 --- a/cli/src/pi/types.ts +++ b/cli/src/pi/types.ts @@ -37,6 +37,11 @@ export interface PiUsage { totalTokens: number; } +export interface PiContextUsage { + tokens: number; + contextWindow?: number; +} + // Individual event types for proper type narrowing export interface PiAgentStartEvent { type: 'agent_start' } export interface PiAgentEndEvent { type: 'agent_end'; messages: unknown[] } @@ -104,7 +109,8 @@ export type PiRpcCommand = | { type: 'set_model'; provider: string; modelId: string } | { type: 'get_available_models' } | { type: 'set_thinking_level'; level: PiThinkingLevel } - | { type: 'get_commands' }; + | { type: 'get_commands' } + | { type: 'get_session_stats' }; // ============================================================================ // Pi RPC Responses (stdout)