diff --git a/apps/desktop/src/main/ipc.ts b/apps/desktop/src/main/ipc.ts index 3df36b1..52d465c 100644 --- a/apps/desktop/src/main/ipc.ts +++ b/apps/desktop/src/main/ipc.ts @@ -244,11 +244,9 @@ function ensurePi(projectDir: string, sessionFile?: string, fresh = false): Prom return startup; } -async function bindPi(projectDir: string, sessionFile: string): Promise { +async function bindPi(projectDir: string, sessionFile: string, mayWrite = true): Promise { const pi = await ensurePi(projectDir, sessionFile); - // Everything reached through bindPi may write the session file (rename, fork, - // clone, compact), so claim the write before it happens. - markBusy(sessionFile, Date.now() + SETTLE_GRACE_MS); + if (mayWrite) markBusy(sessionFile, Date.now() + SETTLE_GRACE_MS); return pi; } @@ -682,6 +680,15 @@ const handlers: HandlerMap = { } }, + getContextInspector: async ({ projectDir, sessionFile }) => { + try { + const pi = sessionFile ? await bindPi(projectDir, sessionFile, false) : await ensurePi(projectDir); + return { inspector: await pi.getContextInspector() }; + } catch (err) { + return { error: errorMessage(err) }; + } + }, + compact: async ({ projectDir, sessionFile }) => { try { const pi = await bindPi(projectDir, sessionFile); diff --git a/apps/desktop/src/main/pi/client.ts b/apps/desktop/src/main/pi/client.ts index bd22ad4..a0d6bfa 100644 --- a/apps/desktop/src/main/pi/client.ts +++ b/apps/desktop/src/main/pi/client.ts @@ -2,6 +2,8 @@ import { fileURLToPath } from "node:url"; import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process"; import { isTuiFrameType, tuiHostFrameSchema, type TuiClientFrame, type TuiHostFrame } from "../../shared/tui-frames.ts"; import type { AuthProviderInfo } from "../../shared/rpc-schema.ts"; +import type { ContextInspector } from "../../shared/pi-types.ts"; +import { contextInspectorSchema } from "../../shared/tui-frames.ts"; import { drainLines, serializeCommand, type PiCommand, type PiMessage } from "./protocol.ts"; /** @@ -176,6 +178,12 @@ export class PiProcess { return this.frameRequest((requestId) => ({ type: "nativepi_tui_get_providers", requestId })); } + getContextInspector(): Promise { + return this.frameRequest((requestId) => ({ type: "nativepi_tui_get_context_inspector", requestId })).then( + (data) => contextInspectorSchema.parse(data), + ); + } + loginProvider(providerId: string, authType: "api_key" | "oauth"): Promise { return this.frameRequest( (requestId) => ({ type: "nativepi_tui_login", requestId, providerId, authType }), diff --git a/apps/desktop/src/main/pi/host/entry.ts b/apps/desktop/src/main/pi/host/entry.ts index f0ce7d0..d04d0f4 100644 --- a/apps/desktop/src/main/pi/host/entry.ts +++ b/apps/desktop/src/main/pi/host/entry.ts @@ -4,6 +4,7 @@ import type { ExtensionUIContext } from "@earendil-works/pi-coding-agent"; import { isTuiFrameType, type TuiClientFrame, type TuiHostFrame } from "../../../shared/tui-frames.ts"; import { toNotice, toPromptRequest } from "../../../shared/providerAuth.ts"; import { shapeProviders } from "../../../shared/providerShape.ts"; +import type { ContextInspector } from "../../../shared/pi-types.ts"; import { hostInternals, withTerminalUi, type HostInternals } from "./uiContext.ts"; /** @@ -91,6 +92,10 @@ function handleClientFrame(frame: TuiClientFrame): void { void respondProviders(frame.requestId); return; } + if (frame.type === "nativepi_tui_get_context_inspector") { + void respondContextInspector(frame.requestId); + return; + } if (frame.type === "nativepi_tui_login") { void respondLogin(frame.requestId, frame.providerId, frame.authType); return; @@ -102,6 +107,21 @@ function handleClientFrame(frame: TuiClientFrame): void { internals?.handle(frame); } +async function respondContextInspector(requestId: string): Promise { + try { + if (!currentSession) throw new Error("No active Pi session"); + const session = currentSession; + const usage = session.getContextUsage(); + const inspector: ContextInspector = { + usedTokens: usage?.tokens ?? null, + contextWindow: usage?.contextWindow ?? session.model?.contextWindow ?? 0, + }; + send({ type: "nativepi_tui_reply", requestId, data: inspector }); + } catch (err) { + send({ type: "nativepi_tui_reply", requestId, error: err instanceof Error ? err.message : String(err) }); + } +} + async function respondProviders(requestId: string): Promise { try { const providers = await shapeProviders(await providerRuntime()); diff --git a/apps/desktop/src/renderer/components/Composer.tsx b/apps/desktop/src/renderer/components/Composer.tsx index 701bba5..b99e76c 100644 --- a/apps/desktop/src/renderer/components/Composer.tsx +++ b/apps/desktop/src/renderer/components/Composer.tsx @@ -1,8 +1,8 @@ import { ArrowBendUpRightIcon, CaretDownIcon, CheckIcon, GitBranchIcon, PaperPlaneRightIcon, PlusIcon, TreeStructureIcon } from "../../shared/icons.ts" import { CircleNotchIcon } from "@phosphor-icons/react/CircleNotch"; import { PaperclipIcon } from "@phosphor-icons/react/Paperclip"; -import { useEffect, useId, useMemo, useRef, useState } from "react"; -import type { AssistantMessage } from "../../shared/pi-types.ts"; +import { useEffect, useMemo, useRef, useState } from "react"; +import type { AssistantMessage, ContextInspector as ContextInspectorData } from "../../shared/pi-types.ts"; import { draftKeyFor } from "../../shared/messages.ts"; import { ACCEPTED_IMAGE_TYPES } from "../lib/attachments.ts"; import { classifyDrop, draggingFiles, mentionPath } from "../lib/drops.ts"; @@ -18,6 +18,7 @@ import { Button } from "@/components/ui/button.tsx"; import { Input } from "@/components/ui/input.tsx"; import { Menu, MenuItem, MenuPopup, MenuTrigger } from "@/components/ui/menu.tsx"; import { Kbd } from "@/components/ui/kbd.tsx"; +import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle } from "@/components/ui/dialog.tsx"; import { SCROLLBAR_GUTTER_OFFSET, cn } from "@/lib/utils.ts"; import ComposerAttachments from "./ComposerAttachments.tsx"; import ComposerInput from "./ComposerInput.tsx"; @@ -625,8 +626,9 @@ function ContextWindow() { const entries = useAppStore((s) => activeConversation(s).entries); const streaming = useAppStore((s) => activeConversation(s).streaming); const model = useAppStore((s) => s.model); - const [pinned, setPinned] = useState(false); - const panelId = useId(); + const projectDir = useAppStore((s) => s.activeProjectPath); + const sessionFile = useAppStore((s) => s.activeSessionFile); + const [open, setOpen] = useState(false); // Walking the whole transcript backwards on every keystroke in the composer is // what this used to do; the answer only changes when the transcript does. @@ -646,7 +648,13 @@ function ContextWindow() { }, [entries, streaming]); const total = model?.contextWindow ?? 0; - const percent = total ? Math.min(100, Math.round((used / total) * 100)) : 0; + const inspection = useRequest( + async () => (open && projectDir ? await rpc.request.getContextInspector({ projectDir, sessionFile: sessionFile ?? undefined }) : null), + [open, projectDir, sessionFile, entries, streaming], + ); + const inspectedUsed = inspection.data?.inspector?.usedTokens ?? used; + const inspectedTotal = inspection.data?.inspector?.contextWindow || total; + const percent = inspectedTotal ? Math.min(100, Math.round((inspectedUsed / inspectedTotal) * 100)) : 0; // A ring with nothing behind it is decoration. Until Pi reports a window for // this model there is no reading to give, so the control is not interactive — @@ -668,12 +676,7 @@ function ContextWindow() { // The figure lives in the accessible name, so the ring is not a // visual-only readout for anyone who can't see it. aria-label={`Context window ${percent}% used, ${formatTokens(used)} of ${formatTokens(total)} tokens`} - aria-expanded={pinned} - aria-controls={panelId} - onClick={() => setPinned((current) => !current)} - onKeyDown={(event) => { - if (event.key === "Escape") setPinned(false); - }} + onClick={() => setOpen(true)} className={cn( "flex h-8 items-center gap-1.5 rounded-lg px-1.5 outline-none hover:bg-muted focus-visible:ring-2 focus-visible:ring-ring", tone, @@ -685,37 +688,66 @@ function ContextWindow() { {tight ? : null} - {/* - Not a live region: this panel is permanently mounted and its numbers - change on every token, so role="status" made screen readers narrate - token counts continuously. It is not aria-hidden either — the button - above claims to expand it, and it holds the only statement anywhere in - the app that Pi compacts before the model limit. - */} - ); } +function ContextInspector({ + open, + onOpenChange, + loading, + inspector, + error, + used, + total, + percent, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; + loading: boolean; + inspector?: ContextInspectorData; + error?: string; + used: number; + total: number; + percent: number; +}) { + return ( + + + + Context window + + Pi’s latest context-window measurement for this conversation. + + + +
+
+

{percent}% used

+

+ {formatTokens(used)} / {formatTokens(total)} tokens +

+
+
+ + {loading ?

Reading the active Pi session…

: null} + {error ?

Could not inspect this context: {error}

: null} + {inspector?.usedTokens === null ?

Pi has not received a context measurement from this model yet.

: null} +
+
+ ); +} + diff --git a/apps/desktop/src/shared/pi-types.ts b/apps/desktop/src/shared/pi-types.ts index e01d3dd..8edf83c 100644 --- a/apps/desktop/src/shared/pi-types.ts +++ b/apps/desktop/src/shared/pi-types.ts @@ -166,6 +166,12 @@ export interface SessionStats { cost: number; } +/** Pi's latest read-only context-window measurement. */ +export interface ContextInspector { + usedTokens: number | null; + contextWindow: number; +} + export interface SessionSummary { path: string; id: string; diff --git a/apps/desktop/src/shared/rpc-schema.ts b/apps/desktop/src/shared/rpc-schema.ts index 7c46001..57f3464 100644 --- a/apps/desktop/src/shared/rpc-schema.ts +++ b/apps/desktop/src/shared/rpc-schema.ts @@ -16,6 +16,7 @@ import type { ResolvedExtension, RpcSessionState, SessionStats, + ContextInspector, SessionSearchResult, SessionSummary, SessionTreeNode, @@ -365,6 +366,10 @@ export type HostRequests = { params: { projectDir: string; sessionFile: string }; response: { stats?: SessionStats; error?: string }; }; + getContextInspector: { + params: { projectDir: string; sessionFile?: string }; + response: { inspector?: ContextInspector; error?: string }; + }; compact: { params: { projectDir: string; sessionFile: string }; response: { ok: boolean; error?: string }; diff --git a/apps/desktop/src/shared/tui-frames.test.ts b/apps/desktop/src/shared/tui-frames.test.ts index f9f88ff..6e6b247 100644 --- a/apps/desktop/src/shared/tui-frames.test.ts +++ b/apps/desktop/src/shared/tui-frames.test.ts @@ -1,5 +1,5 @@ import { expect, test } from "bun:test"; -import { isTuiFrameType, tuiClientFrameSchema, tuiHostFrameSchema } from "./tui-frames.ts"; +import { contextInspectorSchema, isTuiFrameType, tuiClientFrameSchema, tuiHostFrameSchema } from "./tui-frames.ts"; /** * What keeps the side channel and Pi's protocol apart, and what keeps extension @@ -51,3 +51,13 @@ test("a resize from the window is bounded", () => { tuiClientFrameSchema.safeParse({ type: "nativepi_tui_resize", surfaceId: "s1", cols: 0, rows: 24 }).success, ).toBe(false); }); + +test("the context inspector accepts only a non-negative Pi measurement", () => { + const inspector = { + usedTokens: 42, + contextWindow: 200_000, + }; + + expect(contextInspectorSchema.safeParse(inspector).success).toBe(true); + expect(contextInspectorSchema.safeParse({ ...inspector, usedTokens: -1 }).success).toBe(false); +}); diff --git a/apps/desktop/src/shared/tui-frames.ts b/apps/desktop/src/shared/tui-frames.ts index 55ef623..3b32e23 100644 --- a/apps/desktop/src/shared/tui-frames.ts +++ b/apps/desktop/src/shared/tui-frames.ts @@ -1,4 +1,5 @@ import { z } from "zod"; +import type { ContextInspector } from "./pi-types.ts"; const authPromptSchema = z.discriminatedUnion("kind", [ z.object({ kind: z.literal("text"), message: z.string(), placeholder: z.string().optional() }), @@ -104,6 +105,11 @@ const completionItemSchema = z.object({ description: z.string().max(500).optional(), }); +export const contextInspectorSchema = z.object({ + usedTokens: z.number().int().nonnegative().nullable(), + contextWindow: z.number().int().nonnegative(), +}) satisfies z.ZodType; + export type TuiCompletionItem = z.infer; /** An extension autocomplete provider's answer, as `AutocompleteSuggestions`. */ @@ -206,6 +212,7 @@ export const tuiClientFrameSchema = z.discriminatedUnion("type", [ * extensions need this round trip to reach the picker and Settings. */ z.object({ type: z.literal("nativepi_tui_get_providers"), requestId: z.string().min(1).max(64) }), + z.object({ type: z.literal("nativepi_tui_get_context_inspector"), requestId: z.string().min(1).max(64) }), z.object({ type: z.literal("nativepi_tui_login"), requestId: z.string().min(1).max(64),