1import type { TurnTokenUsage } from "@t3tools/contracts";
2import type * as EffectCodexSchema from "effect-codex-app-server/schema";
3
4type CodexCumulativeTokenUsage = {
5 readonly inputTokens: number;
6 readonly cachedInputTokens: number;
7 readonly cacheCreationTokens?: number;
8 readonly outputTokens: number;
9 readonly reasoningTokens: number;
10};
11
12interface CodexTurnTokenUsageAccumulator {
13 inputTokens: number;
14 cachedInputTokens: number;
15 cacheCreationTokens: number | undefined;
16 outputTokens: number;
17 reasoningTokens: number;
18 observed: boolean;
19 hasSubagents: boolean;
20}
21
22export interface CodexTurnTokenUsageState {
23 baseline: CodexCumulativeTokenUsage | undefined;
24 activeTurnId: string | undefined;
25 readonly byTurnId: Map<string, CodexTurnTokenUsageAccumulator>;
26}
27
28function codexTokenUsageBreakdown(
29 usage: EffectCodexSchema.V2ThreadTokenUsageUpdatedNotification__TokenUsageBreakdown,
30): CodexCumulativeTokenUsage {
31 return {
32 inputTokens: usage.inputTokens,
33 cachedInputTokens: usage.cachedInputTokens,
34 ...(usage.cacheWriteInputTokens !== undefined
35 ? { cacheCreationTokens: usage.cacheWriteInputTokens }
36 : {}),
37 outputTokens: usage.outputTokens,
38 reasoningTokens: usage.reasoningOutputTokens,
39 };
40}
41
42export function makeCodexTurnTokenUsageState(): CodexTurnTokenUsageState {
43 return {
44 baseline: undefined,
45 activeTurnId: undefined,
46 byTurnId: new Map(),
47 };
48}
49
50export function getCodexTurnAccumulator(
51 state: CodexTurnTokenUsageState,
52 turnId: string,
53): CodexTurnTokenUsageAccumulator {
54 const existing = state.byTurnId.get(turnId);
55 if (existing) return existing;
56 const created: CodexTurnTokenUsageAccumulator = {
57 inputTokens: 0,
58 cachedInputTokens: 0,
59 cacheCreationTokens: 0,
60 outputTokens: 0,
61 reasoningTokens: 0,
62 observed: false,
63 hasSubagents: false,
64 };
65 state.byTurnId.set(turnId, created);
66 return created;
67}
68
69/**
70 * Usage added by one `thread/tokenUsage/updated` notification. Codex reports a
71 * running `total` for the thread and `last`, the usage of the newest model
72 * response. Within a turn the growth of `total` equals `last`. Without a prior
73 * total (first update after resume or rollback), or when Codex reset the
74 * running total, `last` is the delta.
75 */
76function codexTurnTokenUsageDelta(
77 previous: CodexCumulativeTokenUsage | undefined,
78 current: CodexCumulativeTokenUsage,
79 last: CodexCumulativeTokenUsage,
80): CodexCumulativeTokenUsage {
81 if (
82 previous === undefined ||
83 current.inputTokens < previous.inputTokens ||
84 current.cachedInputTokens < previous.cachedInputTokens ||
85 current.outputTokens < previous.outputTokens ||
86 current.reasoningTokens < previous.reasoningTokens
87 ) {
88 return last;
89 }
90 return {
91 inputTokens: current.inputTokens - previous.inputTokens,
92 cachedInputTokens: current.cachedInputTokens - previous.cachedInputTokens,
93 ...(current.cacheCreationTokens !== undefined &&
94 previous.cacheCreationTokens !== undefined &&
95 current.cacheCreationTokens >= previous.cacheCreationTokens
96 ? { cacheCreationTokens: current.cacheCreationTokens - previous.cacheCreationTokens }
97 : {}),
98 outputTokens: current.outputTokens - previous.outputTokens,
99 reasoningTokens: current.reasoningTokens - previous.reasoningTokens,
100 };
101}
102
103export function accumulateCodexTurnTokenUsage(
104 state: CodexTurnTokenUsageState,
105 turnId: string,
106 usage: EffectCodexSchema.V2ThreadTokenUsageUpdatedNotification["tokenUsage"],
107): void {
108 const current = codexTokenUsageBreakdown(usage.total);
109 if (state.activeTurnId !== turnId) {
110 // The total is thread-wide, so every update moves the baseline. A late
111 // update for a finished turn is not counted toward the live turn.
112 state.baseline = current;
113 return;
114 }
115
116 const accumulator = getCodexTurnAccumulator(state, turnId);
117 const delta = codexTurnTokenUsageDelta(
118 state.baseline,
119 current,
120 codexTokenUsageBreakdown(usage.last),
121 );
122 state.baseline = current;
123
124 if (
125 delta.inputTokens > 0 ||
126 delta.cachedInputTokens > 0 ||
127 delta.outputTokens > 0 ||
128 delta.reasoningTokens > 0
129 ) {
130 accumulator.observed = true;
131 }
132 accumulator.inputTokens += delta.inputTokens;
133 accumulator.cachedInputTokens += delta.cachedInputTokens;
134 accumulator.outputTokens += delta.outputTokens;
135 accumulator.reasoningTokens += delta.reasoningTokens;
136 if (delta.cacheCreationTokens === undefined) {
137 accumulator.cacheCreationTokens = undefined;
138 } else if (accumulator.cacheCreationTokens !== undefined) {
139 accumulator.cacheCreationTokens += delta.cacheCreationTokens;
140 }
141}
142
143export function completeCodexTurnTokenUsage(
144 state: CodexTurnTokenUsageState,
145 turnId: string,
146 completed: boolean,
147): TurnTokenUsage {
148 const usage = state.byTurnId.get(turnId);
149 state.byTurnId.delete(turnId);
150 if (state.activeTurnId === turnId) state.activeTurnId = undefined;
151 if (!usage) {
152 return {
153 usageStatus: "unavailable",
154 usageScope: "main_agent",
155 hasSubagents: false,
156 };
157 }
158
159 if (!usage.observed) {
160 return {
161 usageStatus: "unavailable",
162 usageScope: "main_agent",
163 hasSubagents: usage.hasSubagents,
164 };
165 }
166
167 // Codex counts cache reads and writes inside inputTokens. Clamp the
168 // subsets so the record keeps the documented relationships even if a
169 // counter drifts.
170 return {
171 usageStatus: completed ? "complete" : "partial",
172 usageScope: "main_agent",
173 inputTokens: usage.inputTokens,
174 cachedInputTokens: Math.min(usage.inputTokens, usage.cachedInputTokens),
175 ...(usage.cacheCreationTokens !== undefined
176 ? { cacheCreationTokens: Math.min(usage.inputTokens, usage.cacheCreationTokens) }
177 : {}),
178 outputTokens: usage.outputTokens,
179 reasoningTokens: Math.min(usage.outputTokens, usage.reasoningTokens),
180 hasSubagents: usage.hasSubagents,
181 };
182}