Skip to content

fix: lazy extension load #6353

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
michaelneale wants to merge 2 commits into from
Closed

fix: lazy extension load #6353

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
aa07768
we load extensions when agent starrts so don't do it up front
michaelneale Jan 6, 2026
17b393e
load extensions only when used
michaelneale Jan 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 7 additions & 29 deletions crates/goose-server/src/routes/agent.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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::<Vec<_>>();

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))
Expand Down
36 changes: 36 additions & 0 deletions crates/goose/src/agents/agent.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
54 changes: 53 additions & 1 deletion crates/goose/src/agents/extension_manager.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,8 @@ impl Extension {
/// Manages goose extensions / MCP clients and their interactions
pub struct ExtensionManager {
extensions: Mutex<HashMap<String, Extension>>,
/// Extension configs that are registered but not yet loaded (lazy loading)
pending_configs: Mutex<HashMap<String, ExtensionConfig>>,
context: Mutex<PlatformExtensionContext>,
provider: SharedProvider,
}
Expand Down Expand Up @@ -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,
Expand All @@ -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
Copy link

Copilot AI Jan 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment states "Skip if already loaded or already pending" but the code only checks if the extension is already loaded, not if it's already pending. The code works correctly because HashMap::insert overwrites duplicates and add_extension is idempotent, but the comment should be updated to match the implementation.

Suggested change
// Skip if already loaded or already pending
// Skip if the extension is already loaded

Copilot uses AI. Check for mistakes.
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<ExtensionConfig> = 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;
}
Expand All @@ -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<dyn McpClientTrait> = match &config {
Expand Down Expand Up @@ -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<String>,
) -> ExtensionResult<Vec<Tool>> {
// Lazy load any pending extensions before getting tools
self.ensure_pending_extensions_loaded().await;
self.get_prefixed_tools_impl(extension_name, None).await
}

Expand Down
31 changes: 5 additions & 26 deletions ui/desktop/src/App.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -363,15 +363,14 @@ const ExtensionsRoute = () => {

export function AppInner() {
const [fatalError, setFatalError] = useState<string | null>(null);
const [agentWaitingMessage, setAgentWaitingMessage] = useState<string | null>(null);
const [agentWaitingMessage] = useState<string | null>(null);
const [isLoadingSharedSession, setIsLoadingSharedSession] = useState(false);
const [sharedSessionError, setSharedSessionError] = useState<string | null>(null);
const [isExtensionsLoading, setIsExtensionsLoading] = useState(false);
const [isExtensionsLoading] = useState(false);
const [didSelectProvider, setDidSelectProvider] = useState<boolean>(false);

const navigate = useNavigate();
const setView = useNavigation();
const location = useLocation();

const [chat, setChat] = useState<ChatType>({
sessionId: '',
Expand All @@ -384,7 +383,6 @@ export function AppInner() {
const [activeSessionId, setActiveSessionId] = useState<string | null>(null);

const { addExtension } = useConfig();
const { loadCurrentChat } = useAgent();

useEffect(() => {
console.log('Sending reactReady signal to Electron');
Expand All @@ -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[]) => {
Expand Down
Loading