diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f16cf6c93872..471c46b94eac 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4156,6 +4156,7 @@ dependencies = [ "diffy", "dirs", "dunce", + "futures", "image", "insta", "itertools 0.14.0", @@ -4184,6 +4185,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-stream", + "tokio-tungstenite", "tokio-util", "toml 0.9.11+spec-1.1.0", "tracing", diff --git a/codex-rs/tui/Cargo.toml b/codex-rs/tui/Cargo.toml index 6c0b83a72e98..91f4b991fd8b 100644 --- a/codex-rs/tui/Cargo.toml +++ b/codex-rs/tui/Cargo.toml @@ -155,10 +155,12 @@ core_test_support = { workspace = true } codex-utils-cargo-bin = { workspace = true } assert_matches = { workspace = true } chrono = { workspace = true, features = ["serde"] } +futures = { workspace = true } insta = { workspace = true } pretty_assertions = { workspace = true } rand = { workspace = true } serial_test = { workspace = true } +tokio-tungstenite = { workspace = true } vt100 = { workspace = true } uuid = { workspace = true } wiremock = { workspace = true } diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index 4a75263e1f38..42dbcae5a4ec 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -185,7 +185,6 @@ impl App { match self .replace_chat_widget_with_app_server_thread( tui, - app_server, forked, ThreadAttachPresentation::SessionLineage, /*initial_user_message*/ None, @@ -283,7 +282,6 @@ impl App { match self .replace_chat_widget_with_app_server_thread( tui, - app_server, forked, ThreadAttachPresentation::PromptEdit, /*initial_user_message*/ None, diff --git a/codex-rs/tui/src/app/loaded_threads.rs b/codex-rs/tui/src/app/loaded_threads.rs index 7be4df391e60..c6f24e6a327c 100644 --- a/codex-rs/tui/src/app/loaded_threads.rs +++ b/codex-rs/tui/src/app/loaded_threads.rs @@ -17,6 +17,7 @@ use crate::app_server_session::thread_blocks_direct_input; use codex_app_server_protocol::SessionSource; use codex_app_server_protocol::Thread; +use codex_app_server_protocol::ThreadStatus; use codex_protocol::ThreadId; use codex_protocol::protocol::SubAgentSource; use std::collections::HashMap; @@ -31,6 +32,8 @@ pub(crate) struct LoadedSubagentThread { pub(crate) agent_role: Option, pub(crate) agent_path: Option, pub(crate) blocks_direct_input: bool, + pub(crate) is_running: bool, + pub(crate) is_closed: bool, } /// Walks the spawn tree rooted at `primary_thread_id` and returns every descendant subagent. @@ -87,6 +90,8 @@ pub(crate) fn find_loaded_subagent_threads_for_primary( .remove(&thread_id) .map(|thread| LoadedSubagentThread { blocks_direct_input: thread_blocks_direct_input(&thread), + is_running: matches!(&thread.status, ThreadStatus::Active { .. }), + is_closed: matches!(&thread.status, ThreadStatus::NotLoaded), thread_id, agent_nickname: thread.agent_nickname, agent_role: thread.agent_role, @@ -196,6 +201,9 @@ mod tests { child.agent_nickname = Some("Scout".to_string()); child.agent_role = Some("explorer".to_string()); child.can_accept_direct_input = Some(true); + child.status = ThreadStatus::Active { + active_flags: Vec::new(), + }; let mut grandchild = test_thread( grandchild_thread_id, @@ -204,6 +212,7 @@ mod tests { grandchild.agent_nickname = Some("Atlas".to_string()); grandchild.agent_role = Some("worker".to_string()); grandchild.can_accept_direct_input = Some(false); + grandchild.status = ThreadStatus::NotLoaded; let unrelated_child = test_thread( unrelated_child_id, thread_spawn_source(unrelated_parent_id, /*depth*/ 1, "Other", "researcher"), @@ -228,6 +237,8 @@ mod tests { agent_nickname: Some("Scout".to_string()), agent_role: Some("explorer".to_string()), agent_path: None, + is_running: true, + is_closed: false, }, LoadedSubagentThread { blocks_direct_input: true, @@ -235,6 +246,8 @@ mod tests { agent_nickname: Some("Atlas".to_string()), agent_role: Some("worker".to_string()), agent_path: None, + is_running: false, + is_closed: true, }, ] ); diff --git a/codex-rs/tui/src/app/safety_buffering.rs b/codex-rs/tui/src/app/safety_buffering.rs index 1463f03f1403..d9f223db0e44 100644 --- a/codex-rs/tui/src/app/safety_buffering.rs +++ b/codex-rs/tui/src/app/safety_buffering.rs @@ -135,7 +135,6 @@ impl App { if let Err(err) = self .replace_chat_widget_with_app_server_thread( tui, - app_server, started, ThreadAttachPresentation::SessionLineage, /*initial_user_message*/ None, diff --git a/codex-rs/tui/src/app/session_lifecycle.rs b/codex-rs/tui/src/app/session_lifecycle.rs index e0eee231ef76..ac473d98ae5a 100644 --- a/codex-rs/tui/src/app/session_lifecycle.rs +++ b/codex-rs/tui/src/app/session_lifecycle.rs @@ -8,6 +8,7 @@ use super::*; use crate::app_server_session::source_agent_path; use crate::app_server_session::thread_blocks_direct_input; use codex_config::types::ResumeCwdMode; +use std::collections::HashSet; #[derive(Clone, Copy)] pub(super) enum ThreadAttachPresentation { @@ -15,9 +16,17 @@ pub(super) enum ThreadAttachPresentation { PromptEdit, } +/// Reports whether a loaded-thread backfill completed and which descendants already had their +/// liveness metadata refreshed, allowing the picker to skip duplicate `thread/read` requests. +#[derive(Default)] +pub(super) struct LoadedSubagentBackfill { + pub(super) completed: bool, + pub(super) refreshed_thread_ids: HashSet, +} + impl App { pub(super) async fn open_agent_picker(&mut self, app_server: &mut AppServerSession) { - self.backfill_loaded_subagent_threads(app_server).await; + let backfill = self.backfill_loaded_subagent_threads(app_server).await; // V2 subagents are identified by canonical paths observed from activity events or loaded // thread metadata. A buffered active turn is positive liveness evidence; a completed // snapshot is terminal evidence. An empty store does not clear a successful spawn hint. @@ -46,7 +55,7 @@ impl App { } else if has_terminal_snapshot { self.agent_navigation.mark_stopped(thread_id); } - } else { + } else if !backfill.refreshed_thread_ids.contains(&thread_id) { self.refresh_agent_picker_thread_liveness(app_server, thread_id) .await; } @@ -92,6 +101,7 @@ impl App { for thread_id in thread_ids { if path_backed_thread_ids.contains(&thread_id) || self.side_threads.contains_key(&thread_id) + || backfill.refreshed_thread_ids.contains(&thread_id) { continue; } @@ -631,7 +641,6 @@ impl App { if let Err(err) = self .replace_chat_widget_with_app_server_thread( tui, - app_server, started, ThreadAttachPresentation::SessionLineage, initial_user_message, @@ -666,7 +675,6 @@ impl App { pub(super) async fn replace_chat_widget_with_app_server_thread( &mut self, tui: &mut tui::Tui, - app_server: &mut AppServerSession, started: AppServerStartedThread, presentation: ThreadAttachPresentation, initial_user_message: Option, @@ -690,16 +698,15 @@ impl App { presentation, ) .await?; - self.backfill_loaded_subagent_threads(app_server).await; Ok(()) } /// Fetches all loaded threads from the app server and registers descendants of the primary /// thread in the navigation cache and chat widget metadata. /// - /// Called after `replace_chat_widget_with_app_server_thread` during resume, fork, and new - /// thread creation so that the `/agent` picker and keyboard navigation are pre-populated even - /// if the TUI did not witness the original spawn events. + /// Called when opening the `/agent` picker and after resuming a thread so that the picker and + /// keyboard navigation are pre-populated even if the TUI did not witness the original spawn + /// events. Fresh and forked threads cannot have pre-existing descendants. /// /// The loaded-thread list is fetched in full (no pagination) and the spawn tree is walked /// by `find_loaded_subagent_threads_for_primary`. Each discovered subagent is registered via @@ -708,9 +715,9 @@ impl App { pub(super) async fn backfill_loaded_subagent_threads( &mut self, app_server: &mut AppServerSession, - ) -> bool { + ) -> LoadedSubagentBackfill { let Some(primary_thread_id) = self.primary_thread_id else { - return false; + return LoadedSubagentBackfill::default(); }; let loaded_thread_ids = match app_server @@ -723,7 +730,7 @@ impl App { Ok(response) => response.data, Err(err) => { tracing::warn!(%err, "failed to list loaded threads for subagent backfill"); - return false; + return LoadedSubagentBackfill::default(); } }; @@ -751,8 +758,14 @@ impl App { } } + let mut refreshed_thread_ids = HashSet::new(); for thread in find_loaded_subagent_threads_for_primary(threads, primary_thread_id) { let agent_path = thread.agent_path; + let has_live_channel = self + .thread_event_channels + .get(&thread.thread_id) + .is_some_and(|channel| channel.attachment() == ThreadEventAttachment::Live); + let is_closed = !has_live_channel && thread.is_closed; if thread.blocks_direct_input { self.agent_navigation.mark_parent_owned(thread.thread_id); } @@ -760,14 +773,28 @@ impl App { thread.thread_id, thread.agent_nickname, thread.agent_role, - /*is_closed*/ false, + is_closed, ); self.agent_navigation .set_agent_path(thread.thread_id, agent_path); + // A live channel can have an empty store after a successful spawn. Only apply server + // status for channels that would otherwise need another liveness read. + if !has_live_channel { + if thread.is_running { + self.agent_navigation.mark_running(thread.thread_id); + } else { + self.agent_navigation + .set_running(thread.thread_id, /*is_running*/ false); + } + refreshed_thread_ids.insert(thread.thread_id); + } } self.sync_active_agent_label(); - !had_read_error + LoadedSubagentBackfill { + completed: !had_read_error, + refreshed_thread_ids, + } } /// Returns the adjacent thread id for keyboard navigation, backfilling from the server if the @@ -796,7 +823,11 @@ impl App { return None; } - if self.backfill_loaded_subagent_threads(app_server).await { + if self + .backfill_loaded_subagent_threads(app_server) + .await + .completed + { self.last_subagent_backfill_attempt = Some(primary_thread_id); } self.agent_navigation @@ -932,7 +963,6 @@ impl App { match self .replace_chat_widget_with_app_server_thread( tui, - app_server, resumed, ThreadAttachPresentation::SessionLineage, /*initial_user_message*/ None, @@ -940,6 +970,7 @@ impl App { .await { Ok(()) => { + self.backfill_loaded_subagent_threads(app_server).await; if let Some(summary) = summary { let mut lines: Vec> = Vec::new(); if let Some(usage_line) = summary.usage_line { diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 48efbba2e08a..381df134cda9 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -6,6 +6,8 @@ mod model_catalog; mod plugin_catalog; mod rate_limits; mod safety_buffering; +#[path = "tests/session_lifecycle_requests.rs"] +mod session_lifecycle_requests; mod session_summary; mod startup; @@ -1791,7 +1793,29 @@ fn selected_and_resumed_threads_use_server_capability_for_v1_and_v2_children() - child_thread_ids.push(child_thread_id); } - assert!(app.backfill_loaded_subagent_threads(&mut app_server).await); + app.agent_navigation + .record_sub_agent_activity(SubAgentActivityDisplay { + thread_id: child_thread_ids[0], + agent_path: "/root/child-0".to_string(), + is_running_hint: true, + }); + app.thread_event_channels.remove(&child_thread_ids[1]); + let backfill = app.backfill_loaded_subagent_threads(&mut app_server).await; + assert!(backfill.completed); + assert_eq!( + backfill.refreshed_thread_ids, + [child_thread_ids[1]].into_iter().collect() + ); + assert_eq!( + app.agent_navigation.get(&child_thread_ids[0]), + Some(&AgentPickerThreadEntry { + agent_nickname: Some("child-0".to_string()), + agent_role: Some("worker".to_string()), + agent_path: Some("/root/child-0".to_string()), + is_running: true, + is_closed: false, + }) + ); assert!(!app.agent_navigation.is_parent_owned(child_thread_ids[0])); assert!(app.agent_navigation.is_parent_owned(child_thread_ids[1])); @@ -1833,7 +1857,6 @@ fn selected_and_resumed_threads_use_server_capability_for_v1_and_v2_children() - assert!(resumed.blocks_direct_input); app.replace_chat_widget_with_app_server_thread( &mut tui, - &mut app_server, resumed, crate::app::session_lifecycle::ThreadAttachPresentation::SessionLineage, /*initial_user_message*/ None, diff --git a/codex-rs/tui/src/app/tests/safety_buffering.rs b/codex-rs/tui/src/app/tests/safety_buffering.rs index 751bedd98a11..324cdc2e3a29 100644 --- a/codex-rs/tui/src/app/tests/safety_buffering.rs +++ b/codex-rs/tui/src/app/tests/safety_buffering.rs @@ -289,7 +289,6 @@ goals = true let source_thread_id = started.session.thread_id; app.replace_chat_widget_with_app_server_thread( &mut tui, - &mut app_server, started, ThreadAttachPresentation::SessionLineage, /*initial_user_message*/ None, diff --git a/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs b/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs new file mode 100644 index 000000000000..9c90e23af086 --- /dev/null +++ b/codex-rs/tui/src/app/tests/session_lifecycle_requests.rs @@ -0,0 +1,292 @@ +use super::*; +use app_test_support::create_fake_parented_rollout_with_source; +use app_test_support::create_fake_rollout; +use app_test_support::rollout_path; +use codex_app_server_protocol::ClientNotification; +use codex_app_server_protocol::ClientRequest; +use codex_app_server_protocol::JSONRPCError; +use codex_app_server_protocol::JSONRPCMessage; +use codex_app_server_protocol::JSONRPCResponse; +use codex_protocol::AgentPath; +use futures::SinkExt; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use std::sync::Mutex; +use tokio::net::TcpListener; +use tokio::task::JoinHandle; +use tokio_tungstenite::accept_async; +use tokio_tungstenite::tungstenite::Message; + +/// Returns and resets `(thread/loaded/list, thread/read)` request counts. +fn take_backfill_counts(requests: &Arc>>) -> (usize, usize) { + let requests = std::mem::take(&mut *requests.lock().expect("request recorder lock")); + ( + requests + .iter() + .filter(|method| *method == "thread/loaded/list") + .count(), + requests + .iter() + .filter(|method| *method == "thread/read") + .count(), + ) +} + +/// Starts an embedded app server behind a loopback WebSocket proxy that records JSON-RPC methods. +async fn start_recording_app_server( + config: &Config, +) -> Result<( + AppServerSession, + Arc>>, + JoinHandle>, +)> { + let state_db = + crate::init_state_db_for_app_server_target(config, &crate::AppServerTarget::Embedded) + .await?; + let embedded = crate::start_embedded_app_server( + codex_arg0::Arg0DispatchPaths::default(), + config.clone(), + Vec::new(), + codex_config::LoaderOverrides::default(), + /*strict_config*/ false, + codex_config::CloudConfigBundleLoader::default(), + codex_feedback::CodexFeedback::new(), + /*log_db*/ None, + state_db, + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ) + .await?; + let codex_home = config.codex_home.display().to_string(); + let requests = Arc::new(Mutex::new(Vec::new())); + let request_sink = Arc::clone(&requests); + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let proxy = tokio::spawn(async move { + let (stream, _) = listener.accept().await?; + let mut websocket = accept_async(stream).await?; + while let Some(frame) = websocket.next().await { + let Message::Text(text) = frame? else { + continue; + }; + let message = serde_json::from_str::(&text)?; + match message { + JSONRPCMessage::Request(request) if request.method == "initialize" => { + websocket + .send(Message::Text( + serde_json::to_string(&JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::json!({ + "userAgent": "codex-tui-test", + "codexHome": codex_home, + }), + }))? + .into(), + )) + .await?; + } + JSONRPCMessage::Request(request) => { + request_sink + .lock() + .expect("request recorder lock") + .push(request.method.clone()); + let request_id = request.id.clone(); + let request = + serde_json::from_value::(serde_json::to_value(request)?)?; + let response = match embedded.request(request).await? { + Ok(result) => JSONRPCMessage::Response(JSONRPCResponse { + id: request_id, + result, + }), + Err(error) => JSONRPCMessage::Error(JSONRPCError { + id: request_id, + error, + }), + }; + websocket + .send(Message::Text(serde_json::to_string(&response)?.into())) + .await?; + } + JSONRPCMessage::Notification(notification) + if notification.method == "initialized" => {} + JSONRPCMessage::Notification(notification) => { + embedded + .notify(serde_json::from_value::( + serde_json::to_value(notification)?, + )?) + .await?; + } + JSONRPCMessage::Response(_) | JSONRPCMessage::Error(_) => {} + } + } + embedded.shutdown().await?; + Ok(()) + }); + let app_server = crate::connect_remote_app_server(crate::RemoteAppServerEndpoint::WebSocket { + websocket_url, + auth_token: None, + }) + .await?; + + Ok(( + AppServerSession::new( + app_server, + crate::app_server_session::ThreadParamsMode::Embedded, + ), + requests, + proxy, + )) +} + +#[test] +fn session_lifecycle_avoids_redundant_subagent_metadata_reads() -> Result<()> { + const TEST_STACK_SIZE_BYTES: usize = 8 * 1024 * 1024; + + std::thread::Builder::new() + .name("tui-session-lifecycle-requests".to_string()) + .stack_size(TEST_STACK_SIZE_BYTES) + .spawn(|| { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + runtime.block_on(async { + let mut app = make_test_app().await; + let codex_home = tempdir()?; + app.config.codex_home = codex_home.path().to_path_buf().abs(); + app.config.sqlite_home = codex_home.path().to_path_buf(); + let root_timestamp = "2026-01-01T00-00-00"; + let root_thread_id = ThreadId::from_string( + &create_fake_rollout( + codex_home.path(), + root_timestamp, + "2026-01-01T00:00:00Z", + "Saved user message", + Some(app.config.model_provider_id.as_str()), + /*git_info*/ None, + ) + .expect("create root rollout"), + )?; + let child_thread_id = ThreadId::from_string( + &create_fake_parented_rollout_with_source( + codex_home.path(), + "2026-01-01T00-00-01", + "2026-01-01T00:00:01Z", + "Saved child message", + Some(app.config.model_provider_id.as_str()), + /*git_info*/ None, + RolloutSessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: root_thread_id, + depth: 1, + agent_path: Some( + AgentPath::try_from("/root/worker").expect("valid agent path"), + ), + agent_nickname: Some("worker".to_string()), + agent_role: Some("worker".to_string()), + }), + root_thread_id.into(), + root_thread_id, + ) + .expect("create child rollout"), + )?; + let root_rollout_path = rollout_path( + codex_home.path(), + root_timestamp, + &root_thread_id.to_string(), + ); + let (mut app_server, requests, proxy) = + start_recording_app_server(&app.config).await?; + let root = app_server + .resume_thread( + app.config.clone(), + root_thread_id, + app.resume_model_settings(), + ) + .await?; + app.enqueue_primary_thread_session(root.session, root.turns) + .await?; + app_server + .resume_thread( + app.config.clone(), + child_thread_id, + app.resume_model_settings(), + ) + .await?; + let mut tui = crate::tui::test_support::make_test_tui()?; + take_backfill_counts(&requests); + + let control = Box::pin(app.handle_event( + &mut tui, + &mut app_server, + AppEvent::ForkCurrentSession, + )) + .await?; + + assert!(matches!(control, AppRunControl::Continue)); + assert_ne!(app.chat_widget.thread_id(), Some(root_thread_id)); + // Forking may read the source metadata once when the response includes its parent + // id. It must not scan or backfill loaded threads for the newly created fork. + assert!(matches!(take_backfill_counts(&requests), (0, 0) | (0, 1))); + + app.start_fresh_session_with_summary_hint( + &mut tui, + &mut app_server, + /*session_start_source*/ None, + /*initial_user_message*/ None, + ) + .await; + + assert_ne!(app.chat_widget.thread_id(), Some(root_thread_id)); + assert_eq!(take_backfill_counts(&requests), (0, 0)); + + let loaded_threads = app_server + .thread_loaded_list(ThreadLoadedListParams { + cursor: None, + limit: None, + }) + .await? + .data; + let expected_reads = loaded_threads + .iter() + .filter(|thread_id| *thread_id != &root_thread_id.to_string()) + .count(); + assert!(loaded_threads.contains(&child_thread_id.to_string())); + take_backfill_counts(&requests); + app.harness_overrides.cwd = Some(app.config.cwd.to_path_buf()); + + let control = app + .resume_target_session( + &mut tui, + &mut app_server, + crate::resume_picker::SessionTarget { + path: Some(root_rollout_path), + thread_id: root_thread_id, + }, + ) + .await?; + + assert!(matches!(control, AppRunControl::Continue)); + assert_eq!(app.chat_widget.thread_id(), Some(root_thread_id)); + assert_eq!(take_backfill_counts(&requests), (1, expected_reads)); + assert_eq!( + app.agent_navigation.get(&child_thread_id), + Some(&AgentPickerThreadEntry { + agent_nickname: Some("worker".to_string()), + agent_role: Some("worker".to_string()), + agent_path: Some("/root/worker".to_string()), + is_running: false, + is_closed: false, + }) + ); + + Box::pin(app.open_agent_picker(&mut app_server)).await; + + // The picker refreshes the primary thread once. Discovered children were already + // refreshed by the picker's initial backfill and must not be read a second time. + assert_eq!(take_backfill_counts(&requests), (1, expected_reads + 1)); + app_server.shutdown().await?; + proxy.await??; + Ok(()) + }) + })? + .join() + .expect("session lifecycle request test thread") +}