Skip to content

Emit token usage metadata in ACP notifications #90

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Emit token usage metadata in ACP notifications #90

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 44 additions & 6 deletions src/conversation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,11 +33,11 @@ use codex_core::{
ExecCommandOutputDeltaEvent, ExitedReviewModeEvent, FileChange, ItemCompletedEvent,
ItemStartedEvent, ListCustomPromptsResponseEvent, McpInvocation, McpStartupCompleteEvent,
McpStartupUpdateEvent, McpToolCallBeginEvent, McpToolCallEndEvent, Op,
PatchApplyBeginEvent, PatchApplyEndEvent, ReasoningContentDeltaEvent,
PatchApplyBeginEvent, PatchApplyEndEvent, RateLimitSnapshot, ReasoningContentDeltaEvent,
ReasoningRawContentDeltaEvent, ReviewDecision, ReviewOutputEvent, ReviewRequest,
SandboxPolicy, StreamErrorEvent, TaskCompleteEvent, TaskStartedEvent, TurnAbortedEvent,
UserMessageEvent, ViewImageToolCallEvent, WarningEvent, WebSearchBeginEvent,
WebSearchEndEvent,
SandboxPolicy, StreamErrorEvent, TaskCompleteEvent, TaskStartedEvent, TokenCountEvent,
TokenUsageInfo, TurnAbortedEvent, UserMessageEvent, ViewImageToolCallEvent, WarningEvent,
WebSearchBeginEvent, WebSearchEndEvent,
},
review_format::format_review_findings_block,
};
Expand Down Expand Up @@ -541,11 +541,16 @@ impl PromptState {
"MCP startup complete: ready={ready:?}, failed={failed:?}, cancelled={cancelled:?}"
);
}
EventMsg::TokenCount(TokenCountEvent { info, rate_limits }) => {
if let Some(info) = info {
client
.send_token_usage(info, rate_limits)
.await;
}
}

// Ignore these events
EventMsg::AgentReasoningRawContent(..)
// In the future we can use this to update usage stats
| EventMsg::TokenCount(..)
// we already have a way to diff the turn, so ignore
| EventMsg::TurnDiff(..)
// Revisit when we can emit status updates
Expand Down Expand Up @@ -1396,6 +1401,39 @@ impl SessionClient {
.await;
}

async fn send_token_usage(
&self,
info: TokenUsageInfo,
rate_limits: Option<RateLimitSnapshot>,
) {
let mut meta = serde_json::Map::new();
meta.insert(
"token_usage".to_string(),
serde_json::to_value(info).unwrap_or(serde_json::Value::Null),
);
if let Some(limits) = rate_limits {
meta.insert(
"rate_limits".to_string(),
serde_json::to_value(limits).unwrap_or(serde_json::Value::Null),
);
}
let notification = SessionNotification {
session_id: self.session_id.clone(),
update: SessionUpdate::AgentMessageChunk(ContentChunk {
content: ContentBlock::Text(TextContent {
text: String::new(),
annotations: None,
meta: None,
}),
Comment on lines +1422 to +1427
Copy link
Member

@benbrandt benbrandt Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm we shouldn't need to append an empty chunk ideally...
This is something on the todo list to support at the protocol level in someway, and I think this highlights the issue that there isn't a great way to send it.

Another option might be to send it on the turn end with the stop reason? Like store it on the state of the turn and submit it when we finish?

But I haven't seen how often we get these. Do you have an idea of how often these events happen?

meta: Some(serde_json::Value::Object(meta)),
}),
meta: None,
};
if let Err(e) = self.client.session_notification(notification).await {
error!("Failed to send session notification: {:?}", e);
}
}

async fn update_plan(&self, plan: Vec<PlanItemArg>) {
self.send_notification(SessionUpdate::Plan(Plan {
entries: plan
Expand Down