fix(pi): report authoritative context usage (#1106)

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', () => ({
+vi.mock('./piEventConverter', async (importOriginal) => {
+    const actual = await importOriginal<typeof import('./piEventConverter')>();
+    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<string, unknown> {
+        return (transport.send as ReturnType<typeof vi.fn>).mock.calls[index][0] as Record<string, unknown>;
+    }
+
     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', () => {

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);
+        // 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<void> {
+    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;
         }

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,18 +148,66 @@ 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', () => {
@@ -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'
         });

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.

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 };
+}

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)