From 95e0103a5de64e6a097349c37a3f84c42decc8f4 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Fri, 13 Nov 2026 03:56:56 +0000 Subject: [PATCH] Add bounded streaming discovery for statusline providers (#61250) ## What changed Add a statusline provider discovery stream that publishes the first inventory or named lookup result for each MCP server. Give global inventory a two-second head start, then probe known servers with at most four concurrent lookups while continuing to discover servers from inventory pages. Completed inventory triggers probes for missing known servers immediately. Bound each refresh and each selected-provider lookup to five seconds, with independent deadlines. Report server and inventory failures separately, and cancel outstanding discovery waits when the stream is dropped. ## Testing Add tests with controlled app-server replies and paused time covering inventory/probe races, concurrency limits, timeouts, cancellation and reopening, missing servers, and repeated pagination cursors. GitOrigin-RevId: e2e1af9789bb67236afdbc76ec2d8ee1648d6ac4 --- codex-rs/tui/src/statusline_provider.rs | 28 +- .../tui/src/statusline_provider/discovery.rs | 118 +++++++ .../statusline_provider/discovery_tests.rs | 311 ++++++++++++++++++ 3 files changed, 455 insertions(+), 2 deletions(-) create mode 100644 codex-rs/tui/src/statusline_provider/discovery.rs create mode 100644 codex-rs/tui/src/statusline_provider/discovery_tests.rs diff --git a/codex-rs/tui/src/statusline_provider.rs b/codex-rs/tui/src/statusline_provider.rs index 7ebd18293c..e9112c45f7 100644 --- a/codex-rs/tui/src/statusline_provider.rs +++ b/codex-rs/tui/src/statusline_provider.rs @@ -1,7 +1,7 @@ //! Resolve globally advertised statusline tools and read them on the displayed thread. mod descriptor; -#[cfg_attr(not(test), expect(dead_code, reason = "staged provider discovery"))] +pub(crate) mod discovery; mod inventory; #[cfg(test)] mod test_support; @@ -21,9 +21,20 @@ use color_eyre::eyre::bail; pub(crate) use descriptor::StatuslineProvider; use serde_json::Value; use std::collections::HashMap; +use tokio::time::Duration; use uuid::Uuid; -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] +const LOOKUP_TIMEOUT: Duration = Duration::from_secs(5); + +#[derive(Debug, thiserror::Error)] +pub(crate) enum ProviderLookupError { + #[error("statusline provider lookup timed out after 5 seconds")] + TimedOut, + #[error(transparent)] + Request(#[from] color_eyre::Report), +} + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub(crate) struct McpServerName(String); impl McpServerName { @@ -74,6 +85,19 @@ impl ResolvedStatuslineProvider { thread: ThreadId, cwd: &PathUri, selection: ProviderSelection<'_>, + ) -> Result, ProviderLookupError> { + tokio::select! { + biased; + () = tokio::time::sleep(LOOKUP_TIMEOUT) => Err(ProviderLookupError::TimedOut), + result = Self::resolve(client, thread, cwd, selection) => Ok(result?), + } + } + + async fn resolve( + client: &AppServerRequestHandle, + thread: ThreadId, + cwd: &PathUri, + selection: ProviderSelection<'_>, ) -> Result> { let config = read_effective_config_with_layers(client, cwd.inferred_native_path_string()).await?; diff --git a/codex-rs/tui/src/statusline_provider/discovery.rs b/codex-rs/tui/src/statusline_provider/discovery.rs new file mode 100644 index 0000000000..25adf8cbbb --- /dev/null +++ b/codex-rs/tui/src/statusline_provider/discovery.rs @@ -0,0 +1,118 @@ +//! Stream the first inventory or named lookup result for each server in a picker refresh. +//! Inventory stays open to discover new servers while named lookups are in flight. + +use super::LOOKUP_TIMEOUT; +use super::McpServerName; +use super::ProviderLookupError; +use super::StatuslineProvider; +use super::inventory; +use codex_app_server_client::AppServerRequestHandle; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::stream; +use std::collections::BTreeSet; +use tokio::sync::mpsc; +use tokio::time::Duration; +use tokio::time::Instant; +use tokio_stream::StreamMap; +use tokio_stream::wrappers::UnboundedReceiverStream; + +const INVENTORY_HEAD_START: Duration = Duration::from_secs(2); +const MAX_CONCURRENT_PROBES: usize = 4; + +#[derive(Debug)] +#[cfg_attr(not(test), expect(dead_code, reason = "staged picker"))] +pub(crate) enum DiscoveryUpdate { + Server( + McpServerName, + Result, ProviderLookupError>, + ), + InventoryFailed(ProviderLookupError), +} + +/// Pass the saved and previously observed server names. Dropping the stream stops the TUI's +/// outstanding waits. The host may finish requests it already received. +#[cfg_attr(not(test), expect(dead_code, reason = "staged picker"))] +pub(crate) fn refresh_statusline_providers( + client: AppServerRequestHandle, + known: Vec, +) -> impl futures::Stream + Unpin { + let start = Instant::now(); + let (updates, receiver) = mpsc::unbounded_channel(); + tokio::spawn(async move { + tokio::select! { + biased; + () = updates.closed() => {} + result = discover(&client, start, known, &updates) => { + if let Err(error) = result { + let _ = updates.send(DiscoveryUpdate::InventoryFailed(error)); + } + } + } + }); + UnboundedReceiverStream::new(receiver) +} + +async fn discover( + client: &AppServerRequestHandle, + start: Instant, + known: Vec, + updates: &mpsc::UnboundedSender, +) -> Result<(), ProviderLookupError> { + let mut queued: BTreeSet<_> = known.into_iter().collect(); + let mut published = BTreeSet::new(); + let mut probes = StreamMap::::new(); + let pages = inventory::global_pages(client); + tokio::pin!(pages); + let mut listing = None; + while listing.is_none() || !queued.is_empty() || !probes.is_empty() { + tokio::select! { + biased; + () = tokio::time::sleep_until(start + LOOKUP_TIMEOUT) => break, + Some((name, providers)) = probes.next(), if !probes.is_empty() => { + probes.remove(&name); + published.insert(name.clone()); + let _ = updates.send(DiscoveryUpdate::Server(name, providers)); + } + page = pages.try_next(), if listing.is_none() => match page { + Ok(Some(servers)) => for server in servers { + if published.insert(server.name.clone()) { + queued.remove(&server.name); + // A late catalog frees a named lookup slot even if that server never replies. + probes.remove(&server.name); + let providers = server.providers().map_err(Into::into); + let _ = updates.send(DiscoveryUpdate::Server(server.name, providers)); + } + }, + Ok(None) => listing = Some(Ok(())), + Err(error) => listing = Some(Err(error.into())), + }, + Some(name) = async { + if listing.is_none() { + tokio::time::sleep_until(start + INVENTORY_HEAD_START).await; + } + queued.pop_first() + }, if !queued.is_empty() && probes.len() < MAX_CONCURRENT_PROBES => { + let key = name.clone(); + let lookup = async move { + Ok(match inventory::server(client, &name, /*thread*/ None).await? { + Some(server) => server.providers()?, + None => Vec::new(), + }) + }; + probes.insert(key, Box::pin(stream::once(lookup))); + } + } + } + for name in queued.into_iter().chain(probes.keys().cloned()) { + let _ = updates.send(DiscoveryUpdate::Server( + name, + Err(ProviderLookupError::TimedOut), + )); + } + listing.unwrap_or(Err(ProviderLookupError::TimedOut)) +} + +#[cfg(test)] +#[path = "discovery_tests.rs"] +mod tests; diff --git a/codex-rs/tui/src/statusline_provider/discovery_tests.rs b/codex-rs/tui/src/statusline_provider/discovery_tests.rs new file mode 100644 index 0000000000..e7b258daf2 --- /dev/null +++ b/codex-rs/tui/src/statusline_provider/discovery_tests.rs @@ -0,0 +1,311 @@ +#![cfg(test)] +//! Drive global and selected inventory using independently controlled app-server replies. + +use super::super::ProviderId; +use super::super::ProviderSelection; +use super::super::ResolvedStatuslineProvider; +use super::super::test_support::test_client; +use super::*; +use codex_app_server_client::CloudCommand; +use codex_app_server_client::CloudOperation; +use codex_app_server_client::RequestResult; +use codex_app_server_protocol::ClientRequest; +use codex_protocol::ThreadId; +use codex_utils_path_uri::PathUri; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use std::collections::BTreeMap; +use tokio::sync::mpsc; +use tokio::sync::oneshot; + +type Reply = oneshot::Sender>; +const URI: &str = "ui://test/footer"; + +#[derive(Debug)] +enum Call<'a> { + Page(Option<&'a str>), + Global(&'a str), +} +use Call::Global; +use Call::Page; +use DiscoveryUpdate::InventoryFailed; +use DiscoveryUpdate::Server; +use ProviderLookupError::Request; +use ProviderLookupError::TimedOut; + +struct Peer(mpsc::Receiver); + +impl Peer { + fn connect() -> (AppServerRequestHandle, Self) { + let (client, commands) = test_client(); + (client, Self(commands)) + } + + async fn request(&mut self) -> (ClientRequest, Reply) { + let command = self.0.recv().await.unwrap(); + let CloudOperation::Request(request) = command.operation else { + panic!("expected an app-server request"); + }; + (*request, command.response) + } + + async fn config(&mut self, cwd: &str) -> Reply { + let (ClientRequest::ConfigRead { params, .. }, response) = self.request().await else { + panic!("expected config/read"); + }; + assert_eq!( + (params.include_layers, params.cwd.as_deref()), + (true, Some(cwd)) + ); + response + } + + async fn expect(&mut self, calls: [Call<'_>; N]) -> [Reply; N] { + let mut replies = BTreeMap::new(); + for _ in 0..N { + let (ClientRequest::McpServerStatusList { params, .. }, response) = + self.request().await + else { + panic!("expected mcpServerStatus/list"); + }; + assert_eq!((params.thread_id, params.limit), (None, None)); + let key = (params.server_name, params.cursor); + assert!(replies.insert(key, response).is_none()); + } + calls.map(|call| { + let key = match call { + Page(cursor) => (None, cursor.map(str::to_owned)), + Global(name) => (Some(name.to_owned()), None), + }; + replies + .remove(&key) + .unwrap_or_else(|| panic!("missing {call:?}")) + }) + } +} + +fn reply(to: Reply, data: Value) { + to.send(Ok(Ok(data))).unwrap(); +} + +fn row(name: &str, uri: Option<&str>) -> Value { + let mut row = super::super::test_support::row(name, vec![]); + row["runtimeStatus"] = Value::Null; + row["pluginId"] = json!("daily@local"); + if let Some(uri) = uri { + row["tools"]["footer"] = json!({"name":"footer","inputSchema":{},"_meta":{"openai/ui":{"entrypoints":[{"type":"statusline","resourceUri":uri}]}}}); + } + row +} + +fn reply_page(to: Reply, rows: &[Value], next: Option<&str>) { + reply(to, json!({"data":rows,"nextCursor":next})); +} + +fn catalog(update: Option) -> (McpServerName, Vec) { + let Some(Server(name, Ok(providers))) = update else { + panic!("expected a provider catalog, got {update:?}"); + }; + ( + name, + providers.into_iter().map(|provider| provider.uri).collect(), + ) +} + +fn assert_uri(update: Option, server: &str, uri: &str) { + assert_eq!(catalog(update), (server.into(), vec![uri.to_owned()])); +} + +#[tokio::test(start_paused = true)] +async fn a_fast_inventory_publishes_each_server_without_named_probes() { + let (client, mut peer) = Peer::connect(); + let known = ["ready", "empty", "failed"].map(McpServerName::from).into(); + let mut updates = refresh_statusline_providers(client, known); + let [first] = peer.expect([Page(None)]).await; + tokio::time::advance(Duration::from_millis(1500)).await; + let mut failed = row("failed", /*uri*/ None); + failed["toolsError"] = json!("offline"); + reply_page( + first, + &[row("ready", Some(URI)), row("empty", /*uri*/ None), failed], + Some("next"), + ); + assert_uri(updates.next().await, "ready", URI); + assert_eq!(catalog(updates.next().await), ("empty".into(), vec![])); + assert!( + matches!(updates.next().await, Some(Server(name, Err(Request(e)))) + if name.as_str() == "failed" && e.to_string() == "offline") + ); + let [last] = peer.expect([Page(Some("next"))]).await; + reply_page(last, &[], /*next*/ None); + assert!(updates.next().await.is_none()); +} + +#[tokio::test(start_paused = true)] +async fn late_inventory_frees_probe_slots_and_keeps_discovering_servers() { + let (client, mut peer) = Peer::connect(); + let known = ["a", "b", "c", "d", "e", "queued"] + .map(McpServerName::from) + .into(); + let start = Instant::now(); + let mut updates = refresh_statusline_providers(client, known); + let [first] = peer.expect([Page(None)]).await; + let [a, b, c, d] = peer + .expect([Global("a"), Global("b"), Global("c"), Global("d")]) + .await; + assert_eq!(Instant::now() - start, INVENTORY_HEAD_START); + assert!(matches!( + peer.0.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + + reply_page( + first, + &[row("b", Some(URI)), row("queued", Some(URI))], + Some("next"), + ); + assert_uri(updates.next().await, "b", URI); + assert_uri(updates.next().await, "queued", URI); + let [e, last] = peer.expect([Global("e"), Page(Some("next"))]).await; + assert!(b.is_closed()); + assert_eq!(Instant::now() - start, INVENTORY_HEAD_START); + for (name, probe) in [("a", a), ("c", c), ("d", d), ("e", e)] { + reply_page(probe, &[row(name, Some(URI))], /*next*/ None); + assert_uri(updates.next().await, name, URI); + } + reply_page( + last, + &[row("e", Some("ui://old/e")), row("new-plugin", Some(URI))], + /*next*/ None, + ); + assert_uri(updates.next().await, "new-plugin", URI); + assert!(updates.next().await.is_none()); +} + +#[tokio::test(start_paused = true)] +async fn a_late_page_preserves_completed_probes_and_unanswered_names_time_out() { + let (client, mut peer) = Peer::connect(); + let known = ["empty", "failed", "slow"].map(McpServerName::from).into(); + let start = Instant::now(); + let mut updates = refresh_statusline_providers(client, known); + let [page, empty, failed, _slow] = peer + .expect([ + Page(None), + Global("empty"), + Global("failed"), + Global("slow"), + ]) + .await; + + reply_page(empty, &[], /*next*/ None); + assert_eq!(catalog(updates.next().await), ("empty".into(), vec![])); + let mut error = row("failed", /*uri*/ None); + error["toolsError"] = json!("no catalog"); + reply_page(failed, &[error], /*next*/ None); + assert!( + matches!(updates.next().await, Some(Server(name, Err(Request(e)))) + if name.as_str() == "failed" && e.to_string() == "no catalog") + ); + reply_page( + page, + &[row("empty", Some(URI)), row("failed", Some(URI))], + /*next*/ None, + ); + assert!( + matches!(updates.next().await, Some(Server(name, Err(TimedOut))) if name.as_str() == "slow") + ); + assert!(updates.next().await.is_none()); + assert_eq!(Instant::now() - start, LOOKUP_TIMEOUT); +} + +#[tokio::test(start_paused = true)] +async fn a_cold_picker_and_configured_provider_have_independent_deadlines() { + let (client, mut peer) = Peer::connect(); + let start = Instant::now(); + let mut picker = refresh_statusline_providers(client.clone(), Vec::new()); + let [_inventory] = peer.expect([Page(None)]).await; + tokio::time::advance(Duration::from_secs(1)).await; + let selected_start = Instant::now(); + let selected = tokio::spawn(async move { + let cwd = PathUri::parse("file:///workspace").unwrap(); + let id = ProviderId { + server: "weather".into(), + tool: "footer".into(), + }; + ResolvedStatuslineProvider::lookup( + &client, + ThreadId::new(), + &cwd, + ProviderSelection::Saved(&id), + ) + .await + }); + let config = peer.config("/workspace").await; + tokio::time::advance(Duration::from_secs(2)).await; + reply(config, json!({"config":{},"layers":[]})); + let [_global] = peer.expect([Global("weather")]).await; + assert!(matches!( + picker.next().await, + Some(InventoryFailed(TimedOut)) + )); + assert!(picker.next().await.is_none()); + assert_eq!(Instant::now() - start, LOOKUP_TIMEOUT); + assert!(matches!(selected.await.unwrap(), Err(TimedOut))); + assert_eq!(Instant::now() - selected_start, LOOKUP_TIMEOUT); +} + +#[tokio::test(start_paused = true)] +async fn closing_cancels_requests_and_reopening_reads_the_current_inventory() { + let (client, mut peer) = Peer::connect(); + let mut first = + refresh_statusline_providers(client.clone(), vec!["calendar".into(), "slow".into()]); + let [mut inventory, calendar, mut probe] = peer + .expect([Page(None), Global("calendar"), Global("slow")]) + .await; + reply_page(calendar, &[row("calendar", Some(URI))], /*next*/ None); + assert_uri(first.next().await, "calendar", URI); + drop(first); + tokio::join!(inventory.closed(), probe.closed()); + + let mut reopened = refresh_statusline_providers(client, vec!["calendar".into()]); + let [page] = peer.expect([Page(None)]).await; + let new_uri = "ui://calendar/changed"; + reply_page(page, &[row("calendar", Some(new_uri))], /*next*/ None); + assert_uri(reopened.next().await, "calendar", new_uri); + assert!(reopened.next().await.is_none()); +} + +#[tokio::test(start_paused = true)] +async fn completed_inventory_probes_a_missing_cached_name_immediately() { + let (client, mut peer) = Peer::connect(); + let start = Instant::now(); + let mut updates = refresh_statusline_providers(client, vec!["removed".into()]); + let [page] = peer.expect([Page(None)]).await; + reply_page(page, &[], /*next*/ None); + let [removed] = peer.expect([Global("removed")]).await; + reply_page(removed, &[], /*next*/ None); + assert_eq!(catalog(updates.next().await), ("removed".into(), vec![])); + assert!(updates.next().await.is_none()); + assert_eq!(Instant::now(), start); +} + +#[tokio::test(start_paused = true)] +async fn repeated_inventory_cursor_reports_failure_after_the_first_catalog() { + let (client, mut peer) = Peer::connect(); + let mut updates = refresh_statusline_providers(client, Vec::new()); + let [first] = peer.expect([Page(None)]).await; + reply_page(first, &[row("plugin", Some(URI))], Some("again")); + assert_uri(updates.next().await, "plugin", URI); + let [second] = peer.expect([Page(Some("again"))]).await; + reply_page( + second, + &[row("plugin", Some("ui://different/footer"))], + Some("again"), + ); + assert!( + matches!(updates.next().await, Some(InventoryFailed(Request(error))) + if error.to_string() == "MCP inventory repeated a cursor") + ); + assert!(updates.next().await.is_none()); +}