From 4d251f6cd0adc143f823b77ea19c2adcc12d3e9e Mon Sep 17 00:00:00 2001 From: Jeff <36680501@qq.com> Date: Sat, 22 Aug 2026 17:44:32 +0800 Subject: [PATCH] feat(aionrs): hot-swap model without agent rebuild Same-provider model changes via config-options/model and PATCH now call AgentEngine::apply_config_update instead of killing the task. Enables Auto planner/worker routing (AionUi #4143 Phase 2). Co-authored-by: Cursor --- crates/aionui-ai-agent/src/agent_task.rs | 26 ++++ .../src/manager/aionrs/agent.rs | 117 +++++++++++++++--- crates/aionui-conversation/src/service.rs | 64 ++++++++-- 3 files changed, 184 insertions(+), 23 deletions(-) diff --git a/crates/aionui-ai-agent/src/agent_task.rs b/crates/aionui-ai-agent/src/agent_task.rs index 193ab4c95..9db10a47e 100644 --- a/crates/aionui-ai-agent/src/agent_task.rs +++ b/crates/aionui-ai-agent/src/agent_task.rs @@ -695,6 +695,14 @@ mod aionrs_config_option_tests { .collect::>(), vec!["default", "auto_edit", "yolo"] ); + + let model = response + .config_options + .iter() + .find(|option| option.id == "model") + .expect("aionrs should expose a model config option for hot-swap"); + assert_eq!(model.category.as_deref(), Some("model")); + assert_eq!(model.current_value.as_deref(), Some("claude-sonnet-4-20250514")); } #[tokio::test] @@ -738,6 +746,24 @@ mod aionrs_config_option_tests { "unexpected error: {error:?}" ); } + + #[tokio::test] + async fn aionrs_set_config_option_model_hot_swaps_without_error() { + let instance = aionrs_instance().await; + + let response = instance + .set_config_option("model", "claude-haiku-4-20250514") + .await + .unwrap(); + + assert_eq!(response.confirmation, ConfigOptionConfirmation::Observed); + let options = response.config_options.expect("snapshot"); + let model = options + .iter() + .find(|option| option.id == "model") + .expect("model option"); + assert_eq!(model.current_value.as_deref(), Some("claude-haiku-4-20250514")); + } } #[cfg(test)] diff --git a/crates/aionui-ai-agent/src/manager/aionrs/agent.rs b/crates/aionui-ai-agent/src/manager/aionrs/agent.rs index 3291d92f5..8e56a0b96 100644 --- a/crates/aionui-ai-agent/src/manager/aionrs/agent.rs +++ b/crates/aionui-ai-agent/src/manager/aionrs/agent.rs @@ -119,6 +119,15 @@ pub struct AionrsAgentManager { approval_manager: Arc, confirmations: Arc>>, final_input_dump: Option, + /// Provider type label used when resolving image-input capability after a + /// same-provider model hot-swap (openai / anthropic / …). + provider: String, + base_url: Option, + /// Model id currently applied to the engine (or queued for the next turn). + current_model: Mutex, + /// When a model switch arrives mid-turn, apply it before the next + /// `engine.run_with_blocks` instead of tearing down the agent. + pending_model: Mutex>, /// Signalled by `cancel()` to abort an in-flight `engine.run()` via /// `tokio::select!` in `send_message()`. cancel_notify: Arc, @@ -282,6 +291,10 @@ impl AionrsAgentManager { approval_manager, confirmations, final_input_dump, + provider: config_extra.provider.clone(), + base_url: config_extra.base_url.clone(), + current_model: Mutex::new(config_extra.model.clone()), + pending_model: Mutex::new(None), cancel_notify: Arc::new(Notify::new()), turn_finished_notify: Arc::new(Notify::new()), }) @@ -421,6 +434,7 @@ impl IAgentTask for AionrsAgentManager { .send(AgentStreamEvent::BackendTurnBound(turn_anchor.clone())); let mut engine = self.engine.lock().await; + self.apply_pending_model_locked(&mut engine).await; engine.set_next_turn_id(Some(turn_anchor)); let result = tokio::select! { @@ -586,8 +600,12 @@ impl AionrsAgentManager { } pub async fn config_options(&self) -> Result { + let current_model = self.current_model.lock().await.clone(); Ok(GetConfigOptionsResponse { - config_options: vec![aionrs_mode_config_option(self.approval_manager.current_mode())], + config_options: vec![ + aionrs_mode_config_option(self.approval_manager.current_mode()), + aionrs_model_config_option(¤t_model), + ], }) } @@ -595,30 +613,82 @@ impl AionrsAgentManager { let option_id = option_id.trim(); let value = value.trim(); - if option_id != AIONRS_MODE_OPTION_ID { - return Err(AgentError::bad_request(format!( + match option_id { + AIONRS_MODE_OPTION_ID => { + if !is_aionrs_session_mode(value) { + return Err(AgentError::bad_request(format!( + "Value '{value}' is not selectable for config option '{option_id}'" + ))); + } + self.set_mode(value).await?; + Ok(SetConfigOptionResponse { + confirmation: ConfigOptionConfirmation::Observed, + config_options: Some(self.config_options().await?.config_options), + }) + } + AIONRS_MODEL_OPTION_ID => { + if value.is_empty() { + return Err(AgentError::bad_request(format!( + "Value '{value}' is not selectable for config option '{option_id}'" + ))); + } + let confirmation = self.queue_or_apply_model(value).await?; + Ok(SetConfigOptionResponse { + confirmation, + config_options: Some(self.config_options().await?.config_options), + }) + } + _ => Err(AgentError::bad_request(format!( "Config option '{option_id}' is not available" - ))); - } - if !is_aionrs_session_mode(value) { - return Err(AgentError::bad_request(format!( - "Value '{value}' is not selectable for config option '{option_id}'" - ))); + ))), } - - self.set_mode(value).await?; - Ok(SetConfigOptionResponse { - confirmation: ConfigOptionConfirmation::Observed, - config_options: Some(self.config_options().await?.config_options), - }) } pub async fn get_slash_commands(&self) -> Result, AgentError> { Ok(self.slash_commands.clone()) } + + /// Same-provider hot-swap: updates `AgentEngine.model` without rebuild. + /// Mid-turn requests are queued and applied before the next send. + async fn queue_or_apply_model(&self, model: &str) -> Result { + if self.runtime.status() == Some(ConversationStatus::Running) { + *self.pending_model.lock().await = Some(model.to_owned()); + *self.current_model.lock().await = model.to_owned(); + info!( + conversation_id = %self.runtime.conversation_id(), + model, + "Aionrs model switch queued for next turn" + ); + return Ok(ConfigOptionConfirmation::PendingNextTurn); + } + + let mut engine = self.engine.lock().await; + self.apply_model_locked(&mut engine, model).await; + Ok(ConfigOptionConfirmation::Observed) + } + + async fn apply_pending_model_locked(&self, engine: &mut AgentEngine) { + let pending = self.pending_model.lock().await.take(); + if let Some(model) = pending { + self.apply_model_locked(engine, &model).await; + } + } + + async fn apply_model_locked(&self, engine: &mut AgentEngine, model: &str) { + let image_input = resolve_image_input_capability(&self.provider, self.base_url.as_deref(), model); + let changes = engine.apply_config_update(Some(model.to_owned()), Some(image_input), None, None, None, None); + *self.current_model.lock().await = model.to_owned(); + info!( + conversation_id = %self.runtime.conversation_id(), + model, + ?changes, + "Aionrs model hot-swapped without rebuild" + ); + } } const AIONRS_MODE_OPTION_ID: &str = "mode"; +const AIONRS_MODEL_OPTION_ID: &str = "model"; fn is_aionrs_session_mode(s: &str) -> bool { matches!(s, "default" | "auto_edit" | "yolo") @@ -641,6 +711,23 @@ fn aionrs_mode_config_option(current_value: String) -> AcpConfigOptionDto { } } +fn aionrs_model_config_option(current_value: &str) -> AcpConfigOptionDto { + AcpConfigOptionDto { + id: AIONRS_MODEL_OPTION_ID.to_owned(), + name: Some("Model".to_owned()), + label: None, + description: Some( + "Hot-swap the upstream model on the same aionrs agent (same provider). Used by Auto planner/worker routing." + .to_owned(), + ), + category: Some("model".to_owned()), + option_type: "select".to_owned(), + current_value: Some(current_value.to_owned()), + // Catalog lives in AionUi; Core accepts any non-empty id for BYOK. + options: vec![aionrs_mode_select_option(current_value, current_value)], + } +} + fn aionrs_mode_select_option(value: &str, name: &str) -> AcpConfigSelectOptionDto { AcpConfigSelectOptionDto { value: value.to_owned(), diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index e37b2a119..76f4b458f 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -30,8 +30,8 @@ use aionui_api_types::{ use aionui_api_types::{ChatFileRef, SessionRef}; use aionui_common::{ AgentKillReason, AgentType, ConversationSource, ConversationStatus, ErrorChain, MessageType, OnConversationDelete, - OnConversationTurnCancelled, PaginatedResult, TurnCancelCause, WorkspacePathValidationError, generate_short_id, - now_ms, validate_workspace_path_availability, + OnConversationTurnCancelled, PaginatedResult, ProviderWithModel, TurnCancelCause, WorkspacePathValidationError, + generate_short_id, now_ms, validate_workspace_path_availability, }; use aionui_db::models::{ AssistantDefinitionRow, ConversationAssistantSnapshotRow, ConversationRow, McpServerRow, MessageRow, @@ -2408,12 +2408,41 @@ impl ConversationService { } if model_changed { - info!( - model_changed = true, - "Conversation updated, killing agent task due to model change" - ); - if let Err(e) = task_manager.kill(id, None) { - warn!(error = %ErrorChain(&e), "Failed to kill agent after model change"); + let same_provider = aionrs_model_change_keeps_provider(existing.model.as_deref(), req.model.as_ref()); + if same_provider { + if let Some(model) = req.model.as_ref() { + let selected_model = model.use_model.as_deref().unwrap_or(model.model.as_str()); + if let Some(agent) = task_manager.get_task(id) { + match agent.set_config_option("model", selected_model).await { + Ok(response) => { + info!( + conversation_id = %id, + model = %selected_model, + confirmation = ?response.confirmation, + "Aionrs same-provider model change applied without rebuild" + ); + } + Err(e) => { + warn!( + conversation_id = %id, + error = %ErrorChain(&e), + "Failed to hot-swap aionrs model; falling back to kill/rebuild" + ); + if let Err(kill_err) = task_manager.kill(id, None) { + warn!(error = %ErrorChain(&kill_err), "Failed to kill agent after model hot-swap failure"); + } + } + } + } + } + } else { + info!( + model_changed = true, + "Conversation updated, killing agent task due to model/provider change" + ); + if let Err(e) = task_manager.kill(id, None) { + warn!(error = %ErrorChain(&e), "Failed to kill agent after model change"); + } } } @@ -5930,6 +5959,25 @@ fn log_conversation_created(response: &ConversationResponse, extra: &serde_json: } } +/// Same `provider_id` means the aionrs engine can hot-swap `use_model` without +/// tearing down the agent (Auto planner/worker routing). Provider changes still +/// require kill + rebuild. +fn aionrs_model_change_keeps_provider( + existing_model_json: Option<&str>, + new_model: Option<&ProviderWithModel>, +) -> bool { + let Some(new_model) = new_model else { + return false; + }; + let Some(existing_json) = existing_model_json else { + return false; + }; + let Ok(existing) = serde_json::from_str::(existing_json) else { + return false; + }; + !existing.provider_id.is_empty() && existing.provider_id == new_model.provider_id +} + fn is_tool_message_type(message_type: MessageType) -> bool { matches!( message_type,