diff --git a/crates/goose-server/src/routes/agent.rs b/crates/goose-server/src/routes/agent.rs index 70b871fa7048..f3e56a5d8322 100644 --- a/crates/goose-server/src/routes/agent.rs +++ b/crates/goose-server/src/routes/agent.rs @@ -34,7 +34,7 @@ use std::path::PathBuf; use std::sync::atomic::Ordering; use std::sync::Arc; use tokio_util::sync::CancellationToken; -use tracing::{error, warn}; +use tracing::error; #[derive(Deserialize, utoipa::ToSchema)] pub struct UpdateFromSessionRequest { @@ -298,35 +298,13 @@ async fn resume_agent( }) }; - let extensions_result = async { - let enabled_configs = goose::config::get_enabled_extensions(); - let agent_clone = agent.clone(); - - let extension_futures = enabled_configs - .into_iter() - .map(|config| { - let config_clone = config.clone(); - let agent_ref = agent_clone.clone(); - - async move { - if let Err(e) = agent_ref.add_extension(config_clone.clone()).await { - warn!("Failed to load extension {}: {}", config_clone.name(), e); - goose::posthog::emit_error( - "extension_load_failed", - &format!("{}: {}", config_clone.name(), e), - ); - } - Ok::<_, ErrorResponse>(()) - } - }) - .collect::>(); - - futures::future::join_all(extension_futures).await; - Ok::<(), ErrorResponse>(()) // Fixed type annotation - }; + // Register extensions for lazy loading - they will be loaded on first tool request + let enabled_configs = goose::config::get_enabled_extensions(); + for config in enabled_configs { + agent.register_extension(config).await; + } - let (provider_result, _) = tokio::join!(provider_result, extensions_result); - provider_result?; + provider_result.await?; } Ok(Json(session)) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 1fb2165aca50..26b5ae03ae64 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -602,6 +602,42 @@ impl Agent { Ok(()) } + /// Register an extension for lazy loading. The extension will not be loaded + /// until tools are first requested (e.g., when the user sends a message). + /// Frontend extensions are handled immediately since they don't require MCP connections. + pub async fn register_extension(&self, extension: ExtensionConfig) { + match &extension { + ExtensionConfig::Frontend { + tools, + instructions, + .. + } => { + // For frontend tools, handle immediately since they don't require MCP connections + let mut frontend_tools = self.frontend_tools.lock().await; + for tool in tools { + let frontend_tool = FrontendTool { + name: tool.name.to_string(), + tool: tool.clone(), + }; + frontend_tools.insert(tool.name.to_string(), frontend_tool); + } + let mut frontend_instructions = self.frontend_instructions.lock().await; + if let Some(instructions) = instructions { + *frontend_instructions = Some(instructions.clone()); + } else { + *frontend_instructions = Some( + "The following tools are provided directly by the frontend and will be executed by the frontend when called.".to_string(), + ); + } + } + _ => { + self.extension_manager + .register_extension(extension.clone()) + .await; + } + } + } + pub async fn subagents_enabled(&self) -> bool { let config = crate::config::Config::global(); let is_autonomous = config.get_goose_mode().unwrap_or(GooseMode::Auto) == GooseMode::Auto; diff --git a/crates/goose/src/agents/extension_manager.rs b/crates/goose/src/agents/extension_manager.rs index 206b873c89cb..79b536d995bf 100644 --- a/crates/goose/src/agents/extension_manager.rs +++ b/crates/goose/src/agents/extension_manager.rs @@ -94,6 +94,8 @@ impl Extension { /// Manages goose extensions / MCP clients and their interactions pub struct ExtensionManager { extensions: Mutex>, + /// Extension configs that are registered but not yet loaded (lazy loading) + pending_configs: Mutex>, context: Mutex, provider: SharedProvider, } @@ -445,6 +447,7 @@ impl ExtensionManager { pub fn new(provider: SharedProvider) -> Self { Self { extensions: Mutex::new(HashMap::new()), + pending_configs: Mutex::new(HashMap::new()), context: Mutex::new(PlatformExtensionContext { session_id: None, extension_manager: None, @@ -458,6 +461,47 @@ impl ExtensionManager { Self::new(Arc::new(Mutex::new(None))) } + /// Register an extension config for lazy loading. The extension will not be loaded + /// until `ensure_pending_extensions_loaded` is called or when tools are first requested. + pub async fn register_extension(&self, config: ExtensionConfig) { + let config_name = config.key().to_string(); + let sanitized_name = normalize(config_name); + + // Skip if already loaded or already pending + if self.extensions.lock().await.contains_key(&sanitized_name) { + return; + } + + self.pending_configs + .lock() + .await + .insert(sanitized_name, config); + } + + /// Load all pending extension configs. This is called automatically when tools are requested. + pub async fn ensure_pending_extensions_loaded(&self) { + // Take all pending configs + let pending: Vec = self + .pending_configs + .lock() + .await + .drain() + .map(|(_, v)| v) + .collect(); + + // Load each one + for config in pending { + if let Err(e) = self.add_extension(config.clone()).await { + warn!("Failed to load extension {}: {}", config.name(), e); + } + } + } + + /// Check if there are pending extensions that haven't been loaded yet + pub async fn has_pending_extensions(&self) -> bool { + !self.pending_configs.lock().await.is_empty() + } + pub async fn set_context(&self, context: PlatformExtensionContext) { *self.context.lock().await = context; } @@ -477,6 +521,11 @@ impl ExtensionManager { pub async fn add_extension(&self, config: ExtensionConfig) -> ExtensionResult<()> { let config_name = config.key().to_string(); let sanitized_name = normalize(config_name.clone()); + + if self.extensions.lock().await.contains_key(&sanitized_name) { + return Ok(()); + } + let mut temp_dir = None; let client: Box = match &config { @@ -659,11 +708,14 @@ impl ExtensionManager { .collect() } - /// Get all tools from all clients with proper prefixing + /// Get all tools from all clients with proper prefixing. + /// This will lazily load any pending extensions before returning tools. pub async fn get_prefixed_tools( &self, extension_name: Option, ) -> ExtensionResult> { + // Lazy load any pending extensions before getting tools + self.ensure_pending_extensions_loaded().await; self.get_prefixed_tools_impl(extension_name, None).await } diff --git a/ui/desktop/src/App.tsx b/ui/desktop/src/App.tsx index cffd8c8625ce..023935de739f 100644 --- a/ui/desktop/src/App.tsx +++ b/ui/desktop/src/App.tsx @@ -39,7 +39,7 @@ import PermissionSettingsView from './components/settings/permission/PermissionS import ExtensionsView, { ExtensionsViewOptions } from './components/extensions/ExtensionsView'; import RecipesView from './components/recipes/RecipesView'; import { View, ViewOptions } from './utils/navigationUtils'; -import { NoProviderOrModelError, useAgent } from './hooks/useAgent'; + import { useNavigation } from './hooks/useNavigation'; import { errorMessage } from './utils/conversionUtils'; import { usePageViewTracking } from './hooks/useAnalytics'; @@ -363,15 +363,14 @@ const ExtensionsRoute = () => { export function AppInner() { const [fatalError, setFatalError] = useState(null); - const [agentWaitingMessage, setAgentWaitingMessage] = useState(null); + const [agentWaitingMessage] = useState(null); const [isLoadingSharedSession, setIsLoadingSharedSession] = useState(false); const [sharedSessionError, setSharedSessionError] = useState(null); - const [isExtensionsLoading, setIsExtensionsLoading] = useState(false); + const [isExtensionsLoading] = useState(false); const [didSelectProvider, setDidSelectProvider] = useState(false); const navigate = useNavigate(); const setView = useNavigation(); - const location = useLocation(); const [chat, setChat] = useState({ sessionId: '', @@ -384,7 +383,6 @@ export function AppInner() { const [activeSessionId, setActiveSessionId] = useState(null); const { addExtension } = useConfig(); - const { loadCurrentChat } = useAgent(); useEffect(() => { console.log('Sending reactReady signal to Electron'); @@ -398,27 +396,8 @@ export function AppInner() { } }, []); - // Handle URL parameters and deeplinks on app startup - const loadingHub = location.pathname === '/'; - useEffect(() => { - if (loadingHub) { - (async () => { - try { - const loadedChat = await loadCurrentChat({ - setAgentWaitingMessage, - setIsExtensionsLoading, - }); - setChat(loadedChat); - } catch (e) { - if (e instanceof NoProviderOrModelError) { - // the onboarding flow will trigger - } else { - throw e; - } - } - })(); - } - }, [loadCurrentChat, setAgentWaitingMessage, navigate, loadingHub, setChat]); + // Don't pre-load session/extensions on Hub - wait until user actually starts chatting + // This avoids the slow extension loading on app startup useEffect(() => { const handleOpenSharedSession = async (_event: IpcRendererEvent, ...args: unknown[]) => {