1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
|
import type { Message, Part } from "@opencode-ai/sdk/v2/client"
export type SessionContextBreakdownKey = "system" | "user" | "assistant" | "tool" | "other"
export type SessionContextBreakdownSegment = {
key: SessionContextBreakdownKey
tokens: number
width: number
percent: number
}
const estimateTokens = (chars: number) => Math.ceil(chars / 4)
const toPercent = (tokens: number, input: number) => (tokens / input) * 100
const toPercentLabel = (tokens: number, input: number) => Math.round(toPercent(tokens, input) * 10) / 10
const charsFromUserPart = (part: Part) => {
if (part.type === "text") return part.text.length
if (part.type === "file") return part.source?.text.value.length ?? 0
if (part.type === "agent") return part.source?.value.length ?? 0
return 0
}
const charsFromAssistantPart = (part: Part) => {
if (part.type === "text") return { assistant: part.text.length, tool: 0 }
if (part.type === "reasoning") return { assistant: part.text.length, tool: 0 }
if (part.type !== "tool") return { assistant: 0, tool: 0 }
const input = Object.keys(part.state.input).length * 16
if (part.state.status === "pending") return { assistant: 0, tool: input + part.state.raw.length }
if (part.state.status === "completed") return { assistant: 0, tool: input + part.state.output.length }
if (part.state.status === "error") return { assistant: 0, tool: input + part.state.error.length }
return { assistant: 0, tool: input }
}
const build = (
tokens: { system: number; user: number; assistant: number; tool: number; other: number },
input: number,
) => {
return [
{
key: "system",
tokens: tokens.system,
},
{
key: "user",
tokens: tokens.user,
},
{
key: "assistant",
tokens: tokens.assistant,
},
{
key: "tool",
tokens: tokens.tool,
},
{
key: "other",
tokens: tokens.other,
},
]
.filter((x) => x.tokens > 0)
.map((x) => ({
key: x.key,
tokens: x.tokens,
width: toPercent(x.tokens, input),
percent: toPercentLabel(x.tokens, input),
})) as SessionContextBreakdownSegment[]
}
export function estimateSessionContextBreakdown(args: {
messages: Message[]
parts: Record<string, Part[] | undefined>
input: number
systemPrompt?: string
}) {
if (!args.input) return []
const counts = args.messages.reduce(
(acc, msg) => {
const parts = args.parts[msg.id] ?? []
if (msg.role === "user") {
const user = parts.reduce((sum, part) => sum + charsFromUserPart(part), 0)
return { ...acc, user: acc.user + user }
}
if (msg.role !== "assistant") return acc
const assistant = parts.reduce(
(sum, part) => {
const next = charsFromAssistantPart(part)
return {
assistant: sum.assistant + next.assistant,
tool: sum.tool + next.tool,
}
},
{ assistant: 0, tool: 0 },
)
return {
...acc,
assistant: acc.assistant + assistant.assistant,
tool: acc.tool + assistant.tool,
}
},
{
system: args.systemPrompt?.length ?? 0,
user: 0,
assistant: 0,
tool: 0,
},
)
const tokens = {
system: estimateTokens(counts.system),
user: estimateTokens(counts.user),
assistant: estimateTokens(counts.assistant),
tool: estimateTokens(counts.tool),
}
const estimated = tokens.system + tokens.user + tokens.assistant + tokens.tool
if (estimated <= args.input) {
return build({ ...tokens, other: args.input - estimated }, args.input)
}
const scale = args.input / estimated
const scaled = {
system: Math.floor(tokens.system * scale),
user: Math.floor(tokens.user * scale),
assistant: Math.floor(tokens.assistant * scale),
tool: Math.floor(tokens.tool * scale),
}
const total = scaled.system + scaled.user + scaled.assistant + scaled.tool
return build({ ...scaled, other: Math.max(0, args.input - total) }, args.input)
}
|
