Avoid redundant TUI subagent metadata requests (#34234)

## What changed

- Skip loaded-subagent backfills for fresh and forked threads, which cannot have pre-existing descendants.
- Keep the backfill after resuming a thread so agent navigation is populated immediately.
- Reuse thread status collected during backfill when opening the agent picker, avoiding duplicate `thread/read` requests while preserving liveness from live event channels.

## Testing

- Add a WebSocket-backed session lifecycle test that records app-server requests across fork, fresh-session, resume, and agent-picker flows.

GitOrigin-RevId: 604742b3cb495da12c97908d7a0ffa4ebc6a6da3
This commit is contained in:
Charlie Marsh
2026-07-19 21:39:58 +00:00
committed by copyberry
parent a97ae65362
commit 678157acaa
9 changed files with 380 additions and 21 deletions
+2
View File
@@ -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",
+2
View File
@@ -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 }
-2
View File
@@ -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,
+13
View File
@@ -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<String>,
pub(crate) agent_path: Option<String>,
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,
},
]
);
-1
View File
@@ -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,
+46 -15
View File
@@ -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<ThreadId>,
}
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<crate::chatwidget::UserMessage>,
@@ -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<Line<'static>> = Vec::new();
if let Some(usage_line) = summary.usage_line {
+25 -2
View File
@@ -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,
@@ -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,
@@ -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<Mutex<Vec<String>>>) -> (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<Mutex<Vec<String>>>,
JoinHandle<Result<()>>,
)> {
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::<JSONRPCMessage>(&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::<ClientRequest>(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::<ClientNotification>(
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")
}