fix(agents): persist the catalog the availability probe already fetched

crates/aionui-ai-agent/src/manager/acp/catalog_forwarder.rs

@@ -5,6 +5,7 @@
 //! auth methods) into `AgentHandshake` partials that the registry's
 //! catalog consumer writes to `agent_metadata`.
 
+use agent_client_protocol::schema::v1::{SessionConfigOption, SessionModeState};
 use aionui_api_types::AgentHandshake;
 use aionui_common::normalize_keys_to_snake_case;
 use serde_json::Value;
@@ -91,6 +92,73 @@ fn catalog_partial_from_event(event: &AgentStreamEvent) -> Option<AgentHandshake
     }
 }
 
+/// Project a `session/new` response's catalog straight into an `AgentHandshake`,
+/// bypassing the event stream.
+///
+/// The availability probe opens a real session (that is how it tells "reachable
+/// but unauthorized" apart from other failures) and therefore already holds the
+/// modes/models/config the agent advertises — but it has no `AcpAgentManager`, so
+/// nothing emits the events `catalog_partial_from_event` feeds on. Without this
+/// the catalog is discarded and the picker stays empty until the user opens a
+/// conversation with the agent.
+///
+/// Shapes MUST stay identical to the event path (`catalog_partial_covers_session_fields`
+/// pins mode/model/config; `probe_projection_matches_event_projection` pins that
+/// both routes agree), because both write the same `agent_metadata` columns.
+/// `available_commands` has no counterpart here: it arrives as a session/update
+/// notification, never in the `session/new` response.
+pub fn catalog_partial_from_session_new(
+    modes: Option<&SessionModeState>,
+    models: Option<&super::legacy_session_model::LegacySessionModelState>,
+    config_options: Option<&[SessionConfigOption]>,
+) -> Option<AgentHandshake> {
+    use aionui_api_types::{ModelInfoEntry, ModelInfoPayload};
+
+    let available_modes = modes.and_then(super::agent::sdk_to_snake_value).map(snake_value);
+    let available_models = models
+        .and_then(|models| {
+            let current_id = models.current_model_id.clone();
+            let available: Vec<ModelInfoEntry> = models
+                .available_models
+                .iter()
+                .map(|entry| ModelInfoEntry {
+                    id: entry.model_id.to_string(),
+                    label: entry.name.clone(),
+                })
+                .collect();
+            let current_label = available
+                .iter()
+                .find(|entry| entry.id == current_id)
+                .map(|entry| entry.label.clone())
+                .unwrap_or_else(|| current_id.clone());
+            super::agent::sdk_to_snake_value(&ModelInfoPayload {
+                current_model_id: Some(current_id),
+                current_model_label: Some(current_label),
+                available_models: available,
+            })
+        })
+        .map(snake_value);
+    let config_options = config_options
+        .filter(|options| !options.is_empty())
+        .and_then(|options| super::agent::sdk_to_snake_value(&serde_json::json!({ "config_options": options })))
+        .map(snake_value);
+
+    if available_modes.is_none() && available_models.is_none() && config_options.is_none() {
+        return None;
+    }
+    Some(AgentHandshake {
+        available_modes,
+        available_models,
+        config_options,
+        ..Default::default()
+    })
+}
+
+fn snake_value(mut v: Value) -> Value {
+    normalize_keys_to_snake_case(&mut v);
+    v
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -120,4 +188,54 @@ mod tests {
         // An unrelated event emits no update.
         assert!(catalog_partial_from_event(&AgentStreamEvent::Start(StartEventData { session_id: None })).is_none());
     }
+
+    /// The probe route and the event route write the SAME `agent_metadata`
+    /// columns, so they must produce byte-identical blobs for the same session
+    /// state — otherwise which route ran last would change what the picker reads.
+    /// Feeds one state through both and compares.
+    #[test]
+    fn probe_projection_matches_event_projection() {
+        use crate::manager::acp::legacy_session_model::{LegacyModelEntry, LegacySessionModelState};
+        use agent_client_protocol::schema::v1::{SessionMode, SessionModeState};
+
+        let modes = SessionModeState::new("code", vec![SessionMode::new("code", "Code")]);
+        let models = LegacySessionModelState::new(
+            "gpt-5".to_owned(),
+            vec![LegacyModelEntry {
+                model_id: "gpt-5".to_owned(),
+                name: "GPT-5".to_owned(),
+                description: None,
+            }],
+        );
+
+        let probe = catalog_partial_from_session_new(Some(&modes), Some(&models), None)
+            .expect("a session carrying modes and models projects");
+
+        // Rebuild the events `emit_snapshot_events` would broadcast for the same
+        // state, then run them through the forwarder's own projection.
+        let mode_event = catalog_partial_from_event(&AgentStreamEvent::AcpModeInfo(
+            crate::manager::acp::agent::sdk_to_snake_value(&modes).expect("modes serialize"),
+        ))
+        .expect("mode event projects");
+        assert_eq!(probe.available_modes, mode_event.available_modes, "modes shape");
+
+        let model_payload = aionui_api_types::ModelInfoPayload {
+            current_model_id: Some("gpt-5".to_owned()),
+            current_model_label: Some("GPT-5".to_owned()),
+            available_models: vec![aionui_api_types::ModelInfoEntry {
+                id: "gpt-5".to_owned(),
+                label: "GPT-5".to_owned(),
+            }],
+        };
+        let model_event = catalog_partial_from_event(&AgentStreamEvent::AcpModelInfo(
+            crate::manager::acp::agent::sdk_to_snake_value(&model_payload).expect("models serialize"),
+        ))
+        .expect("model event projects");
+        assert_eq!(probe.available_models, model_event.available_models, "models shape");
+
+        // A session that advertises nothing yields no write at all, so the probe
+        // never blanks a catalog a real conversation had already filled in.
+        assert!(catalog_partial_from_session_new(None, None, None).is_none());
+        assert!(catalog_partial_from_session_new(None, None, Some(&[])).is_none());
+    }
 }

