From e170edfee11bff4cf3ac220dc4585711d8528dec Mon Sep 17 00:00:00 2001 From: li-xiu-qi Date: Tue, 27 Oct 2026 00:40:38 +0800 Subject: [PATCH 1/3] feat(coding-agent): add /context breakdown for context-window usage Report the current request as instructions, tool definitions, and conversation without starting a model request. Split conversation into user input, AI answers, thinking, tool calls, tool output, compaction summaries, and injected context, and show free space plus the configured compaction reserve. Expose an optional command-context snapshot so commands can read the effective system prompt, registered tools, MCP server names, messages, context window, and reserve tokens without re-deriving them. --- packages/coding-agent/docs/extensions.md | 13 +- packages/coding-agent/docs/usage.md | 7 + .../coding-agent/src/core/agent-session.ts | 27 +- .../coding-agent/src/core/extensions/index.ts | 1 + .../src/core/extensions/runner.ts | 7 + .../coding-agent/src/core/extensions/types.ts | 17 + packages/coding-agent/src/core/skills.ts | 16 +- .../coding-agent/src/render/context-panel.ts | 107 +++++ .../coding-agent/src/step/context-report.ts | 364 ++++++++++++++++++ .../coding-agent/src/step/slash-commands.ts | 40 ++ .../coding-agent/test/context-report.test.ts | 236 ++++++++++++ .../test/step-slash-commands.test.ts | 4 +- 12 files changed, 830 insertions(+), 9 deletions(-) create mode 100644 packages/coding-agent/src/render/context-panel.ts create mode 100644 packages/coding-agent/src/step/context-report.ts create mode 100644 packages/coding-agent/test/context-report.test.ts diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index fdfe9f67..894b428a 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -1065,7 +1065,7 @@ pi.on("tool_call", (event, ctx) => { ### ctx.getContextUsage() -Returns current context usage for the active model. Uses last assistant usage when available, then estimates tokens for trailing messages. +Returns current context usage for the active model. Uses last assistant usage when available, then estimates tokens for trailing messages. The optional `reserveTokens` field reports the configured compaction reserve, or zero when automatic compaction is disabled. ```typescript const usage = ctx.getContextUsage(); @@ -1123,6 +1123,17 @@ This has the same shape and mutability as `before_agent_start` `event.systemProm This reports the current base prompt inputs. It does not include per-turn `before_agent_start` chained system-prompt changes, later `context` event message mutations, or `before_provider_request` payload rewrites. +### ctx.getContextReportSnapshot() + +Returns a point-in-time snapshot of the data needed to report context composition, or `undefined` when the host does not provide it. The snapshot contains the effective system prompt, the base system-prompt options, the registered tool catalog, current MCP server names, the active message list, the model context window, and the configured compaction reserve. + +```typescript +const snapshot = ctx.getContextReportSnapshot?.(); +const toolCount = snapshot?.tools.length ?? 0; +``` + +Like `getSystemPromptOptions()`, this can include full prompt, context-file, and message contents. Treat it as sensitive extension-local data and avoid exposing it through command lists, logs, or autocomplete metadata. Prefer it over re-deriving context data when a command only needs a read-only report. + ### ctx.waitForIdle() Wait for the agent to fully settle, including automatic retries, auto-compaction retries, and queued continuations: diff --git a/packages/coding-agent/docs/usage.md b/packages/coding-agent/docs/usage.md index f0a081d5..12be6caf 100644 --- a/packages/coding-agent/docs/usage.md +++ b/packages/coding-agent/docs/usage.md @@ -46,6 +46,7 @@ Type `/` in the editor to open command completion. Extensions can register custo | `/new` | Start a new session | | `/name ` | Set session display name | | `/session` | Show session file, ID, messages, tokens, and cost | +| `/context [all]` | Show the estimated current context breakdown; `all` expands skills, context files, and individual tools | | `/tree` | Jump to any point in the session and continue from there | | `/trust` | Save project trust decision for future sessions | | `/fork` | Create a new session from a previous user message | @@ -94,6 +95,12 @@ Useful session commands: - `/clone` duplicates the current active branch into a new session file. - `/compact` summarizes older messages to free context. +### Current Context Breakdown + +`/context` shows a read-only estimate of the current context window. Instructions, tool definitions, and conversation are grouped separately. Conversation is further divided into user input, AI answers, thinking, tool calls, tool output, compaction summaries, and injected context. The meter also shows free space and the configured compaction reserve. + +Use `/context all`, or press `a` in the panel, to expand individual skills, context files, and tools. The report reads the active session and tool catalog without making a model or network request. Values are estimates of the current request structure, not cumulative billing or cache statistics; `/usage` owns those views. + See [Sessions](sessions.md) and [Compaction](compaction.md) for details. ## Context Files diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 54acd866..466da914 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -51,6 +51,7 @@ import { streamSimple, } from "@step-harness/providers/compat"; import { getAgentDir } from "../config.ts"; +import { getStepMcpStatuses } from "../step/mcp.ts"; import { redactSecretString } from "../step/secret-redaction.ts"; import { getThemeByName, theme } from "../theme/theme.ts"; import { stripFrontmatter } from "../utils/frontmatter.ts"; @@ -85,6 +86,7 @@ import { DEFAULT_THINKING_LEVEL, THINKING_LEVEL_OPTIONS } from "./defaults.ts"; import { exportSessionToHtml, type ToolHtmlRenderer } from "./export-html/index.ts"; import { createToolHtmlRenderer } from "./export-html/tool-renderer.ts"; import { + type ContextReportSnapshot, type ContextUsage, type ExtensionCommandContextActions, type ExtensionErrorListener, @@ -3400,6 +3402,7 @@ export class AgentSession { }, getSystemPrompt: () => this.systemPrompt, getSystemPromptOptions: () => this._baseSystemPromptOptions, + getContextReportSnapshot: () => this.createContextReportSnapshot(), getAutoRetryEnabled: () => this.autoRetryEnabled, setAutoRetryEnabled: (enabled) => this.setAutoRetryEnabled(enabled), }, @@ -4142,13 +4145,14 @@ export class AgentSession { const contextWindow = model.contextWindow ?? 0; if (contextWindow <= 0) return undefined; + const reserveTokens = this.autoCompactionEnabled ? this.settingsManager.getCompactionReserveTokens() : 0; if (this._getLiveContext()) { const tokens = this._workingContextTokens({ systemPrompt: this.systemPrompt, messages: this.agent.state.messages, tools: this.agent.state.tools, }); - return { tokens, contextWindow, percent: (tokens / contextWindow) * 100 }; + return { tokens, contextWindow, percent: (tokens / contextWindow) * 100, reserveTokens }; } // After compaction, the last assistant usage reflects pre-compaction context size. @@ -4176,7 +4180,7 @@ export class AgentSession { } if (!hasPostCompactionUsage) { - return { tokens: null, contextWindow, percent: null }; + return { tokens: null, contextWindow, percent: null, reserveTokens }; } } @@ -4187,6 +4191,25 @@ export class AgentSession { tokens: estimate.tokens, contextWindow, percent, + reserveTokens, + }; + } + + /** Snapshot the inputs a /context report needs, including tools that are registered but not active. */ + private createContextReportSnapshot(): ContextReportSnapshot { + const contextUsage = this.getContextUsage(); + const systemPromptOptions: BuildSystemPromptOptions = { + ...this._baseSystemPromptOptions, + contextFiles: this._baseSystemPromptOptions.contextFiles ?? [], + }; + return { + systemPrompt: this.systemPrompt, + systemPromptOptions, + tools: this.getAllTools(), + mcpServerNames: getStepMcpStatuses().map((status) => status.name), + messages: this.agent.state.messages, + contextWindow: contextUsage?.contextWindow ?? this.model?.contextWindow ?? 0, + reserveTokens: contextUsage?.reserveTokens, }; } diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index d472cbf0..bc735150 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -46,6 +46,7 @@ export type { ContextEvent, // Event Results ContextEventResult, + ContextReportSnapshot, ContextUsage, CustomToolCallEvent, CustomToolResultEvent, diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index e2a7c66c..57603eac 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -20,6 +20,7 @@ import type { CompactOptions, ContextEvent, ContextEventResult, + ContextReportSnapshot, ContextUsage, EntryRenderer, Extension, @@ -288,6 +289,7 @@ export class ExtensionRunner { private compactFn: (options?: CompactOptions) => void = () => {}; private getSystemPromptFn: () => string = () => ""; private getSystemPromptOptionsFn: () => BuildSystemPromptOptions = () => ({ cwd: this.cwd }); + private getContextReportSnapshotFn: () => ContextReportSnapshot | undefined = () => undefined; private getAutoRetryEnabledFn: (() => boolean) | undefined; private setAutoRetryEnabledFn: ((enabled: boolean) => void) | undefined; private newSessionHandler: NewSessionHandler = async () => ({ cancelled: false }); @@ -356,6 +358,7 @@ export class ExtensionRunner { this.compactFn = contextActions.compact; this.getSystemPromptFn = contextActions.getSystemPrompt; this.getSystemPromptOptionsFn = contextActions.getSystemPromptOptions ?? (() => ({ cwd: this.cwd })); + this.getContextReportSnapshotFn = contextActions.getContextReportSnapshot ?? (() => undefined); this.getAutoRetryEnabledFn = contextActions.getAutoRetryEnabled; this.setAutoRetryEnabledFn = contextActions.setAutoRetryEnabled; @@ -861,6 +864,10 @@ export class ExtensionRunner { this.assertActive(); return this.getSystemPromptOptionsFn(); }; + context.getContextReportSnapshot = () => { + this.assertActive(); + return this.getContextReportSnapshotFn(); + }; context.waitForIdle = () => { this.assertActive(); return this.waitForIdleFn(); diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 7fc04617..d93d1f1c 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -356,6 +356,8 @@ export interface ContextUsage { contextWindow: number; /** Context usage as percentage of context window, or null if tokens is unknown. */ percent: number | null; + /** Tokens held back for compaction and the next response; zero when automatic compaction is disabled. */ + reserveTokens?: number; } export interface CompactOptions { @@ -424,6 +426,8 @@ export interface ExtensionContext { export interface ExtensionCommandContext extends ExtensionContext { /** Get the current base system-prompt construction options. */ getSystemPromptOptions(): BuildSystemPromptOptions; + /** Snapshot the current system prompt, messages, and tool catalog for context composition reports. */ + getContextReportSnapshot?(): ContextReportSnapshot | undefined; /** Wait for the agent to finish streaming */ waitForIdle(): Promise; @@ -1739,6 +1743,18 @@ export type ToolInfo = Pick ToolInfo[]; +/** A point-in-time snapshot of the data needed to report context-token composition. */ +export interface ContextReportSnapshot { + systemPrompt: string; + systemPromptOptions: BuildSystemPromptOptions; + tools: readonly ToolInfo[]; + /** Names of MCP servers whose remote tools are currently registered. */ + mcpServerNames: readonly string[]; + messages: readonly AgentMessage[]; + contextWindow: number; + reserveTokens?: number; +} + export type GetCommandsHandler = () => SlashCommandInfo[]; export type SetActiveToolsHandler = (toolNames: string[]) => void; @@ -1818,6 +1834,7 @@ export interface ExtensionContextActions { compact: (options?: CompactOptions) => void; getSystemPrompt: () => string; getSystemPromptOptions?: () => BuildSystemPromptOptions; + getContextReportSnapshot?: () => ContextReportSnapshot | undefined; /** Optional native retry controls for product presets such as Step Autopilot. */ getAutoRetryEnabled?: () => boolean; setAutoRetryEnabled?: (enabled: boolean) => void; diff --git a/packages/coding-agent/src/core/skills.ts b/packages/coding-agent/src/core/skills.ts index 97ab24e6..312d6389 100644 --- a/packages/coding-agent/src/core/skills.ts +++ b/packages/coding-agent/src/core/skills.ts @@ -326,6 +326,16 @@ function loadSkillFromFile( * Skills with disableModelInvocation=true are excluded from the prompt * (they can only be invoked explicitly via /skill:name commands). */ +export function formatSkillForPrompt(skill: Skill): string { + return [ + " ", + ` ${escapeXml(skill.name)}`, + ` ${escapeXml(skill.description)}`, + ` ${escapeXml(skill.filePath)}`, + " ", + ].join("\n"); +} + export function formatSkillsForPrompt(skills: Skill[], readTool: "read" | "read_file" = "read"): string { const visibleSkills = skills.filter((s) => !s.disableModelInvocation); @@ -342,11 +352,7 @@ export function formatSkillsForPrompt(skills: Skill[], readTool: "read" | "read_ ]; for (const skill of visibleSkills) { - lines.push(" "); - lines.push(` ${escapeXml(skill.name)}`); - lines.push(` ${escapeXml(skill.description)}`); - lines.push(` ${escapeXml(skill.filePath)}`); - lines.push(" "); + lines.push(formatSkillForPrompt(skill)); } lines.push(""); diff --git a/packages/coding-agent/src/render/context-panel.ts b/packages/coding-agent/src/render/context-panel.ts new file mode 100644 index 00000000..935b758c --- /dev/null +++ b/packages/coding-agent/src/render/context-panel.ts @@ -0,0 +1,107 @@ +import { + type Component, + getKeybindings, + type KeybindingsManager, + matchesKey, + type TUI, + truncateToWidth, + visibleWidth, + wrapTextWithAnsi, +} from "@step-harness/pi-tui"; +import { + type ContextMeterSegment, + type ContextReport, + contextMeterSegments, + contextReportRows, +} from "../step/context-report.ts"; +import type { Theme } from "../theme/theme.ts"; + +export interface ContextPanelOptions { + report: ContextReport; + showAll?: boolean; + theme: Theme; + tui?: TUI; + onClose: () => void; +} + +const SEGMENT_COLORS: Record[0]> = { + instructions: "accent", + tools: "warning", + conversation: "success", + reserve: "dim", + free: "muted", +}; + +export function formatContextReportText(report: ContextReport, showAll = false): string { + const rows = contextReportRows(report, showAll); + const meter = contextMeterSegments(report) + .map((segment) => (segment === "free" ? "░" : segment === "reserve" ? "▒" : "█")) + .join(""); + if (meter) rows.splice(2, 0, meter); + rows.push("", showAll ? "A summary · Esc close" : "A all details · Esc close"); + return rows.join("\n"); +} + +/** Read-only current-context report. It never starts a model or network request. */ +export class ContextPanel implements Component { + private readonly options: Required> & ContextPanelOptions; + private readonly keybindings: KeybindingsManager; + private showAll: boolean; + private scroll = 0; + + constructor(options: ContextPanelOptions, keybindings: KeybindingsManager = getKeybindings()) { + this.options = options; + this.keybindings = keybindings; + this.showAll = options.showAll ?? false; + } + + invalidate(): void {} + + handleInput(data: string): void { + if (this.keybindings.matches(data, "tui.select.cancel") || data === "\u0003") { + this.options.onClose(); + return; + } + if (data === "a" || data === "A") { + this.showAll = !this.showAll; + this.scroll = 0; + this.options.tui?.requestRender(); + return; + } + if (this.keybindings.matches(data, "tui.select.up") || matchesKey(data, "pageUp")) { + this.scroll = Math.max(0, this.scroll - (matchesKey(data, "pageUp") ? 10 : 1)); + } else if (this.keybindings.matches(data, "tui.select.down") || matchesKey(data, "pageDown")) { + this.scroll += matchesKey(data, "pageDown") ? 10 : 1; + } + this.options.tui?.requestRender(); + } + + private meter(width: number): string { + const theme = this.options.theme; + return contextMeterSegments(this.options.report, Math.min(40, Math.max(1, width))) + .map((segment) => + theme.fg(SEGMENT_COLORS[segment], segment === "free" ? "░" : segment === "reserve" ? "▒" : "█"), + ) + .join(""); + } + + render(width: number): string[] { + const safeWidth = Math.max(1, Math.floor(width)); + const innerWidth = Math.max(1, safeWidth - 2); + const rows = contextReportRows(this.options.report, this.showAll); + if (this.options.report.usedTokens !== null && this.options.report.contextWindow > 0) { + rows.splice(2, 0, this.meter(innerWidth)); + } + rows.push( + "", + this.options.theme.fg("muted", this.showAll ? "A summary · Esc close" : "A all details · Esc close"), + ); + const wrapped = rows.flatMap((row) => wrapTextWithAnsi(row, innerWidth)); + const height = Math.max(5, (this.options.tui?.terminal.rows ?? wrapped.length + 4) - 6); + const maxScroll = Math.max(0, wrapped.length - height); + this.scroll = Math.max(0, Math.min(this.scroll, maxScroll)); + return wrapped + .slice(this.scroll, this.scroll + height) + .map((row) => ` ${visibleWidth(row) > innerWidth ? truncateToWidth(row, innerWidth, "", false) : row}`); + } +} diff --git a/packages/coding-agent/src/step/context-report.ts b/packages/coding-agent/src/step/context-report.ts new file mode 100644 index 00000000..de323df3 --- /dev/null +++ b/packages/coding-agent/src/step/context-report.ts @@ -0,0 +1,364 @@ +import type { AgentMessage } from "@step-harness/agent-core"; +import { estimateTokens } from "../core/compaction/compaction.ts"; +import type { ExtensionCommandContext, ToolInfo } from "../core/extensions/types.ts"; +import { sessionEntryToContextMessages } from "../core/session-manager.ts"; +import { formatSkillForPrompt, formatSkillsForPrompt } from "../core/skills.ts"; +import type { BuildSystemPromptOptions } from "../core/system-prompt.ts"; +import { getStepMcpStatuses } from "./mcp.ts"; + +export interface ContextReportRow { + id: string; + label: string; + tokens: number | null; + detail?: string; + children?: ContextReportRow[]; +} + +export interface ContextReport { + contextWindow: number; + usedTokens: number | null; + percent: number | null; + reserveTokens?: number; + freeTokens: number | null; + groups: ContextReportRow[]; +} + +export interface ContextReportInput { + systemPrompt: string; + systemPromptOptions: BuildSystemPromptOptions; + tools?: readonly ToolInfo[]; + mcpServerNames?: readonly string[]; + messages: readonly AgentMessage[]; + contextWindow: number; + reserveTokens?: number; +} + +export type ContextMeterSegment = "instructions" | "tools" | "conversation" | "reserve" | "free"; + +function estimateText(text: string): number { + return Math.ceil(text.length / 4); +} + +function sumRows(rows: readonly ContextReportRow[]): number | null { + let total = 0; + for (const row of rows) { + if (row.tokens === null) return null; + total += row.tokens; + } + return total; +} + +function nonZero(row: ContextReportRow): boolean { + return row.tokens === null || row.tokens > 0 || (row.children?.length ?? 0) > 0; +} + +function toolKind(tool: ToolInfo, mcpServerNames: readonly string[]): "builtin" | "mcp" | "extension" | "sdk" { + if (mcpServerNames.some((name) => tool.name === name || tool.name.startsWith(`${name}__`))) return "mcp"; + const source = tool.sourceInfo.source; + const sourcePath = tool.sourceInfo.path ?? ""; + if (source === "sdk" || sourcePath.startsWith(" = { + builtin: [], + mcp: [], + extension: [], + sdk: [], + }; + for (const tool of input.tools) { + const serialized = JSON.stringify({ + name: tool.name, + description: tool.description, + parameters: tool.parameters, + }); + groups[toolKind(tool, input.mcpServerNames ?? [])].push({ + id: `tool:${tool.name}`, + label: tool.name, + tokens: estimateText(serialized), + detail: tool.sourceInfo.path, + }); + } + const definitions = [ + { id: "builtin-tools", label: "Built-in tools", rows: groups.builtin }, + { id: "mcp-tools", label: "MCP tools", rows: groups.mcp }, + { id: "extension-tools", label: "Extensions & SDK", rows: groups.extension }, + { id: "sdk-tools", label: "SDK tools", rows: groups.sdk }, + ]; + return definitions.map(({ id, label, rows }) => ({ + id, + label, + tokens: sumRows(rows) ?? 0, + children: rows.sort((a, b) => b.tokens! - a.tokens! || a.label.localeCompare(b.label)), + })); +} + +function instructionRows(input: ContextReportInput): ContextReportRow[] { + const options = input.systemPromptOptions; + const activeTools = options.selectedTools ?? ["read", "bash", "edit", "write"]; + const readTool = activeTools.includes("read") ? "read" : activeTools.includes("read_file") ? "read_file" : undefined; + const skills = options.skills ?? []; + const visibleSkills = readTool ? skills.filter((skill) => !skill.disableModelInvocation) : []; + const skillsSection = readTool ? formatSkillsForPrompt(visibleSkills, readTool) : ""; + const skillsTokens = + skillsSection.length > 0 && input.systemPrompt.includes(skillsSection) ? estimateText(skillsSection) : 0; + const skillCharMap: Record = {}; + for (const skill of visibleSkills) { + skillCharMap[skill.filePath] = formatSkillForPrompt(skill).length; + } + const skillTokenMap = skillsTokens > 0 ? allocateTokenCounts(skillCharMap, skillsTokens) : {}; + const skillRows = visibleSkills.map((skill) => ({ + id: `skill:${skill.filePath}`, + label: skill.name, + tokens: skillTokenMap[skill.filePath] ?? 0, + detail: skill.filePath, + })); + const contextFiles = options.contextFiles ?? []; + const contextBlocks = contextFiles.map( + (file) => `\n${file.content}\n\n\n`, + ); + const contextSection = + contextBlocks.length > 0 + ? `\n\n\n\nProject-specific instructions and guidelines:\n\n${contextBlocks.join("")}\n` + : ""; + const contextFilesTokens = + contextSection.length > 0 && input.systemPrompt.includes(contextSection) ? estimateText(contextSection) : 0; + const contextCharMap: Record = {}; + for (const [index, file] of contextFiles.entries()) { + contextCharMap[file.path] = contextBlocks[index]?.length ?? 0; + } + const contextTokenMap = contextFilesTokens > 0 ? allocateTokenCounts(contextCharMap, contextFilesTokens) : {}; + const contextFileRows = contextFiles.map((file) => ({ + id: `context-file:${file.path}`, + label: file.path.split(/[\\/]/u).at(-1) || file.path, + tokens: contextTokenMap[file.path] ?? 0, + detail: file.path, + })); + const systemTokens = estimateText(input.systemPrompt); + const baseTokens = systemTokens - skillsTokens - contextFilesTokens; + return [ + { id: "system-prompt", label: "System prompt & dynamic", tokens: baseTokens }, + ...(skills.length + ? [{ id: "skills", label: `Skills (${skills.length})`, tokens: skillsTokens, children: skillRows }] + : []), + ...(contextFileRows.length + ? [ + { + id: "project-context", + label: `AGENTS.md & memory (${contextFileRows.length})`, + tokens: contextFilesTokens, + children: contextFileRows, + }, + ] + : []), + ]; +} + +function allocateTokenCounts(chars: Record, total: number): Record { + const entries = Object.entries(chars) as [T, number][]; + const result = Object.fromEntries(entries.map(([id, charCount]) => [id, Math.floor(charCount / 4)])) as Record< + T, + number + >; + let remaining = total - entries.reduce((sum, [id]) => sum + result[id], 0); + const order = entries + .map(([id, charCount], index) => ({ id, index, remainder: charCount % 4 })) + .sort((a, b) => b.remainder - a.remainder || a.index - b.index); + for (let index = 0; remaining > 0 && order.length > 0; index = (index + 1) % order.length) { + result[order[index].id] += 1; + remaining -= 1; + } + return result; +} + +function messageRows(messages: readonly AgentMessage[]): ContextReportRow[] { + const tokens: Record = { + user: 0, + assistant: 0, + thinking: 0, + toolCalls: 0, + toolResults: 0, + compaction: 0, + injected: 0, + other: 0, + }; + for (const message of messages) { + switch (message.role) { + case "user": + tokens.user += estimateTokens(message); + break; + case "assistant": { + const chars = { assistant: 0, thinking: 0, toolCalls: 0 }; + for (const block of message.content) { + if (block.type === "text") chars.assistant += block.text.length; + else if (block.type === "thinking") chars.thinking += block.thinking.length; + else if (block.type === "toolCall") { + const args = JSON.stringify(block.arguments) ?? ""; + chars.toolCalls += block.name.length + args.length; + } + } + const parts = allocateTokenCounts(chars, estimateTokens(message)); + tokens.assistant += parts.assistant; + tokens.thinking += parts.thinking; + tokens.toolCalls += parts.toolCalls; + break; + } + case "toolResult": + tokens.toolResults += estimateTokens(message); + break; + case "bashExecution": { + const parts = allocateTokenCounts( + { toolCalls: message.command.length, toolResults: message.output.length }, + estimateTokens(message), + ); + tokens.toolCalls += parts.toolCalls; + tokens.toolResults += parts.toolResults; + break; + } + case "branchSummary": + case "compactionSummary": + tokens.compaction += estimateTokens(message); + break; + case "custom": + tokens.injected += estimateTokens(message); + break; + default: + tokens.other += estimateTokens(message); + } + } + const definitions = [ + ["user", "User input"], + ["assistant", "AI answers"], + ["thinking", "Thinking"], + ["toolCalls", "Tool calls"], + ["toolResults", "Tool output"], + ["compaction", "Compaction summaries"], + ["injected", "Injected context"], + ["other", "Other messages"], + ] as const; + return definitions.map(([id, label]) => ({ id, label, tokens: tokens[id] })); +} + +export function buildContextReport(input: ContextReportInput): ContextReport { + const instructions = instructionRows(input); + const tools = toolRows(input); + const conversation = messageRows(input.messages); + const groups: ContextReportRow[] = [ + { id: "instructions", label: "Instructions", tokens: sumRows(instructions), children: instructions }, + { id: "tools", label: "Tool definitions", tokens: sumRows(tools), children: tools }, + { id: "conversation", label: "Conversation", tokens: sumRows(conversation), children: conversation }, + ]; + const usedTokens = sumRows(groups); + const percent = usedTokens !== null && input.contextWindow > 0 ? (usedTokens / input.contextWindow) * 100 : null; + const freeTokens = + usedTokens !== null && input.contextWindow > 0 + ? Math.max(0, input.contextWindow - usedTokens - (input.reserveTokens ?? 0)) + : null; + return { + contextWindow: input.contextWindow, + usedTokens, + percent, + reserveTokens: input.reserveTokens, + freeTokens, + groups, + }; +} + +export function captureContextReport(ctx: ExtensionCommandContext): ContextReport { + const snapshot = ctx.getContextReportSnapshot?.(); + const messages = + snapshot?.messages ?? + ctx.sessionManager + .buildContextEntries() + .flatMap(sessionEntryToContextMessages) + .filter((message) => !(message.role === "bashExecution" && message.excludeFromContext)); + const usage = ctx.getContextUsage(); + return buildContextReport({ + systemPrompt: snapshot?.systemPrompt ?? ctx.getSystemPrompt(), + systemPromptOptions: snapshot?.systemPromptOptions ?? ctx.getSystemPromptOptions(), + tools: snapshot?.tools ?? ctx.getToolCatalog?.(), + mcpServerNames: snapshot?.mcpServerNames ?? getStepMcpStatuses().map((status) => status.name), + messages, + contextWindow: snapshot?.contextWindow || usage?.contextWindow || ctx.model?.contextWindow || 0, + reserveTokens: snapshot?.reserveTokens ?? usage?.reserveTokens, + }); +} + +function formatTokens(tokens: number | null): string { + return tokens === null ? "--" : tokens.toLocaleString("en-US"); +} + +function formatRow(row: ContextReportRow, usedTokens: number | null, indent = ""): string { + const share = + row.tokens !== null && usedTokens !== null && usedTokens > 0 + ? `${((row.tokens / usedTokens) * 100).toFixed(1)}%` + : "--"; + return `${indent}${row.label}`.padEnd(32) + formatTokens(row.tokens).padStart(12) + share.padStart(9); +} + +export function contextReportRows(report: ContextReport, showAll = false): string[] { + const used = formatTokens(report.usedTokens); + const window = report.contextWindow > 0 ? report.contextWindow.toLocaleString("en-US") : "--"; + const percent = report.percent === null ? "--" : `${report.percent.toFixed(1)}%`; + const rows = ["Context Usage (estimated)", `${used} / ${window} tokens (${percent})`]; + for (const group of report.groups) { + rows.push("", formatRow(group, report.usedTokens)); + for (const child of group.children ?? []) { + if (!showAll && !nonZero(child)) continue; + rows.push(formatRow(child, report.usedTokens, " ")); + if (showAll) { + for (const detail of child.children ?? []) { + rows.push(formatRow(detail, report.usedTokens, " ") + (detail.detail ? ` ${detail.detail}` : "")); + } + } + } + } + rows.push( + "", + `${"Free after compaction reserve"}`.padEnd(32) + formatTokens(report.freeTokens).padStart(12), + `${"Compaction reserve"}`.padEnd(32) + formatTokens(report.reserveTokens ?? null).padStart(12), + ); + return rows; +} + +export function contextMeterSegments(report: ContextReport, width = 40): ContextMeterSegment[] { + if (report.usedTokens === null || report.contextWindow <= 0 || width <= 0) return []; + const usedCells = Math.min(width, Math.round((report.usedTokens / report.contextWindow) * width)); + const reserveCells = Math.min( + width - usedCells, + Math.round(((report.reserveTokens ?? 0) / report.contextWindow) * width), + ); + const weights = report.groups.map((group) => Math.max(0, group.tokens ?? 0)); + const totalWeight = weights.reduce((sum, weight) => sum + weight, 0); + const allocations = weights.map((weight) => (totalWeight > 0 ? (weight / totalWeight) * usedCells : 0)); + const cells = allocations.map(Math.floor); + let remaining = usedCells - cells.reduce((sum, value) => sum + value, 0); + const order = allocations + .map((value, index) => ({ index, fraction: value - Math.floor(value) })) + .sort((a, b) => b.fraction - a.fraction); + for (let index = 0; remaining > 0 && order.length > 0; index = (index + 1) % order.length) { + cells[order[index].index] += 1; + remaining -= 1; + } + const visibleGroups = weights.filter((weight) => weight > 0).length; + if (usedCells >= visibleGroups) { + for (const [index, weight] of weights.entries()) { + if (weight <= 0 || cells[index] > 0) continue; + const donor = cells.indexOf(Math.max(...cells)); + if (donor < 0 || cells[donor] <= 1) continue; + cells[donor] -= 1; + cells[index] = 1; + } + } + const segments: ContextMeterSegment[] = []; + for (const [index, group] of report.groups.entries()) { + const id = group.id as ContextMeterSegment; + for (let count = 0; count < cells[index]; count++) segments.push(id); + } + for (let count = 0; count < reserveCells; count++) segments.push("reserve"); + while (segments.length < width) segments.push("free"); + return segments; +} diff --git a/packages/coding-agent/src/step/slash-commands.ts b/packages/coding-agent/src/step/slash-commands.ts index 4e847e3b..14274329 100644 --- a/packages/coding-agent/src/step/slash-commands.ts +++ b/packages/coding-agent/src/step/slash-commands.ts @@ -13,8 +13,10 @@ import { type BtwPanelOutcome, createBtwPanel, } from "../render/btw-panel.ts"; +import { ContextPanel, formatContextReportText } from "../render/context-panel.ts"; import { getAvailableThemes } from "../theme/theme.ts"; import { copyToClipboard } from "../utils/clipboard.ts"; +import { captureContextReport } from "./context-report.ts"; import { type FeedbackSubmitResult, formatFeedbackBundleFailureDetails, @@ -130,6 +132,18 @@ export function registerStepPiCommandAdapters( telemetry, ); + registerTrackedCommand( + pi, + "context", + { + description: "Show the current context-window breakdown", + getArgumentCompletions: (prefix) => + ["all"].filter((value) => value.startsWith(prefix)).map((value) => ({ value, label: value })), + handler: handleStepContextCommand, + }, + telemetry, + ); + registerTrackedCommand( pi, "btw", @@ -543,6 +557,32 @@ async function handleStepThemeCommand(args: string, ctx: ExtensionCommandContext ctx.ui.notify(result.error ? `Failed to set theme: ${result.error}` : `Unknown theme "${requested}".`, "warning"); } +async function handleStepContextCommand(args: string, ctx: ExtensionCommandContext): Promise { + const requested = args.trim(); + if (requested !== "" && requested !== "all") { + ctx.ui.notify("Usage: /context [all]", "warning"); + return; + } + const report = captureContextReport(ctx); + const showAll = requested === "all"; + if (ctx.mode !== "tui" || !ctx.hasUI) { + ctx.ui.notify(formatContextReportText(report, showAll), "info"); + return; + } + await ctx.ui.custom((tui, theme, keybindings, done) => { + return new ContextPanel( + { + report, + showAll, + theme, + tui, + onClose: () => done(), + }, + keybindings, + ); + }); +} + /** Format a secret-free status snapshot from Pi's public command context. */ export function formatStepStatus(ctx: ExtensionCommandContext): string { const model = ctx.model; diff --git a/packages/coding-agent/test/context-report.test.ts b/packages/coding-agent/test/context-report.test.ts new file mode 100644 index 00000000..676028f6 --- /dev/null +++ b/packages/coding-agent/test/context-report.test.ts @@ -0,0 +1,236 @@ +import type { TUI } from "@step-harness/pi-tui"; +import type { AssistantMessage } from "@step-harness/providers/compat"; +import { beforeAll, describe, expect, test, vi } from "vitest"; +import { estimateTokens } from "../src/core/compaction/compaction.ts"; +import type { ToolInfo } from "../src/core/extensions/types.ts"; +import { KeybindingsManager } from "../src/core/keybindings.ts"; +import { type BuildSystemPromptOptions, buildSystemPrompt } from "../src/core/system-prompt.ts"; +import { ContextPanel, formatContextReportText } from "../src/render/context-panel.ts"; +import { + buildContextReport, + type ContextReportInput, + type ContextReportRow, + contextMeterSegments, + contextReportRows, +} from "../src/step/context-report.ts"; +import { initTheme, theme } from "../src/theme/theme.ts"; + +function assistant(content: AssistantMessage["content"]): AssistantMessage { + return { + role: "assistant", + content, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 1, + api: "anthropic-messages", + provider: "test", + model: "test-model", + }; +} + +function tool(name: string, source = "builtin"): ToolInfo { + return { + name, + description: `${name} description`, + parameters: { type: "object", properties: {} }, + sourceInfo: { path: `<${source}:${name}>`, source, scope: "temporary", origin: "top-level" }, + }; +} + +function fixture(): ContextReportInput { + const systemPromptOptions: BuildSystemPromptOptions = { + cwd: "/repo", + selectedTools: ["read"], + contextFiles: [{ path: "/repo/AGENTS.md", content: "a".repeat(40) }], + skills: [ + { + name: "alpha-skill", + description: "Alpha skill description", + filePath: "/repo/.step-code/skills/alpha/SKILL.md", + baseDir: "/repo/.step-code/skills/alpha", + sourceInfo: { + path: "/repo/.step-code/skills/alpha/SKILL.md", + source: "project", + scope: "project", + origin: "top-level", + }, + disableModelInvocation: false, + }, + ], + }; + return { + systemPrompt: buildSystemPrompt(systemPromptOptions), + systemPromptOptions, + tools: [tool("read"), tool("docs__search", "extension")], + mcpServerNames: ["docs"], + messages: [ + { role: "user", content: "u".repeat(40), timestamp: 1 }, + assistant([ + { type: "text", text: "a".repeat(40) }, + { type: "thinking", thinking: "t".repeat(20) }, + { type: "toolCall", id: "call-1", name: "read", arguments: { path: "x" } }, + ]), + { + role: "toolResult", + toolCallId: "call-1", + toolName: "read", + content: [{ type: "text", text: "o".repeat(80) }], + isError: false, + timestamp: 2, + }, + { role: "compactionSummary", summary: "s".repeat(40), tokensBefore: 100, timestamp: 3 }, + { role: "custom", customType: "notice", content: "i".repeat(20), display: false, timestamp: 4 }, + ], + contextWindow: 100_000, + reserveTokens: 1_000, + }; +} + +function group(report: ReturnType, id: string) { + const result = report.groups.find((candidate) => candidate.id === id); + if (!result) throw new Error(`Missing group ${id}`); + return result; +} + +function child(rows: ReturnType["groups"], id: string): ContextReportRow | undefined { + for (const row of rows) { + if (row.id === id) return row; + const found = child(row.children ?? [], id); + if (found) return found; + } + return undefined; +} + +beforeAll(() => { + initTheme("dark"); +}); + +describe("current context report", () => { + test("separates message roles and reconciles every category with the total", () => { + const input = fixture(); + const report = buildContextReport(input); + const conversation = group(report, "conversation"); + const toolCalls = child(report.groups, "toolCalls")!.tokens!; + const expectedMessages = input.messages.reduce((total, message) => total + estimateTokens(message), 0); + + expect(child(report.groups, "user")!.tokens).toBe(10); + expect(child(report.groups, "assistant")!.tokens).toBe(10); + expect(child(report.groups, "thinking")!.tokens).toBe(5); + expect(child(report.groups, "toolResults")!.tokens).toBe(20); + expect(child(report.groups, "compaction")!.tokens).toBe(10); + expect(child(report.groups, "injected")!.tokens).toBe(5); + expect(toolCalls).toBeGreaterThan(0); + expect(conversation.tokens).toBe(expectedMessages); + expect(report.usedTokens).toBe( + group(report, "instructions").tokens! + group(report, "tools").tokens! + conversation.tokens!, + ); + expect(report.freeTokens).toBe(100_000 - report.usedTokens! - 1_000); + }); + + test("keeps category rounding reconciled with the session token estimator", () => { + const messages: ContextReportInput["messages"] = [ + assistant([ + { type: "text", text: "a" }, + { type: "thinking", thinking: "t" }, + { type: "toolCall", id: "odd-call", name: "r", arguments: {} }, + ]), + { + role: "bashExecution", + command: "x", + output: "yy", + cancelled: false, + truncated: false, + exitCode: 0, + timestamp: 2, + }, + ]; + const report = buildContextReport({ + ...fixture(), + systemPrompt: "", + systemPromptOptions: { cwd: "/repo" }, + tools: [], + messages, + reserveTokens: 0, + }); + const conversation = group(report, "conversation"); + const expected = messages.reduce((total, message) => total + estimateTokens(message), 0); + const categoryTotal = conversation.children!.reduce((total, row) => total + row.tokens!, 0); + expect(conversation.tokens).toBe(expected); + expect(categoryTotal).toBe(expected); + expect(report.usedTokens).toBe(expected); + }); + + test("separates MCP tools and expands resources only in the detailed view", () => { + const report = buildContextReport(fixture()); + expect(child(report.groups, "builtin-tools")!.children!.map((row) => row.label)).toEqual(["read"]); + expect(child(report.groups, "mcp-tools")!.children!.map((row) => row.label)).toEqual(["docs__search"]); + + const summary = contextReportRows(report).join("\n"); + const detailed = contextReportRows(report, true).join("\n"); + expect(summary).not.toContain("docs__search"); + expect(summary).not.toContain("alpha-skill"); + expect(detailed).toContain("docs__search"); + expect(detailed).toContain("alpha-skill"); + expect(detailed).toContain("AGENTS.md"); + expect(formatContextReportText(report, true)).toContain("A summary"); + }); + + test("separates SDK tools and reconciles instruction children with the system prompt", () => { + const input = fixture(); + const report = buildContextReport({ ...input, tools: [...input.tools!, tool("sdk_tool", "sdk")] }); + expect(child(report.groups, "builtin-tools")!.children!.map((row) => row.label)).toEqual(["read"]); + expect(child(report.groups, "sdk-tools")!.children!.map((row) => row.label)).toEqual(["sdk_tool"]); + + const instructions = group(report, "instructions"); + const childrenTotal = instructions.children!.reduce((total, row) => total + row.tokens!, 0); + expect(childrenTotal).toBe(instructions.tokens); + }); + + test("marks an unavailable tool catalog as unknown instead of zero", () => { + const report = buildContextReport({ ...fixture(), tools: undefined }); + expect(report.usedTokens).toBeNull(); + expect(report.freeTokens).toBeNull(); + expect(contextReportRows(report).join("\n")).toContain("--"); + expect(contextMeterSegments(report)).toEqual([]); + }); + + test("allocates the meter across used, reserve, and free space", () => { + const base = buildContextReport(fixture()); + const report = { + ...base, + contextWindow: base.usedTokens! + 1_000, + reserveTokens: 200, + freeTokens: 800, + }; + const segments = contextMeterSegments(report); + expect(segments).toHaveLength(40); + expect(segments).toContain("instructions"); + expect(segments).toContain("tools"); + expect(segments).toContain("conversation"); + expect(segments).toContain("reserve"); + expect(segments).toContain("free"); + }); + + test("panel toggles details, renders within the terminal, and closes", () => { + const report = buildContextReport(fixture()); + const onClose = vi.fn(); + const tui = { terminal: { rows: 24 }, requestRender: vi.fn() } as unknown as TUI; + const panel = new ContextPanel({ report, theme, tui, onClose }, new KeybindingsManager()); + const summary = panel.render(80).join("\n"); + expect(summary).toContain("Context Usage"); + expect(summary).not.toContain("docs__search"); + expect(summary.split("\n").length).toBeLessThanOrEqual(18); + + panel.handleInput("a"); + expect(panel.render(80).join("\n")).toContain("docs__search"); + panel.handleInput("\x1b"); + expect(onClose).toHaveBeenCalledOnce(); + }); +}); diff --git a/packages/coding-agent/test/step-slash-commands.test.ts b/packages/coding-agent/test/step-slash-commands.test.ts index 6aa19dde..480ede0d 100644 --- a/packages/coding-agent/test/step-slash-commands.test.ts +++ b/packages/coding-agent/test/step-slash-commands.test.ts @@ -115,7 +115,9 @@ describe("Step Pi slash command adapters", () => { test("registers product adapters without replacing Pi's native commands", () => { const { commands } = registerCommands(); - expect([...commands.keys()]).toEqual(expect.arrayContaining(["plugin", "clear", "theme", "status", "btw"])); + expect([...commands.keys()]).toEqual( + expect.arrayContaining(["plugin", "clear", "theme", "status", "context", "btw"]), + ); expect(commands.has("effort")).toBe(false); }); From d47dc98a792fa35ce4ab68896d6b11c996b18f41 Mon Sep 17 00:00:00 2001 From: li-xiu-qi Date: Tue, 27 Oct 2026 01:23:42 +0800 Subject: [PATCH 2/3] feat(coding-agent): count /context segments through the tokenizer endpoint Ask the Step token-count endpoint for instructions, tool definitions, the full conversation, and each conversation category. Allocate structural children from those counts and fall back to local estimates when the endpoint is unavailable. --- packages/coding-agent/docs/usage.md | 4 +- .../coding-agent/src/render/context-panel.ts | 2 +- .../coding-agent/src/step/context-report.ts | 516 +++++++++++++++++- .../coding-agent/src/step/slash-commands.ts | 2 +- packages/coding-agent/src/step/token-count.ts | 100 ++++ .../coding-agent/test/context-report.test.ts | 93 +++- .../coding-agent/test/token-count.test.ts | 170 ++++++ 7 files changed, 848 insertions(+), 39 deletions(-) create mode 100644 packages/coding-agent/src/step/token-count.ts create mode 100644 packages/coding-agent/test/token-count.test.ts diff --git a/packages/coding-agent/docs/usage.md b/packages/coding-agent/docs/usage.md index 12be6caf..51c2c109 100644 --- a/packages/coding-agent/docs/usage.md +++ b/packages/coding-agent/docs/usage.md @@ -99,7 +99,9 @@ Useful session commands: `/context` shows a read-only estimate of the current context window. Instructions, tool definitions, and conversation are grouped separately. Conversation is further divided into user input, AI answers, thinking, tool calls, tool output, compaction summaries, and injected context. The meter also shows free space and the configured compaction reserve. -Use `/context all`, or press `a` in the panel, to expand individual skills, context files, and tools. The report reads the active session and tool catalog without making a model or network request. Values are estimates of the current request structure, not cumulative billing or cache statistics; `/usage` owns those views. +Use `/context all`, or press `a` in the panel, to expand individual skills, context files, and tools. The report reads the active session and tool catalog without starting a model turn. When the active model is a Step endpoint, it asks the platform token-count endpoint for segmented totals: instructions, tool definitions, the full conversation, and each conversation category. Those requests send request-shaped payloads to the tokenizer only; they do not generate a model response. If the endpoint is unavailable, unauthenticated, rate-limited, or the active provider is not Step, `/context` falls back to local structural estimates. The header shows whether the totals came from the token-count API, a mix of both, or local estimates. + +Local fallback values are estimates of the current request structure, not cumulative billing or cache statistics; `/usage` owns those views. Local fallback counts ASCII text at about four characters per token, non-ASCII text at about one token per code point, and images as fixed blocks. See [Sessions](sessions.md) and [Compaction](compaction.md) for details. diff --git a/packages/coding-agent/src/render/context-panel.ts b/packages/coding-agent/src/render/context-panel.ts index 935b758c..d3ea22b2 100644 --- a/packages/coding-agent/src/render/context-panel.ts +++ b/packages/coding-agent/src/render/context-panel.ts @@ -42,7 +42,7 @@ export function formatContextReportText(report: ContextReport, showAll = false): return rows.join("\n"); } -/** Read-only current-context report. It never starts a model or network request. */ +/** Read-only current-context report. It never starts a model turn; token totals may come from the provider tokenizer endpoint. */ export class ContextPanel implements Component { private readonly options: Required> & ContextPanelOptions; private readonly keybindings: KeybindingsManager; diff --git a/packages/coding-agent/src/step/context-report.ts b/packages/coding-agent/src/step/context-report.ts index de323df3..c297f910 100644 --- a/packages/coding-agent/src/step/context-report.ts +++ b/packages/coding-agent/src/step/context-report.ts @@ -1,10 +1,16 @@ import type { AgentMessage } from "@step-harness/agent-core"; -import { estimateTokens } from "../core/compaction/compaction.ts"; +import type { Api, ImageContent, Model, TextContent } from "@step-harness/providers"; import type { ExtensionCommandContext, ToolInfo } from "../core/extensions/types.ts"; import { sessionEntryToContextMessages } from "../core/session-manager.ts"; import { formatSkillForPrompt, formatSkillsForPrompt } from "../core/skills.ts"; import type { BuildSystemPromptOptions } from "../core/system-prompt.ts"; import { getStepMcpStatuses } from "./mcp.ts"; +import { + countStepTokens, + type StepTokenCountApi, + type StepTokenCountRequest, + supportsStepTokenCount, +} from "./token-count.ts"; export interface ContextReportRow { id: string; @@ -21,6 +27,15 @@ export interface ContextReport { reserveTokens?: number; freeTokens: number | null; groups: ContextReportRow[]; + /** Where the displayed totals come from when the host can count tokens. */ + tokenSource: "estimate" | "provider" | "mixed"; +} + +export type ContextGroupId = "instructions" | "tools" | "conversation"; + +export interface ContextTokenCounts { + groups?: Partial>; + rows?: Record; } export interface ContextReportInput { @@ -31,12 +46,78 @@ export interface ContextReportInput { messages: readonly AgentMessage[]; contextWindow: number; reserveTokens?: number; + /** Host-reported current context tokens when available; anchors the estimated breakdown. */ + reportedTokens?: number | null; + /** Provider tokenizer counts for the report's structural segments. */ + tokenCounts?: ContextTokenCounts; } export type ContextMeterSegment = "instructions" | "tools" | "conversation" | "reserve" | "free"; +/** + * Estimate text tokens by code-point class. ASCII is roughly four characters per + * token; CJK, full-width, and emoji are closer to one token per code point. This + * is still an estimate, but it avoids the worst failure of a flat chars/4 rule on + * Chinese prompts and tool output. + */ function estimateText(text: string): number { - return Math.ceil(text.length / 4); + let ascii = 0; + let nonAscii = 0; + for (const character of text) { + if ((character.codePointAt(0) ?? 0) < 128) ascii += 1; + else nonAscii += 1; + } + return Math.ceil(ascii / 4) + nonAscii; +} + +function estimateJson(value: unknown): number { + return estimateText(JSON.stringify(value) ?? ""); +} + +/** Image blocks are not tokenized from base64; they are counted as a fixed block. */ +const ESTIMATED_IMAGE_TOKENS = 1500; + +function estimateAgentMessage(message: AgentMessage): number { + switch (message.role) { + case "user": { + if (typeof message.content === "string") return estimateText(message.content); + let tokens = 0; + for (const block of message.content) { + if (block.type === "text") tokens += estimateText(block.text); + else if (block.type === "image") tokens += ESTIMATED_IMAGE_TOKENS; + else tokens += estimateJson(block); + } + return tokens; + } + case "assistant": { + let tokens = 0; + for (const block of message.content) { + if (block.type === "text") tokens += estimateText(block.text); + else if (block.type === "thinking") tokens += estimateText(block.thinking); + else if (block.type === "toolCall") tokens += estimateText(block.name) + estimateJson(block.arguments); + else tokens += estimateJson(block); + } + return tokens; + } + case "toolResult": { + let tokens = 0; + for (const block of message.content) { + if (block.type === "text") tokens += estimateText(block.text); + else if (block.type === "image") tokens += ESTIMATED_IMAGE_TOKENS; + else tokens += estimateJson(block); + } + return tokens; + } + case "bashExecution": + return estimateText(message.command) + estimateText(message.output); + case "branchSummary": + case "compactionSummary": + return estimateText(message.summary); + case "custom": + return typeof message.content === "string" ? estimateText(message.content) : estimateJson(message.content); + default: + return estimateJson(message); + } } function sumRows(rows: readonly ContextReportRow[]): number | null { @@ -188,45 +269,39 @@ function messageRows(messages: readonly AgentMessage[]): ContextReportRow[] { for (const message of messages) { switch (message.role) { case "user": - tokens.user += estimateTokens(message); + tokens.user += estimateAgentMessage(message); break; case "assistant": { - const chars = { assistant: 0, thinking: 0, toolCalls: 0 }; + const weights = { assistant: 0, thinking: 0, toolCalls: 0 }; for (const block of message.content) { - if (block.type === "text") chars.assistant += block.text.length; - else if (block.type === "thinking") chars.thinking += block.thinking.length; + if (block.type === "text") weights.assistant += estimateText(block.text); + else if (block.type === "thinking") weights.thinking += estimateText(block.thinking); else if (block.type === "toolCall") { - const args = JSON.stringify(block.arguments) ?? ""; - chars.toolCalls += block.name.length + args.length; - } + weights.toolCalls += estimateText(block.name) + estimateJson(block.arguments); + } else weights.assistant += estimateJson(block); } - const parts = allocateTokenCounts(chars, estimateTokens(message)); - tokens.assistant += parts.assistant; - tokens.thinking += parts.thinking; - tokens.toolCalls += parts.toolCalls; + tokens.assistant += weights.assistant; + tokens.thinking += weights.thinking; + tokens.toolCalls += weights.toolCalls; break; } case "toolResult": - tokens.toolResults += estimateTokens(message); + tokens.toolResults += estimateAgentMessage(message); break; case "bashExecution": { - const parts = allocateTokenCounts( - { toolCalls: message.command.length, toolResults: message.output.length }, - estimateTokens(message), - ); - tokens.toolCalls += parts.toolCalls; - tokens.toolResults += parts.toolResults; + tokens.toolCalls += estimateText(message.command); + tokens.toolResults += estimateText(message.output); break; } case "branchSummary": case "compactionSummary": - tokens.compaction += estimateTokens(message); + tokens.compaction += estimateAgentMessage(message); break; case "custom": - tokens.injected += estimateTokens(message); + tokens.injected += estimateAgentMessage(message); break; default: - tokens.other += estimateTokens(message); + tokens.other += estimateAgentMessage(message); } } const definitions = [ @@ -242,6 +317,334 @@ function messageRows(messages: readonly AgentMessage[]): ContextReportRow[] { return definitions.map(([id, label]) => ({ id, label, tokens: tokens[id] })); } +function scaleRows(rows: readonly ContextReportRow[], factor: number): ContextReportRow[] { + if (factor === 1) return [...rows]; + return rows.map((row) => { + if (row.children?.length) { + const children = scaleRows(row.children, factor); + const tokens = sumRows(children) ?? (row.tokens === null ? null : Math.round(row.tokens * factor)); + return { ...row, tokens, children }; + } + return { ...row, tokens: row.tokens === null ? null : Math.round(row.tokens * factor) }; + }); +} + +const CONVERSATION_CATEGORY_IDS = [ + "user", + "assistant", + "thinking", + "toolCalls", + "toolResults", + "compaction", + "injected", + "other", +] as const; + +type ConversationCategoryId = (typeof CONVERSATION_CATEGORY_IDS)[number]; + +export interface ContextTokenCountSegment { + id: string; + request: StepTokenCountRequest; +} + +function serializeToolDefinitions(tools: readonly ToolInfo[]): unknown[] { + return tools.map((tool) => ({ + type: "function", + function: { + name: tool.name, + description: tool.description, + parameters: tool.parameters, + }, + })); +} + +function toolDefinitionsAsJson(tools: readonly ToolInfo[]): string { + return JSON.stringify( + tools.map((tool) => ({ + name: tool.name, + description: tool.description, + parameters: tool.parameters, + })), + ); +} + +function textBlocksToOpenAiContent(content: string | readonly (TextContent | ImageContent)[]): unknown { + if (typeof content === "string") return content; + const parts: unknown[] = []; + for (const block of content) { + if (block.type === "text") parts.push({ type: "text", text: block.text }); + else if (block.type === "image") { + parts.push({ type: "image_url", image_url: { url: `data:${block.mimeType};base64,${block.data}` } }); + } + } + if (parts.length === 1 && (parts[0] as { type?: string }).type === "text") { + return (parts[0] as { text: string }).text; + } + return parts; +} + +function assistantBlockText(content: AssistantMessageLike["content"], type: "text" | "thinking" | "toolCall"): string { + const parts: string[] = []; + for (const block of content) { + if (type === "toolCall" && block.type === "toolCall") { + parts.push(JSON.stringify({ name: block.name, arguments: block.arguments })); + } else if (type === "thinking" && block.type === "thinking") { + parts.push(block.thinking); + } else if (type === "text" && block.type === "text") { + parts.push(block.text); + } + } + return parts.join("\n\n"); +} + +type AssistantMessageLike = Extract; + +function categoryMessages( + messages: readonly AgentMessage[], + category: ConversationCategoryId, + api: StepTokenCountApi, +): unknown[] { + const result: unknown[] = []; + const pushText = (role: "user" | "assistant" | "system", text: string) => { + if (text.length === 0) return; + result.push({ role, content: text }); + }; + for (const message of messages) { + switch (category) { + case "user": + if (message.role === "user") { + result.push({ role: "user", content: textBlocksToOpenAiContent(message.content) }); + } + break; + case "assistant": + if (message.role === "assistant") { + const text = assistantBlockText(message.content, "text"); + if (api === "anthropic-messages") pushText("assistant", text); + else if (text.length > 0) result.push({ role: "assistant", content: text }); + } + break; + case "thinking": + if (message.role === "assistant") { + const text = assistantBlockText(message.content, "thinking"); + if (api === "anthropic-messages") pushText("assistant", text); + else if (text.length > 0) result.push({ role: "assistant", content: text }); + } + break; + case "toolCalls": + if (message.role === "assistant") { + const text = assistantBlockText(message.content, "toolCall"); + if (api === "anthropic-messages") pushText("assistant", text); + else if (text.length > 0) result.push({ role: "assistant", content: text }); + } + break; + case "toolResults": + if (message.role === "toolResult") { + const text = message.content + .filter((block): block is TextContent => block.type === "text") + .map((block) => block.text) + .join("\n"); + if (api === "anthropic-messages") pushText("user", text); + else if (text.length > 0) result.push({ role: "tool", content: text, tool_call_id: message.toolCallId }); + } + break; + case "compaction": + if (message.role === "compactionSummary" || message.role === "branchSummary") { + pushText("user", message.summary); + } + break; + case "injected": + if (message.role === "custom") { + result.push({ role: "user", content: textBlocksToOpenAiContent(message.content) }); + } + break; + case "other": + if (message.role === "bashExecution") { + pushText("user", `${message.command}\n${message.output}`); + } else if (message.role !== "user" && message.role !== "assistant" && message.role !== "toolResult") { + pushText("user", JSON.stringify(message)); + } + break; + } + } + return result; +} + +function allConversationMessages(messages: readonly AgentMessage[], api: StepTokenCountApi): unknown[] { + const result: unknown[] = []; + for (const message of messages) { + if (message.role === "user") { + result.push({ role: "user", content: textBlocksToOpenAiContent(message.content) }); + continue; + } + if (message.role === "assistant") { + const text = assistantBlockText(message.content, "text"); + const thinking = assistantBlockText(message.content, "thinking"); + const toolCalls = assistantBlockText(message.content, "toolCall"); + const combined = [text, thinking, toolCalls].filter((part) => part.length > 0).join("\n\n"); + if (combined.length > 0) result.push({ role: "assistant", content: combined }); + continue; + } + if (message.role === "toolResult") { + const text = message.content + .filter((block): block is TextContent => block.type === "text") + .map((block) => block.text) + .join("\n"); + if (text.length > 0) { + result.push( + api === "anthropic-messages" + ? { role: "user", content: text } + : { role: "tool", content: text, tool_call_id: message.toolCallId }, + ); + } + continue; + } + if (message.role === "compactionSummary" || message.role === "branchSummary") { + result.push({ role: "user", content: message.summary }); + continue; + } + if (message.role === "custom") { + result.push({ role: "user", content: textBlocksToOpenAiContent(message.content) }); + continue; + } + result.push({ role: "user", content: JSON.stringify(message) }); + } + return result; +} + +/** + * Build the segmented token-count requests for one report. + * + * The platform token-count endpoint returns a total per request, so the report + * asks for the three top-level groups, a tool-definition probe, and each + * conversation category. Structural children are then allocated from these + * counts instead of being sent as one unbounded request per row. + */ +export function buildContextTokenCountSegments( + input: ContextReportInput, + api: StepTokenCountApi, +): ContextTokenCountSegment[] { + const segments: ContextTokenCountSegment[] = []; + const tools = input.tools ?? []; + const systemMessage = (content: string): unknown[] => [{ role: "system", content }]; + const anthropicSystem = (content: string): StepTokenCountRequest => + api === "anthropic-messages" ? { messages: [], system: content } : { messages: systemMessage(content) }; + + if (input.systemPrompt.length > 0) { + segments.push({ id: "instructions", request: anthropicSystem(input.systemPrompt) }); + } + segments.push({ id: "empty", request: { messages: [] } }); + if (tools.length > 0) { + segments.push({ id: "tools-field", request: { messages: [], tools: serializeToolDefinitions(tools) } }); + segments.push({ id: "tools-json", request: anthropicSystem(toolDefinitionsAsJson(tools)) }); + } + const conversation = allConversationMessages(input.messages, api); + if (conversation.length > 0) { + segments.push({ id: "conversation", request: { messages: conversation } }); + for (const category of CONVERSATION_CATEGORY_IDS) { + const messages = categoryMessages(input.messages, category, api); + if (messages.length > 0) segments.push({ id: category, request: { messages } }); + } + } + return segments; +} + +/** Count the report segments through the Step tokenizer endpoint. */ +export async function countContextTokenCounts(options: { + model: Model; + auth: Parameters[0]["auth"]; + segments: readonly ContextTokenCountSegment[]; + signal?: AbortSignal; +}): Promise { + if (!supportsStepTokenCount(options.model) || !options.auth.ok) return {}; + const entries = await Promise.all( + options.segments.map( + async (segment) => [segment.id, await countStepTokens({ ...options, request: segment.request })] as const, + ), + ); + const counts = new Map(entries); + const empty = counts.get("empty") ?? 0; + const toolsField = counts.get("tools-field"); + const toolsJson = counts.get("tools-json"); + let toolsTokens: number | null = null; + if (toolsField !== null && toolsField !== undefined) toolsTokens = Math.max(0, toolsField - empty); + if ((toolsTokens === null || toolsTokens === 0) && toolsJson !== null && toolsJson !== undefined) { + const jsonTokens = Math.max(0, toolsJson - empty); + if (jsonTokens > 0) toolsTokens = jsonTokens; + } + const rows: Record = {}; + for (const category of CONVERSATION_CATEGORY_IDS) { + const count = counts.get(category); + if (count !== null && count !== undefined) rows[category] = count; + } + const conversation = counts.get("conversation") ?? null; + const groups: ContextTokenCounts["groups"] = { + instructions: counts.get("instructions") ?? null, + tools: toolsTokens, + conversation, + }; + if (groups.conversation === null) { + const rowSum = CONVERSATION_CATEGORY_IDS.reduce((total, category) => total + (rows[category] ?? 0), 0); + if (rowSum > 0) groups.conversation = rowSum; + } + return { groups, rows }; +} + +function allocateChildrenToTotal(row: ContextReportRow, total: number, weights: readonly number[]): ContextReportRow { + const children = row.children ?? []; + if (children.length === 0) return { ...row, tokens: total }; + const weightSum = weights.reduce((sum, weight) => sum + Math.max(0, weight), 0); + if (weightSum <= 0) { + return { + ...row, + tokens: total, + children: children.map((child) => ({ ...child, tokens: 0 })), + }; + } + const factor = total / weightSum; + const scaledChildren = children.map((child, index) => ({ + ...child, + tokens: Math.round(Math.max(0, weights[index] ?? 0) * factor), + })); + const sum = sumRows(scaledChildren) ?? 0; + const difference = total - sum; + if (difference !== 0) { + let targetIndex = -1; + let maxTokens = -1; + scaledChildren.forEach((child, index) => { + if ((child.tokens ?? 0) > maxTokens) { + maxTokens = child.tokens ?? 0; + targetIndex = index; + } + }); + if (targetIndex >= 0) { + const target = scaledChildren[targetIndex]; + if (target) scaledChildren[targetIndex] = { ...target, tokens: (target.tokens ?? 0) + difference }; + } + } + return { ...row, tokens: sumRows(scaledChildren), children: scaledChildren }; +} + +function applyGroupTokenCount( + group: ContextReportRow, + exactTotal: number | null | undefined, + rowCounts: Record, +): ContextReportRow { + const children = group.children ?? []; + if (children.length === 0) { + return exactTotal === null || exactTotal === undefined ? group : { ...group, tokens: exactTotal }; + } + if (exactTotal === null || exactTotal === undefined) { + const hasRowCounts = children.some((child) => rowCounts[child.id] !== undefined); + if (!hasRowCounts) return group; + const nextChildren = children.map((child) => + rowCounts[child.id] === undefined ? child : { ...child, tokens: rowCounts[child.id] }, + ); + return { ...group, tokens: sumRows(nextChildren), children: nextChildren }; + } + const weights = children.map((child) => rowCounts[child.id] ?? child.tokens ?? 0); + return allocateChildrenToTotal(group, exactTotal, weights); +} + export function buildContextReport(input: ContextReportInput): ContextReport { const instructions = instructionRows(input); const tools = toolRows(input); @@ -251,23 +654,59 @@ export function buildContextReport(input: ContextReportInput): ContextReport { { id: "tools", label: "Tool definitions", tokens: sumRows(tools), children: tools }, { id: "conversation", label: "Conversation", tokens: sumRows(conversation), children: conversation }, ]; - const usedTokens = sumRows(groups); + + const tokenCounts = input.tokenCounts; + const groupCounts = tokenCounts?.groups ?? {}; + const rowCounts = tokenCounts?.rows ?? {}; + const hasGroupCount = Object.values(groupCounts).some((value) => value !== null && value !== undefined); + const hasRowCount = Object.values(rowCounts).some((value) => value !== null && value !== undefined); + const hasProviderCounts = hasGroupCount || hasRowCount; + const countedGroups = hasProviderCounts + ? groups.map((group) => applyGroupTokenCount(group, groupCounts[group.id as ContextGroupId] ?? null, rowCounts)) + : groups; + + // Without provider tokenizer counts, the host may still know the real request + // size from provider usage or its calibrated live-context estimator. Anchor the + // estimated rows to that total when it exists; keep raw estimates before the + // first response or when the host cannot report a total. + const estimatedUsed = sumRows(groups); + const reportedTokens = input.reportedTokens; + const shouldAnchor = + !hasProviderCounts && + reportedTokens !== undefined && + reportedTokens !== null && + reportedTokens > 0 && + estimatedUsed !== null && + estimatedUsed > 0; + const finalGroups = shouldAnchor + ? scaleRows(groups, (reportedTokens as number) / (estimatedUsed as number)) + : countedGroups; + const usedTokens = sumRows(finalGroups); const percent = usedTokens !== null && input.contextWindow > 0 ? (usedTokens / input.contextWindow) * 100 : null; const freeTokens = usedTokens !== null && input.contextWindow > 0 ? Math.max(0, input.contextWindow - usedTokens - (input.reserveTokens ?? 0)) : null; + const hasAllGroupCounts = (["instructions", "tools", "conversation"] as ContextGroupId[]).every( + (id) => groupCounts[id] !== null && groupCounts[id] !== undefined, + ); + const tokenSource: ContextReport["tokenSource"] = !hasProviderCounts + ? "estimate" + : hasAllGroupCounts + ? "provider" + : "mixed"; return { contextWindow: input.contextWindow, usedTokens, percent, reserveTokens: input.reserveTokens, freeTokens, - groups, + groups: finalGroups, + tokenSource, }; } -export function captureContextReport(ctx: ExtensionCommandContext): ContextReport { +export async function captureContextReport(ctx: ExtensionCommandContext): Promise { const snapshot = ctx.getContextReportSnapshot?.(); const messages = snapshot?.messages ?? @@ -276,7 +715,7 @@ export function captureContextReport(ctx: ExtensionCommandContext): ContextRepor .flatMap(sessionEntryToContextMessages) .filter((message) => !(message.role === "bashExecution" && message.excludeFromContext)); const usage = ctx.getContextUsage(); - return buildContextReport({ + const input: ContextReportInput = { systemPrompt: snapshot?.systemPrompt ?? ctx.getSystemPrompt(), systemPromptOptions: snapshot?.systemPromptOptions ?? ctx.getSystemPromptOptions(), tools: snapshot?.tools ?? ctx.getToolCatalog?.(), @@ -284,7 +723,20 @@ export function captureContextReport(ctx: ExtensionCommandContext): ContextRepor messages, contextWindow: snapshot?.contextWindow || usage?.contextWindow || ctx.model?.contextWindow || 0, reserveTokens: snapshot?.reserveTokens ?? usage?.reserveTokens, - }); + reportedTokens: usage?.tokens ?? null, + }; + + const model = ctx.model; + if (model && supportsStepTokenCount(model)) { + const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model); + if (auth.ok) { + const segments = buildContextTokenCountSegments(input, model.api); + if (segments.length > 0) { + input.tokenCounts = await countContextTokenCounts({ model, auth, segments, signal: ctx.signal }); + } + } + } + return buildContextReport(input); } function formatTokens(tokens: number | null): string { @@ -303,7 +755,13 @@ export function contextReportRows(report: ContextReport, showAll = false): strin const used = formatTokens(report.usedTokens); const window = report.contextWindow > 0 ? report.contextWindow.toLocaleString("en-US") : "--"; const percent = report.percent === null ? "--" : `${report.percent.toFixed(1)}%`; - const rows = ["Context Usage (estimated)", `${used} / ${window} tokens (${percent})`]; + const sourceLabel = + report.tokenSource === "provider" + ? "Context Usage (token-count API)" + : report.tokenSource === "mixed" + ? "Context Usage (mixed)" + : "Context Usage (estimated)"; + const rows = [sourceLabel, `${used} / ${window} tokens (${percent})`]; for (const group of report.groups) { rows.push("", formatRow(group, report.usedTokens)); for (const child of group.children ?? []) { diff --git a/packages/coding-agent/src/step/slash-commands.ts b/packages/coding-agent/src/step/slash-commands.ts index 14274329..b0910fee 100644 --- a/packages/coding-agent/src/step/slash-commands.ts +++ b/packages/coding-agent/src/step/slash-commands.ts @@ -563,7 +563,7 @@ async function handleStepContextCommand(args: string, ctx: ExtensionCommandConte ctx.ui.notify("Usage: /context [all]", "warning"); return; } - const report = captureContextReport(ctx); + const report = await captureContextReport(ctx); const showAll = requested === "all"; if (ctx.mode !== "tui" || !ctx.hasUI) { ctx.ui.notify(formatContextReportText(report, showAll), "info"); diff --git a/packages/coding-agent/src/step/token-count.ts b/packages/coding-agent/src/step/token-count.ts new file mode 100644 index 00000000..ad8ca3a6 --- /dev/null +++ b/packages/coding-agent/src/step/token-count.ts @@ -0,0 +1,100 @@ +import type { Api, Model } from "@step-harness/providers"; +import type { ResolvedRequestAuth } from "../core/model-registry.ts"; + +/** Provider wire formats whose Step token-count contract is known. */ +export type StepTokenCountApi = "openai-completions" | "anthropic-messages"; + +export interface StepTokenCountRequest { + messages: readonly unknown[]; + system?: string; + tools?: readonly unknown[]; +} + +export interface StepTokenCountParams { + model: Model; + auth: ResolvedRequestAuth; + request: StepTokenCountRequest; + signal?: AbortSignal; +} + +const STEP_TOKEN_COUNT_HOSTS = new Set(["api.stepfun.com", "api.stepfun.ai"]); + +/** + * Token counting sends the current request shape to the provider's tokenizer + * endpoint, so it is restricted to Step endpoints. Custom providers keep the + * local structural estimate instead of leaking prompts to an unrelated host. + */ +export function supportsStepTokenCount(model: Model | undefined): model is Model { + if (!model) return false; + if (model.api !== "openai-completions" && model.api !== "anthropic-messages") return false; + try { + return STEP_TOKEN_COUNT_HOSTS.has(new URL(model.baseUrl).hostname); + } catch { + return false; + } +} + +function joinUrl(baseUrl: string, path: string): string { + return `${baseUrl.replace(/\/+$/, "")}/${path.replace(/^\/+/, "")}`; +} + +function readPositiveInt(value: unknown): number | null { + return typeof value === "number" && Number.isFinite(value) && value >= 0 ? Math.round(value) : null; +} + +/** + * Count tokens for one request payload through the Step platform tokenizer. + * + * Returns null for any transport, auth, or shape failure so callers can fall + * back to local estimates without turning `/context` into a failing command. + */ +export async function countStepTokens(params: StepTokenCountParams): Promise { + const { model, auth, request, signal } = params; + if (!supportsStepTokenCount(model)) return null; + if (!auth.ok) return null; + + const url = + model.api === "openai-completions" + ? joinUrl(model.baseUrl, "token/count") + : joinUrl(model.baseUrl, "messages/count_tokens"); + const body = + model.api === "openai-completions" + ? { + model: model.id, + messages: request.messages, + ...(request.tools ? { tools: request.tools } : {}), + } + : { + model: model.id, + messages: request.messages, + ...(request.system ? { system: request.system } : {}), + ...(request.tools ? { tools: request.tools } : {}), + }; + + const headers = new Headers({ + "Content-Type": "application/json", + ...model.headers, + ...auth.headers, + }); + if (auth.apiKey && !headers.has("Authorization")) { + headers.set("Authorization", `Bearer ${auth.apiKey}`); + } + + try { + const response = await fetch(url, { + method: "POST", + headers, + body: JSON.stringify(body), + signal: signal ?? AbortSignal.timeout(10_000), + }); + if (!response.ok) return null; + const payload: unknown = await response.json(); + if (model.api === "openai-completions") { + const data = (payload as { data?: { total_tokens?: unknown } }).data; + return readPositiveInt(data?.total_tokens); + } + return readPositiveInt((payload as { input_tokens?: unknown }).input_tokens); + } catch { + return null; + } +} diff --git a/packages/coding-agent/test/context-report.test.ts b/packages/coding-agent/test/context-report.test.ts index 676028f6..3bf7249e 100644 --- a/packages/coding-agent/test/context-report.test.ts +++ b/packages/coding-agent/test/context-report.test.ts @@ -1,7 +1,6 @@ import type { TUI } from "@step-harness/pi-tui"; import type { AssistantMessage } from "@step-harness/providers/compat"; import { beforeAll, describe, expect, test, vi } from "vitest"; -import { estimateTokens } from "../src/core/compaction/compaction.ts"; import type { ToolInfo } from "../src/core/extensions/types.ts"; import { KeybindingsManager } from "../src/core/keybindings.ts"; import { type BuildSystemPromptOptions, buildSystemPrompt } from "../src/core/system-prompt.ts"; @@ -118,7 +117,14 @@ describe("current context report", () => { const report = buildContextReport(input); const conversation = group(report, "conversation"); const toolCalls = child(report.groups, "toolCalls")!.tokens!; - const expectedMessages = input.messages.reduce((total, message) => total + estimateTokens(message), 0); + const expectedMessages = + (child(report.groups, "user")!.tokens ?? 0) + + (child(report.groups, "assistant")!.tokens ?? 0) + + (child(report.groups, "thinking")!.tokens ?? 0) + + toolCalls + + (child(report.groups, "toolResults")!.tokens ?? 0) + + (child(report.groups, "compaction")!.tokens ?? 0) + + (child(report.groups, "injected")!.tokens ?? 0); expect(child(report.groups, "user")!.tokens).toBe(10); expect(child(report.groups, "assistant")!.tokens).toBe(10); @@ -126,7 +132,7 @@ describe("current context report", () => { expect(child(report.groups, "toolResults")!.tokens).toBe(20); expect(child(report.groups, "compaction")!.tokens).toBe(10); expect(child(report.groups, "injected")!.tokens).toBe(5); - expect(toolCalls).toBeGreaterThan(0); + expect(toolCalls).toBe(4); expect(conversation.tokens).toBe(expectedMessages); expect(report.usedTokens).toBe( group(report, "instructions").tokens! + group(report, "tools").tokens! + conversation.tokens!, @@ -160,11 +166,39 @@ describe("current context report", () => { reserveTokens: 0, }); const conversation = group(report, "conversation"); - const expected = messages.reduce((total, message) => total + estimateTokens(message), 0); const categoryTotal = conversation.children!.reduce((total, row) => total + row.tokens!, 0); - expect(conversation.tokens).toBe(expected); - expect(categoryTotal).toBe(expected); - expect(report.usedTokens).toBe(expected); + expect(child(report.groups, "assistant")!.tokens).toBe(1); + expect(child(report.groups, "thinking")!.tokens).toBe(1); + expect(child(report.groups, "toolCalls")!.tokens).toBe(3); + expect(child(report.groups, "toolResults")!.tokens).toBe(1); + expect(conversation.tokens).toBe(6); + expect(categoryTotal).toBe(6); + expect(report.usedTokens).toBe(6); + }); + + test("counts non-ASCII text closer to one token per code point", () => { + const report = buildContextReport({ + ...fixture(), + systemPrompt: "中文", + systemPromptOptions: { cwd: "/repo" }, + tools: [], + messages: [], + reserveTokens: 0, + }); + expect(child(report.groups, "system-prompt")!.tokens).toBe(2); + expect(report.usedTokens).toBe(2); + }); + + test("anchors the estimated breakdown to host-reported context tokens", () => { + const base = fixture(); + const unanchored = buildContextReport(base); + const reportedTokens = (unanchored.usedTokens ?? 0) * 2; + const report = buildContextReport({ ...base, reportedTokens }); + expect(report.usedTokens).toBe(reportedTokens); + expect(group(report, "instructions").tokens).toBe(group(unanchored, "instructions").tokens! * 2); + expect(group(report, "tools").tokens).toBe(group(unanchored, "tools").tokens! * 2); + expect(group(report, "conversation").tokens).toBe(group(unanchored, "conversation").tokens! * 2); + expect(report.freeTokens).toBe(100_000 - reportedTokens - 1_000); }); test("separates MCP tools and expands resources only in the detailed view", () => { @@ -218,6 +252,51 @@ describe("current context report", () => { expect(segments).toContain("free"); }); + test("uses provider token counts for group totals and keeps children reconciled", () => { + const base = fixture(); + const report = buildContextReport({ + ...base, + tokenCounts: { + groups: { instructions: 1_000, tools: 200, conversation: 5_000 }, + }, + }); + expect(report.tokenSource).toBe("provider"); + expect(report.usedTokens).toBe(6_200); + expect(group(report, "instructions").tokens).toBe(1_000); + expect(group(report, "tools").tokens).toBe(200); + expect(group(report, "conversation").tokens).toBe(5_000); + for (const groupRow of report.groups) { + const childrenTotal = groupRow.children!.reduce((total, row) => total + row.tokens!, 0); + expect(childrenTotal).toBe(groupRow.tokens); + } + expect(contextReportRows(report).join("\n")).toContain("Context Usage (token-count API)"); + }); + + test("uses provider category counts as conversation weights", () => { + const base = fixture(); + const report = buildContextReport({ + ...base, + tokenCounts: { + groups: { conversation: 800 }, + rows: { + user: 100, + assistant: 200, + thinking: 50, + toolCalls: 150, + toolResults: 300, + compaction: 0, + injected: 0, + other: 0, + }, + }, + }); + expect(report.tokenSource).toBe("mixed"); + expect(group(report, "conversation").tokens).toBe(800); + expect(child(report.groups, "user")!.tokens).toBe(100); + expect(child(report.groups, "toolResults")!.tokens).toBe(300); + expect(report.usedTokens).toBe(group(report, "instructions").tokens! + group(report, "tools").tokens! + 800); + }); + test("panel toggles details, renders within the terminal, and closes", () => { const report = buildContextReport(fixture()); const onClose = vi.fn(); diff --git a/packages/coding-agent/test/token-count.test.ts b/packages/coding-agent/test/token-count.test.ts new file mode 100644 index 00000000..72eb5c38 --- /dev/null +++ b/packages/coding-agent/test/token-count.test.ts @@ -0,0 +1,170 @@ +import type { Api, Model } from "@step-harness/providers"; +import { afterEach, describe, expect, test, vi } from "vitest"; +import type { ToolInfo } from "../src/core/extensions/types.ts"; +import { + buildContextTokenCountSegments, + type ContextReportInput, + type ContextTokenCountSegment, + countContextTokenCounts, +} from "../src/step/context-report.ts"; +import { countStepTokens, supportsStepTokenCount } from "../src/step/token-count.ts"; + +function model(api: Api, baseUrl: string): Model { + return { + id: "step-5-preview", + name: "Step 5 Preview", + api, + provider: "step", + baseUrl, + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1_000_000, + maxTokens: 64_000, + }; +} + +const auth = { ok: true, apiKey: "test-key" } as const; + +function tool(name: string): ToolInfo { + return { + name, + description: `${name} description`, + parameters: { type: "object", properties: {} }, + sourceInfo: { path: ``, source: "builtin", scope: "temporary", origin: "top-level" }, + }; +} + +function input(): ContextReportInput { + return { + systemPrompt: "system prompt", + systemPromptOptions: { cwd: "/repo" }, + tools: [tool("read")], + messages: [ + { role: "user", content: "hello", timestamp: 1 }, + { + role: "assistant", + content: [{ type: "text", text: "hi" }], + api: "openai-completions", + provider: "step", + model: "step-5-preview", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 2, + }, + ], + contextWindow: 1_000_000, + }; +} + +afterEach(() => { + vi.unstubAllGlobals(); +}); + +describe("step token count endpoint", () => { + test("only counts Step endpoints for supported provider APIs", () => { + expect(supportsStepTokenCount(model("openai-completions", "https://api.stepfun.com/v1"))).toBe(true); + expect(supportsStepTokenCount(model("anthropic-messages", "https://api.stepfun.ai/v1"))).toBe(true); + expect(supportsStepTokenCount(model("openai-completions", "https://api.openai.com/v1"))).toBe(false); + expect(supportsStepTokenCount(model("openai-responses", "https://api.stepfun.com/v1"))).toBe(false); + expect(supportsStepTokenCount(undefined)).toBe(false); + }); + + test("parses the OpenAI-compatible token count response", async () => { + const fetchMock = vi.fn(async () => ({ ok: true, json: async () => ({ data: { total_tokens: 42 } }) })); + vi.stubGlobal("fetch", fetchMock); + + const tokens = await countStepTokens({ + model: model("openai-completions", "https://api.stepfun.com/v1"), + auth, + request: { messages: [{ role: "user", content: "hello" }] }, + }); + + expect(tokens).toBe(42); + expect(fetchMock).toHaveBeenCalledWith( + "https://api.stepfun.com/v1/token/count", + expect.objectContaining({ method: "POST" }), + ); + }); + + test("parses the Anthropic messages count response", async () => { + const fetchMock = vi.fn(async () => ({ ok: true, json: async () => ({ input_tokens: 17 }) })); + vi.stubGlobal("fetch", fetchMock); + + const tokens = await countStepTokens({ + model: model("anthropic-messages", "https://api.stepfun.ai/v1"), + auth, + request: { messages: [{ role: "user", content: "hello" }], system: "system" }, + }); + + expect(tokens).toBe(17); + expect(fetchMock).toHaveBeenCalledWith( + "https://api.stepfun.ai/v1/messages/count_tokens", + expect.objectContaining({ method: "POST" }), + ); + }); + + test("returns null instead of throwing when the endpoint fails", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => ({ ok: false, json: async () => ({}) })), + ); + await expect( + countStepTokens({ + model: model("openai-completions", "https://api.stepfun.com/v1"), + auth, + request: { messages: [] }, + }), + ).resolves.toBeNull(); + }); +}); + +describe("context token-count segments", () => { + test("builds segmented requests for instructions, tools, and conversation categories", () => { + const segments = buildContextTokenCountSegments(input(), "openai-completions"); + expect(segments.map((segment) => segment.id)).toEqual([ + "instructions", + "empty", + "tools-field", + "tools-json", + "conversation", + "user", + "assistant", + ]); + }); + + test("prefers the tools field and subtracts the empty-request baseline", async () => { + const segments: ContextTokenCountSegment[] = [ + { id: "empty", request: { messages: [] } }, + { id: "tools-field", request: { messages: [], tools: [{ type: "function" }] } }, + { id: "tools-json", request: { messages: [{ role: "system", content: "tools" }] } }, + { id: "conversation", request: { messages: [{ role: "user", content: "hello" }] } }, + { id: "user", request: { messages: [{ role: "user", content: "hello" }] } }, + ]; + vi.stubGlobal( + "fetch", + vi.fn(async (_url: string, init: { body: string }) => { + const body = JSON.parse(init.body) as { messages: unknown[]; tools?: unknown[] }; + if (body.tools) return { ok: true, json: async () => ({ data: { total_tokens: 102 } }) }; + if (body.messages.length === 0) return { ok: true, json: async () => ({ data: { total_tokens: 2 } }) }; + return { ok: true, json: async () => ({ data: { total_tokens: body.messages.length > 1 ? 500 : 100 } }) }; + }), + ); + + const counts = await countContextTokenCounts({ + model: model("openai-completions", "https://api.stepfun.com/v1"), + auth, + segments, + }); + + expect(counts.groups).toMatchObject({ instructions: null, tools: 100, conversation: 100 }); + expect(counts.rows).toMatchObject({ user: 100 }); + }); +}); From c0bb5c6701a2dc4a12ae69a179a6f6f386d500ca Mon Sep 17 00:00:00 2001 From: li-xiu-qi Date: Tue, 27 Oct 2026 01:56:51 +0800 Subject: [PATCH 3/3] fix(coding-agent): keep /context token counts all-or-nothing Borrow the Step 3.7 tokenizer for non-Step models and fall back to local structural estimates whenever Step tokenization is unavailable or any segment fails. Remove the mixed source label so the report header is either token-count API or estimated. --- .../coding-agent/src/step/context-report.ts | 38 +++++------- packages/coding-agent/src/step/token-count.ts | 36 ++++++++++-- .../coding-agent/test/context-report.test.ts | 2 +- .../coding-agent/test/token-count.test.ts | 58 +++++++++++++++++-- 4 files changed, 99 insertions(+), 35 deletions(-) diff --git a/packages/coding-agent/src/step/context-report.ts b/packages/coding-agent/src/step/context-report.ts index c297f910..7e0fecbd 100644 --- a/packages/coding-agent/src/step/context-report.ts +++ b/packages/coding-agent/src/step/context-report.ts @@ -7,6 +7,7 @@ import type { BuildSystemPromptOptions } from "../core/system-prompt.ts"; import { getStepMcpStatuses } from "./mcp.ts"; import { countStepTokens, + resolveStepTokenCountModel, type StepTokenCountApi, type StepTokenCountRequest, supportsStepTokenCount, @@ -28,7 +29,7 @@ export interface ContextReport { freeTokens: number | null; groups: ContextReportRow[]; /** Where the displayed totals come from when the host can count tokens. */ - tokenSource: "estimate" | "provider" | "mixed"; + tokenSource: "estimate" | "provider"; } export type ContextGroupId = "instructions" | "tools" | "conversation"; @@ -556,11 +557,15 @@ export async function countContextTokenCounts(options: { signal?: AbortSignal; }): Promise { if (!supportsStepTokenCount(options.model) || !options.auth.ok) return {}; + if (options.segments.length === 0) return {}; const entries = await Promise.all( options.segments.map( async (segment) => [segment.id, await countStepTokens({ ...options, request: segment.request })] as const, ), ); + // Tokenizer counts are all-or-nothing: a partial response would mix exact and + // estimated totals, so fall back to the local structural estimator instead. + if (entries.some(([, count]) => count === null)) return {}; const counts = new Map(entries); const empty = counts.get("empty") ?? 0; const toolsField = counts.get("tools-field"); @@ -576,16 +581,12 @@ export async function countContextTokenCounts(options: { const count = counts.get(category); if (count !== null && count !== undefined) rows[category] = count; } - const conversation = counts.get("conversation") ?? null; + const rowSum = CONVERSATION_CATEGORY_IDS.reduce((total, category) => total + (rows[category] ?? 0), 0); const groups: ContextTokenCounts["groups"] = { - instructions: counts.get("instructions") ?? null, - tools: toolsTokens, - conversation, + instructions: counts.get("instructions") ?? 0, + tools: toolsTokens ?? 0, + conversation: counts.get("conversation") ?? rowSum, }; - if (groups.conversation === null) { - const rowSum = CONVERSATION_CATEGORY_IDS.reduce((total, category) => total + (rows[category] ?? 0), 0); - if (rowSum > 0) groups.conversation = rowSum; - } return { groups, rows }; } @@ -687,14 +688,7 @@ export function buildContextReport(input: ContextReportInput): ContextReport { usedTokens !== null && input.contextWindow > 0 ? Math.max(0, input.contextWindow - usedTokens - (input.reserveTokens ?? 0)) : null; - const hasAllGroupCounts = (["instructions", "tools", "conversation"] as ContextGroupId[]).every( - (id) => groupCounts[id] !== null && groupCounts[id] !== undefined, - ); - const tokenSource: ContextReport["tokenSource"] = !hasProviderCounts - ? "estimate" - : hasAllGroupCounts - ? "provider" - : "mixed"; + const tokenSource: ContextReport["tokenSource"] = hasProviderCounts ? "provider" : "estimate"; return { contextWindow: input.contextWindow, usedTokens, @@ -726,8 +720,8 @@ export async function captureContextReport(ctx: ExtensionCommandContext): Promis reportedTokens: usage?.tokens ?? null, }; - const model = ctx.model; - if (model && supportsStepTokenCount(model)) { + const model = resolveStepTokenCountModel(ctx.modelRegistry, ctx.model); + if (model) { const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model); if (auth.ok) { const segments = buildContextTokenCountSegments(input, model.api); @@ -756,11 +750,7 @@ export function contextReportRows(report: ContextReport, showAll = false): strin const window = report.contextWindow > 0 ? report.contextWindow.toLocaleString("en-US") : "--"; const percent = report.percent === null ? "--" : `${report.percent.toFixed(1)}%`; const sourceLabel = - report.tokenSource === "provider" - ? "Context Usage (token-count API)" - : report.tokenSource === "mixed" - ? "Context Usage (mixed)" - : "Context Usage (estimated)"; + report.tokenSource === "provider" ? "Context Usage (token-count API)" : "Context Usage (estimated)"; const rows = [sourceLabel, `${used} / ${window} tokens (${percent})`]; for (const group of report.groups) { rows.push("", formatRow(group, report.usedTokens)); diff --git a/packages/coding-agent/src/step/token-count.ts b/packages/coding-agent/src/step/token-count.ts index ad8ca3a6..f6a4eb4a 100644 --- a/packages/coding-agent/src/step/token-count.ts +++ b/packages/coding-agent/src/step/token-count.ts @@ -1,9 +1,16 @@ import type { Api, Model } from "@step-harness/providers"; -import type { ResolvedRequestAuth } from "../core/model-registry.ts"; +import type { ModelRegistry, ResolvedRequestAuth } from "../core/model-registry.ts"; /** Provider wire formats whose Step token-count contract is known. */ export type StepTokenCountApi = "openai-completions" | "anthropic-messages"; +/** + * Tokenizer model used when the active model belongs to another provider. Step + * exposes one tokenizer contract for these models, so `/context` can still count + * third-party prompts without a provider-specific tokenizer. + */ +export const STEP_TOKENIZER_MODEL_ID = "step-3.7-flash"; + export interface StepTokenCountRequest { messages: readonly unknown[]; system?: string; @@ -20,9 +27,11 @@ export interface StepTokenCountParams { const STEP_TOKEN_COUNT_HOSTS = new Set(["api.stepfun.com", "api.stepfun.ai"]); /** - * Token counting sends the current request shape to the provider's tokenizer - * endpoint, so it is restricted to Step endpoints. Custom providers keep the - * local structural estimate instead of leaking prompts to an unrelated host. + * Token counting sends the current request shape to the Step tokenizer endpoint. + * The active provider may be a third-party model, but tokenization still uses a + * Step tokenizer model so `/context` does not need a provider-specific counter. + * Local estimation remains the fallback when Step auth or the endpoint is + * unavailable. */ export function supportsStepTokenCount(model: Model | undefined): model is Model { if (!model) return false; @@ -34,6 +43,25 @@ export function supportsStepTokenCount(model: Model | undefined): model is } } +/** + * Pick the model whose tokenizer endpoint should count the report. + * + * A Step model keeps its own tokenizer. Any other active model borrows a Step + * tokenizer model, reusing the discovered Step endpoint and auth. When no Step + * model is registered, callers fall back to the local structural estimator. + */ +export function resolveStepTokenCountModel( + registry: ModelRegistry, + activeModel: Model | undefined, +): Model | undefined { + if (activeModel && supportsStepTokenCount(activeModel)) return activeModel; + const stepModel = registry + .getAll() + .find((model): model is Model => model.provider === "step" && supportsStepTokenCount(model)); + if (!stepModel) return undefined; + return { ...stepModel, id: STEP_TOKENIZER_MODEL_ID, name: STEP_TOKENIZER_MODEL_ID }; +} + function joinUrl(baseUrl: string, path: string): string { return `${baseUrl.replace(/\/+$/, "")}/${path.replace(/^\/+/, "")}`; } diff --git a/packages/coding-agent/test/context-report.test.ts b/packages/coding-agent/test/context-report.test.ts index 3bf7249e..54d41e64 100644 --- a/packages/coding-agent/test/context-report.test.ts +++ b/packages/coding-agent/test/context-report.test.ts @@ -290,7 +290,7 @@ describe("current context report", () => { }, }, }); - expect(report.tokenSource).toBe("mixed"); + expect(report.tokenSource).toBe("provider"); expect(group(report, "conversation").tokens).toBe(800); expect(child(report.groups, "user")!.tokens).toBe(100); expect(child(report.groups, "toolResults")!.tokens).toBe(300); diff --git a/packages/coding-agent/test/token-count.test.ts b/packages/coding-agent/test/token-count.test.ts index 72eb5c38..9dc1070e 100644 --- a/packages/coding-agent/test/token-count.test.ts +++ b/packages/coding-agent/test/token-count.test.ts @@ -7,14 +7,14 @@ import { type ContextTokenCountSegment, countContextTokenCounts, } from "../src/step/context-report.ts"; -import { countStepTokens, supportsStepTokenCount } from "../src/step/token-count.ts"; +import { countStepTokens, resolveStepTokenCountModel, supportsStepTokenCount } from "../src/step/token-count.ts"; -function model(api: Api, baseUrl: string): Model { +function model(api: Api, baseUrl: string, provider = "step", id = "step-5-preview"): Model { return { - id: "step-5-preview", - name: "Step 5 Preview", + id, + name: provider === "step" ? "Step 5 Preview" : "Other model", api, - provider: "step", + provider, baseUrl, reasoning: false, input: ["text"], @@ -77,6 +77,27 @@ describe("step token count endpoint", () => { expect(supportsStepTokenCount(undefined)).toBe(false); }); + test("borrows the Step tokenizer for non-Step providers", () => { + const registry = { + getAll: () => [ + model("openai-completions", "https://api.openai.com/v1", "openai", "gpt-5"), + model("openai-completions", "https://api.stepfun.com/v1", "step", "step-5-preview"), + ], + } as unknown as Parameters[0]; + + const active = model("openai-completions", "https://api.openai.com/v1", "openai", "gpt-5"); + const resolved = resolveStepTokenCountModel(registry, active); + + expect(resolved?.id).toBe("step-3.7-flash"); + expect(resolved?.baseUrl).toBe("https://api.stepfun.com/v1"); + expect( + resolveStepTokenCountModel( + { getAll: () => [] } as unknown as Parameters[0], + active, + ), + ).toBeUndefined(); + }); + test("parses the OpenAI-compatible token count response", async () => { const fetchMock = vi.fn(async () => ({ ok: true, json: async () => ({ data: { total_tokens: 42 } }) })); vi.stubGlobal("fetch", fetchMock); @@ -164,7 +185,32 @@ describe("context token-count segments", () => { segments, }); - expect(counts.groups).toMatchObject({ instructions: null, tools: 100, conversation: 100 }); + expect(counts.groups).toMatchObject({ instructions: 0, tools: 100, conversation: 100 }); expect(counts.rows).toMatchObject({ user: 100 }); }); + + test("falls back to local estimation when any segment fails", async () => { + const segments: ContextTokenCountSegment[] = [ + { id: "empty", request: { messages: [] } }, + { id: "instructions", request: { messages: [{ role: "system", content: "system" }] } }, + { id: "conversation", request: { messages: [{ role: "user", content: "hello" }] } }, + { id: "user", request: { messages: [{ role: "user", content: "hello" }] } }, + ]; + vi.stubGlobal( + "fetch", + vi.fn(async (_url: string, init: { body: string }) => { + const body = JSON.parse(init.body) as { messages: Array<{ content?: string }> }; + if (body.messages[0]?.content === "system") return { ok: false, json: async () => ({}) }; + return { ok: true, json: async () => ({ data: { total_tokens: 10 } }) }; + }), + ); + + const counts = await countContextTokenCounts({ + model: model("openai-completions", "https://api.stepfun.com/v1"), + auth, + segments, + }); + + expect(counts).toEqual({}); + }); });