fix(gateway): reconnect after config save and hold workspace pin across initialize

Evict-only left enabled servers bound-but-offline after a definition write,
and stdio OAuth reconnect still built an HTTP transport. Hold a non-empty
X-Mcpmux-Workspace until mcp-session-id exists so multi-root PendingRoots
does not win after a gateway reload.

Signed-off-by: crimsonsunset <jsangio1@gmail.com>

Author: crimsonsunset

crates/mcpmux-gateway/src/consumers/server_config_handler.rs

@@ -1,31 +1,38 @@
-//! Server Config Event Handler - Evicts pool instances on config changes
+//! Server Config Event Handler - Reconnects pool instances on config changes
 //!
-//! Listens for `ServerConfigUpdated` and removes the pooled instance for
-//! enabled servers so the next connect rebuilds transport from DB.
+//! Listens for `ServerConfigUpdated` and runs `reconnect_fresh` for
+//! enabled servers so the next call uses transport rebuilt from DB.
 
 use mcpmux_core::{DomainEvent, InstalledServerRepository};
+use std::path::PathBuf;
 use std::sync::Arc;
+use std::time::Instant;
 use tokio::sync::broadcast;
 use tracing::{debug, info, warn};
 use uuid::Uuid;
 
-use crate::pool::PoolService;
+use crate::pool::transport::resolution::resolve_auto_connection_context;
+use crate::pool::{ConnectionResult, PoolService};
 
-/// Evicts pooled server instances when configuration changes.
+/// Evicts and reconnects pooled server instances when configuration changes.
 pub struct ServerConfigUpdatedHandler {
     installed_server_repo: Arc<dyn InstalledServerRepository + Send + Sync>,
     pool_service: Arc<PoolService>,
+    state_dir: Option<PathBuf>,
 }
 
 impl ServerConfigUpdatedHandler {
-    /// Create a handler wired to the installed-server repo and connection pool.
+    /// Create a handler wired to the installed-server repo, connection pool,
+    /// and optional MCP state dir used when re-resolving transport.
     pub fn new(
         installed_server_repo: Arc<dyn InstalledServerRepository + Send + Sync>,
         pool_service: Arc<PoolService>,
+        state_dir: Option<PathBuf>,
     ) -> Self {
         Self {
             installed_server_repo,
             pool_service,
+            state_dir,
         }
     }
 
@@ -53,7 +60,7 @@ impl ServerConfigUpdatedHandler {
         });
     }
 
