Expose assistant content token counts in completion notifications (#59041)

## What changed

Add optional `contentTokenCount` metadata to `item/completed` notifications, carrying the producer-supplied token count for each completed assistant message in default and plan modes. Counts include any sampled footer and are `null` when unavailable.

Keep this telemetry separate from stored thread items and omit it from subsequent model requests. Update the protocol schemas and generated TypeScript and Python types.

## Testing

Add coverage for per-message counts, including zero, in default and plan modes; missing and malformed metadata; backward-compatible notification decoding; and exclusion from subsequent model requests.

GitOrigin-RevId: 35973037a42650fa2d7569da8e88cb81215cf971
This commit is contained in:
aojiao-oai
2026-11-03 02:26:29 +00:00
committed by copyberry
parent e839b828fb
commit 16ddd538bd
48 changed files with 443 additions and 39 deletions
@@ -2619,6 +2619,16 @@
"format": "int64",
"type": "integer"
},
"contentTokenCount": {
"default": null,
"description": "Native content tokens for the completed assistant message, including a sampled footer. Null when unavailable. Completion telemetry is not part of the stored thread item.",
"format": "uint64",
"minimum": 0.0,
"type": [
"integer",
"null"
]
},
"item": {
"$ref": "#/definitions/ThreadItem"
},
@@ -14396,6 +14396,16 @@
"format": "int64",
"type": "integer"
},
"contentTokenCount": {
"default": null,
"description": "Native content tokens for the completed assistant message, including a sampled footer. Null when unavailable. Completion telemetry is not part of the stored thread item.",
"format": "uint64",
"minimum": 0.0,
"type": [
"integer",
"null"
]
},
"item": {
"$ref": "#/definitions/v2/ThreadItem"
},
@@ -10357,6 +10357,16 @@
"format": "int64",
"type": "integer"
},
"contentTokenCount": {
"default": null,
"description": "Native content tokens for the completed assistant message, including a sampled footer. Null when unavailable. Completion telemetry is not part of the stored thread item.",
"format": "uint64",
"minimum": 0.0,
"type": [
"integer",
"null"
]
},
"item": {
"$ref": "#/definitions/ThreadItem"
},
@@ -2065,6 +2065,16 @@
"format": "int64",
"type": "integer"
},
"contentTokenCount": {
"default": null,
"description": "Native content tokens for the completed assistant message, including a sampled footer. Null when unavailable. Completion telemetry is not part of the stored thread item.",
"format": "uint64",
"minimum": 0.0,
"type": [
"integer",
"null"
]
},
"item": {
"$ref": "#/definitions/ThreadItem"
},
@@ -4,6 +4,11 @@
import type { ThreadItem } from "./ThreadItem";
export type ItemCompletedNotification = { item: ThreadItem, threadId: string, turnId: string,
/**
* Native content tokens for the completed assistant message, including a sampled footer.
* Null when unavailable. Completion telemetry is not part of the stored thread item.
*/
contentTokenCount?: number | null,
/**
* Unix timestamp (in milliseconds) when this item lifecycle completed.
*/
+3 -1
View File
@@ -2425,7 +2425,9 @@ mod tests {
|| (path == Path::new("v2/ProcessExitedNotification.ts")
&& field_prefix.trim() == "outputDrainTimedOut?")
|| (path == Path::new("v2/MisalignmentErrorDetails.ts")
&& field_prefix.trim() == "reviewTarget?");
&& field_prefix.trim() == "reviewTarget?")
|| (path == Path::new("v2/ItemCompletedNotification.ts")
&& field_prefix.trim() == "contentTokenCount?");
if field_prefix.chars().rev().find(|c| !c.is_whitespace()) == Some('?')
&& !allow_optional_nullable
&& !legacy_optional_nullable_field
@@ -419,13 +419,15 @@ pub fn item_event_to_server_notification(
))
}
EventMsg::ItemCompleted(item_completed_event) => {
ServerNotification::ItemCompleted(ItemCompletedNotification::new(
let mut notification = ItemCompletedNotification::new(
thread_id,
turn_id,
ThreadItemPayload::from(item_completed_event.item),
item_completed_event.started_at_ms,
item_completed_event.completed_at_ms,
))
);
notification.content_token_count = item_completed_event.content_token_count;
ServerNotification::ItemCompleted(notification)
}
EventMsg::PatchApplyUpdated(event) => {
ServerNotification::FileChangePatchUpdated(FileChangePatchUpdatedNotification {
@@ -487,6 +489,56 @@ mod tests {
use codex_protocol::protocol::ExecOutputStream;
use pretty_assertions::assert_eq;
#[test]
fn assistant_completion_count_is_optional_envelope_metadata() {
let item = codex_protocol::items::TurnItem::AgentMessage(
codex_protocol::items::AgentMessageItem {
id: "answer".into(),
content: vec![codex_protocol::items::AgentMessageContent::Text {
text: "Answer".into(),
}],
phase: None,
memory_citation: None,
delivery: None,
questions: None,
},
);
let payload = ThreadItemPayload::from(item.clone());
for count in [None, Some(0), Some(23)] {
let event = codex_protocol::protocol::ItemCompletedEvent {
thread_id: ThreadId::new(),
turn_id: "turn".into(),
item: item.clone(),
content_token_count: count,
started_at_ms: Some(1),
completed_at_ms: 2,
};
let mut expected = ItemCompletedNotification::new(
"thread".into(),
"turn".into(),
payload.clone(),
/*started_at_ms*/ Some(1),
/*completed_at_ms*/ 2,
);
expected.content_token_count = count;
assert_item_completed_server_notification(
item_event_to_server_notification(EventMsg::ItemCompleted(event), "thread", "turn"),
expected.clone(),
);
let mut wire = serde_json::to_value(&expected).unwrap();
assert_eq!(
wire.get("contentTokenCount"),
Some(&serde_json::json!(count))
);
wire.as_object_mut().unwrap().remove("contentTokenCount");
expected.content_token_count = None;
assert_eq!(
serde_json::from_value::<ItemCompletedNotification>(wire).unwrap(),
expected,
);
}
}
fn assert_item_started_server_notification(
notification: ServerNotification,
expected: ItemStartedNotification,
@@ -2135,6 +2135,7 @@ mod tests {
};
builder.handle_event_for_history_mode(&EventMsg::ItemStarted(started), mode);
let completed = ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::default(),
turn_id: turn_id.clone(),
item,
@@ -2336,6 +2337,7 @@ mod tests {
});
let events = vec![
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "turn-1".into(),
item: entered,
@@ -2343,6 +2345,7 @@ mod tests {
completed_at_ms: 0,
}),
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "turn-1".into(),
item: exited,
@@ -2594,6 +2597,7 @@ mod tests {
collaboration_mode_kind: Default::default(),
}),
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: turn_id.to_string(),
item: sleep_item,
@@ -2679,6 +2683,7 @@ mod tests {
collaboration_mode_kind: Default::default(),
}),
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: turn_id.to_string(),
item: CoreTurnItem::Extension(CoreExtensionItem::ImageGeneration(
@@ -2797,6 +2802,7 @@ mod tests {
interaction_input: None,
}),
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: turn_id.to_string(),
item: command_item,
@@ -4073,6 +4079,7 @@ mod tests {
use codex_protocol::protocol::HasLegacyEvent;
let completion = ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-a".into(),
started_at_ms: None,
@@ -4388,6 +4395,7 @@ mod tests {
("c", "latest"),
] {
builder.handle_live_paginated_event(&EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::default(),
turn_id: turn_id.clone(),
item: CoreTurnItem::AgentMessage(codex_protocol::items::AgentMessageItem {
@@ -5530,6 +5538,7 @@ mod tests {
}));
let changes =
builder.handle_event_with_changes(&EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-a".into(),
item: hook_prompt,
@@ -5551,6 +5560,7 @@ mod tests {
.join("worker")
.expect("worker path");
let spawn = ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-a".into(),
item: CoreTurnItem::SubAgentActivity(CoreSubAgentActivityItem {
@@ -5601,6 +5611,7 @@ mod tests {
time_to_first_token_ms: None,
})),
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-a".into(),
item: CoreTurnItem::SubAgentActivity(CoreSubAgentActivityItem {
@@ -178,6 +178,7 @@ fn projects_optional_completed_item_lifecycle_timestamps() {
{
let changes = project(RolloutItem::EventMsg(EventMsg::ItemCompleted(
ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::default(),
turn_id: "turn-1".to_string(),
item: item.clone(),
@@ -283,6 +284,7 @@ fn project(item: RolloutItem) -> ThreadHistoryChangeSet {
fn item_completed(thread_id: ThreadId, turn_id: &str, item: TurnItem) -> RolloutItem {
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: turn_id.to_string(),
item,
@@ -1505,6 +1505,12 @@ pub struct ItemCompletedNotification {
pub item: ThreadItem,
pub thread_id: String,
pub turn_id: String,
/// Native content tokens for the completed assistant message, including a sampled footer.
/// Null when unavailable. Completion telemetry is not part of the stored thread item.
#[serde(default)]
#[ts(optional)]
#[ts(type = "number | null")]
pub content_token_count: Option<u64>,
/// Unix timestamp (in milliseconds) when this item lifecycle completed.
#[ts(type = "number")]
pub completed_at_ms: i64,
@@ -1520,6 +1526,7 @@ impl ItemCompletedNotification {
Self {
thread_id,
turn_id,
content_token_count: None,
completed_at_ms,
item: ThreadItem::completed(payload, started_at_ms, completed_at_ms),
}
@@ -2122,6 +2122,17 @@ mod tests {
"failed to connect app-server remote control websocket `{}`: HTTP error: 503 Service Unavailable, request-id: <none>, cf-ray: <none>, body: upstream unavailable",
remote_control_target.websocket_url
);
let codex_home = TempDir::new().expect("temp dir should create");
let state_db = remote_control_state_runtime(&codex_home).await;
let auth_manager = remote_control_auth_manager();
let session_auth = RemoteControlAuth::capture(auth_manager.clone()).0;
let mut auth_recovery = session_auth.unauthorized_recovery();
let mut auth_change_rx = auth_manager.auth_change_receiver();
let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some(
TEST_REMOTE_CONTROL_SERVER_TOKEN,
))));
let (status_publisher, status_rx) = remote_control_status_channel();
let server_task = tokio::spawn(async move {
let (stream, request_line) = accept_http_request(&listener).await;
assert_eq!(
@@ -2136,16 +2147,6 @@ mod tests {
)
.await;
});
let codex_home = TempDir::new().expect("temp dir should create");
let state_db = remote_control_state_runtime(&codex_home).await;
let auth_manager = remote_control_auth_manager();
let session_auth = RemoteControlAuth::capture(auth_manager.clone()).0;
let mut auth_recovery = session_auth.unauthorized_recovery();
let mut auth_change_rx = auth_manager.auth_change_receiver();
let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some(
TEST_REMOTE_CONTROL_SERVER_TOKEN,
))));
let (status_publisher, status_rx) = remote_control_status_channel();
let err = match connect_remote_control_websocket(
&remote_control_target,
@@ -2266,14 +2267,6 @@ mod tests {
let remote_control_target =
normalize_remote_control_url(&remote_control_url).expect("target should parse");
let enroll_url = remote_control_target.enroll_url.clone();
let server_task = tokio::spawn(async move {
let (stream, request_line) = accept_http_request(&listener).await;
assert_eq!(
request_line,
"POST /backend-api/wham/remote/control/server/enroll HTTP/1.1"
);
respond_with_status_and_headers(stream, "401 Unauthorized", &[], "unauthorized").await;
});
let codex_home = TempDir::new().expect("temp dir should create");
save_auth(
codex_home.path(),
@@ -2306,6 +2299,15 @@ mod tests {
)
.expect("fresh auth should save");
let server_task = tokio::spawn(async move {
let (stream, request_line) = accept_http_request(&listener).await;
assert_eq!(
request_line,
"POST /backend-api/wham/remote/control/server/enroll HTTP/1.1"
);
respond_with_status_and_headers(stream, "401 Unauthorized", &[], "unauthorized").await;
});
let err = connect_remote_control_websocket(
&remote_control_target,
Some(state_db.as_ref()),
@@ -2366,14 +2368,6 @@ mod tests {
let remote_control_target =
normalize_remote_control_url(&remote_control_url).expect("target should parse");
let refresh_url = remote_control_target.refresh_url.clone();
let server_task = tokio::spawn(async move {
let (stream, request_line) = accept_http_request(&listener).await;
assert_eq!(
request_line,
"POST /backend-api/wham/remote/control/server/refresh HTTP/1.1"
);
respond_with_status_and_headers(stream, "401 Unauthorized", &[], "unauthorized").await;
});
let codex_home = TempDir::new().expect("temp dir should create");
save_auth(
codex_home.path(),
@@ -2411,6 +2405,15 @@ mod tests {
)
.expect("fresh auth should save");
let server_task = tokio::spawn(async move {
let (stream, request_line) = accept_http_request(&listener).await;
assert_eq!(
request_line,
"POST /backend-api/wham/remote/control/server/refresh HTTP/1.1"
);
respond_with_status_and_headers(stream, "401 Unauthorized", &[], "unauthorized").await;
});
let err = connect_remote_control_websocket(
&remote_control_target,
Some(state_db.as_ref()),
@@ -3331,6 +3331,7 @@ mod tests {
Event {
id: "turn-1".to_string(),
msg: EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: conversation_id,
turn_id: "turn-1".to_string(),
item: CoreTurnItem::SubAgentActivity(SubAgentActivityItem {
@@ -3503,6 +3504,7 @@ mod tests {
state.track_current_turn_event(
&event_turn_id,
&EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: conversation_id,
turn_id: event_turn_id.clone(),
item: CoreTurnItem::AgentMessage(CoreAgentMessageItem {
@@ -3527,6 +3529,7 @@ mod tests {
state.track_current_turn_event(
&event_turn_id,
&EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: conversation_id,
turn_id: event_turn_id.clone(),
item: CoreTurnItem::AgentMessage(CoreAgentMessageItem {
@@ -11,6 +11,7 @@ use app_test_support::write_mock_responses_config_toml_with_chatgpt_base_url;
use codex_app_server_protocol::ApprovalsReviewer;
use codex_app_server_protocol::AskForApproval;
use codex_app_server_protocol::ClientRequest;
use codex_app_server_protocol::ItemCompletedNotification;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::SandboxPolicy;
@@ -25,6 +26,7 @@ use codex_app_server_protocol::TurnSettingsUpdateResponse;
use codex_app_server_protocol::TurnSettingsUpdateStatus;
use codex_app_server_protocol::TurnStartParams;
use codex_app_server_protocol::TurnStartResponse;
use codex_app_server_protocol::TurnStatus;
use codex_app_server_protocol::UserInput;
use codex_config::types::AuthCredentialsStoreMode;
use codex_config::types::OtelExporterKind;
@@ -35,12 +37,18 @@ use codex_core_plugins::store::PluginStore;
use codex_features::Feature;
use codex_models_manager::bundled_models_response;
use codex_plugin::PluginId;
use codex_protocol::config_types::CollaborationMode;
use codex_protocol::config_types::ModeKind;
use codex_protocol::config_types::Settings;
use codex_protocol::models::MessagePhase;
use codex_protocol::openai_models::ModelsResponse;
use codex_protocol::openai_models::ReasoningEffort;
use core_test_support::responses;
use core_test_support::skip_if_no_network;
use core_test_support::skip_if_remote;
use core_test_support::skip_if_wine_exec;
use core_test_support::startup::STARTUP_TIMEOUT;
use core_test_support::startup::expect_startup;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
@@ -59,6 +67,123 @@ use wiremock::matchers::path_regex;
const SERVICE_VERSION: &str = "0.0.0-test";
// Completion metadata follows its source message without entering later model requests.
#[test_case::test_case(ModeKind::Default; "normal")]
#[test_case::test_case(ModeKind::Plan; "plan")]
#[tokio::test]
async fn native_counts_follow_completed_messages_without_request_replay(
mode: ModeKind,
) -> Result<()> {
#[cfg(windows)]
const READ_TIMEOUT: Duration = Duration::from_secs(25);
#[cfg(not(windows))]
const READ_TIMEOUT: Duration = Duration::from_secs(10);
let server = responses::start_mock_server().await;
let mut events = vec![responses::ev_response_created("response-1")];
for (id, text, phase, count) in [
("partial", "Partial.", "partial_answer", 0),
("final", "Done.", "final_answer", 23),
] {
let mut added = responses::ev_message_item_added(id, "");
added["item"]["phase"] = json!(phase);
let mut done = responses::ev_assistant_message(id, text);
done["item"]["phase"] = json!(phase);
done["item"]["internal_chat_message_metadata_passthrough"] =
json!({"content_token_count": count});
events.extend([added, responses::ev_output_text_delta(text), done]);
}
events.push(responses::ev_completed("response-1"));
let codex_home = TempDir::new()?;
MockResponsesConfig::new(&server.uri()).write(codex_home.path())?;
let mut app = expect_startup(
TestAppServer::builder()
.with_codex_home(codex_home.path())
.build_initialized_with_timeout(STARTUP_TIMEOUT),
)
.await;
let response_mock = responses::mount_sse_sequence(
&server,
vec![
responses::sse(events),
responses::sse(vec![
responses::ev_response_created("response-2"),
responses::ev_assistant_message("next", "Next."),
responses::ev_completed("response-2"),
]),
],
)
.await;
let thread = app
.start_thread(ThreadStartParams {
model: Some("mock-model".into()),
..Default::default()
})
.await?
.thread;
timeout(READ_TIMEOUT, async {
for prompt in ["First turn", "Next turn"] {
let completed = app
.start_turn_and_wait_for_completion(TurnStartParams {
thread_id: thread.id.clone(),
input: vec![UserInput::Text {
text: prompt.into(),
text_elements: Vec::new(),
}],
collaboration_mode: Some(CollaborationMode {
mode,
settings: Settings {
model: "mock-model".into(),
reasoning_effort: None,
developer_instructions: None,
},
}),
..Default::default()
})
.await?;
assert_eq!(completed.turn.status, TurnStatus::Completed);
}
let mut counts = Vec::new();
while counts.len() < 2 {
let completed: ItemCompletedNotification =
app.read_notification("item/completed").await?;
if let ThreadItemPayload::AgentMessage { id, phase, .. } = completed.item.payload {
counts.push((id, phase, completed.content_token_count));
}
}
assert_eq!(
counts,
vec![
("partial".into(), Some(MessagePhase::PartialAnswer), Some(0)),
("final".into(), Some(MessagePhase::FinalAnswer), Some(23)),
]
);
app.shutdown_gracefully().await?;
anyhow::Ok(())
})
.await??;
let requests = response_mock.requests();
assert_eq!(requests.len(), 2);
let assistant_input: Vec<_> = requests[1]
.inputs_of_type("message")
.into_iter()
.filter(|item| item["role"] == "assistant")
.map(|item| {
(
item["content"][0]["text"].clone(),
item["internal_chat_message_metadata_passthrough"]
.get("content_token_count")
.cloned(),
)
})
.collect();
assert_eq!(
assistant_input,
vec![(json!("Partial."), None), (json!("Done."), None)]
);
Ok(())
}
#[tokio::test]
async fn usage_limit_window_reaches_turn_analytics() -> Result<()> {
skip_if_no_network!(Ok(()));
@@ -1924,6 +1924,7 @@ async fn thread_fork_freezes_active_paginated_turn_as_interrupted(
};
let completed_user_item = |id: &str, completed_at_ms| {
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: source_id,
turn_id: "active-turn".to_string(),
item: CoreTurnItem::UserMessage(UserMessageItem {
@@ -2761,6 +2761,7 @@ fn paginated_completed_item(
item: CoreTurnItem,
) -> RolloutItem {
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: turn_id.to_string(),
item,
@@ -108,6 +108,7 @@ async fn timeline_pages_mix_items_and_resolve_the_opening_realtime_session() ->
collaboration_mode_kind: Default::default(),
})),
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "turn-1".to_string(),
item: TurnItem::UserMessage(UserMessageItem {
+1
View File
@@ -213,6 +213,7 @@ impl LocalAgentControl {
.await;
let completed_at_ms = now_unix_timestamp_ms();
let completed = ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: turn_id.clone(),
item,
+1
View File
@@ -2435,6 +2435,7 @@ async fn spawn_agent_fork_from_paginated_parent_uses_model_context_prefix(
internal_chat_message_metadata_passthrough: None,
}),
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: parent_thread_id,
turn_id: "parent-turn".to_string(),
item: TurnItem::UserMessage(UserMessageItem {
+1 -3
View File
@@ -3557,11 +3557,9 @@ async fn guardian_review_routes_required_actions(
#[tokio::test]
async fn guardian_ephemeral_retry_preserves_parallel_trunk_and_fork_history() -> anyhow::Result<()>
{
const TEST_STACK_SIZE_BYTES: usize = 4 * 1024 * 1024;
// Inherit the test runner's stack budget, including the larger Windows CI setting.
let handle = std::thread::Builder::new()
.name("guardian_ephemeral_retry_preserves_parallel_trunk_and_fork_history".to_string())
.stack_size(TEST_STACK_SIZE_BYTES)
.spawn(|| -> anyhow::Result<()> {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
@@ -60,6 +60,7 @@ pub(super) async fn record(session: &Session, text: String) -> std::io::Result<(
// ItemCompleted is the existing source for thread/items/list. Keep the legacy
// equivalent too; the store selects the appropriate format for this thread.
let completed = EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: session.thread_id,
turn_id: turn.sub_id.clone(),
item: TurnItem::UserMessage(UserMessageItem::new(&content)),
@@ -66,6 +66,7 @@ fn assistant_delta(item_id: &str, delta: &str) -> EventMsg {
fn completed_item(item: TurnItem) -> EventMsg {
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-1".to_string(),
item,
+15
View File
@@ -2908,6 +2908,20 @@ impl Session {
&self,
turn_context: &TurnContext,
item: TurnItem,
) {
self.emit_turn_item_completed_with_token_count(
turn_context,
item,
/*content_token_count*/ None,
)
.await;
}
pub(crate) async fn emit_turn_item_completed_with_token_count(
&self,
turn_context: &TurnContext,
item: TurnItem,
content_token_count: Option<u64>,
) {
record_turn_ttfm_metric(turn_context, &item).await;
for contributor in self.services.extensions.turn_lifecycle_contributors() {
@@ -2940,6 +2954,7 @@ impl Session {
thread_id: self.thread_id,
turn_id: turn_context.sub_id.clone(),
item,
content_token_count,
started_at_ms: Some(started_at_ms),
completed_at_ms,
}),
+17 -3
View File
@@ -2565,6 +2565,7 @@ async fn emit_agent_message_in_plan_mode(
sess: &Session,
turn_context: &TurnContext,
agent_message: codex_protocol::items::AgentMessageItem,
content_token_count: Option<u64>,
state: &mut PlanModeStreamState,
) {
let agent_message_id = agent_message.id.clone();
@@ -2600,8 +2601,12 @@ async fn emit_agent_message_in_plan_mode(
.insert(agent_message_id.clone());
}
sess.emit_turn_item_completed(turn_context, TurnItem::AgentMessage(agent_message))
.await;
sess.emit_turn_item_completed_with_token_count(
turn_context,
TurnItem::AgentMessage(agent_message),
content_token_count,
)
.await;
state.started_agent_message_items.remove(&agent_message_id);
}
@@ -2610,12 +2615,20 @@ async fn emit_turn_item_in_plan_mode(
sess: &Session,
turn_context: &TurnContext,
turn_item: TurnItem,
content_token_count: Option<u64>,
previously_active_item: Option<&TurnItem>,
state: &mut PlanModeStreamState,
) {
match turn_item {
TurnItem::AgentMessage(agent_message) => {
emit_agent_message_in_plan_mode(sess, turn_context, agent_message, state).await;
emit_agent_message_in_plan_mode(
sess,
turn_context,
agent_message,
content_token_count,
state,
)
.await;
}
_ => {
if previously_active_item.is_none() {
@@ -2656,6 +2669,7 @@ async fn handle_assistant_item_done_in_plan_mode(
sess,
turn_context,
finalized_turn_item.turn_item,
finalized_turn_item.content_token_count,
previously_active_item,
state,
)
+16 -1
View File
@@ -258,6 +258,7 @@ pub(crate) enum TurnItemContributorPolicy<'a> {
pub(crate) struct FinalizedTurnItem {
pub(crate) turn_item: TurnItem,
pub(crate) facts: FinalizedTurnItemFacts,
pub(crate) content_token_count: Option<u64>,
}
#[derive(Clone, Default)]
@@ -303,6 +304,16 @@ pub(crate) async fn finalize_non_tool_response_item(
_ => (None, None, false),
};
Some(FinalizedTurnItem {
content_token_count: match item {
ResponseItem::Message {
role,
internal_chat_message_metadata_passthrough,
..
} if role == "assistant" => internal_chat_message_metadata_passthrough
.as_ref()
.and_then(|metadata| metadata.content_token_count),
_ => None,
},
turn_item,
facts: FinalizedTurnItemFacts {
memory_citation,
@@ -395,7 +406,11 @@ pub(crate) async fn handle_output_item_done(
}
ctx.sess
.emit_turn_item_completed(&ctx.step_context.turn, finalized_turn_item.turn_item)
.emit_turn_item_completed_with_token_count(
&ctx.step_context.turn,
finalized_turn_item.turn_item,
finalized_turn_item.content_token_count,
)
.await;
}
record_completed_response_item_with_finalized_facts(
+16 -3
View File
@@ -466,20 +466,32 @@ async fn direct_results_keep_their_own_records_when_call_ids_repeat() {
assert!(result.item.executed_tool_call_metadata().is_none());
}
#[test_case::test_case(false; "default_mode")]
#[test_case::test_case(true; "plan_mode")]
#[tokio::test]
async fn finalized_turn_item_defers_mailbox_for_contributed_visible_text() {
async fn finalized_turn_item_defers_mailbox_for_contributed_visible_text(plan_mode: bool) {
let (mut session, turn_context) = make_session_and_context().await;
let mut builder = codex_extension_api::ExtensionRegistryBuilder::new();
builder.turn_item_contributor(Arc::new(RewriteAgentMessageContributor));
session.services.extensions = Arc::new(builder.build());
let turn_store = ExtensionData::new(turn_context.sub_id.clone());
let item = assistant_output_text("<oai-mem-citation>hidden only</oai-mem-citation>");
let mut item = assistant_output_text("<oai-mem-citation>hidden only</oai-mem-citation>");
let ResponseItem::Message {
internal_chat_message_metadata_passthrough,
..
} = &mut item
else {
panic!("expected assistant message");
};
internal_chat_message_metadata_passthrough
.get_or_insert_default()
.content_token_count = Some(23);
let finalized = finalize_non_tool_response_item(
&session,
TurnItemContributorPolicy::Run(&turn_store),
&item,
/*plan_mode*/ false,
plan_mode,
)
.await
.expect("assistant message should parse");
@@ -489,6 +501,7 @@ async fn finalized_turn_item_defers_mailbox_for_contributed_visible_text() {
Some("contributed assistant text")
);
assert!(finalized.facts.defers_mailbox_delivery_to_next_turn);
assert_eq!(finalized.content_token_count, Some(23));
}
#[test_case::test_case(MessagePhase::Commentary; "commentary")]
+7
View File
@@ -1009,6 +1009,13 @@ pub struct InternalChatMessageMetadataPassthrough {
#[schemars(skip)]
#[ts(skip)]
pub create_time: Option<serde_json::Number>,
/// Native content-token count for this sampled message, including its sampled footer.
/// Read from the producer only; do not replay telemetry into subsequent model requests.
#[serde_as(deserialize_as = "serde_with::DefaultOnError")]
#[serde(default, skip_serializing)]
#[schemars(skip)]
#[ts(skip)]
pub content_token_count: Option<u64>,
/// Harness-owned classifications aligned with the item's content entries.
#[serde_as(deserialize_as = "serde_with::DefaultOnError")]
#[serde(default, skip_serializing_if = "Option::is_none")]
@@ -57,6 +57,41 @@ fn malformed_content_kinds_do_not_prevent_loading_response_item() -> Result<()>
Ok(())
}
#[test]
fn malformed_content_token_counts_do_not_prevent_loading_response_item() -> Result<()> {
for value in [json!(-1), json!(true), json!(1.5), json!("4"), json!(null)] {
let item: ResponseItem = serde_json::from_value(response_item_json(json!({
"content_token_count": value,
})))?;
assert_eq!(
item,
response_item(InternalChatMessageMetadataPassthrough::default())
);
}
Ok(())
}
#[test]
fn content_token_counts_are_not_replayed_into_requests() -> Result<()> {
let item: ResponseItem = serde_json::from_value(response_item_json(json!({
"turn_id": "turn-1", "content_token_count": 23,
})))?;
let ResponseItem::Message {
internal_chat_message_metadata_passthrough: Some(ref metadata),
..
} = item
else {
panic!("expected message metadata");
};
assert_eq!(metadata.content_token_count, Some(23));
let encoded = serde_json::to_value(item)?;
assert_eq!(
encoded["internal_chat_message_metadata_passthrough"],
json!({"turn_id": "turn-1"})
);
Ok(())
}
fn response_item(metadata: InternalChatMessageMetadataPassthrough) -> ResponseItem {
ResponseItem::Message {
status: None,
+13
View File
@@ -2194,6 +2194,11 @@ pub struct ItemCompletedEvent {
pub thread_id: ThreadId,
pub turn_id: String,
pub item: TurnItem,
/// Native tokens in the completed assistant message, including any sampled footer.
/// Completion telemetry only; absent for items without a producer-supplied count.
#[serde(default, skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub content_token_count: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[ts(optional)]
pub started_at_ms: Option<i64>,
@@ -5763,6 +5768,7 @@ mod tests {
#[test]
fn item_completed_event_from_image_generation_emits_end_event() {
let event = ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-1".into(),
item: TurnItem::ImageGeneration(ImageGenerationItem {
@@ -5801,6 +5807,7 @@ mod tests {
#[test]
fn item_completed_event_from_file_change_emits_patch_end_event() {
let event = ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-1".into(),
started_at_ms: Some(0),
@@ -5840,6 +5847,7 @@ mod tests {
#[test]
fn item_completed_event_from_mcp_tool_call_emits_end_event() {
let event = ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-1".into(),
started_at_ms: Some(0),
@@ -5921,6 +5929,7 @@ mod tests {
}),
};
let completed = ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-1".into(),
started_at_ms: Some(10),
@@ -5999,6 +6008,7 @@ mod tests {
}),
};
let completed = ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-1".into(),
started_at_ms: Some(10),
@@ -6044,6 +6054,7 @@ mod tests {
#[test]
fn review_mode_item_completion_emits_legacy_events_with_ids() {
let entered = ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-1".into(),
started_at_ms: Some(0),
@@ -6057,6 +6068,7 @@ mod tests {
}),
};
let exited = ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-1".into(),
started_at_ms: Some(0),
@@ -6117,6 +6129,7 @@ mod tests {
#[test]
fn item_completed_event_defaults_missing_completed_at_ms() {
let mut value = serde_json::to_value(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-1".into(),
item: TurnItem::UserMessage(UserMessageItem::new(&[])),
@@ -267,6 +267,7 @@ fn invalid_turn_boundaries_reset_partial_measurements() {
#[test]
fn item_completion_persistence_depends_on_history_mode() {
let item = RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::default(),
turn_id: "turn".to_string(),
item: TurnItem::UserMessage(UserMessageItem {
@@ -307,6 +308,7 @@ fn item_completion_persistence_depends_on_history_mode() {
fn projected_items_report_post_projection_bytes() {
let original_output = "x".repeat(PERSISTED_COMMAND_OUTPUT_MAX_BYTES * 2);
let item = RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::default(),
turn_id: "turn".to_string(),
item: TurnItem::CommandExecution(CommandExecutionItem {
@@ -364,6 +366,7 @@ fn projected_items_report_post_projection_bytes() {
fn review_mode_persistence_depends_on_history_mode() {
let completed_items = vec![
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::default(),
turn_id: "turn".to_string(),
item: TurnItem::EnteredReviewMode(EnteredReviewModeItem {
@@ -377,6 +380,7 @@ fn review_mode_persistence_depends_on_history_mode() {
completed_at_ms: 0,
})),
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::default(),
turn_id: "turn".to_string(),
item: TurnItem::ExitedReviewMode(ExitedReviewModeItem {
+2
View File
@@ -28,6 +28,7 @@ fn image_item(image: GeneratedImageRef) -> ImageGenerationItem {
#[test]
fn paginated_image_generation_preserves_file_reference() {
let event = EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-1".to_string(),
item: TurnItem::Extension(ExtensionItem::ImageGeneration(image_item(
@@ -89,6 +90,7 @@ fn legacy_image_generation_preserves_file_reference() {
#[test]
fn inline_image_generation_result_is_preserved() {
let event = EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: "turn-1".to_string(),
item: TurnItem::Extension(ExtensionItem::ImageGeneration(image_item(
+3
View File
@@ -313,6 +313,7 @@ mod tests {
let mut expected = metadata.clone();
for text in [" ", "delegated task", "later output"] {
let item = RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: metadata.id,
turn_id: "turn-1".to_string(),
item: TurnItem::FunctionCallOutput(codex_protocol::items::FunctionCallOutputItem {
@@ -345,6 +346,7 @@ mod tests {
fn completed_user_message_items_set_title_and_first_user_message() {
let mut metadata = metadata_for_test();
let item = RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::default(),
turn_id: "turn-1".to_string(),
item: TurnItem::UserMessage(UserMessageItem::new(&[UserInput::Text {
@@ -395,6 +397,7 @@ mod tests {
for event in [
EventMsg::UserMessage(user.as_legacy_user_message_event()),
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "turn-1".to_string(),
item: TurnItem::UserMessage(user),
+1
View File
@@ -2329,6 +2329,7 @@ mod tests {
.await
.expect("create paginated thread");
let paginated_item = RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "turn-1".to_string(),
item: TurnItem::UserMessage(UserMessageItem {
@@ -720,6 +720,7 @@ fn contextual_user_message() -> RolloutItem {
fn completed_user_message(turn_id: &str, message: &str) -> RolloutItem {
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: codex_protocol::ThreadId::from_string("00000000-0000-0000-0000-000000000000")
.expect("fixture thread id"),
turn_id: turn_id.to_string(),
@@ -936,6 +936,7 @@ mod tests {
codex_rollout::append_rollout_item_to_path(
&rollout_path,
&RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "turn-1".to_string(),
item: TurnItem::UserMessage(UserMessageItem::new(&[UserInput::Text {
@@ -422,6 +422,7 @@ impl LegacyRolloutCanonicalizer {
writer,
timestamp,
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: self.thread_id,
turn_id,
item,
@@ -222,6 +222,7 @@ fn exec_completion(turn_id: &str, call_id: &str, output: &str) -> RolloutItem {
fn item_completed(turn_id: &str, item_id: &str) -> RolloutItem {
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::new(),
turn_id: turn_id.to_string(),
item: TurnItem::Reasoning(ReasoningItem {
@@ -29,6 +29,7 @@ async fn delegated_thread_is_listed_before_user_followup() {
live_thread
.append_items(&[RolloutItem::EventMsg(EventMsg::ItemCompleted(
ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "delegated-turn".to_string(),
item: TurnItem::FunctionCallOutput(FunctionCallOutputItem {
@@ -1573,6 +1573,7 @@ async fn projection_preserves_exact_lifecycle_timestamps() {
for (index, (started_at_ms, completed_at_ms)) in timestamps.iter().enumerate() {
items.push(RolloutItem::EventMsg(EventMsg::ItemCompleted(
ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "turn-1".to_string(),
item: TurnItem::UserMessage(UserMessageItem {
@@ -3171,6 +3172,7 @@ fn contains_user_message(items: &[RolloutItem], expected: &str) -> bool {
fn completed_item(thread_id: ThreadId, turn_id: &str, item: TurnItem) -> RolloutItem {
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: turn_id.to_string(),
item,
@@ -538,6 +538,7 @@ mod tests {
thread_id,
vec![RolloutItem::EventMsg(EventMsg::ItemCompleted(
ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "turn-1".to_string(),
item: TurnItem::UserMessage(UserMessageItem::new(&[UserInput::Text {
@@ -650,6 +651,7 @@ mod tests {
let update = sync
.observe_appended_items(&[RolloutItem::EventMsg(EventMsg::ItemCompleted(
ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "turn-1".to_string(),
item: TurnItem::UserMessage(UserMessageItem::new(&[UserInput::Text {
@@ -27,6 +27,7 @@ fn delegated_output_emits_only_first_preview_in_live_patch() {
);
for (index, text) in ["first task", "later task"].into_iter().enumerate() {
let item = RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "turn".to_string(),
item: TurnItem::FunctionCallOutput(FunctionCallOutputItem {
+1
View File
@@ -9293,6 +9293,7 @@ async fn prompt_edit_reverts_earlier_and_first_visible_prompts_in_place_scenario
collaboration_mode_kind: ModeKind::default(),
})),
RolloutItem::EventMsg(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id: ThreadId::from_string(&source_thread_id)?,
turn_id: turn_id.to_string(),
item: TurnItem::UserMessage(UserMessageItem {
@@ -65,6 +65,7 @@ async fn history_fixture(
})];
for item in 0..*count {
events.push(EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: turn_id.clone(),
started_at_ms: None,
@@ -87,6 +87,7 @@ pub(super) async fn completed_history_app(
];
events.extend(items.into_iter().map(|item| {
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: turn_id.clone(),
item,
@@ -2517,6 +2517,7 @@ async fn older_pagination_reconciles_review_prompts_across_page_boundaries() ->
}))
.chain(items.into_iter().map(|item| {
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "cross-page-review-turn".to_string(),
item,
@@ -2681,6 +2682,7 @@ async fn transcript_alt_beginning_loads_every_older_history_page() -> Result<()>
}))
.chain((0..305).map(|index| {
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "multi-page-turn".to_string(),
item: TurnItem::AgentMessage(AgentMessageItem {
@@ -3097,6 +3099,7 @@ async fn underfilled_scrollback_fetches_older_pages_without_opening_the_transcri
}))
.chain((0..120).map(|index| {
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: "scrollback-pagination-turn".to_string(),
item: TurnItem::AgentMessage(AgentMessageItem {
@@ -5345,6 +5348,7 @@ async fn command_center_read_only_open_requests_and_failure_preservation() -> Re
collaboration_mode_kind: Default::default(),
}),
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: format!("saved-turn-{index}"),
item: TurnItem::UserMessage(UserMessageItem::new(&[CoreUserInput::Text {
@@ -5355,6 +5359,7 @@ async fn command_center_read_only_open_requests_and_failure_preservation() -> Re
completed_at_ms: 0,
}),
EventMsg::ItemCompleted(ItemCompletedEvent {
content_token_count: None,
thread_id,
turn_id: format!("saved-turn-{index}"),
item: TurnItem::AgentMessage(AgentMessageItem {
+1
View File
@@ -1010,6 +1010,7 @@ mod tests {
let mut expected = serde_json::json!({
"method": "item/completed", "params": {
"threadId": "thread", "turnId": "turn", "completedAtMs": 20,
"contentTokenCount": null,
"item": {
"type": "agentMessage", "id": "question", "text": "already in the snapshot",
"startedAtMs": 10, "completedAtMs": 20,