diff --git a/server/src/session/codex/translate.rs b/server/src/session/codex/translate.rs index edf5a0a..ebc1eb4 100644 --- a/server/src/session/codex/translate.rs +++ b/server/src/session/codex/translate.rs @@ -17,7 +17,6 @@ pub(super) struct Translator { pub(super) thread_id: Option, completed: bool, limited: bool, - pending_usage: Option, subagents: Option>, children: HashMap, prompts: HashMap, @@ -25,11 +24,6 @@ pub(super) struct Translator { in_turn: bool, } -struct Usage { - tokens: u64, - context: Option, -} - impl Translator { pub(super) fn new(subagents: Arc, thread_id: Option, in_turn: bool) -> Self { Self { @@ -130,7 +124,6 @@ impl Translator { Some("turn.started") | Some("turn/started") => { self.completed = false; self.limited = false; - self.pending_usage = None; vec![Event::Status { state: SessionStatus::Running, }] @@ -191,27 +184,31 @@ impl Translator { output: output.to_string(), }] } + // One of these per *model request*, not per turn: app-server's `last` is the + // request that just finished, and a turn is as many requests as it made tool + // calls. Reported as each arrives, so `tokens` adds up to what the turn cost -- + // held to the end of the turn it was the last request's cost alone, which on a + // two-request turn measured 28,878 against the 51,399 actually spent -- and so + // the context figure moves while a long turn is still running rather than + // standing at what it was before the turn began. Some("thread/tokenUsage/updated") => { let last = &body["tokenUsage"]["last"]; - self.pending_usage = - last.get("totalTokens") - .and_then(Value::as_u64) - .map(|tokens| Usage { - tokens, - // Cached input is a subset of this figure, not an additional count. - context: last.get("inputTokens").and_then(Value::as_u64), - }); - Vec::new() + last.get("totalTokens") + .and_then(Value::as_u64) + .map(|tokens| Event::UsageDelta { + tokens, + // Cached input is a subset of this figure, not an additional count. + context: last.get("inputTokens").and_then(Value::as_u64), + }) + .into_iter() + .collect() } Some("turn.completed") | Some("turn/completed") => { self.completed = true; let mut events = Vec::new(); - if let Some(usage) = self.pending_usage.take() { - events.push(Event::UsageDelta { - tokens: usage.tokens, - context: usage.context, - }); - } else if let Some(usage) = line.get("usage") { + // The old `codex exec --json` dialect reports the turn's usage here and + // sends no `thread/tokenUsage/updated` at all. + if let Some(usage) = line.get("usage") { let input = number(usage, "input_tokens"); let output = number(usage, "output_tokens"); if input.is_some() || output.is_some() { @@ -1069,26 +1066,33 @@ mod tests { text: "hello, revised".to_string() }] ); - assert!( - translator - .translate(&line( - r#"{"method":"thread/tokenUsage/updated","params":{"tokenUsage":{"last":{"inputTokens":39,"cachedInputTokens":30,"outputTokens":3,"reasoningOutputTokens":1,"totalTokens":42},"total":{"inputTokens":100,"cachedInputTokens":80,"outputTokens":9,"reasoningOutputTokens":2,"totalTokens":109},"modelContextWindow":258400}}}"# - )) - .is_empty() + // One per model request, as it arrives. A turn that made two of them costs both, + // and the context figure moves while the turn is still running. + assert_eq!( + translator.translate(&line( + r#"{"method":"thread/tokenUsage/updated","params":{"tokenUsage":{"last":{"inputTokens":39,"cachedInputTokens":30,"outputTokens":3,"reasoningOutputTokens":1,"totalTokens":42},"total":{"inputTokens":100,"cachedInputTokens":80,"outputTokens":9,"reasoningOutputTokens":2,"totalTokens":109},"modelContextWindow":258400}}}"# + )), + vec![Event::UsageDelta { + tokens: 42, + context: Some(39) + }] + ); + assert_eq!( + translator.translate(&line( + r#"{"method":"thread/tokenUsage/updated","params":{"tokenUsage":{"last":{"inputTokens":61,"cachedInputTokens":39,"outputTokens":4,"reasoningOutputTokens":1,"totalTokens":65},"total":{"inputTokens":161,"cachedInputTokens":119,"outputTokens":13,"reasoningOutputTokens":3,"totalTokens":174},"modelContextWindow":258400}}}"# + )), + vec![Event::UsageDelta { + tokens: 65, + context: Some(61) + }] ); assert_eq!( translator.translate(&line( r#"{"method":"turn/completed","params":{"turn":{"id":"turn-1","status":"completed","error":null}}}"# )), - vec![ - Event::UsageDelta { - tokens: 42, - context: Some(39) - }, - Event::Status { - state: SessionStatus::Idle - } - ] + vec![Event::Status { + state: SessionStatus::Idle + }] ); }