-    /// Handle one domain event, evicting the pool instance when applicable.
+    /// Handle one domain event, reconnecting the pool instance when applicable.
     async fn handle_event(&self, event: DomainEvent) -> anyhow::Result<()> {
         let DomainEvent::ServerConfigUpdated {
             space_id,
@@ -66,7 +73,10 @@ impl ServerConfigUpdatedHandler {
         self.handle_config_updated(space_id, &server_id).await
     }
 
-    /// Remove a pooled instance for an enabled server after a config/definition write.
+    /// Reconnect an enabled server from current DB config after a write.
+    ///
+    /// Resolve failure falls back to evict-only so a stale instance cannot
+    /// keep serving the pre-save transport.
     async fn handle_config_updated(&self, space_id: Uuid, server_id: &str) -> anyhow::Result<()> {
         let space_id_str = space_id.to_string();
         let Some(installed) = self
@@ -87,10 +97,251 @@ impl ServerConfigUpdatedHandler {
             return Ok(());
         }
 
-        info!(
-            "[ServerConfigHandler] Evicting pool instance for enabled server {space_id}/{server_id} after config update"
-        );
-        self.pool_service.remove_instance(space_id, server_id);
+        let started = Instant::now();
+        match resolve_auto_connection_context(
+            self.installed_server_repo.as_ref(),
+            self.state_dir.as_deref(),
+            space_id,
+            server_id,
+        )
+        .await
+        {
+            Ok(ctx) => {
+                let result = self.pool_service.reconnect_fresh(&ctx).await;
+                info!(
+                    server_id = %server_id,
+                    space_id = %space_id,
+                    ok = result.is_connected(),
+                    duration_ms = started.elapsed().as_millis(),
+                    "[ServerConfigHandler] reconnect_fresh after config update"
+                );
+                if let ConnectionResult::Failed { error } = result {
+                    warn!(
+                        server_id = %server_id,
+                        space_id = %space_id,
+                        error = %error,
+                        "[ServerConfigHandler] reconnect_fresh failed after config update"
+                    );
+                }
+            }
+            Err(error) => {
+                warn!(
+                    server_id = %server_id,
+                    space_id = %space_id,
+                    error = %error,
+                    "[ServerConfigHandler] re-resolve failed, evicting only"
+                );
+                self.pool_service.remove_instance(space_id, server_id);
+            }
+        }
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::pool::PoolService;
+    use async_trait::async_trait;
+    use chrono::Utc;
+    use mcpmux_core::{InstalledServer, ServerDefinition, TransportConfig};
+    use std::collections::HashMap;
+    use std::sync::Mutex;
+
+    struct MapInstalledRepo {
+        servers: Mutex<HashMap<(String, String), InstalledServer>>,
+    }
+
+    impl MapInstalledRepo {
+        /// Store one installed-server row keyed by `(space_id, server_id)`.
+        fn with(server: InstalledServer) -> Arc<Self> {
+            let mut servers = HashMap::new();
+            servers.insert((server.space_id.clone(), server.server_id.clone()), server);
+            Arc::new(Self {
+                servers: Mutex::new(servers),
+            })
+        }
+    }
+
+    #[async_trait]
+    impl InstalledServerRepository for MapInstalledRepo {
+        async fn list(&self) -> mcpmux_core::repository::RepoResult<Vec<InstalledServer>> {
+            Ok(self.servers.lock().unwrap().values().cloned().collect())
+        }
+        async fn list_for_space(
+            &self,
+            _space_id: &str,
+        ) -> mcpmux_core::repository::RepoResult<Vec<InstalledServer>> {
+            Ok(vec![])
+        }
+        async fn list_by_source_file(
+            &self,
+            _file_path: &std::path::Path,
+        ) -> mcpmux_core::repository::RepoResult<Vec<InstalledServer>> {
+            Ok(vec![])
+        }
+        async fn get(
+            &self,
+            _id: &Uuid,
+        ) -> mcpmux_core::repository::RepoResult<Option<InstalledServer>> {
+            Ok(None)
+        }
+        async fn get_by_server_id(
+            &self,
+            space_id: &str,
+            server_id: &str,
+        ) -> mcpmux_core::repository::RepoResult<Option<InstalledServer>> {
+            Ok(self
+                .servers
+                .lock()
+                .unwrap()
+                .get(&(space_id.to_string(), server_id.to_string()))
+                .cloned())
+        }
+        async fn install(
+            &self,
+            _server: &InstalledServer,
+        ) -> mcpmux_core::repository::RepoResult<()> {
+            Ok(())
+        }
+        async fn update(
+            &self,
+            _server: &InstalledServer,
+        ) -> mcpmux_core::repository::RepoResult<()> {
+            Ok(())
+        }
+        async fn uninstall(&self, _id: &Uuid) -> mcpmux_core::repository::RepoResult<()> {
+            Ok(())
+        }
+        async fn list_enabled(
+            &self,
+            _space_id: &str,
+        ) -> mcpmux_core::repository::RepoResult<Vec<InstalledServer>> {
+            Ok(vec![])
+        }
+        async fn list_enabled_all(
+            &self,
+        ) -> mcpmux_core::repository::RepoResult<Vec<InstalledServer>> {
+            Ok(vec![])
+        }
+        async fn set_enabled(
+            &self,
+            _id: &Uuid,
+            _enabled: bool,
+        ) -> mcpmux_core::repository::RepoResult<()> {
+            Ok(())
+        }
+        async fn set_oauth_connected(
+            &self,
+            _id: &Uuid,
+            _connected: bool,
+        ) -> mcpmux_core::repository::RepoResult<()> {
+            Ok(())
+        }
+        async fn update_inputs(
+            &self,
+            _id: &Uuid,
+            _input_values: HashMap<String, String>,
+        ) -> mcpmux_core::repository::RepoResult<()> {
+            Ok(())
+        }
+        async fn update_cached_definition(
+            &self,
+            _id: &Uuid,
+            _server_name: Option<String>,
+            _cached_definition: Option<String>,
+        ) -> mcpmux_core::repository::RepoResult<()> {
+            Ok(())
+        }
+        async fn set_display_name_override(
+            &self,
+            _id: &Uuid,
+            _value: Option<String>,
+        ) -> mcpmux_core::repository::RepoResult<()> {
+            Ok(())
+        }
+        async fn update_version_cache(
+            &self,
+            _id: &Uuid,
+            _latest_available_version: Option<String>,
+            _current_version: Option<String>,
+            _version_checked_at: chrono::DateTime<Utc>,
+        ) -> mcpmux_core::repository::RepoResult<()> {
+            Ok(())
+        }
+    }
+
+    /// Minimal stdio definition whose command does not exist on disk.
+    fn stdio_definition(server_id: &str) -> ServerDefinition {
+        ServerDefinition {
+            id: server_id.to_string(),
+            name: server_id.to_string(),
+            description: None,
+            alias: None,
+            auth: None,
+            icon: None,
+            transport: TransportConfig::Stdio {
+                command: "/nonexistent/mcpmux-config-handler-test".into(),
+                args: vec![],
+                env: HashMap::new(),
+                metadata: Default::default(),
+            },
+            categories: vec![],
+            publisher: None,
+            source: Default::default(),
+            badges: vec![],
+            hosting_type: Default::default(),
+            license: None,
+            license_url: None,
+            installation: None,
+            capabilities: None,
+            sponsored: None,
+            media: None,
+            changelog_url: None,
+        }
+    }
+
+    #[tokio::test]
+    async fn enabled_server_reconnects_and_evicts_stale_instance() {
+        let space_id = Uuid::new_v4();
+        let server_id = "cfg-enabled";
+        let installed = InstalledServer::new(space_id.to_string(), server_id)
+            .with_definition(&stdio_definition(server_id))
+            .with_enabled(true);
+        let repo = MapInstalledRepo::with(installed);
+        let pool = Arc::new(PoolService::new_test_with_repo(repo.clone()));
+        pool.insert_test_instance(space_id, server_id);
+        assert_eq!(pool.stats().total_instances, 1);
+
+        let handler = ServerConfigUpdatedHandler::new(repo, pool.clone(), None);
+        handler
+            .handle_config_updated(space_id, server_id)
+            .await
+            .expect("handler");
+
+        assert!(
+            pool.get_instance(space_id, server_id).is_none(),
+            "stale instance must be gone after reconnect_fresh"
+        );
+    }
+
+    #[tokio::test]
+    async fn disabled_server_is_left_in_the_pool() {
+        let space_id = Uuid::new_v4();
+        let server_id = "cfg-disabled";
+        let installed = InstalledServer::new(space_id.to_string(), server_id)
+            .with_definition(&stdio_definition(server_id))
+            .with_enabled(false);
+        let repo = MapInstalledRepo::with(installed);
+        let pool = Arc::new(PoolService::new_test_with_repo(repo.clone()));
+        pool.insert_test_instance(space_id, server_id);
+
+        let handler = ServerConfigUpdatedHandler::new(repo, pool.clone(), None);
+        handler
+            .handle_config_updated(space_id, server_id)
+            .await
+            .expect("handler");
+
+        assert!(pool.get_instance(space_id, server_id).is_some());
+    }
+}

crates/mcpmux-gateway/src/mcp/oauth_middleware.rs

@@ -271,11 +271,19 @@ pub async fn mcp_oauth_middleware(
             services.session_roots.set_pinned(sid, ws);
         }
         (None, Some(ws)) => {
-            warn!(
+            info!(
                 trace_id = %trace_id,
                 workspace_header = %ws,
-                "[SessionRoots] X-Mcpmux-Workspace present without mcp-session-id — pin skipped",
+                "[SessionRoots] X-Mcpmux-Workspace held until mcp-session-id exists",
             );
+            services
+                .session_roots
+                .remember_pending_workspace(&client_id, ws);
+        }
+        (Some(sid), None) => {
+            services
+                .session_roots
+                .apply_pending_workspace(&client_id, sid);
         }
         _ => {}
     }
@@ -327,6 +335,20 @@ pub async fn mcp_oauth_middleware(
 
     let response = next.run(request).await;
 
+    if let Some(ws) = workspace_header
+        .as_deref()
+        .filter(|value| !value.trim().is_empty())
+    {
+        if let Some(sid) = response
+            .headers()
+            .get("mcp-session-id")
+            .and_then(|value| value.to_str().ok())
+            .filter(|value| !value.is_empty())
+        {
+            services.session_roots.set_pinned(sid, ws);
+        }
+    }
+
     // Log errors only — except two rmcp spec-correct shapes that are not
     // gateway problems: a GET without Mcp-Session-Id (client opening the SSE
     // stream before initialize) and any request against a session rmcp has