public/galaxy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix reasoning model identity and restored selections

Browse Source
This commit is contained in:
Ryan Ward
2026-10-25 07:46:35 +00:00
parent 6655fe7822
commit 2e74b00d5d
13 changed files with 571 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/src/ai/runtime/mod.rs
+1 -1
View File
@@ -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};
app/src/ai/runtime/rig.rs
+96 -2
View File
@@ -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<crate::ai::provider::ProviderConfig> {
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<ProviderRunProfile>,
@@ -264,7 +357,7 @@ async fn prepare_provider_profile(
mode: Option<RigRequestMode>,
) -> anyhow::Result<(Arc<dyn AgentRuntime>, 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))
}
app/src/ai/runtime/rig_request_tests.rs
+1
View File
@@ -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,
app/src/ai/runtime/rig_tests.rs
+183 -8
View File
@@ -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");