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
This commit is contained in:
Owen Lin authored and copyberry committed 2026-08-31 21:15:38 +00:00
1 parent a7913390f7
commit 5f79a92e39
67 files changed
+932 -162

No files matched your search

+10 -17
View File
@@ -169,16 +169,16 @@ impl InputQueue {
.and_then(|id| id.filter(|id| !id.trim().is_empty()).map(str::to_string));
start_options.root_turn_id = pending_mails
.iter()
.filter(|mail| mail.communication.trigger_turn)
.map(|mail| {
.find(|mail| mail.communication.trigger_turn)
.and_then(|mail| {
mail.start_options
.parent_turn_id
.as_deref()
.filter(|id| !id.trim().is_empty())
.and(mail.start_options.root_turn_id.as_deref())
.filter(|id| !id.trim().is_empty())
})
.reduce(|expected, candidate| expected.filter(|id| candidate == Some(*id)))
.and_then(|id| id.filter(|id| !id.trim().is_empty()).map(str::to_string));
.map(str::to_string);
let items = pending_mails
.into_iter()
.map(|mail| TurnInput::InterAgentCommunication(mail.communication))
@@ -320,17 +320,10 @@ impl InputQueue {
}
let (mailbox_items, start_options) = self.drain_mailbox_input_items().await;
if let Some(active_turn_metadata) = active_turn_metadata
&& mailbox_items.iter().any(|item| {
matches!(
item,
TurnInput::InterAgentCommunication(communication)
if communication.trigger_turn
)
})
&& (start_options.root_turn_id.is_none()
|| active_turn_metadata.root_turn_id() != start_options.root_turn_id)
&& active_turn_metadata.root_turn_id().is_none()
&& let Some(root_turn_id) = start_options.root_turn_id.as_ref()
{
active_turn_metadata.mark_root_turn_ambiguous();
active_turn_metadata.set_root_turn_id(root_turn_id.clone());
}
if pending_input.is_empty() {
(mailbox_items, start_options)
@@ -561,7 +554,7 @@ mod tests {
}
#[tokio::test]
async fn input_queue_requires_one_unambiguous_trigger_parent() {
async fn input_queue_uses_unambiguous_trigger_parent_and_first_root() {
let (parent, peer, root, root2) = (Some("a"), Some("b"), Some("r"), Some("s"));
for (pending_mails, expected_parent_turn_id, expected_root_turn_id) in [
(Vec::new(), None, None),
@@ -572,8 +565,8 @@ mod tests {
(vec![(true, parent, None)], parent, None),
(vec![(true, parent, Some(""))], parent, None),
(vec![(true, parent, root), (true, peer, root)], None, root),
(vec![(true, parent, root), (true, peer, root2)], None, None),
(vec![(true, parent, root), (true, None, root)], None, None),
(vec![(true, parent, root), (true, peer, root2)], None, root),
(vec![(true, parent, root), (true, None, root)], None, root),
(
vec![(true, parent, root), (true, parent, root)],
parent,