From 29c5868a7cbbd3dd9220445c33223f75e6238403 Mon Sep 17 00:00:00 2001 From: nonlooped Date: Fri, 31 Jul 2026 11:33:50 +0300 Subject: [PATCH 1/2] feat(desktop): add context window inspector --- apps/desktop/src/main/ipc.ts | 9 + apps/desktop/src/main/pi/client.ts | 8 + apps/desktop/src/main/pi/host/entry.ts | 88 +++++++- .../src/renderer/components/Composer.tsx | 196 ++++++++++++++---- apps/desktop/src/shared/pi-types.ts | 16 ++ apps/desktop/src/shared/rpc-schema.ts | 5 + apps/desktop/src/shared/tui-frames.test.ts | 23 +- apps/desktop/src/shared/tui-frames.ts | 21 ++ 8 files changed, 326 insertions(+), 40 deletions(-) diff --git a/apps/desktop/src/main/ipc.ts b/apps/desktop/src/main/ipc.ts index 2898140..e9293ae 100644 --- a/apps/desktop/src/main/ipc.ts +++ b/apps/desktop/src/main/ipc.ts @@ -662,6 +662,15 @@ const handlers: HandlerMap = { } }, + getContextInspector: async ({ projectDir, sessionFile }) => { + try { + const pi = await bindPi(projectDir, sessionFile); + 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..1a612b9 100644 --- a/apps/desktop/src/main/pi/host/entry.ts +++ b/apps/desktop/src/main/pi/host/entry.ts @@ -1,9 +1,11 @@ import { PassThrough } from "node:stream"; -import { AgentSession, main, ModelRuntime } from "@earendil-works/pi-coding-agent"; +import { AgentSession, estimateTokens, formatSkillsForPrompt, main, ModelRuntime } from "@earendil-works/pi-coding-agent"; +import type { Skill } from "@earendil-works/pi-coding-agent"; 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 +93,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 +108,86 @@ function handleClientFrame(frame: TuiClientFrame): void { internals?.handle(frame); } +function textTokens(text: string): number { + return estimateTokens({ role: "user", content: [{ type: "text", text }], timestamp: 0 }); +} + +async function respondContextInspector(requestId: string): Promise { + try { + if (!currentSession) throw new Error("No active Pi session"); + const session = currentSession; + // These are the prompt inputs Pi assembled for this session. They remain + // private in Pi's API, so read them defensively and keep the inspector + // unavailable rather than presenting a made-up breakdown after an upgrade. + const options = (session as unknown as { + _baseSystemPromptOptions?: { + contextFiles?: { path: string; content: string }[]; + skills?: Skill[]; + }; + })._baseSystemPromptOptions; + if (!options) throw new Error("Pi has not prepared this session's prompt yet"); + + const contextFiles = options.contextFiles ?? []; + const skills = options.skills ?? []; + const contextFileTokens = contextFiles.map((file) => ({ path: file.path, tokens: textTokens(file.content) })); + const skillsPrompt = formatSkillsForPrompt(skills); + const skillTokens = skills.map((skill) => ({ + name: skill.name, + tokens: textTokens(formatSkillsForPrompt([skill])), + })); + const activeNames = new Set(session.getActiveToolNames()); + const tools = session.getAllTools() + .filter((tool) => activeNames.has(tool.name)) + .map((tool) => ({ name: tool.name, tokens: textTokens(JSON.stringify(tool.parameters)) })); + const contextTokens = contextFileTokens.reduce((sum, file) => sum + file.tokens, 0); + const skillTotal = textTokens(skillsPrompt); + const systemTokens = Math.max(0, textTokens(session.systemPrompt) - contextTokens - skillTotal); + const historyTokens = session.messages.reduce((sum, message) => sum + estimateTokens(message), 0); + // Pi does not export this helper, but it is the same internal module its + // session uses to decide the boundary. Resolving it beside the published + // entry mirrors the existing host integrations above without duplicating + // compaction logic in NativePi. + const index = import.meta.resolve("@earendil-works/pi-coding-agent"); + const compactionModule = (await import(new URL("./core/compaction/index.js", index).href)) as { + prepareCompaction?: ( + entries: unknown[], + settings: { enabled: boolean; reserveTokens: number; keepRecentTokens: number }, + ) => { + messagesToSummarize: Parameters[0][]; + turnPrefixMessages: Parameters[0][]; + settings: { keepRecentTokens: number }; + } | undefined; + }; + if (!compactionModule.prepareCompaction) throw new Error("This Pi version cannot inspect the next compaction"); + const plan = compactionModule.prepareCompaction(session.sessionManager.getBranch(), session.settingsManager.getCompactionSettings()); + const compacted = plan ? [...plan.messagesToSummarize, ...plan.turnPrefixMessages] : []; + const usage = session.getContextUsage(); + const inspector: ContextInspector = { + usedTokens: usage?.tokens ?? null, + contextWindow: usage?.contextWindow ?? session.model?.contextWindow ?? 0, + categories: [ + { kind: "system", tokens: systemTokens }, + { kind: "context", tokens: contextTokens, count: contextFiles.length }, + { kind: "skills", tokens: skillTotal, count: skillTokens.length }, + { kind: "tools", tokens: tools.reduce((sum, tool) => sum + tool.tokens, 0), count: tools.length }, + { kind: "history", tokens: historyTokens, count: session.messages.length }, + ], + contextFiles: contextFileTokens, + skills: skillTokens, + tools, + compaction: plan ? { + tokens: compacted.reduce((sum, message) => sum + estimateTokens(message), 0), + messages: plan.messagesToSummarize.length, + turnPrefixMessages: plan.turnPrefixMessages.length, + keepRecentTokens: plan.settings.keepRecentTokens, + } : undefined, + }; + 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 2f56a4a..82931b0 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"; @@ -624,8 +625,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,6 +648,10 @@ function ContextWindow() { const total = model?.contextWindow ?? 0; const percent = total ? Math.min(100, Math.round((used / total) * 100)) : 0; + const inspection = useRequest( + async () => (open && projectDir && sessionFile ? await rpc.request.getContextInspector({ projectDir, sessionFile }) : null), + [open, projectDir, sessionFile], + ); // 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 — @@ -667,12 +673,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, @@ -684,37 +685,156 @@ 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. - */} -