From 723a825fb9ff1f744dbed49d4516319e2b4e1c15 Mon Sep 17 00:00:00 2001 From: Douwe Osinga Date: Wed, 7 Oct 2026 19:21:55 +0000 Subject: [PATCH] Add ProviderManager and route session providers through it (#12745) Co-authored-by: Claude Opus 5.5 --- crates/goose-cli/src/commands/configure.rs | 4 +- crates/goose-cli/src/session/builder.rs | 40 +- crates/goose-cli/src/session/mod.rs | 57 ++- crates/goose/examples/agent.rs | 4 +- crates/goose/examples/databricks_oauth.rs | 4 +- crates/goose/examples/image_tool.rs | 8 +- crates/goose/src/acp/response_builder.rs | 7 +- crates/goose/src/acp/server.rs | 76 ++-- crates/goose/src/acp/server/dispatch.rs | 2 +- crates/goose/src/acp/server/fork_session.rs | 4 +- crates/goose/src/acp/server/load_session.rs | 4 +- .../goose/src/acp/server/manage_sessions.rs | 1 + crates/goose/src/acp/server/new_session.rs | 2 +- crates/goose/src/acp/server/providers.rs | 13 +- crates/goose/src/acp/server_factory.rs | 27 +- crates/goose/src/agents/agent.rs | 391 +++++++----------- crates/goose/src/agents/execute_commands.rs | 6 +- .../goose/src/agents/extension_manager/mod.rs | 41 +- crates/goose/src/agents/mod.rs | 1 + crates/goose/src/agents/moim.rs | 54 ++- .../agents/platform_extensions/analyze/mod.rs | 2 +- .../src/agents/platform_extensions/apps.rs | 19 +- .../agents/platform_extensions/chatrecall.rs | 2 +- .../platform_extensions/code_execution.rs | 18 +- .../platform_extensions/developer/mod.rs | 2 +- .../agents/platform_extensions/ext_manager.rs | 8 +- .../src/agents/platform_extensions/mod.rs | 17 +- .../agents/platform_extensions/summarize.rs | 15 +- .../src/agents/platform_extensions/summon.rs | 20 +- crates/goose/src/agents/prompt_manager.rs | 2 +- crates/goose/src/agents/provider_manager.rs | 109 +++++ .../agents/state_machine/tests/pipeline.rs | 7 +- crates/goose/src/agents/subagent_handler.rs | 12 +- crates/goose/src/agents/types.rs | 5 - .../goose/src/bin/build_canonical_models.rs | 4 +- crates/goose/src/config/base.rs | 14 + crates/goose/src/execution/manager.rs | 128 +----- .../src/permission/permission_inspector.rs | 17 +- crates/goose/src/providers/amp_acp.rs | 1 + crates/goose/src/providers/base.rs | 4 + crates/goose/src/providers/claude_acp.rs | 1 + crates/goose/src/providers/claude_code.rs | 1 + crates/goose/src/providers/codex.rs | 1 + crates/goose/src/providers/codex_acp.rs | 1 + crates/goose/src/providers/copilot_acp.rs | 1 + crates/goose/src/providers/gemini_cli.rs | 1 + crates/goose/src/providers/init.rs | 17 - crates/goose/src/providers/mod.rs | 5 +- crates/goose/src/providers/pi_acp.rs | 1 + .../goose/src/providers/provider_registry.rs | 7 + crates/goose/src/scheduler/full.rs | 4 +- .../goose/src/security/adversary_inspector.rs | 52 +-- crates/goose/src/tool_call_labels.rs | 4 +- crates/goose/tests/acp_common_tests/mod.rs | 2 +- .../goose/tests/acp_custom_requests_test.rs | 31 +- crates/goose/tests/acp_fixtures/mod.rs | 28 +- .../acp_secret_cache_invalidation_test.rs | 16 +- .../goose/tests/adversary_inspector_tests.rs | 10 +- crates/goose/tests/mcp_integration_test.rs | 4 +- crates/goose/tests/providers.rs | 4 +- 60 files changed, 596 insertions(+), 747 deletions(-) create mode 100644 crates/goose/src/agents/provider_manager.rs diff --git a/crates/goose-cli/src/commands/configure.rs b/crates/goose-cli/src/commands/configure.rs index a437255de1..f9fef71d55 100644 --- a/crates/goose-cli/src/commands/configure.rs +++ b/crates/goose-cli/src/commands/configure.rs @@ -1819,10 +1819,8 @@ pub async fn configure_tool_permissions_dialog() -> anyhow::Result<()> { return Ok(()); } - let extensions = extension_config.into_iter().collect::>(); - let new_provider = create(&provider_name, extensions).await?; agent - .update_provider(new_provider, model_config, &session.id) + .switch_provider(&session.id, &provider_name, model_config) .await?; let permission_manager = PermissionManager::instance(); diff --git a/crates/goose-cli/src/session/builder.rs b/crates/goose-cli/src/session/builder.rs index 35ad6b90e9..e383c4d4c4 100644 --- a/crates/goose-cli/src/session/builder.rs +++ b/crates/goose-cli/src/session/builder.rs @@ -1,7 +1,9 @@ use crate::cli::StreamableHttpOptions; use super::output; -use super::{derive_extension_name_from_command, split_extension_name_prefix, CliSession}; +use super::{ + derive_extension_name_from_command, session_provider, split_extension_name_prefix, CliSession, +}; use console::style; use goose::agents::final_output_tool::FinalOutputTool; use goose::agents::{Agent, Container, ExtensionError}; @@ -9,7 +11,6 @@ use goose::config::extensions::name_to_key; use goose::config::resolve_extensions_for_new_session; use goose::config::{Config, ExtensionConfig, GooseMode}; use goose::model_config::model_config_from_user_config; -use goose::providers::create; use goose::recipe::Recipe; use goose::session::session_manager::SessionType; use goose::session::{EnabledExtensionsState, SessionManager}; @@ -715,9 +716,16 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession { process::exit(1); } }; + if let Err(e) = agent + .persist_extension_configs(&session_id, extensions_for_provider.clone()) + .await + { + output::render_error(&format!("Failed to save session extensions: {}", e)); + process::exit(1); + } let (new_provider, effective_provider_name, effective_model_name, effective_model_config) = - match create(&resolved.provider_name, extensions_for_provider.clone()).await { + match session_provider(&agent, &session_id, &resolved.provider_name).await { Ok(provider) => ( provider, resolved.provider_name.clone(), @@ -758,7 +766,7 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession { if !session_config.interactive { fallback_model_config = fallback_model_config.with_cache_ttl_clamped(); } - match create(&fallback_provider, extensions_for_provider.clone()).await { + match session_provider(&agent, &session_id, &fallback_provider).await { Ok(provider) => ( provider, fallback_provider, @@ -804,7 +812,11 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession { }); agent - .update_provider(new_provider, effective_model_config, &session_id) + .switch_provider( + &session_id, + &effective_provider_name, + effective_model_config, + ) .await .unwrap_or_else(|e| { output::render_error(&format!("Failed to initialize agent: {}", e)); @@ -837,19 +849,11 @@ pub async fn build_session(session_config: SessionBuilderConfig) -> CliSession { // Extensions are loaded after session creation because we may change // directory when resuming. let agent_ptr = Arc::new(agent); - let loading_handle = match agent_ptr - .persist_extension_configs(&session_id, extensions_for_provider.clone()) - .await - { - Ok(()) => AbortOnDropHandle::new(tokio::spawn({ - let agent = agent_ptr.clone(); - let sid = session_id.clone(); - async move { load_extensions(agent, extensions_for_provider, &sid).await } - })), - Err(error) => AbortOnDropHandle::new(tokio::spawn(async move { - vec![ExtensionFailure { label: None, error }] - })), - }; + let loading_handle = AbortOnDropHandle::new(tokio::spawn({ + let agent = agent_ptr.clone(); + let sid = session_id.clone(); + async move { load_extensions(agent, extensions_for_provider, &sid).await } + })); let edit_mode = config .get_param::("EDIT_MODE") diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 86843a8ac3..fbf30f0bee 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -27,7 +27,7 @@ use goose::agents::platform_extensions::developer::shell::{ use goose::agents::AgentEvent; use goose::agents::SUBAGENT_TOOL_REQUEST_TYPE; use goose::permission::Permission; -use goose::providers::base::ProviderUsage; +use goose::providers::base::{Provider, ProviderUsage}; use goose::utils::safe_truncate; use anyhow::Result; @@ -786,7 +786,7 @@ impl CliSession { history.save(editor); self.push_message(Message::user().with_text(content)); - let _provider = self.agent.provider().await?; + let _provider = self.agent.provider(&self.session_id).await?; println!(); output::run_status_hook("thinking"); @@ -880,7 +880,7 @@ impl CliSession { } async fn handle_model(&mut self, options: input::ModelCommandOptions) -> Result<()> { - let provider = self.agent.provider().await?; + let provider = self.agent.provider(&self.session_id).await?; let current_provider_name = provider.get_name().to_string(); let current_model_config = self .agent @@ -992,8 +992,13 @@ impl CliSession { ) .await?; - let extensions = self.agent.get_extension_configs().await; - let new_provider = match goose::providers::create(target_provider_name, extensions).await { + let new_provider = match session_provider( + &self.agent, + &self.session_id, + target_provider_name, + ) + .await + { Ok(p) => p, Err(e) => { output::render_error(&format!( @@ -1031,12 +1036,9 @@ impl CliSession { } self.agent - .update_provider(new_provider, new_model_config, &self.session_id) + .switch_provider(&self.session_id, target_provider_name, new_model_config) .await?; - let mode = self.agent.goose_mode().await; - self.agent.update_goose_mode(mode, &self.session_id).await?; - self.update_completion_cache().await?; if provider_unchanged { @@ -1054,7 +1056,7 @@ impl CliSession { } async fn handle_clear(&mut self) -> Result<()> { - let provider = self.agent.provider().await?; + let provider = self.agent.provider(&self.session_id).await?; if provider.manages_own_context() { output::render_error(&context_management_unsupported_message( "clear", @@ -1101,7 +1103,7 @@ impl CliSession { } async fn handle_new(&mut self) -> Result<()> { - let provider = self.agent.provider().await?; + let provider = self.agent.provider(&self.session_id).await?; if provider.manages_own_context() { output::render_error(&format!( "Starting a new session is not supported for provider '{}' because it manages its own conversation context.", @@ -1130,6 +1132,7 @@ impl CliSession { .await; self.agent.discard_pending_steers(&self.session_id).await; + self.agent.config.providers.release(&self.session_id); self.session_id = new_session_id; self.messages.clear(); @@ -1250,7 +1253,7 @@ impl CliSession { } async fn handle_compact(&mut self) -> Result<()> { - let provider = self.agent.provider().await?; + let provider = self.agent.provider(&self.session_id).await?; if provider.manages_own_context() { output::render_error(&context_management_unsupported_message( "compact", @@ -1766,7 +1769,7 @@ impl CliSession { ) -> Result<()> { let prompts = agent.list_extension_prompts(session_id).await; let all_providers = goose::providers::providers().await; - let session_provider = agent.provider().await?.get_name().to_string(); + let session_provider = agent.provider(session_id).await?.get_name().to_string(); let provider_ids: Vec = all_providers.iter().map(|(m, _)| m.name.clone()).collect(); let inventory_models: HashMap> = { @@ -1896,7 +1899,7 @@ impl CliSession { /// Display enhanced context usage with session totals pub async fn display_context_usage(&self) -> Result<()> { - let provider = self.agent.provider().await?; + let provider = self.agent.provider(&self.session_id).await?; let model_config = self .agent .model_config_for_session(&self.session_id) @@ -2015,6 +2018,22 @@ impl CliSession { } } +/// The provider `session_id` would get with `provider_name`, without switching +/// the session to it. +pub(crate) async fn session_provider( + agent: &Agent, + session_id: &str, + provider_name: &str, +) -> anyhow::Result> { + let mut session = agent + .config + .session_manager + .get_session(session_id, false) + .await?; + session.provider_name = Some(provider_name.to_string()); + agent.config.providers.provider_for(&session).await +} + async fn create_successor_session( session_manager: &SessionManager, old_session: &goose::session::Session, @@ -3177,11 +3196,11 @@ mod tests { extension_loading: Option>>, refresh_completions: bool, ) -> CliSession { - let temp_dir = tempfile::TempDir::new().unwrap(); - let session_manager = SessionManager::new(temp_dir.path().to_path_buf()); + let data_dir = tempfile::TempDir::new().unwrap().keep(); + let session_manager = SessionManager::new(data_dir.clone()); let session = session_manager .create_session( - temp_dir.path().to_path_buf(), + data_dir.clone(), "Loading gate test".to_string(), goose::session::SessionType::User, GooseMode::default(), @@ -3191,9 +3210,7 @@ mod tests { let agent = goose::agents::Agent::with_config(goose::agents::AgentConfig::new( Arc::new(session_manager), - Arc::new(goose::config::PermissionManager::new( - temp_dir.path().to_path_buf(), - )), + Arc::new(goose::config::PermissionManager::new(data_dir.clone())), None, GooseMode::default(), // Disable background session naming so the test agent starts no diff --git a/crates/goose/examples/agent.rs b/crates/goose/examples/agent.rs index bba1eab1b1..fbb43007b8 100644 --- a/crates/goose/examples/agent.rs +++ b/crates/goose/examples/agent.rs @@ -3,7 +3,7 @@ use futures::StreamExt; use goose::agents::{Agent, AgentEvent, ExtensionConfig, SessionConfig}; use goose::config::{GooseMode, DEFAULT_EXTENSION_DESCRIPTION, DEFAULT_EXTENSION_TIMEOUT}; use goose::conversation::message::Message; -use goose::providers::create_with_named_model; +use goose::providers::create; use goose::session::session_manager::SessionType; use goose_providers::databricks::DATABRICKS_DEFAULT_MODEL; use std::path::PathBuf; @@ -12,7 +12,7 @@ use std::path::PathBuf; async fn main() -> anyhow::Result<()> { let _ = dotenv(); - let provider = create_with_named_model("databricks", Vec::new()).await?; + let provider = create("databricks", Vec::new()).await?; let model_config = goose::model_config::model_config_from_user_config("databricks", DATABRICKS_DEFAULT_MODEL)?; diff --git a/crates/goose/examples/databricks_oauth.rs b/crates/goose/examples/databricks_oauth.rs index 904829939e..8e93a49a47 100644 --- a/crates/goose/examples/databricks_oauth.rs +++ b/crates/goose/examples/databricks_oauth.rs @@ -1,7 +1,7 @@ use anyhow::Result; use dotenvy::dotenv; use goose::conversation::message::Message; -use goose::providers::create_with_named_model; +use goose::providers::create; use goose_providers::databricks::DATABRICKS_DEFAULT_MODEL; #[tokio::main] @@ -10,7 +10,7 @@ async fn main() -> Result<()> { std::env::remove_var("DATABRICKS_TOKEN"); - let provider = create_with_named_model("databricks", Vec::new()).await?; + let provider = create("databricks", Vec::new()).await?; let message = Message::user().with_text("Tell me a short joke about programming."); diff --git a/crates/goose/examples/image_tool.rs b/crates/goose/examples/image_tool.rs index df330317b9..dfc674d061 100644 --- a/crates/goose/examples/image_tool.rs +++ b/crates/goose/examples/image_tool.rs @@ -3,7 +3,7 @@ use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _}; use dotenvy::dotenv; use goose::conversation::message::Message; use goose::providers::anthropic::ANTHROPIC_DEFAULT_MODEL; -use goose::providers::create_with_named_model; +use goose::providers::create; use goose::providers::openai::OPEN_AI_DEFAULT_MODEL; use goose_providers::databricks::DATABRICKS_DEFAULT_MODEL; use rmcp::model::{CallToolRequestParams, ContentBlock, Tool}; @@ -22,18 +22,18 @@ async fn main() -> Result<()> { goose_providers::model::ModelConfig, )> = vec![ ( - create_with_named_model("databricks", Vec::new()).await?, + create("databricks", Vec::new()).await?, goose::model_config::model_config_from_user_config( "databricks", DATABRICKS_DEFAULT_MODEL, )?, ), ( - create_with_named_model("openai", Vec::new()).await?, + create("openai", Vec::new()).await?, goose::model_config::model_config_from_user_config("openai", OPEN_AI_DEFAULT_MODEL)?, ), ( - create_with_named_model("anthropic", Vec::new()).await?, + create("anthropic", Vec::new()).await?, goose::model_config::model_config_from_user_config( "anthropic", ANTHROPIC_DEFAULT_MODEL, diff --git a/crates/goose/src/acp/response_builder.rs b/crates/goose/src/acp/response_builder.rs index c7cf96b6fb..2e3cfeb29a 100644 --- a/crates/goose/src/acp/response_builder.rs +++ b/crates/goose/src/acp/response_builder.rs @@ -235,8 +235,11 @@ pub(super) fn build_mode_state( /// The provider decides whether goose owns the effort menu; a session without a /// live provider keeps the model-name based path. -pub(super) async fn agent_thinking_effort_support(agent: &Agent) -> ThinkingEffortSupport { - match agent.provider().await { +pub(super) async fn agent_thinking_effort_support( + agent: &Agent, + session_id: &str, +) -> ThinkingEffortSupport { + match agent.provider(session_id).await { Ok(provider) => provider.thinking_effort_support(), Err(_) => ThinkingEffortSupport::Unspecified, } diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index f6ee78a78c..645cc86dff 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -130,16 +130,8 @@ mod tool_calls; mod tool_notifications; mod tools; -pub type AcpProviderFactory = Arc< - dyn Fn( - String, - Vec, - Option, - bool, - ) -> BoxFuture<'static, Result>> - + Send - + Sync, ->; +pub type AcpProviderFactory = + Arc BoxFuture<'static, Result>> + Send + Sync>; const ACP_VISIBLE_SESSION_TYPES: [SessionType; 3] = [SessionType::User, SessionType::Scheduled, SessionType::Acp]; @@ -894,7 +886,7 @@ impl GooseAcpAgent { let provider_name = session.provider_name.clone(); let agent = self.get_session_agent(session_id).await?; let provider = agent - .provider() + .provider(session_id) .await .internal_err_ctx("Failed to resolve session provider")?; let context_limit = @@ -907,7 +899,7 @@ impl GooseAcpAgent { .await .internal_err_ctx("Failed to refresh session for setup notifications")?; let current_provider = agent - .provider() + .provider(session_id) .await .internal_err_ctx("Failed to refresh session provider")?; let refreshed_model_name = session @@ -1007,20 +999,8 @@ impl GooseAcpAgent { Ok(Config::global()) } - async fn create_provider( - &self, - provider_name: &str, - extensions: Vec, - working_dir: Option, - use_default_model: bool, - ) -> Result> { - (self.provider_factory)( - provider_name.to_string(), - extensions, - working_dir, - use_default_model, - ) - .await + async fn create_provider(&self, provider_name: &str) -> Result> { + (self.provider_factory)(provider_name.to_string()).await } /// Warm the provider model-list cache after session creation. @@ -1047,7 +1027,7 @@ impl GooseAcpAgent { if !should_refresh_inventory_for_session_init(&inventory) { return; } - let provider = match agent.provider().await { + let provider = match agent.provider(&session_id).await { Ok(provider) => provider, Err(error) => { warn!( @@ -1253,7 +1233,7 @@ impl GooseAcpAgent { } async fn subscribe_thinking_effort_updates(&self, session_id: &str, agent: &Arc) { - let Ok(provider) = agent.provider().await else { + let Ok(provider) = agent.provider(session_id).await else { return; }; let Some(mut updates) = provider.subscribe_thinking_effort_support() else { @@ -2036,7 +2016,7 @@ impl GooseAcpAgent { session_id: &str, agent: &Arc, ) -> Result<(), agent_client_protocol::Error> { - let Ok(provider) = agent.provider().await else { + let Ok(provider) = agent.provider(session_id).await else { return Ok(()); }; if provider.get_name() != "local" { @@ -2072,7 +2052,7 @@ impl GooseAcpAgent { agent: &Arc, ) -> Result { let provider = agent - .provider() + .provider(&session.id) .await .internal_err_ctx("Failed to resolve session provider")?; let model = session.model_config.as_ref().ok_or_else(|| { @@ -2443,7 +2423,7 @@ impl GooseAcpAgent { ) -> Result<(), agent_client_protocol::Error> { let agent = self.get_session_agent(session_id).await?; let current_provider = agent - .provider() + .provider(session_id) .await .internal_err_ctx("Failed to get provider")?; let provider_name = current_provider.get_name().to_string(); @@ -2461,12 +2441,9 @@ impl GooseAcpAgent { ) .invalid_params_err_ctx("Invalid model config")?; agent - .recreate_provider_for_session(session_id, &provider_name, model_config) + .switch_provider(session_id, &provider_name, model_config) .await - .internal_err_ctx("Failed to recreate provider")?; - self.subscribe_thinking_effort_updates(session_id, &agent) - .await; - // model_config is already updated on the session by the agent's update_provider call. + .internal_err_ctx("Failed to switch provider")?; Ok(()) } @@ -2481,7 +2458,7 @@ impl GooseAcpAgent { .internal_err()?; let agent = self.get_session_agent(&session_id.0).await?; let provider = agent - .provider() + .provider(&session_id.0) .await .internal_err_ctx("Failed to get provider")?; let provider_name = provider.get_name().to_string(); @@ -2564,7 +2541,7 @@ impl GooseAcpAgent { let config = self.config()?; let agent = self.get_session_agent(session_id).await?; let current_provider = agent - .provider() + .provider(session_id) .await .internal_err_ctx("Failed to get provider")?; let current_provider_name = current_provider.get_name(); @@ -2608,14 +2585,13 @@ impl GooseAcpAgent { ) .invalid_params_err_ctx("Invalid model config")?; + agent.config.providers.release(session_id); agent - .recreate_provider_for_session(session_id, &resolved_provider_name, model_config) + .switch_provider(session_id, &resolved_provider_name, model_config) .await - .internal_err_ctx("Failed to recreate provider")?; + .internal_err_ctx("Failed to switch provider")?; self.subscribe_thinking_effort_updates(session_id, &agent) .await; - - // provider_name is already updated on the session by the agent's update_provider call. Ok(()) } @@ -3658,11 +3634,9 @@ print(\"hello, world\") let root = tempfile::tempdir().unwrap(); let active_runs = Arc::new(ActiveRunRegistry::default()); let live_voice = Arc::new(LiveVoiceService::from_config(active_runs.clone())); - let provider_factory: AcpProviderFactory = Arc::new( - |_provider_name, _extensions, _working_dir, _use_default_model| { - Box::pin(async { Err(anyhow::anyhow!("unused provider factory")) }) - }, - ); + let provider_factory: AcpProviderFactory = Arc::new(|_provider_name| { + Box::pin(async { Err(anyhow::anyhow!("unused provider factory")) }) + }); let agent = GooseAcpAgent::new(GooseAcpAgentOptions { provider_factory, builtin_selection: AcpBuiltinSelection::default(), @@ -3791,11 +3765,9 @@ print(\"hello, world\") let root = tempfile::tempdir().unwrap(); let active_runs = Arc::new(ActiveRunRegistry::default()); let live_voice = Arc::new(LiveVoiceService::from_config(active_runs.clone())); - let provider_factory: AcpProviderFactory = Arc::new( - |_provider_name, _extensions, _working_dir, _use_default_model| { - Box::pin(async { Err(anyhow::anyhow!("unused provider factory")) }) - }, - ); + let provider_factory: AcpProviderFactory = Arc::new(|_provider_name| { + Box::pin(async { Err(anyhow::anyhow!("unused provider factory")) }) + }); let server = Arc::new( GooseAcpAgent::new(GooseAcpAgentOptions { provider_factory, diff --git a/crates/goose/src/acp/server/dispatch.rs b/crates/goose/src/acp/server/dispatch.rs index f4c348a45c..f8be1ffc56 100644 --- a/crates/goose/src/acp/server/dispatch.rs +++ b/crates/goose/src/acp/server/dispatch.rs @@ -244,7 +244,7 @@ impl HandleDispatchFrom for GooseAcpHandler { let session_agent = agent_bg.get_session_agent(&session_id_bg.0).await?; let provider = session_agent - .provider() + .provider(&session_id_bg.0) .await .map_err(|e| anyhow::anyhow!(e.to_string()))?; let provider_name = provider.get_name().to_string(); diff --git a/crates/goose/src/acp/server/fork_session.rs b/crates/goose/src/acp/server/fork_session.rs index ce585b701a..8efd070f3a 100644 --- a/crates/goose/src/acp/server/fork_session.rs +++ b/crates/goose/src/acp/server/fork_session.rs @@ -57,11 +57,11 @@ impl GooseAcpAgent { self.register_acp_session(goose_session.id.clone(), agent.clone()) .await; let provider = agent - .provider() + .provider(&goose_session.id) .await .internal_err_ctx("Failed to get provider while forking ACP session")?; resume_saved_provider_session(&provider, goose_session.conversation.as_ref()).await; - let effort_support = agent_thinking_effort_support(&agent).await; + let effort_support = agent_thinking_effort_support(&agent, &goose_session.id).await; let acp_session_id = SessionId::new(new_session_id.clone()); let mut meta = session_meta(&goose_session); diff --git a/crates/goose/src/acp/server/load_session.rs b/crates/goose/src/acp/server/load_session.rs index d0b7f182b3..890896ebb9 100644 --- a/crates/goose/src/acp/server/load_session.rs +++ b/crates/goose/src/acp/server/load_session.rs @@ -416,7 +416,7 @@ impl GooseAcpAgent { self.register_acp_session(session_id_str.clone(), agent.clone()) .await; let provider = agent - .provider() + .provider(&session.id) .await .internal_err_ctx("Failed to get provider while loading ACP session")?; resume_saved_provider_session(&provider, session.conversation.as_ref()).await; @@ -434,7 +434,7 @@ impl GooseAcpAgent { let (mode_state, config_options) = build_session_setup_config( &self.provider_inventory, &session, - &agent_thinking_effort_support(&agent).await, + &agent_thinking_effort_support(&agent, &session.id).await, ) .await?; diff --git a/crates/goose/src/acp/server/manage_sessions.rs b/crates/goose/src/acp/server/manage_sessions.rs index ded22ec606..63a441f75f 100644 --- a/crates/goose/src/acp/server/manage_sessions.rs +++ b/crates/goose/src/acp/server/manage_sessions.rs @@ -34,6 +34,7 @@ impl GooseAcpAgent { .get_session(session_id, false) .await .internal_err_ctx("Failed to reload session")?; + agent.config.providers.release(&session.id); agent .restore_provider_from_session(&session) .await diff --git a/crates/goose/src/acp/server/new_session.rs b/crates/goose/src/acp/server/new_session.rs index 9bb45cbbc5..dd716e895a 100644 --- a/crates/goose/src/acp/server/new_session.rs +++ b/crates/goose/src/acp/server/new_session.rs @@ -92,7 +92,7 @@ impl GooseAcpAgent { .build_new_session_response( &reloaded_session, &extension_results, - &super::agent_thinking_effort_support(&agent).await, + &super::agent_thinking_effort_support(&agent, &reloaded_session.id).await, ) .await?; Ok(response) diff --git a/crates/goose/src/acp/server/providers.rs b/crates/goose/src/acp/server/providers.rs index df9bf85858..3514cfe9f7 100644 --- a/crates/goose/src/acp/server/providers.rs +++ b/crates/goose/src/acp/server/providers.rs @@ -504,7 +504,7 @@ impl GooseAcpAgent { req: ProviderSupportedModelsListRequest, ) -> Result { let provider = self - .create_provider(&req.provider_id, Vec::new(), None, true) + .create_provider(&req.provider_id) .await .internal_err_ctx("Failed to initialize provider")?; let models = match provider.fetch_supported_models().await { @@ -546,7 +546,7 @@ impl GooseAcpAgent { let result = tokio::time::timeout( ACP_READINESS_TIMEOUT, - self.create_provider(&req.provider_id, Vec::new(), None, true), + self.create_provider(&req.provider_id), ) .await; let (ready, error) = match result { @@ -865,11 +865,10 @@ impl GooseAcpAgent { let toolshim = refresh_job.toolshim; tokio::spawn(async move { let mut refresh_guard = provider_inventory.refresh_guard(&identity); - let provider_result = AssertUnwindSafe(async { - provider_factory(provider_id.clone(), Vec::new(), None, true).await - }) - .catch_unwind() - .await; + let provider_result = + AssertUnwindSafe(async { provider_factory(provider_id.clone()).await }) + .catch_unwind() + .await; let fetch_result: Result> = match provider_result { Ok(Ok(provider)) => { diff --git a/crates/goose/src/acp/server_factory.rs b/crates/goose/src/acp/server_factory.rs index 1f6c1ff1c4..e7bdfd4ace 100644 --- a/crates/goose/src/acp/server_factory.rs +++ b/crates/goose/src/acp/server_factory.rs @@ -106,28 +106,11 @@ impl AcpServer { scheduler.list_scheduled_jobs().await; } - let provider_factory: AcpProviderFactory = Arc::new( - move |provider_name, extensions, working_dir, use_default_model| { - Box::pin(async move { - if use_default_model { - crate::providers::create_with_default_model(&provider_name, extensions) - .await - } else { - match working_dir { - Some(working_dir) => { - crate::providers::create_with_working_dir( - &provider_name, - extensions, - working_dir, - ) - .await - } - None => crate::providers::create(&provider_name, extensions).await, - } - } - }) - }, - ); + let provider_factory: AcpProviderFactory = Arc::new(|provider_name| { + Box::pin(async move { + crate::providers::create_with_default_model(&provider_name, Vec::new()).await + }) + }); let agent = GooseAcpAgent::new(GooseAcpAgentOptions { provider_factory, diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 4e15e62d5e..025f7391bd 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -30,6 +30,7 @@ use crate::agents::final_output_tool::{ FINAL_OUTPUT_TOOL_NAME, }; use crate::agents::prompt_manager::PromptManager; +use crate::agents::provider_manager::ProviderManager; use crate::agents::retry::{RetryManager, RetryResult}; use crate::agents::state_machine::ops_recipe; use crate::agents::state_machine::{ @@ -44,8 +45,7 @@ use crate::agents::state_machine::{ }; use crate::agents::subagent_handler::ForegroundSubagentRunner; use crate::agents::types::{ - SessionConfig, SharedProvider, DEFAULT_ON_FAILURE_TIMEOUT_SECONDS, - DEFAULT_RETRY_TIMEOUT_SECONDS, + SessionConfig, DEFAULT_ON_FAILURE_TIMEOUT_SECONDS, DEFAULT_RETRY_TIMEOUT_SECONDS, }; use crate::agents::AgentEvent; use crate::config::extensions::name_to_key; @@ -226,6 +226,7 @@ pub struct AgentConfig { pub session_name_update_tx: Option>, pub use_login_shell_path: Option, pub is_subagent: bool, + pub providers: Arc, } impl AgentConfig { @@ -250,6 +251,7 @@ impl AgentConfig { session_name_update_tx: None, use_login_shell_path: None, is_subagent: false, + providers: Arc::default(), } } @@ -282,7 +284,6 @@ fn resolve_use_login_shell_path(explicit: Option, platform: &GoosePlatform /// The main goose Agent pub struct Agent { - pub(super) provider: SharedProvider, pub config: AgentConfig, pub(super) current_goose_mode: Mutex, @@ -424,7 +425,7 @@ impl Agent { } pub fn with_config(config: AgentConfig) -> Self { - let provider = Arc::new(Mutex::new(None)); + let providers = config.providers.clone(); let goose_platform = config.goose_platform.clone(); let initial_mode = config.goose_mode; @@ -454,11 +455,10 @@ impl Agent { let use_login_shell_path = config.resolve_use_login_shell_path(); let is_subagent = config.is_subagent; Self { - provider: provider.clone(), config, current_goose_mode: Mutex::new(initial_mode), extension_manager: Arc::new(ExtensionManager::new( - provider.clone(), + providers.clone(), session_manager, scheduler, client_name, @@ -472,7 +472,7 @@ impl Agent { retry_manager: RetryManager::new(), tool_inspection_manager: Self::create_tool_inspection_manager( permission_manager, - provider.clone(), + providers, inspection_session_manager, ), hook_manager: if is_subagent { @@ -795,7 +795,7 @@ impl Agent { /// Create a tool inspection manager with default inspectors fn create_tool_inspection_manager( permission_manager: Arc, - provider: SharedProvider, + providers: Arc, session_manager: Arc, ) -> ToolInspectionManager { let mut tool_inspection_manager = ToolInspectionManager::new(); @@ -806,14 +806,14 @@ impl Agent { // Add adversary inspector (LLM-based review, enabled by ~/.config/goose/adversary.md) tool_inspection_manager.add_inspector(Box::new(AdversaryInspector::new( - provider.clone(), + providers.clone(), session_manager.clone(), ))); // Add permission inspector (medium-high priority) tool_inspection_manager.add_inspector(Box::new(PermissionInspector::new( permission_manager, - provider, + providers, session_manager, ))); @@ -895,7 +895,7 @@ impl Agent { { Ok(v) => v, Err(_) => { - let context_limit = match self.provider().await { + let context_limit = match self.config.providers.provider_for(&session).await { Ok(provider) => crate::context_limit::get_context_limit( provider.as_ref(), &model_config.model_name, @@ -990,44 +990,25 @@ impl Agent { } } - /// Get a reference count clone to the provider - pub async fn provider(&self) -> Result, anyhow::Error> { - match &*self.provider.lock().await { - Some(provider) => Ok(Arc::clone(provider)), - None => Err(anyhow!("Provider not set")), - } + pub async fn provider(&self, session_id: &str) -> Result> { + let session = self + .config + .session_manager + .get_session(session_id, false) + .await?; + self.config.providers.provider_for(&session).await } - /// Resolve the active model config for a session. - /// - /// The session is the source of truth for the selected model and its - /// settings. When the session has no stored config (e.g. before the - /// provider has been persisted), fall back to the configured provider - /// defaults. pub async fn model_config_for_session( &self, session_id: &str, ) -> Result { - if let Ok(session) = self + let session = self .config .session_manager .get_session(session_id, false) - .await - { - if let Some(model_config) = session.model_config { - return Ok(model_config); - } - } - - let config = Config::global(); - let provider_name = config - .get_goose_provider() - .map_err(|_| anyhow!("Could not resolve model config: missing provider"))?; - let model_name = config - .get_goose_model() - .map_err(|_| anyhow!("Could not resolve model config: missing model"))?; - crate::model_config::model_config_from_user_config(&provider_name, &model_name) - .map_err(|e| anyhow!("Could not resolve model config: {e}")) + .await?; + crate::agents::provider_manager::model_config_for(&session) } pub(super) async fn effective_model_config_for_session( @@ -1035,7 +1016,7 @@ impl Agent { session_id: &str, ) -> Result { let model_config = self.model_config_for_session(session_id).await?; - let provider_name = self.provider().await?.get_name().to_string(); + let provider_name = self.provider(session_id).await?.get_name().to_string(); match crate::providers::get_from_registry(&provider_name).await { Ok(entry) => Ok(entry .normalize_model_config(model_config.clone()) @@ -1254,7 +1235,9 @@ impl Agent { }; let manages_own_context = self - .provider() + .config + .providers + .provider_for(session) .await .map(|p| p.manages_own_context()) .unwrap_or(false); @@ -1577,7 +1560,7 @@ impl Agent { permission: permission.clone(), }; if self - .try_route_tool_confirmation_to_provider(request_id, &confirmation) + .try_route_tool_confirmation_to_provider(session_id, request_id, &confirmation) .await { if state.contains_request(request_id) { @@ -1619,11 +1602,11 @@ impl Agent { async fn try_route_tool_confirmation_to_provider( &self, + session_id: &str, request_id: &str, confirmation: &PermissionConfirmation, ) -> bool { - let provider = self.provider.lock().await.clone(); - if let Some(provider) = provider.as_ref() { + if let Ok(provider) = self.provider(session_id).await { if provider.permission_routing() == PermissionRouting::ActionRequired && provider .handle_permission_confirmation(request_id, confirmation) @@ -1642,7 +1625,7 @@ impl Agent { confirmation: PermissionConfirmation, ) { if self - .try_route_tool_confirmation_to_provider(&request_id, &confirmation) + .try_route_tool_confirmation_to_provider(session_id, &request_id, &confirmation) .await { return; @@ -1656,11 +1639,10 @@ impl Agent { } } - pub async fn supports_action_required_permissions(&self) -> bool { - if let Some(provider) = self.provider.lock().await.as_ref() { - return provider.permission_routing() == PermissionRouting::ActionRequired; - } - false + pub async fn supports_action_required_permissions(&self, session_id: &str) -> bool { + self.provider(session_id).await.is_ok_and(|provider| { + provider.permission_routing() == PermissionRouting::ActionRequired + }) } pub(super) async fn create_state_machine( @@ -1850,12 +1832,7 @@ impl Agent { .await?; if !self.config.disable_session_naming { - let provider = self - .provider - .lock() - .await - .clone() - .ok_or_else(|| anyhow!("Provider not set"))?; + let provider = self.provider(&session_config.id).await?; let manager = session_manager.clone(); let tx = self.config.session_name_update_tx.clone(); let id = session_id.clone(); @@ -2060,12 +2037,7 @@ impl Agent { ) -> Result>> { let session_manager = self.config.session_manager.clone(); let session_id = session_config.id.clone(); - let provider = self - .provider - .lock() - .await - .clone() - .ok_or_else(|| anyhow!("Provider not set"))?; + let provider = self.provider(&session_id).await?; let model_config = self.effective_model_config_for_session(&session_id).await?; let context_limit = @@ -2228,27 +2200,6 @@ impl Agent { } } - // Doctor repairs the provider by writing it to the session, not to this agent, and must - // run even when restoring the session's provider is what fails. - let is_doctor = - crate::agents::execute_commands::parse_slash_command(&message_text_for_trace) - .is_some_and(|parsed| parsed.command == "doctor"); - let session = session_manager - .get_session(&session_config.id, false) - .await?; - let live_provider_name = self - .provider - .lock() - .await - .as_ref() - .map(|provider| provider.get_name().to_string()); - if !is_doctor - && session.provider_name.is_some() - && session.provider_name != live_provider_name - { - self.restore_provider_from_session(&session).await?; - } - if use_state_machine { tracing::info!("dispatching reply via experimental state machine"); return self @@ -2428,7 +2379,7 @@ impl Agent { let has_final_output_tool = final_output_tool.is_some(); *self.final_output_tool.lock().await = final_output_tool; if has_final_output_tool { - let provider = self.provider().await?; + let provider = self.provider(&session_config.id).await?; if !provider.supports_builtin_tools() { let provider_name = provider.get_name(); warn!( @@ -2452,7 +2403,7 @@ impl Agent { } let needs_auto_compact = check_if_compaction_needed( - self.provider().await?.as_ref(), + self.provider(&session_config.id).await?.as_ref(), &conversation, None, &session, @@ -2498,7 +2449,7 @@ impl Agent { let compact_model_config = self.model_config_for_session(&session_config.id).await?; match compact_messages( - self.provider().await?.as_ref(), + self.provider(&session_config.id).await?.as_ref(), &compact_model_config, &session_config.id, &conversation_to_compact, @@ -2569,7 +2520,7 @@ impl Agent { self.reset_retry_attempts().await; - let provider = self.provider().await?; + let provider = self.provider(&session_config.id).await?; let provider_name = provider.get_name().to_string(); let saved_provider_session_id = super::latest_provider_session_id(conversation.messages(), &provider_name); @@ -2839,7 +2790,7 @@ impl Agent { } let mut stream = crate::agents::reply_parts::stream_response_from_provider( - self.provider().await?, + self.provider(&session_config.id).await?, model_config.clone(), &session_config.id, &system_prompt, @@ -2859,7 +2810,7 @@ impl Agent { None } else { crate::context_mgmt::maybe_summarize_tool_pairs( - self.provider().await?, + self.provider(&session_config.id).await?, model_config.clone(), session_config.id.clone(), conversation.clone(), @@ -3294,7 +3245,7 @@ impl Agent { ); match compact_messages( - self.provider().await?.as_ref(), + self.provider(&session_config.id).await?.as_ref(), &model_config, &session_config.id, &conversation, @@ -3695,6 +3646,8 @@ impl Agent { Ok(inner) } + /// Pins `provider` to the session instead of letting the provider manager + /// build one. pub async fn update_provider( &self, provider: Arc, @@ -3702,18 +3655,48 @@ impl Agent { session_id: &str, ) -> Result<()> { let provider_name = provider.get_name().to_string(); - let registry_entry = crate::providers::get_from_registry(&provider_name) + self.config + .providers + .set_provider(session_id, provider) + .await; + self.switch_provider(session_id, &provider_name, model_config) + .await + } + + pub async fn switch_provider( + &self, + session_id: &str, + provider_name: &str, + model_config: goose_providers::model::ModelConfig, + ) -> Result<()> { + let mut session = self + .config + .session_manager + .get_session(session_id, false) + .await?; + let registry_entry = crate::providers::get_from_registry(provider_name) .await .ok(); - let model_config = if registry_entry.is_some() { - crate::model_config::materialize_model_config(&provider_name, model_config.clone()) + crate::model_config::materialize_model_config(provider_name, model_config.clone()) .unwrap_or(model_config) } else { model_config }; - let effort_support = provider.thinking_effort_support(); - let model_config = normalize_legacy_provider_thinking_effort(model_config, &effort_support); + + session.provider_name = Some(provider_name.to_string()); + session.model_config = Some(model_config.clone()); + let provider = self + .config + .providers + .provider_for(&session) + .await + .map_err(|error| provider_creation_error(error, "Could not create provider"))?; + + let model_config = normalize_legacy_provider_thinking_effort( + model_config, + &provider.thinking_effort_support(), + ); let effective_model_config = match registry_entry { Some(entry) => entry .normalize_model_config(model_config.clone()) @@ -3721,15 +3704,9 @@ impl Agent { None => model_config.clone(), }; - { - let mut current_provider = self.provider.lock().await; - *current_provider = Some(Arc::clone(&provider)); - } - - // A freshly created provider that manages its own model starts on its - // own default, so the session's selection has to be pushed to it before - // the next config snapshot is built. Failures are not fatal here: the - // selection is re-applied at stream time. + // A provider that manages its own model has to be told about the + // selection before the next config snapshot is built. Failures are not + // fatal here: the selection is re-applied at stream time. if let Err(e) = provider .apply_model_selection(&effective_model_config) .await @@ -3741,7 +3718,7 @@ impl Agent { .session_manager .clone() .update(session_id) - .provider_name(&provider_name) + .provider_name(provider_name) .model_config(model_config) .apply() .await @@ -3749,7 +3726,7 @@ impl Agent { } pub async fn update_goose_mode(&self, mode: GooseMode, session_id: &str) -> Result<()> { - if let Some(provider) = self.provider.lock().await.as_ref() { + if let Ok(provider) = self.provider(session_id).await { provider .update_mode(session_id, mode) .await @@ -3770,44 +3747,11 @@ impl Agent { *self.current_goose_mode.lock().await } - pub async fn recreate_provider_for_session( - &self, - session_id: &str, - provider_name: &str, - model_config: goose_providers::model::ModelConfig, - ) -> Result<()> { - let session = self - .config - .session_manager - .get_session(session_id, false) - .await - .context("Failed to get session")?; - - let extensions = EnabledExtensionsState::extensions_or_default( - Some(&session.extension_data), - Config::global(), - ); - - let provider = crate::providers::create_with_working_dir( - provider_name, - extensions, - session.working_dir.clone(), - ) - .await - .map_err(|error| provider_creation_error(error, "Could not create provider"))?; - - self.update_provider(provider, model_config, session_id) - .await?; - - let mode = self.goose_mode().await; - self.update_goose_mode(mode, session_id).await - } - /// Apply a thinking-effort selection. `effort` is the raw option value: a /// provider that manages effort through a harness has its own vocabulary, /// which is not always a `ThinkingEffort` member. pub async fn update_thinking_effort(&self, session_id: &str, effort: &str) -> Result<()> { - let current_provider = self.provider().await?; + let current_provider = self.provider(session_id).await?; // Context rather than a formatted string: the caller distinguishes a // value rejection from an operational failure by downcasting to // `ProviderError`, which stringifying would destroy. @@ -3819,8 +3763,6 @@ impl Agent { let model_config = self.model_config_for_session(session_id).await?; if provider_handled { - // The provider applied the value live; recreating it would discard - // the very session state we just configured. let model_config = model_config.with_merged_request_params(HashMap::from([( "thinking_effort".to_string(), Value::String(effort.to_string()), @@ -3842,7 +3784,7 @@ impl Agent { ))) })?; let provider_name = current_provider.get_name().to_string(); - self.recreate_provider_for_session( + self.switch_provider( session_id, &provider_name, model_config.with_thinking_effort(effort), @@ -3850,17 +3792,12 @@ impl Agent { .await } - /// Restore the provider from session data or fall back to global config - /// This is used when resuming a session to restore the provider state - /// Returns true if the session's provider was replaced with a fallback. + /// Points the session at its stored provider, falling back to the global + /// provider when the stored one no longer exists. Returns true if it fell + /// back. pub async fn restore_provider_from_session(&self, session: &Session) -> Result { let config = Config::global(); - - let provider_name = session - .provider_name - .clone() - .or_else(|| config.get_goose_provider().ok()) - .ok_or_else(|| anyhow!("Could not configure agent: missing provider"))?; + let provider_name = crate::agents::provider_manager::provider_name_for(session)?; let mut model_config = match session.model_config.clone() { Some(saved_config) => crate::model_config::with_rederived_cache_ttl(saved_config) @@ -3889,92 +3826,55 @@ impl Agent { } } - let extensions = - EnabledExtensionsState::extensions_or_default(Some(&session.extension_data), config); - - let (provider, active_model_config, provider_changed) = - if crate::providers::get_from_registry(&provider_name) - .await - .is_ok() - { - let p = crate::providers::create_with_working_dir( - &provider_name, - extensions, - session.working_dir.clone(), - ) - .await - .map_err(|error| provider_creation_error(error, "Could not create provider"))?; - (p, model_config, false) - } else { - let fallback_provider_name = config - .get_goose_provider() - .ok() - .filter(|name| name != &provider_name) - .ok_or_else(|| { - anyhow!( - "Could not create provider: provider '{}' not found", - provider_name - ) - })?; - - tracing::warn!( - "Session provider '{}' unavailable, falling back to '{}'", - provider_name, - fallback_provider_name - ); - - let fallback_model_name = config.get_goose_model().ok().ok_or_else(|| { - anyhow!("Could not configure fallback provider: missing model") - })?; - let fallback_model_config = crate::model_config::model_config_from_user_config( - &fallback_provider_name, - &fallback_model_name, - ) - .map_err(|e| { - anyhow!("Could not configure fallback provider: invalid model {}", e) - })?; - - let fallback_provider = crate::providers::create_with_working_dir( - &fallback_provider_name, - extensions, - session.working_dir.clone(), - ) - .await - .map_err(|error| { - provider_creation_error( - error, - format!( - "Could not create provider '{provider_name}' or fallback '{fallback_provider_name}'" - ), - ) - })?; - - if let Err(e) = self - .config - .session_manager - .update(&session.id) - .provider_name(&fallback_provider_name) - .model_config(fallback_model_config.clone()) - .apply() - .await - { - tracing::warn!("Failed to update session provider: {}", e); - } - - (fallback_provider, fallback_model_config, true) - }; - - self.update_provider(provider, active_model_config, &session.id) - .await?; - // Propagate session mode to the new provider - if let Some(provider) = self.provider.lock().await.as_ref() { - provider - .update_mode(&session.id, session.goose_mode) - .await - .map_err(|e| anyhow!("Failed to propagate mode to provider: {}", e))?; - } *self.current_goose_mode.lock().await = session.goose_mode; - Ok(provider_changed) + + if crate::providers::get_from_registry(&provider_name) + .await + .is_ok() + { + self.switch_provider(&session.id, &provider_name, model_config) + .await?; + return Ok(false); + } + + let fallback_provider_name = config + .get_goose_provider() + .ok() + .filter(|name| name != &provider_name) + .ok_or_else(|| { + anyhow!( + "Could not create provider: provider '{}' not found", + provider_name + ) + })?; + + tracing::warn!( + "Session provider '{}' unavailable, falling back to '{}'", + provider_name, + fallback_provider_name + ); + + let fallback_model_name = config + .get_goose_model() + .ok() + .ok_or_else(|| anyhow!("Could not configure fallback provider: missing model"))?; + let fallback_model_config = crate::model_config::model_config_from_user_config( + &fallback_provider_name, + &fallback_model_name, + ) + .map_err(|e| anyhow!("Could not configure fallback provider: invalid model {}", e))?; + + self.switch_provider(&session.id, &fallback_provider_name, fallback_model_config) + .await + .map_err(|error| { + provider_creation_error( + error, + format!( + "Could not create provider '{provider_name}' or fallback '{fallback_provider_name}'" + ), + ) + })?; + Ok(true) } pub async fn list_extension_prompts(&self, session_id: &str) -> HashMap> { @@ -4853,8 +4753,14 @@ mod tests { async fn test_submit_tool_confirmation_routes_to_provider() { let (agent, session, _data_dir) = tracing_test_agent_and_session().await; let provider = Arc::new(ActionRequiredProvider::new()); - *agent.provider.lock().await = - Some(provider.clone() as Arc); + agent + .update_provider( + provider.clone(), + goose_providers::model::ModelConfig::new("test-model"), + &session.id, + ) + .await + .unwrap(); // Known request_id → provider handles it, confirmation_router NOT called agent @@ -5178,7 +5084,10 @@ mod tests { Some("xhigh") ); // The unregistered test provider was not respawned. - assert_eq!(agent.provider().await.unwrap().get_name(), "test-effort"); + assert_eq!( + agent.provider(&session_id).await.unwrap().get_name(), + "test-effort" + ); } #[tokio::test] diff --git a/crates/goose/src/agents/execute_commands.rs b/crates/goose/src/agents/execute_commands.rs index 45cdce6290..7edd91faa3 100644 --- a/crates/goose/src/agents/execute_commands.rs +++ b/crates/goose/src/agents/execute_commands.rs @@ -184,7 +184,7 @@ impl Agent { } async fn handle_compact_command(&self, session_id: &str) -> Result> { - let provider = self.provider().await?; + let provider = self.provider(session_id).await?; if provider.manages_own_context() { return Err(anyhow!(context_management_unsupported_message( "compact", @@ -226,7 +226,7 @@ impl Agent { async fn handle_clear_command(&self, session_id: &str) -> Result> { use crate::conversation::Conversation; - let provider = self.provider().await?; + let provider = self.provider(session_id).await?; if provider.manages_own_context() { return Err(anyhow!(context_management_unsupported_message( "clear", @@ -265,7 +265,7 @@ impl Agent { } async fn handle_status_command(&self, session_id: &str) -> Result> { - let provider = self.provider().await?; + let provider = self.provider(session_id).await?; let model_config = self.model_config_for_session(session_id).await?; let context_limit = crate::context_limit::get_context_limit(provider.as_ref(), &model_config.model_name) diff --git a/crates/goose/src/agents/extension_manager/mod.rs b/crates/goose/src/agents/extension_manager/mod.rs index c9d032cd38..02787fa55a 100644 --- a/crates/goose/src/agents/extension_manager/mod.rs +++ b/crates/goose/src/agents/extension_manager/mod.rs @@ -18,11 +18,11 @@ use super::extension::{ ExtensionConfig, ExtensionError, ExtensionResult, PlatformExtensionContext, PLATFORM_EXTENSIONS, }; use super::tool_execution::ToolCallResult; -use super::types::SharedProvider; use crate::action_required_manager::ActionRequiredManager; use crate::agents::mcp_client::{ ConnectContext, GooseMcpClientCapabilities, GooseMcpHostInfo, McpClientTrait, }; +use crate::agents::provider_manager::ProviderManager; use crate::config::extensions::name_to_key; use crate::config::{get_extension_by_name, Config}; use crate::oauth::GooseCredentialStore; @@ -288,7 +288,6 @@ pub struct ExtensionManager { mutation_lock: Mutex<()>, directory_lock: RwLock<()>, context: PlatformExtensionContext, - provider: SharedProvider, client_name: String, capabilities: ExtensionManagerCapabilities, } @@ -501,7 +500,7 @@ impl ExtensionManager { } pub fn new( - provider: SharedProvider, + providers: Arc, session_manager: Arc, scheduler: Option>, client_name: String, @@ -514,21 +513,20 @@ impl ExtensionManager { directory_lock: RwLock::new(()), context: PlatformExtensionContext { extension_manager: None, - provider: provider.clone(), + providers, session_manager, scheduler, use_login_shell_path, }, - provider, client_name, capabilities, } } - pub fn new_without_provider(data_dir: std::path::PathBuf) -> Self { + pub fn with_data_dir(data_dir: std::path::PathBuf) -> Self { let session_manager = Arc::new(crate::session::SessionManager::new(data_dir)); Self::new( - Arc::new(Mutex::new(None)), + Default::default(), session_manager, None, "goose-cli".to_string(), @@ -546,10 +544,6 @@ impl ExtensionManager { &self.context } - pub fn get_provider(&self) -> &SharedProvider { - &self.provider - } - fn hydrate_mcp_apps(&self) -> bool { match &self.capabilities.host_info { Some(host_info) if host_info.explicit_extensions => host_info.mcpui_enabled(), @@ -1256,8 +1250,7 @@ mod tests { async fn dispatch_notification_methods(ctx: ToolCallContext) -> Vec { let temp_dir = tempfile::tempdir().unwrap(); - let extension_manager = - ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); + let extension_manager = ExtensionManager::with_data_dir(temp_dir.path().to_path_buf()); extension_manager .add_mock_extension( "notifications".to_string(), @@ -1312,8 +1305,7 @@ mod tests { #[tokio::test] async fn dispatch_reuses_existing_notification_emitter() { let temp_dir = tempfile::tempdir().unwrap(); - let extension_manager = - ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); + let extension_manager = ExtensionManager::with_data_dir(temp_dir.path().to_path_buf()); extension_manager .add_mock_extension( "notifications".to_string(), @@ -1453,8 +1445,7 @@ mod tests { #[tokio::test] async fn app_dispatch_rejects_colliding_flattened_name_from_sibling_owner() { let temp_dir = tempfile::tempdir().unwrap(); - let extension_manager = - ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); + let extension_manager = ExtensionManager::with_data_dir(temp_dir.path().to_path_buf()); extension_manager .add_mock_extension( "ext_a__ext_b".to_string(), @@ -1561,8 +1552,7 @@ mod tests { #[tokio::test] async fn resolve_leaves_out_an_extension_running_under_a_different_config() { let temp_dir = tempfile::tempdir().unwrap(); - let extension_manager = - ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); + let extension_manager = ExtensionManager::with_data_dir(temp_dir.path().to_path_buf()); extension_manager .add_mock_extension("ext_a".to_string(), Arc::new(MockClient {})) .await; @@ -1586,7 +1576,7 @@ mod tests { let data_dir = tempfile::tempdir().unwrap(); let old_working_dir = tempfile::tempdir().unwrap(); let new_working_dir = tempfile::tempdir().unwrap(); - let manager = Arc::new(ExtensionManager::new_without_provider( + let manager = Arc::new(ExtensionManager::with_data_dir( data_dir.path().to_path_buf(), )); let session = manager @@ -1647,7 +1637,7 @@ mod tests { async fn stale_working_dir_reconnect_does_not_restore_removed_extension() { let data_dir = tempfile::tempdir().unwrap(); let new_working_dir = tempfile::tempdir().unwrap(); - let extension_manager = Arc::new(ExtensionManager::new_without_provider( + let extension_manager = Arc::new(ExtensionManager::with_data_dir( data_dir.path().to_path_buf(), )); let config = ExtensionConfig::Platform { @@ -1703,7 +1693,7 @@ mod tests { #[tokio::test] async fn extension_manager_tools_follow_resource_support() { let temp_dir = tempfile::tempdir().unwrap(); - let extension_manager = Arc::new(ExtensionManager::new_without_provider( + let extension_manager = Arc::new(ExtensionManager::with_data_dir( temp_dir.path().to_path_buf(), )); extension_manager @@ -1770,7 +1760,7 @@ mod tests { #[tokio::test] async fn extension_manager_resource_tools_use_the_calling_lease() { let temp_dir = tempfile::tempdir().unwrap(); - let extension_manager = Arc::new(ExtensionManager::new_without_provider( + let extension_manager = Arc::new(ExtensionManager::with_data_dir( temp_dir.path().to_path_buf(), )); extension_manager @@ -1859,8 +1849,7 @@ mod tests { #[tokio::test] async fn tool_list_changed_during_fetch_prevents_stale_cache() { let temp_dir = tempfile::tempdir().unwrap(); - let extension_manager = - ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); + let extension_manager = ExtensionManager::with_data_dir(temp_dir.path().to_path_buf()); let tools_client = Arc::new(BlockingToolsClient { calls: AtomicUsize::new(0), first_fetch_started: Semaphore::new(0), @@ -1904,7 +1893,7 @@ mod tests { #[tokio::test] async fn successful_mutation_is_persisted_when_another_mutation_fails() { let temp_dir = tempfile::tempdir().unwrap(); - let manager = Arc::new(ExtensionManager::new_without_provider( + let manager = Arc::new(ExtensionManager::with_data_dir( temp_dir.path().to_path_buf(), )); let session = manager diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index 65ac50dedf..f5f01efb05 100644 --- a/crates/goose/src/agents/mod.rs +++ b/crates/goose/src/agents/mod.rs @@ -13,6 +13,7 @@ pub mod platform_extensions; #[cfg(feature = "scheduler")] pub mod platform_tools; pub mod prompt_manager; +pub mod provider_manager; pub mod reply_parts; pub mod retry; #[cfg(feature = "scheduler")] diff --git a/crates/goose/src/agents/moim.rs b/crates/goose/src/agents/moim.rs index 7496e3d9b8..7b6b47dfe1 100644 --- a/crates/goose/src/agents/moim.rs +++ b/crates/goose/src/agents/moim.rs @@ -47,18 +47,24 @@ pub(super) async fn compute_compaction_info( let session_model_config = session .as_ref() .and_then(|session| session.model_config.clone()); - let context_limit = if let Some(model_config) = session_model_config.as_ref() { - let provider = extension_manager.get_provider().lock().await.clone(); - match provider { - Some(provider) => { - crate::context_limit::get_context_limit(provider.as_ref(), &model_config.model_name) - .await - .ok() + let context_limit = match (session.as_ref(), session_model_config.as_ref()) { + (Some(session), Some(model_config)) => { + match extension_manager + .get_context() + .providers + .provider_for(session) + .await + { + Ok(provider) => crate::context_limit::get_context_limit( + provider.as_ref(), + &model_config.model_name, + ) + .await + .ok(), + Err(_) => None, } - None => None, } - } else { - None + _ => None, }; let total_tokens = session .as_ref() @@ -91,18 +97,24 @@ pub async fn turn_context_message( let session_model_config = session .as_ref() .and_then(|session| session.model_config.clone()); - let context_limit = if let Some(model_config) = session_model_config.as_ref() { - let provider = extension_manager.get_provider().lock().await.clone(); - match provider { - Some(provider) => { - crate::context_limit::get_context_limit(provider.as_ref(), &model_config.model_name) - .await - .ok() + let context_limit = match (session.as_ref(), session_model_config.as_ref()) { + (Some(session), Some(model_config)) => { + match extension_manager + .get_context() + .providers + .provider_for(session) + .await + { + Ok(provider) => crate::context_limit::get_context_limit( + provider.as_ref(), + &model_config.model_name, + ) + .await + .ok(), + Err(_) => None, } - None => None, } - } else { - None + _ => None, }; if should_skip_moim(context_limit) { return None; @@ -278,7 +290,7 @@ mod tests { async fn session_and_manager() -> (String, ExtensionManager, tempfile::TempDir) { let temp_dir = tempfile::tempdir().unwrap(); - let em = ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); + let em = ExtensionManager::with_data_dir(temp_dir.path().to_path_buf()); let session = em .get_context() .session_manager diff --git a/crates/goose/src/agents/platform_extensions/analyze/mod.rs b/crates/goose/src/agents/platform_extensions/analyze/mod.rs index d9d204d2f7..8b3a798980 100644 --- a/crates/goose/src/agents/platform_extensions/analyze/mod.rs +++ b/crates/goose/src/agents/platform_extensions/analyze/mod.rs @@ -278,7 +278,7 @@ mod tests { fn ctx() -> PlatformExtensionContext { PlatformExtensionContext { extension_manager: None, - provider: Arc::new(tokio::sync::Mutex::new(None)), + providers: Default::default(), session_manager: Arc::new(SessionManager::new(std::env::temp_dir())), scheduler: None, use_login_shell_path: false, diff --git a/crates/goose/src/agents/platform_extensions/apps.rs b/crates/goose/src/agents/platform_extensions/apps.rs index e4312df2d7..663aabcc50 100644 --- a/crates/goose/src/agents/platform_extensions/apps.rs +++ b/crates/goose/src/agents/platform_extensions/apps.rs @@ -262,19 +262,6 @@ impl AppsManagerClient { .result_with_platform_notification(result, EXTENSION_NAME, event_type, params) } - async fn get_provider(&self) -> Result, String> { - let provider = self - .context - .provider - .lock() - .await - .as_ref() - .ok_or("Provider not available")? - .clone(); - - Ok(provider) - } - async fn effective_model_config( &self, session_id: &str, @@ -320,7 +307,7 @@ impl AppsManagerClient { session_id: &str, prd: &str, ) -> Result { - let provider = self.get_provider().await?; + let provider = self.context.provider_for_session(session_id).await?; let existing_apps = self.list_stored_apps().unwrap_or_default(); let existing_names = existing_apps.join(", "); @@ -371,7 +358,7 @@ impl AppsManagerClient { existing_prd: &str, feedback: &str, ) -> Result { - let provider = self.get_provider().await?; + let provider = self.context.provider_for_session(session_id).await?; let context: HashMap<&str, &str> = HashMap::new(); let system_prompt = render_template("apps_iterate.md", &context) @@ -828,7 +815,7 @@ mod tests { info: AppsManagerClient::create_info(), context: PlatformExtensionContext { extension_manager: None, - provider: Arc::new(tokio::sync::Mutex::new(None)), + providers: Default::default(), session_manager: Arc::new(SessionManager::new(apps_dir.join("sessions"))), scheduler: None, use_login_shell_path: false, diff --git a/crates/goose/src/agents/platform_extensions/chatrecall.rs b/crates/goose/src/agents/platform_extensions/chatrecall.rs index 8c12be01e0..e4c359c5fb 100644 --- a/crates/goose/src/agents/platform_extensions/chatrecall.rs +++ b/crates/goose/src/agents/platform_extensions/chatrecall.rs @@ -446,7 +446,7 @@ mod tests { let client = ChatRecallClient::new(PlatformExtensionContext { extension_manager: None, - provider: Arc::new(tokio::sync::Mutex::new(None)), + providers: Default::default(), session_manager, scheduler: None, use_login_shell_path: false, diff --git a/crates/goose/src/agents/platform_extensions/code_execution.rs b/crates/goose/src/agents/platform_extensions/code_execution.rs index eddf452d3c..3131e4861e 100644 --- a/crates/goose/src/agents/platform_extensions/code_execution.rs +++ b/crates/goose/src/agents/platform_extensions/code_execution.rs @@ -801,9 +801,7 @@ mod tests { #[tokio::test] async fn callback_configs_exclude_tools_hidden_from_model() { let temp = tempfile::tempdir().unwrap(); - let manager = Arc::new(ExtensionManager::new_without_provider( - temp.path().join("manager"), - )); + let manager = Arc::new(ExtensionManager::with_data_dir(temp.path().join("manager"))); manager .add_client( ExtensionConfig::Builtin { @@ -836,9 +834,7 @@ mod tests { #[tokio::test] async fn callback_configs_leave_delegate_to_direct_tool_calls() { let temp = tempfile::tempdir().unwrap(); - let manager = Arc::new(ExtensionManager::new_without_provider( - temp.path().join("manager"), - )); + let manager = Arc::new(ExtensionManager::with_data_dir(temp.path().join("manager"))); let session = manager .get_context() .session_manager @@ -881,9 +877,7 @@ mod tests { #[tokio::test] async fn callback_uses_lease_and_applies_extension_mutation() { let temp = tempfile::tempdir().unwrap(); - let manager = Arc::new(ExtensionManager::new_without_provider( - temp.path().join("manager"), - )); + let manager = Arc::new(ExtensionManager::with_data_dir(temp.path().join("manager"))); let session = manager .get_context() .session_manager @@ -1173,9 +1167,7 @@ mod tests { #[tokio::test] async fn moim_uses_the_supplied_tool_snapshot() { let temp = tempfile::tempdir().unwrap(); - let manager = Arc::new(ExtensionManager::new_without_provider( - temp.path().join("manager"), - )); + let manager = Arc::new(ExtensionManager::with_data_dir(temp.path().join("manager"))); let mut context = manager.get_context().clone(); context.extension_manager = Some(Arc::downgrade(&manager)); let client = CodeExecutionClient::new(context, ToolDisclosure::Catalog).unwrap(); @@ -1198,7 +1190,7 @@ mod tests { let client = CodeExecutionClient::new( PlatformExtensionContext { extension_manager: None, - provider: Arc::new(tokio::sync::Mutex::new(None)), + providers: Default::default(), session_manager: Arc::new(crate::session::SessionManager::new( temp.path().join("sessions"), )), diff --git a/crates/goose/src/agents/platform_extensions/developer/mod.rs b/crates/goose/src/agents/platform_extensions/developer/mod.rs index c09baa5217..e6ef5412c9 100644 --- a/crates/goose/src/agents/platform_extensions/developer/mod.rs +++ b/crates/goose/src/agents/platform_extensions/developer/mod.rs @@ -294,7 +294,7 @@ mod tests { fn test_context(data_dir: std::path::PathBuf) -> PlatformExtensionContext { PlatformExtensionContext { extension_manager: None, - provider: Arc::new(tokio::sync::Mutex::new(None)), + providers: Default::default(), session_manager: Arc::new(SessionManager::new(data_dir)), scheduler: None, use_login_shell_path: false, diff --git a/crates/goose/src/agents/platform_extensions/ext_manager.rs b/crates/goose/src/agents/platform_extensions/ext_manager.rs index 0ff28bf8ab..81a597af6e 100644 --- a/crates/goose/src/agents/platform_extensions/ext_manager.rs +++ b/crates/goose/src/agents/platform_extensions/ext_manager.rs @@ -512,7 +512,7 @@ mod tests { fn client_for(manager: &Arc) -> ExtensionManagerClient { ExtensionManagerClient::new(PlatformExtensionContext { extension_manager: None, - provider: manager.get_provider().clone(), + providers: manager.get_context().providers.clone(), session_manager: manager.get_context().session_manager.clone(), scheduler: None, use_login_shell_path: false, @@ -568,7 +568,7 @@ mod tests { #[tokio::test] async fn manage_extensions_emits_a_mutation_and_refuses_subagents() { let temp_dir = tempfile::tempdir().unwrap(); - let manager = Arc::new(ExtensionManager::new_without_provider( + let manager = Arc::new(ExtensionManager::with_data_dir( temp_dir.path().to_path_buf(), )); let client = client_for(&manager); @@ -608,7 +608,7 @@ mod tests { #[tokio::test] async fn extension_manager_cannot_disable_itself() { let temp_dir = tempfile::tempdir().unwrap(); - let manager = Arc::new(ExtensionManager::new_without_provider( + let manager = Arc::new(ExtensionManager::with_data_dir( temp_dir.path().to_path_buf(), )); let client = client_for(&manager); @@ -635,7 +635,7 @@ mod tests { #[tokio::test] async fn subagent_and_unknown_callers_are_not_offered_extension_management() { let temp_dir = tempfile::tempdir().unwrap(); - let manager = Arc::new(ExtensionManager::new_without_provider( + let manager = Arc::new(ExtensionManager::with_data_dir( temp_dir.path().to_path_buf(), )); let client = client_for(&manager); diff --git a/crates/goose/src/agents/platform_extensions/mod.rs b/crates/goose/src/agents/platform_extensions/mod.rs index 78fc99144f..a685c75045 100644 --- a/crates/goose/src/agents/platform_extensions/mod.rs +++ b/crates/goose/src/agents/platform_extensions/mod.rs @@ -223,13 +223,28 @@ pub struct PlatformExtensionContext { /// below. pub extension_manager: Option>, - pub provider: crate::agents::types::SharedProvider, + pub providers: std::sync::Arc, pub session_manager: std::sync::Arc, pub scheduler: Option>, pub use_login_shell_path: bool, } impl PlatformExtensionContext { + pub async fn provider_for_session( + &self, + session_id: &str, + ) -> Result, String> { + let session = self + .session_manager + .get_session(session_id, false) + .await + .map_err(|e| e.to_string())?; + self.providers + .provider_for(&session) + .await + .map_err(|e| e.to_string()) + } + pub async fn model_config_for_session( &self, session_id: &str, diff --git a/crates/goose/src/agents/platform_extensions/summarize.rs b/crates/goose/src/agents/platform_extensions/summarize.rs index 1d16f25e02..f6e025c87e 100644 --- a/crates/goose/src/agents/platform_extensions/summarize.rs +++ b/crates/goose/src/agents/platform_extensions/summarize.rs @@ -48,19 +48,6 @@ impl SummarizeClient { Ok(Self { info, context }) } - async fn get_provider(&self) -> Result, String> { - let provider = self - .context - .provider - .lock() - .await - .as_ref() - .ok_or("Provider not available")? - .clone(); - - Ok(provider) - } - fn get_tools() -> Vec { let schema = schema_for!(SummarizeParams); let schema_value = @@ -134,7 +121,7 @@ impl McpClientTrait for SummarizeClient { )])); } - let provider = match self.get_provider().await { + let provider = match self.context.provider_for_session(&ctx.session_id).await { Ok(p) => p, Err(e) => { return Ok(CallToolResult::error(vec![ContentBlock::text(format!( diff --git a/crates/goose/src/agents/platform_extensions/summon.rs b/crates/goose/src/agents/platform_extensions/summon.rs index 151a88f5d2..80b4db3942 100644 --- a/crates/goose/src/agents/platform_extensions/summon.rs +++ b/crates/goose/src/agents/platform_extensions/summon.rs @@ -1273,8 +1273,8 @@ impl SummonClient { .await?; let provider = match providers::get_from_registry(&config.provider_name).await { Ok(entry) => entry.create(config.extensions.clone()).await?, - Err(error) => match self.context.provider.lock().await.clone() { - Some(provider) + Err(error) => match self.context.providers.provider_for(session).await { + Ok(provider) if provider.get_name() == config.provider_name && !provider.manages_own_context() => { @@ -1640,7 +1640,7 @@ mod tests { ) -> PlatformExtensionContext { PlatformExtensionContext { extension_manager: None, - provider: Arc::new(tokio::sync::Mutex::new(None)), + providers: Default::default(), session_manager, scheduler: None, use_login_shell_path: false, @@ -1782,7 +1782,7 @@ mod tests { .await .unwrap(); assert_eq!( - reloaded_agent.provider().await.unwrap().get_name(), + reloaded_agent.provider(&child.id).await.unwrap().get_name(), "openai" ); } @@ -2462,15 +2462,19 @@ You review code."#; ) .unwrap(), ); - let mut context = create_test_context(); - context.provider = Arc::new(tokio::sync::Mutex::new(Some(Arc::clone(&parent_provider)))); + let context = create_test_context(); + let providers = context.providers.clone(); let client = SummonClient::new(context).unwrap(); let session = crate::session::Session { + id: "unregistered-parent".to_string(), provider_name: Some(parent_provider.get_name().to_string()), model_config: Some(goose_providers::model::ModelConfig::new("test-model")), working_dir: temp_dir.path().to_path_buf(), ..Default::default() }; + providers + .set_provider(&session.id, Arc::clone(&parent_provider)) + .await; let params = DelegateParams { instructions: Some("Review the change".to_string()), @@ -2497,9 +2501,7 @@ You review code."#; async fn test_build_task_config_recreates_registered_parent_provider() { let temp_dir = TempDir::new().unwrap(); let parent_provider = providers::create("openai", Vec::new()).await.unwrap(); - let mut context = create_test_context(); - context.provider = Arc::new(tokio::sync::Mutex::new(Some(Arc::clone(&parent_provider)))); - let client = SummonClient::new(context).unwrap(); + let client = SummonClient::new(create_test_context()).unwrap(); let session = crate::session::Session { provider_name: Some(parent_provider.get_name().to_string()), model_config: Some(goose_providers::model::ModelConfig::new("test-model")), diff --git a/crates/goose/src/agents/prompt_manager.rs b/crates/goose/src/agents/prompt_manager.rs index 82b662601c..bb8e5afcd3 100644 --- a/crates/goose/src/agents/prompt_manager.rs +++ b/crates/goose/src/agents/prompt_manager.rs @@ -518,7 +518,7 @@ mod tests { .unwrap(); let context = PlatformExtensionContext { extension_manager: None, - provider: Arc::new(tokio::sync::Mutex::new(None)), + providers: Default::default(), session_manager, scheduler: Some(scheduler), use_login_shell_path: false, diff --git a/crates/goose/src/agents/provider_manager.rs b/crates/goose/src/agents/provider_manager.rs new file mode 100644 index 0000000000..8b24f76112 --- /dev/null +++ b/crates/goose/src/agents/provider_manager.rs @@ -0,0 +1,109 @@ +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use anyhow::{anyhow, Result}; +use goose_providers::model::ModelConfig; +use tracing::warn; + +use crate::config::Config; +use crate::providers::base::Provider; +use crate::providers::provider_registry::ProviderEntry; +use crate::session::extension_data::EnabledExtensionsState; +use crate::session::Session; + +type SessionSlot = Arc>>>; +type SharedProviders = HashMap)>; + +#[derive(Default)] +pub struct ProviderManager { + shared: tokio::sync::Mutex, + sessions: Mutex>, +} + +impl ProviderManager { + pub async fn provider_for(&self, session: &Session) -> Result> { + let name = provider_name_for(session)?; + let slot = self.slot(&session.id); + let mut slot = slot.lock().await; + if let Some(provider) = slot.as_ref().filter(|p| p.get_name() == name) { + return Ok(provider.clone()); + } + + let entry = crate::providers::get_from_registry(&name).await?; + if !entry.session_bound() { + let provider = self.shared(&name, &entry).await?; + *slot = None; + return Ok(provider); + } + + let extensions = EnabledExtensionsState::extensions_or_default( + Some(&session.extension_data), + Config::global(), + ); + let provider = entry + .create_with_working_dir(extensions, session.working_dir.clone()) + .await?; + let model_config = entry.normalize_model_config(model_config_for(session)?)?; + if let Err(e) = provider.apply_model_selection(&model_config).await { + warn!("Failed to apply model selection to provider: {e}"); + } + provider + .update_mode(&session.id, session.goose_mode) + .await + .map_err(|e| anyhow!("Failed to propagate mode to provider: {e}"))?; + *slot = Some(provider.clone()); + Ok(provider) + } + + pub async fn set_provider(&self, session_id: &str, provider: Arc) { + *self.slot(session_id).lock().await = Some(provider); + } + + pub fn release(&self, session_id: &str) { + self.sessions.lock().unwrap().remove(session_id); + } + + fn slot(&self, session_id: &str) -> SessionSlot { + self.sessions + .lock() + .unwrap() + .entry(session_id.to_string()) + .or_default() + .clone() + } + + async fn shared(&self, name: &str, entry: &ProviderEntry) -> Result> { + let generation = Config::global().generation(); + let mut shared = self.shared.lock().await; + if let Some((cached_generation, provider)) = shared.get(name) { + if *cached_generation == generation { + return Ok(provider.clone()); + } + } + let provider = entry.create(Vec::new()).await?; + shared.insert(name.to_string(), (generation, provider.clone())); + Ok(provider) + } +} + +pub fn provider_name_for(session: &Session) -> Result { + match &session.provider_name { + Some(name) => Ok(name.clone()), + None => Config::global() + .get_goose_provider() + .map_err(|_| anyhow!("Could not configure agent: missing provider")), + } +} + +pub fn model_config_for(session: &Session) -> Result { + if let Some(model_config) = &session.model_config { + return Ok(model_config.clone()); + } + let config = Config::global(); + let provider_name = provider_name_for(session)?; + let model_name = config + .get_goose_model() + .map_err(|_| anyhow!("Could not resolve model config: missing model"))?; + crate::model_config::model_config_from_user_config(&provider_name, &model_name) + .map_err(|e| anyhow!("Could not resolve model config: {e}")) +} diff --git a/crates/goose/src/agents/state_machine/tests/pipeline.rs b/crates/goose/src/agents/state_machine/tests/pipeline.rs index fd62923c4d..f395ea87cb 100644 --- a/crates/goose/src/agents/state_machine/tests/pipeline.rs +++ b/crates/goose/src/agents/state_machine/tests/pipeline.rs @@ -781,9 +781,10 @@ async fn build_test_pipeline( } else { provider }; - let shared_provider = Arc::new(TokioMutex::new(Some(provider.clone()))); + let providers: Arc = Default::default(); + providers.set_provider(&session.id, provider.clone()).await; let extension_manager = Arc::new(ExtensionManager::new( - shared_provider.clone(), + providers.clone(), session_manager.clone(), scheduler .clone() @@ -802,7 +803,7 @@ async fn build_test_pipeline( tool_inspection_manager.add_inspector(Box::new(SecurityInspector::enabled())); tool_inspection_manager.add_inspector(Box::new(PermissionInspector::new( permission_manager.clone(), - shared_provider, + providers, session_manager.clone(), ))); let model_config = session diff --git a/crates/goose/src/agents/subagent_handler.rs b/crates/goose/src/agents/subagent_handler.rs index 74836bd441..c359a6d937 100644 --- a/crates/goose/src/agents/subagent_handler.rs +++ b/crates/goose/src/agents/subagent_handler.rs @@ -100,14 +100,6 @@ pub(crate) async fn from_foreground_subagent_session( .ok_or_else(|| anyhow!("Subagent {session_id} has no saved extension selection"))?; let extensions = EnabledExtensionsState::from_value(saved_extensions)?.extensions; - let provider = crate::providers::create_with_working_dir( - provider_name, - extensions.clone(), - session.working_dir.clone(), - ) - .await?; - provider.apply_model_selection(model_config).await?; - let mut config = AgentConfig::new( session_manager, PermissionManager::instance(), @@ -119,7 +111,9 @@ pub(crate) async fn from_foreground_subagent_session( .with_use_login_shell_path(use_login_shell_path); config.is_subagent = true; let agent = Agent::with_config(config); - *agent.provider.lock().await = Some(provider); + agent + .switch_provider(session_id, provider_name, model_config.clone()) + .await?; for extension in extensions { let name = extension.name(); if let Err(e) = agent.add_extension_inner(extension, session_id).await { diff --git a/crates/goose/src/agents/types.rs b/crates/goose/src/agents/types.rs index 684b7c290a..998079fdf0 100644 --- a/crates/goose/src/agents/types.rs +++ b/crates/goose/src/agents/types.rs @@ -1,9 +1,4 @@ -use crate::providers::base::Provider; use serde::{Deserialize, Serialize}; -use std::sync::Arc; -use tokio::sync::Mutex; - -pub type SharedProvider = Arc>>>; /// Default timeout for retry operations (5 minutes) pub const DEFAULT_RETRY_TIMEOUT_SECONDS: u64 = 300; diff --git a/crates/goose/src/bin/build_canonical_models.rs b/crates/goose/src/bin/build_canonical_models.rs index 9445d60921..d7ff1d8f60 100644 --- a/crates/goose/src/bin/build_canonical_models.rs +++ b/crates/goose/src/bin/build_canonical_models.rs @@ -9,7 +9,7 @@ /// use anyhow::{Context, Result}; use clap::Parser; -use goose::providers::create_with_named_model; +use goose::providers::create; use goose_providers::canonical::ModelMapping; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -433,7 +433,7 @@ async fn check_provider( ) -> Result<(Vec, Vec, Vec)> { println!("Checking provider: {}", provider_name); - let provider = match create_with_named_model(provider_name, Vec::new()).await { + let provider = match create(provider_name, Vec::new()).await { Ok(p) => p, Err(e) => { println!(" ⚠ Failed to create provider: {}", e); diff --git a/crates/goose/src/config/base.rs b/crates/goose/src/config/base.rs index ba53968b69..4b914b313e 100644 --- a/crates/goose/src/config/base.rs +++ b/crates/goose/src/config/base.rs @@ -14,6 +14,7 @@ use std::env; use std::fs::OpenOptions; use std::io::Write; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use thiserror::Error; @@ -124,6 +125,7 @@ pub struct Config { secrets: SecretStorage, guard: Mutex<()>, secrets_cache: Arc>>>, + generation: AtomicU64, } enum SecretStorage { @@ -211,6 +213,7 @@ impl Default for Config { }, guard: Mutex::new(()), secrets_cache: Arc::new(Mutex::new(None)), + generation: AtomicU64::new(0), }; let keyring_disabled = env::var("GOOSE_DISABLE_KEYRING").is_ok() @@ -223,6 +226,7 @@ impl Default for Config { secrets, guard: Mutex::new(()), secrets_cache: Arc::new(Mutex::new(None)), + generation: AtomicU64::new(0), } } } @@ -441,6 +445,7 @@ impl Config { secrets, guard: Mutex::new(()), secrets_cache: Arc::new(Mutex::new(None)), + generation: AtomicU64::new(0), }) } @@ -459,6 +464,7 @@ impl Config { }, guard: Mutex::new(()), secrets_cache: Arc::new(Mutex::new(None)), + generation: AtomicU64::new(0), }) } @@ -473,6 +479,7 @@ impl Config { }, guard: Mutex::new(()), secrets_cache: Arc::new(Mutex::new(None)), + generation: AtomicU64::new(0), }) } @@ -663,6 +670,7 @@ impl Config { // Atomically replace the original file std::fs::rename(&temp_path, &target_path)?; + self.generation.fetch_add(1, Ordering::Relaxed); Ok(()) } @@ -1186,6 +1194,12 @@ impl Config { pub fn invalidate_secrets_cache(&self) { let mut cache = self.secrets_cache.lock().unwrap(); *cache = None; + self.generation.fetch_add(1, Ordering::Relaxed); + } + + /// Changes whenever this process writes config or secrets. + pub fn generation(&self) -> u64 { + self.generation.load(Ordering::Relaxed) } /// Check if an error string indicates a keyring availability issue that should trigger fallback diff --git a/crates/goose/src/execution/manager.rs b/crates/goose/src/execution/manager.rs index a2b9fe2c20..d2d583ff06 100644 --- a/crates/goose/src/execution/manager.rs +++ b/crates/goose/src/execution/manager.rs @@ -11,7 +11,7 @@ use std::num::NonZeroUsize; use std::sync::Arc; use tokio::sync::{mpsc, Mutex, OnceCell, RwLock}; use tokio_util::sync::CancellationToken; -use tracing::{debug, info}; +use tracing::info; const DEFAULT_MAX_SESSION: usize = 100; @@ -33,7 +33,6 @@ pub struct AgentManagerGetResult { pub struct AgentManager { sessions: Arc>>>, agent_config: AgentConfig, - default_provider: Arc>>>, cancel_tokens: Arc>>, /// Per-session creation locks. When `get_or_create_agent` misses the /// `sessions` cache it acquires the per-session lock before doing the @@ -54,7 +53,6 @@ impl AgentManager { let manager = Self { sessions: Arc::new(RwLock::new(LruCache::new(capacity))), agent_config, - default_provider: Arc::new(RwLock::new(None)), cancel_tokens: Arc::new(RwLock::new(HashMap::new())), creation_locks: Arc::new(Mutex::new(HashMap::new())), }; @@ -95,11 +93,6 @@ impl AgentManager { self.agent_config.session_manager.as_ref() } - pub async fn set_default_provider(&self, provider: Arc) { - debug!("Setting default provider on AgentManager"); - *self.default_provider.write().await = Some(provider); - } - pub async fn get_or_create_agent(&self, session_id: String) -> Result> { Ok(self .get_or_create_agent_with_runtime_context(session_id, RuntimeContext::default()) @@ -141,8 +134,8 @@ impl AgentManager { // Funnel the fallible work through a helper so we can prune the // per-session creation lock on every error exit. Without this - // the provider-setup path (update_provider / update_mode) could - // bail out via `?`, leaving a permanent `creation_locks` entry + // provider restore or recipe setup could bail out via `?`, + // leaving a permanent `creation_locks` entry // for a session that never made it into the LRU cache and that // no one will ever call `remove_session` on. let result = self.create_agent_locked(&session_id, runtime_context).await; @@ -228,31 +221,6 @@ impl AgentManager { extension_results = agent.load_extensions_from_session(&session).await; } - if agent.provider().await.is_err() { - if let Some(provider) = &*self.default_provider.read().await { - let config = crate::config::Config::global(); - let model_config = config - .get_goose_provider() - .ok() - .zip(config.get_goose_model().ok()) - .and_then(|(provider_name, model_name)| { - crate::model_config::model_config_from_user_config( - &provider_name, - &model_name, - ) - .ok() - }) - .unwrap_or_else(|| goose_providers::model::ModelConfig::new("unknown")); - agent - .update_provider(Arc::clone(provider), model_config, session_id) - .await?; - provider - .update_mode(session_id, mode) - .await - .map_err(|e| anyhow::anyhow!("Failed to propagate mode to provider: {}", e))?; - } - } - let mut sessions = self.sessions.write().await; if let Some(existing) = sessions.get(session_id) { return Ok(AgentManagerGetResult { @@ -266,12 +234,15 @@ impl AgentManager { // key so we can also drop its creation lock below, otherwise the // `creation_locks` HashMap would grow without bound in long-lived // processes that churn through many sessions. - let evicted = sessions - .push(session_id.to_string(), agent.clone()) - .map(|(k, _)| k); + let evicted = sessions.push(session_id.to_string(), agent.clone()); drop(sessions); - if let Some(evicted_id) = evicted { + if let Some((evicted_id, evicted_agent)) = evicted { + // A run may still hold the evicted agent; its session-bound + // provider has to outlive that run. + if Arc::strong_count(&evicted_agent) == 1 { + self.agent_config.providers.release(&evicted_id); + } self.prune_creation_lock(&evicted_id).await; } @@ -311,6 +282,7 @@ impl AgentManager { .pop(session_id) .ok_or_else(|| anyhow::anyhow!("Session {} not found", session_id))?; drop(sessions); + self.agent_config.providers.release(session_id); // Best-effort prune of the per-session creation lock so the // HashMap doesn't grow unbounded. Any caller still holding a // clone of the Arc keeps the underlying Mutex alive until it @@ -325,11 +297,11 @@ impl AgentManager { if let Some(token) = self.cancel_tokens.write().await.remove(session_id) { token.cancel(); } - let mut sessions = self.sessions.write().await; - if sessions.pop(session_id).is_none() { + let removed = self.sessions.write().await.pop(session_id); + self.agent_config.providers.release(session_id); + if removed.is_none() { return Ok(()); } - drop(sessions); self.prune_creation_lock(session_id).await; info!("Removed session {}", session_id); Ok(()) @@ -625,78 +597,6 @@ mod tests { ); } - #[tokio::test] - async fn test_failed_creation_prunes_creation_lock() { - // Regression test for the Codex review note on PR #9357: when the - // provider-setup path in `create_agent_locked` returns Err, the - // outer `get_or_create_agent` must also drop its local Arc clone - // of the creation lock before pruning. Otherwise - // `Arc::strong_count` stays > 1 and the failed session leaks a - // permanent entry in `creation_locks`. - use async_trait::async_trait; - use rmcp::model::Tool; - - use crate::conversation::message::Message; - use crate::providers::base::{MessageStream, Provider}; - use goose_providers::conversation::token_usage::{ProviderUsage, Usage}; - use goose_providers::errors::ProviderError; - use goose_providers::model::ModelConfig; - - struct FailingProvider; - - #[async_trait] - impl Provider for FailingProvider { - fn get_name(&self) -> &str { - "failing-test-provider" - } - - async fn stream( - &self, - _model_config: &ModelConfig, - _system: &str, - _messages: &[Message], - _tools: &[Tool], - ) -> std::result::Result { - Ok(crate::providers::base::stream_from_single_message( - Message::assistant().with_text("unused"), - ProviderUsage::new("failing-test-provider".into(), Usage::default()), - )) - } - - async fn update_mode( - &self, - _session_id: &str, - _mode: GooseMode, - ) -> std::result::Result<(), ProviderError> { - Err(ProviderError::ExecutionError( - "intentional failure for test".into(), - )) - } - } - - let temp_dir = TempDir::new().unwrap(); - let manager = create_test_manager(&temp_dir).await; - manager - .set_default_provider(Arc::new(FailingProvider)) - .await; - - let session_id = String::from("failed-creation-test"); - let result = manager.get_or_create_agent(session_id.clone()).await; - - assert!( - result.is_err(), - "expected provider mode-update failure to propagate" - ); - assert!( - manager.creation_locks.lock().await.is_empty(), - "creation_locks must be empty after a failed agent creation" - ); - assert!( - !manager.has_session(&session_id).await, - "failed creation must not insert into the LRU cache" - ); - } - #[tokio::test] async fn test_lru_eviction_prunes_creation_lock() { // Sessions can disappear from the LRU cache without going through diff --git a/crates/goose/src/permission/permission_inspector.rs b/crates/goose/src/permission/permission_inspector.rs index 94e11d3dbc..c97e581f8f 100644 --- a/crates/goose/src/permission/permission_inspector.rs +++ b/crates/goose/src/permission/permission_inspector.rs @@ -1,5 +1,5 @@ use crate::agents::platform_extensions::MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE; -use crate::agents::types::SharedProvider; +use crate::agents::provider_manager::ProviderManager; use crate::config::permission::PermissionLevel; use crate::config::{GooseMode, PermissionManager}; use crate::conversation::message::{Message, ToolRequest}; @@ -14,7 +14,7 @@ use std::sync::{Arc, RwLock}; /// Permission Inspector that handles tool permission checking pub struct PermissionInspector { pub permission_manager: Arc, - provider: SharedProvider, + providers: Arc, session_manager: Arc, readonly_tools: RwLock>, } @@ -36,12 +36,12 @@ fn cache_non_readonly_decision( impl PermissionInspector { pub fn new( permission_manager: Arc, - provider: SharedProvider, + providers: Arc, session_manager: Arc, ) -> Self { Self { permission_manager, - provider, + providers, session_manager, readonly_tools: RwLock::new(HashSet::new()), } @@ -228,7 +228,11 @@ impl ToolInspector for PermissionInspector { // LLM-based read-only detection for deferred SmartApprove candidates if !llm_detect_candidates.is_empty() { - let detected_request_ids: HashSet = match self.provider.lock().await.clone() { + let provider = match self.session_manager.get_session(session_id, false).await { + Ok(session) => self.providers.provider_for(&session).await.ok(), + Err(_) => None, + }; + let detected_request_ids: HashSet = match provider { Some(provider) => detect_read_only_requests( provider, &self.session_manager, @@ -276,7 +280,6 @@ mod tests { use rmcp::object; use std::sync::Arc; use test_case::test_case; - use tokio::sync::Mutex; async fn inspect_tool( mode: GooseMode, @@ -295,7 +298,7 @@ mod tests { tempfile::tempdir().unwrap().keep(), )); let inspector = - PermissionInspector::new(Arc::clone(&pm), Arc::new(Mutex::new(None)), session_manager); + PermissionInspector::new(Arc::clone(&pm), Default::default(), session_manager); if smart_approved { *inspector.readonly_tools.write().unwrap() = ["tool".to_string()].into_iter().collect(); } diff --git a/crates/goose/src/providers/amp_acp.rs b/crates/goose/src/providers/amp_acp.rs index 3a99f6d88b..84898a5180 100644 --- a/crates/goose/src/providers/amp_acp.rs +++ b/crates/goose/src/providers/amp_acp.rs @@ -46,6 +46,7 @@ impl goose_providers::base::ProviderDescriptor for AmpAcpProvider { impl ProviderDef for AmpAcpProvider { type Provider = AcpProvider; + const SESSION_BOUND: bool = true; fn from_env( extensions: Vec, diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index 2a5cfcfb52..aa8fe1eddf 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -31,6 +31,10 @@ pub(crate) fn current_working_dir() -> PathBuf { pub trait ProviderDef: ProviderDescriptor + Send + Sync { type Provider: Provider + 'static; + /// Instances hold a remote session or are built from one session's + /// extensions and working directory, so they cannot be shared. + const SESSION_BOUND: bool = false; + fn from_env( extensions: Vec, tls_config: Option, diff --git a/crates/goose/src/providers/claude_acp.rs b/crates/goose/src/providers/claude_acp.rs index fe4680ad65..275a6107c8 100644 --- a/crates/goose/src/providers/claude_acp.rs +++ b/crates/goose/src/providers/claude_acp.rs @@ -48,6 +48,7 @@ impl goose_providers::base::ProviderDescriptor for ClaudeAcpProvider { impl ProviderDef for ClaudeAcpProvider { type Provider = AcpProvider; + const SESSION_BOUND: bool = true; fn from_env( extensions: Vec, diff --git a/crates/goose/src/providers/claude_code.rs b/crates/goose/src/providers/claude_code.rs index 2fc3d19343..d29c9005ab 100644 --- a/crates/goose/src/providers/claude_code.rs +++ b/crates/goose/src/providers/claude_code.rs @@ -640,6 +640,7 @@ impl goose_providers::base::ProviderDescriptor for ClaudeCodeProvider { impl ProviderDef for ClaudeCodeProvider { type Provider = Self; + const SESSION_BOUND: bool = true; fn from_env( extensions: Vec, diff --git a/crates/goose/src/providers/codex.rs b/crates/goose/src/providers/codex.rs index 2e5b17717a..2317892cb0 100644 --- a/crates/goose/src/providers/codex.rs +++ b/crates/goose/src/providers/codex.rs @@ -637,6 +637,7 @@ impl goose_providers::base::ProviderDescriptor for CodexProvider { impl ProviderDef for CodexProvider { type Provider = Self; + const SESSION_BOUND: bool = true; fn from_env( extensions: Vec, diff --git a/crates/goose/src/providers/codex_acp.rs b/crates/goose/src/providers/codex_acp.rs index 4a084ab2b9..6deb6f8588 100644 --- a/crates/goose/src/providers/codex_acp.rs +++ b/crates/goose/src/providers/codex_acp.rs @@ -56,6 +56,7 @@ impl goose_providers::base::ProviderDescriptor for CodexAcpProvider { impl ProviderDef for CodexAcpProvider { type Provider = AcpProvider; + const SESSION_BOUND: bool = true; fn from_env( extensions: Vec, diff --git a/crates/goose/src/providers/copilot_acp.rs b/crates/goose/src/providers/copilot_acp.rs index 543f98f404..fcfdc8cdae 100644 --- a/crates/goose/src/providers/copilot_acp.rs +++ b/crates/goose/src/providers/copilot_acp.rs @@ -105,6 +105,7 @@ impl CopilotAcpProvider { impl ProviderDef for CopilotAcpProvider { type Provider = AcpProvider; + const SESSION_BOUND: bool = true; fn from_env( extensions: Vec, diff --git a/crates/goose/src/providers/gemini_cli.rs b/crates/goose/src/providers/gemini_cli.rs index 844978363b..df3d2a3783 100644 --- a/crates/goose/src/providers/gemini_cli.rs +++ b/crates/goose/src/providers/gemini_cli.rs @@ -183,6 +183,7 @@ impl goose_providers::base::ProviderDescriptor for GeminiCliProvider { impl ProviderDef for GeminiCliProvider { type Provider = Self; + const SESSION_BOUND: bool = true; fn from_env( _extensions: Vec, diff --git a/crates/goose/src/providers/init.rs b/crates/goose/src/providers/init.rs index 04702c1d37..08bf5ac9c1 100644 --- a/crates/goose/src/providers/init.rs +++ b/crates/goose/src/providers/init.rs @@ -1,4 +1,3 @@ -use std::path::PathBuf; use std::sync::{Arc, RwLock}; #[cfg(feature = "aws-providers")] @@ -277,15 +276,6 @@ pub async fn create(name: &str, extensions: Vec) -> Result, - working_dir: PathBuf, -) -> Result> { - let entry = get_from_registry(name).await?; - entry.create_with_working_dir(extensions, working_dir).await -} - pub async fn create_with_default_model( name: impl AsRef, extensions: Vec, @@ -310,13 +300,6 @@ pub async fn cleanup_provider(name: &str) -> Result<()> { Ok(()) } -pub async fn create_with_named_model( - provider_name: &str, - extensions: Vec, -) -> Result> { - create(provider_name, extensions).await -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 28710b43f3..4f935fd307 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -100,8 +100,7 @@ pub mod xai; pub mod xai_oauth; pub use init::{ - cleanup_provider, create, create_with_default_model, create_with_named_model, - create_with_working_dir, get_from_registry, inventory_identity, providers, - refresh_custom_providers, + cleanup_provider, create, create_with_default_model, get_from_registry, inventory_identity, + providers, refresh_custom_providers, }; pub use retry::{retry_operation, RetryConfig}; diff --git a/crates/goose/src/providers/pi_acp.rs b/crates/goose/src/providers/pi_acp.rs index cd290adc2f..188aaca54f 100644 --- a/crates/goose/src/providers/pi_acp.rs +++ b/crates/goose/src/providers/pi_acp.rs @@ -88,6 +88,7 @@ impl PiAcpProvider { impl ProviderDef for PiAcpProvider { type Provider = AcpProvider; + const SESSION_BOUND: bool = true; fn from_env( extensions: Vec, diff --git a/crates/goose/src/providers/provider_registry.rs b/crates/goose/src/providers/provider_registry.rs index ed15656112..d297c0eea5 100644 --- a/crates/goose/src/providers/provider_registry.rs +++ b/crates/goose/src/providers/provider_registry.rs @@ -33,6 +33,7 @@ pub struct ProviderEntry { supports_inventory_refresh: bool, tls_config: Option, toolshim: bool, + session_bound: bool, } impl ProviderEntry { @@ -44,6 +45,10 @@ impl ProviderEntry { self.provider_type } + pub fn session_bound(&self) -> bool { + self.session_bound + } + pub fn supports_inventory_refresh(&self) -> bool { self.supports_inventory_refresh } @@ -154,6 +159,7 @@ impl ProviderRegistry { supports_inventory_refresh: inventory.supports_refresh, tls_config: self.tls_config.clone(), toolshim: false, + session_bound: F::SESSION_BOUND, }, ); } @@ -319,6 +325,7 @@ impl ProviderRegistry { supports_inventory_refresh, tls_config: self.tls_config.clone(), toolshim: config.toolshim, + session_bound: false, }, ); } diff --git a/crates/goose/src/scheduler/full.rs b/crates/goose/src/scheduler/full.rs index e1e94a2a68..d28c3d33bd 100644 --- a/crates/goose/src/scheduler/full.rs +++ b/crates/goose/src/scheduler/full.rs @@ -18,7 +18,6 @@ use crate::conversation::message::Message; use crate::conversation::Conversation; #[cfg(feature = "telemetry")] use crate::posthog; -use crate::providers::create; use crate::recipe::build_recipe::build_recipe_from_template; use crate::recipe::validate_recipe::{ recipe_file_format, validate_recipe_for_scheduling, SchedulerRecipeError, @@ -895,9 +894,8 @@ async fn execute_job( agent.add_extension(ext.clone(), &session.id).await?; } - let agent_provider = create(&provider_name, extensions).await?; agent - .update_provider(agent_provider, model_config, &session.id) + .switch_provider(&session.id, &provider_name, model_config) .await?; agent .update_goose_mode(GooseMode::Auto, &session.id) diff --git a/crates/goose/src/security/adversary_inspector.rs b/crates/goose/src/security/adversary_inspector.rs index c2309fc5d1..9311eabb9d 100644 --- a/crates/goose/src/security/adversary_inspector.rs +++ b/crates/goose/src/security/adversary_inspector.rs @@ -3,7 +3,7 @@ use async_trait::async_trait; use chrono::Utc; use std::sync::{Arc, OnceLock}; -use crate::agents::types::SharedProvider; +use crate::agents::provider_manager::ProviderManager; use crate::config::paths::Paths; use crate::config::GooseMode; use crate::conversation::message::{Message, MessageContent, ToolRequest}; @@ -13,28 +13,6 @@ use crate::utils::safe_truncate; const DEFAULT_TOOLS: &[&str] = &["shell"]; -async fn resolve_model_config( - session_manager: &crate::session::SessionManager, - session_id: &str, -) -> Result { - if !session_id.is_empty() { - if let Ok(session) = session_manager.get_session(session_id, false).await { - if let Some(model_config) = session.model_config { - return Ok(model_config); - } - } - } - - let config = crate::config::Config::global(); - let provider_name = config - .get_goose_provider() - .map_err(|_| anyhow::anyhow!("missing provider"))?; - let model_name = config - .get_goose_model() - .map_err(|_| anyhow::anyhow!("missing model"))?; - crate::model_config::model_config_from_user_config(&provider_name, &model_name) -} - const DEFAULT_RULES: &str = r#"BLOCK if the command: - Exfiltrates data (curl/wget posting to unknown URLs, piping secrets out) - Is destructive beyond the project scope (rm -rf /, modifying system files) @@ -71,7 +49,7 @@ struct AdversaryConfig { /// If the file is absent, this inspector is disabled. /// If the review fails, the inspector fails open (allows the tool call). pub struct AdversaryInspector { - provider: SharedProvider, + providers: Arc, session_manager: Arc, config: OnceLock>, config_path: Option, @@ -79,11 +57,11 @@ pub struct AdversaryInspector { impl AdversaryInspector { pub fn new( - provider: SharedProvider, + providers: Arc, session_manager: Arc, ) -> Self { Self { - provider, + providers, session_manager, config: OnceLock::new(), config_path: None, @@ -91,12 +69,12 @@ impl AdversaryInspector { } pub fn with_config_dir( - provider: SharedProvider, + providers: Arc, session_manager: Arc, config_dir: std::path::PathBuf, ) -> Self { Self { - provider, + providers, session_manager, config: OnceLock::new(), config_path: Some(config_dir.join("adversary.md")), @@ -276,12 +254,12 @@ impl AdversaryInspector { recent_messages: &[String], rules: &str, ) -> Result<(bool, String)> { - let provider_guard = self.provider.lock().await; - let provider = match provider_guard.clone() { - Some(p) => p, - None => return Ok((true, "No provider available".to_string())), + let Ok(session) = self.session_manager.get_session(session_id, false).await else { + return Ok((true, "No provider available".to_string())); + }; + let Ok(provider) = self.providers.provider_for(&session).await else { + return Ok((true, "No provider available".to_string())); }; - drop(provider_guard); let history_section = if !recent_messages.is_empty() { let mut s = String::from("Recent user messages (oldest first):\n"); @@ -321,9 +299,7 @@ impl AdversaryInspector { )]; let conversation = Conversation::new_unvalidated(check_messages); - let model_config = resolve_model_config(&self.session_manager, session_id) - .await - .map_err(|e| anyhow::anyhow!("Could not resolve model config: {}", e))?; + let model_config = crate::agents::provider_manager::model_config_for(&session)?; let (response, _usage) = crate::session_context::with_session_id( Some(session_id.to_string()), provider.complete(&model_config, system_prompt, conversation.messages(), &[]), @@ -499,7 +475,6 @@ mod tests { use rmcp::model::CallToolRequestParams; use rmcp::object; use std::sync::Arc; - use tokio::sync::Mutex; #[test] fn test_parse_with_tools_frontmatter() { @@ -683,12 +658,11 @@ mod tests { async fn test_disabled_when_no_adversary_md() { let tmp = tempfile::tempdir().unwrap(); - let provider: SharedProvider = Arc::new(Mutex::new(None)); let session_manager = Arc::new(crate::session::SessionManager::new( tmp.path().to_path_buf(), )); let inspector = AdversaryInspector::with_config_dir( - provider, + Default::default(), session_manager, tmp.path().to_path_buf(), ); diff --git a/crates/goose/src/tool_call_labels.rs b/crates/goose/src/tool_call_labels.rs index ca1fd5987b..47b1544b10 100644 --- a/crates/goose/src/tool_call_labels.rs +++ b/crates/goose/src/tool_call_labels.rs @@ -33,7 +33,7 @@ pub(crate) async fn generate_tool_title( session_id: &str, tool_request: &ToolRequest, ) -> Option { - let provider = agent.provider().await.ok()?; + let provider = agent.provider(session_id).await.ok()?; if provider.manages_own_context() { return None; } @@ -72,7 +72,7 @@ pub(crate) async fn generate_tool_chain_summary( return None; } - let provider = agent.provider().await.ok()?; + let provider = agent.provider(session_id).await.ok()?; if provider.manages_own_context() { return None; } diff --git a/crates/goose/tests/acp_common_tests/mod.rs b/crates/goose/tests/acp_common_tests/mod.rs index 823d26117c..ea5cbbc9bd 100644 --- a/crates/goose/tests/acp_common_tests/mod.rs +++ b/crates/goose/tests/acp_common_tests/mod.rs @@ -448,7 +448,7 @@ pub async fn run_fs_write_text_file_true() { pub async fn run_initialize_doesnt_hit_provider() { let provider_factory: AcpProviderFactory = - Arc::new(|_, _, _, _| Box::pin(async { Err(anyhow::anyhow!("no provider configured")) })); + Arc::new(|_| Box::pin(async { Err(anyhow::anyhow!("no provider configured")) })); let openai = OpenAiFixture::new(vec![], C::expected_session_id()).await; let config = TestConnectionConfig { diff --git a/crates/goose/tests/acp_custom_requests_test.rs b/crates/goose/tests/acp_custom_requests_test.rs index 0d8bbe597a..c4275be1cc 100644 --- a/crates/goose/tests/acp_custom_requests_test.rs +++ b/crates/goose/tests/acp_custom_requests_test.rs @@ -1136,21 +1136,18 @@ fn test_custom_provider_supported_models_lists_raw_provider_models() { write_acp_global_config(DEFAULT_ACP_TEST_CONFIG); run_test(async move { let openai = OpenAiFixture::new(vec![], Arc::new(EnforceSessionId::default())).await; - let provider_factory: AcpProviderFactory = Arc::new( - |provider_name, _extensions, _working_dir, use_default_model| { - assert!(use_default_model); - Box::pin(async move { - Ok(Arc::new(MockProvider { - name: provider_name, - recommended_models: vec!["canonical-filtered-model".to_string()], - supported_models: Ok(vec![ - "goose-claude-opus-4-8".to_string(), - "raw-databricks-endpoint".to_string(), - ]), - }) as Arc) - }) - }, - ); + let provider_factory: AcpProviderFactory = Arc::new(|provider_name| { + Box::pin(async move { + Ok(Arc::new(MockProvider { + name: provider_name, + recommended_models: vec!["canonical-filtered-model".to_string()], + supported_models: Ok(vec![ + "goose-claude-opus-4-8".to_string(), + "raw-databricks-endpoint".to_string(), + ]), + }) as Arc) + }) + }); let conn = AcpServerConnection::new( TestConnectionConfig { provider_factory: Some(provider_factory), @@ -1188,7 +1185,7 @@ fn test_custom_provider_supported_models_maps_not_configured_error() { write_acp_global_config(DEFAULT_ACP_TEST_CONFIG); run_test(async move { let openai = OpenAiFixture::new(vec![], Arc::new(EnforceSessionId::default())).await; - let provider_factory: AcpProviderFactory = Arc::new(|provider_name, _, _, _| { + let provider_factory: AcpProviderFactory = Arc::new(|provider_name| { Box::pin(async move { Ok(Arc::new(MockProvider { name: provider_name, @@ -1225,7 +1222,7 @@ fn test_custom_provider_supported_models_maps_authentication_error() { write_acp_global_config(DEFAULT_ACP_TEST_CONFIG); run_test(async move { let openai = OpenAiFixture::new(vec![], Arc::new(EnforceSessionId::default())).await; - let provider_factory: AcpProviderFactory = Arc::new(|provider_name, _, _, _| { + let provider_factory: AcpProviderFactory = Arc::new(|provider_name| { Box::pin(async move { Ok(Arc::new(MockProvider { name: provider_name, diff --git a/crates/goose/tests/acp_fixtures/mod.rs b/crates/goose/tests/acp_fixtures/mod.rs index 4111cf5605..2df7ca8389 100644 --- a/crates/goose/tests/acp_fixtures/mod.rs +++ b/crates/goose/tests/acp_fixtures/mod.rs @@ -382,21 +382,19 @@ pub async fn spawn_acp_server_in_process( write_global_test_config(&config_path, openai_base_url); let provider_factory = provider_factory.unwrap_or_else(|| { let base_url = openai_base_url.to_string(); - Arc::new( - move |_provider_name, _extensions, _working_dir, _use_default_model| { - let base_url = base_url.clone(); - Box::pin(async move { - let api_client = ApiClient::new_with_tls( - base_url, - ApiAuthMethod::BearerToken("test-key".to_string()), - None, - ) - .unwrap(); - let provider: Arc = Arc::new(OpenAiProvider::new(api_client)); - Ok(provider) - }) - }, - ) + Arc::new(move |_provider_name| { + let base_url = base_url.clone(); + Box::pin(async move { + let api_client = ApiClient::new_with_tls( + base_url, + ApiAuthMethod::BearerToken("test-key".to_string()), + None, + ) + .unwrap(); + let provider: Arc = Arc::new(OpenAiProvider::new(api_client)); + Ok(provider) + }) + }) }); let active_runs = Arc::new(goose::acp::server::ActiveRunRegistry::default()); diff --git a/crates/goose/tests/acp_secret_cache_invalidation_test.rs b/crates/goose/tests/acp_secret_cache_invalidation_test.rs index b53a0ab920..f424869503 100644 --- a/crates/goose/tests/acp_secret_cache_invalidation_test.rs +++ b/crates/goose/tests/acp_secret_cache_invalidation_test.rs @@ -45,15 +45,13 @@ impl Provider for MockProvider { } fn mock_provider_factory() -> goose::acp::server::AcpProviderFactory { - Arc::new( - |provider_name, _extensions, _working_dir, _use_default_model| { - Box::pin(async move { - Ok(Arc::new(MockProvider { - name: provider_name, - }) as Arc) - }) - }, - ) + Arc::new(|provider_name| { + Box::pin(async move { + Ok(Arc::new(MockProvider { + name: provider_name, + }) as Arc) + }) + }) } fn write_config(config_dir: &std::path::Path) { diff --git a/crates/goose/tests/adversary_inspector_tests.rs b/crates/goose/tests/adversary_inspector_tests.rs index 1a5bc6aa17..db4b7a8601 100644 --- a/crates/goose/tests/adversary_inspector_tests.rs +++ b/crates/goose/tests/adversary_inspector_tests.rs @@ -5,7 +5,6 @@ use goose::tool_inspection::ToolInspector; use rmcp::model::CallToolRequestParams; use rmcp::object; use std::sync::Arc; -use tokio::sync::Mutex; fn make_request( id: &str, @@ -29,9 +28,8 @@ fn write_adversary_md(dir: &std::path::Path, content: &str) { async fn test_adversary_disabled_without_config_file() { let tmp = tempfile::tempdir().unwrap(); - let provider = Arc::new(Mutex::new(None)); let inspector = AdversaryInspector::with_config_dir( - provider, + Default::default(), Arc::new(goose::session::SessionManager::new( tmp.path().to_path_buf(), )), @@ -63,9 +61,8 @@ async fn test_adversary_enabled_default_tools() { let tmp = tempfile::tempdir().unwrap(); write_adversary_md(tmp.path(), "BLOCK everything for testing"); - let provider = Arc::new(Mutex::new(None)); let inspector = AdversaryInspector::with_config_dir( - provider, + Default::default(), Arc::new(goose::session::SessionManager::new( tmp.path().to_path_buf(), )), @@ -127,9 +124,8 @@ async fn test_adversary_custom_tool_filter() { "tools: shell, developer__shell\n---\nBLOCK bad stuff", ); - let provider = Arc::new(Mutex::new(None)); let inspector = AdversaryInspector::with_config_dir( - provider, + Default::default(), Arc::new(goose::session::SessionManager::new( tmp.path().to_path_buf(), )), diff --git a/crates/goose/tests/mcp_integration_test.rs b/crates/goose/tests/mcp_integration_test.rs index 937f64d084..8289746ba5 100644 --- a/crates/goose/tests/mcp_integration_test.rs +++ b/crates/goose/tests/mcp_integration_test.rs @@ -95,7 +95,7 @@ async fn fixture(mcpui: bool, protocol_version: Option) -> Fixt temp_dir.path().to_path_buf(), )); let manager = Arc::new(ExtensionManager::new( - Arc::new(tokio::sync::Mutex::new(None)), + Default::default(), session_manager.clone(), None, "extension-contract".to_string(), @@ -902,7 +902,7 @@ async fn test_replayed_session( temp_dir.path().to_path_buf(), )); let extension_manager = Arc::new(ExtensionManager::new( - Arc::new(tokio::sync::Mutex::new(None)), + Default::default(), session_manager, None, GoosePlatform::GooseDesktop.to_string(), diff --git a/crates/goose/tests/providers.rs b/crates/goose/tests/providers.rs index 4875c4d1db..602ba032a9 100644 --- a/crates/goose/tests/providers.rs +++ b/crates/goose/tests/providers.rs @@ -14,7 +14,7 @@ use goose::providers::base::Provider; use goose::providers::bedrock::BEDROCK_DEFAULT_MODEL; use goose::providers::claude_code::CLAUDE_CODE_DEFAULT_MODEL; use goose::providers::codex::CODEX_DEFAULT_MODEL; -use goose::providers::create_with_named_model; +use goose::providers::create; use goose::providers::google::GOOGLE_DEFAULT_MODEL; use goose::providers::litellm::LITELLM_DEFAULT_MODEL; use goose::providers::openai::OPEN_AI_DEFAULT_MODEL; @@ -241,7 +241,7 @@ impl ProviderFixture { available_tools: vec![], }; - let provider = create_with_named_model( + let provider = create( &config.name.to_lowercase(), vec![mcp_extension.clone(), developer_extension.clone()], )