From c91c63f81001363c74ff102449546e24e486bc0f Mon Sep 17 00:00:00 2001
From: Coffee <zjchen1234@foxmail.com>
Date: Thu, 19 Mar 2026 15:24:10 +0800
Subject: [PATCH] Activate NEAR AI companion MCP when auth becomes available

---
 src/agent/dispatcher.rs   |  9 ++++++
 src/extensions/manager.rs | 63 +++++++++++++++++++++++++++++++++++++--
 2 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs
index 9be0d654..0fb20a2b 100644
--- a/src/agent/dispatcher.rs
+++ b/src/agent/dispatcher.rs
@@ -46,6 +46,15 @@ impl Agent {
         thread_id: Uuid,
         initial_messages: Vec<ChatMessage>,
     ) -> Result<AgenticLoopResult, Error> {
+        if let Some(ext_mgr) = self.deps.extension_manager.as_ref()
+            && let Err(e) = ext_mgr.ensure_nearai_companion_active_if_ready().await
+        {
+            tracing::debug!(
+                "Failed to auto-activate NEAR AI companion MCP before turn: {}",
+                e
+            );
+        }
+
         // Detect group chat from channel metadata (needed before loading system prompt)
         let is_group_chat = message
             .metadata
diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs
index 823aa8b3..157aff16 100644
--- a/src/extensions/manager.rs
+++ b/src/extensions/manager.rs
@@ -1009,6 +1009,32 @@ impl ExtensionManager {
         }
     }
 
+    /// Activate the derived NEAR AI companion MCP server if auth is already
+    /// available and the companion is not active yet.
+    ///
+    /// Returns `Ok(true)` only when this call performed an activation.
+    pub async fn ensure_nearai_companion_active_if_ready(&self) -> Result<bool, ExtensionError> {
+        let Some(companion) = self.companion_mcp_server.as_ref() else {
+            return Ok(false);
+        };
+
+        let companion_name = companion.name.clone();
+
+        {
+            let clients = self.mcp_clients.read().await;
+            if clients.contains_key(&companion_name) {
+                return Ok(false);
+            }
+        }
+
+        if !self.is_runtime_authenticated(companion).await {
+            return Ok(false);
+        }
+
+        self.activate(&companion_name).await?;
+        Ok(true)
+    }
+
     /// List extensions with their status.
     ///
     /// When `include_available` is `true`, registry entries that are not yet
@@ -5556,6 +5582,7 @@ mod tests {
         wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
         tools_dir: std::path::PathBuf,
         channels_dir: std::path::PathBuf,
+        companion_mcp_server: Option<crate::tools::mcp::config::McpServerConfig>,
     ) -> crate::extensions::manager::ExtensionManager {
         use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
         use crate::tools::mcp::process::McpProcessManager;
@@ -5585,7 +5612,7 @@ mod tests {
             None, // tunnel_url
             "test".to_string(),
             None, // db
-            None, // companion MCP
+            companion_mcp_server,
             vec![],
         )
     }
@@ -5594,7 +5621,7 @@ mod tests {
         wasm_runtime: Option<Arc<crate::tools::wasm::WasmToolRuntime>>,
         tools_dir: std::path::PathBuf,
     ) -> crate::extensions::manager::ExtensionManager {
-        make_test_manager_with_dirs(wasm_runtime, tools_dir.clone(), tools_dir)
+        make_test_manager_with_dirs(wasm_runtime, tools_dir.clone(), tools_dir, None)
     }
 
     #[tokio::test]
@@ -5664,6 +5691,36 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_ensure_nearai_companion_active_if_ready_skips_without_auth() {
+        let dir = tempfile::tempdir().expect("temp dir");
+        let companion = crate::tools::mcp::config::McpServerConfig::new(
+            crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME,
+            "https://private.near.ai/mcp",
+        )
+        .with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi);
+        let manager = make_test_manager_with_dirs(
+            None,
+            dir.path().join("tools"),
+            dir.path().join("channels"),
+            Some(companion),
+        );
+
+        let activated = manager
+            .ensure_nearai_companion_active_if_ready()
+            .await
+            .expect("helper should not fail when auth is missing");
+
+        assert!(!activated, "companion should not activate without auth");
+        assert!(
+            !manager
+                .mcp_clients
+                .read()
+                .await
+                .contains_key(crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME)
+        );
+    }
+
     #[test]
     fn test_capabilities_files_also_separate() {
         // capabilities.json files for tools and channels should also be separate.
@@ -6638,7 +6695,7 @@ mod tests {
         let dir = tempfile::tempdir().expect("temp dir");
         let tools_dir = dir.path().join("tools");
         let channels_dir = dir.path().join("channels");
-        let mgr = make_test_manager_with_dirs(None, tools_dir, channels_dir.clone());
+        let mgr = make_test_manager_with_dirs(None, tools_dir, channels_dir.clone(), None);
 
         let wasm_path = channels_dir.join("telegram.wasm");
         let cap_path = channels_dir.join("telegram.capabilities.json");