diff --git a/codex-rs/tui/src/app/app_server_event_targets.rs b/codex-rs/tui/src/app/app_server_event_targets.rs index 5e5fe181a..291082000 100644 --- a/codex-rs/tui/src/app/app_server_event_targets.rs +++ b/codex-rs/tui/src/app/app_server_event_targets.rs @@ -39,6 +39,19 @@ pub(super) enum ServerNotificationThreadTarget { Global, } +impl ServerNotificationThreadTarget { + /// The thread this notification targets, if it is addressed to a specific + /// thread. Non-thread targets (app-scoped, global, invalid) return `None`. + pub(super) fn thread_id(&self) -> Option { + match self { + ServerNotificationThreadTarget::Thread(thread_id) => Some(*thread_id), + ServerNotificationThreadTarget::InvalidThreadId(_) + | ServerNotificationThreadTarget::AppScoped + | ServerNotificationThreadTarget::Global => None, + } + } +} + pub(super) fn server_notification_thread_target( notification: &ServerNotification, ) -> ServerNotificationThreadTarget { diff --git a/codex-rs/tui/src/app/app_server_events.rs b/codex-rs/tui/src/app/app_server_events.rs index 8c18ef3f9..b4d7ad8f8 100644 --- a/codex-rs/tui/src/app/app_server_events.rs +++ b/codex-rs/tui/src/app/app_server_events.rs @@ -127,10 +127,21 @@ impl App { self.note_session_recap_turn_completed(thread_id, ¬ification.turn); } - let result = if self.primary_thread_id == Some(thread_id) - || self.primary_thread_id.is_none() - { - self.enqueue_primary_thread_notification(notification).await + let result = if self.primary_thread_id == Some(thread_id) { + self.enqueue_thread_notification(thread_id, notification) + .await + } else if self.primary_thread_id.is_none() { + // Startup pre-primary window: the primary thread has not + // finished starting yet. Only buffer as a pending primary + // event when the target thread is not already tracked; a + // known loaded/scheduled thread routes to its own channel + // so its events never leak into the primary buffer. + if self.is_tracked_thread(thread_id) { + self.enqueue_thread_notification(thread_id, notification) + .await + } else { + self.enqueue_primary_thread_notification(notification).await + } } else { self.enqueue_thread_notification(thread_id, notification) .await @@ -202,12 +213,20 @@ impl App { return; }; - let result = - if self.primary_thread_id == Some(thread_id) || self.primary_thread_id.is_none() { - self.enqueue_primary_thread_request(request).await - } else { + let result = if self.primary_thread_id == Some(thread_id) { + self.enqueue_thread_request(thread_id, request).await + } else if self.primary_thread_id.is_none() { + // Startup pre-primary window: only buffer requests for threads that + // are not already tracked. Known loaded/scheduled threads route to + // their own channel instead of the pending primary buffer. + if self.is_tracked_thread(thread_id) { self.enqueue_thread_request(thread_id, request).await - }; + } else { + self.enqueue_primary_thread_request(request).await + } + } else { + self.enqueue_thread_request(thread_id, request).await + }; if let Err(err) = result { tracing::warn!("failed to enqueue app-server request: {err}"); } diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 5b72dda2d..a7f99335c 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -315,6 +315,142 @@ async fn resolved_buffered_approval_does_not_become_actionable_after_drain() -> Ok(()) } +#[tokio::test] +async fn enqueue_primary_thread_session_drops_unrelated_pre_primary_notifications() -> Result<()> { + let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; + let primary_thread_id = ThreadId::new(); + let unrelated_thread_id = ThreadId::new(); + + // Pre-primary startup window: no primary thread is set yet, so events are + // buffered as pending primary events. + assert!(app.primary_thread_id.is_none()); + app.enqueue_primary_thread_notification(agent_message_delta_notification( + unrelated_thread_id, + "turn-x", + "item-x", + "unrelated", + )) + .await?; + app.enqueue_primary_thread_notification(agent_message_delta_notification( + primary_thread_id, + "turn-1", + "item-1", + "primary", + )) + .await?; + + app.enqueue_primary_thread_session( + test_thread_session(primary_thread_id, test_path_buf("/tmp/project")), + Vec::new(), + ) + .await?; + while app_event_rx.try_recv().is_ok() {} + + // The unrelated thread must not leak a channel out of the primary buffer flush. + assert!( + !app.thread_event_channels.contains_key(&unrelated_thread_id), + "unrelated thread should not receive a channel from the primary flush" + ); + + let rx = app + .active_thread_rx + .as_mut() + .expect("primary thread receiver should be active"); + + let event = time::timeout(Duration::from_millis(50), rx.recv()) + .await + .expect("timed out waiting for buffered primary notification") + .expect("channel closed unexpectedly"); + assert!( + matches!( + &event, + ThreadBufferedEvent::Notification(ServerNotification::AgentMessageDelta(delta)) + if delta.thread_id == primary_thread_id.to_string() && delta.delta == "primary" + ), + "expected the primary thread's buffered notification to replay first, got {event:?}" + ); + + // The unrelated-thread notification must have been dropped, not replayed onto + // the primary thread's stream. + let extra = time::timeout(Duration::from_millis(50), rx.recv()).await; + assert!( + extra.is_err(), + "unrelated pre-primary notification should be dropped, not replayed onto the primary" + ); + + Ok(()) +} + +#[tokio::test] +async fn enqueue_primary_thread_session_routes_known_thread_pre_primary_notifications() -> Result<()> +{ + let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; + let primary_thread_id = ThreadId::new(); + let known_thread_id = ThreadId::new(); + + // A previously loaded/scheduled thread is already tracked before the primary + // thread finishes starting. + app.thread_event_channels.insert( + known_thread_id, + ThreadEventChannel::new(THREAD_EVENT_CHANNEL_CAPACITY), + ); + // Mark the known thread active so routed events are delivered on its channel + // receiver rather than parked in its store buffer. + app.set_thread_active(known_thread_id, /*active*/ true) + .await; + + assert!(app.primary_thread_id.is_none()); + app.enqueue_primary_thread_notification(agent_message_delta_notification( + known_thread_id, + "turn-k", + "item-k", + "known", + )) + .await?; + + app.enqueue_primary_thread_session( + test_thread_session(primary_thread_id, test_path_buf("/tmp/project")), + Vec::new(), + ) + .await?; + while app_event_rx.try_recv().is_ok() {} + + // The buffered event for the already-tracked thread is delivered on that + // thread's own channel, not on the primary thread's stream. + let mut known_rx = app + .thread_event_channels + .get_mut(&known_thread_id) + .expect("known thread channel should still exist") + .receiver + .take() + .expect("known thread receiver should be available"); + let known_event = time::timeout(Duration::from_millis(50), known_rx.recv()) + .await + .expect("timed out waiting for known-thread notification") + .expect("known thread channel closed unexpectedly"); + assert!( + matches!( + &known_event, + ThreadBufferedEvent::Notification(ServerNotification::AgentMessageDelta(delta)) + if delta.thread_id == known_thread_id.to_string() && delta.delta == "known" + ), + "expected known-thread notification to route to its own channel, got {known_event:?}" + ); + + // Nothing from the known thread should have leaked onto the primary stream. + let rx = app + .active_thread_rx + .as_mut() + .expect("primary thread receiver should be active"); + let leaked = time::timeout(Duration::from_millis(50), rx.recv()).await; + assert!( + leaked.is_err(), + "known-thread notification must not be replayed onto the primary stream" + ); + + Ok(()) +} + #[tokio::test] async fn enqueue_primary_thread_session_replays_turns_before_initial_prompt_submit() -> Result<()> { let (mut app, mut app_event_rx, _op_rx) = make_test_app_with_channels().await; diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 3f3d8b5b7..51da4b3a7 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -5,6 +5,8 @@ //! when the visible thread changes. use super::*; +use crate::app::app_server_event_targets::server_notification_thread_target; +use crate::app::app_server_event_targets::server_request_thread_id; use crate::app::app_server_requests::AppServerRequestResolution; use crate::session_resume::read_session_model; use codex_app_server_protocol::ThreadExternalAgentStartStatus; @@ -1194,11 +1196,37 @@ impl App { for pending_event in pending { match pending_event { ThreadBufferedEvent::Notification(notification) => { - self.enqueue_thread_notification(thread_id, notification) - .await?; + let event_thread_id = + server_notification_thread_target(¬ification).thread_id(); + match self.resolve_pending_primary_flush_target(event_thread_id, thread_id) { + Some(target_thread_id) => { + self.enqueue_thread_notification(target_thread_id, notification) + .await?; + } + None => { + tracing::debug!( + ?event_thread_id, + %thread_id, + "dropping buffered notification for unrelated untracked thread while starting primary thread" + ); + } + } } ThreadBufferedEvent::Request(request) => { - self.enqueue_thread_request(thread_id, request).await?; + let event_thread_id = server_request_thread_id(&request); + match self.resolve_pending_primary_flush_target(event_thread_id, thread_id) { + Some(target_thread_id) => { + self.enqueue_thread_request(target_thread_id, request) + .await?; + } + None => { + tracing::debug!( + ?event_thread_id, + %thread_id, + "dropping buffered request for unrelated untracked thread while starting primary thread" + ); + } + } } ThreadBufferedEvent::HistoryEntryResponse(event) => { self.enqueue_thread_history_entry_response(thread_id, event) @@ -1215,6 +1243,42 @@ impl App { Ok(()) } + /// Whether `thread_id` is already tracked by the TUI (has an event channel + /// or is a registered side thread). Used during the startup pre-primary + /// window to distinguish known loaded/scheduled threads from the primary + /// thread that is still starting. + pub(super) fn is_tracked_thread(&self, thread_id: ThreadId) -> bool { + self.thread_event_channels.contains_key(&thread_id) + || self.side_threads.contains_key(&thread_id) + } + + /// Resolve which thread a buffered pre-primary event should flush onto once + /// the primary thread (`primary_thread_id`) has started. + /// + /// - Events that target the started primary thread flush onto the primary. + /// - Events for another already-tracked thread flush onto that thread's own + /// channel so they are not misrouted onto the primary's stream. + /// - Events for an unrelated, untracked thread return `None` and are dropped + /// by the caller; they never belonged to the primary being started. + /// - Events with no resolvable thread target flush onto the primary, matching + /// the historical behavior (such events are not buffered here in practice). + fn resolve_pending_primary_flush_target( + &self, + event_thread_id: Option, + primary_thread_id: ThreadId, + ) -> Option { + match event_thread_id { + None => Some(primary_thread_id), + Some(event_thread_id) if event_thread_id == primary_thread_id => { + Some(primary_thread_id) + } + Some(event_thread_id) if self.is_tracked_thread(event_thread_id) => { + Some(event_thread_id) + } + Some(_) => None, + } + } + pub(super) async fn enqueue_primary_thread_notification( &mut self, notification: ServerNotification,