Skip to content

Commit 17bc214

Browse files
TarquinenTarquinen
authored
Merge pull request #292 from Opencode-DCP/feature/context-token-calculation-optimization
refactor: optimize token calculation in context command
2 parents d476e03 + 08e537d commit 17bc214

File tree

1 file changed

+86
-73
lines changed

1 file changed

+86
-73
lines changed

‎lib/commands/context.ts‎

Lines changed: 86 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,40 @@
11
/**
2-
* DCP Context command handler.
2+
* DCP Context Command
33
* Shows a visual breakdown of token usage in the current session.
4+
*
5+
* TOKEN CALCULATION STRATEGY
6+
* ==========================
7+
* We minimize tokenizer estimation by leveraging API-reported values wherever possible.
8+
*
9+
* WHAT WE GET FROM THE API (exact):
10+
* - tokens.input : Input tokens for each assistant response
11+
* - tokens.output : Output tokens generated (includes text + tool calls)
12+
* - tokens.reasoning: Reasoning tokens used
13+
* - tokens.cache : Cache read/write tokens
14+
*
15+
* HOW WE CALCULATE EACH CATEGORY:
16+
*
17+
* SYSTEM = firstAssistant.input + cache.read - tokenizer(firstUserMessage)
18+
* The first response's input contains system + first user message.
19+
*
20+
* TOOLS = tokenizer(toolInputs + toolOutputs) - prunedTokens
21+
* We must tokenize tools anyway for pruning decisions.
22+
*
23+
* USER = tokenizer(all user messages)
24+
* User messages are typically small, so estimation is acceptable.
25+
*
26+
* ASSISTANT = total - system - user - tools
27+
* Calculated as residual. This absorbs:
28+
* - Assistant text output tokens
29+
* - Reasoning tokens (if persisted by the model)
30+
* - Any estimation errors
31+
*
32+
* TOTAL = input + output + reasoning + cache.read + cache.write
33+
* Matches opencode's UI display.
34+
*
35+
* WHY ASSISTANT IS THE RESIDUAL:
36+
* If reasoning tokens persist in context (model-dependent), they semantically
37+
* belong with "Assistant" since reasoning IS assistant-generated content.
438
*/
539

640
import type { Logger } from "../logger"
@@ -24,7 +58,6 @@ interface TokenBreakdown {
2458
system: number
2559
user: number
2660
assistant: number
27-
reasoning: number
2861
tools: number
2962
prunedTokens: number
3063
prunedCount: number
@@ -36,7 +69,6 @@ function analyzeTokens(state: SessionState, messages: WithParts[]): TokenBreakdo
3669
system: 0,
3770
user: 0,
3871
assistant: 0,
39-
reasoning: 0,
4072
tools: 0,
4173
prunedTokens: state.stats.totalPruneTokens,
4274
prunedCount: state.prune.toolIds.length,
@@ -54,26 +86,6 @@ function analyzeTokens(state: SessionState, messages: WithParts[]): TokenBreakdo
5486
}
5587
}
5688

57-
let firstUserTokens = 0
58-
for (const msg of messages) {
59-
if (msg.info.role === "user" && !isIgnoredUserMessage(msg)) {
60-
for (const part of msg.parts) {
61-
if (part.type === "text") {
62-
const textPart = part as TextPart
63-
firstUserTokens += countTokens(textPart.text || "")
64-
}
65-
}
66-
break
67-
}
68-
}
69-
70-
// Calculate system tokens: first response's total input minus first user message
71-
if (firstAssistant) {
72-
const firstInput =
73-
(firstAssistant.tokens?.input || 0) + (firstAssistant.tokens?.cache?.read || 0)
74-
breakdown.system = Math.max(0, firstInput - firstUserTokens)
75-
}
76-
7789
let lastAssistant: AssistantMessage | undefined
7890
for (let i = messages.length - 1; i >= 0; i--) {
7991
const msg = messages[i]
@@ -86,71 +98,73 @@ function analyzeTokens(state: SessionState, messages: WithParts[]): TokenBreakdo
8698
}
8799
}
88100

