Skip to content

Commit eced14e

Browse files
committed
Fixed formulat for token estimations
1 parent 96ffb0b commit eced14e

10 files changed

Lines changed: 187 additions & 108 deletions

File tree

packages/core/src/host/session-manager.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@ import type { PromptAttachment, PromptImage, SessionCommandInfo } from "../proto
1111
import {
1212
EVENT_HOST_ERROR,
1313
EVENT_SESSION_AGENT_EVENT,
14+
EVENT_SESSION_CONTEXT_COST,
1415
EVENT_SESSION_HISTORY_LOADED,
15-
EVENT_SESSION_MCP_USAGE,
1616
EVENT_SESSION_MESSAGE_DELTA,
1717
EVENT_SESSION_TOOL_APPROVAL_REQUESTED,
1818
EVENT_SESSION_TOOL_CALL_END,
@@ -78,7 +78,7 @@ const WORKER_TOPIC_MAP: Record<string, EventTopic> = {
7878
"tool.call.update": EVENT_SESSION_TOOL_CALL_UPDATE,
7979
"tool.call.end": EVENT_SESSION_TOOL_CALL_END,
8080
"turn.end": EVENT_SESSION_TURN_END,
81-
"session.mcp.usage": EVENT_SESSION_MCP_USAGE,
81+
"session.context.cost": EVENT_SESSION_CONTEXT_COST,
8282
"agent.event": EVENT_SESSION_AGENT_EVENT,
8383
// The agent-mode plugin emits this topic verbatim from the worker; we expose it under the
8484
// canonical EventTopic name so renderer code subscribes via the same constant.

packages/core/src/protocol/events.ts

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ export const EVENT_GIT_TURN_TOUCHES_CHANGED = "git.turnTouches.changed" as const
2626
export const EVENT_REVIEW_AVAILABLE = "review.available" as const;
2727
export const EVENT_REVIEW_CLEARED = "review.cleared" as const;
2828
export const EVENT_SESSION_ARTEFACTS_CHANGED = "session.artefacts.changed" as const;
29-
export const EVENT_SESSION_MCP_USAGE = "session.mcp.usage" as const;
29+
export const EVENT_SESSION_CONTEXT_COST = "session.context.cost" as const;
3030
export const EVENT_PLAN_FILE_CHANGED = "plan.file.changed" as const;
3131
export const EVENT_FS_TREE_CHANGED = "fs.tree.changed" as const;
3232
export const EVENT_TERMINAL_OUTPUT = "terminal.output" as const;
@@ -109,19 +109,27 @@ export const SessionWorkerExitPayload = z.object({
109109
});
110110

111111
/**
112-
* Estimated context cost of the MCP tools registered for this session. Emitted by the worker once
113-
* its extensions are bound (and again on every respawn, e.g. after toggling a server), so the
114-
* Context tab and composer ring can attribute MCP's slice of the context window. It's a chars/4
115-
* estimate over the real registered tool definitions — pi's aggregate can't be decomposed.
112+
* Estimated breakdown of this session's fixed context overhead — the system prompt and the tool
113+
* definitions (built-in + MCP) pi sends every turn. Emitted by the worker once its extensions are
114+
* bound (and again on every respawn, e.g. after toggling a server), so the Context tab and composer
115+
* ring can attribute each slice of the context window instead of guessing with constants.
116+
*
117+
* pi only reports an aggregate token count, so these are derived from the real artefacts pi holds:
118+
* `session.systemPrompt` (chars/4) and `session.getAllTools()` (chars/4 per definition). Stable for
119+
* the worker's lifetime — the conversation is everything in `used` beyond this overhead.
116120
*/
117-
export const SessionMcpUsagePayload = z.object({
121+
export const SessionContextCostPayload = z.object({
118122
sessionId: z.string(),
119-
/** Estimated tokens consumed by MCP tool definitions (proxy tool + any direct-exposed tools). */
120-
tokens: z.number(),
123+
/** Estimated tokens for the assembled system prompt text (`session.systemPrompt`). */
124+
systemPrompt: z.number(),
125+
/** Estimated tokens for built-in (non-MCP) tool definitions. */
126+
builtinTools: z.number(),
127+
/** Estimated tokens for MCP tool definitions (the `mcp` proxy + any direct-exposed tools). */
128+
mcp: z.number(),
121129
/** Number of MCP-origin tools registered (the `mcp` proxy counts as one). */
122-
toolCount: z.number(),
130+
mcpToolCount: z.number(),
123131
});
124-
export type SessionMcpUsage = z.infer<typeof SessionMcpUsagePayload>;
132+
export type SessionContextCost = z.infer<typeof SessionContextCostPayload>;
125133

126134
/**
127135
* Snapshot of past session messages + tool calls, broadcast by the host after `activate`
@@ -350,7 +358,7 @@ export const EventSchemas = {
350358
[EVENT_REVIEW_AVAILABLE]: ReviewAvailablePayload,
351359
[EVENT_REVIEW_CLEARED]: ReviewClearedPayload,
352360
[EVENT_SESSION_ARTEFACTS_CHANGED]: SessionArtefactsChangedPayload,
353-
[EVENT_SESSION_MCP_USAGE]: SessionMcpUsagePayload,
361+
[EVENT_SESSION_CONTEXT_COST]: SessionContextCostPayload,
354362
[EVENT_FS_TREE_CHANGED]: FsTreeChangedPayload,
355363
[EVENT_PLAN_FILE_CHANGED]: PlanFileChangedPayload,
356364
[EVENT_TERMINAL_OUTPUT]: TerminalOutputPayload,

packages/core/src/worker/agent-bridge.ts

Lines changed: 31 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,10 @@ import type { AgentMode, SessionModelRef, ThinkingLevel } from "../domain/sessio
1919
import { type ApprovalDecision, createAgentModeExtension } from "../extensions/agent-mode/index.js";
2020
import { createAttachmentsExtension } from "../extensions/attachments/index.js";
2121
import { listProjectFiles } from "../git/files.js";
22-
import { estimateToolsTokens, isMcpTool } from "../host/mcp-tokens.js";
22+
import { estimateToolTokens, isMcpTool } from "../host/mcp-tokens.js";
2323
import type { PromptAttachment, PromptImage, SessionCommandInfo } from "../protocol/commands.js";
2424
import {
25-
EVENT_SESSION_MCP_USAGE,
25+
EVENT_SESSION_CONTEXT_COST,
2626
EVENT_SESSION_TOOL_APPROVAL_REQUESTED,
2727
} from "../protocol/events.js";
2828
import { validateAndChdir } from "./cwd.js";
@@ -138,23 +138,38 @@ function toPiThinkingLevel(level: ThinkingLevel | undefined): PiThinkingLevel |
138138
}
139139

140140
/**
141-
* Estimate the MCP tools' context cost and emit it to the host. Filters pi's registered tools to
142-
* those the MCP adapter contributed (matched by name / source path), then sums a chars/4 estimate
143-
* over their definitions.
141+
* Estimate this session's fixed context overhead and emit it to the host: the system prompt text,
142+
* the built-in tool definitions, and the MCP tool definitions. pi only reports an aggregate token
143+
* count, so we derive these from the real artefacts it holds — `session.systemPrompt` and
144+
* `session.getAllTools()` — with the same chars/4 heuristic pi uses internally. This lets the
145+
* Context tab attribute each slice instead of guessing with constants. Best-effort: a failure here
146+
* must never break session startup.
144147
*/
145-
function emitMcpUsage(
148+
function emitContextCost(
146149
session: AgentSession,
147150
adapterPath: string | undefined,
148151
emit: EventEmitter,
149152
): void {
150153
try {
151-
const mcpTools = session.getAllTools().filter((tool) => isMcpTool(tool, adapterPath));
152-
emit(EVENT_SESSION_MCP_USAGE, {
153-
tokens: estimateToolsTokens(mcpTools),
154-
toolCount: mcpTools.length,
154+
let builtinTools = 0;
155+
let mcp = 0;
156+
let mcpToolCount = 0;
157+
for (const tool of session.getAllTools()) {
158+
if (isMcpTool(tool, adapterPath)) {
159+
mcp += estimateToolTokens(tool);
160+
mcpToolCount += 1;
161+
} else {
162+
builtinTools += estimateToolTokens(tool);
163+
}
164+
}
165+
emit(EVENT_SESSION_CONTEXT_COST, {
166+
systemPrompt: Math.ceil(session.systemPrompt.length / 4),
167+
builtinTools,
168+
mcp,
169+
mcpToolCount,
155170
});
156171
} catch {
157-
// Non-fatal: the Context tab simply shows no MCP slice until the next successful emit.
172+
// Non-fatal: the Context tab falls back to its floor estimates until the next successful emit.
158173
}
159174
}
160175

@@ -268,11 +283,11 @@ export async function initBridge(params: InitParams, emit: EventEmitter): Promis
268283
});
269284
}
270285

271-
// Now that extensions are bound, pi has registered the MCP adapter's tools (the `mcp` proxy +
272-
// any direct-exposed tools). Estimate their context cost and push it to the renderer so the
273-
// Context tab / composer ring can show MCP's slice of the window. Fires on every (re)spawn, so
274-
// toggling a server refreshes the figure.
275-
emitMcpUsage(session, mcpAdapterPath, emit);
286+
// Now that extensions are bound, pi has assembled the system prompt and registered every tool
287+
// (built-in + the MCP adapter's proxy/direct tools). Estimate that fixed overhead and push it to
288+
// the renderer so the Context tab / composer ring can attribute each slice of the window. Fires
289+
// on every (re)spawn, so toggling a server refreshes the figure.
290+
emitContextCost(session, mcpAdapterPath, emit);
276291

277292
return {
278293
session,

packages/ui/src/features/chat/composer/ContextUsageIndicator.tsx

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import {
66
formatTokens,
77
} from "../../context/contextBreakdown.js";
88
import { selectMessages, useMessagesStore } from "../useMessagesStore.js";
9-
import { selectSessionMcp, selectSessionUsage, useUsageStore } from "../useUsageStore.js";
9+
import { selectSessionCost, selectSessionUsage, useUsageStore } from "../useUsageStore.js";
1010

1111
interface ContextUsageIndicatorProps {
1212
sessionId: string;
@@ -26,12 +26,12 @@ interface ContextUsageIndicatorProps {
2626
*/
2727
export function ContextUsageIndicator({ sessionId }: ContextUsageIndicatorProps) {
2828
const usage = useUsageStore(selectSessionUsage(sessionId));
29-
const mcp = useUsageStore(selectSessionMcp(sessionId));
29+
const cost = useUsageStore(selectSessionCost(sessionId));
3030
const messages = useMessagesStore(selectMessages(sessionId));
3131

3232
const breakdown = useMemo(
33-
() => computeContextBreakdown(usage?.context, messages, mcp?.tokens ?? 0),
34-
[usage?.context, messages, mcp?.tokens],
33+
() => computeContextBreakdown(usage?.context, messages, cost),
34+
[usage?.context, messages, cost],
3535
);
3636

3737
// No turn has happened yet — show a placeholder. Once usage data lands, the ring fills.

packages/ui/src/features/chat/useUsageStore.ts

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,31 @@
11
import type { ContextUsage, TokenUsage } from "@pi-deck/core/protocol/events.js";
22
import { create } from "zustand";
33

4-
/** Estimated context cost of this session's MCP tools, pushed by the worker after extensions bind. */
5-
export interface McpUsage {
6-
tokens: number;
7-
toolCount: number;
4+
/**
5+
* This session's fixed context overhead, pushed by the worker once its extensions bind. Estimated
6+
* from the real system prompt + registered tool definitions (chars/4), so the breakdown can
7+
* attribute each slice instead of guessing with constants.
8+
*/
9+
export interface ContextCost {
10+
systemPrompt: number;
11+
builtinTools: number;
12+
mcp: number;
13+
mcpToolCount: number;
814
}
915

1016
interface SessionUsage {
1117
/** Token counts from the most recent turn. */
1218
lastTurn: TokenUsage;
1319
/** Aggregate context window state at the end of that turn. */
1420
context: ContextUsage | undefined;
15-
/** MCP tools' estimated slice of the context window, when the worker has reported it. */
16-
mcp: McpUsage | undefined;
21+
/** Per-category context overhead (system prompt + tool defs), when the worker has reported it. */
22+
cost: ContextCost | undefined;
1723
}
1824

1925
interface UsageStoreState {
2026
bySession: Record<string, SessionUsage>;
2127
setTurnUsage: (sessionId: string, usage: TokenUsage, context: ContextUsage | undefined) => void;
22-
setMcpUsage: (sessionId: string, mcp: McpUsage) => void;
28+
setContextCost: (sessionId: string, cost: ContextCost) => void;
2329
clearSession: (sessionId: string) => void;
2430
}
2531

@@ -33,15 +39,15 @@ export const useUsageStore = create<UsageStoreState>((set) => ({
3339
bySession: {
3440
...state.bySession,
3541
[sessionId]: {
36-
// Preserve the MCP estimate across turns — it's pushed once on worker spawn, not per turn.
37-
mcp: state.bySession[sessionId]?.mcp,
42+
// Preserve the cost estimate across turns — it's pushed once on worker spawn, not per turn.
43+
cost: state.bySession[sessionId]?.cost,
3844
lastTurn: usage,
3945
context,
4046
},
4147
},
4248
})),
4349

44-
setMcpUsage: (sessionId, mcp) =>
50+
setContextCost: (sessionId, cost) =>
4551
set((state) => {
4652
const prev = state.bySession[sessionId];
4753
return {
@@ -50,7 +56,7 @@ export const useUsageStore = create<UsageStoreState>((set) => ({
5056
[sessionId]: {
5157
lastTurn: prev?.lastTurn ?? EMPTY_TURN,
5258
context: prev?.context,
53-
mcp,
59+
cost,
5460
},
5561
},
5662
};
@@ -70,7 +76,7 @@ export function selectSessionUsage(sessionId: string | undefined) {
7076
sessionId ? state.bySession[sessionId] : undefined;
7177
}
7278

73-
export function selectSessionMcp(sessionId: string | undefined) {
74-
return (state: UsageStoreState): McpUsage | undefined =>
75-
sessionId ? state.bySession[sessionId]?.mcp : undefined;
79+
export function selectSessionCost(sessionId: string | undefined) {
80+
return (state: UsageStoreState): ContextCost | undefined =>
81+
sessionId ? state.bySession[sessionId]?.cost : undefined;
7682
}

packages/ui/src/features/context/PidContextPane.tsx

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -73,8 +73,8 @@ export function PidContextPane({ sessionId }: PidContextPaneProps) {
7373
}, [sessionId, client]);
7474

7575
const breakdown = useMemo(
76-
() => computeContextBreakdown(usage?.context, messages, usage?.mcp?.tokens ?? 0),
77-
[usage?.context, messages, usage?.mcp?.tokens],
76+
() => computeContextBreakdown(usage?.context, messages, usage?.cost),
77+
[usage?.context, messages, usage?.cost],
7878
);
7979

8080
const scope = useMemo(() => collectScope(messages), [messages]);
@@ -99,7 +99,7 @@ export function PidContextPane({ sessionId }: PidContextPaneProps) {
9999
breakdown={breakdown}
100100
percent={percent}
101101
active={hasData}
102-
mcpToolCount={usage?.mcp?.toolCount ?? 0}
102+
mcpToolCount={usage?.cost?.mcpToolCount ?? 0}
103103
/>
104104
<ScopeSection entries={scope} />
105105
<ArtefactsSection entries={artefactRows} />

packages/ui/src/features/context/contextBreakdown.ts

Lines changed: 44 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -22,42 +22,63 @@ export interface ContextBreakdown {
2222
free: number;
2323
}
2424

25+
/**
26+
* This session's fixed context overhead, estimated worker-side from the real system prompt + tool
27+
* definitions (see `ContextCost` / the `session.context.cost` event). Optional: until the worker
28+
* has reported it we fall back to the floor constants below.
29+
*/
30+
export interface ContextOverhead {
31+
systemPrompt?: number;
32+
builtinTools?: number;
33+
mcp?: number;
34+
}
35+
2536
const DEFAULT_CONTEXT_WINDOW = 200_000;
26-
const SYSTEM_PROMPT_FLOOR = 1_500; // ballpark — pi's base prompt + envelopes.
27-
const BUILTIN_TOOLS_FLOOR = 4_500; // ballpark — pi's built-in tool definitions (excludes MCP).
37+
// Fallbacks used only until the worker reports the real per-session overhead. Ballparks: pi's base
38+
// prompt + envelopes, and its built-in tool definitions (excludes MCP).
39+
const SYSTEM_PROMPT_FLOOR = 1_500;
40+
const BUILTIN_TOOLS_FLOOR = 4_500;
2841

2942
/**
30-
* Derive a per-category breakdown from `ContextUsage` (aggregate, from pi), the visible messages
31-
* (used to estimate the messages bucket), and the worker's MCP-tools estimate. When pi hasn't
32-
* reported usage yet (`ctx` is undefined) we still surface a floor estimate — including the MCP
33-
* figure when known — so the Context tab paints something useful before the first turn.
43+
* Derive a per-category breakdown from `ContextUsage` (aggregate, from pi), the visible messages,
44+
* and the worker's overhead estimate. pi only reports an aggregate `used`, so the breakdown is an
45+
* *attribution* of that single number — not an independent recount.
3446
*
35-
* pi only reports an aggregate `used`, so we carve the buckets out of it in priority order
36-
* (messages → system prompt → MCP → built-in tools) and they always sum to `used`. Post-turn the
37-
* aggregate already includes MCP (those tools are in the real payload), so carving is correct;
38-
* pre-turn we fold the MCP estimate into the floor so the bar reflects it immediately.
47+
* The overhead buckets (system prompt + built-in tools + MCP tools) are fixed for the worker's
48+
* lifetime and known precisely, so they're laid down first; **messages is the residual** — whatever
49+
* of `used` is left once the overhead is accounted for. That's the honest split: `used` is
50+
* dominated by the conversation, and the residual captures history + tool results + attachments the
51+
* visible message text alone can't see. The four buckets always sum to `used`.
52+
*
53+
* Before the first turn (`ctx` undefined) there's no aggregate, so `used` is the overhead plus a
54+
* chars/4 estimate of the visible messages. When the worker hasn't reported overhead yet, the floor
55+
* constants stand in.
3956
*/
4057
export function computeContextBreakdown(
4158
ctx: ContextUsage | undefined,
4259
messages: ReadonlyArray<{ text?: string }>,
43-
mcpTokens = 0,
60+
overhead: ContextOverhead = {},
4461
): ContextBreakdown {
45-
const messagesTokens = estimateMessagesTokens(messages);
4662
const contextWindow = ctx?.contextWindow ?? DEFAULT_CONTEXT_WINDOW;
47-
const mcp = Math.max(0, mcpTokens);
63+
const systemFixed = Math.max(0, overhead.systemPrompt ?? SYSTEM_PROMPT_FLOOR);
64+
const toolsFixed = Math.max(0, overhead.builtinTools ?? BUILTIN_TOOLS_FLOOR);
65+
const mcpFixed = Math.max(0, overhead.mcp ?? 0);
66+
67+
const messagesTokens = estimateMessagesTokens(messages);
4868
const used =
4969
typeof ctx?.tokens === "number" && ctx.tokens > 0
5070
? ctx.tokens
51-
: messagesTokens + SYSTEM_PROMPT_FLOOR + BUILTIN_TOOLS_FLOOR + mcp;
71+
: systemFixed + toolsFixed + mcpFixed + messagesTokens;
5272

53-
const messagesBucket = Math.min(messagesTokens, used);
54-
const afterMessages = Math.max(0, used - messagesBucket);
55-
const systemPrompt = Math.min(SYSTEM_PROMPT_FLOOR, afterMessages);
56-
const afterSystem = Math.max(0, afterMessages - systemPrompt);
57-
// MCP claims its estimate next, clamped so it never exceeds the real aggregate; built-in tools
58-
// take whatever remains.
59-
const mcpBucket = Math.min(mcp, afterSystem);
60-
const tools = Math.max(0, afterSystem - mcpBucket);
73+
// Lay down the fixed overhead in priority order, each clamped to what's left of `used` (so a
74+
// small real aggregate can't push the buckets past it), then hand the remainder to messages.
75+
const systemPrompt = Math.min(systemFixed, used);
76+
let remaining = used - systemPrompt;
77+
const tools = Math.min(toolsFixed, remaining);
78+
remaining -= tools;
79+
const mcp = Math.min(mcpFixed, remaining);
80+
remaining -= mcp;
81+
const messagesBucket = remaining; // residual — the conversation fills whatever overhead doesn't.
6182
const free = Math.max(0, contextWindow - used);
6283

6384
return {
@@ -66,7 +87,7 @@ export function computeContextBreakdown(
6687
messages: messagesBucket,
6788
systemPrompt,
6889
tools,
69-
mcp: mcpBucket,
90+
mcp,
7091
free,
7192
};
7293
}

0 commit comments

Comments
 (0)