From 57a3ffa745aed3d0e92bed87ca17370764d21d85 Mon Sep 17 00:00:00 2001 From: Inference_ <68734681+Inference1@users.noreply.github.com> Date: Tue, 14 Jul 2026 01:04:41 +0800 Subject: [PATCH] fix(team): persist models and queue background messages --- crates/aionui-app/tests/team_model_e2e.rs | 203 ++++++++++++++++++ crates/aionui-team/src/routes.rs | 22 +- .../src/scheduler/agent_lifecycle.rs | 10 + crates/aionui-team/src/service.rs | 75 +++++-- crates/aionui-team/src/session.rs | 4 + crates/aionui-team/tests/e2e_smoke.rs | 2 +- .../tests/session_service_integration.rs | 87 +++++++- 7 files changed, 378 insertions(+), 25 deletions(-) create mode 100644 crates/aionui-app/tests/team_model_e2e.rs diff --git a/crates/aionui-app/tests/team_model_e2e.rs b/crates/aionui-app/tests/team_model_e2e.rs new file mode 100644 index 000000000..ae565c4f2 --- /dev/null +++ b/crates/aionui-app/tests/team_model_e2e.rs @@ -0,0 +1,203 @@ +mod common; + +use axum::http::StatusCode; +use serde_json::{Value, json}; +use tower::ServiceExt; + +use common::{body_json, build_app, json_with_token, setup_and_login}; + +const TEAM_ASSISTANT_ID: &str = "team-model-e2e-assistant"; +const TEAM_AGENT_ID: &str = "2d23ff1c"; + +async fn create_team(app: &mut axum::Router, services: &aionui_app::AppServices, token: &str, csrf: &str) -> Value { + let command = std::env::current_exe() + .expect("test executable path") + .to_string_lossy() + .to_string(); + let source_info = json!({ "binary_name": command }).to_string(); + sqlx::query( + "UPDATE agent_metadata \ + SET agent_source = 'custom', agent_source_info = ?, command = ?, args = '[]', env = '[]', \ + updated_at = unixepoch('now','subsec') * 1000 \ + WHERE id = ?", + ) + .bind(&source_info) + .bind(&command) + .bind(TEAM_AGENT_ID) + .execute(services.database.pool()) + .await + .expect("seed deterministic team agent"); + services + .agent_registry + .reload_one(TEAM_AGENT_ID) + .await + .expect("reload deterministic team agent"); + + let req = json_with_token( + "POST", + "/api/assistants", + json!({ + "id": TEAM_ASSISTANT_ID, + "name": "Team Model E2E Assistant", + "agent_id": TEAM_AGENT_ID + }), + token, + csrf, + ); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::CREATED); + + let req = json_with_token( + "POST", + "/api/teams", + json!({ + "name": "Model Team", + "agents": [ + { + "name": "Lead", + "role": "lead", + "model": "gpt-5.5", + "assistant_id": TEAM_ASSISTANT_ID + }, + { + "name": "Worker", + "role": "teammate", + "model": "gpt-5.5", + "assistant_id": TEAM_ASSISTANT_ID + } + ] + }), + token, + csrf, + ); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::CREATED); + body_json(resp).await["data"].clone() +} + +#[tokio::test] +async fn update_agent_model_persists_the_team_roster_value() { + let (mut app, services) = build_app().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let team = create_team(&mut app, &services, &token, &csrf).await; + let team_id = team["id"].as_str().unwrap(); + let slot_id = team["assistants"][1]["slot_id"].as_str().unwrap(); + + let req = json_with_token( + "PATCH", + &format!("/api/teams/{team_id}/agents/{slot_id}/model"), + json!({ "model_id": "gpt-5.6-sol" }), + &token, + &csrf, + ); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let req = common::get_with_token(&format!("/api/teams/{team_id}"), &token); + let resp = app.oneshot(req).await.unwrap(); + let body = body_json(resp).await; + let worker = body["data"]["assistants"] + .as_array() + .unwrap() + .iter() + .find(|agent| agent["slot_id"] == slot_id) + .unwrap(); + assert_eq!(worker["model"], "gpt-5.6-sol"); +} + +#[tokio::test] +async fn update_agent_model_requires_authentication() { + let (app, _services) = build_app().await; + let csrf = "csrf-test"; + let req = axum::http::Request::builder() + .method("PATCH") + .uri("/api/teams/team-1/agents/worker-1/model") + .header("content-type", "application/json") + .header("x-csrf-token", csrf) + .header("cookie", format!("aionui-csrf-token={csrf}")) + .body(axum::body::Body::from(r#"{"model_id":"gpt-5.6-sol"}"#)) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + assert_eq!(body_json(resp).await["code"], "UNAUTHORIZED"); +} + +#[tokio::test] +async fn update_agent_model_requires_csrf() { + let (mut app, services) = build_app().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let team = create_team(&mut app, &services, &token, &csrf).await; + let team_id = team["id"].as_str().unwrap(); + let slot_id = team["assistants"][1]["slot_id"].as_str().unwrap(); + let req = axum::http::Request::builder() + .method("PATCH") + .uri(format!("/api/teams/{team_id}/agents/{slot_id}/model")) + .header("authorization", format!("Bearer {token}")) + .header("content-type", "application/json") + .body(axum::body::Body::from(r#"{"model_id":"gpt-5.6-sol"}"#)) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + assert_eq!(body_json(resp).await["code"], "CSRF_INVALID"); +} + +#[tokio::test] +async fn update_agent_model_rejects_cross_user_access() { + let (mut app, services) = build_app().await; + let (owner_token, owner_csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let team = create_team(&mut app, &services, &owner_token, &owner_csrf).await; + let team_id = team["id"].as_str().unwrap(); + let slot_id = team["assistants"][1]["slot_id"].as_str().unwrap(); + let (other_token, other_csrf) = setup_and_login(&mut app, &services, "other", "StrongP@ss2").await; + let req = json_with_token( + "PATCH", + &format!("/api/teams/{team_id}/agents/{slot_id}/model"), + json!({ "model_id": "gpt-5.6-sol" }), + &other_token, + &other_csrf, + ); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + assert_eq!(body_json(resp).await["code"], "FORBIDDEN"); +} + +#[tokio::test] +async fn update_agent_model_rejects_empty_models() { + let (mut app, services) = build_app().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let team = create_team(&mut app, &services, &token, &csrf).await; + let team_id = team["id"].as_str().unwrap(); + let slot_id = team["assistants"][1]["slot_id"].as_str().unwrap(); + + let empty = json_with_token( + "PATCH", + &format!("/api/teams/{team_id}/agents/{slot_id}/model"), + json!({ "model_id": " " }), + &token, + &csrf, + ); + let resp = app.clone().oneshot(empty).await.unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + assert_eq!(body_json(resp).await["code"], "BAD_REQUEST"); +} + +#[tokio::test] +async fn update_agent_model_rejects_missing_agents() { + let (mut app, services) = build_app().await; + let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let team = create_team(&mut app, &services, &token, &csrf).await; + let team_id = team["id"].as_str().unwrap(); + let missing = json_with_token( + "PATCH", + &format!("/api/teams/{team_id}/agents/missing/model"), + json!({ "model_id": "gpt-5.6-sol" }), + &token, + &csrf, + ); + let resp = app.oneshot(missing).await.unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + assert_eq!(body_json(resp).await["code"], "NOT_FOUND"); +} diff --git a/crates/aionui-team/src/routes.rs b/crates/aionui-team/src/routes.rs index 0bab888ff..9f9f5ceaa 100644 --- a/crates/aionui-team/src/routes.rs +++ b/crates/aionui-team/src/routes.rs @@ -12,8 +12,8 @@ use aionui_ai_agent::ActiveLeaseRegistry; use aionui_api_types::{ AddAgentRequest, ApiResponse, CancelTeamChildTurnRequest, CancelTeamRunRequest, CreateTeamRequest, GetConfigOptionsResponse, PauseTeamSlotRequest, RenameAgentRequest, RenameTeamRequest, SendAgentMessageRequest, - SendTeamMessageRequest, SetModeRequest, TeamAgentResponse, TeamListResponse, TeamResponse, TeamRunAckResponse, - TeamRunStateResponse, + SendTeamMessageRequest, SetModeRequest, SetModelRequest, TeamAgentResponse, TeamListResponse, TeamResponse, + TeamRunAckResponse, TeamRunStateResponse, }; use aionui_auth::CurrentUser; use aionui_common::ApiError; @@ -99,6 +99,10 @@ pub fn team_routes(state: TeamRouterState) -> Router { "/api/teams/{id}/agents/{slot_id}/name", axum::routing::patch(rename_agent), ) + .route( + "/api/teams/{id}/agents/{slot_id}/model", + axum::routing::patch(update_agent_model), + ) .route("/api/teams/{id}/messages", post(send_message)) .route("/api/teams/{id}/agents/{slot_id}/messages", post(send_message_to_agent)) .route( @@ -232,6 +236,20 @@ async fn rename_agent( Ok(Json(ApiResponse::success())) } +async fn update_agent_model( + State(state): State, + Extension(user): Extension, + Path(params): Path, + body: Result, JsonRejection>, +) -> Result>, ApiError> { + let Json(req) = body.map_err(ApiError::from)?; + state + .service + .update_agent_model(&user.id, ¶ms.id, ¶ms.slot_id, &req.model_id) + .await?; + Ok(Json(ApiResponse::success())) +} + async fn send_message( State(state): State, Extension(user): Extension, diff --git a/crates/aionui-team/src/scheduler/agent_lifecycle.rs b/crates/aionui-team/src/scheduler/agent_lifecycle.rs index ee8d9aeb1..5bc4a7f7e 100644 --- a/crates/aionui-team/src/scheduler/agent_lifecycle.rs +++ b/crates/aionui-team/src/scheduler/agent_lifecycle.rs @@ -75,4 +75,14 @@ impl TeammateManager { debug!(team_id = %self.team_id, slot_id, new_name, "agent renamed"); Ok(()) } + + pub async fn update_agent_model(&self, slot_id: &str, model: &str) -> Result<(), TeamError> { + let mut slots = self.slots.lock().await; + let slot = slots + .get_mut(slot_id) + .ok_or_else(|| TeamError::AgentNotFound(slot_id.to_owned()))?; + slot.agent.model = model.to_owned(); + debug!(team_id = %self.team_id, slot_id, model, "agent model updated"); + Ok(()) + } } diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 858b38b83..c7b29e2dc 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -792,6 +792,63 @@ impl TeamSessionService { Ok(()) } + pub async fn update_agent_model( + &self, + user_id: &str, + team_id: &str, + slot_id: &str, + model: &str, + ) -> Result<(), TeamError> { + let model = model.trim(); + if model.is_empty() { + return Err(TeamError::InvalidRequest("model must not be empty".into())); + } + + let lock = self + .add_agent_locks + .entry(team_id.to_owned()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone(); + let _guard = lock.lock().await; + let mut team = self.load_owned_team(user_id, team_id).await?; + let conversation_id = team + .agents + .iter() + .find(|agent| agent.slot_id == slot_id) + .map(|agent| agent.conversation_id.clone()) + .ok_or_else(|| TeamError::AgentNotFound(slot_id.to_owned()))?; + + self.conversation_port + .patch_runtime_config(&conversation_id, serde_json::json!({ "current_model_id": model })) + .await?; + let agent = team + .agents + .iter_mut() + .find(|agent| agent.slot_id == slot_id) + .expect("validated team agent"); + agent.model = model.to_owned(); + self.repo + .update_team( + team_id, + &UpdateTeamParams { + agents: Some(serde_json::to_string(&team.agents)?), + ..Default::default() + }, + ) + .await?; + + if let Some(session) = self.sessions.get(team_id).map(|entry| Arc::clone(&entry.session)) + && let Err(error) = session.update_agent_model(slot_id, model).await + { + warn!(team_id, slot_id, error = %error, "persisted team model but live session roster update failed"); + } + info!( + team_id, + slot_id, conversation_id, model, "team agent model preference persisted" + ); + Ok(()) + } + /// Start the team's MCP server and rebuild every agent process so it /// carries a fresh `team_mcp_stdio_config` pointing at the new server. /// @@ -1938,7 +1995,6 @@ impl TeamSessionService { to_slot_id: &str, content: &str, ) -> Result { - self.require_active_team_run_for_team_work(team_id).await?; let session = { let entry = self .sessions @@ -1968,23 +2024,6 @@ impl TeamSessionService { session.shutdown_agent(caller_slot_id, target_slot_id, reason).await } - /// Friendly pre-check used before invoking run-scoped team tools. This is - /// not a concurrency guarantee; any operation - /// that writes mailbox, projection, scheduler, spawn, shutdown, or wake state - /// must still acquire a TeamRun operation lease in TeamSession/TeamRunManager. - pub(crate) async fn require_active_team_run_for_team_work(&self, team_id: &str) -> Result<(), TeamError> { - let entry = self - .sessions - .get(team_id) - .ok_or_else(|| TeamError::SessionNotFound(team_id.into()))?; - if entry.session.team_run_manager().current_active_run_id().is_some() { - return Ok(()); - } - Err(TeamError::InvalidRequest( - "no active team run for run-scoped wake".into(), - )) - } - pub(crate) async fn wake_leader_after_recovery_message( &self, team_id: &str, diff --git a/crates/aionui-team/src/session.rs b/crates/aionui-team/src/session.rs index 7a3e0452e..790622819 100644 --- a/crates/aionui-team/src/session.rs +++ b/crates/aionui-team/src/session.rs @@ -1294,6 +1294,10 @@ impl TeamSession { self.scheduler.rename_agent(slot_id, new_name).await } + pub async fn update_agent_model(&self, slot_id: &str, model: &str) -> Result<(), TeamError> { + self.scheduler.update_agent_model(slot_id, model).await + } + /// Spawn a new teammate at the Lead's request (backing of `team_spawn_agent`). /// /// Validation chain mirrors the assistant-first team contract: diff --git a/crates/aionui-team/tests/e2e_smoke.rs b/crates/aionui-team/tests/e2e_smoke.rs index ddfc18a8e..6bbf56c14 100644 --- a/crates/aionui-team/tests/e2e_smoke.rs +++ b/crates/aionui-team/tests/e2e_smoke.rs @@ -277,7 +277,7 @@ async fn smoke_mcp_tool_execution_not_noop() { let env = setup_team_with_lead().await; let mut stream = mcp_connect(&env, &env.lead_slot_id).await; - // --- team_send_message requires a live active Team Run ---------------- + // --- team_send_message requires a live TeamSessionService ------------- let msg_resp = mcp_call( &mut stream, 10, diff --git a/crates/aionui-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index e8d8ba223..5ace39867 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -4696,6 +4696,85 @@ async fn an1_rename_agent() { assert_eq!(agent.name, "Senior Worker"); } +#[tokio::test] +async fn observed_model_switch_updates_team_roster_runtime_seed_and_live_session() { + let (svc, _, conv_repo) = + setup_with_factory_and_metadata_and_conversation_repo(success_factory(), seeded_agent_metadata_repo()); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + svc.ensure_session("user1", &created.id).await.unwrap(); + let worker = created + .assistants + .iter() + .find(|agent| agent.role == "teammate") + .unwrap(); + + svc.update_agent_model("user1", &created.id, &worker.slot_id, "gpt-5.6-sol") + .await + .unwrap(); + + let persisted = svc.get_team("user1", &created.id).await.unwrap(); + let persisted_worker = persisted + .assistants + .iter() + .find(|agent| agent.slot_id == worker.slot_id) + .unwrap(); + assert_eq!(persisted_worker.model, "gpt-5.6-sol"); + let extra = conv_repo.get_extra(&worker.conversation_id).unwrap(); + assert_eq!(extra["current_model_id"], "gpt-5.6-sol"); + let live_worker = svc + .get_session_scheduler(&created.id) + .unwrap() + .get_agent(&worker.slot_id) + .await + .unwrap(); + assert_eq!(live_worker.model, "gpt-5.6-sol"); +} + +#[tokio::test] +async fn update_agent_model_rejects_an_empty_model_without_changing_the_roster() { + let svc = setup(); + let created = svc + .create_team( + "user1", + CreateTeamRequest { + name: "T".into(), + agents: two_agent_input(), + workspace: None, + }, + ) + .await + .unwrap(); + let worker = created + .assistants + .iter() + .find(|agent| agent.role == "teammate") + .unwrap(); + + let error = svc + .update_agent_model("user1", &created.id, &worker.slot_id, " ") + .await + .expect_err("empty model must be rejected"); + + assert!(error.to_string().contains("model must not be empty")); + let persisted = svc.get_team("user1", &created.id).await.unwrap(); + let persisted_worker = persisted + .assistants + .iter() + .find(|agent| agent.slot_id == worker.slot_id) + .unwrap(); + assert_eq!(persisted_worker.model, worker.model); +} + #[tokio::test] async fn an3_rename_nonexistent_agent() { let svc = setup(); @@ -4911,7 +4990,7 @@ async fn leader_spawn_then_immediate_ensure_joins_the_same_attach_operation() { } #[tokio::test] -async fn lead_send_agent_message_in_session_requires_active_team_run() { +async fn lead_send_agent_message_in_session_queues_without_active_team_run() { let svc = setup(); let created = svc .create_team( @@ -4940,11 +5019,11 @@ async fn lead_send_agent_message_in_session_requires_active_team_run() { .map(|agent| agent.slot_id.clone()) .expect("seeded teammate slot"); - let err = svc + let queued = svc .send_agent_message_from_agent(&created.id, &lead_slot_id, &worker_slot_id, "Do this") .await - .expect_err("leader direct message should require active Team Run"); - assert!(err.to_string().contains("no active team run")); + .expect("leader direct message should queue outside an active Team Run"); + assert_eq!(queued.team_run_id, None); } #[tokio::test]