diff --git a/app/src/ai/blocklist/controller.rs b/app/src/ai/blocklist/controller.rs index 46d948b9..fda103e0 100644 --- a/app/src/ai/blocklist/controller.rs +++ b/app/src/ai/blocklist/controller.rs @@ -83,8 +83,8 @@ use crate::ai::provider::ProviderConfig; use crate::ai::remote_logging::{self, RemoteLogLevel, RemoteLogRecord}; use crate::ai::runtime::{ prepare_provider_run, provider_context_window_tokens, provider_runtime_for_request, - OrchestrationModelOption, PreparedProviderRun, ProviderActionContext, ProviderRunBlock, - ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection, + restore_provider_config, OrchestrationModelOption, PreparedProviderRun, ProviderActionContext, + ProviderRunBlock, ProviderRunCoordinator, ProviderRunProfile, ProviderRunProjection, ProviderRunResponseProjector, ProviderToolExecutionRef, ProviderToolLifecycleOutcome, RuntimeResponseConfig, BASE_PROVIDER_PROFILE, CLI_MONITOR_PROVIDER_PROFILE, COMPUTER_USE_PROVIDER_PROFILE, @@ -5891,8 +5891,34 @@ impl BlocklistAIController { self.fail_restored_provider_run(conversation_id, error, ctx); return; } - let base_provider_config = - ResponseStream::resolve_provider_config(snapshot.base_request.model.as_str(), ctx); + let provider_configs = (|| { + let resolve = |request: &TurnRequest| { + restore_provider_config(request, |model| { + ResponseStream::resolve_provider_config(model, ctx) + }) + }; + Ok::<_, anyhow::Error>(( + resolve(&snapshot.base_request)?, + snapshot + .cli_monitor_request + .as_ref() + .map(&resolve) + .transpose()?, + snapshot + .computer_use_request + .as_ref() + .map(&resolve) + .transpose()?, + )) + })(); + let (base_provider_config, cli_provider_config, computer_use_provider_config) = + match provider_configs { + Ok(configs) => configs, + Err(error) => { + self.fail_restored_provider_run(conversation_id, error.to_string(), ctx); + return; + } + }; // A restored checkpoint may contain an older, larger catalog budget. // Reconcile it before any resumed model call and persist the corrected limit. snapshot.response_config.max_context_tokens = provider_context_window_tokens( @@ -5905,14 +5931,6 @@ impl BlocklistAIController { self.fail_restored_provider_run(conversation_id, error, ctx); return; } - let cli_provider_config = snapshot - .cli_monitor_request - .as_ref() - .map(|request| ResponseStream::resolve_provider_config(request.model.as_str(), ctx)); - let computer_use_provider_config = snapshot - .computer_use_request - .as_ref() - .map(|request| ResponseStream::resolve_provider_config(request.model.as_str(), ctx)); let _ = ctx.spawn( async move { let base_runtime = diff --git a/app/src/ai/blocklist/controller/response_stream.rs b/app/src/ai/blocklist/controller/response_stream.rs index ea77f7d9..d28e819b 100644 --- a/app/src/ai/blocklist/controller/response_stream.rs +++ b/app/src/ai/blocklist/controller/response_stream.rs @@ -170,6 +170,7 @@ impl ResponseStream { if let Some(client_config) = llm_prefs.openai_client_config_for_model(model_id) { return ProviderConfig::OpenAI(OpenAIClientConfig { kind: client_config.kind, + connection_identity: client_config.connection_identity.clone(), base_url: client_config.base_url.clone(), api_key: client_config.api_key.clone(), project_id: client_config.project_id.clone(), diff --git a/app/src/ai/crosscheck/reviewer.rs b/app/src/ai/crosscheck/reviewer.rs index 1ed7a181..234e60af 100644 --- a/app/src/ai/crosscheck/reviewer.rs +++ b/app/src/ai/crosscheck/reviewer.rs @@ -164,6 +164,7 @@ impl CrosscheckReviewer { if let Some(client_config) = llm_prefs.openai_client_config_for_model(model_id) { return ProviderConfig::OpenAI(OpenAIClientConfig { kind: client_config.kind, + connection_identity: client_config.connection_identity.clone(), base_url: client_config.base_url.clone(), api_key: client_config.api_key.clone(), project_id: client_config.project_id.clone(), diff --git a/app/src/ai/custom_model_routers.rs b/app/src/ai/custom_model_routers.rs index c00a9ccf..d5b9063f 100644 --- a/app/src/ai/custom_model_routers.rs +++ b/app/src/ai/custom_model_routers.rs @@ -126,6 +126,7 @@ impl CustomModelRouter { base_model_name: name, id: config_key.into(), reasoning_level: None, + reasoning_model_id: None, usage_metadata: LLMUsageMetadata { request_multiplier: 1, credit_multiplier: None, diff --git a/app/src/ai/execution_profiles/editor/mod_tests.rs b/app/src/ai/execution_profiles/editor/mod_tests.rs index 82330030..401f57c6 100644 --- a/app/src/ai/execution_profiles/editor/mod_tests.rs +++ b/app/src/ai/execution_profiles/editor/mod_tests.rs @@ -34,6 +34,7 @@ fn configurable_model(provider: LLMProvider) -> LLMInfo { base_model_name: "test model".to_string(), id: "test-model".into(), reasoning_level: None, + reasoning_model_id: None, usage_metadata: LLMUsageMetadata { request_multiplier: 1, credit_multiplier: None, diff --git a/app/src/ai/llms.rs b/app/src/ai/llms.rs index 990a4a9e..58589d17 100644 --- a/app/src/ai/llms.rs +++ b/app/src/ai/llms.rs @@ -204,6 +204,10 @@ pub struct LLMInfo { pub base_model_name: String, pub id: LLMId, pub reasoning_level: Option, + /// Exact provider model identity for grouping locally discovered reasoning variants. + /// Older catalogs and ACP retain their existing display-name grouping. + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning_model_id: Option, pub usage_metadata: LLMUsageMetadata, pub description: Option, pub disable_reason: Option, @@ -244,6 +248,8 @@ impl<'de> Deserialize<'de> for LLMInfo { id: LLMId, #[serde(default)] reasoning_level: Option, + #[serde(default)] + reasoning_model_id: Option, usage_metadata: LLMUsageMetadata, #[serde(default)] description: Option, @@ -288,6 +294,7 @@ impl<'de> Deserialize<'de> for LLMInfo { display_name: wire.display_name, id: wire.id, reasoning_level: wire.reasoning_level, + reasoning_model_id: wire.reasoning_model_id, usage_metadata: wire.usage_metadata, description: wire.description, disable_reason: wire.disable_reason, @@ -337,9 +344,21 @@ impl LLMInfo { /// Reasoning menus must never merge models belonging to different connections. pub fn reasoning_group_key(&self) -> String { - let connection = self.id.as_str().rsplit_once("::connection::") - .map(|(_, connection)| connection).unwrap_or_default(); - format!("{}\u{1f}{}\u{1f}{:?}\u{1f}{connection}", self.base_model_name, self.description.as_deref().unwrap_or_default(), self.provider) + let connection = self + .id + .as_str() + .rsplit_once("::connection::") + .map(|(_, connection)| connection) + .unwrap_or_default(); + let model_identity = self + .reasoning_model_id + .as_deref() + .unwrap_or(&self.base_model_name); + format!( + "{model_identity}\u{1f}{}\u{1f}{:?}\u{1f}{connection}", + self.description.as_deref().unwrap_or_default(), + self.provider + ) } /// Returns the given model's base name. @@ -366,6 +385,7 @@ impl LLMInfo { base_model_name: llm_name.to_string(), id: llm_name.into(), reasoning_level: None, + reasoning_model_id: None, usage_metadata: LLMUsageMetadata { request_multiplier: 1, credit_multiplier: None, @@ -457,6 +477,7 @@ impl AvailableLLMs { base_model_name: "No models configured".to_owned(), id: "none".to_owned().into(), reasoning_level: None, + reasoning_model_id: None, usage_metadata: LLMUsageMetadata { request_multiplier: 1, credit_multiplier: None, @@ -531,6 +552,7 @@ fn default_computer_use_llms() -> AvailableLLMs { base_model_name: "auto".to_owned(), id: "computer-use-agent-auto".to_owned().into(), reasoning_level: None, + reasoning_model_id: None, usage_metadata: LLMUsageMetadata { request_multiplier: 1, credit_multiplier: None, @@ -562,6 +584,7 @@ impl Default for ModelsByFeature { base_model_name: "No models configured".to_owned(), id: "placeholder".to_owned().into(), reasoning_level: None, + reasoning_model_id: None, usage_metadata: LLMUsageMetadata { request_multiplier: 1, credit_multiplier: None, @@ -1010,6 +1033,7 @@ impl LLMPreferences { display_name: model.display_name.clone(), base_model_name: model.display_name.clone(), reasoning_level: None, + reasoning_model_id: Some(model_id.clone()), usage_metadata: LLMUsageMetadata { request_multiplier: 1, credit_multiplier: None, @@ -1317,6 +1341,7 @@ impl LLMPreferences { // actual model ID while using their synthetic key only for selection. let client_config = OpenAIClientConfig { kind: provider_kind, + connection_identity: Some(connection_identity.clone()), base_url: base_url.clone(), api_key: api_key.clone(), project_id: provider_project_id.clone(), @@ -1339,6 +1364,7 @@ impl LLMPreferences { display_name, base_model_name: model.display_name.clone(), reasoning_level: reasoning_effort, + reasoning_model_id: Some(model.model_id.clone()), usage_metadata: LLMUsageMetadata { request_multiplier: 1, credit_multiplier: None, @@ -1505,6 +1531,7 @@ impl LLMPreferences { display_name, base_model_name: base_model_name.to_owned(), reasoning_level: reasoning.map(|reasoning| reasoning.name.clone()), + reasoning_model_id: None, usage_metadata: LLMUsageMetadata { request_multiplier: 1, credit_multiplier: None, @@ -2239,13 +2266,23 @@ impl LLMPreferences { .data() .computer_use_model .clone() - .and_then(|id| { - self.get_computer_use_llm_choices(app) - .find(|info| info.id == id) - }) + .and_then(|id| self.computer_use_llm_info_for_id(&id, app)) .unwrap_or_else(|| self.get_default_computer_use_model()) } + /// Resolve migrated selection IDs before checking computer-use eligibility. + /// Looking up the canonical ID in these choices preserves the vision filter + /// and any computer-use-specific metadata, including disable reasons. + fn computer_use_llm_info_for_id<'a>( + &'a self, + id: &LLMId, + app: &'a AppContext, + ) -> Option<&'a LLMInfo> { + let canonical_id = self.get_llm_info(id).map_or(id, |info| &info.id); + self.get_computer_use_llm_choices(app) + .find(|info| &info.id == canonical_id) + } + /// Returns the default computer use model as a fallback. pub fn get_default_computer_use_model(&self) -> &LLMInfo { self.get_computer_use_available().default_llm_info() @@ -2748,8 +2785,7 @@ impl LLMPreferences { } if let Some(preferred_llm_id) = &profile.data().computer_use_model { if self - .get_computer_use_llm_choices(ctx) - .find(|info| info.id == *preferred_llm_id) + .computer_use_llm_info_for_id(preferred_llm_id, ctx) .filter(|info| { let has_byok_key = is_using_api_key_for_provider(&info.provider, ctx); diff --git a/app/src/ai/llms_tests.rs b/app/src/ai/llms_tests.rs index 4c243a0e..6f0534d6 100644 --- a/app/src/ai/llms_tests.rs +++ b/app/src/ai/llms_tests.rs @@ -2,6 +2,9 @@ use serde_json::json; use warpui::App; use super::*; +use crate::ai::execution_profiles::model_menu_items::{ + collapse_model_reasoning_choices, reasoning_choices_for_model, +}; use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel; use crate::ai::mcp::TemplatableMCPServerManager; use crate::auth::auth_manager::AuthManager; @@ -120,6 +123,7 @@ fn llm_info_round_trip_serializes_and_deserializes() { let wire_json = r#"{ "display_name": "claude-3", "base_model_name": "claude-3", + "reasoning_model_id": "provider/claude-3", "id": "claude-3", "usage_metadata": { "request_multiplier": 2, "credit_multiplier": 1.5 }, "description": "A powerful model", @@ -1574,6 +1578,208 @@ fn overlapping_connections_keep_all_reasoning_variants_in_one_family() { // A disabled second connection still reserves its legacy IDs: they // must not silently migrate to the third connection. assert!(preferences.get_llm_info(&old_id).is_none()); - assert!(preferences.openai_client_config_for_model(old_id.as_str()).is_none()); + assert!(preferences + .openai_client_config_for_model(old_id.as_str()) + .is_none()); + }); +} + +#[test] +fn discovered_models_with_the_same_label_keep_separate_reasoning_choices() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + let models = ["route-a", "route-b"].map(|id| { + let mut model = openai_model(id); + model.display_name = "Shared model name".to_owned(); + model.reasoning_efforts = vec!["low".to_owned(), "high".to_owned()]; + model + }); + AISettings::handle(&app).update(&mut app, |settings, ctx| { + settings.openai_enabled.set_value(true, ctx).unwrap(); + settings + .openai_providers + .set_value( + vec![OpenAIProviderConfig { + kind: OpenAIProviderKind::LiteLLM, + name: "Shared connection".to_owned(), + base_url: "https://models.test/v1".to_owned(), + models: models.into(), + enabled: true, + api_key: None, + project_id: None, + location: None, + }], + ctx, + ) + .unwrap(); + }); + let mut preferences = empty_preferences(); + app.read(|ctx| preferences.inject_openai_models(ctx)); + let choices: Vec<_> = preferences + .models_by_feature + .agent_mode + .choices + .iter() + .collect(); + let selected_id = LLMId::from("route-b::reasoning::high"); + let collapsed = collapse_model_reasoning_choices(choices.clone(), Some(&selected_id)); + assert_eq!(collapsed.len(), 2); + assert_eq!(collapsed[0].id.as_str(), "route-a"); + assert_eq!(collapsed[1].id, selected_id); + let levels = reasoning_choices_for_model(&selected_id, &choices); + assert_eq!(levels.len(), 3); + for model in levels { + let route = preferences + .openai_client_config_for_model(model.id.as_str()) + .unwrap(); + assert_eq!(route.model.as_deref(), Some("route-b")); + assert_eq!(route.reasoning_effort, model.reasoning_level); + } + }); +} + +#[test] +fn migrated_computer_use_reasoning_selection_survives_reconciliation() { + App::test((), |mut app| async move { + initialize_settings_for_tests(&mut app); + app.add_singleton_model(|_| ServerApiProvider::new_for_test()); + app.add_singleton_model(|_| AuthStateProvider::new_for_test()); + app.add_singleton_model(AuthManager::new_for_test); + app.add_singleton_model(|_| NetworkStatus::new()); + app.add_singleton_model(UserWorkspaces::default_mock); + app.add_singleton_model(CloudModel::mock); + app.add_singleton_model(|ctx| { + crate::local_object_repository::LocalObjectRepository::new(None, None, ctx) + }); + app.add_singleton_model(TeamTesterStatus::mock); + app.add_singleton_model(SyncQueue::mock); + app.add_singleton_model(UpdateManager::mock); + app.add_singleton_model(|_| TemplatableMCPServerManager::default()); + let profiles = app.add_singleton_model(|ctx| { + AIExecutionProfilesModel::new(&LaunchMode::new_for_unit_test(), ctx) + }); + AISettings::handle(&app).update(&mut app, |settings, ctx| { + settings.openai_enabled.set_value(true, ctx).unwrap(); + let mut visual = openai_model("shared-model"); + visual.vision_supported = true; + visual.reasoning_efforts = vec!["high".to_owned()]; + settings + .openai_providers + .set_value( + vec![ + OpenAIProviderConfig { + kind: OpenAIProviderKind::LiteLLM, + name: "First connection".to_owned(), + base_url: "https://first.test/v1".to_owned(), + models: vec![openai_model("shared-model")], + enabled: true, + api_key: None, + project_id: None, + location: None, + }, + OpenAIProviderConfig { + kind: OpenAIProviderKind::ChatGPTSubscription, + name: "Second connection".to_owned(), + base_url: "https://second.test/v1".to_owned(), + models: vec![visual], + enabled: true, + api_key: None, + project_id: None, + location: None, + }, + ], + ctx, + ) + .unwrap(); + }); + let old_id = LLMId::from("shared-model::reasoning::high"); + let mut preferences = empty_preferences(); + preferences.models_by_feature.computer_use = Some(default_computer_use_llms()); + app.read(|ctx| preferences.inject_openai_models(ctx)); + let canonical_id = preferences.get_llm_info(&old_id).unwrap().id.clone(); + assert_ne!(canonical_id, old_id); + preferences.models_by_feature.agent_mode.default_id = "shared-model".into(); + let preferences = app.add_singleton_model(|_| preferences); + let profile_id = profiles.as_ref(&app).default_profile_id(); + profiles.update(&mut app, |profiles, ctx| { + profiles.set_computer_use_model(profile_id, Some(old_id.clone()), ctx); + }); + app.read(|ctx| { + let selected = preferences + .as_ref(ctx) + .get_active_computer_use_model(ctx, None); + assert_eq!(selected.id, canonical_id); + assert_eq!(selected.reasoning_level.as_deref(), Some("high")); + }); + preferences.update(&mut app, |preferences, ctx| { + preferences.reconcile_disabled_model_preferences(ctx); + }); + app.read(|ctx| { + assert_eq!( + profiles + .as_ref(ctx) + .default_profile(ctx) + .data() + .computer_use_model, + Some(old_id.clone()) + ); + }); + + // A migrated ID must still obey the normal permanent-disable policy. + preferences.update(&mut app, |preferences, ctx| { + preferences + .models_by_feature + .agent_mode + .choices + .iter_mut() + .find(|model| model.id == canonical_id) + .unwrap() + .disable_reason = Some(DisableReason::AdminDisabled); + preferences.reconcile_disabled_model_preferences(ctx); + }); + app.read(|ctx| { + assert_eq!( + profiles + .as_ref(ctx) + .default_profile(ctx) + .data() + .computer_use_model, + None + ); + }); + + // Losing screenshot support must also invalidate the migrated selection. + profiles.update(&mut app, |profiles, ctx| { + profiles.set_computer_use_model(profile_id, Some(old_id.clone()), ctx); + }); + preferences.update(&mut app, |preferences, ctx| { + let model = preferences + .models_by_feature + .agent_mode + .choices + .iter_mut() + .find(|model| model.id == canonical_id) + .unwrap(); + model.disable_reason = None; + model.vision_supported = false; + assert_eq!( + preferences + .get_active_computer_use_model(ctx, None) + .id + .as_str(), + "computer-use-agent-auto" + ); + preferences.reconcile_disabled_model_preferences(ctx); + }); + app.read(|ctx| { + assert_eq!( + profiles + .as_ref(ctx) + .default_profile(ctx) + .data() + .computer_use_model, + None + ); + }); }); } diff --git a/app/src/ai/openai/client.rs b/app/src/ai/openai/client.rs index 3e00eb7a..43ef274e 100644 --- a/app/src/ai/openai/client.rs +++ b/app/src/ai/openai/client.rs @@ -9,6 +9,8 @@ use crate::settings::OpenAIProviderKind; #[derive(Clone, Debug)] pub struct OpenAIClientConfig { pub kind: OpenAIProviderKind, + /// Stable registry connection identity, excluding credentials. + pub connection_identity: Option, pub base_url: String, pub api_key: Option, pub project_id: Option, diff --git a/app/src/ai/runtime/mod.rs b/app/src/ai/runtime/mod.rs index 7abefc26..0485bdfb 100644 --- a/app/src/ai/runtime/mod.rs +++ b/app/src/ai/runtime/mod.rs @@ -18,7 +18,7 @@ pub(crate) use provider_run_coordinator::{ }; pub(crate) use rig::{ prepare_provider_run, provider_context_window_tokens, provider_runtime_for_request, - PreparedProviderRun, ProviderActionContext, + restore_provider_config, PreparedProviderRun, ProviderActionContext, }; pub(crate) use rig_request::OrchestrationModelOption; pub(crate) use task_progress::{restored_task_list, task_list_from_api_tasks}; diff --git a/app/src/ai/runtime/rig.rs b/app/src/ai/runtime/rig.rs index 1b21e252..a30bbe07 100644 --- a/app/src/ai/runtime/rig.rs +++ b/app/src/ai/runtime/rig.rs @@ -9,6 +9,7 @@ use galaxy_agent_rig::{ VertexAiRuntimeConfig, }; use galaxy_bedrock_model_catalog::model_metadata; +use sha2::{Digest, Sha256}; use uuid::Uuid; use warp_multi_agent_api::ToolType; @@ -21,6 +22,7 @@ use super::rig_tool::action_from_tool_call; use super::ProviderRunProfile; use crate::ai::agent::api::RequestParams; use crate::ai::agent::AIAgentAction; +use crate::ai::openai::client::OpenAIClientConfig; use crate::ai::provider::client::BedrockClient; use crate::ai::provider::convert::CachingConfig; use crate::ai::provider::external_config::ExternalBedrockConfig; @@ -28,6 +30,97 @@ use crate::ai::provider::types::ConversationMessage; use crate::ai::runtime::RuntimeResponseConfig; use crate::settings::OpenAIProviderKind; +const PROVIDER_SELECTION_METADATA: &str = "galaxy_provider_selection"; + +/// Keep routing identity separate from the wire model name. This metadata is +/// checkpointed with each request, but is never included in provider payloads. +#[derive(serde::Serialize, serde::Deserialize)] +#[serde(deny_unknown_fields)] +struct PersistedProviderSelection { + model_id: String, + connection_fingerprint: String, + reasoning: PersistedReasoningSelection, +} + +#[derive(serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +enum PersistedReasoningSelection { + ProviderDefault, + Explicit(String), +} + +fn provider_connection_fingerprint(config: &OpenAIClientConfig) -> String { + // Credentials are resolved afresh on restore. Hash only connection metadata + // so endpoint URLs do not appear in checkpoints either. + let identity = serde_json::json!([ + config.connection_identity, + format!("{:?}", config.kind), + config.base_url, + config.project_id, + config.location, + ]); + format!("{:x}", Sha256::digest(identity.to_string().as_bytes())) +} + +fn remember_provider_selection( + request: &mut TurnRequest, + model_id: String, + provider_config: &crate::ai::provider::ProviderConfig, +) -> anyhow::Result<()> { + if let crate::ai::provider::ProviderConfig::OpenAI(config) = provider_config { + request.metadata.insert( + PROVIDER_SELECTION_METADATA.to_string(), + serde_json::to_value(PersistedProviderSelection { + model_id, + connection_fingerprint: provider_connection_fingerprint(config), + reasoning: config.reasoning_effort.clone().map_or( + PersistedReasoningSelection::ProviderDefault, + PersistedReasoningSelection::Explicit, + ), + })?, + ); + } + Ok(()) +} + +pub(crate) fn restore_provider_config( + request: &TurnRequest, + resolve: impl FnOnce(&str) -> crate::ai::provider::ProviderConfig, +) -> anyhow::Result { + let Some(selection) = request.metadata.get(PROVIDER_SELECTION_METADATA) else { + // Older checkpoints only contain the wire model ID. + return Ok(resolve(request.model.as_str())); + }; + let selection: PersistedProviderSelection = serde_json::from_value(selection.clone()) + .map_err(|_| anyhow::anyhow!("invalid persisted provider selection"))?; + anyhow::ensure!( + !selection.model_id.trim().is_empty(), + "empty persisted provider selection" + ); + let crate::ai::provider::ProviderConfig::OpenAI(mut config) = resolve(&selection.model_id) + else { + anyhow::bail!("the saved provider model is no longer available"); + }; + anyhow::ensure!( + provider_connection_fingerprint(&config) == selection.connection_fingerprint, + "the saved provider connection has changed" + ); + let wire_model = config + .model + .as_deref() + .filter(|model| !model.is_empty() && *model != "auto") + .unwrap_or(&selection.model_id); + anyhow::ensure!( + wire_model == request.model.as_str(), + "the saved provider model no longer matches its request" + ); + config.reasoning_effort = match selection.reasoning { + PersistedReasoningSelection::ProviderDefault => None, + PersistedReasoningSelection::Explicit(effort) => Some(effort), + }; + Ok(crate::ai::provider::ProviderConfig::OpenAI(config)) +} + pub(crate) struct PreparedProviderRun { pub(crate) base_profile: ProviderRunProfile, pub(crate) cli_monitor_profile: Option, @@ -264,7 +357,7 @@ async fn prepare_provider_profile( mode: Option, ) -> anyhow::Result<(Arc, PreparedRigTurn)> { let model = params.model.as_str().to_string(); - let prepared = match &provider_config { + let mut prepared = match &provider_config { crate::ai::provider::ProviderConfig::OpenAI(config) => match mode { Some(mode) => prepare_rig_turn_for_mode( config, @@ -278,7 +371,7 @@ async fn prepare_provider_profile( crate::ai::provider::ProviderConfig::Bedrock(_) => { let max_output_tokens = Some(bedrock_max_output_tokens(&model)); prepare_bedrock_rig_turn_for_mode( - model, + model.clone(), max_output_tokens, params, supported_tools, @@ -292,6 +385,7 @@ async fn prepare_provider_profile( ); } }; + remember_provider_selection(&mut prepared.request, model, &provider_config)?; let runtime = provider_runtime_for_request(provider_config, &prepared.request).await?; Ok((runtime, prepared)) } diff --git a/app/src/ai/runtime/rig_request_tests.rs b/app/src/ai/runtime/rig_request_tests.rs index 830fb161..54a46fc6 100644 --- a/app/src/ai/runtime/rig_request_tests.rs +++ b/app/src/ai/runtime/rig_request_tests.rs @@ -24,6 +24,7 @@ use crate::ai::skills::SkillDescriptor; fn config() -> OpenAIClientConfig { OpenAIClientConfig { kind: crate::settings::OpenAIProviderKind::LiteLLM, + connection_identity: None, base_url: "http://localhost:4000/v1".to_string(), api_key: None, project_id: None, diff --git a/app/src/ai/runtime/rig_tests.rs b/app/src/ai/runtime/rig_tests.rs index b157252a..1b78f0df 100644 --- a/app/src/ai/runtime/rig_tests.rs +++ b/app/src/ai/runtime/rig_tests.rs @@ -1,4 +1,4 @@ -use galaxy_agent_core::ToolCall; +use galaxy_agent_core::{ToolCall, TurnRequest}; use super::{prepare_provider_run, ProviderActionContext}; use crate::ai::agent::api::RequestParams; @@ -10,6 +10,7 @@ use crate::settings::OpenAIProviderKind; fn openai_config(model: &str) -> OpenAIClientConfig { OpenAIClientConfig { kind: OpenAIProviderKind::LiteLLM, + connection_identity: Some("LiteLLM test connection".to_string()), base_url: "http://localhost:4000/v1".to_string(), api_key: None, project_id: None, @@ -42,6 +43,11 @@ async fn missing_cli_provider_route_falls_back_to_base_provider_profile() { .expect("base provider should supply the CLI monitor fallback"); assert_eq!(cli_profile.request.model.as_str(), "base-provider-model"); + super::restore_provider_config(&cli_profile.request, |model| { + assert_eq!(model, "base-selection"); + ProviderConfig::OpenAI(openai_config("base-provider-model")) + }) + .unwrap(); assert!(cli_profile .request .system_prompt @@ -96,18 +102,187 @@ async fn automatic_and_same_model_computer_use_keep_a_dedicated_profile() { params.computer_use_enabled = true; let prepared = prepare_provider_run( ProviderConfig::OpenAI(openai_config("base-provider-model")), - ProviderConfig::None, ProviderConfig::None, params, Vec::new(), - ).await.unwrap(); - assert!(prepared.base_profile.request.tools.iter().any(|tool| tool.name == "request_computer_use")); - assert!(!prepared.base_profile.request.tools.iter().any(|tool| tool.name == "use_computer")); + ProviderConfig::None, + ProviderConfig::None, + params, + Vec::new(), + ) + .await + .unwrap(); + assert!(prepared + .base_profile + .request + .tools + .iter() + .any(|tool| tool.name == "request_computer_use")); + assert!(!prepared + .base_profile + .request + .tools + .iter() + .any(|tool| tool.name == "use_computer")); let computer = prepared.computer_use_profile.unwrap(); assert_eq!(computer.request.model.as_str(), "base-provider-model"); - assert!(computer.request.tools.iter().any(|tool| tool.name == "use_computer")); - assert!(!computer.request.tools.iter().any(|tool| tool.name == "request_computer_use")); - assert!(computer.request.system_prompt.unwrap().contains("Inspect the latest returned screenshot before acting")); + super::restore_provider_config(&computer.request, |model| { + assert_eq!(model, "base-selection"); + ProviderConfig::OpenAI(openai_config("base-provider-model")) + }) + .unwrap(); + assert!(computer + .request + .tools + .iter() + .any(|tool| tool.name == "use_computer")); + assert!(!computer + .request + .tools + .iter() + .any(|tool| tool.name == "request_computer_use")); + assert!(computer + .request + .system_prompt + .unwrap() + .contains("Inspect the latest returned screenshot before acting")); } } +#[tokio::test] +async fn restored_profiles_keep_their_selection_and_effort_with_refreshed_credentials() { + let mut params = RequestParams::new_for_test(); + params.model = LLMId::from("base-selection::reasoning::high"); + params.cli_agent_model = LLMId::from("cli-selection::reasoning::low"); + params.computer_use_model = LLMId::from("computer-selection::reasoning::medium"); + params.computer_use_enabled = true; + let config = |model, effort: &str| { + let mut config = openai_config(model); + config.reasoning_effort = Some(effort.to_string()); + config.api_key = Some("old-secret".to_string()); + ProviderConfig::OpenAI(config) + }; + let prepared = prepare_provider_run( + config("base-provider-model", "high"), + config("cli-provider-model", "low"), + config("computer-provider-model", "medium"), + params, + Vec::new(), + ) + .await + .unwrap(); + + for (profile, selection, model, effort) in [ + ( + prepared.base_profile, + "base-selection::reasoning::high", + "base-provider-model", + "high", + ), + ( + prepared.cli_monitor_profile.unwrap(), + "cli-selection::reasoning::low", + "cli-provider-model", + "low", + ), + ( + prepared.computer_use_profile.unwrap(), + "computer-selection::reasoning::medium", + "computer-provider-model", + "medium", + ), + ] { + let serialized = serde_json::to_string(&profile.request).unwrap(); + assert!(!serialized.contains("old-secret")); + assert!(!serialized.contains("http://localhost:4000")); + let request: TurnRequest = serde_json::from_str(&serialized).unwrap(); + assert_eq!(request.model.as_str(), model); + let restored = super::restore_provider_config(&request, |id| { + assert_eq!(id, selection); + let mut current = openai_config(model); + current.api_key = Some("new-secret".to_string()); + current.reasoning_effort = Some("changed-after-checkpoint".to_string()); + ProviderConfig::OpenAI(current) + }) + .unwrap(); + let ProviderConfig::OpenAI(restored) = restored else { + panic!("expected the original provider"); + }; + assert_eq!(restored.reasoning_effort.as_deref(), Some(effort)); + assert_eq!(restored.api_key.as_deref(), Some("new-secret")); + } +} + +#[test] +fn restored_provider_default_clears_a_later_effort_setting() { + let mut request = TurnRequest::new("wire-model", Vec::new()); + super::remember_provider_selection( + &mut request, + "model-selection".to_string(), + &ProviderConfig::OpenAI(openai_config("wire-model")), + ) + .unwrap(); + let ProviderConfig::OpenAI(restored) = super::restore_provider_config(&request, |_| { + let mut current = openai_config("wire-model"); + current.reasoning_effort = Some("high".to_string()); + ProviderConfig::OpenAI(current) + }) + .unwrap() else { + panic!("expected an OpenAI-compatible provider"); + }; + assert_eq!(restored.reasoning_effort, None); +} + +#[test] +fn restored_selection_rejects_unavailable_changed_and_malformed_routes() { + let mut request = TurnRequest::new("wire-model", Vec::new()); + super::remember_provider_selection( + &mut request, + "model-selection".to_string(), + &ProviderConfig::OpenAI(openai_config("wire-model")), + ) + .unwrap(); + assert!(super::restore_provider_config(&request, |_| ProviderConfig::None).is_err()); + assert!(super::restore_provider_config(&request, |_| { + let mut current = openai_config("wire-model"); + current.base_url = "https://another-provider.example/v1".to_string(); + ProviderConfig::OpenAI(current) + }) + .is_err()); + assert!(super::restore_provider_config(&request, |_| { + let mut current = openai_config("wire-model"); + current.connection_identity = Some("different-account-at-same-endpoint".to_string()); + ProviderConfig::OpenAI(current) + }) + .is_err()); + assert!(super::restore_provider_config(&request, |_| { + ProviderConfig::OpenAI(openai_config("different-wire-model")) + }) + .is_err()); + + request + .metadata + .get_mut(super::PROVIDER_SELECTION_METADATA) + .unwrap() + .as_object_mut() + .unwrap() + .remove("reasoning"); + assert!(super::restore_provider_config(&request, |_| { + panic!("malformed selection must fail before looking up a provider") + }) + .is_err()); +} + +#[test] +fn legacy_request_without_selection_metadata_keeps_raw_model_lookup() { + let request = TurnRequest::new("legacy-model", Vec::new()); + let ProviderConfig::OpenAI(restored) = super::restore_provider_config(&request, |model| { + assert_eq!(model, "legacy-model"); + ProviderConfig::OpenAI(openai_config("legacy-model")) + }) + .unwrap() else { + panic!("expected legacy provider lookup to succeed"); + }; + assert_eq!(restored.model.as_deref(), Some("legacy-model")); +} + #[test] fn subscription_context_budget_uses_catalog_limit_and_clamps_stale_overrides() { let mut config = openai_config("subscription-model"); diff --git a/app/src/server/server_api/ai.rs b/app/src/server/server_api/ai.rs index ca8ebfc9..7c75cc7b 100644 --- a/app/src/server/server_api/ai.rs +++ b/app/src/server/server_api/ai.rs @@ -3026,6 +3026,7 @@ impl From for LLMIn display_name: value.display_name, base_model_name: value.base_model_name, reasoning_level: value.reasoning_level, + reasoning_model_id: None, usage_metadata: value.usage_metadata.into(), description: value.description, disable_reason: value.disable_reason.map(DisableReason::from), @@ -3065,6 +3066,7 @@ impl From for LLMInfo { display_name: value.display_name, base_model_name: value.base_model_name, reasoning_level: value.reasoning_level, + reasoning_model_id: None, usage_metadata: value.usage_metadata.into(), description: value.description, disable_reason: value.disable_reason.map(DisableReason::from),