diff --git a/codex-rs/core/src/context_manager/context_breakdown.rs b/codex-rs/core/src/context_manager/context_breakdown.rs index 1f94df3fff..9f48c41e54 100644 --- a/codex-rs/core/src/context_manager/context_breakdown.rs +++ b/codex-rs/core/src/context_manager/context_breakdown.rs @@ -728,4 +728,53 @@ mod tests { conversation.tokens ); } + + #[test] + fn developer_message_content_items_are_classified_individually() { + let breakdown = build_context_window_breakdown( + &[ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ + ContentItem::InputText { + text: "\npolicy".to_string(), + }, + ContentItem::InputText { + text: "\n### Available skills\n- notify: send a notification\n### How to use skills\n" + .to_string(), + }, + ], + end_turn: None, + phase: None, + }], + &BaseInstructions { + text: String::new(), + }, + /*model_context_window*/ None, + /*verbose*/ false, + ); + + let mut sections = breakdown + .sections + .iter() + .map(|section| (section.label.clone(), section.tokens)) + .collect::>(); + sections.sort_by(|left, right| left.0.cmp(&right.0)); + + assert_eq!( + sections + .iter() + .map(|(label, _)| label.clone()) + .collect::>(), + vec!["Built-in".to_string(), "Skills".to_string()] + ); + assert_eq!( + breakdown.total_tokens, + breakdown + .sections + .iter() + .map(|section| section.tokens) + .sum::() + ); + } } diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 58531681e9..63290c1902 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -505,6 +505,7 @@ enum ThreadBufferedEvent { Request(ServerRequest), HistoryEntryResponse(GetHistoryEntryResponseEvent), FeedbackSubmission(FeedbackThreadEvent), + ContextWindowBreakdown(ContextWindowBreakdownThreadEvent), } #[derive(Debug, Clone, PartialEq, Eq)] @@ -515,6 +516,11 @@ struct FeedbackThreadEvent { result: Result, } +#[derive(Debug, Clone)] +struct ContextWindowBreakdownThreadEvent { + result: Result, +} + #[derive(Debug)] struct ThreadEventStore { session: Option, @@ -535,6 +541,7 @@ impl ThreadEventStore { | ThreadBufferedEvent::Notification(ServerNotification::HookStarted(_)) | ThreadBufferedEvent::Notification(ServerNotification::HookCompleted(_)) | ThreadBufferedEvent::FeedbackSubmission(_) + | ThreadBufferedEvent::ContextWindowBreakdown(_) ) } @@ -640,7 +647,8 @@ impl ThreadEventStore { .should_replay_snapshot_request(request), ThreadBufferedEvent::Notification(_) | ThreadBufferedEvent::HistoryEntryResponse(_) - | ThreadBufferedEvent::FeedbackSubmission(_) => true, + | ThreadBufferedEvent::FeedbackSubmission(_) + | ThreadBufferedEvent::ContextWindowBreakdown(_) => true, }) .cloned() .collect(), @@ -1917,7 +1925,7 @@ impl App { .await .map(|response: ThreadContextReadResponse| response.context) .map_err(|err| err.to_string()); - app_event_tx.send(AppEvent::ContextWindowBreakdownLoaded { result }); + app_event_tx.send(AppEvent::ContextWindowBreakdownLoaded { thread_id, result }); }); } @@ -2108,6 +2116,49 @@ impl App { } } + async fn enqueue_thread_context_window_breakdown_event( + &mut self, + thread_id: ThreadId, + event: ContextWindowBreakdownThreadEvent, + ) { + let (sender, store) = { + let channel = self.ensure_thread_channel(thread_id); + (channel.sender.clone(), Arc::clone(&channel.store)) + }; + + let should_send = { + let mut guard = store.lock().await; + guard + .buffer + .push_back(ThreadBufferedEvent::ContextWindowBreakdown(event.clone())); + if guard.buffer.len() > guard.capacity + && let Some(removed) = guard.buffer.pop_front() + && let ThreadBufferedEvent::Request(request) = &removed + { + guard + .pending_interactive_replay + .note_evicted_server_request(request); + } + guard.active + }; + + if should_send { + match sender.try_send(ThreadBufferedEvent::ContextWindowBreakdown(event)) { + Ok(()) => {} + Err(TrySendError::Full(event)) => { + tokio::spawn(async move { + if let Err(err) = sender.send(event).await { + tracing::warn!("thread {thread_id} event channel closed: {err}"); + } + }); + } + Err(TrySendError::Closed(_)) => { + tracing::warn!("thread {thread_id} event channel closed"); + } + } + } + } + async fn handle_feedback_submitted( &mut self, origin_thread_id: Option, @@ -2798,6 +2849,10 @@ impl App { ThreadBufferedEvent::FeedbackSubmission(event) => { self.enqueue_thread_feedback_event(thread_id, event).await; } + ThreadBufferedEvent::ContextWindowBreakdown(event) => { + self.enqueue_thread_context_window_breakdown_event(thread_id, event) + .await; + } } } self.chat_widget @@ -4495,8 +4550,12 @@ impl App { )); } } - AppEvent::ContextWindowBreakdownLoaded { result } => { - self.handle_context_window_breakdown_result(result); + AppEvent::ContextWindowBreakdownLoaded { thread_id, result } => { + self.enqueue_thread_context_window_breakdown_event( + thread_id, + ContextWindowBreakdownThreadEvent { result }, + ) + .await; } AppEvent::StartFileSearch(query) => { self.file_search.on_user_query(query); @@ -5688,6 +5747,9 @@ impl App { ThreadBufferedEvent::FeedbackSubmission(event) => { self.handle_feedback_thread_event(event); } + ThreadBufferedEvent::ContextWindowBreakdown(event) => { + self.handle_context_window_breakdown_result(event.result); + } } if needs_refresh { self.refresh_status_line(); @@ -5708,6 +5770,9 @@ impl App { ThreadBufferedEvent::FeedbackSubmission(event) => { self.handle_feedback_thread_event(event); } + ThreadBufferedEvent::ContextWindowBreakdown(event) => { + self.handle_context_window_breakdown_result(event.result); + } } } @@ -9643,6 +9708,86 @@ guardian_approval = true })); } + #[tokio::test] + async fn context_window_breakdown_for_inactive_thread_replays_into_origin_thread() { + let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; + let origin_thread_id = ThreadId::new(); + let active_thread_id = ThreadId::new(); + let origin_session = test_thread_session(origin_thread_id, PathBuf::from("/tmp/origin")); + let active_session = test_thread_session(active_thread_id, PathBuf::from("/tmp/active")); + app.thread_event_channels.insert( + origin_thread_id, + ThreadEventChannel::new_with_session( + THREAD_EVENT_CHANNEL_CAPACITY, + origin_session.clone(), + Vec::new(), + ), + ); + app.thread_event_channels.insert( + active_thread_id, + ThreadEventChannel::new_with_session( + THREAD_EVENT_CHANNEL_CAPACITY, + active_session.clone(), + Vec::new(), + ), + ); + app.activate_thread_channel(active_thread_id).await; + app.chat_widget.handle_thread_session(active_session); + while app_event_rx.try_recv().is_ok() {} + + app.enqueue_thread_context_window_breakdown_event( + origin_thread_id, + ContextWindowBreakdownThreadEvent { + result: Ok(ThreadContextWindowBreakdown { + model_context_window: Some(1000), + total_tokens: 250, + sections: vec![codex_app_server_protocol::ThreadContextWindowSection { + label: "Conversation".to_string(), + tokens: 250, + details: vec![codex_app_server_protocol::ThreadContextWindowDetail { + label: "User message".to_string(), + tokens: 250, + }], + }], + }), + }, + ) + .await; + + assert_matches!( + app_event_rx.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + ); + + let snapshot = { + let channel = app + .thread_event_channels + .get(&origin_thread_id) + .expect("origin thread channel should exist"); + let store = channel.store.lock().await; + assert!(matches!( + store.buffer.back(), + Some(ThreadBufferedEvent::ContextWindowBreakdown(_)) + )); + store.snapshot() + }; + + app.replay_thread_snapshot(snapshot, /*resume_restored_queue*/ false); + + let mut rendered_cells = Vec::new(); + while let Ok(event) = app_event_rx.try_recv() { + if let AppEvent::InsertHistoryCell(cell) = event { + rendered_cells.push(lines_to_single_string(&cell.display_lines(/*width*/ 120))); + } + } + assert!( + rendered_cells + .iter() + .any(|cell| cell.contains("/context") && cell.contains("Context map")), + "expected replayed context breakdown output, got {rendered_cells:?}" + ); + } + fn next_user_turn_op(op_rx: &mut tokio::sync::mpsc::UnboundedReceiver) -> Op { let mut seen = Vec::new(); while let Ok(op) = op_rx.try_recv() { diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index a655522392..b975db38fe 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -272,6 +272,7 @@ pub(crate) enum AppEvent { /// Result of fetching the current thread's context-window breakdown. ContextWindowBreakdownLoaded { + thread_id: ThreadId, result: Result, },