Repository navigation
fix: report subagent token usage in _meta.quota.model_usage #706
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | |||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,133 @@ | |||||||||||||||||
| import type {ServerNotification} from "./app-server"; | |||||||||||||||||
| import {isModelOutputItem} from "./CodexEventHandler"; | |||||||||||||||||
| import type {ModelTokenCount} from "./QuotaMeta"; | |||||||||||||||||
| import {isRootAgentPath} from "./subagents/CodexAgentPath"; | |||||||||||||||||
| import {type ChildModel, preferredModel} from "./subagents/CodexSubagentEventRouter"; | |||||||||||||||||
| import { | |||||||||||||||||
| addTokenCounts, | |||||||||||||||||
| sameTokenCount, | |||||||||||||||||
| subtractTokenCounts, | |||||||||||||||||
| toTokenCount, | |||||||||||||||||
| type TokenCount, | |||||||||||||||||
| ZERO_TOKEN_COUNT, | |||||||||||||||||
| } from "./TokenCount"; | |||||||||||||||||
|
|
|||||||||||||||||
| type ChildThread = { | |||||||||||||||||
| model?: ChildModel; | |||||||||||||||||
| /** The last total of the thread, or null when the total it started from is not known. */ | |||||||||||||||||
| previousTotal: TokenCount | null; | |||||||||||||||||
| modelOutputSeen: boolean; | |||||||||||||||||
| /** The sum of the model requests of the thread. */ | |||||||||||||||||
| used: TokenCount; | |||||||||||||||||
| /** The part of `used` that a prompt reported. */ | |||||||||||||||||
| reported: TokenCount; | |||||||||||||||||
| /** The part of `used` that the prompts before the current one reported. */ | |||||||||||||||||
| reportedBeforePrompt: TokenCount; | |||||||||||||||||
| }; | |||||||||||||||||
|
|
|||||||||||||||||
| /** | |||||||||||||||||
| * Accumulates the token usage of the child threads of a session: its subagents and their subagents. Codex reports | |||||||||||||||||
| * the usage of every thread on its own, and the usage of the root thread does not include its children, so a prompt | |||||||||||||||||
| * reports theirs in `_meta.quota.model_usage` next to its own. | |||||||||||||||||
| * | |||||||||||||||||
| * It reads the notifications as the app-server sends them, before the session routes them: a client without | |||||||||||||||||
| * subagent sessions gets the notifications of a child attributed to the root thread. A child that is still working | |||||||||||||||||
| * when a prompt ends adds its later usage to the next prompt that starts model work. | |||||||||||||||||
| * | |||||||||||||||||
| * Codex 0.165 sends no `thread/started` for a subagent: the tracker then learns the child from its `subAgentActivity` | |||||||||||||||||
| * item and its own notifications, with an unknown start, so the usage of a child counts from its first model output. | |||||||||||||||||
| */ | |||||||||||||||||
| export class SubagentTokenUsage { | |||||||||||||||||
| private readonly threads = new Map<string, ChildThread>(); | |||||||||||||||||
|
|
|||||||||||||||||
| constructor(private readonly rootThreadId: string) {} | |||||||||||||||||
|
|
|||||||||||||||||
| observe(notification: ServerNotification): void { | |||||||||||||||||
| switch (notification.method) { | |||||||||||||||||
| case "thread/started": { | |||||||||||||||||
| const thread = notification.params.thread; | |||||||||||||||||
| const parent = thread.parentThreadId; | |||||||||||||||||
| // Every session gets every thread start, so only the children of this root count. | |||||||||||||||||
| if (thread.id === this.rootThreadId || parent === null | |||||||||||||||||
| || (parent !== this.rootThreadId && !this.threads.has(parent))) return; | |||||||||||||||||
| // A new thread starts at zero, while a forked one starts from the usage of its history. | |||||||||||||||||
| const child = this.child(thread.id, thread.forkedFromId === null ? ZERO_TOKEN_COUNT : null); | |||||||||||||||||
| if (thread.model) child.model = preferredModel(child.model, {model: thread.model, source: "threadStarted"}); | |||||||||||||||||
| return; | |||||||||||||||||
| } | |||||||||||||||||
| case "thread/settings/updated": { | |||||||||||||||||
| const threadId = notification.params.threadId; | |||||||||||||||||
| if (threadId === this.rootThreadId) return; | |||||||||||||||||
| const child = this.child(threadId, null); | |||||||||||||||||
| child.model = preferredModel(child.model, {model: notification.params.threadSettings.model, source: "settings"}); | |||||||||||||||||
| return; | |||||||||||||||||
| } | |||||||||||||||||
| case "item/started": | |||||||||||||||||
| case "item/completed": { | |||||||||||||||||
| const {threadId, item} = notification.params; | |||||||||||||||||
| if (item.type === "subAgentActivity" && item.model && !isRootAgentPath(item.agentPath)) { | |||||||||||||||||
| const child = this.child(item.agentThreadId, null); | |||||||||||||||||
| child.model = preferredModel(child.model, {model: item.model, source: "activity"}); | |||||||||||||||||
| } | |||||||||||||||||
| if (notification.method === "item/started" && threadId !== this.rootThreadId && isModelOutputItem(item)) { | |||||||||||||||||
| this.child(threadId, null).modelOutputSeen = true; | |||||||||||||||||
| } | |||||||||||||||||
| return; | |||||||||||||||||
| } | |||||||||||||||||
| case "thread/tokenUsage/updated": { | |||||||||||||||||
|
Collaborator
There was a problem hiding this comment.
The comment below even says "as in PromptTokenUsage". These rules work around Codex quirks: Codex sends the current usage again when a request starts, and a resumed or forked thread's first update includes its history. When Codex changes one of these, someone has to remember to fix both places, and nothing will fail if only one gets fixed. The two copies already differ in one place. With a known start, Suggestion: move the shared part into one per-thread class in /** The usage of one thread from its `thread/tokenUsage/updated` notifications. */
export class ThreadTokenUsage {
private modelOutputSeen = false;
/** The `last` values of the updates that were counted. */
summedLast: TokenCount = ZERO_TOKEN_COUNT;
/** @param previousTotal the thread total at the start, or null when it is not known. */
constructor(private previousTotal: TokenCount | null) {}
observeModelOutput(): void { this.modelOutputSeen = true; }
/** "repeat" and "history" updates are not counted. */
observe(total: TokenCount, last: TokenCount): "repeat" | "history" | "counted" { /* the rules above */ }
}
|
|||||||||||||||||
| const threadId = notification.params.threadId; | |||||||||||||||||
| if (threadId === this.rootThreadId) return; | |||||||||||||||||
| const child = this.child(threadId, null); | |||||||||||||||||
| const total = toTokenCount(notification.params.tokenUsage.total); | |||||||||||||||||
| const previous = child.previousTotal; | |||||||||||||||||
| // Codex sends the current usage again when a model request starts; it is not a new request. | |||||||||||||||||
| if (previous !== null && sameTokenCount(total, previous)) return; | |||||||||||||||||
| child.previousTotal = total; | |||||||||||||||||
| // Before any output, an update of a thread with an unknown start carries its history, as in | |||||||||||||||||
| // PromptTokenUsage. | |||||||||||||||||
| if (previous === null && !child.modelOutputSeen) return; | |||||||||||||||||
| child.used = addTokenCounts(child.used, toTokenCount(notification.params.tokenUsage.last)); | |||||||||||||||||
| return; | |||||||||||||||||
| } | |||||||||||||||||
| } | |||||||||||||||||
| } | |||||||||||||||||
|
|
|||||||||||||||||
| /** Starts a prompt that does model work: the child usage that no prompt reported yet belongs to it. */ | |||||||||||||||||
| beginPrompt(): void { | |||||||||||||||||
| for (const thread of this.threads.values()) thread.reportedBeforePrompt = thread.reported; | |||||||||||||||||
| } | |||||||||||||||||
|
|
|||||||||||||||||
| /** | |||||||||||||||||
| * The `model_usage` rows of the current prompt: the usage of the root thread, then the usage of the children by | |||||||||||||||||
| * model, with the same model in one row. The child usage counts as reported. A child whose model Codex did not | |||||||||||||||||
| * report is counted for the root model, which a spawn that names no model inherits. | |||||||||||||||||
| */ | |||||||||||||||||
| reportModelUsage(rootModel: string, rootUsage: TokenCount | null): ModelTokenCount[] { | |||||||||||||||||
| const byModel = new Map<string, TokenCount>(); | |||||||||||||||||
| if (rootUsage !== null) byModel.set(rootModel, rootUsage); | |||||||||||||||||
| for (const thread of this.threads.values()) { | |||||||||||||||||
| const usage = subtractTokenCounts(thread.used, thread.reportedBeforePrompt); | |||||||||||||||||
| thread.reported = thread.used; | |||||||||||||||||
| if (sameTokenCount(usage, ZERO_TOKEN_COUNT)) continue; | |||||||||||||||||
| const model = thread.model?.model ?? rootModel; | |||||||||||||||||
| const current = byModel.get(model); | |||||||||||||||||
| byModel.set(model, current === undefined ? usage : addTokenCounts(current, usage)); | |||||||||||||||||
| } | |||||||||||||||||
| return [...byModel].map(([model, token_count]) => ({model, token_count})); | |||||||||||||||||
| } | |||||||||||||||||
|
|
|||||||||||||||||
| private child(threadId: string, startTotal: TokenCount | null): ChildThread { | |||||||||||||||||
| let child = this.threads.get(threadId); | |||||||||||||||||
| if (child === undefined) { | |||||||||||||||||
| child = { | |||||||||||||||||
| previousTotal: startTotal, | |||||||||||||||||
| modelOutputSeen: false, | |||||||||||||||||
| used: ZERO_TOKEN_COUNT, | |||||||||||||||||
| reported: ZERO_TOKEN_COUNT, | |||||||||||||||||
| reportedBeforePrompt: ZERO_TOKEN_COUNT, | |||||||||||||||||
| }; | |||||||||||||||||
| this.threads.set(threadId, child); | |||||||||||||||||
| } | |||||||||||||||||
| return child; | |||||||||||||||||
| } | |||||||||||||||||
| } | |||||||||||||||||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,49 @@ | ||
| { | ||
| "stopReason": "end_turn", | ||
| "usage": { | ||
| "totalTokens": 1200, | ||
| "inputTokens": 600, | ||
| "outputTokens": 200, | ||
| "thoughtTokens": 0, | ||
| "cachedReadTokens": 400, | ||
| "cachedWriteTokens": 0 | ||
| }, | ||
| "_meta": { | ||
| "quota": { | ||
| "token_count": { | ||
| "totalTokens": 1200, | ||
| "inputTokens": 600, | ||
| "cachedInputTokens": 400, | ||
| "cacheWriteInputTokens": 0, | ||
| "outputTokens": 200, | ||
| "reasoningOutputTokens": 0 | ||
| }, | ||
| "model_usage": [ | ||
| { | ||
| "model": "gpt-5", | ||
| "token_count": { | ||
| "totalTokens": 1200, | ||
| "inputTokens": 600, | ||
| "cachedInputTokens": 400, | ||
| "cacheWriteInputTokens": 0, | ||
| "outputTokens": 200, | ||
| "reasoningOutputTokens": 0 | ||
| } | ||
| }, | ||
| { | ||
| "model": "gpt-5.6-luna", | ||
| "token_count": { | ||
| "totalTokens": 30000, | ||
| "inputTokens": 4500, | ||
| "cachedInputTokens": 25000, | ||
| "cacheWriteInputTokens": 0, | ||
| "outputTokens": 500, | ||
| "reasoningOutputTokens": 0 | ||
| } | ||
| } | ||
| ] | ||
| } | ||
| }, | ||
| "state": "idle", | ||
| "sessionUpdate": "state_update" | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,47 @@ | ||
| { | ||
| "stopReason": "end_turn", | ||
| "usage": { | ||
| "totalTokens": 1000, | ||
| "inputTokens": 800, | ||
| "cachedReadTokens": 0, | ||
| "cachedWriteTokens": 0, | ||
| "outputTokens": 200, | ||
| "thoughtTokens": 0 | ||
| }, | ||
| "_meta": { | ||
| "quota": { | ||
| "token_count": { | ||
| "totalTokens": 1000, | ||
| "inputTokens": 800, | ||
| "cachedInputTokens": 0, | ||
| "cacheWriteInputTokens": 0, | ||
| "outputTokens": 200, | ||
| "reasoningOutputTokens": 0 | ||
| }, | ||
| "model_usage": [ | ||
| { | ||
| "model": "gpt-5.6-sol", | ||
| "token_count": { | ||
| "totalTokens": 1000, | ||
| "inputTokens": 800, | ||
| "cachedInputTokens": 0, | ||
| "cacheWriteInputTokens": 0, | ||
| "outputTokens": 200, | ||
| "reasoningOutputTokens": 0 | ||
| } | ||
| }, | ||
| { | ||
| "model": "gpt-5.6-luna", | ||
| "token_count": { | ||
| "totalTokens": 50000, | ||
| "inputTokens": 9000, | ||
| "cachedInputTokens": 40000, | ||
| "cacheWriteInputTokens": 0, | ||
| "outputTokens": 1000, | ||
| "reasoningOutputTokens": 0 | ||
| } | ||
| } | ||
| ] | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This row should spell out a few more things a client needs:
state_updatethat ends the turn. The row only namesPromptResponse._meta.quota.model_usage.nullwith rows.token_countcan benullwhilemodel_usageis not empty. A child that kept working after the last prompt adds a row to a prompt that was cancelled before the root thread reported any usage. The last sentence currently tiesnulland[]together.Possible wording: