Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src-rust/crates/cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -740,6 +740,7 @@ async fn main() -> anyhow::Result<()> {

// Build query config
let mut query_config = claurst_query::QueryConfig::from_config_with_registry(&config, &model_registry);
query_config.preserve_selected_model = cli.model.is_some() || cli.provider.is_some();
query_config.model_registry = Some(model_registry.clone());
query_config.max_turns = cli.max_turns;
query_config.system_prompt = Some(system_prompt);
Expand Down
1 change: 1 addition & 0 deletions src-rust/crates/query/src/agent_tool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -367,6 +367,7 @@ impl Tool for AgentTool {
agent_definition: None,
model_registry: Some(model_registry),
managed_agents: None,
preserve_selected_model: false,
};
// -----------------------------------------------------------------------
// Background mode: spawn and return agent_id immediately.
Expand Down
75 changes: 63 additions & 12 deletions src-rust/crates/query/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,9 @@ pub struct QueryConfig {
pub model_registry: Option<std::sync::Arc<claurst_api::ModelRegistry>>,
/// Managed agent (manager-executor) configuration.
pub managed_agents: Option<claurst_core::ManagedAgentConfig>,
/// Preserve an explicit caller-selected model instead of allowing managed-agent
/// manager_model to replace it at query runtime.
pub preserve_selected_model: bool,
}

impl Default for QueryConfig {
Expand All @@ -152,6 +155,7 @@ impl Default for QueryConfig {
agent_definition: None,
model_registry: None,
managed_agents: None,
preserve_selected_model: false,
}
}
}
Expand Down Expand Up @@ -697,6 +701,26 @@ fn should_emit_turn_complete(stop: &str, max_tokens_recovery_count: u32) -> bool
}
}

fn selected_model_for_query(config: &QueryConfig) -> String {
let selected_model = if let Some(ref agent) = config.agent_definition {
agent.model.clone().unwrap_or_else(|| config.model.clone())
} else {
config.model.clone()
};

if config.preserve_selected_model {
return selected_model;
}

if let Some(ref ma_config) = config.managed_agents {
if ma_config.enabled && !ma_config.manager_model.is_empty() {
return ma_config.manager_model.clone();
}
}

selected_model
}

// Spinner verbs are imported from claurst_core::spinner

/// Run the agentic query loop.
Expand Down Expand Up @@ -726,18 +750,7 @@ pub async fn run_query_loop(
let mut max_tokens_recovery_count: u32 = 0;
// Active model — may switch to fallback on overloaded errors.
// Agent model override takes priority over the session model when set.
let mut effective_model = if let Some(ref agent) = config.agent_definition {
agent.model.clone().unwrap_or_else(|| config.model.clone())
} else {
config.model.clone()
};

// If managed-agent mode is active, override the model to the manager model.
if let Some(ref ma_config) = config.managed_agents {
if ma_config.enabled && !ma_config.manager_model.is_empty() {
effective_model = ma_config.manager_model.clone();
}
}
let mut effective_model = selected_model_for_query(config);

let mut used_fallback = false;
// How many automatic retries remain when a stream stalls (no data for 45s).
Expand Down Expand Up @@ -2296,6 +2309,7 @@ mod tests {
agent_definition: None,
model_registry: None,
managed_agents: None,
preserve_selected_model: false,
}
}

Expand Down Expand Up @@ -2411,6 +2425,43 @@ mod tests {
assert_eq!(cloned.system_prompt, Some("test".to_string()));
}

#[test]
fn test_managed_agent_manager_model_overrides_by_default() {
let mut cfg = make_config(None, None);
cfg.managed_agents = Some(claurst_core::ManagedAgentConfig {
enabled: true,
manager_model: "openai/gpt-4o".to_string(),
executor_model: "anthropic/claude-sonnet-4-6".to_string(),
executor_max_turns: 10,
max_concurrent_executors: 4,
budget_split: claurst_core::BudgetSplitPolicy::SharedPool,
total_budget_usd: None,
preset_name: None,
executor_isolation: false,
});

assert_eq!(selected_model_for_query(&cfg), "openai/gpt-4o");
}

#[test]
fn test_explicit_model_selection_blocks_managed_agent_override() {
let mut cfg = make_config(None, None);
cfg.preserve_selected_model = true;
cfg.managed_agents = Some(claurst_core::ManagedAgentConfig {
enabled: true,
manager_model: "openai/gpt-4o".to_string(),
executor_model: "anthropic/claude-sonnet-4-6".to_string(),
executor_max_turns: 10,
max_concurrent_executors: 4,
budget_split: claurst_core::BudgetSplitPolicy::SharedPool,
total_budget_usd: None,
preset_name: None,
executor_isolation: false,
});

assert_eq!(selected_model_for_query(&cfg), "claude-sonnet-4-6");
}

// ---- QueryOutcome variant tests -----------------------------------------

#[test]
Expand Down