89-
// Get total from API
90-
// Total = input + output + reasoning + cache.read + cache.write
91101
const apiInput = lastAssistant?.tokens?.input || 0
92102
const apiOutput = lastAssistant?.tokens?.output || 0
93103
const apiReasoning = lastAssistant?.tokens?.reasoning || 0
94104
const apiCacheRead = lastAssistant?.tokens?.cache?.read || 0
95105
const apiCacheWrite = lastAssistant?.tokens?.cache?.write || 0
96-
const apiTotal = apiInput + apiOutput + apiReasoning + apiCacheRead + apiCacheWrite
97-
98-
for (const msg of messages) {
99-
if (isMessageCompacted(state, msg)) {
100-
continue
101-
}
106+
breakdown.total = apiInput + apiOutput + apiReasoning + apiCacheRead + apiCacheWrite
102107

103-
if (msg.info.role === "user" && isIgnoredUserMessage(msg)) {
104-
continue
105-
}
108+
const userTextParts: string[] = []
109+
const toolInputParts: string[] = []
110+
const toolOutputParts: string[] = []
111+
let firstUserText = ""
112+
let foundFirstUser = false
106113

107-
const info = msg.info
108-
const role = info.role
114+
for (const msg of messages) {
115+
if (isMessageCompacted(state, msg)) continue
116+
if (msg.info.role === "user" && isIgnoredUserMessage(msg)) continue
109117

110118
for (const part of msg.parts) {
111-
switch (part.type) {
112-
case "text": {
113-
const textPart = part as TextPart
114-
const tokens = countTokens(textPart.text || "")
115-
if (role === "user") {
116-
breakdown.user += tokens
117-
} else {
118-
breakdown.assistant += tokens
119-
}
120-
break
119+
if (part.type === "text" && msg.info.role === "user") {
120+
const textPart = part as TextPart
121+
const text = textPart.text || ""
122+
userTextParts.push(text)
123+
if (!foundFirstUser) {
124+
firstUserText += text
121125
}
122-
case "tool": {
123-
const toolPart = part as ToolPart
124-
125-
if (toolPart.state?.input) {
126-
const inputStr =
127-
typeof toolPart.state.input === "string"
128-
? toolPart.state.input
129-
: JSON.stringify(toolPart.state.input)
130-
breakdown.tools += countTokens(inputStr)
131-
}
132-
133-
if (toolPart.state?.status === "completed" && toolPart.state?.output) {
134-
const outputStr =
135-
typeof toolPart.state.output === "string"
136-
? toolPart.state.output
137-
: JSON.stringify(toolPart.state.output)
138-
breakdown.tools += countTokens(outputStr)
139-
}
140-
break
126+
} else if (part.type === "tool") {
127+
const toolPart = part as ToolPart
128+
129+
if (toolPart.state?.input) {
130+
const inputStr =
131+
typeof toolPart.state.input === "string"
132+
? toolPart.state.input
133+
: JSON.stringify(toolPart.state.input)
134+
toolInputParts.push(inputStr)
135+
}
136+
137+
if (toolPart.state?.status === "completed" && toolPart.state?.output) {
138+
const outputStr =
139+
typeof toolPart.state.output === "string"
140+
? toolPart.state.output
141+
: JSON.stringify(toolPart.state.output)
142+
toolOutputParts.push(outputStr)
141143
}
142144
}
143145
}
146+
147+
if (msg.info.role === "user" && !isIgnoredUserMessage(msg) && !foundFirstUser) {
148+
foundFirstUser = true
149+
}
144150
}
145151

146-
breakdown.tools = Math.max(0, breakdown.tools - breakdown.prunedTokens)
152+
const firstUserTokens = countTokens(firstUserText)
153+
breakdown.user = countTokens(userTextParts.join("\n"))
154+
const toolInputTokens = countTokens(toolInputParts.join("\n"))
155+
const toolOutputTokens = countTokens(toolOutputParts.join("\n"))
147156

148-
// Calculate reasoning as the difference between API total and our counted parts
149-
// This handles both interleaved thinking and non-interleaved models correctly
150-
const countedParts = breakdown.system + breakdown.user + breakdown.assistant + breakdown.tools
151-
breakdown.reasoning = Math.max(0, apiTotal - countedParts)
157+
if (firstAssistant) {
158+
const firstInput =
159+
(firstAssistant.tokens?.input || 0) + (firstAssistant.tokens?.cache?.read || 0)
160+
breakdown.system = Math.max(0, firstInput - firstUserTokens)
161+
}
152162

153-
breakdown.total = apiTotal
163+
breakdown.tools = Math.max(0, toolInputTokens + toolOutputTokens - breakdown.prunedTokens)
164+
breakdown.assistant = Math.max(
165+
0,
166+
breakdown.total - breakdown.system - breakdown.user - breakdown.tools,
167+
)
154168

155169
return breakdown
156170
}
@@ -170,8 +184,7 @@ function formatContextMessage(breakdown: TokenBreakdown): string {
170184
{ label: "System", value: breakdown.system, char: "█" },
171185
{ label: "User", value: breakdown.user, char: "▓" },
172186
{ label: "Assistant", value: breakdown.assistant, char: "▒" },
173-
{ label: "Reasoning", value: breakdown.reasoning, char: "░" },
174-
{ label: "Tools", value: breakdown.tools, char: "⣿" },
187+
{ label: "Tools", value: breakdown.tools, char: "░" },
175188
] as const
176189

177190
lines.push("╭───────────────────────────────────────────────────────────╮")

0 commit comments

Comments
 (0)