File size: 7,277 Bytes
afa0cbf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | //! Replays persisted token usage snapshots when a client attaches to an existing thread.
//!
//! The message processor decides when replay is allowed and preserves JSON-RPC response
//! ordering. This module owns notification construction and the attribution rules that
//! map the latest persisted `TokenCount` back to a v2 turn id.
//!
//! Rollout histories can contain explicit turn ids or generated turn ids. When explicit
//! ids do not match the rebuilt thread, replay falls back to the active turn position at
//! the time the `TokenCount` was persisted so the notification still targets the
//! corresponding rebuilt turn.
use std::sync::Arc;
use codex_app_server_protocol::ServerNotification;
use codex_app_server_protocol::ThreadHistoryBuilder;
use codex_app_server_protocol::ThreadTokenUsage;
use codex_app_server_protocol::ThreadTokenUsageUpdatedNotification;
use codex_app_server_protocol::Turn;
use codex_app_server_protocol::TurnStatus;
use codex_core::CodexThread;
use codex_protocol::ThreadId;
use codex_protocol::protocol::EventMsg;
use codex_rollout::RolloutItem;
use crate::outgoing_message::ConnectionId;
use crate::outgoing_message::OutgoingMessageSender;
/// Sends a restored token usage update to the connection that attached to a thread.
///
/// This is lifecycle replay rather than a model event: the rollout already contains
/// the original `TokenCount`, and emitting through `send_event` here would duplicate
/// persisted usage records. Keeping replay connection-scoped also avoids
/// surprising other subscribers with a historical usage update while they may be
/// rendering live turn events.
pub(super) async fn send_thread_token_usage_update_to_connection(
outgoing: &Arc<OutgoingMessageSender>,
connection_id: ConnectionId,
thread_id: ThreadId,
conversation: &CodexThread,
token_usage_turn_id: String,
) {
let Some(info) = conversation.token_usage_info().await else {
return;
};
let notification = ThreadTokenUsageUpdatedNotification {
thread_id: thread_id.to_string(),
turn_id: token_usage_turn_id,
token_usage: ThreadTokenUsage::from(info),
};
outgoing
.send_server_notification_to_connections(
&[connection_id],
ServerNotification::ThreadTokenUsageUpdated(notification),
)
.await;
}
pub(super) fn restored_token_usage_turn_id(
rollout_items: &[RolloutItem],
turns: &[Turn],
) -> String {
latest_token_usage_turn_id_from_rollout_items(rollout_items, turns)
.unwrap_or_else(|| latest_token_usage_turn_id(turns))
}
/// Identifies the turn that was active when the latest `TokenCount` record appeared.
///
/// The id is preferred when it still appears in the rebuilt thread. The position is a
/// fallback for histories whose implicit turn ids are regenerated during reconstruction.
fn latest_token_usage_turn_id_from_rollout_items(
rollout_items: &[RolloutItem],
turns: &[Turn],
) -> Option<String> {
let token_count_index = rollout_items
.iter()
.rposition(|item| matches!(item, RolloutItem::EventMsg(EventMsg::TokenCount(_))))?;
let mut builder = ThreadHistoryBuilder::new();
for item in &rollout_items[..token_count_index] {
builder.handle_rollout_item(item);
}
if turns.is_empty() {
return builder.active_turn_id_if_explicit();
}
let active_turn_id = builder.active_turn_id()?;
if turns.iter().any(|turn| turn.id == active_turn_id) {
Some(active_turn_id.to_string())
} else {
builder
.active_turn_position()
.and_then(|position| turns.get(position))
.map(|turn| turn.id.clone())
}
}
/// Chooses a fallback turn id that should own a replayed token usage update.
///
/// Normal replay derives the owner from the rollout position of the latest
/// `TokenCount` event. This fallback only preserves a stable wire shape for
/// unusual histories where that rollout information cannot be read.
fn latest_token_usage_turn_id(turns: &[Turn]) -> String {
turns
.iter()
.rev()
.find(|turn| matches!(turn.status, TurnStatus::Completed | TurnStatus::Failed))
.or_else(|| turns.last())
.map(|turn| turn.id.clone())
.unwrap_or_default()
}
#[cfg(test)]
mod tests {
use super::*;
use codex_app_server_protocol::build_turns_from_rollout_items;
use codex_protocol::protocol::AgentMessageEvent;
use codex_protocol::protocol::TokenCountEvent;
use codex_protocol::protocol::UserMessageEvent;
use pretty_assertions::assert_eq;
#[test]
fn replay_attribution_uses_already_loaded_history() {
let rollout_items = token_usage_history();
let turns = build_turns_from_rollout_items(&rollout_items);
assert_eq!(
latest_token_usage_turn_id_from_rollout_items(&rollout_items, turns.as_slice()),
Some(turns[0].id.clone())
);
}
#[test]
fn replay_attribution_falls_back_to_rebuilt_turn_position() {
let rollout_items = token_usage_history();
let mut turns = build_turns_from_rollout_items(&rollout_items);
turns[0].id = "rebuilt-turn-id".to_string();
assert_eq!(
latest_token_usage_turn_id_from_rollout_items(&rollout_items, turns.as_slice()),
Some("rebuilt-turn-id".to_string())
);
}
#[test]
fn replay_attribution_rejects_suffix_generated_turn_ids() {
let rollout_items = token_usage_history();
assert_eq!(
latest_token_usage_turn_id_from_rollout_items(&rollout_items, /*turns*/ &[]),
None
);
}
#[test]
fn replay_attribution_uses_latest_token_count_and_ignores_tail_turn() {
let mut rollout_items = token_usage_history();
rollout_items.extend(token_usage_history());
let turns = build_turns_from_rollout_items(&rollout_items);
assert_eq!(
latest_token_usage_turn_id_from_rollout_items(&rollout_items, turns.as_slice()),
Some(turns[2].id.clone())
);
}
fn token_usage_history() -> Vec<RolloutItem> {
vec![
RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
client_id: None,
message: "first turn".to_string(),
images: None,
local_images: Vec::new(),
text_elements: Vec::new(),
..Default::default()
})),
RolloutItem::EventMsg(EventMsg::AgentMessage(AgentMessageEvent {
message: "first answer".to_string(),
phase: None,
memory_citation: None,
delivery: None,
questions: None,
})),
RolloutItem::EventMsg(EventMsg::TokenCount(TokenCountEvent {
info: None,
rate_limits: None,
})),
RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
client_id: None,
message: "second turn".to_string(),
images: None,
local_images: Vec::new(),
text_elements: Vec::new(),
..Default::default()
})),
]
}
}
|