diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 18ec20187c1c..a33c74005ecf 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -678,7 +678,7 @@ impl Agent { } match agent_ref - .add_extension(config_clone, &session_id_clone) + .add_extension_inner(config_clone, &session_id_clone) .await { Ok(_) => ExtensionLoadResult { @@ -700,13 +700,43 @@ impl Agent { }) .collect::>(); - futures::future::join_all(extension_futures).await + let results = futures::future::join_all(extension_futures).await; + + // Persist once after all extensions are loaded + if results.iter().any(|r| r.success) { + if let Err(e) = self.persist_extension_state(&session_id).await { + warn!("Failed to persist extension state after bulk load: {}", e); + } + } + + results } pub async fn add_extension( &self, extension: ExtensionConfig, session_id: &str, + ) -> ExtensionResult<()> { + self.add_extension_inner(extension, session_id).await?; + + // Persist extension state after successful add + self.persist_extension_state(session_id) + .await + .map_err(|e| { + error!("Failed to persist extension state: {}", e); + crate::agents::extension::ExtensionError::SetupError(format!( + "Failed to persist extension state: {}", + e + )) + })?; + + Ok(()) + } + + async fn add_extension_inner( + &self, + extension: ExtensionConfig, + session_id: &str, ) -> ExtensionResult<()> { let session = self .config @@ -760,17 +790,6 @@ impl Agent { } } - // Persist extension state after successful add - self.persist_extension_state(session_id) - .await - .map_err(|e| { - error!("Failed to persist extension state: {}", e); - crate::agents::extension::ExtensionError::SetupError(format!( - "Failed to persist extension state: {}", - e - )) - })?; - Ok(()) } diff --git a/ui/desktop/src/components/bottom_menu/BottomMenuExtensionSelection.tsx b/ui/desktop/src/components/bottom_menu/BottomMenuExtensionSelection.tsx index 468121086f61..221baab1a9ac 100644 --- a/ui/desktop/src/components/bottom_menu/BottomMenuExtensionSelection.tsx +++ b/ui/desktop/src/components/bottom_menu/BottomMenuExtensionSelection.tsx @@ -33,6 +33,11 @@ export const BottomMenuExtensionSelection = ({ sessionId }: BottomMenuExtensionS const { extensionsList: allExtensions } = useConfig(); const isHubView = !sessionId; + useEffect(() => { + setIsSessionExtensionsLoaded(false); + setSessionExtensions([]); + }, [sessionId]); + useEffect(() => { const handleExtensionsLoaded = () => { setRefreshTrigger((prev) => prev + 1); @@ -53,8 +58,11 @@ export const BottomMenuExtensionSelection = ({ sessionId }: BottomMenuExtensionS }; }, []); - // Fetch session-specific extensions or use global defaults useEffect(() => { + if (refreshTrigger === 0 && !isOpen) { + return; + } + const fetchExtensions = async () => { if (!sessionId) { return; @@ -75,7 +83,6 @@ export const BottomMenuExtensionSelection = ({ sessionId }: BottomMenuExtensionS } }; - setIsSessionExtensionsLoaded(false); fetchExtensions(); }, [sessionId, isOpen, refreshTrigger]); diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts index aa0d1497b68a..ef78641c7736 100644 --- a/ui/desktop/src/hooks/useChatStream.ts +++ b/ui/desktop/src/hooks/useChatStream.ts @@ -433,6 +433,7 @@ export function useChatStream({ }, }, }); + window.dispatchEvent(new CustomEvent(AppEvents.SESSION_EXTENSIONS_LOADED)); onSessionLoaded?.(); return; }