Files
codex/codex-rs/core/src/session/inject.rs
T
Owen Lin 5f79a92e39 Persist response token usage in rollout history (#41912)
## Why

Resumed threads need to continue per-turn and per-thread token totals without scanning arbitrarily far beyond the latest compaction checkpoint.

## What changed

- Add durable `TokenUsageRecord` rollout items with response, turn, thread, session, and root-turn attribution.
- Restore accumulated usage on resume and snapshot the latest record plus the compaction response ID in compaction checkpoints.
- Preserve root-turn lineage in persisted turn context, while ensuring forked child threads start with their own usage totals.

## Testing

- Cover usage accumulation across multiple responses and resumed turns.
- Cover local and remote compaction checkpoints, invalid remote compaction output, rollout reconstruction, and fork isolation.

GitOrigin-RevId: ef9e0c4a9102a08a2c382be4cdac68c84353c90a
2026-08-31 21:15:38 +00:00

165 lines
5.8 KiB
Rust

use super::TurnInput as PendingTurnInput;
use super::session::Session;
use super::turn_context::TurnContext;
use codex_features::Feature;
use codex_history::CodexHarnessMetadata;
use codex_history::ResponseItemEnvelope;
use codex_protocol::models::ResponseItem;
impl Session {
/// Returns the input if there is no active turn to inject into.
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state updates must remain atomic"
)]
pub async fn inject_if_running(
&self,
input: Vec<ResponseItem>,
) -> Result<(), Vec<ResponseItem>> {
let mut active = self.active_turn.lock().await;
match active.as_mut() {
Some(active_turn) => {
self.input_queue
.extend_pending_input_and_accept_mailbox_delivery_for_turn_state(
active_turn.turn_state.as_ref(),
input
.into_iter()
.map(ResponseItemEnvelope::new)
.map(PendingTurnInput::ResponseItem)
.collect(),
)
.await;
Ok(())
}
None => Err(input),
}
}
/// Injects hook context into the running turn atomically.
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn provenance and turn state updates must remain atomic"
)]
pub(crate) async fn inject_hook_context_if_running(
&self,
input: Vec<ResponseItem>,
) -> Result<(), Vec<ResponseItem>> {
let mut active = self.active_turn.lock().await;
let Some(active_turn) = active.as_mut() else {
return Err(input);
};
if active_turn.task.is_none() {
return Err(input);
}
self.input_queue
.extend_pending_input_and_accept_mailbox_delivery_for_turn_state(
active_turn.turn_state.as_ref(),
input
.into_iter()
.map(ResponseItemEnvelope::new)
.map(PendingTurnInput::ResponseItem)
.collect(),
)
.await;
Ok(())
}
/// Preserves trusted client provenance while items wait for an active turn.
#[expect(
clippy::await_holding_invalid_type,
reason = "active turn checks and turn state updates must remain atomic"
)]
pub(crate) async fn inject_client_response_items(
&self,
items: Vec<ResponseItem>,
turn_context: &TurnContext,
) {
let items = items
.into_iter()
.map(|item| self.annotate_client_response_item(item))
.collect::<Vec<_>>();
let mut active = self.active_turn.lock().await;
if let Some(active_turn) = active.as_mut() {
self.input_queue
.extend_pending_input_and_accept_mailbox_delivery_for_turn_state(
active_turn.turn_state.as_ref(),
items
.into_iter()
.map(PendingTurnInput::ResponseItem)
.collect(),
)
.await;
return;
}
drop(active);
self.record_annotated_conversation_items(turn_context, items)
.await;
}
pub(crate) fn annotate_client_response_item(&self, item: ResponseItem) -> ResponseItemEnvelope {
let metadata = (self.enabled(Feature::RetainClientDeveloperMessages)
&& matches!(&item, ResponseItem::Message { role, .. } if role == "developer"))
.then_some(CodexHarnessMetadata {
client_authored: true,
..Default::default()
});
ResponseItemEnvelope { item, metadata }
}
pub(crate) async fn record_annotated_conversation_items(
&self,
turn_context: &TurnContext,
items: Vec<ResponseItemEnvelope>,
) {
if items.iter().all(|item| item.metadata.is_none()) {
let items = items
.into_iter()
.map(ResponseItemEnvelope::into_item)
.collect::<Vec<_>>();
self.record_conversation_items(turn_context, &items).await;
return;
}
let mut annotated_items = Vec::with_capacity(items.len());
let mut image_preparations = Vec::new();
for envelope in items {
let (prepared_items, prepared_images) = self.prepare_conversation_items_for_history(
turn_context,
std::slice::from_ref(&envelope.item),
);
image_preparations.extend(prepared_images);
let mut metadata = envelope.metadata;
annotated_items.extend(prepared_items.into_owned().into_iter().map(|item| {
ResponseItemEnvelope {
item,
metadata: metadata.take(),
}
}));
}
self.record_prepared_conversation_items(turn_context, annotated_items, image_preparations)
.await;
}
/// Injects items into active work, or records them without starting a turn.
pub(crate) async fn inject_no_new_turn(
&self,
items: Vec<ResponseItem>,
current_turn_context: Option<&TurnContext>,
) {
let Err(items) = self.inject_if_running(items).await else {
return;
};
let default_turn_context;
let turn_context = match current_turn_context {
Some(turn_context) => turn_context,
None => {
default_turn_context = self.new_default_turn().await;
default_turn_context.as_ref()
}
};
self.record_conversation_items(turn_context, &items).await;
}
}