diff --git a/crates/goose/src/providers/gemini_oauth.rs b/crates/goose/src/providers/gemini_oauth.rs index a872356d6104..ce2c11ae9db7 100644 --- a/crates/goose/src/providers/gemini_oauth.rs +++ b/crates/goose/src/providers/gemini_oauth.rs @@ -144,7 +144,7 @@ struct SetupData { } #[derive(Debug, Clone)] -struct TokenCache { +pub(crate) struct TokenCache { cache_path: PathBuf, } @@ -153,7 +153,7 @@ fn get_cache_path() -> PathBuf { } impl TokenCache { - fn new() -> Self { + pub(crate) fn new() -> Self { let cache_path = get_cache_path(); if let Some(parent) = cache_path.parent() { let _ = std::fs::create_dir_all(parent); @@ -167,6 +167,10 @@ impl TokenCache { .and_then(|contents| serde_json::from_str(&contents).ok()) } + pub(crate) fn has_token(&self) -> bool { + self.load().is_some() + } + fn save(&self, data: &SetupData) -> Result<()> { if let Some(parent) = self.cache_path.parent() { std::fs::create_dir_all(parent)?; @@ -1096,8 +1100,15 @@ mod tests { } #[test] + #[serial_test::serial] fn test_token_cache_roundtrip() { + let root = tempfile::tempdir().unwrap(); + let root_path = root.path().to_string_lossy().to_string(); + let _guard = env_lock::lock_env([("GOOSE_PATH_ROOT", Some(root_path.as_str()))]); + let cache = TokenCache::new(); + cache.clear(); + assert!(!cache.has_token()); let data = SetupData { project_id: "test-project".to_string(), token: TokenData { @@ -1111,7 +1122,9 @@ mod tests { assert_eq!(loaded.project_id, "test-project"); assert_eq!(loaded.token.access_token, "test-access"); assert_eq!(loaded.token.refresh_token, "test-refresh"); + assert!(cache.has_token()); cache.clear(); assert!(cache.load().is_none()); + assert!(!cache.has_token()); } } diff --git a/crates/goose/src/providers/init.rs b/crates/goose/src/providers/init.rs index 2c59ea209c3f..459d78031411 100644 --- a/crates/goose/src/providers/init.rs +++ b/crates/goose/src/providers/init.rs @@ -102,7 +102,10 @@ async fn init_registry() -> RwLock { ); registry.register::(false); registry.register::(false); - registry.register::(true); + registry.register_with_inventory::( + true, + Some(registrations::gemini_oauth_inventory()), + ); registry.register::(false); registry.register_with_inventory::( true, diff --git a/crates/goose/src/providers/inventory/registrations.rs b/crates/goose/src/providers/inventory/registrations.rs index b3852cf5062b..2c3ef5488e67 100644 --- a/crates/goose/src/providers/inventory/registrations.rs +++ b/crates/goose/src/providers/inventory/registrations.rs @@ -11,6 +11,7 @@ use crate::providers::claude_acp::{CLAUDE_ACP_BINARY, CLAUDE_ACP_PROVIDER_NAME}; use crate::providers::codex_acp::CODEX_ACP_PROVIDER_NAME; use crate::providers::copilot_acp::{COPILOT_ACP_BINARY, COPILOT_ACP_PROVIDER_NAME}; use crate::providers::formats::anthropic::ANTHROPIC_PROVIDER_NAME; +use crate::providers::gemini_oauth::TokenCache as GeminiOAuthTokenCache; use crate::providers::google::{GOOGLE_API_HOST, GOOGLE_PROVIDER_NAME}; use crate::providers::huggingface::HuggingFaceProvider; use crate::providers::huggingface_auth; @@ -155,6 +156,15 @@ pub fn chatgpt_codex_inventory() -> InventoryRegistration { .with_configured(|| ChatGptCodexTokenCache::new().has_token()) } +pub fn gemini_oauth_inventory() -> InventoryRegistration { + InventoryRegistration { + supports_refresh: false, + identity: default_inventory_identity_resolver(), + configured: None, + } + .with_configured(|| GeminiOAuthTokenCache::new().has_token()) +} + pub fn xai_oauth_inventory() -> InventoryRegistration { InventoryRegistration { supports_refresh: false, @@ -196,3 +206,43 @@ pub fn copilot_acp_inventory() -> InventoryRegistration { pub fn pi_acp_inventory() -> InventoryRegistration { acp_inventory(PI_ACP_PROVIDER_NAME, PI_ACP_BINARY, false) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::paths::Paths; + use chrono::Utc; + + #[test] + #[serial_test::serial] + fn gemini_oauth_inventory_configured_uses_token_cache() { + let root = tempfile::tempdir().unwrap(); + let root_path = root.path().to_string_lossy().to_string(); + let _guard = env_lock::lock_env([("GOOSE_PATH_ROOT", Some(root_path.as_str()))]); + + let registration = gemini_oauth_inventory(); + let configured = registration + .configured + .expect("Gemini OAuth should define configured resolver"); + + assert!(!configured()); + + let cache_path = Paths::in_config_dir("gemini_oauth/tokens.json"); + std::fs::create_dir_all(cache_path.parent().unwrap()).unwrap(); + std::fs::write( + cache_path, + serde_json::to_string(&serde_json::json!({ + "project_id": "test-project", + "token": { + "access_token": "access", + "refresh_token": "refresh", + "expires_at": (Utc::now() + chrono::Duration::hours(1)).to_rfc3339(), + }, + })) + .unwrap(), + ) + .unwrap(); + + assert!(configured()); + } +} diff --git a/ui/desktop/src/components/onboarding/LocalModelPicker.tsx b/ui/desktop/src/components/onboarding/LocalModelPicker.tsx index 17a3c7441d3b..b29aa4de7bc3 100644 --- a/ui/desktop/src/components/onboarding/LocalModelPicker.tsx +++ b/ui/desktop/src/components/onboarding/LocalModelPicker.tsx @@ -9,6 +9,7 @@ import { } from '../../acp/local-inference'; import { trackOnboardingSetupFailed } from '../../utils/analytics'; import { defineMessages, useIntl } from '../../i18n'; +import { errorMessage as formatErrorMessage } from '../../utils/conversionUtils'; const i18n = defineMessages({ checkingModels: { @@ -61,7 +62,8 @@ const i18n = defineMessages({ }, localModelsNote: { id: 'localModelPicker.localModelsNote', - defaultMessage: 'Local models keep everything on your machine for full privacy. Performance and context window size may vary compared to cloud providers depending on your hardware and model size.', + defaultMessage: + 'Local models keep everything on your machine for full privacy. Performance and context window size may vary compared to cloud providers depending on your hardware and model size.', }, failedToLoad: { id: 'localModelPicker.failedToLoad', @@ -82,7 +84,7 @@ const i18n = defineMessages({ }); interface LocalModelPickerProps { - onConfigured: (providerName: string, modelId: string) => void; + onConfigured: (providerName: string, modelId: string) => void | Promise; } const formatBytes = (bytes: number): string => { @@ -146,8 +148,15 @@ export default function LocalModelPicker({ onConfigured }: LocalModelPickerProps load(); }, [intl]); - const finishSetup = (modelId: string) => { - onConfigured(LOCAL_PROVIDER, modelId); + const finishSetup = async (modelId: string) => { + try { + await onConfigured(LOCAL_PROVIDER, modelId); + } catch (error) { + console.error('Failed to finish local model setup:', error); + setErrorMessage(formatErrorMessage(error)); + trackOnboardingSetupFailed(LOCAL_PROVIDER, 'save_defaults_failed'); + setPhase('error'); + } }; const startDownload = async (modelId: string) => { @@ -186,7 +195,23 @@ export default function LocalModelPicker({ onConfigured }: LocalModelPickerProps setDownloadProgress(progress); if (progress.status === 'completed') { cleanup(); - finishSetup(modelId); + setModels((previousModels) => + previousModels.map((model) => + model.id === modelId + ? { + ...model, + status: { + ...model.status, + state: 'Downloaded', + progressPercent: 100, + bytesDownloaded: model.sizeBytes, + totalBytes: model.sizeBytes, + }, + } + : model + ) + ); + await finishSetup(modelId); } else if (progress.status === 'failed') { cleanup(); setErrorMessage(progress.error || 'Download failed.'); @@ -223,7 +248,7 @@ export default function LocalModelPicker({ onConfigured }: LocalModelPickerProps const model = models.find((m) => m.id === selectedModelId); if (!model) return; if (model.status.state === 'Downloaded') { - finishSetup(model.id); + await finishSetup(model.id); } else { await startDownload(model.id); } @@ -310,7 +335,9 @@ export default function LocalModelPicker({ onConfigured }: LocalModelPickerProps onClick={() => setShowAllModels(!showAllModels)} className="text-sm text-blue-500 hover:text-blue-400 transition-colors flex items-center gap-1" > - {showAllModels ? intl.formatMessage(i18n.hideOtherSizes) : intl.formatMessage(i18n.showOtherSizes, { count: otherModels.length })} + {showAllModels + ? intl.formatMessage(i18n.hideOtherSizes) + : intl.formatMessage(i18n.showOtherSizes, { count: otherModels.length })} - )} @@ -427,7 +456,9 @@ export default function LocalModelPicker({ onConfigured }: LocalModelPickerProps ) : (
- {intl.formatMessage(i18n.startingDownload)} + + {intl.formatMessage(i18n.startingDownload)} +
)} diff --git a/ui/desktop/src/components/onboarding/ProviderConfigForm.tsx b/ui/desktop/src/components/onboarding/ProviderConfigForm.tsx index a926c45f2352..4cfd3ccddc1b 100644 --- a/ui/desktop/src/components/onboarding/ProviderConfigForm.tsx +++ b/ui/desktop/src/components/onboarding/ProviderConfigForm.tsx @@ -10,6 +10,9 @@ import { SecureStorageNotice } from '../settings/providers/modal/subcomponents/S import { Button } from '../ui/button'; import { LogIn, ChevronRight } from 'lucide-react'; import { defineMessages, useIntl } from '../../i18n'; +import { errorMessage } from '../../utils/conversionUtils'; + +type OnConfigured = (name: string) => void | Promise; const i18n = defineMessages({ browserWindowOpen: { @@ -69,7 +72,7 @@ function OAuthForm({ onError, }: { provider: ProviderDetails; - onConfigured: (name: string) => void; + onConfigured: OnConfigured; onError: (msg: string) => void; }) { const intl = useIntl(); @@ -79,9 +82,9 @@ function OAuthForm({ setIsLoading(true); try { await acpAuthenticateProvider(provider.name); - onConfigured(provider.name); + await onConfigured(provider.name); } catch (err) { - onError(`Sign-in failed: ${err instanceof Error ? err.message : String(err)}`); + onError(`Setup failed: ${errorMessage(err)}`); } finally { setIsLoading(false); } @@ -117,7 +120,7 @@ function ApiKeyForm({ onError, }: { provider: ProviderDetails; - onConfigured: (name: string) => void; + onConfigured: OnConfigured; onError: (msg: string) => void; }) { const intl = useIntl(); @@ -157,15 +160,9 @@ function ApiKeyForm({ setIsSubmitting(true); try { await providerConfigSubmitHandler(provider, toSubmit); - onConfigured(provider.name); + await onConfigured(provider.name); } catch (err) { - const msg = - err instanceof Error - ? err.message - : typeof err === 'object' && err !== null && 'message' in err - ? String((err as Record).message) - : JSON.stringify(err); - onError(msg); + onError(errorMessage(err)); } finally { setIsSubmitting(false); } @@ -214,7 +211,7 @@ function ApiKeyForm({ interface ProviderConfigFormProps { provider: ProviderDetails; - onConfigured: (providerName: string) => void; + onConfigured: OnConfigured; } export default function ProviderConfigForm({ provider, onConfigured }: ProviderConfigFormProps) { diff --git a/ui/desktop/src/components/onboarding/ProviderSelector.tsx b/ui/desktop/src/components/onboarding/ProviderSelector.tsx index 449a24b24410..d484ed350d07 100644 --- a/ui/desktop/src/components/onboarding/ProviderSelector.tsx +++ b/ui/desktop/src/components/onboarding/ProviderSelector.tsx @@ -53,7 +53,7 @@ interface ProviderOption { } interface ProviderSelectorProps { - onConfigured: (providerName: string, modelId?: string) => void; + onConfigured: (providerName: string, modelId?: string) => void | Promise; onFirstSelection?: () => void; } @@ -123,7 +123,7 @@ export default function ProviderSelector({ const result = await acpCreateCustomProviderFromRequest(data); setShowCustomModal(false); if (result.provider_name) { - onConfigured(result.provider_name); + await onConfigured(result.provider_name); } };