crates/aionui-ai-agent/src/protocol/custom_agent_probe.rs

@@ -15,7 +15,7 @@
 use std::collections::HashMap;
 use std::time::Duration;
 
-use aionui_api_types::TryConnectCustomAgentResponse;
+use aionui_api_types::{AgentHandshake, TryConnectCustomAgentResponse};
 use aionui_common::{CommandSpec, EnvVar};
 use aionui_runtime::{NodeRuntimeProgressReporter, ResolvedCommand, ensure_runtime_command_with_reporter};
 use tokio::sync::{broadcast, mpsc};
@@ -48,29 +48,59 @@ pub async fn try_connect_custom_agent(
     env: &HashMap<String, String>,
     reporter: Option<&dyn NodeRuntimeProgressReporter>,
 ) -> TryConnectCustomAgentResponse {
+    try_connect_custom_agent_with_catalog(command, args, env, reporter)
+        .await
+        .0
+}
+
+/// As [`try_connect_custom_agent`], plus the catalog the probe's `session/new`
+/// advertised. The probe already pays for a full spawn + `initialize` +
+/// `session/new`; callers that persist agent metadata use this variant so that
+/// data is stored instead of thrown away. The partial is `Some` only on a
+/// successful session whose agent advertised at least one of modes / models /
+/// config options — an auth-gated or failing probe yields `None`, never an
+/// empty write that would blank an existing catalog.
+pub async fn try_connect_custom_agent_with_catalog(
+    command: &str,
+    args: &[String],
+    env: &HashMap<String, String>,
+    reporter: Option<&dyn NodeRuntimeProgressReporter>,
+) -> (TryConnectCustomAgentResponse, Option<Box<AgentHandshake>>) {
     // ── Step 1 — which check ────────────────────────────────────────
     let head = first_token(command);
     let resolved = match ensure_runtime_command_with_reporter(head, reporter).await {
         Ok(resolved) => resolved,
         Err(error) => {
-            return TryConnectCustomAgentResponse::FailCli {
-                error: error.to_string(),
-            };
+            return (
+                TryConnectCustomAgentResponse::FailCli {
+                    error: error.to_string(),
+                },
+                None,
+            );
         }
     };
     debug!(program = %resolved.program.display(), "probe step 1 ok");
 
     // ── Step 2 — spawn + ACP initialize ─────────────────────────────
     let proc = match spawn_probe_process(resolved, args, env).await {
         Ok(proc) => proc,
-        Err(msg) => return TryConnectCustomAgentResponse::FailAcp { error: msg },
+        Err(msg) => return (TryConnectCustomAgentResponse::FailAcp { error: msg }, None),
     };
 
     let outcome = match tokio::time::timeout(STEP2_TIMEOUT, run_handshake(&proc)).await {
-        Ok(outcome) => outcome.into_response(),
-        Err(_) => TryConnectCustomAgentResponse::FailAcp {
-            error: format!("ACP handshake did not complete within {}s", STEP2_TIMEOUT.as_secs()),
-        },
+        Ok(outcome) => {
+            let catalog = match &outcome {
+                ProbeOutcome::Ok(catalog) => catalog.clone(),
+                _ => None,
+            };
+            (outcome.into_response(), catalog)
+        }
+        Err(_) => (
+            TryConnectCustomAgentResponse::FailAcp {
+                error: format!("ACP handshake did not complete within {}s", STEP2_TIMEOUT.as_secs()),
+            },
+            None,
+        ),
     };
 
     // Always tear down the whole process group. `kill_on_drop(true)` only
@@ -131,15 +161,19 @@ async fn spawn_probe_process(
 /// alone returns `authMethods` even for already-authorized agents and cannot
 /// make this distinction.
 enum ProbeOutcome {
-    Ok,
+    /// Carries the catalog the successful `session/new` advertised (modes /
+    /// models / config options), already projected into the same shape the live
+    /// session path persists. `None` when the agent advertised nothing. Boxed to
+    /// keep the enum small: the success payload dwarfs the error strings.
+    Ok(Option<Box<AgentHandshake>>),
     Auth(String),
     Fail(String),
 }
 
 impl ProbeOutcome {
     fn into_response(self) -> TryConnectCustomAgentResponse {
         match self {
-            ProbeOutcome::Ok => TryConnectCustomAgentResponse::Success,
+            ProbeOutcome::Ok(_) => TryConnectCustomAgentResponse::Success,
             ProbeOutcome::Auth(error) => TryConnectCustomAgentResponse::FailAuth { error },
             ProbeOutcome::Fail(error) => TryConnectCustomAgentResponse::FailAcp { error },
         }
@@ -187,8 +221,26 @@ async fn run_handshake(proc: &CliAgentProcess) -> ProbeOutcome {
     // `initialize` only proves the agent speaks ACP, not that it is usable.
     // Open a real session (no prompt) so an auth-gated agent surfaces its
     // `auth_required` error here instead of silently appearing "online".
+    // Keep the response: it carries the agent's advertised modes / models /
+    // config options, which the caller persists into `agent_metadata` so the
+    // picker is populated before the user ever opens a conversation. Discarding
+    // it meant a probed-online agent still showed an empty picker.
     let outcome = match protocol.new_session(NewSessionRequest::new(std::env::temp_dir())).await {
-        Ok(_) => ProbeOutcome::Ok,
+        Ok((response, legacy_models)) => {
+            // Same extraction the live session path performs (`agent_session_flow`):
+            // models ride beside the response because the SDK dropped the field.
+            let models = legacy_models
+                .as_ref()
+                .and_then(crate::manager::acp::legacy_session_model::LegacySessionModelState::from_state_value);
+            ProbeOutcome::Ok(
+                crate::manager::acp::catalog_forwarder::catalog_partial_from_session_new(
+                    response.modes.as_ref(),
+                    models.as_ref(),
+                    response.config_options.as_deref(),
+                )
+                .map(Box::new),
+            )
+        }
         Err(AcpError::AuthRequired) => {
             ProbeOutcome::Auth("Agent reachable but requires login/authorization".to_string())
         }

crates/aionui-ai-agent/src/services/availability/mod.rs

@@ -166,7 +166,7 @@ impl AgentAvailabilityService {
 }
 
 async fn run_probe(
-    _registry: &Arc<AgentRegistry>,
+    registry: &Arc<AgentRegistry>,
     provider_repo: &Arc<dyn IProviderRepository>,
     meta: &AgentMetadata,
     user_id: &str,
@@ -202,26 +202,43 @@ async fn run_probe(
                 Some("package_lock_invalid".to_owned()),
                 Some(error),
             ),
-            Ok(args) => match custom_agent_probe::try_connect_custom_agent(command, &args, &env, None).await {
-                TryConnectCustomAgentResponse::Success => (AgentSnapshotCheckStatus::Online, None, None),
-                TryConnectCustomAgentResponse::FailCli { error } => (
-                    AgentSnapshotCheckStatus::Offline,
-                    Some("command_not_found".to_owned()),
-                    Some(error),
-                ),
-                TryConnectCustomAgentResponse::FailAcp { error } => (
-                    AgentSnapshotCheckStatus::Offline,
-                    Some("acp_init_failed".to_owned()),
-                    Some(error),
-                ),
-                // Reachable but not authorized: still offline (unusable), but a
-                // dedicated code lets the UI guide the user to log in.
-                TryConnectCustomAgentResponse::FailAuth { error } => (
-                    AgentSnapshotCheckStatus::Offline,
-                    Some("auth_required".to_owned()),
-                    Some(error),
-                ),
-            },
+            Ok(args) => {
+                match custom_agent_probe::try_connect_custom_agent_with_catalog(command, &args, &env, None).await {
+                    // The probe opened a real session to reach this verdict, so its
+                    // `session/new` already carried whatever modes / models / config
+                    // options the agent advertises. Persist them through the same
+                    // channel a live conversation uses, so the pickers are populated
+                    // before the user ever opens one. Best-effort and additive:
+                    // `apply_handshake` skips `None` fields, so this never blanks a
+                    // catalog a real session had filled in, and an agent that
+                    // advertises nothing sends nothing.
+                    (TryConnectCustomAgentResponse::Success, catalog) => {
+                        if let Some(partial) = catalog {
+                            registry
+                                .catalog_sender()
+                                .send_partial(user_id.to_owned(), meta.id.clone(), *partial);
+                        }
+                        (AgentSnapshotCheckStatus::Online, None, None)
+                    }
+                    (TryConnectCustomAgentResponse::FailCli { error }, _) => (
+                        AgentSnapshotCheckStatus::Offline,
+                        Some("command_not_found".to_owned()),
+                        Some(error),
+                    ),
+                    (TryConnectCustomAgentResponse::FailAcp { error }, _) => (
+                        AgentSnapshotCheckStatus::Offline,
+                        Some("acp_init_failed".to_owned()),
+                        Some(error),
+                    ),
+                    // Reachable but not authorized: still offline (unusable), but a
+                    // dedicated code lets the UI guide the user to log in.
+                    (TryConnectCustomAgentResponse::FailAuth { error }, _) => (
+                        AgentSnapshotCheckStatus::Offline,
+                        Some("auth_required".to_owned()),
+                        Some(error),
+                    ),
+                }
+            }
         }
     } else if meta.backend.is_some() {
         // Commandless builtin fallback: same PATH + `--version` treatment as