diff --git a/Cargo.lock b/Cargo.lock index 8ba8603e3..9c6b6bd9e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -521,6 +521,7 @@ version = "0.1.51" dependencies = [ "aionui-ai-agent", "aionui-api-types", + "aionui-auth", "aionui-common", "aionui-conversation", "aionui-db", @@ -583,6 +584,7 @@ dependencies = [ "regex", "serde", "serde_json", + "sqlx", "tempfile", "thiserror 2.0.18", "tokio", @@ -611,6 +613,7 @@ dependencies = [ "iana-time-zone", "serde", "serde_json", + "sqlx", "tempfile", "thiserror 2.0.18", "tokio", @@ -639,6 +642,7 @@ name = "aionui-extension" version = "0.1.51" dependencies = [ "aionui-api-types", + "aionui-auth", "aionui-common", "aionui-db", "aionui-realtime", @@ -671,6 +675,7 @@ name = "aionui-file" version = "0.1.51" dependencies = [ "aionui-api-types", + "aionui-auth", "aionui-common", "aionui-realtime", "async-trait", @@ -699,6 +704,7 @@ name = "aionui-mcp" version = "0.1.51" dependencies = [ "aionui-api-types", + "aionui-auth", "aionui-common", "aionui-db", "aionui-realtime", @@ -847,6 +853,7 @@ name = "aionui-shell" version = "0.1.51" dependencies = [ "aionui-api-types", + "aionui-auth", "aionui-common", "aionui-db", "aionui-runtime", diff --git a/crates/aionui-ai-agent/src/capability/first_message_injector.rs b/crates/aionui-ai-agent/src/capability/first_message_injector.rs index 3883a0a3e..180d354b8 100644 --- a/crates/aionui-ai-agent/src/capability/first_message_injector.rs +++ b/crates/aionui-ai-agent/src/capability/first_message_injector.rs @@ -11,6 +11,8 @@ use crate::capability::skill_manager::{AcpSkillManager, prepare_first_message_wi /// Configuration for the first-message injector. pub struct InjectionConfig<'a> { + /// Core user that owns the conversation and its user-scoped skills. + pub user_id: &'a str, /// Preset context (assistant-level system prompt injection). pub preset_context: Option<&'a str>, /// Resolved skill names (snapshot from `conversation.extra.skills`). @@ -42,7 +44,7 @@ pub async fn inject_first_message_prefix( }; } - let skills = manager.discover_by_names(config.skills).await; + let skills = manager.discover_by_names_for_user(config.user_id, config.skills).await; let has_context = config.preset_context.is_some_and(|s| !s.is_empty()); if skills.is_empty() && !has_context { return content.to_string(); @@ -89,6 +91,7 @@ mod tests { "Hello", &mgr, InjectionConfig { + user_id: "system_default_user", preset_context: Some("Be concise."), skills: &[], native_skill_support: true, @@ -110,6 +113,7 @@ mod tests { "Hello", &mgr, InjectionConfig { + user_id: "system_default_user", preset_context: None, skills: &[], native_skill_support: true, @@ -129,6 +133,7 @@ mod tests { "Hello", &mgr, InjectionConfig { + user_id: "system_default_user", preset_context: None, skills: &[], native_skill_support: false, @@ -148,6 +153,7 @@ mod tests { "Go.", &mgr, InjectionConfig { + user_id: "system_default_user", preset_context: Some("Rule 1."), skills: &[], native_skill_support: false, @@ -184,6 +190,7 @@ mod tests { "Hello", &mgr, InjectionConfig { + user_id: "system_default_user", preset_context: None, skills: &["cron".to_owned()], native_skill_support: false, @@ -205,6 +212,7 @@ mod tests { "Do stuff", &mgr, InjectionConfig { + user_id: "system_default_user", preset_context: Some("Custom rule"), skills: &["cron".to_owned()], native_skill_support: true, diff --git a/crates/aionui-ai-agent/src/capability/skill_manager/mod.rs b/crates/aionui-ai-agent/src/capability/skill_manager/mod.rs index 596383234..aa307b563 100644 --- a/crates/aionui-ai-agent/src/capability/skill_manager/mod.rs +++ b/crates/aionui-ai-agent/src/capability/skill_manager/mod.rs @@ -92,7 +92,17 @@ impl AcpSkillManager { enabled_skills: Option<&[String]>, exclude_builtin_skills: Option<&[String]>, ) -> Vec { - let items = match self.list_available_skills().await { + self.discover_skills_for_user("system_default_user", enabled_skills, exclude_builtin_skills) + .await + } + + pub async fn discover_skills_for_user( + &self, + user_id: &str, + enabled_skills: Option<&[String]>, + exclude_builtin_skills: Option<&[String]>, + ) -> Vec { + let items = match self.list_available_skills_for_user(user_id).await { Ok(v) => v, Err(e) => { warn!(error = %e, "Failed to list skills via extension service"); @@ -159,6 +169,10 @@ impl AcpSkillManager { /// auto-inject/opt-in). Returns the resulting index. Used by the /// snapshot-driven first-message injector. pub async fn discover_by_names(&self, names: &[String]) -> Vec { + self.discover_by_names_for_user("system_default_user", names).await + } + + pub async fn discover_by_names_for_user(&self, user_id: &str, names: &[String]) -> Vec { // Always reset state so repeated calls produce a deterministic cache. if names.is_empty() { let mut cache = self.cache.write().await; @@ -167,7 +181,7 @@ impl AcpSkillManager { *discovered = true; return Vec::new(); } - let items = match self.list_available_skills().await { + let items = match self.list_available_skills_for_user(user_id).await { Ok(v) => v, Err(e) => { warn!(error = %e, "discover_by_names: list_available_skills failed"); @@ -205,11 +219,12 @@ impl AcpSkillManager { .collect() } - async fn list_available_skills( + async fn list_available_skills_for_user( &self, + user_id: &str, ) -> Result, aionui_extension::ExtensionError> { if let Some(repo) = &self.skill_repo { - aionui_extension::list_available_skills_with_repo(&self.paths, repo.as_ref()).await + aionui_extension::list_available_skills_with_repo_for_user(&self.paths, repo.as_ref(), user_id).await } else { aionui_extension::list_available_skills(&self.paths).await } diff --git a/crates/aionui-ai-agent/src/factory/acp.rs b/crates/aionui-ai-agent/src/factory/acp.rs index 1ee93585d..180bbf8f5 100644 --- a/crates/aionui-ai-agent/src/factory/acp.rs +++ b/crates/aionui-ai-agent/src/factory/acp.rs @@ -7,9 +7,10 @@ use crate::factory::acp_assembler::{WorkspaceInfo, assemble_acp_params}; use crate::factory::acp_launch_policy::{AcpLaunchPolicyInput, apply_acp_launch_policy}; use crate::factory::context::FactoryContext; use crate::manager::acp::{AcpAgentManager, CatalogForwarder}; +use crate::registry::AgentRegistry; use crate::session_context::AcpSessionBuildContext; use agent_client_protocol::schema::{EnvVariable, HttpHeader, McpServer, McpServerHttp, McpServerSse, McpServerStdio}; -use aionui_api_types::{SessionMcpServer, SessionMcpTransport}; +use aionui_api_types::{AgentMetadata, SessionMcpServer, SessionMcpTransport}; use aionui_common::CommandSpec; use aionui_db::IMcpServerRepository; use aionui_db::models::McpServerRow; @@ -28,14 +29,7 @@ pub(super) async fn build( // Resolve the catalog row — prefer explicit agent_id, fall // back to a vendor-label match for legacy payloads. - let meta = if let Some(ref agent_id) = config.agent_id { - deps.agent_registry.get(agent_id).await - } else if let Some(ref vendor) = config.backend { - deps.agent_registry.find_builtin_by_backend(vendor).await - } else { - None - } - .ok_or_else(|| AgentError::bad_request("ACP agent requires either agent_id or backend in extra"))?; + let meta = resolve_catalog_metadata(&deps.agent_registry, &config, &ctx.user_id).await?; // Trust the catalog row over the client-supplied `backend` when an // `agent_id` was provided. The frontend collapses row-scoped rows @@ -60,6 +54,7 @@ pub(super) async fn build( backend_label, crate::session_agent::SessionBuildInputs { conversation_id: ctx.conversation_id.clone(), + user_id: ctx.user_id.clone(), workspace: ctx.workspace.clone(), config: &config, metadata: &meta, @@ -101,8 +96,14 @@ pub(super) async fn build( return Ok(instance); } - let mut command_spec = - resolve_agent_command_spec(&meta, &ctx.workspace, &ctx.conversation_id, deps.broadcaster.clone()).await?; + let mut command_spec = resolve_agent_command_spec( + &meta, + &ctx.user_id, + &ctx.workspace, + &ctx.conversation_id, + deps.broadcaster.clone(), + ) + .await?; apply_acp_launch_policy( &mut command_spec, AcpLaunchPolicyInput { @@ -130,6 +131,7 @@ pub(super) async fn build( load_user_mcp_servers( repo.as_ref(), config.mcp_server_ids.as_deref(), + &ctx.user_id, &ctx.conversation_id, &mcp_capabilities, ) @@ -165,6 +167,7 @@ pub(super) async fn build( let params = Arc::new( assemble_acp_params( ctx.conversation_id.clone(), + ctx.user_id.clone(), WorkspaceInfo { path: ctx.workspace, is_custom: ctx.is_custom_workspace, @@ -189,6 +192,7 @@ pub(super) async fn build( arc.start_permission_handler(); arc.start_session_event_tracker(notification_rx); CatalogForwarder::spawn( + ctx.user_id.clone(), arc.agent_id().to_owned(), crate::IAgentTask::subscribe(arc.as_ref()), catalog_tx, @@ -213,13 +217,40 @@ pub(super) async fn build( // Hand the service the domain event receiver so it can // persist user intent changes without reverse-engineering // them from CLI observations. - deps.acp_agent_service.attach(ctx.conversation_id, domain_rx).await; + deps.acp_agent_service + .attach(ctx.user_id, ctx.conversation_id, domain_rx) + .await; Ok(instance) } +async fn resolve_catalog_metadata( + registry: &Arc, + config: &aionui_api_types::AcpBuildExtra, + user_id: &str, +) -> Result { + if let Some(ref agent_id) = config.agent_id { + return registry + .get_for_user(user_id, agent_id) + .await? + .ok_or_else(|| AgentError::bad_request("ACP agent_id is not available for this user")); + } + + if let Some(ref vendor) = config.backend { + return registry + .find_builtin_by_backend_for_user(user_id, vendor) + .await + .ok_or_else(|| AgentError::bad_request("ACP backend is not available for this user")); + } + + Err(AgentError::bad_request( + "ACP agent requires either agent_id or backend in extra", + )) +} + async fn resolve_agent_command_spec( meta: &aionui_api_types::AgentMetadata, + user_id: &str, workspace: &str, conversation_id: &str, broadcaster: Arc, @@ -229,7 +260,7 @@ async fn resolve_agent_command_spec( .as_deref() .filter(|value| !value.is_empty()) .ok_or_else(|| AgentError::bad_request(format!("Agent '{}' has no spawn command configured", meta.name)))?; - let reporter = conversation_runtime_reporter(broadcaster, conversation_id.to_owned()); + let reporter = conversation_runtime_reporter(broadcaster, user_id.to_owned(), conversation_id.to_owned()); let resolved = ensure_runtime_command_with_reporter(command, Some(reporter.as_ref())) .await .map_err(|error| AgentError::bad_request(format!("Agent '{}' CLI unavailable: {error}", meta.name)))?; @@ -283,12 +314,13 @@ async fn resolve_agent_command_spec( async fn load_user_mcp_servers( repo: &dyn IMcpServerRepository, selected_ids: Option<&[String]>, + user_id: &str, conversation_id: &str, capabilities: &AcpMcpCapabilities, ) -> Vec { let rows_result = match selected_ids { - Some(ids) => repo.list_by_ids_any(ids).await, - None => repo.list().await, + Some(ids) => repo.list_by_ids_any(user_id, ids).await, + None => repo.list(user_id).await, }; let rows = match rows_result { Ok(r) => r, @@ -503,6 +535,10 @@ fn session_server_supported_by_capabilities(server: &SessionMcpServer, capabilit #[cfg(test)] mod tests { use super::*; + use aionui_api_types::AcpBuildExtra; + use aionui_db::{ + IAgentMetadataRepository, SqliteAgentMetadataRepository, UpsertAgentMetadataParams, init_database_memory, + }; use aionui_realtime::BroadcastEventBus; use aionui_runtime::{ManagedResourcesMode, init as init_runtime, set_managed_resources_mode}; use std::sync::OnceLock; @@ -511,6 +547,8 @@ mod tests { path::{Path, PathBuf}, }; + const TEST_USER_ID: &str = "user-1"; + fn make_row( name: &str, transport_type: &str, @@ -520,6 +558,7 @@ mod tests { ) -> McpServerRow { McpServerRow { id: format!("mcp_{name}"), + user_id: TEST_USER_ID.to_owned(), name: name.to_owned(), description: None, enabled, @@ -558,6 +597,84 @@ mod tests { command == "npx" || command.ends_with("/npx") || command.ends_with("\\npx.cmd") } + async fn seed_user(pool: &sqlx::SqlitePool, user_id: &str) { + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, 0, 0)", + ) + .bind(user_id) + .bind(user_id) + .execute(pool) + .await + .unwrap(); + } + + fn custom_agent_params<'a>(id: &'a str, name: &'a str, command: &'a str) -> UpsertAgentMetadataParams<'a> { + UpsertAgentMetadataParams { + id, + icon: None, + name, + name_i18n: None, + description: None, + description_i18n: None, + backend: Some("custom"), + agent_type: "acp", + agent_source: "custom", + agent_source_info: Some(r#"{"binary_name":"custom"}"#), + enabled: true, + command: Some(command), + args: Some("[]"), + env: Some("[]"), + native_skills_dirs: None, + behavior_policy: None, + yolo_id: None, + agent_capabilities: None, + auth_methods: None, + config_options: None, + available_modes: None, + available_models: None, + available_commands: None, + sort_order: 100, + } + } + + #[tokio::test] + async fn resolve_catalog_metadata_rejects_other_users_custom_agent_id() { + let db = init_database_memory().await.unwrap(); + seed_user(db.pool(), "user-a").await; + seed_user(db.pool(), "user-b").await; + + let repo: Arc = Arc::new(SqliteAgentMetadataRepository::new(db.pool().clone())); + let command = std::env::current_exe() + .expect("current test executable") + .to_string_lossy() + .into_owned(); + repo.upsert_for_user( + "user-b", + &custom_agent_params("custom-agent-b", "User B Agent", &command), + ) + .await + .unwrap(); + + let registry = AgentRegistry::new(repo); + registry.hydrate().await.unwrap(); + let config = AcpBuildExtra { + agent_id: Some("custom-agent-b".to_owned()), + ..Default::default() + }; + + let err = resolve_catalog_metadata(®istry, &config, "user-a") + .await + .unwrap_err(); + assert!( + err.to_string().contains("not available for this user"), + "unexpected error: {err}" + ); + + let meta = resolve_catalog_metadata(®istry, &config, "user-b").await.unwrap(); + assert_eq!(meta.id, "custom-agent-b"); + } + #[cfg(unix)] fn test_runtime_data_dir() -> &'static PathBuf { static DIR: OnceLock = OnceLock::new(); @@ -722,6 +839,7 @@ mod tests { let spec = resolve_agent_command_spec( &meta, + "user-acp", "/tmp/workspace", "conv-acp", Arc::new(BroadcastEventBus::new(16)), @@ -743,6 +861,7 @@ mod tests { meta.args = vec!["-y".into(), "pi-acp".into()]; let spec = resolve_agent_command_spec( &meta, + "user-acp", "/tmp/workspace", "conv-acp", Arc::new(BroadcastEventBus::new(16)), @@ -838,26 +957,35 @@ mod tests { #[async_trait] impl IMcpServerRepository for MockRepo { - async fn list(&self) -> Result, aionui_db::DbError> { + async fn list(&self, user_id: &str) -> Result, aionui_db::DbError> { if self.fail { Err(aionui_db::DbError::Init("simulated".into())) } else { - Ok(self.rows.clone()) + Ok(self.rows.iter().filter(|row| row.user_id == user_id).cloned().collect()) } } - async fn find_by_id(&self, _id: &str) -> Result, aionui_db::DbError> { + async fn find_by_id(&self, _user_id: &str, _id: &str) -> Result, aionui_db::DbError> { unimplemented!() } - async fn find_by_name(&self, _name: &str) -> Result, aionui_db::DbError> { + async fn find_by_name(&self, _user_id: &str, _name: &str) -> Result, aionui_db::DbError> { unimplemented!() } - async fn list_by_ids_any(&self, ids: &[String]) -> Result, aionui_db::DbError> { + async fn list_by_ids_any( + &self, + user_id: &str, + ids: &[String], + ) -> Result, aionui_db::DbError> { if self.fail { return Err(aionui_db::DbError::Init("simulated".into())); } Ok(ids .iter() - .filter_map(|id| self.rows.iter().find(|row| row.id == *id).cloned()) + .filter_map(|id| { + self.rows + .iter() + .find(|row| row.user_id == user_id && row.id == *id) + .cloned() + }) .collect()) } async fn create( @@ -868,29 +996,37 @@ mod tests { } async fn update( &self, + _user_id: &str, _id: &str, _params: aionui_db::UpdateMcpServerParams<'_>, ) -> Result { unimplemented!() } - async fn delete(&self, _id: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), aionui_db::DbError> { unimplemented!() } async fn batch_upsert( &self, + _user_id: &str, _servers: &[aionui_db::CreateMcpServerParams<'_>], ) -> Result, aionui_db::DbError> { unimplemented!() } async fn update_status( &self, + _user_id: &str, _id: &str, _status: &str, _last_connected: Option, ) -> Result<(), aionui_db::DbError> { unimplemented!() } - async fn update_tools(&self, _id: &str, _tools: Option<&str>) -> Result<(), aionui_db::DbError> { + async fn update_tools( + &self, + _user_id: &str, + _id: &str, + _tools: Option<&str>, + ) -> Result<(), aionui_db::DbError> { unimplemented!() } } @@ -917,7 +1053,7 @@ mod tests { ], fail: false, }); - let servers = load_user_mcp_servers(repo.as_ref(), None, "conv-1", &caps).await; + let servers = load_user_mcp_servers(repo.as_ref(), None, TEST_USER_ID, "conv-1", &caps).await; assert_eq!(servers.len(), 1); match &servers[0] { McpServer::Stdio(s) => assert_eq!(s.name, "user-enabled"), @@ -936,7 +1072,7 @@ mod tests { rows: vec![], fail: true, }); - let servers = load_user_mcp_servers(repo.as_ref(), None, "conv-1", &caps).await; + let servers = load_user_mcp_servers(repo.as_ref(), None, TEST_USER_ID, "conv-1", &caps).await; assert!(servers.is_empty()); } @@ -955,7 +1091,7 @@ mod tests { ], fail: false, }); - let servers = load_user_mcp_servers(repo.as_ref(), None, "conv-1", &caps).await; + let servers = load_user_mcp_servers(repo.as_ref(), None, TEST_USER_ID, "conv-1", &caps).await; assert_eq!(servers.len(), 1); match &servers[0] { McpServer::Stdio(s) => assert_eq!(s.name, "good"), @@ -980,7 +1116,7 @@ mod tests { }); let selected = vec!["mcp_disabled-picked".to_owned()]; - let servers = load_user_mcp_servers(repo.as_ref(), Some(&selected), "conv-1", &caps).await; + let servers = load_user_mcp_servers(repo.as_ref(), Some(&selected), TEST_USER_ID, "conv-1", &caps).await; assert_eq!(servers.len(), 1); match &servers[0] { @@ -1007,7 +1143,7 @@ mod tests { fail: false, }); - let servers = load_user_mcp_servers(repo.as_ref(), None, "conv-1", &caps).await; + let servers = load_user_mcp_servers(repo.as_ref(), None, TEST_USER_ID, "conv-1", &caps).await; assert!(servers.is_empty()); } } diff --git a/crates/aionui-ai-agent/src/factory/acp_assembler.rs b/crates/aionui-ai-agent/src/factory/acp_assembler.rs index 0bb8f8c1d..ab23eab6e 100644 --- a/crates/aionui-ai-agent/src/factory/acp_assembler.rs +++ b/crates/aionui-ai-agent/src/factory/acp_assembler.rs @@ -21,6 +21,7 @@ pub struct WorkspaceInfo { #[derive(Debug, Clone)] pub struct AcpSessionParams { pub conversation_id: String, + pub user_id: String, pub workspace: WorkspaceInfo, pub metadata: AgentMetadata, pub command_spec: CommandSpec, @@ -59,6 +60,7 @@ impl AcpSessionParams { #[allow(clippy::too_many_arguments)] pub async fn assemble_acp_params( conversation_id: String, + user_id: String, workspace: WorkspaceInfo, metadata: AgentMetadata, command_spec: CommandSpec, @@ -73,6 +75,7 @@ pub async fn assemble_acp_params( AcpSessionParams { conversation_id, + user_id, workspace, metadata, command_spec, @@ -200,6 +203,7 @@ mod tests { let params = assemble_acp_params( "conv-1".into(), + "user-1".into(), WorkspaceInfo { path: "/tmp/workspace".into(), is_custom: false, diff --git a/crates/aionui-ai-agent/src/factory/aionrs.rs b/crates/aionui-ai-agent/src/factory/aionrs.rs index 07d568f46..da9fba085 100644 --- a/crates/aionui-ai-agent/src/factory/aionrs.rs +++ b/crates/aionui-ai-agent/src/factory/aionrs.rs @@ -53,6 +53,7 @@ pub(super) async fn build( for (name, config) in load_user_mcp_servers( repo.as_ref(), overrides.mcp_server_ids.as_deref(), + &ctx.user_id, &ctx.conversation_id, deps.broadcaster.clone(), ) @@ -64,6 +65,7 @@ pub(super) async fn build( merge_session_snapshot_mcp_servers( &mut extra_mcp_servers, &overrides.session_mcp_servers, + &ctx.user_id, &ctx.conversation_id, deps.broadcaster.clone(), ) @@ -81,7 +83,7 @@ pub(super) async fn build( let provider_id = &model.provider_id; let row = deps .provider_repo - .find_by_id(provider_id) + .find_by_id(&ctx.user_id, provider_id) .await .map_err(|e| AgentError::internal(format!("Failed to load provider config: {e}")))? .ok_or_else(|| AgentError::bad_request(format!("Provider '{provider_id}' not found")))?; @@ -442,12 +444,13 @@ pub(crate) fn resolve_bedrock_config(json: Option<&str>) -> Option, + user_id: &str, conversation_id: &str, broadcaster: Arc, ) -> HashMap { let rows_result = match selected_ids { - Some(ids) => repo.list_by_ids_any(ids).await, - None => repo.list().await, + Some(ids) => repo.list_by_ids_any(user_id, ids).await, + None => repo.list(user_id).await, }; let rows = match rows_result { Ok(r) => r, @@ -470,7 +473,7 @@ async fn load_user_mcp_servers( continue; } - match row_to_mcp_server_config(&row, conversation_id, broadcaster.clone()).await { + match row_to_mcp_server_config(&row, user_id, conversation_id, broadcaster.clone()).await { Ok(config) => { servers.insert(row.name.clone(), config); } @@ -491,6 +494,7 @@ async fn load_user_mcp_servers( async fn row_to_mcp_server_config( row: &McpServerRow, + user_id: &str, conversation_id: &str, broadcaster: Arc, ) -> Result { @@ -518,7 +522,7 @@ async fn row_to_mcp_server_config( }) .unwrap_or_default(); let (resolved_command, args, env) = - ensure_stdio_launch(command, &args, &env_entries, conversation_id, broadcaster).await?; + ensure_stdio_launch(command, &args, &env_entries, user_id, conversation_id, broadcaster).await?; Ok(McpServerConfig { transport: TransportType::Stdio, @@ -589,6 +593,7 @@ async fn row_to_mcp_server_config( async fn session_server_to_mcp_server_config( server: &SessionMcpServer, + user_id: &str, conversation_id: &str, broadcaster: Arc, ) -> Result { @@ -599,7 +604,7 @@ async fn session_server_to_mcp_server_config( } let entries: Vec<(String, String)> = env.iter().map(|(k, v)| (k.clone(), v.clone())).collect(); let (command, args, env) = - ensure_stdio_launch(command, args, &entries, conversation_id, broadcaster).await?; + ensure_stdio_launch(command, args, &entries, user_id, conversation_id, broadcaster).await?; Ok(McpServerConfig { transport: TransportType::Stdio, command: Some(command), @@ -662,11 +667,12 @@ async fn session_server_to_mcp_server_config( async fn merge_session_snapshot_mcp_servers( extra_mcp_servers: &mut HashMap, session_mcp_servers: &[SessionMcpServer], + user_id: &str, conversation_id: &str, broadcaster: Arc, ) { for server in session_mcp_servers { - match session_server_to_mcp_server_config(server, conversation_id, broadcaster.clone()).await { + match session_server_to_mcp_server_config(server, user_id, conversation_id, broadcaster.clone()).await { Ok(config) => { if extra_mcp_servers.insert(server.name.clone(), config).is_some() { debug!( @@ -693,10 +699,11 @@ async fn ensure_stdio_launch( command: &str, args: &[String], env: &[(String, String)], + user_id: &str, conversation_id: &str, broadcaster: Arc, ) -> Result<(String, Vec, HashMap), String> { - let reporter = conversation_runtime_reporter(broadcaster, conversation_id.to_owned()); + let reporter = conversation_runtime_reporter(broadcaster, user_id.to_owned(), conversation_id.to_owned()); let resolved = ensure_runtime_command_with_reporter(command, Some(reporter.as_ref())) .await .map_err(|error| error.to_string())?; @@ -761,6 +768,8 @@ mod tests { path::{Path, PathBuf}, }; + const TEST_USER_ID: &str = "user-1"; + fn path_test_lock() -> &'static tokio::sync::Mutex<()> { static LOCK: OnceLock> = OnceLock::new(); LOCK.get_or_init(|| tokio::sync::Mutex::new(())) @@ -860,6 +869,7 @@ mod tests { ) -> McpServerRow { McpServerRow { id: format!("mcp_{name}"), + user_id: TEST_USER_ID.to_owned(), name: name.to_owned(), description: None, enabled, @@ -882,16 +892,24 @@ mod tests { #[async_trait::async_trait] impl IMcpServerRepository for MockMcpRepo { - async fn list(&self) -> Result, aionui_db::DbError> { - Ok(self.rows.clone()) + async fn list(&self, user_id: &str) -> Result, aionui_db::DbError> { + Ok(self.rows.iter().filter(|row| row.user_id == user_id).cloned().collect()) } - async fn find_by_id(&self, id: &str) -> Result, aionui_db::DbError> { - Ok(self.rows.iter().find(|row| row.id == id).cloned()) + async fn find_by_id(&self, user_id: &str, id: &str) -> Result, aionui_db::DbError> { + Ok(self + .rows + .iter() + .find(|row| row.user_id == user_id && row.id == id) + .cloned()) } - async fn find_by_name(&self, name: &str) -> Result, aionui_db::DbError> { - Ok(self.rows.iter().find(|row| row.name == name).cloned()) + async fn find_by_name(&self, user_id: &str, name: &str) -> Result, aionui_db::DbError> { + Ok(self + .rows + .iter() + .find(|row| row.user_id == user_id && row.name == name) + .cloned()) } async fn create( @@ -903,18 +921,20 @@ mod tests { async fn update( &self, + _user_id: &str, _id: &str, _params: aionui_db::UpdateMcpServerParams<'_>, ) -> Result { unimplemented!("not needed for factory tests") } - async fn delete(&self, _id: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), aionui_db::DbError> { unimplemented!("not needed for factory tests") } async fn batch_upsert( &self, + _user_id: &str, _servers: &[aionui_db::CreateMcpServerParams<'_>], ) -> Result, aionui_db::DbError> { unimplemented!("not needed for factory tests") @@ -922,6 +942,7 @@ mod tests { async fn update_status( &self, + _user_id: &str, _id: &str, _status: &str, _last_connected: Option, @@ -929,7 +950,12 @@ mod tests { unimplemented!("not needed for factory tests") } - async fn update_tools(&self, _id: &str, _tools: Option<&str>) -> Result<(), aionui_db::DbError> { + async fn update_tools( + &self, + _user_id: &str, + _id: &str, + _tools: Option<&str>, + ) -> Result<(), aionui_db::DbError> { unimplemented!("not needed for factory tests") } } @@ -951,8 +977,14 @@ mod tests { let repo = MockMcpRepo { rows: vec![row] }; let selected = vec!["mcp-docs".to_owned()]; - let extra_mcp_servers = - load_user_mcp_servers(&repo, Some(&selected), "conv-frozen-mcp", test_broadcaster()).await; + let extra_mcp_servers = load_user_mcp_servers( + &repo, + Some(&selected), + TEST_USER_ID, + "conv-frozen-mcp", + test_broadcaster(), + ) + .await; assert!(extra_mcp_servers.contains_key("mcp-docs")); assert_eq!(extra_mcp_servers["mcp-docs"].transport, TransportType::StreamableHttp); @@ -974,7 +1006,7 @@ mod tests { false, ); - let config = row_to_mcp_server_config(&row, "conv-row", test_broadcaster()) + let config = row_to_mcp_server_config(&row, "user-row", "conv-row", test_broadcaster()) .await .expect("convert"); let command = config.command.as_deref().expect("resolved command"); @@ -1807,7 +1839,14 @@ mod tests { }, }]; - merge_session_snapshot_mcp_servers(&mut servers, &snapshot, "conv-override", test_broadcaster()).await; + merge_session_snapshot_mcp_servers( + &mut servers, + &snapshot, + "user-override", + "conv-override", + test_broadcaster(), + ) + .await; let server = servers.get("demo-mcp").expect("snapshot should remain"); assert_eq!(server.transport, TransportType::Stdio); diff --git a/crates/aionui-ai-agent/src/factory/context.rs b/crates/aionui-ai-agent/src/factory/context.rs index dc119762e..8069a9a33 100644 --- a/crates/aionui-ai-agent/src/factory/context.rs +++ b/crates/aionui-ai-agent/src/factory/context.rs @@ -7,6 +7,7 @@ use crate::session_context::AgentSessionContext; pub(super) struct FactoryContext { pub conversation_id: String, + pub user_id: String, pub workspace: String, pub is_custom_workspace: bool, pub runtime_env: Vec<(String, String)>, @@ -16,6 +17,7 @@ impl FactoryContext { pub async fn resolve(context: &AgentSessionContext) -> Result { Ok(Self { conversation_id: context.conversation.conversation_id.clone(), + user_id: context.conversation.user_id.clone(), workspace: context.workspace.path.clone(), is_custom_workspace: context.workspace.is_custom, runtime_env: context.runtime_env.clone(), diff --git a/crates/aionui-ai-agent/src/manager/acp/agent.rs b/crates/aionui-ai-agent/src/manager/acp/agent.rs index 20415606e..940531294 100644 --- a/crates/aionui-ai-agent/src/manager/acp/agent.rs +++ b/crates/aionui-ai-agent/src/manager/acp/agent.rs @@ -611,7 +611,11 @@ impl AcpAgentManager { ..Default::default() }; if init_handshake.agent_capabilities.is_some() || init_handshake.auth_methods.is_some() { - catalog_tx.send_partial(self.params.metadata.id.clone(), init_handshake); + catalog_tx.send_partial( + self.params.user_id.clone(), + self.params.metadata.id.clone(), + init_handshake, + ); } // Seed the observed/advertised layers (observed mode/model, cached diff --git a/crates/aionui-ai-agent/src/manager/acp/catalog_forwarder.rs b/crates/aionui-ai-agent/src/manager/acp/catalog_forwarder.rs index 3e6318a2e..b90bcbeb5 100644 --- a/crates/aionui-ai-agent/src/manager/acp/catalog_forwarder.rs +++ b/crates/aionui-ai-agent/src/manager/acp/catalog_forwarder.rs @@ -28,6 +28,7 @@ impl CatalogForwarder { /// awaited — callers drop it and rely on the broadcast channel /// closing to terminate the task. pub fn spawn( + user_id: String, agent_id: String, mut event_rx: broadcast::Receiver, catalog_tx: CatalogSender, @@ -37,14 +38,14 @@ impl CatalogForwarder { match event_rx.recv().await { Ok(event) => { if let Some(partial) = catalog_partial_from_event(&event) { - catalog_tx.send_partial(agent_id.clone(), partial); + catalog_tx.send_partial(user_id.clone(), agent_id.clone(), partial); } } Err(broadcast::error::RecvError::Closed) => break, Err(broadcast::error::RecvError::Lagged(_)) => continue, } } - debug!(agent_id, "CatalogForwarder exiting"); + debug!(user_id, agent_id, "CatalogForwarder exiting"); }) } } diff --git a/crates/aionui-ai-agent/src/manager/acp/hooks.rs b/crates/aionui-ai-agent/src/manager/acp/hooks.rs index 9b48209f0..19aa822e2 100644 --- a/crates/aionui-ai-agent/src/manager/acp/hooks.rs +++ b/crates/aionui-ai-agent/src/manager/acp/hooks.rs @@ -22,6 +22,7 @@ impl PreSendHook for SessionNewPreludeHook { let metadata = &ctx.params.metadata; let config = InjectionConfig { + user_id: &ctx.params.user_id, preset_context: ctx.params.preset_context.as_deref(), skills: &ctx.params.config.skills, native_skill_support: metadata diff --git a/crates/aionui-ai-agent/src/mcp_resolve.rs b/crates/aionui-ai-agent/src/mcp_resolve.rs index 82fbca3f4..5991072d9 100644 --- a/crates/aionui-ai-agent/src/mcp_resolve.rs +++ b/crates/aionui-ai-agent/src/mcp_resolve.rs @@ -35,13 +35,14 @@ use tracing::{info, warn}; /// reporting) and reserved for that use. pub async fn resolve_session_mcp_servers( repo: &dyn IMcpServerRepository, + user_id: &str, selected_ids: Option<&[String]>, conversation_id: &str, _broadcaster: Arc, ) -> Vec { let rows_result = match selected_ids { - Some(ids) => repo.list_by_ids_any(ids).await, - None => repo.list().await, + Some(ids) => repo.list_by_ids_any(user_id, ids).await, + None => repo.list(user_id).await, }; let rows = match rows_result { Ok(r) => r, diff --git a/crates/aionui-ai-agent/src/persistence/acp_session_sync.rs b/crates/aionui-ai-agent/src/persistence/acp_session_sync.rs index 0d96c4bfe..7d7a3fa1b 100644 --- a/crates/aionui-ai-agent/src/persistence/acp_session_sync.rs +++ b/crates/aionui-ai-agent/src/persistence/acp_session_sync.rs @@ -51,11 +51,16 @@ impl AcpSessionSyncService { } /// Read the persisted per-session state for `conversation_id`. - pub async fn load_persisted(&self, conversation_id: &str) -> Option { - match self.repo.load_runtime_state(conversation_id).await { + pub async fn load_persisted( + &self, + user_id: &str, + conversation_id: &str, + ) -> Option { + match self.repo.load_runtime_state_for_user(user_id, conversation_id).await { Ok(state) => state, Err(err) => { warn!( + user_id, conversation_id, error = %err, "AcpSessionSyncService::load_persisted failed" @@ -69,12 +74,13 @@ impl AcpSessionSyncService { /// aggregate's value-object shape. Returns `None` when the row does /// not exist or the JSON payload is empty. Errors are logged and /// swallowed so the caller can proceed with a fresh session. - pub async fn load_snapshot_state(&self, conversation_id: &str) -> Option { - let row = match self.repo.load_runtime_state(conversation_id).await { + pub async fn load_snapshot_state(&self, user_id: &str, conversation_id: &str) -> Option { + let row = match self.repo.load_runtime_state_for_user(user_id, conversation_id).await { Ok(Some(row)) => row, Ok(None) => return None, Err(err) => { warn!( + user_id, conversation_id, error = %err, "load_snapshot_state: repository failed; skipping preload" @@ -107,12 +113,13 @@ impl AcpSessionSyncService { /// Read the persisted CLI-assigned session id, if any. /// Used by the factory on resume paths to seed the aggregate before /// the first prompt. - pub async fn load_session_id(&self, conversation_id: &str) -> Option { - match self.repo.get(conversation_id).await { + pub async fn load_session_id(&self, user_id: &str, conversation_id: &str) -> Option { + match self.repo.get_for_user(user_id, conversation_id).await { Ok(Some(row)) => row.session_id, Ok(None) => None, Err(err) => { warn!( + user_id, conversation_id, error = %err, "load_session_id: repository failed" @@ -126,10 +133,10 @@ impl AcpSessionSyncService { /// the per-conversation persistence consumer. Lifetime of the /// spawned task is tied to the sender being dropped when the manager /// is destroyed. - pub async fn attach(&self, conversation_id: String, domain_rx: mpsc::Receiver) { + pub async fn attach(&self, user_id: String, conversation_id: String, domain_rx: mpsc::Receiver) { let repo = self.repo.clone(); let cid = conversation_id.clone(); - let task = tokio::spawn(domain_event_consumer(cid, domain_rx, repo)); + let task = tokio::spawn(domain_event_consumer(user_id, cid, domain_rx, repo)); let mut guard = self.active.write().await; if let Some(prev) = guard.insert(conversation_id, task) { @@ -199,6 +206,7 @@ impl PendingUpdate { /// written immediately so the next turn can take the resume path even /// if the process crashes before any other event fires. async fn domain_event_consumer( + user_id: String, conversation_id: String, mut rx: mpsc::Receiver, repo: Arc, @@ -213,7 +221,7 @@ async fn domain_event_consumer( biased; maybe_event = rx.recv() => maybe_event, () = sleep_until(deadline.into()) => { - flush(&repo, &conversation_id, &mut pending).await; + flush(&repo, &user_id, &conversation_id, &mut pending).await; flush_at = None; continue; } @@ -225,13 +233,17 @@ async fn domain_event_consumer( match recv { Some(event) => { if let AcpSessionEvent::SessionAssigned { session_id } = &event { - match repo.update_session_id(&conversation_id, session_id.as_str()).await { + match repo + .update_session_id_for_user(&user_id, &conversation_id, session_id.as_str()) + .await + { Ok(true) => {} Ok(false) => debug!( - conversation_id, - "session-sync: acp_session row missing; session_id not written" + user_id, + conversation_id, "session-sync: acp_session row missing; session_id not written" ), Err(err) => warn!( + user_id, conversation_id, error = %err, "session-sync: update_session_id failed" @@ -244,26 +256,38 @@ async fn domain_event_consumer( } } None => { - flush(&repo, &conversation_id, &mut pending).await; - debug!(conversation_id, "session-sync domain consumer exiting"); + flush(&repo, &user_id, &conversation_id, &mut pending).await; + debug!(user_id, conversation_id, "session-sync domain consumer exiting"); return; } } } } -async fn flush(repo: &Arc, conversation_id: &str, pending: &mut PendingUpdate) { +async fn flush( + repo: &Arc, + user_id: &str, + conversation_id: &str, + pending: &mut PendingUpdate, +) { if pending.is_empty() { return; } let params = pending.as_save_params(); - match repo.save_runtime_state(conversation_id, ¶ms).await { + match repo + .save_runtime_state_for_user(user_id, conversation_id, ¶ms) + .await + { Ok(true) => {} Ok(false) => { - debug!(conversation_id, "session sync: acp_session row missing; update dropped"); + debug!( + user_id, + conversation_id, "session sync: acp_session row missing; update dropped" + ); } Err(err) => { warn!( + user_id, conversation_id, error = %err, "session sync: save_runtime_state failed" @@ -282,8 +306,26 @@ mod tests { async fn setup() -> (Arc, Arc) { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT OR IGNORE INTO users \ + (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES ('user-1', 'local', 'user-1', 'hash', 'active', 0, 1, 1)", + ) + .execute(db.pool()) + .await + .unwrap(); + sqlx::query( + "INSERT INTO conversations \ + (id, user_id, name, type, extra, model, status, source, \ + channel_chat_id, pinned, pinned_at, created_at, updated_at) \ + VALUES ('conv-1', 'user-1', 'Test', 'acp', '{}', NULL, 'pending', NULL, NULL, 0, NULL, 1, 1)", + ) + .execute(db.pool()) + .await + .unwrap(); let repo: Arc = Arc::new(SqliteAcpSessionRepository::new(db.pool().clone())); repo.create(&CreateAcpSessionParams { + user_id: "user-1", conversation_id: "conv-1", agent_source: "builtin", agent_id: "2d23ff1c", @@ -297,7 +339,8 @@ mod tests { #[tokio::test] async fn load_persisted_round_trips() { let (svc, repo) = setup().await; - repo.save_runtime_state( + repo.save_runtime_state_for_user( + "user-1", "conv-1", &SaveRuntimeStateParams { current_mode_id: Some(Some("plan")), @@ -307,7 +350,7 @@ mod tests { .await .unwrap(); - let state = svc.load_persisted("conv-1").await.unwrap(); + let state = svc.load_persisted("user-1", "conv-1").await.unwrap(); assert_eq!(state.current_mode_id.as_deref(), Some("plan")); } @@ -318,18 +361,26 @@ mod tests { let (tx, rx) = mpsc::channel(64); let cid = "conv-1".to_owned(); - tokio::spawn(domain_event_consumer(cid, rx, repo.clone())); + tokio::spawn(domain_event_consumer("user-1".to_owned(), cid, rx, repo.clone())); tx.send(AcpSessionEvent::ObservedModeSynced { mode: "plan".into() }) .await .unwrap(); sleep(Duration::from_millis(200)).await; - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .unwrap(); assert!(state.current_mode_id.is_none(), "debounce not yet elapsed"); sleep(Duration::from_millis(400)).await; - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .unwrap(); assert_eq!(state.current_mode_id.as_deref(), Some("plan")); } @@ -340,7 +391,7 @@ mod tests { let (tx, rx) = mpsc::channel(64); let cid = "conv-1".to_owned(); - tokio::spawn(domain_event_consumer(cid, rx, repo.clone())); + tokio::spawn(domain_event_consumer("user-1".to_owned(), cid, rx, repo.clone())); for label in ["code", "plan", "ask"] { tx.send(AcpSessionEvent::ObservedModeSynced { mode: label.into() }) @@ -350,7 +401,11 @@ mod tests { } sleep(Duration::from_millis(600)).await; - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .unwrap(); assert_eq!(state.current_mode_id.as_deref(), Some("ask")); } @@ -361,12 +416,16 @@ mod tests { let (tx, rx) = mpsc::channel(64); let cid = "conv-1".to_owned(); - tokio::spawn(domain_event_consumer(cid, rx, repo.clone())); + tokio::spawn(domain_event_consumer("user-1".to_owned(), cid, rx, repo.clone())); tx.send(AcpSessionEvent::SessionOpened).await.unwrap(); sleep(Duration::from_millis(600)).await; - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .unwrap(); assert!(state.current_mode_id.is_none()); } @@ -377,7 +436,7 @@ mod tests { let (tx, rx) = mpsc::channel(64); let cid = "conv-1".to_owned(); - tokio::spawn(domain_event_consumer(cid, rx, repo.clone())); + tokio::spawn(domain_event_consumer("user-1".to_owned(), cid, rx, repo.clone())); tx.send(AcpSessionEvent::ObservedModeSynced { mode: "plan".into() }) .await @@ -385,7 +444,11 @@ mod tests { drop(tx); sleep(Duration::from_millis(50)).await; - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .unwrap(); assert_eq!( state.current_mode_id.as_deref(), Some("plan"), @@ -404,7 +467,7 @@ mod tests { let (tx, rx) = mpsc::channel(64); let cid = "conv-1".to_owned(); - tokio::spawn(domain_event_consumer(cid, rx, repo.clone())); + tokio::spawn(domain_event_consumer("user-1".to_owned(), cid, rx, repo.clone())); tx.send(AcpSessionEvent::ObservedModelSynced { model: "claude-opus-4".into(), @@ -413,7 +476,11 @@ mod tests { .unwrap(); sleep(Duration::from_millis(700)).await; - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .unwrap(); assert_eq!(state.current_model_id.as_deref(), Some("claude-opus-4")); } @@ -423,7 +490,7 @@ mod tests { let (tx, rx) = mpsc::channel(64); let cid = "conv-1".to_owned(); - tokio::spawn(domain_event_consumer(cid, rx, repo.clone())); + tokio::spawn(domain_event_consumer("user-1".to_owned(), cid, rx, repo.clone())); tx.send(AcpSessionEvent::DesiredModelChanged { model: "claude-opus-4".into(), @@ -432,7 +499,11 @@ mod tests { .unwrap(); sleep(Duration::from_millis(700)).await; - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .unwrap(); assert!( state.current_model_id.is_none(), "DesiredModelChanged is reconcile/UI-only; persistence only follows Observed*", @@ -445,14 +516,18 @@ mod tests { let (tx, rx) = mpsc::channel(64); let cid = "conv-1".to_owned(); - tokio::spawn(domain_event_consumer(cid, rx, repo.clone())); + tokio::spawn(domain_event_consumer("user-1".to_owned(), cid, rx, repo.clone())); tx.send(AcpSessionEvent::DesiredModeChanged { mode: "plan".into() }) .await .unwrap(); sleep(Duration::from_millis(700)).await; - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .unwrap(); assert!( state.current_mode_id.is_none(), "DesiredModeChanged is reconcile/UI-only; persistence only follows Observed*", @@ -468,7 +543,7 @@ mod tests { let (tx, rx) = mpsc::channel(64); let cid = "conv-1".to_owned(); - tokio::spawn(domain_event_consumer(cid, rx, repo.clone())); + tokio::spawn(domain_event_consumer("user-1".to_owned(), cid, rx, repo.clone())); tx.send(AcpSessionEvent::ObservedContextUsageChanged { usage_json: r#"{"used":12345,"size":200000}"#.to_owned(), @@ -477,7 +552,11 @@ mod tests { .unwrap(); sleep(Duration::from_millis(700)).await; - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .unwrap(); let raw = state.context_usage_json.expect("usage must be persisted"); let parsed: serde_json::Value = serde_json::from_str(&raw).unwrap(); assert_eq!(parsed["used"], 12345); @@ -492,7 +571,7 @@ mod tests { let (tx, rx) = mpsc::channel(64); let cid = "conv-1".to_owned(); - tokio::spawn(domain_event_consumer(cid, rx, repo.clone())); + tokio::spawn(domain_event_consumer("user-1".to_owned(), cid, rx, repo.clone())); tx.send(AcpSessionEvent::SessionAssigned { session_id: SessionId::new("sess-42"), @@ -503,7 +582,9 @@ mod tests { // Well under the debounce window — the event must have already // been written. sleep(Duration::from_millis(100)).await; - let row = repo.get("conv-1").await.unwrap().unwrap(); + let row = repo.get_for_user("user-1", "conv-1").await.unwrap().unwrap(); assert_eq!(row.session_id.as_deref(), Some("sess-42")); + drop(tx); + sleep(Duration::from_millis(50)).await; } } diff --git a/crates/aionui-ai-agent/src/protocol/cli_detect.rs b/crates/aionui-ai-agent/src/protocol/cli_detect.rs index 6cf4f5756..d40ada5cf 100644 --- a/crates/aionui-ai-agent/src/protocol/cli_detect.rs +++ b/crates/aionui-ai-agent/src/protocol/cli_detect.rs @@ -13,10 +13,10 @@ pub(crate) struct CliHealthCheckResult { /// Perform a health check for an ACP backend. /// /// Checks CLI availability and returns an availability/error pair. -pub(crate) async fn health_check(registry: &Arc, backend: &str) -> CliHealthCheckResult { +pub(crate) async fn health_check(registry: &Arc, user_id: &str, backend: &str) -> CliHealthCheckResult { let start = Instant::now(); - let Some(meta) = registry.find_builtin_by_backend(backend).await else { + let Some(meta) = registry.find_builtin_by_backend_for_user(user_id, backend).await else { return CliHealthCheckResult { available: false, error: Some(format!("No agent_metadata row for backend '{backend}'")), diff --git a/crates/aionui-ai-agent/src/protocol/error.rs b/crates/aionui-ai-agent/src/protocol/error.rs index c6f579c00..e22910f17 100644 --- a/crates/aionui-ai-agent/src/protocol/error.rs +++ b/crates/aionui-ai-agent/src/protocol/error.rs @@ -63,6 +63,7 @@ impl CloseReason { Some(AgentKillReason::RuntimeCapabilityChanged) => { "Agent killed: runtime capability changed".to_owned() } + Some(AgentKillReason::SessionRevoked) => "Agent killed: session revoked".to_owned(), None => "Agent killed".to_owned(), }, CloseReason::ProcessExited { diff --git a/crates/aionui-ai-agent/src/registry.rs b/crates/aionui-ai-agent/src/registry.rs index 590520650..774e63915 100644 --- a/crates/aionui-ai-agent/src/registry.rs +++ b/crates/aionui-ai-agent/src/registry.rs @@ -39,10 +39,12 @@ use crate::manager::acp::config_option_catalog::{ /// before producers start to back off. const CATALOG_SYNC_CHANNEL_CAPACITY: usize = 256; const CLI_PROBE_CONCURRENCY: usize = 8; +const SYSTEM_DEFAULT_USER_ID: &str = "system_default_user"; /// One unit of work submitted to the catalog sync consumer task. #[derive(Debug)] struct CatalogSyncMessage { + user_id: String, agent_metadata_id: String, handshake: AgentHandshake, } @@ -86,8 +88,12 @@ impl AgentRegistry { fn spawn_catalog_consumer(self: Arc, mut rx: mpsc::Receiver) { tokio::spawn(async move { while let Some(msg) = rx.recv().await { - if let Err(err) = self.apply_handshake_inner(&msg.agent_metadata_id, &msg.handshake).await { + if let Err(err) = self + .apply_handshake_inner(&msg.user_id, &msg.agent_metadata_id, &msg.handshake) + .await + { warn!( + user_id = %msg.user_id, agent_metadata_id = %msg.agent_metadata_id, error = %err, "Catalog sync: apply_handshake failed" @@ -105,13 +111,15 @@ impl AgentRegistry { /// tests and the consumer itself. /// /// `None` fields are left untouched (partial update). - async fn apply_handshake_inner(&self, id: &str, snapshot: &AgentHandshake) -> Result<(), AgentError> { + async fn apply_handshake_inner( + &self, + user_id: &str, + id: &str, + snapshot: &AgentHandshake, + ) -> Result<(), AgentError> { let mut snapshot = snapshot.clone(); if let Some(incoming_config_options) = snapshot.config_options.as_ref() { - let existing_config_options = { - let guard = self.by_id.read().await; - guard.get(id).and_then(|meta| meta.handshake.config_options.clone()) - }; + let existing_config_options = self.existing_config_options_for_user(user_id, id).await?; if let Some(merged_config_options) = merge_config_option_values(existing_config_options.as_ref(), incoming_config_options) { @@ -136,16 +144,17 @@ impl AgentRegistry { available_commands: available_commands.as_deref().map(Some), }; - let Some(row) = self - .repo - .apply_handshake(id, ¶ms) - .await - .map_err(|e| AgentError::internal(format!("apply_handshake: {e}")))? - else { + let row = if user_id == SYSTEM_DEFAULT_USER_ID { + self.repo.apply_handshake(id, ¶ms).await + } else { + self.repo.apply_handshake_for_user(user_id, id, ¶ms).await + }; + let Some(row) = row.map_err(|e| AgentError::internal(format!("apply_handshake: {e}")))? else { return Ok(()); }; - if let Some((mut meta, reason)) = decode_row(row, AvailabilityProjection::Cached) { + let can_update_global_cache = row.user_id.is_none() || row.user_id.as_deref() == Some(SYSTEM_DEFAULT_USER_ID); + if can_update_global_cache && let Some((mut meta, reason)) = decode_row(row, AvailabilityProjection::Cached) { let existing_availability = self .by_id .read() @@ -169,6 +178,18 @@ impl AgentRegistry { Ok(()) } + async fn existing_config_options_for_user(&self, user_id: &str, id: &str) -> Result, AgentError> { + let row = if user_id == SYSTEM_DEFAULT_USER_ID { + self.repo.get(id).await + } else { + self.repo.get_for_user(user_id, id).await + } + .map_err(|e| AgentError::internal(format!("load agent_metadata '{id}' for handshake merge: {e}")))?; + Ok(row + .and_then(|row| decode_row(row, AvailabilityProjection::Cached)) + .and_then(|(meta, _)| meta.handshake.config_options)) + } + async fn update_cached_unavailable_reason(&self, id: &str, reason: Option) { let mut guard = self.unavailable_reasons.write().await; if let Some(reason) = reason { @@ -298,6 +319,17 @@ impl AgentRegistry { .cloned() } + /// User-scoped builtin lookup for runtime paths that must honor per-user + /// agent overrides while still seeing global builtin rows. + pub async fn find_builtin_by_backend_for_user(&self, user_id: &str, vendor: &str) -> Option { + let row = self + .repo + .find_builtin_by_backend_for_user(user_id, vendor) + .await + .ok()??; + decode_row(row, AvailabilityProjection::Cached).map(|(meta, _)| meta) + } + /// Every enabled, installed row whose `agent_type` matches, /// sorted by `sort_order`. See [`Self::list_all`] for the filter /// semantics. @@ -352,100 +384,96 @@ impl AgentRegistry { .values() .cloned() .map(|meta| { - let reason = reasons.get(&meta.id); - let status = derive_management_status(&meta, reason); - let diagnostics = derive_management_diagnostics(&meta, status, reason); - let handshake = meta.handshake; - AgentManagementRow { - id: meta.id, - icon: meta.icon, - name: meta.name, - name_i18n: meta.name_i18n, - description: meta.description, - description_i18n: meta.description_i18n, - backend: meta.backend, - agent_type: meta.agent_type, - agent_source: meta.agent_source, - agent_source_info: meta.agent_source_info, - enabled: meta.enabled, - installed: meta.available, - command: meta.command, - args: meta.args, - env: Vec::new(), - native_skills_dirs: meta.native_skills_dirs, - behavior_policy: meta.behavior_policy, - yolo_id: meta.yolo_id, - config_options: handshake.config_options.clone(), - available_modes: handshake.available_modes.clone(), - available_models: handshake.available_models.clone(), - available_commands: handshake.available_commands.clone(), - sort_order: meta.sort_order, - team_capable: meta.team_capable, - status, - last_check_status: meta.last_check_status, - last_check_kind: meta.last_check_kind, - last_check_error_code: diagnostics.error_code, - last_check_error_message: diagnostics.error_message, - last_check_error_details: diagnostics.details, - last_check_guidance: diagnostics.guidance, - last_check_latency_ms: meta.last_check_latency_ms, - last_check_at: meta.last_check_at, - last_success_at: meta.last_success_at, - last_failure_at: meta.last_failure_at, - has_command_override: meta.has_command_override, - env_override_key_count: meta.env_override_key_count, - } + let reason = reasons.get(&meta.id).cloned(); + agent_management_row(meta, reason.as_ref()) }) .collect(); rows.sort_by(|a, b| a.sort_order.cmp(&b.sort_order).then_with(|| a.name.cmp(&b.name))); rows } + pub async fn list_management_rows_for_user(&self, user_id: &str) -> Result, AgentError> { + let rows = self + .repo + .list_all_for_user(user_id) + .await + .map_err(|e| AgentError::internal(format!("load agent_metadata for user: {e}")))?; + let cached_rows = self.by_id.read().await.clone(); + let cached_reasons = self.unavailable_reasons.read().await.clone(); + let mut rows: Vec = rows + .into_iter() + .filter_map(|row| { + let can_use_global_cache = + row.user_id.is_none() || row.user_id.as_deref() == Some(SYSTEM_DEFAULT_USER_ID); + decode_row(row, AvailabilityProjection::Cached).map(|decoded| (decoded, can_use_global_cache)) + }) + .map(|((mut meta, mut reason), can_use_global_cache)| { + let cached_meta = can_use_global_cache.then(|| cached_rows.get(&meta.id)).flatten(); + let cached_reason = can_use_global_cache.then(|| cached_reasons.get(&meta.id)).flatten(); + if cached_meta.is_some() { + (meta, reason) = overlay_hydrated_availability(meta, reason, cached_meta, cached_reason); + } else if !has_availability_snapshot(&meta) { + reason = apply_probe_availability(&mut meta); + } + agent_management_row(meta, reason.as_ref()) + }) + .collect(); + rows.sort_by(|a, b| a.sort_order.cmp(&b.sort_order).then_with(|| a.name.cmp(&b.name))); + Ok(rows) + } + pub async fn management_row_by_id(&self, id: &str) -> Option { let reason = self.unavailable_reasons.read().await.get(id).cloned(); let meta = self.by_id.read().await.get(id).cloned()?; - let status = derive_management_status(&meta, reason.as_ref()); - let diagnostics = derive_management_diagnostics(&meta, status, reason.as_ref()); - let handshake = meta.handshake.clone(); - Some(AgentManagementRow { - id: meta.id, - icon: meta.icon, - name: meta.name, - name_i18n: meta.name_i18n, - description: meta.description, - description_i18n: meta.description_i18n, - backend: meta.backend, - agent_type: meta.agent_type, - agent_source: meta.agent_source, - agent_source_info: meta.agent_source_info, - enabled: meta.enabled, - installed: meta.available, - command: meta.command, - args: meta.args, - env: Vec::new(), - native_skills_dirs: meta.native_skills_dirs, - behavior_policy: meta.behavior_policy, - yolo_id: meta.yolo_id, - config_options: handshake.config_options.clone(), - available_modes: handshake.available_modes.clone(), - available_models: handshake.available_models.clone(), - available_commands: handshake.available_commands.clone(), - sort_order: meta.sort_order, - team_capable: meta.team_capable, - status, - last_check_status: meta.last_check_status, - last_check_kind: meta.last_check_kind, - last_check_error_code: diagnostics.error_code, - last_check_error_message: diagnostics.error_message, - last_check_error_details: diagnostics.details, - last_check_guidance: diagnostics.guidance, - last_check_latency_ms: meta.last_check_latency_ms, - last_check_at: meta.last_check_at, - last_success_at: meta.last_success_at, - last_failure_at: meta.last_failure_at, - has_command_override: meta.has_command_override, - env_override_key_count: meta.env_override_key_count, - }) + Some(agent_management_row(meta, reason.as_ref())) + } + + pub async fn get_for_user(&self, user_id: &str, id: &str) -> Result, AgentError> { + let row = self + .repo + .get_for_user(user_id, id) + .await + .map_err(|e| AgentError::internal(format!("load agent_metadata '{id}' for user: {e}")))?; + let Some((meta, _)) = row.and_then(|row| decode_row(row, AvailabilityProjection::Probe)) else { + return Ok(None); + }; + let (meta, _) = validate_cli_availability(meta, None).await; + Ok(Some(meta)) + } + + pub async fn management_row_by_id_for_user( + &self, + user_id: &str, + id: &str, + ) -> Result, AgentError> { + let row = self + .repo + .get_for_user(user_id, id) + .await + .map_err(|e| AgentError::internal(format!("load agent_metadata '{id}' for user: {e}")))?; + let Some(row) = row else { + return Ok(None); + }; + let can_use_global_cache = row.user_id.is_none() || row.user_id.as_deref() == Some(SYSTEM_DEFAULT_USER_ID); + let Some((mut meta, mut reason)) = decode_row(row, AvailabilityProjection::Cached) else { + return Ok(None); + }; + let cached_meta = if can_use_global_cache { + self.by_id.read().await.get(&meta.id).cloned() + } else { + None + }; + let cached_reason = if can_use_global_cache { + self.unavailable_reasons.read().await.get(&meta.id).cloned() + } else { + None + }; + if cached_meta.is_some() { + (meta, reason) = overlay_hydrated_availability(meta, reason, cached_meta.as_ref(), cached_reason.as_ref()); + } else if !has_availability_snapshot(&meta) { + reason = apply_probe_availability(&mut meta); + } + Ok(Some(agent_management_row(meta, reason.as_ref()))) } /// Like [`Self::list_all_including_hidden`] but pairs every row @@ -493,6 +521,91 @@ fn is_visible(meta: &AgentMetadata) -> bool { meta.enabled && matches!(derive_management_status(meta, None), AgentManagementStatus::Online) } +fn agent_management_row(meta: AgentMetadata, reason: Option<&UnavailableReason>) -> AgentManagementRow { + let status = derive_management_status(&meta, reason); + let diagnostics = derive_management_diagnostics(&meta, status, reason); + let handshake = meta.handshake; + AgentManagementRow { + id: meta.id, + icon: meta.icon, + name: meta.name, + name_i18n: meta.name_i18n, + description: meta.description, + description_i18n: meta.description_i18n, + backend: meta.backend, + agent_type: meta.agent_type, + agent_source: meta.agent_source, + agent_source_info: meta.agent_source_info, + enabled: meta.enabled, + installed: meta.available, + command: meta.command, + args: meta.args, + env: Vec::new(), + native_skills_dirs: meta.native_skills_dirs, + behavior_policy: meta.behavior_policy, + yolo_id: meta.yolo_id, + config_options: handshake.config_options.clone(), + available_modes: handshake.available_modes.clone(), + available_models: handshake.available_models.clone(), + available_commands: handshake.available_commands.clone(), + sort_order: meta.sort_order, + team_capable: meta.team_capable, + status, + last_check_status: meta.last_check_status, + last_check_kind: meta.last_check_kind, + last_check_error_code: diagnostics.error_code, + last_check_error_message: diagnostics.error_message, + last_check_error_details: diagnostics.details, + last_check_guidance: diagnostics.guidance, + last_check_latency_ms: meta.last_check_latency_ms, + last_check_at: meta.last_check_at, + last_success_at: meta.last_success_at, + last_failure_at: meta.last_failure_at, + has_command_override: meta.has_command_override, + env_override_key_count: meta.env_override_key_count, + } +} + +fn overlay_hydrated_availability( + mut meta: AgentMetadata, + reason: Option, + cached_meta: Option<&AgentMetadata>, + cached_reason: Option<&UnavailableReason>, +) -> (AgentMetadata, Option) { + let Some(cached_meta) = cached_meta else { + return (meta, reason); + }; + if has_availability_snapshot(&meta) || !same_runtime_availability_inputs(&meta, cached_meta) { + return (meta, reason); + } + + meta.available = cached_meta.available; + meta.resolved_command = cached_meta.resolved_command.clone(); + let reason = if meta.available { + None + } else { + cached_reason.cloned().or(reason) + }; + (meta, reason) +} + +fn same_runtime_availability_inputs(a: &AgentMetadata, b: &AgentMetadata) -> bool { + a.enabled == b.enabled + && a.agent_source == b.agent_source + && a.agent_type == b.agent_type + && a.backend == b.backend + && a.command == b.command + && a.args == b.args + && json_values_equal(&a.env, &b.env) + && a.native_skills_dirs == b.native_skills_dirs + && json_values_equal(&a.behavior_policy, &b.behavior_policy) + && a.yolo_id == b.yolo_id +} + +fn json_values_equal(a: &T, b: &T) -> bool { + serde_json::to_value(a).ok() == serde_json::to_value(b).ok() +} + /// Extract and trim a command override, filtering out empty strings. fn meta_command_override(raw: &Option) -> Option { raw.as_deref() @@ -1137,8 +1250,9 @@ impl CatalogSender { /// Submit a partial handshake update. Returns without error when the /// channel is closed (only happens at shutdown) or full — callers do /// not need to care because the consumer is best-effort. - pub fn send_partial(&self, agent_metadata_id: String, handshake: AgentHandshake) { + pub fn send_partial(&self, user_id: String, agent_metadata_id: String, handshake: AgentHandshake) { let msg = CatalogSyncMessage { + user_id, agent_metadata_id, handshake, }; @@ -1273,7 +1387,9 @@ fn probe_command_candidate(command: &str) -> Option { #[cfg(test)] mod tests { use super::*; - use aionui_db::{SqliteAgentMetadataRepository, init_database_memory}; + use aionui_db::{ + IAgentMetadataRepository, SqliteAgentMetadataRepository, UpsertAgentMetadataParams, init_database_memory, + }; async fn registry() -> Arc { let db = init_database_memory().await.unwrap(); @@ -1283,6 +1399,35 @@ mod tests { reg } + fn custom_params<'a>(id: &'a str, name: &'a str) -> UpsertAgentMetadataParams<'a> { + UpsertAgentMetadataParams { + id, + icon: None, + name, + name_i18n: None, + description: Some("test custom agent"), + description_i18n: None, + backend: None, + agent_type: "acp", + agent_source: "custom", + agent_source_info: Some(r#"{"binary_name":"sh"}"#), + enabled: true, + command: Some("sh"), + args: Some("[]"), + env: Some("[]"), + native_skills_dirs: None, + behavior_policy: None, + yolo_id: None, + agent_capabilities: None, + auth_methods: None, + config_options: None, + available_modes: None, + available_models: None, + available_commands: None, + sort_order: 1500, + } + } + #[tokio::test] async fn hydrate_loads_seed_rows() { // `list_all_including_hidden` bypasses the available/enabled @@ -1432,13 +1577,136 @@ mod tests { ])), ..Default::default() }; - reg.apply_handshake_inner(&claude.id, &snapshot).await.unwrap(); + reg.apply_handshake_inner(SYSTEM_DEFAULT_USER_ID, &claude.id, &snapshot) + .await + .unwrap(); let refreshed = reg.get(&claude.id).await.unwrap(); let methods = refreshed.handshake.auth_methods.unwrap(); assert_eq!(methods.as_array().unwrap().len(), 1); } + #[tokio::test] + async fn apply_handshake_scopes_custom_agent_by_user() { + let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES ('user-b', 'local', 'user-b', 'hash', 'active', 0, 1, 1)", + ) + .execute(db.pool()) + .await + .unwrap(); + let repo = Arc::new(SqliteAgentMetadataRepository::new(db.pool().clone())); + repo.upsert_for_user( + SYSTEM_DEFAULT_USER_ID, + &custom_params("custom-shared", "Default Custom"), + ) + .await + .unwrap(); + repo.upsert_for_user("user-b", &custom_params("custom-shared", "User B Custom")) + .await + .unwrap(); + + let reg = AgentRegistry::new(repo.clone()); + reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, + "custom-shared", + &AgentHandshake { + config_options: Some(serde_json::json!({ + "config_options": [ + { + "id": "mode", + "name": "Mode", + "type": "select", + "category": "mode", + "current_value": "default-mode", + "options": [{"value": "default-mode", "name": "Default Mode"}] + } + ] + })), + ..Default::default() + }, + ) + .await + .unwrap(); + reg.apply_handshake_inner( + "user-b", + "custom-shared", + &AgentHandshake { + auth_methods: Some(serde_json::json!([{"type":"agent","id":"oauth"}])), + config_options: Some(serde_json::json!({ + "config_options": [ + { + "id": "model", + "name": "Model", + "type": "select", + "category": "model", + "current_value": "user-b-model", + "options": [{"value": "user-b-model", "name": "User B Model"}] + } + ] + })), + ..Default::default() + }, + ) + .await + .unwrap(); + + let default_row = repo + .get_for_user(SYSTEM_DEFAULT_USER_ID, "custom-shared") + .await + .unwrap() + .unwrap(); + let user_b_row = repo.get_for_user("user-b", "custom-shared").await.unwrap().unwrap(); + assert!(default_row.auth_methods.is_none()); + assert!( + default_row + .config_options + .as_deref() + .is_some_and(|value| value.contains(r#""id":"mode""#)) + ); + assert!( + default_row + .config_options + .as_deref() + .is_none_or(|value| !value.contains(r#""id":"model""#)) + ); + assert_eq!( + user_b_row.auth_methods.as_deref(), + Some(r#"[{"id":"oauth","type":"agent"}]"#) + ); + assert!( + user_b_row + .config_options + .as_deref() + .is_some_and(|value| value.contains(r#""id":"model""#)) + ); + assert!( + user_b_row + .config_options + .as_deref() + .is_none_or(|value| !value.contains(r#""id":"mode""#)) + ); + + let default_management = reg + .management_row_by_id_for_user(SYSTEM_DEFAULT_USER_ID, "custom-shared") + .await + .unwrap() + .unwrap(); + let default_options = default_management.config_options.unwrap(); + assert!(default_options.to_string().contains(r#""mode""#)); + assert!(!default_options.to_string().contains(r#""model""#)); + + let user_b_management = reg + .management_row_by_id_for_user("user-b", "custom-shared") + .await + .unwrap() + .unwrap(); + let user_b_options = user_b_management.config_options.unwrap(); + assert!(user_b_options.to_string().contains(r#""model""#)); + assert!(!user_b_options.to_string().contains(r#""mode""#)); + } + /// Partial updates must leave unrelated columns untouched. /// /// Three consecutive writes target three different columns — each @@ -1455,6 +1723,7 @@ mod tests { // Write #1: agent_capabilities only. reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &claude.id, &AgentHandshake { agent_capabilities: Some(serde_json::json!({"load_session": true})), @@ -1466,6 +1735,7 @@ mod tests { // Write #2: auth_methods only. Capabilities must survive. reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &claude.id, &AgentHandshake { auth_methods: Some(serde_json::json!([{"type": "agent", "id": "oauth"}])), @@ -1477,6 +1747,7 @@ mod tests { // Write #3: available_modes only. Capabilities + auth_methods must survive. reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &claude.id, &AgentHandshake { available_modes: Some(serde_json::json!([{"id": "code", "name": "Code"}])), @@ -1544,6 +1815,7 @@ mod tests { let claude = reg.find_builtin_by_backend("claude").await.unwrap(); reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &claude.id, &AgentHandshake { agent_capabilities: Some(serde_json::json!({"x": 1})), @@ -1553,7 +1825,7 @@ mod tests { .await .unwrap(); - reg.apply_handshake_inner(&claude.id, &AgentHandshake::default()) + reg.apply_handshake_inner(SYSTEM_DEFAULT_USER_ID, &claude.id, &AgentHandshake::default()) .await .unwrap(); @@ -1589,6 +1861,7 @@ mod tests { use aionui_db::AgentMetadataRow; let row = AgentMetadataRow { id: "test-agent".to_string(), + user_id: None, icon: None, name: "Test Agent".to_string(), name_i18n: None, @@ -1635,6 +1908,7 @@ mod tests { use aionui_db::AgentMetadataRow; let row = AgentMetadataRow { id: "632f31d2".to_string(), + user_id: None, icon: None, name: "Aion CLI".to_string(), name_i18n: None, @@ -1690,6 +1964,7 @@ mod tests { use aionui_db::AgentMetadataRow; let row = AgentMetadataRow { id: "test-agent-2".to_string(), + user_id: None, icon: None, name: "Test Agent 2".to_string(), name_i18n: None, diff --git a/crates/aionui-ai-agent/src/registry_config_option_tests.rs b/crates/aionui-ai-agent/src/registry_config_option_tests.rs index 4ddec5a21..397124148 100644 --- a/crates/aionui-ai-agent/src/registry_config_option_tests.rs +++ b/crates/aionui-ai-agent/src/registry_config_option_tests.rs @@ -5,7 +5,7 @@ use aionui_db::{SqliteAgentMetadataRepository, init_database_memory}; use crate::manager::acp::config_option_catalog::extract_config_options_from_value; -use super::AgentRegistry; +use super::{AgentRegistry, SYSTEM_DEFAULT_USER_ID}; async fn registry() -> Arc { let db = init_database_memory().await.unwrap(); @@ -21,6 +21,7 @@ async fn apply_handshake_derives_catalogs_from_config_options_before_persisting( let opencode = reg.find_builtin_by_backend("opencode").await.unwrap(); reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &opencode.id, &AgentHandshake { config_options: Some(serde_json::json!({ @@ -90,6 +91,7 @@ async fn apply_handshake_falls_back_to_available_catalogs_when_config_options_ha }); reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &opencode.id, &AgentHandshake { config_options: Some(serde_json::json!([ @@ -127,6 +129,7 @@ async fn apply_handshake_prefers_config_options_over_available_catalogs() { }); reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &opencode.id, &AgentHandshake { available_modes: Some(existing_modes), @@ -138,6 +141,7 @@ async fn apply_handshake_prefers_config_options_over_available_catalogs() { .unwrap(); reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &opencode.id, &AgentHandshake { config_options: Some(serde_json::json!({ @@ -206,6 +210,7 @@ async fn apply_handshake_config_only_partial_prefers_config_options_over_existin }); reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &opencode.id, &AgentHandshake { available_modes: Some(explicit_modes.clone()), @@ -217,6 +222,7 @@ async fn apply_handshake_config_only_partial_prefers_config_options_over_existin .unwrap(); reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &opencode.id, &AgentHandshake { config_options: Some(serde_json::json!({ @@ -267,6 +273,7 @@ async fn apply_handshake_merges_partial_config_option_updates_before_persisting( let opencode = reg.find_builtin_by_backend("opencode").await.unwrap(); reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &opencode.id, &AgentHandshake { config_options: Some(serde_json::json!({ @@ -307,6 +314,7 @@ async fn apply_handshake_merges_partial_config_option_updates_before_persisting( .unwrap(); reg.apply_handshake_inner( + SYSTEM_DEFAULT_USER_ID, &opencode.id, &AgentHandshake { config_options: Some(serde_json::json!({ diff --git a/crates/aionui-ai-agent/src/routes/agent.rs b/crates/aionui-ai-agent/src/routes/agent.rs index ad6614cd4..b50c5ddb4 100644 --- a/crates/aionui-ai-agent/src/routes/agent.rs +++ b/crates/aionui-ai-agent/src/routes/agent.rs @@ -55,12 +55,12 @@ async fn list_agent_logos( async fn list_management_agents( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, ) -> Result>>, ApiError> { Ok(Json(ApiResponse::ok( state .service - .list_management_agents() + .list_management_agents(&user.id) .await .map_err(agent_error_to_api_error)?, ))) @@ -68,13 +68,13 @@ async fn list_management_agents( async fn health_check_by_id( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { Ok(Json(ApiResponse::ok( state .service - .health_check_agent_by_id(&id) + .health_check_agent_by_id(&user.id, &id) .await .map_err(agent_error_to_api_error)?, ))) @@ -82,14 +82,14 @@ async fn health_check_by_id( async fn provider_health_check( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; Ok(Json(ApiResponse::ok( state .service - .provider_health_check(req) + .provider_health_check(&user.id, req) .await .map_err(agent_error_to_api_error)?, ))) @@ -97,14 +97,14 @@ async fn provider_health_check( async fn try_connect_custom( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; Ok(Json(ApiResponse::ok( state .service - .try_connect_custom_agent(req) + .try_connect_custom_agent(&user.id, req) .await .map_err(agent_error_to_api_error)?, ))) @@ -112,14 +112,14 @@ async fn try_connect_custom( async fn create_custom( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; Ok(Json(ApiResponse::ok( state .service - .create_custom_agent(req) + .create_custom_agent(&user.id, req) .await .map_err(agent_error_to_api_error)?, ))) @@ -127,7 +127,7 @@ async fn create_custom( async fn update_custom( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { @@ -135,7 +135,7 @@ async fn update_custom( Ok(Json(ApiResponse::ok( state .service - .update_custom_agent(&id, req) + .update_custom_agent(&user.id, &id, req) .await .map_err(agent_error_to_api_error)?, ))) @@ -143,12 +143,12 @@ async fn update_custom( async fn delete_custom( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { state .service - .delete_custom_agent(&id) + .delete_custom_agent(&user.id, &id) .await .map_err(agent_error_to_api_error)?; Ok(Json(ApiResponse::ok(DeleteCustomAgentResponse { deleted: true }))) @@ -156,7 +156,7 @@ async fn delete_custom( async fn set_agent_enabled( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { @@ -164,7 +164,7 @@ async fn set_agent_enabled( Ok(Json(ApiResponse::ok( state .service - .set_agent_enabled(&id, req.enabled) + .set_agent_enabled(&user.id, &id, req.enabled) .await .map_err(agent_error_to_api_error)?, ))) @@ -172,13 +172,13 @@ async fn set_agent_enabled( async fn get_agent_overrides( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { Ok(Json(ApiResponse::ok( state .service - .get_agent_overrides(&id) + .get_agent_overrides(&user.id, &id) .await .map_err(agent_error_to_api_error)?, ))) @@ -186,7 +186,7 @@ async fn get_agent_overrides( async fn set_agent_overrides( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { @@ -194,7 +194,7 @@ async fn set_agent_overrides( Ok(Json(ApiResponse::ok( state .service - .set_agent_overrides(&id, req) + .set_agent_overrides(&user.id, &id, req) .await .map_err(agent_error_to_api_error)?, ))) diff --git a/crates/aionui-ai-agent/src/routes/remote.rs b/crates/aionui-ai-agent/src/routes/remote.rs index 3da7b37de..4bdaf320c 100644 --- a/crates/aionui-ai-agent/src/routes/remote.rs +++ b/crates/aionui-ai-agent/src/routes/remote.rs @@ -42,48 +42,64 @@ pub fn remote_agent_routes(state: RemoteAgentRouterState) -> Router { async fn list( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, ) -> Result>>, ApiError> { - let items = state.service.list().await.map_err(agent_error_to_api_error)?; + let items = state.service.list(&user.id).await.map_err(agent_error_to_api_error)?; Ok(Json(ApiResponse::ok(items))) } async fn get_one( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - let agent = state.service.get(&id).await.map_err(agent_error_to_api_error)?; + let agent = state + .service + .get(&user.id, &id) + .await + .map_err(agent_error_to_api_error)?; Ok(Json(ApiResponse::ok(agent))) } async fn create( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result<(StatusCode, Json>), ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let agent = state.service.create(req).await.map_err(agent_error_to_api_error)?; + let agent = state + .service + .create(&user.id, req) + .await + .map_err(agent_error_to_api_error)?; Ok((StatusCode::CREATED, Json(ApiResponse::ok(agent)))) } async fn update( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let agent = state.service.update(&id, req).await.map_err(agent_error_to_api_error)?; + let agent = state + .service + .update(&user.id, &id, req) + .await + .map_err(agent_error_to_api_error)?; Ok(Json(ApiResponse::ok(agent))) } async fn delete_one( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - state.service.delete(&id).await.map_err(agent_error_to_api_error)?; + state + .service + .delete(&user.id, &id) + .await + .map_err(agent_error_to_api_error)?; Ok(Json(ApiResponse::success())) } @@ -103,9 +119,13 @@ async fn test_connection( async fn handshake( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - let resp = state.service.handshake(&id).await.map_err(agent_error_to_api_error)?; + let resp = state + .service + .handshake(&user.id, &id) + .await + .map_err(agent_error_to_api_error)?; Ok(Json(ApiResponse::ok(resp))) } diff --git a/crates/aionui-ai-agent/src/runtime_status.rs b/crates/aionui-ai-agent/src/runtime_status.rs index 09fe195de..c017b46fc 100644 --- a/crates/aionui-ai-agent/src/runtime_status.rs +++ b/crates/aionui-ai-agent/src/runtime_status.rs @@ -9,10 +9,12 @@ use aionui_runtime::{NodeRuntimeFailureKind, NodeRuntimeProgress, SharedNodeRunt pub(crate) fn conversation_runtime_reporter( broadcaster: Arc, + user_id: impl Into, conversation_id: impl Into, ) -> SharedNodeRuntimeProgressReporter { node_runtime_reporter( broadcaster, + Some(user_id.into()), RuntimeStatusScope { kind: RuntimeStatusScopeKind::Conversation, id: conversation_id.into(), @@ -22,10 +24,12 @@ pub(crate) fn conversation_runtime_reporter( pub(crate) fn custom_agent_runtime_reporter( broadcaster: Arc, + user_id: impl Into, scope_id: impl Into, ) -> SharedNodeRuntimeProgressReporter { node_runtime_reporter( broadcaster, + Some(user_id.into()), RuntimeStatusScope { kind: RuntimeStatusScopeKind::CustomAgent, id: scope_id.into(), @@ -35,10 +39,12 @@ pub(crate) fn custom_agent_runtime_reporter( fn node_runtime_reporter( broadcaster: Arc, + user_id: Option, scope: RuntimeStatusScope, ) -> SharedNodeRuntimeProgressReporter { Arc::new(move |update: NodeRuntimeProgress| { let payload = RuntimeStatusPayload { + user_id: user_id.clone(), resource: RuntimeResourceKind::Node, resource_id: None, scope: scope.clone(), @@ -76,3 +82,53 @@ fn map_failure_kind(kind: NodeRuntimeFailureKind) -> RuntimeFailureKind { NodeRuntimeFailureKind::Unknown => RuntimeFailureKind::Unknown, } } + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use aionui_runtime::{NodeRuntimeProgress, NodeRuntimeProgressPhase}; + + use super::*; + + struct RecordingBroadcaster { + events: Mutex>>, + } + + impl RecordingBroadcaster { + fn new() -> Self { + Self { + events: Mutex::new(Vec::new()), + } + } + + fn events(&self) -> Vec> { + self.events.lock().unwrap().clone() + } + } + + impl EventBroadcaster for RecordingBroadcaster { + fn broadcast(&self, event: WebSocketMessage) { + self.events.lock().unwrap().push(event); + } + } + + #[test] + fn conversation_runtime_reporter_scopes_event_to_user() { + let broadcaster = Arc::new(RecordingBroadcaster::new()); + let reporter = conversation_runtime_reporter(broadcaster.clone(), "user-1", "conv-1"); + + reporter.report(NodeRuntimeProgress { + phase: NodeRuntimeProgressPhase::Ready, + failure_kind: None, + message: None, + status_code: None, + }); + + let events = broadcaster.events(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].name, "runtime.statusChanged"); + assert_eq!(events[0].data["user_id"], "user-1"); + assert_eq!(events[0].data["scope"]["id"], "conv-1"); + } +} diff --git a/crates/aionui-ai-agent/src/services/agent.rs b/crates/aionui-ai-agent/src/services/agent.rs index c0789d81b..322e0a541 100644 --- a/crates/aionui-ai-agent/src/services/agent.rs +++ b/crates/aionui-ai-agent/src/services/agent.rs @@ -66,8 +66,8 @@ impl AgentService { // Agent operations impl AgentService { - pub async fn list_management_agents(&self) -> Result, AgentError> { - Ok(self.availability.list_management_rows().await) + pub async fn list_management_agents(&self, user_id: &str) -> Result, AgentError> { + self.availability.list_management_rows(user_id).await } /// Backend → logo URL catalog for business surfaces. @@ -102,27 +102,29 @@ impl AgentService { Ok(entries) } - pub async fn health_check_agent_by_id(&self, id: &str) -> Result { - self.availability.run_manual_health_check(id).await + pub async fn health_check_agent_by_id(&self, user_id: &str, id: &str) -> Result { + self.availability.run_manual_health_check(user_id, id).await } pub async fn provider_health_check( &self, + user_id: &str, req: ProviderHealthCheckRequest, ) -> Result { - self.provider_health.health_check(req).await + self.provider_health.health_check(user_id, req).await } pub async fn set_agent_overrides( &self, + user_id: &str, id: &str, req: aionui_api_types::SetAgentOverridesRequest, ) -> Result { let repo = self.registry.repo_handle(); let row = repo - .get(id) + .get_for_user(user_id, id) .await - .map_err(|e| AgentError::internal(format!("repo.get: {e}")))? + .map_err(|e| AgentError::internal(format!("repo.get_for_user: {e}")))? .ok_or_else(|| AgentError::not_found(format!("Agent '{id}' not found")))?; let command_override = req @@ -159,20 +161,24 @@ impl AgentService { _ => None, }; - repo.update_agent_overrides(id, command_override.as_deref(), env_json.as_deref()) + repo.update_agent_overrides_for_user(user_id, id, command_override.as_deref(), env_json.as_deref()) .await - .map_err(|e| AgentError::internal(format!("repo.update_agent_overrides: {e}")))?; + .map_err(|e| AgentError::internal(format!("repo.update_agent_overrides_for_user: {e}")))?; - self.availability.run_manual_health_check(id).await + self.availability.run_manual_health_check(user_id, id).await } - pub async fn get_agent_overrides(&self, id: &str) -> Result { + pub async fn get_agent_overrides( + &self, + user_id: &str, + id: &str, + ) -> Result { let row = self .registry .repo_handle() - .get(id) + .get_for_user(user_id, id) .await - .map_err(|e| AgentError::internal(format!("repo.get: {e}")))? + .map_err(|e| AgentError::internal(format!("repo.get_for_user: {e}")))? .ok_or_else(|| AgentError::not_found(format!("Agent '{id}' not found")))?; let env_override = row diff --git a/crates/aionui-ai-agent/src/services/availability/mod.rs b/crates/aionui-ai-agent/src/services/availability/mod.rs index 9ee755271..5e905f215 100644 --- a/crates/aionui-ai-agent/src/services/availability/mod.rs +++ b/crates/aionui-ai-agent/src/services/availability/mod.rs @@ -16,8 +16,14 @@ use crate::registry::{AgentRegistry, guidance_for_snapshot_error_code}; #[async_trait::async_trait] pub trait AgentAvailabilityFeedbackPort: Send + Sync { - async fn record_session_success(&self, agent_id: &str) -> Result<(), AgentError>; - async fn record_session_failure(&self, agent_id: &str, code: &str, message: &str) -> Result<(), AgentError>; + async fn record_session_success(&self, user_id: &str, agent_id: &str) -> Result<(), AgentError>; + async fn record_session_failure( + &self, + user_id: &str, + agent_id: &str, + code: &str, + message: &str, + ) -> Result<(), AgentError>; } struct AvailabilitySnapshot { @@ -45,21 +51,21 @@ impl AgentAvailabilityService { } } - pub async fn list_management_rows(&self) -> Vec { - self.registry.list_management_rows().await + pub async fn list_management_rows(&self, user_id: &str) -> Result, AgentError> { + self.registry.list_management_rows_for_user(user_id).await } - pub async fn run_manual_health_check(&self, id: &str) -> Result { + pub async fn run_manual_health_check(&self, user_id: &str, id: &str) -> Result { let meta = self .registry - .reload_one(id) - .await - .and_then(|row| row.ok_or_else(|| AgentError::not_found(format!("Agent '{id}' not found"))))?; + .get_for_user(user_id, id) + .await? + .ok_or_else(|| AgentError::not_found(format!("Agent '{id}' not found")))?; if !meta.available { return self - .management_row_by_id(id) - .await + .management_row_by_id(user_id, id) + .await? .ok_or_else(|| AgentError::not_found(format!("Agent '{id}' not found"))); } @@ -67,16 +73,23 @@ impl AgentAvailabilityService { &self.registry, &self.provider_repo, &meta, + user_id, AgentSnapshotCheckKind::Manual, ) .await; - self.persist_snapshot(id, &snapshot).await?; - self.management_row_by_id(id) - .await + self.persist_snapshot_for_user(user_id, id, &snapshot).await?; + self.management_row_by_id(user_id, id) + .await? .ok_or_else(|| AgentError::not_found(format!("Agent '{id}' not found"))) } - pub async fn record_session_failure(&self, agent_id: &str, code: &str, message: &str) -> Result<(), AgentError> { + pub async fn record_session_failure( + &self, + user_id: &str, + agent_id: &str, + code: &str, + message: &str, + ) -> Result<(), AgentError> { let checked_at = now_ms(); let snapshot = AvailabilitySnapshot { status: "offline", @@ -86,10 +99,10 @@ impl AgentAvailabilityService { latency_ms: 0, checked_at, }; - self.persist_snapshot(agent_id, &snapshot).await + self.persist_snapshot_for_user(user_id, agent_id, &snapshot).await } - pub async fn record_session_success(&self, agent_id: &str) -> Result<(), AgentError> { + pub async fn record_session_success(&self, user_id: &str, agent_id: &str) -> Result<(), AgentError> { let checked_at = now_ms(); let snapshot = AvailabilitySnapshot { status: "online", @@ -99,20 +112,29 @@ impl AgentAvailabilityService { latency_ms: 0, checked_at, }; - self.persist_snapshot(agent_id, &snapshot).await + self.persist_snapshot_for_user(user_id, agent_id, &snapshot).await } - pub async fn management_row_by_id(&self, id: &str) -> Option { - self.registry.management_row_by_id(id).await + pub async fn management_row_by_id( + &self, + user_id: &str, + id: &str, + ) -> Result, AgentError> { + self.registry.management_row_by_id_for_user(user_id, id).await } - async fn persist_snapshot(&self, id: &str, snapshot: &AvailabilitySnapshot) -> Result<(), AgentError> { + async fn persist_snapshot_for_user( + &self, + user_id: &str, + id: &str, + snapshot: &AvailabilitySnapshot, + ) -> Result<(), AgentError> { let existing = self .registry .repo_handle() - .get(id) + .get_for_user(user_id, id) .await - .map_err(|error| AgentError::internal(format!("repo.get: {error}")))? + .map_err(|error| AgentError::internal(format!("repo.get_for_user: {error}")))? .ok_or_else(|| AgentError::not_found(format!("Agent '{id}' not found")))?; let params = UpdateAgentAvailabilitySnapshotParams { @@ -139,10 +161,9 @@ impl AgentAvailabilityService { }; self.registry .repo_handle() - .update_availability_snapshot(id, ¶ms) + .update_availability_snapshot_for_user(user_id, id, ¶ms) .await - .map_err(|error| AgentError::internal(format!("repo.update_availability_snapshot: {error}")))?; - self.registry.reload_one(id).await?; + .map_err(|error| AgentError::internal(format!("repo.update_availability_snapshot_for_user: {error}")))?; Ok(()) } } @@ -151,6 +172,7 @@ async fn run_probe( registry: &Arc, provider_repo: &Arc, meta: &AgentMetadata, + user_id: &str, kind: AgentSnapshotCheckKind, ) -> AvailabilitySnapshot { let started_at = now_ms(); @@ -211,7 +233,7 @@ async fn run_probe( }, } } else if let Some(backend) = meta.backend.as_deref() { - let result = cli_detect::health_check(registry, backend).await; + let result = cli_detect::health_check(registry, user_id, backend).await; if result.available { (AgentSnapshotCheckStatus::Online, None, None) } else { @@ -226,7 +248,7 @@ async fn run_probe( // so its usability hinges entirely on having a configured model. It is // online only when at least one model provider is enabled — otherwise // it cannot run a single turn. - probe_aionrs_provider_readiness(provider_repo).await + probe_aionrs_provider_readiness(provider_repo, user_id).await } else { (AgentSnapshotCheckStatus::Online, None, None) }; @@ -272,8 +294,9 @@ fn explicit_probe_args(meta: &AgentMetadata) -> Result, String> { /// `no_provider` code the UI maps to "configure a model" guidance. async fn probe_aionrs_provider_readiness( provider_repo: &Arc, + user_id: &str, ) -> (AgentSnapshotCheckStatus, Option, Option) { - match provider_repo.list().await { + match provider_repo.list(user_id).await { Ok(providers) if providers.iter().any(|p| p.enabled) => (AgentSnapshotCheckStatus::Online, None, None), Ok(_) => ( AgentSnapshotCheckStatus::Offline, @@ -290,12 +313,18 @@ async fn probe_aionrs_provider_readiness( #[async_trait::async_trait] impl AgentAvailabilityFeedbackPort for AgentAvailabilityService { - async fn record_session_success(&self, agent_id: &str) -> Result<(), AgentError> { - AgentAvailabilityService::record_session_success(self, agent_id).await + async fn record_session_success(&self, user_id: &str, agent_id: &str) -> Result<(), AgentError> { + AgentAvailabilityService::record_session_success(self, user_id, agent_id).await } - async fn record_session_failure(&self, agent_id: &str, code: &str, message: &str) -> Result<(), AgentError> { - AgentAvailabilityService::record_session_failure(self, agent_id, code, message).await + async fn record_session_failure( + &self, + user_id: &str, + agent_id: &str, + code: &str, + message: &str, + ) -> Result<(), AgentError> { + AgentAvailabilityService::record_session_failure(self, user_id, agent_id, code, message).await } } @@ -315,9 +344,12 @@ mod tests { SqliteProviderRepository, UpsertAgentMetadataParams, init_database_memory, }; + const TEST_USER_ID: &str = "system_default_user"; + fn enabled_provider_params() -> CreateProviderParams<'static> { CreateProviderParams { id: None, + user_id: TEST_USER_ID, platform: "openai", name: "OpenAI", base_url: "https://api.openai.com", @@ -340,7 +372,7 @@ mod tests { let db = init_database_memory().await.unwrap(); let provider_repo: Arc = Arc::new(SqliteProviderRepository::new(db.pool().clone())); - let (status, code, _msg) = probe_aionrs_provider_readiness(&provider_repo).await; + let (status, code, _msg) = probe_aionrs_provider_readiness(&provider_repo, TEST_USER_ID).await; assert_eq!(status, AgentSnapshotCheckStatus::Offline); assert_eq!(code.as_deref(), Some("no_provider")); @@ -352,7 +384,7 @@ mod tests { let provider_repo: Arc = Arc::new(SqliteProviderRepository::new(db.pool().clone())); provider_repo.create(enabled_provider_params()).await.unwrap(); - let (status, code, _msg) = probe_aionrs_provider_readiness(&provider_repo).await; + let (status, code, _msg) = probe_aionrs_provider_readiness(&provider_repo, TEST_USER_ID).await; assert_eq!(status, AgentSnapshotCheckStatus::Online); assert!(code.is_none()); @@ -399,6 +431,7 @@ mod tests { let service = AgentAvailabilityService::new(registry.clone(), provider_repo); service .record_session_failure( + TEST_USER_ID, "agent-session-failure", "session_send_failed", "provider returned 401 invalid api key", @@ -407,8 +440,9 @@ mod tests { .unwrap(); let row = service - .list_management_rows() + .list_management_rows(TEST_USER_ID) .await + .unwrap() .into_iter() .find(|item| item.id == "agent-session-failure") .unwrap(); @@ -471,6 +505,7 @@ mod tests { let service = AgentAvailabilityService::new(registry.clone(), provider_repo); service .record_session_failure( + TEST_USER_ID, "agent-session-success", "session_send_failed", "provider returned 401 invalid api key", @@ -478,11 +513,15 @@ mod tests { .await .unwrap(); - service.record_session_success("agent-session-success").await.unwrap(); + service + .record_session_success(TEST_USER_ID, "agent-session-success") + .await + .unwrap(); let row = service - .list_management_rows() + .list_management_rows(TEST_USER_ID) .await + .unwrap() .into_iter() .find(|item| item.id == "agent-session-success") .unwrap(); @@ -547,7 +586,14 @@ mod tests { env_override_key_count: 0, }; - let snapshot = run_probe(®istry, &provider_repo, &meta, AgentSnapshotCheckKind::Manual).await; + let snapshot = run_probe( + ®istry, + &provider_repo, + &meta, + "system_default_user", + AgentSnapshotCheckKind::Manual, + ) + .await; assert_eq!(snapshot.status, "offline"); assert_eq!(snapshot.error_code.as_deref(), Some("command_not_found")); diff --git a/crates/aionui-ai-agent/src/services/custom.rs b/crates/aionui-ai-agent/src/services/custom.rs index 5b3bc2a62..d25b7f119 100644 --- a/crates/aionui-ai-agent/src/services/custom.rs +++ b/crates/aionui-ai-agent/src/services/custom.rs @@ -26,6 +26,7 @@ use crate::protocol::custom_agent_probe::try_connect_custom_agent as probe; use crate::runtime_status::custom_agent_runtime_reporter; const CUSTOM_SORT_ORDER_DEFAULT: i64 = 1500; +const SYSTEM_DEFAULT_USER_ID: &str = "system_default_user"; impl AgentService { /// Public accessor for the probe — powers both @@ -33,28 +34,34 @@ impl AgentService { /// below. pub async fn try_connect_custom_agent( &self, + user_id: &str, req: TryConnectCustomAgentRequest, ) -> Result { if req.command.trim().is_empty() { return Err(AgentError::bad_request("command must not be empty")); } - let reporter = req - .runtime_scope_id - .as_ref() - .map(|scope_id| custom_agent_runtime_reporter(self.broadcaster().clone(), scope_id.clone())); + let reporter = req.runtime_scope_id.as_ref().map(|scope_id| { + custom_agent_runtime_reporter(self.broadcaster().clone(), user_id.to_owned(), scope_id.clone()) + }); Ok(probe(&req.command, &req.acp_args, &req.env, reporter.as_deref()).await) } - pub async fn create_custom_agent(&self, req: CustomAgentUpsertRequest) -> Result { + pub async fn create_custom_agent( + &self, + user_id: &str, + req: CustomAgentUpsertRequest, + ) -> Result { validate_upsert(&req)?; probe_or_reject(&req).await?; let id = generate_short_id(); - self.upsert_custom_row(&id, &req, /* keep_enabled = */ true).await + self.upsert_custom_row(user_id, &id, &req, /* keep_enabled = */ true) + .await } pub async fn update_custom_agent( &self, + user_id: &str, id: &str, req: CustomAgentUpsertRequest, ) -> Result { @@ -62,9 +69,9 @@ impl AgentService { let existing = self .registry() .repo_handle() - .get(id) + .get_for_user(user_id, id) .await - .map_err(|e| AgentError::internal(format!("repo.get: {e}")))? + .map_err(|e| AgentError::internal(format!("repo.get_for_user: {e}")))? .ok_or_else(|| AgentError::not_found(format!("Agent '{id}' not found")))?; if existing.agent_source != "custom" { return Err(AgentError::forbidden( @@ -74,16 +81,16 @@ impl AgentService { probe_or_reject(&req).await?; let keep_enabled = existing.enabled; - self.upsert_custom_row(id, &req, keep_enabled).await + self.upsert_custom_row(user_id, id, &req, keep_enabled).await } - pub async fn delete_custom_agent(&self, id: &str) -> Result<(), AgentError> { + pub async fn delete_custom_agent(&self, user_id: &str, id: &str) -> Result<(), AgentError> { let existing = self .registry() .repo_handle() - .get(id) + .get_for_user(user_id, id) .await - .map_err(|e| AgentError::internal(format!("repo.get: {e}")))? + .map_err(|e| AgentError::internal(format!("repo.get_for_user: {e}")))? .ok_or_else(|| AgentError::not_found(format!("Agent '{id}' not found")))?; if existing.agent_source != "custom" { return Err(AgentError::forbidden( @@ -93,39 +100,44 @@ impl AgentService { let removed = self .registry() .repo_handle() - .delete(id) + .delete_for_user(user_id, id) .await - .map_err(|e| AgentError::internal(format!("repo.delete: {e}")))?; + .map_err(|e| AgentError::internal(format!("repo.delete_for_user: {e}")))?; if !removed { return Err(AgentError::not_found(format!("Agent '{id}' not found"))); } - if let Err(err) = self.registry().reload_one(id).await { + if user_id == SYSTEM_DEFAULT_USER_ID + && let Err(err) = self.registry().reload_one(id).await + { warn!(agent_id = %id, error = %err, "registry reload failed after delete_custom_agent"); } Ok(()) } - pub async fn set_agent_enabled(&self, id: &str, enabled: bool) -> Result { + pub async fn set_agent_enabled(&self, user_id: &str, id: &str, enabled: bool) -> Result { let updated = self .registry() .repo_handle() - .set_enabled(id, enabled) + .set_enabled_for_user(user_id, id, enabled) .await - .map_err(|e| AgentError::internal(format!("repo.set_enabled: {e}")))?; + .map_err(|e| AgentError::internal(format!("repo.set_enabled_for_user: {e}")))?; if !updated { return Err(AgentError::not_found(format!("Agent '{id}' not found"))); } - if let Err(err) = self.registry().reload_one(id).await { + if user_id == SYSTEM_DEFAULT_USER_ID + && let Err(err) = self.registry().reload_one(id).await + { warn!(agent_id = %id, error = %err, "registry reload failed after set_agent_enabled"); } self.registry() - .get(id) - .await + .get_for_user(user_id, id) + .await? .ok_or_else(|| AgentError::internal(format!("Agent '{id}' not visible after enable toggle"))) } async fn upsert_custom_row( &self, + user_id: &str, id: &str, req: &CustomAgentUpsertRequest, enabled: bool, @@ -182,18 +194,20 @@ impl AgentService { self.registry() .repo_handle() - .upsert(¶ms) + .upsert_for_user(user_id, ¶ms) .await - .map_err(|e| AgentError::internal(format!("repo.upsert: {e}")))?; + .map_err(|e| AgentError::internal(format!("repo.upsert_for_user: {e}")))?; - self.registry() - .reload_one(id) - .await - .map_err(|e| AgentError::internal(format!("registry reload: {e}")))?; + if user_id == SYSTEM_DEFAULT_USER_ID { + self.registry() + .reload_one(id) + .await + .map_err(|e| AgentError::internal(format!("registry reload: {e}")))?; + } self.registry() - .get(id) - .await + .get_for_user(user_id, id) + .await? .ok_or_else(|| AgentError::internal(format!("Agent '{id}' not visible after upsert"))) } } diff --git a/crates/aionui-ai-agent/src/services/provider_health.rs b/crates/aionui-ai-agent/src/services/provider_health.rs index 40692283e..3453f1202 100644 --- a/crates/aionui-ai-agent/src/services/provider_health.rs +++ b/crates/aionui-ai-agent/src/services/provider_health.rs @@ -44,6 +44,7 @@ impl ProviderHealthCheckService { pub async fn health_check( &self, + user_id: &str, req: ProviderHealthCheckRequest, ) -> Result { if req.provider_id.trim().is_empty() { @@ -57,7 +58,7 @@ impl ProviderHealthCheckService { let model = req.model.trim(); let row = self .provider_repo - .find_by_id(provider_id) + .find_by_id(user_id, provider_id) .await .map_err(|e| AgentError::internal(format!("Failed to load provider config: {e}")))? .ok_or_else(|| AgentError::bad_request(format!("Provider '{provider_id}' not found")))?; @@ -338,16 +339,17 @@ mod tests { use aionui_db::{CreateProviderParams, DbError, UpdateProviderParams}; const TEST_KEY: [u8; 32] = [0xAB; 32]; + const TEST_USER_ID: &str = "user-1"; struct UnusedProviderRepository; #[async_trait::async_trait] impl IProviderRepository for UnusedProviderRepository { - async fn list(&self) -> Result, DbError> { + async fn list(&self, _user_id: &str) -> Result, DbError> { unreachable!("provider repo is not used by resolve_probe_config") } - async fn find_by_id(&self, _id: &str) -> Result, DbError> { + async fn find_by_id(&self, _user_id: &str, _id: &str) -> Result, DbError> { unreachable!("provider repo is not used by resolve_probe_config") } @@ -355,11 +357,16 @@ mod tests { unreachable!("provider repo is not used by resolve_probe_config") } - async fn update(&self, _id: &str, _params: UpdateProviderParams<'_>) -> Result { + async fn update( + &self, + _user_id: &str, + _id: &str, + _params: UpdateProviderParams<'_>, + ) -> Result { unreachable!("provider repo is not used by resolve_probe_config") } - async fn delete(&self, _id: &str) -> Result<(), DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), DbError> { unreachable!("provider repo is not used by resolve_probe_config") } } @@ -375,6 +382,7 @@ mod tests { fn test_provider() -> Provider { Provider { id: "provider-1".to_owned(), + user_id: TEST_USER_ID.to_owned(), platform: "anthropic".to_owned(), name: "Test Anthropic".to_owned(), base_url: "https://api.anthropic.com".to_owned(), diff --git a/crates/aionui-ai-agent/src/services/remote.rs b/crates/aionui-ai-agent/src/services/remote.rs index 31b1b2122..e9acaa993 100644 --- a/crates/aionui-ai-agent/src/services/remote.rs +++ b/crates/aionui-ai-agent/src/services/remote.rs @@ -29,16 +29,16 @@ impl RemoteAgentService { } /// List all remote agents (auth_token omitted). - pub async fn list(&self) -> Result, AgentError> { - let rows = self.repo.list().await.map_err(db_err)?; + pub async fn list(&self, user_id: &str) -> Result, AgentError> { + let rows = self.repo.list(user_id).await.map_err(db_err)?; rows.into_iter().map(|r| self.row_to_list_item(r)).collect() } /// Get a single remote agent by ID (auth_token masked). - pub async fn get(&self, id: &str) -> Result { + pub async fn get(&self, user_id: &str, id: &str) -> Result { let row = self .repo - .find_by_id(id) + .find_by_id(user_id, id) .await .map_err(db_err)? .ok_or_else(|| AgentError::not_found(format!("Remote agent '{id}' not found")))?; @@ -46,7 +46,11 @@ impl RemoteAgentService { } /// Create a new remote agent. OpenClaw protocol auto-generates Ed25519 keys. - pub async fn create(&self, req: CreateRemoteAgentRequest) -> Result { + pub async fn create( + &self, + user_id: &str, + req: CreateRemoteAgentRequest, + ) -> Result { validate_create_request(&req)?; let encrypted_token = req @@ -65,6 +69,7 @@ impl RemoteAgentService { let row = self .repo .create(aionui_db::CreateRemoteAgentParams { + user_id, name: &req.name, protocol: &enum_to_str(&req.protocol), url: &req.url, @@ -85,7 +90,12 @@ impl RemoteAgentService { } /// Update an existing remote agent. - pub async fn update(&self, id: &str, req: UpdateRemoteAgentRequest) -> Result { + pub async fn update( + &self, + user_id: &str, + id: &str, + req: UpdateRemoteAgentRequest, + ) -> Result { let encrypted_token = match &req.auth_token { Some(Some(t)) => Some(Some( encrypt_string(t, &self.encryption_key).map_err(|e| AgentError::internal(e.to_string()))?, @@ -108,7 +118,7 @@ impl RemoteAgentService { description: req.description.as_ref().map(|o| o.as_deref()), }; - let row = self.repo.update(id, params).await.map_err(|e| match e { + let row = self.repo.update(user_id, id, params).await.map_err(|e| match e { aionui_db::DbError::NotFound(msg) => AgentError::not_found(msg), other => AgentError::internal(other.to_string()), })?; @@ -117,8 +127,8 @@ impl RemoteAgentService { } /// Delete a remote agent. - pub async fn delete(&self, id: &str) -> Result<(), AgentError> { - self.repo.delete(id).await.map_err(|e| match e { + pub async fn delete(&self, user_id: &str, id: &str) -> Result<(), AgentError> { + self.repo.delete(user_id, id).await.map_err(|e| match e { aionui_db::DbError::NotFound(msg) => AgentError::not_found(msg), other => AgentError::internal(other.to_string()), }) @@ -148,10 +158,10 @@ impl RemoteAgentService { } /// OpenClaw device handshake (15s timeout). - pub async fn handshake(&self, id: &str) -> Result { + pub async fn handshake(&self, user_id: &str, id: &str) -> Result { let row = self .repo - .find_by_id(id) + .find_by_id(user_id, id) .await .map_err(db_err)? .ok_or_else(|| AgentError::not_found(format!("Remote agent '{id}' not found")))?; @@ -180,17 +190,17 @@ impl RemoteAgentService { match connect_result { Ok(Ok(_)) => { let now = aionui_common::now_ms(); - let _ = self.repo.update_status(id, "connected", Some(now)).await; + let _ = self.repo.update_status(user_id, id, "connected", Some(now)).await; Ok(HandshakeResponse { status: "ok".to_string(), }) } Ok(Err(e)) => { - let _ = self.repo.update_status(id, "error", None).await; + let _ = self.repo.update_status(user_id, id, "error", None).await; Err(e) } Err(_) => { - let _ = self.repo.update_status(id, "error", None).await; + let _ = self.repo.update_status(user_id, id, "error", None).await; Err(AgentError::timeout("Handshake timed out after 15 seconds")) } } diff --git a/crates/aionui-ai-agent/src/session_agent.rs b/crates/aionui-ai-agent/src/session_agent.rs index 532458514..6e5d13556 100644 --- a/crates/aionui-ai-agent/src/session_agent.rs +++ b/crates/aionui-ai-agent/src/session_agent.rs @@ -268,6 +268,9 @@ fn session_content_blocks_to_json(content: &[ContentBlock]) -> Vec, runtime: Arc, @@ -308,6 +311,7 @@ impl SessionAgentTask { pub fn new( agent_type: AgentType, conversation_id: String, + user_id: String, workspace: String, backend: Arc, session_repo: Option>, @@ -315,6 +319,7 @@ impl SessionAgentTask { Self::build( agent_type, conversation_id, + user_id, workspace, backend, session_repo, @@ -332,6 +337,7 @@ impl SessionAgentTask { pub fn new_with_preload( agent_type: AgentType, conversation_id: String, + user_id: String, workspace: String, backend: Arc, session_repo: Option>, @@ -341,6 +347,7 @@ impl SessionAgentTask { Self::build( agent_type, conversation_id, + user_id, workspace, backend, session_repo, @@ -349,9 +356,11 @@ impl SessionAgentTask { ) } + #[allow(clippy::too_many_arguments)] fn build( agent_type: AgentType, conversation_id: String, + user_id: String, workspace: String, backend: Arc, session_repo: Option>, @@ -372,10 +381,17 @@ impl SessionAgentTask { // stream to the pump — never a backend Arc (see `spawn_event_pump` for why // capturing a backend Arc there would leak the child process). let events = backend.events(); - spawn_event_pump(events, runtime.clone(), conversation_id.clone(), session_repo.clone()); + spawn_event_pump( + events, + runtime.clone(), + conversation_id.clone(), + user_id.clone(), + session_repo.clone(), + ); Arc::new(Self { agent_type, conversation_id, + user_id, workspace, backend, runtime, @@ -826,7 +842,7 @@ impl SessionAgentTask { }; // Merge into the existing selection map (preserve unrelated keys). let mut selections: std::collections::HashMap = match repo - .load_runtime_state(&self.conversation_id) + .load_runtime_state_for_user(&self.user_id, &self.conversation_id) .await { Ok(Some(state)) => state @@ -852,7 +868,10 @@ impl SessionAgentTask { config_selections_json: Some(Some(&json)), ..Default::default() }; - if let Err(err) = repo.save_runtime_state(&self.conversation_id, ¶ms).await { + if let Err(err) = repo + .save_runtime_state_for_user(&self.user_id, &self.conversation_id, ¶ms) + .await + { tracing::warn!(conversation_id = %self.conversation_id, error = %err, "persist_effort: save_runtime_state failed"); } } @@ -994,6 +1013,9 @@ impl IAgentTask for SessionAgentTask { pub struct SessionBuildInputs<'a> { /// The conversation this session belongs to (the clean-slate `session_id`). pub conversation_id: String, + /// Owner Core user. Scopes every acp_session persistence write, the MCP + /// server resolution, and the catalog writeback (multi-account boundary). + pub user_id: String, /// The resolved workspace path (`SessionConfig.cwd`). pub workspace: String, /// The conversation's persisted build `extra` (mode/model/mcp/preset/skills). @@ -1118,6 +1140,7 @@ pub async fn build_session_instance( let SessionBuildInputs { conversation_id, + user_id, workspace, config, metadata, @@ -1143,6 +1166,7 @@ pub async fn build_session_instance( Some(repo) => { crate::mcp_resolve::resolve_session_mcp_servers( repo.as_ref(), + &user_id, config.mcp_server_ids.as_deref(), &conversation_id, broadcaster, @@ -1306,7 +1330,7 @@ pub async fn build_session_instance( // GAP #7 (G5): project the backend's discovered catalog back into agent_metadata // so the cold-start picker stays fresh. Best-effort, detached, off the open path. if let Some((agent_id, catalog_tx)) = catalog_writeback { - spawn_catalog_writeback(agent_id, backend.clone(), catalog_tx); + spawn_catalog_writeback(agent_id, user_id.clone(), backend.clone(), catalog_tx); } let prompt_dump = prompt_dump_dir.map(|dir| SessionPromptDump { @@ -1319,6 +1343,7 @@ pub async fn build_session_instance( let task = SessionAgentTask::new_with_preload( AgentType::Acp, conversation_id, + user_id, workspace, backend, acp_session_repo, @@ -1483,6 +1508,7 @@ fn team_mcp_server_spec(cfg: &aionui_api_types::TeamMcpStdioConfig) -> aionui_se /// forwarding the best model-less partial only if the window elapses. pub fn spawn_catalog_writeback( agent_id: String, + user_id: String, backend: Arc, catalog_tx: crate::registry::CatalogSender, ) { @@ -1493,7 +1519,7 @@ pub fn spawn_catalog_writeback( if let Some(partial) = catalog_partial_from_caps(&caps) { if !caps.available_models.is_empty() { // Complete enough — models present → commit the full catalog. - catalog_tx.send_partial(agent_id, partial); + catalog_tx.send_partial(user_id.clone(), agent_id, partial); return; } // Modes/commands only so far — remember it, keep waiting for models. @@ -1502,7 +1528,7 @@ pub fn spawn_catalog_writeback( tokio::time::sleep(std::time::Duration::from_millis(50)).await; } if let Some(partial) = best_partial { - catalog_tx.send_partial(agent_id, partial); + catalog_tx.send_partial(user_id, agent_id, partial); } }); } @@ -1682,6 +1708,7 @@ fn spawn_event_pump( mut events: BoxStream<'static, SessionEnvelope>, runtime: Arc, conversation_id: String, + user_id: String, session_repo: Option>, ) { use futures_util::StreamExt as _; @@ -2020,7 +2047,7 @@ fn spawn_event_pump( // AcpSessionSyncService (which this direct-CLI path bypasses). Best-effort: // a repo error is warn-logged, never fatal to the stream. if let Some(repo) = session_repo.as_ref() { - persist_side_effects(repo.as_ref(), &conversation_id, &env.event).await; + persist_side_effects(repo.as_ref(), &user_id, &conversation_id, &env.event).await; } for mut ev in translate_event(env.event, &conversation_id, terminal_result_seen) { // Keep the tool name alive across a call's multi-frame lifecycle (see @@ -2123,7 +2150,12 @@ fn is_dead_resume_anchor(event: &SessionEvent) -> bool { /// for the ACP-manager path. Without this the resume anchor /// (`build_session_instance` GAP #1) and the mode/model precedence source (GAP #2) /// are never written, so a restart always loses continuity. -async fn persist_side_effects(repo: &dyn IAcpSessionRepository, conversation_id: &str, event: &SessionEvent) { +async fn persist_side_effects( + repo: &dyn IAcpSessionRepository, + user_id: &str, + conversation_id: &str, + event: &SessionEvent, +) { // Self-heal a dead resume anchor: a turn that failed *because* the stored // backend session id no longer resolves must null that id, or every subsequent // send re-resumes the same dead session and the conversation wedges forever. @@ -2132,7 +2164,7 @@ async fn persist_side_effects(repo: &dyn IAcpSessionRepository, conversation_id: // direct-CLI path dropped: clean-slate `Orchestrator` emits `BackendBound{None}` // and legacy ACP does `rebuild_after_session_not_found` → `clear_session_id`. if is_dead_resume_anchor(event) { - match repo.clear_session_id(conversation_id).await { + match repo.clear_session_id_for_user(user_id, conversation_id).await { Ok(_) => tracing::info!( conversation_id, "session-sync: cleared dead resume anchor (unrecoverable resume error) — next turn opens Fresh" @@ -2150,7 +2182,7 @@ async fn persist_side_effects(repo: &dyn IAcpSessionRepository, conversation_id: SessionEvent::BackendBound { backend_session_id: Some(bid), } => { - if let Err(err) = repo.update_session_id(conversation_id, bid).await { + if let Err(err) = repo.update_session_id_for_user(user_id, conversation_id, bid).await { tracing::warn!(conversation_id, error = %err, "session-sync: update_session_id failed"); } } @@ -2163,7 +2195,10 @@ async fn persist_side_effects(repo: &dyn IAcpSessionRepository, conversation_id: config_selections_json: None, context_usage_json: None, }; - if let Err(err) = repo.save_runtime_state(conversation_id, ¶ms).await { + if let Err(err) = repo + .save_runtime_state_for_user(user_id, conversation_id, ¶ms) + .await + { tracing::warn!(conversation_id, error = %err, "session-sync: save_runtime_state failed"); } } @@ -3582,8 +3617,25 @@ mod persist_tests { // for the test's lifetime (the cloned SqlitePool keeps the in-memory DB alive). async fn seeded_repo() -> (Arc, aionui_db::Database) { let db = init_database_memory().await.unwrap(); + // The scoped repo authorizes through the conversations parent chain, so + // seed the owning user + conversation the acp_session row hangs off. + sqlx::query( + "INSERT INTO users (id, username, password_hash, created_at, updated_at) \ + VALUES ('user-1', 'user-1', 'hash', 0, 0)", + ) + .execute(db.pool()) + .await + .unwrap(); + sqlx::query( + "INSERT INTO conversations (id, user_id, name, type, status, created_at, updated_at, extra) \ + VALUES ('conv-1', 'user-1', 'conv-1', 'acp', 'pending', 0, 0, '{}')", + ) + .execute(db.pool()) + .await + .unwrap(); let repo: Arc = Arc::new(SqliteAcpSessionRepository::new(db.pool().clone())); repo.create(&CreateAcpSessionParams { + user_id: "user-1", conversation_id: "conv-1", agent_source: "builtin", agent_id: "claude", @@ -3598,13 +3650,18 @@ mod persist_tests { let (repo, _db) = seeded_repo().await; persist_side_effects( repo.as_ref(), + "user-1", "conv-1", &SessionEvent::BackendBound { backend_session_id: Some("bsid-abc".into()), }, ) .await; - let row = repo.get("conv-1").await.unwrap().expect("row exists"); + let row = repo + .get_for_user("user-1", "conv-1") + .await + .unwrap() + .expect("row exists"); assert_eq!( row.session_id.as_deref(), Some("bsid-abc"), @@ -3615,16 +3672,23 @@ mod persist_tests { #[tokio::test] async fn backend_bound_none_does_not_clobber_anchor() { let (repo, _db) = seeded_repo().await; - repo.update_session_id("conv-1", "bsid-existing").await.unwrap(); + repo.update_session_id_for_user("user-1", "conv-1", "bsid-existing") + .await + .unwrap(); persist_side_effects( repo.as_ref(), + "user-1", "conv-1", &SessionEvent::BackendBound { backend_session_id: None, }, ) .await; - let row = repo.get("conv-1").await.unwrap().expect("row exists"); + let row = repo + .get_for_user("user-1", "conv-1") + .await + .unwrap() + .expect("row exists"); assert_eq!( row.session_id.as_deref(), Some("bsid-existing"), @@ -3637,6 +3701,7 @@ mod persist_tests { let (repo, _db) = seeded_repo().await; persist_side_effects( repo.as_ref(), + "user-1", "conv-1", &SessionEvent::ConfigChanged { mode: Some("plan".into()), @@ -3644,7 +3709,11 @@ mod persist_tests { }, ) .await; - let state = repo.load_runtime_state("conv-1").await.unwrap().expect("runtime state"); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .expect("runtime state"); assert_eq!(state.current_mode_id.as_deref(), Some("plan")); assert_eq!(state.current_model_id.as_deref(), Some("claude-opus-4-8")); } @@ -3662,6 +3731,7 @@ mod persist_tests { let task = SessionAgentTask::new( AgentType::Acp, "conv-1".into(), + "user-1".into(), "/w".into(), backend, Some(repo.clone()), @@ -3677,7 +3747,11 @@ mod persist_tests { ); // Persisted under the effort key so build_session_instance can re-apply it. - let state = repo.load_runtime_state("conv-1").await.unwrap().expect("runtime state"); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .expect("runtime state"); let selections: std::collections::HashMap = serde_json::from_str( state .config_selections_json @@ -3748,7 +3822,14 @@ mod persist_tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn get_config_options_surfaces_reasoning_effort_for_current_model() { let backend: Arc = Arc::new(EffortCapsBackend); - let task = SessionAgentTask::new(AgentType::Acp, "conv-1".into(), "/w".into(), backend, None); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-1".into(), + "user-1".into(), + "/w".into(), + backend, + None, + ); let snapshot = task.get_config_options().await.unwrap(); let effort = snapshot .config_options @@ -3801,7 +3882,14 @@ mod persist_tests { } } let backend: Arc = Arc::new(HaikuCurrentBackend); - let task = SessionAgentTask::new(AgentType::Acp, "conv-1".into(), "/w".into(), backend, None); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-1".into(), + "user-1".into(), + "/w".into(), + backend, + None, + ); let snapshot = task.get_config_options().await.unwrap(); assert!( snapshot @@ -3820,7 +3908,14 @@ mod persist_tests { async fn set_config_option_effort_returns_observed_via_override() { let (repo, _db) = seeded_repo().await; let backend: Arc = Arc::new(EffortCapsBackend); - let task = SessionAgentTask::new(AgentType::Acp, "conv-1".into(), "/w".into(), backend, Some(repo)); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-1".into(), + "user-1".into(), + "/w".into(), + backend, + Some(repo), + ); let resp = task.set_config_option("reasoning_effort", "high").await.unwrap(); assert!( matches!(resp.confirmation, aionui_api_types::ConfigOptionConfirmation::Observed), @@ -3861,14 +3956,21 @@ mod persist_tests { #[tokio::test] async fn no_conversation_found_clears_dead_anchor() { let (repo, _db) = seeded_repo().await; - repo.update_session_id("conv-1", "dead-sid").await.unwrap(); + repo.update_session_id_for_user("user-1", "conv-1", "dead-sid") + .await + .unwrap(); persist_side_effects( repo.as_ref(), + "user-1", "conv-1", &errored_turn("No conversation found with session ID dead-sid"), ) .await; - let row = repo.get("conv-1").await.unwrap().expect("row exists"); + let row = repo + .get_for_user("user-1", "conv-1") + .await + .unwrap() + .expect("row exists"); assert_eq!( row.session_id, None, "an unrecoverable resume error must null the dead anchor so the next turn opens Fresh" @@ -3878,9 +3980,21 @@ mod persist_tests { #[tokio::test] async fn error_during_execution_clears_dead_anchor() { let (repo, _db) = seeded_repo().await; - repo.update_session_id("conv-1", "dead-sid").await.unwrap(); - persist_side_effects(repo.as_ref(), "conv-1", &errored_turn("error_during_execution")).await; - let row = repo.get("conv-1").await.unwrap().expect("row exists"); + repo.update_session_id_for_user("user-1", "conv-1", "dead-sid") + .await + .unwrap(); + persist_side_effects( + repo.as_ref(), + "user-1", + "conv-1", + &errored_turn("error_during_execution"), + ) + .await; + let row = repo + .get_for_user("user-1", "conv-1") + .await + .unwrap() + .expect("row exists"); assert_eq!( row.session_id, None, "error_during_execution is a structural resume failure" @@ -3890,15 +4004,22 @@ mod persist_tests { #[tokio::test] async fn ordinary_error_keeps_anchor() { let (repo, _db) = seeded_repo().await; - repo.update_session_id("conv-1", "live-sid").await.unwrap(); + repo.update_session_id_for_user("user-1", "conv-1", "live-sid") + .await + .unwrap(); // A normal tool/turn error is NOT a resume failure — the anchor is still good. persist_side_effects( repo.as_ref(), + "user-1", "conv-1", &errored_turn("the Bash tool exited with code 1"), ) .await; - let row = repo.get("conv-1").await.unwrap().expect("row exists"); + let row = repo + .get_for_user("user-1", "conv-1") + .await + .unwrap() + .expect("row exists"); assert_eq!( row.session_id.as_deref(), Some("live-sid"), @@ -3910,11 +4031,14 @@ mod persist_tests { async fn cancelled_turn_keeps_anchor_even_with_matching_text() { use aionui_session::{CancelReason, TurnOutcome}; let (repo, _db) = seeded_repo().await; - repo.update_session_id("conv-1", "live-sid").await.unwrap(); + repo.update_session_id_for_user("user-1", "conv-1", "live-sid") + .await + .unwrap(); // claude reports a user interrupt as is_error with cancel-noise text; the // anchor is still good, so a cancel must never trigger the self-heal. persist_side_effects( repo.as_ref(), + "user-1", "conv-1", &SessionEvent::TurnResult { is_error: true, @@ -3927,7 +4051,11 @@ mod persist_tests { }, ) .await; - let row = repo.get("conv-1").await.unwrap().expect("row exists"); + let row = repo + .get_for_user("user-1", "conv-1") + .await + .unwrap() + .expect("row exists"); assert_eq!( row.session_id.as_deref(), Some("live-sid"), @@ -4073,7 +4201,14 @@ mod pump_tests { script, gate: gate.clone(), }); - let task = SessionAgentTask::new(AgentType::Acp, "conv-1".into(), "/w".into(), backend, None); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-1".into(), + "user-1".into(), + "/w".into(), + backend, + None, + ); let mut rx = crate::agent_task::IAgentTask::subscribe(task.as_ref()); // Release only AFTER subscribing: the pump cannot emit before we listen. gate.notify_one(); @@ -4368,7 +4503,14 @@ mod pump_tests { let backend: Arc = Arc::new(ScriptBackend(vec![env(SessionEvent::BackendBound { backend_session_id: Some("sid-abc".into()), })])); - let task = SessionAgentTask::new(AgentType::Acp, "conv-1".into(), "/w".into(), backend, None); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-1".into(), + "user-1".into(), + "/w".into(), + backend, + None, + ); // Let the pump process the BackendBound so session_id is known. tokio::time::sleep(std::time::Duration::from_millis(100)).await; let mut rx = crate::agent_task::IAgentTask::subscribe(task.as_ref()); @@ -4414,6 +4556,7 @@ mod pump_tests { let task = SessionAgentTask::build( AgentType::Acp, "conv-1".into(), + "user-1".into(), "/w".into(), backend, None, @@ -4455,6 +4598,7 @@ mod pump_tests { let task = SessionAgentTask::build( AgentType::Acp, "conv-1".into(), + "user-1".into(), "/w".into(), backend, None, @@ -4492,6 +4636,7 @@ mod pump_tests { let task = SessionAgentTask::build( AgentType::Acp, "conv-1".into(), + "user-1".into(), "/w".into(), backend, None, @@ -4689,7 +4834,14 @@ mod pump_tests { }] })), })); - let task = SessionAgentTask::new(AgentType::Acp, "conv-1".into(), "/w".into(), backend, None); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-1".into(), + "user-1".into(), + "/w".into(), + backend, + None, + ); let confs = task.get_confirmations(); assert_eq!(confs.len(), 1, "the pending permission must be recovered"); assert_eq!( @@ -4712,7 +4864,14 @@ mod pump_tests { tool_name: "Bash".into(), questions: None, })); - let task = SessionAgentTask::new(AgentType::Acp, "conv-1".into(), "/w".into(), backend, None); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-1".into(), + "user-1".into(), + "/w".into(), + backend, + None, + ); let confs = task.get_confirmations(); assert_eq!(confs.len(), 1); let vals: Vec = confs[0] @@ -4785,7 +4944,14 @@ mod pump_tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn set_config_option_mode_returns_observed_via_override() { let backend: Arc = Arc::new(StaticCapsBackend); - let task = SessionAgentTask::new(AgentType::Acp, "conv-1".into(), "/w".into(), backend, None); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-1".into(), + "user-1".into(), + "/w".into(), + backend, + None, + ); let resp = task.set_config_option("mode", "plan").await.unwrap(); assert!( matches!(resp.confirmation, aionui_api_types::ConfigOptionConfirmation::Observed), @@ -4806,7 +4972,14 @@ mod pump_tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn set_config_option_model_returns_observed_via_override() { let backend: Arc = Arc::new(StaticCapsBackend); - let task = SessionAgentTask::new(AgentType::Acp, "conv-1".into(), "/w".into(), backend, None); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-1".into(), + "user-1".into(), + "/w".into(), + backend, + None, + ); let resp = task.set_config_option("model", "sonnet").await.unwrap(); assert!( matches!(resp.confirmation, aionui_api_types::ConfigOptionConfirmation::Observed), @@ -4824,7 +4997,14 @@ mod pump_tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn set_config_option_rejects_invalid_mode_and_model() { let backend: Arc = Arc::new(StaticCapsBackend); - let task = SessionAgentTask::new(AgentType::Acp, "conv-1".into(), "/w".into(), backend, None); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-1".into(), + "user-1".into(), + "/w".into(), + backend, + None, + ); let mode_err = task .set_config_option("mode", "no-such-mode") @@ -4944,7 +5124,14 @@ mod pump_tests { // `_keep` is dropped here, so the ONLY remaining Sender is the backend's field // — the reap now hinges purely on the backend being dropped. drop(_keep); - let task = SessionAgentTask::new(AgentType::Acp, "conv-1".into(), "/w".into(), backend, None); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-1".into(), + "user-1".into(), + "/w".into(), + backend, + None, + ); // Let the pump subscribe and settle into its await. tokio::time::sleep(std::time::Duration::from_millis(50)).await; @@ -4992,6 +5179,7 @@ mod pump_tests { let task = SessionAgentTask::new_with_preload( AgentType::Acp, "conv-1".into(), + "user-1".into(), "/w".into(), backend, None, @@ -5050,6 +5238,7 @@ mod pump_tests { let task = SessionAgentTask::new_with_preload( AgentType::Acp, "conv-1".into(), + "user-1".into(), "/w".into(), backend, None, @@ -5104,6 +5293,7 @@ mod pump_tests { let task = SessionAgentTask::new_with_preload( AgentType::Acp, "conv-1".into(), + "user-1".into(), "/w".into(), backend, None, diff --git a/crates/aionui-ai-agent/src/task_manager.rs b/crates/aionui-ai-agent/src/task_manager.rs index b51b2778f..73ffb9449 100644 --- a/crates/aionui-ai-agent/src/task_manager.rs +++ b/crates/aionui-ai-agent/src/task_manager.rs @@ -62,6 +62,12 @@ pub trait IWorkerTaskManager: Send + Sync { /// Number of active tasks (useful for diagnostics). fn active_count(&self) -> usize; + /// Conversation IDs for active tasks, used by callers that need to apply + /// owner-aware filtering outside the task manager. + fn active_conversation_ids(&self) -> Vec { + Vec::new() + } + /// Collect tasks eligible for idle cleanup. /// /// Returns conversation IDs of tasks that: @@ -283,6 +289,13 @@ impl IWorkerTaskManager for WorkerTaskManagerImpl { self.tasks.iter().filter(|entry| entry.value().get().is_some()).count() } + fn active_conversation_ids(&self) -> Vec { + self.tasks + .iter() + .filter_map(|entry| entry.value().get().map(|_| entry.key().clone())) + .collect() + } + fn collect_idle(&self, idle_threshold_ms: TimestampMs) -> Vec { let now = now_ms(); self.tasks @@ -336,7 +349,7 @@ impl IWorkerTaskManager for WorkerTaskManagerImpl { /// (Sentry ELECTRON-1BD). #[async_trait] impl OnConversationDelete for WorkerTaskManagerImpl { - async fn on_conversation_deleted(&self, conversation_id: &str) { + async fn on_conversation_deleted(&self, _user_id: &str, conversation_id: &str) { if let Err(e) = self.kill(conversation_id, Some(AgentKillReason::ConversationDeleted)) { warn!( conversation_id, @@ -573,6 +586,18 @@ mod tests { assert_eq!(mgr.active_count(), 1); } + #[tokio::test] + async fn active_conversation_ids_returns_initialized_tasks() { + let mgr = make_manager(); + mgr.get_or_build_task("conv-1", make_options("conv-1")).await.unwrap(); + mgr.get_or_build_task("conv-2", make_options("conv-2")).await.unwrap(); + + let mut ids = mgr.active_conversation_ids(); + ids.sort(); + + assert_eq!(ids, vec!["conv-1".to_owned(), "conv-2".to_owned()]); + } + #[tokio::test] async fn get_or_build_returns_existing() { let mgr = make_manager(); diff --git a/crates/aionui-ai-agent/tests/acp_agent_integration.rs b/crates/aionui-ai-agent/tests/acp_agent_integration.rs index 975622cd6..004b5a834 100644 --- a/crates/aionui-ai-agent/tests/acp_agent_integration.rs +++ b/crates/aionui-ai-agent/tests/acp_agent_integration.rs @@ -98,6 +98,7 @@ async fn make_mock_agent(script: &str, backend: &str) -> (Arc, let params = Arc::new( assemble_acp_params( "test-conv-1".into(), + "user-acp-test".into(), WorkspaceInfo { path: "/tmp".into(), is_custom: true, diff --git a/crates/aionui-ai-agent/tests/agent_availability_integration.rs b/crates/aionui-ai-agent/tests/agent_availability_integration.rs index e799e0d8e..639a5f217 100644 --- a/crates/aionui-ai-agent/tests/agent_availability_integration.rs +++ b/crates/aionui-ai-agent/tests/agent_availability_integration.rs @@ -8,6 +8,8 @@ use aionui_db::{ }; use aionui_realtime::EventBroadcaster; +const SYSTEM_DEFAULT_USER_ID: &str = "system_default_user"; + struct NoopBroadcaster; impl EventBroadcaster for NoopBroadcaster { @@ -281,7 +283,7 @@ async fn management_list_keeps_hydrated_installation_without_reprobing_path() { std::fs::remove_file(&command_path).unwrap(); let service = agent_service(registry, provider_repo, temp.path().to_path_buf()); - let rows = service.list_management_agents().await.unwrap(); + let rows = service.list_management_agents(SYSTEM_DEFAULT_USER_ID).await.unwrap(); let cached = rows.iter().find(|row| row.id == "agent-cached").unwrap(); assert_eq!(cached.status, AgentManagementStatus::Unchecked); @@ -321,7 +323,10 @@ async fn manual_health_check_does_not_refresh_unrelated_agents() { std::fs::remove_file(&unrelated_path).unwrap(); let service = agent_service(registry.clone(), provider_repo, temp.path().to_path_buf()); - service.health_check_agent_by_id("agent-target-missing").await.unwrap(); + service + .health_check_agent_by_id(SYSTEM_DEFAULT_USER_ID, "agent-target-missing") + .await + .unwrap(); let rows = registry.list_management_rows().await; let unrelated = rows.iter().find(|row| row.id == "agent-unrelated").unwrap(); @@ -366,7 +371,10 @@ async fn custom_enabled_toggle_does_not_refresh_unrelated_agents() { std::fs::remove_file(&unrelated_path).unwrap(); let service = agent_service(registry.clone(), provider_repo, temp.path().to_path_buf()); - service.set_agent_enabled("agent-target-toggle", false).await.unwrap(); + service + .set_agent_enabled(SYSTEM_DEFAULT_USER_ID, "agent-target-toggle", false) + .await + .unwrap(); let rows = registry.list_management_rows().await; let unrelated = rows.iter().find(|row| row.id == "agent-unrelated").unwrap(); @@ -411,7 +419,10 @@ async fn custom_delete_does_not_refresh_unrelated_agents() { std::fs::remove_file(&unrelated_path).unwrap(); let service = agent_service(registry.clone(), provider_repo, temp.path().to_path_buf()); - service.delete_custom_agent("agent-target-delete").await.unwrap(); + service + .delete_custom_agent(SYSTEM_DEFAULT_USER_ID, "agent-target-delete") + .await + .unwrap(); let rows = registry.list_management_rows().await; let unrelated = rows.iter().find(|row| row.id == "agent-unrelated").unwrap(); diff --git a/crates/aionui-ai-agent/tests/factory_provider_integration.rs b/crates/aionui-ai-agent/tests/factory_provider_integration.rs index 2fd342757..76feba1b4 100644 --- a/crates/aionui-ai-agent/tests/factory_provider_integration.rs +++ b/crates/aionui-ai-agent/tests/factory_provider_integration.rs @@ -17,6 +17,8 @@ use aionui_db::{ }; use aionui_realtime::BroadcastEventBus; +const TEST_USER_ID: &str = "system_default_user"; + fn test_encryption_key() -> [u8; 32] { [0xABu8; 32] } @@ -42,6 +44,7 @@ async fn insert_test_provider(repo: &dyn IProviderRepository, id: &str, platform let encrypted_api_key = encrypt_string("sk-test-key-12345", &key).unwrap(); repo.create(CreateProviderParams { id: Some(id), + user_id: TEST_USER_ID, platform, name: "Test Provider", base_url: "https://api.example.com/v1", @@ -101,7 +104,7 @@ fn make_aionrs_options( BuildTaskOptions::new(AgentSessionContext { conversation: ConversationContext { conversation_id: conversation_id.to_owned(), - user_id: "user-1".to_owned(), + user_id: "system_default_user".to_owned(), agent_type: AgentType::Aionrs, source: None, }, diff --git a/crates/aionui-ai-agent/tests/prompt_pipeline_integration.rs b/crates/aionui-ai-agent/tests/prompt_pipeline_integration.rs index c98f95cdf..751283dfa 100644 --- a/crates/aionui-ai-agent/tests/prompt_pipeline_integration.rs +++ b/crates/aionui-ai-agent/tests/prompt_pipeline_integration.rs @@ -56,6 +56,7 @@ async fn fixture_params( Arc::new( assemble_acp_params( "conv-pp-test".into(), + "user-pp-test".into(), WorkspaceInfo { path: "/tmp".into(), is_custom: is_custom_workspace, diff --git a/crates/aionui-ai-agent/tests/skill_manager_integration.rs b/crates/aionui-ai-agent/tests/skill_manager_integration.rs index 061b67b6e..36f0d8067 100644 --- a/crates/aionui-ai-agent/tests/skill_manager_integration.rs +++ b/crates/aionui-ai-agent/tests/skill_manager_integration.rs @@ -19,6 +19,7 @@ use aionui_ai_agent::{ AcpSkillManager, build_skills_index_text, build_system_instructions, detect_skill_load_request, prepare_first_message, prepare_first_message_with_skills_index, }; +use aionui_db::{ISkillRepository, SqliteSkillRepository, UpsertSkillParams, init_database_memory}; use aionui_extension::{BUILTIN_SKILLS_ENV_VAR, resolve_skill_paths}; use tempfile::TempDir; /// Serialize env var mutations across tests — `BUILTIN_SKILLS_ENV_VAR` is @@ -156,6 +157,107 @@ async fn get_skill_loads_custom_body_via_fs_read() { assert_eq!(skill.body.as_deref(), Some("Custom body here")); } +#[tokio::test] +async fn discover_by_names_for_user_ignores_other_users_skills() { + let _guard = ENV_MUTEX.lock().unwrap(); + let tmp = TempDir::new().unwrap(); + let data_dir = tmp.path().join("data"); + fs::create_dir_all(&data_dir).unwrap(); + + let current_dir = data_dir.join("skills").join("current-skill"); + fs::create_dir_all(¤t_dir).unwrap(); + fs::write( + current_dir.join("SKILL.md"), + "---\nname: current-skill\ndescription: Current user skill\n---\nCurrent body", + ) + .unwrap(); + + let foreign_dir = data_dir.join("skills").join("foreign-skill"); + fs::create_dir_all(&foreign_dir).unwrap(); + fs::write( + foreign_dir.join("SKILL.md"), + "---\nname: foreign-skill\ndescription: Foreign user skill\n---\nForeign body", + ) + .unwrap(); + + let builtin_src = tmp.path().join("builtin"); + let builtin_dir = builtin_src.join("global-skill"); + fs::create_dir_all(&builtin_dir).unwrap(); + fs::write( + builtin_dir.join("SKILL.md"), + "---\nname: global-skill\ndescription: Global skill\n---\nGlobal body", + ) + .unwrap(); + unsafe { + std::env::set_var(BUILTIN_SKILLS_ENV_VAR, &builtin_src); + } + + let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users \ + (id, user_type, external_user_id, username, email, password_hash, avatar_path, jwt_secret, \ + status, session_generation, created_at, updated_at, last_login) \ + VALUES ('user-b', 'aionpro', 'external-user-b', 'User B', NULL, NULL, NULL, NULL, 'active', 0, 1, 1, NULL)", + ) + .execute(db.pool()) + .await + .unwrap(); + + let repo = Arc::new(SqliteSkillRepository::new(db.pool().clone())); + repo.upsert_for_user( + "system_default_user", + UpsertSkillParams { + name: "current-skill", + description: Some("Current user skill"), + path: ¤t_dir.to_string_lossy(), + source: "user", + enabled: true, + }, + ) + .await + .unwrap(); + repo.upsert_for_user( + "user-b", + UpsertSkillParams { + name: "foreign-skill", + description: Some("Foreign user skill"), + path: &foreign_dir.to_string_lossy(), + source: "user", + enabled: true, + }, + ) + .await + .unwrap(); + repo.upsert_global(UpsertSkillParams { + name: "global-skill", + description: Some("Global skill"), + path: "global-skill", + source: "builtin", + enabled: true, + }) + .await + .unwrap(); + + let paths = Arc::new(resolve_skill_paths(tmp.path(), &data_dir)); + let mgr = AcpSkillManager::new_with_repo(paths, repo); + let requested = vec![ + "current-skill".to_owned(), + "foreign-skill".to_owned(), + "global-skill".to_owned(), + ]; + + let idx = mgr.discover_by_names_for_user("system_default_user", &requested).await; + let names: std::collections::HashSet<&str> = idx.iter().map(|s| s.name.as_str()).collect(); + + assert!(names.contains("current-skill")); + assert!(names.contains("global-skill")); + assert!(!names.contains("foreign-skill")); + + unsafe { + std::env::remove_var(BUILTIN_SKILLS_ENV_VAR); + } +} + #[tokio::test] async fn discover_skills_respects_exclude_builtin() { let _guard = ENV_MUTEX.lock().unwrap(); diff --git a/crates/aionui-api-types/src/auth.rs b/crates/aionui-api-types/src/auth.rs index 6d069ac12..4f6b31b06 100644 --- a/crates/aionui-api-types/src/auth.rs +++ b/crates/aionui-api-types/src/auth.rs @@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize}; /// Public user info returned in API responses. /// /// Contains only the fields safe to expose to clients. -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct PublicUser { pub id: String, pub username: String, @@ -86,6 +86,74 @@ pub struct WsTokenResponse { pub expires_in: u64, } +// --------------------------------------------------------------------------- +// Internal AionPro user provisioning endpoints +// --------------------------------------------------------------------------- + +/// Request body for `PUT /api/auth/internal/external-users/{external_user_id}`. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct EnsureExternalUserRequest { + pub user_type: ExternalUserType, + pub username: Option, + pub email: Option, + pub avatar_path: Option, +} + +/// Supported external identity projection sources. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum ExternalUserType { + Aionpro, +} + +/// Response for successful internal external-user provisioning. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct EnsureExternalUserResponse { + pub user_id: String, + pub user_type: ExternalUserType, + pub external_user_id: String, + pub session_generation: i64, +} + +/// Request body for `POST /api/auth/internal/external-sessions`. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct EnsureExternalSessionRequest { + pub user_type: ExternalUserType, + pub external_user_id: String, +} + +/// Response for successful internal external-session exchange. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct EnsureExternalSessionResponse { + pub user: PublicUser, + pub session_generation: i64, +} + +/// Request body for `POST /api/auth/internal/external-sessions/revoke`. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct RevokeExternalSessionRequest { + pub user_type: ExternalUserType, + pub external_user_id: String, +} + +/// Response for successful internal external-session revocation. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct RevokeExternalSessionResponse { + pub user_id: String, + pub session_generation: i64, +} + +/// Stable internal auth error codes for AionPro/Core session exchange. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum InternalAuthErrorCode { + BootstrapSecretRequired, + InvalidBootstrapSecret, + UserContextRequired, + UserDisabled, + ExternalUserConflict, +} + // --------------------------------------------------------------------------- // WebUI admin credential endpoints (local-only) // --------------------------------------------------------------------------- @@ -192,6 +260,55 @@ mod tests { assert_eq!(json["token"], "eyJhbGciOi"); } + #[test] + fn test_internal_external_user_contract_serialization() { + let req = EnsureExternalUserRequest { + user_type: ExternalUserType::Aionpro, + username: Some("Pro User".into()), + email: Some("pro@example.com".into()), + avatar_path: None, + }; + let json = serde_json::to_value(&req).unwrap(); + assert_eq!(json["user_type"], "aionpro"); + assert_eq!(json["username"], "Pro User"); + assert_eq!(json["email"], "pro@example.com"); + assert_eq!(json["avatar_path"], serde_json::Value::Null); + + let resp = EnsureExternalUserResponse { + user_id: "user_123".into(), + user_type: ExternalUserType::Aionpro, + external_user_id: "external_123".into(), + session_generation: 2, + }; + let json = serde_json::to_value(&resp).unwrap(); + assert_eq!(json["user_id"], "user_123"); + assert_eq!(json["user_type"], "aionpro"); + assert_eq!(json["external_user_id"], "external_123"); + assert_eq!(json["session_generation"], 2); + + let code = serde_json::to_value(InternalAuthErrorCode::UserContextRequired).unwrap(); + assert_eq!(code, "USER_CONTEXT_REQUIRED"); + } + + #[test] + fn test_internal_external_session_revoke_contract_serialization() { + let req = RevokeExternalSessionRequest { + user_type: ExternalUserType::Aionpro, + external_user_id: "external_123".into(), + }; + let json = serde_json::to_value(&req).unwrap(); + assert_eq!(json["user_type"], "aionpro"); + assert_eq!(json["external_user_id"], "external_123"); + + let resp = RevokeExternalSessionResponse { + user_id: "user_123".into(), + session_generation: 3, + }; + let json = serde_json::to_value(&resp).unwrap(); + assert_eq!(json["user_id"], "user_123"); + assert_eq!(json["session_generation"], 3); + } + #[test] fn test_change_password_request_snake_case() { let raw = r#"{"current_password":"old123","new_password":"new456"}"#; diff --git a/crates/aionui-api-types/src/channel.rs b/crates/aionui-api-types/src/channel.rs index df84a53fd..a403bf36f 100644 --- a/crates/aionui-api-types/src/channel.rs +++ b/crates/aionui-api-types/src/channel.rs @@ -254,6 +254,7 @@ pub struct ChannelSessionResponse { /// pairing authorization flow. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct PairingRequestedPayload { + pub user_id: String, pub code: String, pub platform_user_id: String, pub platform_type: String, @@ -267,6 +268,7 @@ pub struct PairingRequestedPayload { /// Pushed when a plugin starts, stops, or encounters an error. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct PluginStatusChangedPayload { + pub user_id: String, pub plugin_id: String, pub status: PluginStatusResponse, } @@ -276,6 +278,7 @@ pub struct PluginStatusChangedPayload { /// Pushed after a pairing code is approved and the user record is created. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct UserAuthorizedPayload { + pub user_id: String, pub id: String, pub platform_user_id: String, pub platform_type: String, @@ -684,6 +687,7 @@ mod tests { #[test] fn test_pairing_requested_payload_serde() { let payload = PairingRequestedPayload { + user_id: "user-1".into(), code: "123456".into(), platform_user_id: "tg_42".into(), platform_type: "telegram".into(), @@ -691,6 +695,7 @@ mod tests { expires_at: 1700000600000, }; let json = serde_json::to_value(&payload).unwrap(); + assert_eq!(json["user_id"], "user-1"); assert_eq!(json["code"], "123456"); assert_eq!(json["platform_user_id"], "tg_42"); assert_eq!(json["platform_type"], "telegram"); @@ -701,6 +706,7 @@ mod tests { #[test] fn test_pairing_requested_payload_no_display_name() { let payload = PairingRequestedPayload { + user_id: "user-1".into(), code: "000001".into(), platform_user_id: "u1".into(), platform_type: "dingtalk".into(), @@ -714,6 +720,7 @@ mod tests { #[test] fn test_plugin_status_changed_payload_serde() { let payload = PluginStatusChangedPayload { + user_id: "user-1".into(), plugin_id: "telegram".into(), status: PluginStatusResponse { plugin_id: "telegram".into(), @@ -731,6 +738,7 @@ mod tests { }, }; let json = serde_json::to_value(&payload).unwrap(); + assert_eq!(json["user_id"], "user-1"); assert_eq!(json["plugin_id"], "telegram"); assert_eq!(json["status"]["type"], "telegram"); assert_eq!(json["status"]["status"], "running"); @@ -740,12 +748,14 @@ mod tests { #[test] fn test_user_authorized_payload_serde() { let payload = UserAuthorizedPayload { + user_id: "user-1".into(), id: "usr_1".into(), platform_user_id: "tg_42".into(), platform_type: "telegram".into(), display_name: Some("Alice".into()), }; let json = serde_json::to_value(&payload).unwrap(); + assert_eq!(json["user_id"], "user-1"); assert_eq!(json["id"], "usr_1"); assert_eq!(json["platform_user_id"], "tg_42"); assert_eq!(json["platform_type"], "telegram"); @@ -755,6 +765,7 @@ mod tests { #[test] fn test_user_authorized_payload_no_display_name() { let payload = UserAuthorizedPayload { + user_id: "user-1".into(), id: "usr_2".into(), platform_user_id: "lk_1".into(), platform_type: "lark".into(), diff --git a/crates/aionui-api-types/src/lib.rs b/crates/aionui-api-types/src/lib.rs index 3175f1279..356ec79b6 100644 --- a/crates/aionui-api-types/src/lib.rs +++ b/crates/aionui-api-types/src/lib.rs @@ -61,8 +61,10 @@ pub use assistant::{ is_local_avatar_value, }; pub use auth::{ - AuthStatusResponse, ChangePasswordRequest, LoginRequest, LoginResponse, PublicUser, QrLoginRequest, - RefreshResponse, RefreshTokenRequest, UserInfoResponse, WebuiChangePasswordRequest, WebuiChangeUsernameRequest, + AuthStatusResponse, ChangePasswordRequest, EnsureExternalSessionRequest, EnsureExternalSessionResponse, + EnsureExternalUserRequest, EnsureExternalUserResponse, ExternalUserType, InternalAuthErrorCode, LoginRequest, + LoginResponse, PublicUser, QrLoginRequest, RefreshResponse, RefreshTokenRequest, RevokeExternalSessionRequest, + RevokeExternalSessionResponse, UserInfoResponse, WebuiChangePasswordRequest, WebuiChangeUsernameRequest, WebuiChangeUsernameResponse, WebuiGenerateQrTokenResponse, WebuiResetPasswordResponse, WsTokenResponse, }; pub use channel::{ diff --git a/crates/aionui-api-types/src/runtime.rs b/crates/aionui-api-types/src/runtime.rs index b1dddc7a9..d04e78483 100644 --- a/crates/aionui-api-types/src/runtime.rs +++ b/crates/aionui-api-types/src/runtime.rs @@ -47,6 +47,8 @@ pub enum RuntimeFailureKind { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct RuntimeStatusPayload { + #[serde(skip_serializing_if = "Option::is_none")] + pub user_id: Option, pub resource: RuntimeResourceKind, #[serde(skip_serializing_if = "Option::is_none")] pub resource_id: Option, @@ -77,6 +79,7 @@ mod tests { #[test] fn runtime_status_payload_serializes() { let payload = RuntimeStatusPayload { + user_id: Some("user-1".into()), resource: RuntimeResourceKind::Node, resource_id: None, scope: RuntimeStatusScope { @@ -91,6 +94,7 @@ mod tests { let json = serde_json::to_value(&payload).unwrap(); assert_eq!(json["resource"], "node"); + assert_eq!(json["user_id"], "user-1"); assert_eq!(json["scope"]["kind"], "conversation"); assert_eq!(json["phase"], "downloading"); assert_eq!(json["message"], "downloading"); diff --git a/crates/aionui-app/src/bootstrap/environment.rs b/crates/aionui-app/src/bootstrap/environment.rs index c9817e3e7..720705a79 100644 --- a/crates/aionui-app/src/bootstrap/environment.rs +++ b/crates/aionui-app/src/bootstrap/environment.rs @@ -4,7 +4,7 @@ use std::time::Instant; use tracing::info; -use aionui_app::AppConfig; +use aionui_app::{AppConfig, IdentityMode}; use aionui_db::Database; use crate::cli::Cli; @@ -42,20 +42,36 @@ pub fn init_environment(cli: &Cli, merged_path: &str) -> Result Result, +) -> Result<(), BootstrapError> { + if identity_mode == IdentityMode::AionPro && bootstrap_secret.is_none() { + return Err(BootstrapError::new( + BootstrapErrorCode::ConfigInvalid, + "config.identity_mode", + "AionPro identity mode requires AIONCORE_BOOTSTRAP_SECRET", + )); + } + + Ok(()) +} + /// Layer 2: Materialize builtin skills + initialize the database. /// /// Requires only `data_dir`. Subcommands that need persistent state @@ -120,6 +151,8 @@ pub async fn init_data_layer(config: &AppConfig) -> Result, @@ -71,6 +75,23 @@ pub(crate) enum ManagedResourcesModeArg { Download, } +#[derive(ValueEnum, Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum IdentityModeArg { + Local, + Webui, + Aionpro, +} + +impl From for aionui_app::IdentityMode { + fn from(value: IdentityModeArg) -> Self { + match value { + IdentityModeArg::Local => Self::Local, + IdentityModeArg::Webui => Self::WebUi, + IdentityModeArg::Aionpro => Self::AionPro, + } + } +} + impl From for aionui_runtime::ManagedResourcesMode { fn from(value: ManagedResourcesModeArg) -> Self { match value { diff --git a/crates/aionui-app/src/commands/cmd_config.rs b/crates/aionui-app/src/commands/cmd_config.rs index 4adb64e6c..325a80dce 100644 --- a/crates/aionui-app/src/commands/cmd_config.rs +++ b/crates/aionui-app/src/commands/cmd_config.rs @@ -1318,7 +1318,7 @@ async fn resolve_top_level_selectors( object.insert("conversation_id".into(), Value::String(env.conversation_id.clone())); selectors.insert("conversation_id", env.conversation_id.clone()); } - if is_current_selector(object.get("user_id")) { + if object.contains_key("user_id") { object.insert("user_id".into(), Value::String(env.user_id.clone())); selectors.insert("user_id", env.user_id.clone()); } @@ -1799,4 +1799,26 @@ mod tests { fn path_segments_are_percent_encoded() { assert_eq!(encode_path_segment("a/b c"), "a%2Fb%20c"); } + + #[tokio::test] + async fn top_level_user_id_is_always_bound_to_config_env() { + let client = reqwest::Client::new(); + let env = ConfigEnv { + base_url: "http://127.0.0.1".into(), + conversation_id: "conv-current".into(), + user_id: "user-current".into(), + }; + let mut payload = json!({ + "user_id": "user-other", + "name": "Task" + }); + let mut selectors = SelectorMeta::default(); + + resolve_top_level_selectors(&client, &env, "config cron jobs update", &mut payload, &mut selectors) + .await + .unwrap(); + + assert_eq!(payload["user_id"], "user-current"); + assert_eq!(selectors.into_map()["resolved_selectors"]["user_id"], "user-current"); + } } diff --git a/crates/aionui-app/src/commands/config_capabilities.rs b/crates/aionui-app/src/commands/config_capabilities.rs index aa250383f..30b13bf12 100644 --- a/crates/aionui-app/src/commands/config_capabilities.rs +++ b/crates/aionui-app/src/commands/config_capabilities.rs @@ -23,7 +23,7 @@ pub(crate) fn data() -> Value { }, "user_id": { "current": "resolve from AIONUI_USER_ID", - "literal": "treat as user id" + "literal": "ignored and replaced with AIONUI_USER_ID" } } }, diff --git a/crates/aionui-app/src/config.rs b/crates/aionui-app/src/config.rs index e99c2e2a9..1798778f4 100644 --- a/crates/aionui-app/src/config.rs +++ b/crates/aionui-app/src/config.rs @@ -4,6 +4,27 @@ use std::path::PathBuf; use sha2::{Digest, Sha256}; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum IdentityMode { + Local, + WebUi, + AionPro, +} + +impl IdentityMode { + pub fn auth_label(self) -> &'static str { + match self { + Self::Local => "local", + Self::WebUi => "webui", + Self::AionPro => "aionpro", + } + } + + pub fn is_local(self) -> bool { + self == Self::Local + } +} + /// Application configuration parsed from CLI arguments. #[derive(Debug, Clone)] pub struct AppConfig { @@ -14,6 +35,8 @@ pub struct AppConfig { pub app_version: String, /// Run in local embedded mode (skip authentication, use system_default_user). pub local: bool, + pub identity_mode: IdentityMode, + pub bootstrap_secret: Option, /// Dump prompt diagnostics under `data_dir/prompt-dumps`. pub dump_prompts: bool, /// Explicitly authorize backup and rebuild for corruption-like local databases. @@ -21,6 +44,14 @@ pub struct AppConfig { } impl AppConfig { + pub fn effective_identity_mode(&self) -> IdentityMode { + if self.local { + IdentityMode::Local + } else { + self.identity_mode + } + } + /// Format as `host:port` for socket binding. pub fn socket_addr(&self) -> String { format!("{}:{}", self.host, self.port) @@ -50,6 +81,8 @@ impl Default for AppConfig { work_dir: PathBuf::from("data"), app_version: env!("CARGO_PKG_VERSION").to_string(), local: false, + identity_mode: IdentityMode::WebUi, + bootstrap_secret: None, dump_prompts: false, recover_corrupted_database: false, } @@ -75,6 +108,8 @@ mod tests { assert_eq!(config.port, 25808); assert_eq!(config.data_dir, PathBuf::from("data")); assert_eq!(config.app_version, env!("CARGO_PKG_VERSION")); + assert_eq!(config.identity_mode, IdentityMode::WebUi); + assert!(config.bootstrap_secret.is_none()); assert!(!config.dump_prompts); assert!(!config.recover_corrupted_database); } diff --git a/crates/aionui-app/src/lib.rs b/crates/aionui-app/src/lib.rs index e09e9349f..533b0a1b0 100644 --- a/crates/aionui-app/src/lib.rs +++ b/crates/aionui-app/src/lib.rs @@ -9,7 +9,7 @@ mod config; mod router; mod services; -pub use config::{AppConfig, derive_encryption_key}; +pub use config::{AppConfig, IdentityMode, derive_encryption_key}; pub use router::{ ChannelOrchestratorComponents, ModuleStates, RouterBuildError, RouterRuntime, build_assistant_state, build_conversation_state, build_extension_states, build_module_states, build_ws_state, create_router, diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index ed88d8b7d..74bd08832 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -6,19 +6,20 @@ use std::time::Instant; use axum::Json; use axum::extract::DefaultBodyLimit; use axum::extract::Request; -use axum::http::{Method, StatusCode, header}; +use axum::http::{HeaderName, Method, StatusCode, header}; use axum::middleware::{Next, from_fn_with_state}; use axum::response::{IntoResponse, Response}; use axum::routing::get; use axum::{Router, middleware}; -use tower_http::cors::{Any, CorsLayer}; +use tower_http::cors::{AllowOrigin, Any, CorsLayer}; use aionui_ai_agent::{agent_routes, remote_agent_routes}; use aionui_api_types::ErrorResponse; use aionui_assets::{AssetRouterState, asset_routes}; use aionui_assistant::assistant_routes; use aionui_auth::{ - AuthRouterState, AuthState, auth_middleware, auth_routes, csrf_middleware, security_headers_middleware, + AuthIdentityMode, AuthRouterState, AuthState, auth_middleware, auth_routes, csrf_middleware, + security_headers_middleware, }; use aionui_channel::channel_routes; #[cfg(feature = "weixin")] @@ -70,7 +71,21 @@ pub async fn create_router_with_runtime(services: &AppServices) -> Result<(Route let ws_manager = services.ws_manager.clone(); tokio::spawn(async move { while let Ok(event) = event_rx.recv().await { - ws_manager.broadcast_all(event); + if let Some(user_id) = event + .data + .get("user_id") + .and_then(|value| value.as_str()) + .map(str::to_owned) + { + ws_manager.broadcast_to_user(&user_id, event); + } else if is_global_websocket_event(&event.name) { + ws_manager.broadcast_all(event); + } else { + tracing::warn!( + event_name = %event.name, + "dropping websocket event without user_id; add user_id to payload or whitelist explicit global event" + ); + } } }); @@ -93,13 +108,22 @@ pub async fn create_router_with_runtime(services: &AppServices) -> Result<(Route // Restore enabled channel plugins (starts receiving IM messages) let chan_mgr = channel_components.manager; let chan_factory = channel_components.plugin_factory; + let chan_owner_user_id = channel_components.owner_user_id; tokio::spawn(async move { - if let Err(e) = chan_mgr.restore_plugins(&chan_factory).await { - tracing::warn!( - code = "BOOTSTRAP_DEGRADED_CHANNEL_RESTORE", + if let Some(chan_owner_user_id) = chan_owner_user_id { + if let Err(e) = chan_mgr.restore_plugins(&chan_owner_user_id, &chan_factory).await { + tracing::warn!( + code = "BOOTSTRAP_DEGRADED_CHANNEL_RESTORE", + stage = "channel.restore", + owner_user_id = %chan_owner_user_id, + error = %e, + "failed to restore channel plugins" + ); + } + } else { + tracing::info!( stage = "channel.restore", - error = %e, - "failed to restore channel plugins" + "skipping channel plugin restore until an owner user is available" ); } }); @@ -148,13 +172,74 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates user_repo: services.user_repo.clone(), cookie_config: services.cookie_config.clone(), qr_token_store: services.qr_token_store.clone(), + identity_mode: auth_identity_mode(services.identity_mode), + bootstrap_secret: services.bootstrap_secret.clone(), + session_revoked_hook: { + let ws_manager = services.ws_manager.clone(); + let conversation_service = states.conversation.service.clone(); + let team_service = states.team.service.clone(); + let channel_manager = states.channel.manager.clone(); + let channel_session_manager = states.channel.session_manager.clone(); + let file_watch_service = states.file.watch_service.clone(); + let office_watch_manager = states.office.watch_manager.clone(); + Some(Arc::new(move |user_id: &str| { + ws_manager.disconnect_user(user_id, "session revoked"); + let stopped_team_sessions = team_service.stop_sessions_for_user(user_id); + if stopped_team_sessions > 0 { + tracing::info!( + user_id = %user_id, + stopped_team_sessions, + "stopped team sessions after session revocation" + ); + } + let user_id = user_id.to_owned(); + let conversation_service = conversation_service.clone(); + let channel_manager = channel_manager.clone(); + let channel_session_manager = channel_session_manager.clone(); + let file_watch_service = file_watch_service.clone(); + let office_watch_manager = office_watch_manager.clone(); + tokio::spawn(async move { + channel_manager.shutdown_for_user(&user_id).await; + office_watch_manager.stop_all_for_user(&user_id); + if let Err(err) = channel_session_manager.clear_all_sessions(&user_id).await { + tracing::warn!( + user_id = %user_id, + error = %err, + "failed to clear channel sessions after session revocation" + ); + } + if let Err(err) = conversation_service.terminate_runtime_for_user(&user_id).await { + tracing::warn!( + user_id = %user_id, + error = %err, + "failed to terminate runtimes after session revocation" + ); + } + if let Err(err) = file_watch_service.stop_all_watches_for_user(&user_id).await { + tracing::warn!( + user_id = %user_id, + error = %err, + "failed to stop file watches after session revocation" + ); + } + if let Err(err) = file_watch_service.stop_all_office_watches_for_user(&user_id).await { + tracing::warn!( + user_id = %user_id, + error = %err, + "failed to stop office file watches after session revocation" + ); + } + }); + })) + }, local: services.local, + aionpro_mode: services.identity_mode == crate::config::IdentityMode::AionPro, }; let auth_mw_state = AuthState { jwt_service: services.jwt_service.clone(), user_repo: services.user_repo.clone(), - local: services.local, + identity_mode: auth_identity_mode(services.identity_mode), }; // System routes protected by auth middleware @@ -228,8 +313,10 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates let assistant_authenticated = assistant_routes(states.assistant).route_layer(from_fn_with_state(auth_mw_state.clone(), auth_middleware)); - // Office proxy routes — exempt from auth (serve iframe content) - let office_proxy = office_proxy_routes(states.office); + // Office proxy routes serve iframe content but still require auth so + // preview ports remain scoped to the active Core user. + let office_proxy = + office_proxy_routes(states.office).route_layer(from_fn_with_state(auth_mw_state.clone(), auth_middleware)); let public_assets = asset_routes(AssetRouterState::default()); // WebSocket upgrade route — exempt from CSRF (no cookie-based @@ -266,7 +353,7 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates #[cfg(feature = "weixin")] let router = router.merge(weixin_login_authenticated); - let router = if services.local { + let router = if services.identity_mode.is_local() { router } else { router.layer(middleware::from_fn_with_state( @@ -292,7 +379,7 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates "startup: route tree build with states completed" ); - if services.local { + if services.identity_mode.is_local() { let cors = CorsLayer::new() .allow_origin(Any) .allow_methods([ @@ -306,10 +393,47 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates .allow_headers(Any); router.layer(cors) } else { - router + // Non-local (external identity) mode: the desktop renderer is a + // cross-origin browser context (localhost:5173 in dev, file:// when + // packaged) authenticating with the session cookie, so responses must + // opt in to credentialed CORS. Credentialed mode forbids wildcards: + // reflect the request origin and enumerate headers explicitly + // (x-csrf-token is required by the CSRF double-submit middleware). + let cors = CorsLayer::new() + .allow_origin(AllowOrigin::mirror_request()) + .allow_credentials(true) + .allow_methods([ + Method::GET, + Method::POST, + Method::PUT, + Method::PATCH, + Method::DELETE, + Method::OPTIONS, + ]) + .allow_headers([ + header::CONTENT_TYPE, + header::AUTHORIZATION, + HeaderName::from_static("x-csrf-token"), + ]); + router.layer(cors) + } +} + +fn auth_identity_mode(identity_mode: crate::config::IdentityMode) -> AuthIdentityMode { + match identity_mode { + crate::config::IdentityMode::Local => AuthIdentityMode::Local, + crate::config::IdentityMode::WebUi => AuthIdentityMode::UserSession, + crate::config::IdentityMode::AionPro => AuthIdentityMode::AionPro, } } +fn is_global_websocket_event(event_name: &str) -> bool { + matches!( + event_name, + "runtime.statusChanged" | "extensions.lifecycle" | "hub.state-changed" + ) +} + async fn normalize_boundary_error_response(request: Request, next: Next) -> Response { let response = next.run(request).await; if response.status().is_success() || response_has_json_content_type(&response) { @@ -367,7 +491,7 @@ fn boundary_error_for_status(status: StatusCode) -> Option<(&'static str, &'stat mod tests { use axum::http::StatusCode; - use super::{boundary_error_for_status, create_router_with_runtime}; + use super::{boundary_error_for_status, create_router_with_runtime, is_global_websocket_event}; use crate::config::AppConfig; use crate::services::AppServices; @@ -396,6 +520,12 @@ mod tests { } } + #[test] + fn extension_enablement_events_are_user_scoped() { + assert!(!is_global_websocket_event("extensions.state-changed")); + assert!(is_global_websocket_event("extensions.lifecycle")); + } + #[tokio::test] async fn create_router_with_runtime_exposes_team_service_for_background_coordinators() { let db = aionui_db::init_database_memory().await.unwrap(); diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index a2461348b..9ee09e7f8 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -36,7 +36,7 @@ use aionui_mcp::{ use aionui_office::{ ConversionService, OfficeRouterState, OfficecliWatchManager, ProxyService, SnapshotService as OfficeSnapshotService, }; -use aionui_realtime::{NoopMessageRouter, WsHandlerState}; +use aionui_realtime::{NoopMessageRouter, TokenUserResolver, WsHandlerState}; use aionui_shell::ShellRouterState; use aionui_system::{ ClientPrefService, ConnectionTestRouterState, ConnectionTestService, FeedbackDiagnosticsService, ModelFetchService, @@ -48,7 +48,7 @@ use aionui_team::{ TeamConversationProvisioningPort, TeamProjectionMessageStore, TeamRouterState, TeamSessionService, }; -use crate::config::derive_encryption_key; +use crate::config::{IdentityMode, derive_encryption_key}; use crate::router::team_conversation_adapters::TeamConversationAdapters; use crate::services::AppServices; @@ -185,6 +185,7 @@ pub struct ChannelOrchestratorComponents { pub confirm_rx: tokio::sync::mpsc::Receiver<(String, String)>, pub manager: Arc, pub plugin_factory: Arc, + pub owner_user_id: Option, } /// Build all default `ModuleStates` from application services. @@ -328,8 +329,14 @@ pub fn build_assistant_state(services: &AppServices) -> AssistantRouterState { #[async_trait::async_trait] impl AssistantAgentCatalogPort for RegistryAssistantAgentCatalog { - async fn list_management_agents(&self) -> Result, AssistantError> { - Ok(self.registry.list_management_rows().await) + async fn list_management_agents( + &self, + user_id: &str, + ) -> Result, AssistantError> { + self.registry + .list_management_rows_for_user(user_id) + .await + .map_err(|error| AssistantError::Internal(format!("agent catalog unavailable: {error}"))) } } @@ -379,12 +386,17 @@ pub fn build_system_state(services: &AppServices) -> SystemRouterState { let provider_repo = Arc::new(SqliteProviderRepository::new(pool.clone())); let http_client = reqwest::Client::new(); + let client_pref_repo = Arc::new(SqliteClientPreferenceRepository::new(pool.clone())); + let keep_awake_controller = Arc::new(aionui_system::SystemKeepAwakeController::new()); + let client_pref_service = if services.identity_mode.is_local() { + ClientPrefService::with_keep_awake_controller(client_pref_repo, keep_awake_controller, "system_default_user") + } else { + ClientPrefService::with_keep_awake_controller_without_restore(client_pref_repo, keep_awake_controller) + }; + SystemRouterState { settings_service: SettingsService::new(Arc::new(SqliteSettingsRepository::new(pool.clone()))), - client_pref_service: ClientPrefService::with_keep_awake_controller( - Arc::new(SqliteClientPreferenceRepository::new(pool.clone())), - Arc::new(aionui_system::SystemKeepAwakeController::new()), - ), + client_pref_service, provider_service: ProviderService::new(provider_repo.clone(), encryption_key), model_fetch_service: ModelFetchService::new(provider_repo, encryption_key, http_client.clone()), protocol_detection_service: ProtocolDetectionService::new(http_client.clone()), @@ -507,6 +519,18 @@ async fn build_channel_message_service( services: &AppServices, channel_settings: Arc, ) -> Arc { + Arc::new(aionui_channel::message_service::ChannelMessageService::new( + Arc::new(services.conversation_service.clone()), + services.worker_task_manager.clone(), + channel_settings, + )) +} + +async fn startup_channel_owner_user_id(services: &AppServices) -> Option { + if services.identity_mode == IdentityMode::AionPro { + return None; + } + let owner_user_id = services .user_repo .get_primary_webui_user() @@ -515,13 +539,7 @@ async fn build_channel_message_service( .flatten() .map(|u| u.id) .unwrap_or_else(|| "system_default_user".to_string()); - - Arc::new(aionui_channel::message_service::ChannelMessageService::new( - Arc::new(services.conversation_service.clone()), - services.worker_task_manager.clone(), - channel_settings, - owner_user_id, - )) + Some(owner_user_id) } /// Build the default `ChannelRouterState` and orchestrator components. @@ -556,12 +574,14 @@ pub async fn build_channel_state( // Build channel settings service for per-plugin agent/model configuration. let channel_settings = build_channel_settings_service(services); + let startup_owner_user_id = startup_channel_owner_user_id(services).await; // Build orchestrator dependencies let action_executor = Arc::new(aionui_channel::action::ActionExecutor::new( Arc::clone(&pairing_service), Arc::clone(&session_manager), Arc::clone(&channel_settings), + startup_owner_user_id.clone(), )); let message_service = build_channel_message_service(services, Arc::clone(&channel_settings)).await; @@ -589,6 +609,7 @@ pub async fn build_channel_state( confirm_rx, manager, plugin_factory, + owner_user_id: startup_owner_user_id, }; (state, components) @@ -614,8 +635,9 @@ pub fn build_team_state( impl TeamAssistantCatalogPort for AssistantServiceTeamCatalog { async fn list_team_selectable_assistants( &self, + user_id: &str, ) -> Result, aionui_team::TeamError> { - let assistants = self.assistant_service.list().await.map_err(|error| { + let assistants = self.assistant_service.list_for_user(user_id).await.map_err(|error| { aionui_team::TeamError::InvalidRequest(format!("assistant catalog unavailable: {error}")) })?; @@ -751,6 +773,7 @@ pub fn build_cron_state(services: &AppServices) -> CronRouterState { agent_metadata_repo, assistant_definition_repo, assistant_overlay_repo, + skill_repo: services.skill_repo.clone(), scheduler, executor, emitter, @@ -849,12 +872,28 @@ pub fn build_ws_state(services: &AppServices) -> WsHandlerState { manager: services.ws_manager.clone(), router: Arc::new(NoopMessageRouter), token_validator: Arc::new(|_| true), + token_user_resolver: Arc::new(|_| Box::pin(async { Some("system_default_user".to_owned()) })), token_extractor: Arc::new(|_| Some("local".into())), }; } let jwt_service = services.jwt_service.clone(); let token_validator = Arc::new(move |token: &str| jwt_service.verify(token).is_ok()); + let jwt_service = services.jwt_service.clone(); + let user_repo = services.user_repo.clone(); + let identity_mode = services.identity_mode; + let token_user_resolver: TokenUserResolver = Arc::new(move |token: String| { + let jwt_service = jwt_service.clone(); + let user_repo = user_repo.clone(); + Box::pin(async move { + let payload = jwt_service.verify(&token).ok()?; + let user = user_repo.find_active_by_id(&payload.user_id).await.ok()??; + if identity_mode == IdentityMode::AionPro && user.user_type != aionui_db::UserType::Aionpro { + return None; + } + (payload.session_generation == user.session_generation).then_some(user.id) + }) + }); let token_extractor = Arc::new(|headers: &axum::http::HeaderMap| extract_token_from_ws_headers(headers)); @@ -862,6 +901,7 @@ pub fn build_ws_state(services: &AppServices) -> WsHandlerState { manager: services.ws_manager.clone(), router: Arc::new(NoopMessageRouter), token_validator, + token_user_resolver, token_extractor, } } @@ -948,6 +988,8 @@ mod tests { impl IMockAgent for ChannelStateNoopAgent {} + type CapturedRuntimeEnv = Arc>>>; + fn mock_worker_task_manager() -> Arc { let factory = Arc::new(|opts: BuildTaskOptions| { Box::pin(async move { @@ -961,9 +1003,7 @@ mod tests { Arc::new(WorkerTaskManagerImpl::new(factory)) } - fn capturing_worker_task_manager( - captured_env: Arc>>>, - ) -> Arc { + fn capturing_worker_task_manager(captured_env: CapturedRuntimeEnv) -> Arc { let factory = Arc::new(move |opts: BuildTaskOptions| { let captured_env = captured_env.clone(); Box::pin(async move { @@ -980,7 +1020,7 @@ mod tests { Arc::new(WorkerTaskManagerImpl::new(factory)) } - async fn wait_for_captured_env(captured_env: &Arc>>>) -> Vec<(String, String)> { + async fn wait_for_captured_env(captured_env: &CapturedRuntimeEnv) -> Vec<(String, String)> { for _ in 0..50 { if let Some(env) = captured_env.lock().unwrap().first().cloned() { return env; @@ -1030,6 +1070,52 @@ mod tests { } } + #[tokio::test] + async fn build_ws_state_rejects_stale_session_generation() { + let db = aionui_db::init_database_memory().await.unwrap(); + let services = AppServices::from_config(db, &AppConfig::default()).await.unwrap(); + let ws_state = build_ws_state(&services); + let token = services + .jwt_service + .sign_with_session_generation("system_default_user", "admin", 0) + .unwrap(); + + let resolved = (ws_state.token_user_resolver)(token.clone()).await; + assert_eq!(resolved.as_deref(), Some("system_default_user")); + + services + .user_repo + .increment_session_generation("system_default_user") + .await + .unwrap(); + + let resolved = (ws_state.token_user_resolver)(token).await; + assert_eq!(resolved, None); + + services.database.close().await; + } + + #[tokio::test] + async fn build_ws_state_aionpro_rejects_local_user_token() { + let db = aionui_db::init_database_memory().await.unwrap(); + let config = AppConfig { + identity_mode: IdentityMode::AionPro, + bootstrap_secret: Some("bootstrap-secret".to_owned()), + ..AppConfig::default() + }; + let services = AppServices::from_config(db, &config).await.unwrap(); + let ws_state = build_ws_state(&services); + let token = services + .jwt_service + .sign_with_session_generation("system_default_user", "admin", 0) + .unwrap(); + + let resolved = (ws_state.token_user_resolver)(token).await; + assert_eq!(resolved, None); + + services.database.close().await; + } + #[tokio::test] async fn build_channel_message_service_uses_app_conversation_service_for_assistant_bindings() { let db = aionui_db::init_database_memory().await.unwrap(); @@ -1047,10 +1133,13 @@ mod tests { let pref_repo = SqliteClientPreferenceRepository::new(pool.clone()); pref_repo - .upsert_batch(&[( - "assistant.weixin.agent", - r#"{"assistant_id":"bare-channel-aionrs","name":"Weixin Aionrs"}"#, - )]) + .upsert_batch( + "system_default_user", + &[( + "assistant.weixin.agent", + r#"{"assistant_id":"bare-channel-aionrs","name":"Weixin Aionrs"}"#, + )], + ) .await .unwrap(); @@ -1068,18 +1157,23 @@ mod tests { }; let first = message_service - .send_to_agent(&session, "hello", PluginType::Weixin) + .send_to_agent("system_default_user", &session, "hello", PluginType::Weixin) .await .unwrap(); let conversation_repo = SqliteConversationRepository::new(pool); + let user_id = conversation_repo + .owner_user_id(&first.conversation_id) + .await + .unwrap() + .expect("channel-created conversation should have an owner"); let snapshot = conversation_repo - .get_assistant_snapshot(&first.conversation_id) + .get_assistant_snapshot(&user_id, &first.conversation_id) .await .unwrap() .expect("channel-created conversation should persist assistant snapshot"); let conversation = conversation_repo - .get(&first.conversation_id) + .get(&user_id, &first.conversation_id) .await .unwrap() .expect("channel-created conversation should be persisted"); @@ -1094,7 +1188,7 @@ mod tests { ..session }; let second = message_service - .send_to_agent(&second_session, "again", PluginType::Weixin) + .send_to_agent("system_default_user", &second_session, "again", PluginType::Weixin) .await .unwrap(); assert_eq!(second.conversation_id, first.conversation_id); diff --git a/crates/aionui-app/src/router/team_conversation_adapters.rs b/crates/aionui-app/src/router/team_conversation_adapters.rs index 4e6ea5662..4a5317b75 100644 --- a/crates/aionui-app/src/router/team_conversation_adapters.rs +++ b/crates/aionui-app/src/router/team_conversation_adapters.rs @@ -35,6 +35,16 @@ impl TeamConversationAdapters { task_manager, } } + + async fn owner_user_id(&self, conversation_id: &str) -> Result, TeamError> { + Ok(self.conversation_repo.owner_user_id(conversation_id).await?) + } + + async fn require_owner_user_id(&self, conversation_id: &str) -> Result { + self.owner_user_id(conversation_id) + .await? + .ok_or_else(|| TeamError::InvalidRequest(format!("conversation {conversation_id} not found"))) + } } #[async_trait] @@ -152,13 +162,18 @@ impl TeamProjectionMessageStore for TeamConversationAdapters { ) -> Result, TeamError> { Ok(self .conversation_repo - .get_message_by_msg_id(conversation_id, msg_id, msg_type) + .get_message_by_msg_id( + &self.require_owner_user_id(conversation_id).await?, + conversation_id, + msg_id, + msg_type, + ) .await?) } async fn insert_projected_message(&self, row: &MessageRow) -> Result<(), TeamError> { self.conversation_service - .insert_raw_message(row) + .insert_raw_message(&self.require_owner_user_id(&row.conversation_id).await?, row) .await .map_err(map_conversation_update_error) } @@ -204,7 +219,10 @@ impl TeamConversationProvisioningPort for TeamConversationAdapters { } async fn conversation_workspace(&self, conversation_id: &str) -> Result, TeamError> { - let Some(row) = self.conversation_repo.get(conversation_id).await? else { + let Some(user_id) = self.owner_user_id(conversation_id).await? else { + return Ok(None); + }; + let Some(row) = self.conversation_repo.get(&user_id, conversation_id).await? else { return Ok(None); }; let extra: serde_json::Value = serde_json::from_str(&row.extra).unwrap_or(serde_json::Value::Null); @@ -217,14 +235,21 @@ impl TeamConversationProvisioningPort for TeamConversationAdapters { } async fn conversation_assistant_id(&self, conversation_id: &str) -> Result, TeamError> { - if let Some(snapshot) = self.conversation_repo.get_assistant_snapshot(conversation_id).await? { + let Some(user_id) = self.owner_user_id(conversation_id).await? else { + return Ok(None); + }; + if let Some(snapshot) = self + .conversation_repo + .get_assistant_snapshot(&user_id, conversation_id) + .await? + { let assistant_id = snapshot.assistant_id.trim(); if !assistant_id.is_empty() { return Ok(Some(assistant_id.to_owned())); } } - let Some(row) = self.conversation_repo.get(conversation_id).await? else { + let Some(row) = self.conversation_repo.get(&user_id, conversation_id).await? else { return Ok(None); }; @@ -246,21 +271,27 @@ impl TeamConversationProvisioningPort for TeamConversationAdapters { async fn patch_runtime_config(&self, conversation_id: &str, patch: serde_json::Value) -> Result<(), TeamError> { self.conversation_service - .update_extra(conversation_id, patch) + .update_extra( + &self.require_owner_user_id(conversation_id).await?, + conversation_id, + patch, + ) .await .map_err(map_conversation_update_error) } async fn save_acp_runtime_mode(&self, conversation_id: &str, mode: &str) -> Result<(), TeamError> { + let user_id = self.require_owner_user_id(conversation_id).await?; self.conversation_service - .save_acp_runtime_mode(conversation_id, mode) + .save_acp_runtime_mode(&user_id, conversation_id, mode) .await .map_err(map_conversation_update_error) } async fn get_config_options(&self, conversation_id: &str) -> Result { + let user_id = self.require_owner_user_id(conversation_id).await?; self.conversation_service - .get_config_options(conversation_id) + .get_config_options(&user_id, conversation_id) .await .map_err(map_conversation_update_error) } @@ -288,7 +319,10 @@ impl TeamConversationProvisioningPort for TeamConversationAdapters { &self, conversation_id: &str, ) -> Result, TeamError> { - let Some(row) = self.conversation_repo.get(conversation_id).await? else { + let Some(user_id) = self.owner_user_id(conversation_id).await? else { + return Ok(None); + }; + let Some(row) = self.conversation_repo.get(&user_id, conversation_id).await? else { return Ok(None); }; let extra: serde_json::Value = serde_json::from_str(&row.extra).unwrap_or(serde_json::Value::Null); diff --git a/crates/aionui-app/src/services.rs b/crates/aionui-app/src/services.rs index ac4daf65e..7cfd7a7a3 100644 --- a/crates/aionui-app/src/services.rs +++ b/crates/aionui-app/src/services.rs @@ -3,7 +3,7 @@ use std::path::PathBuf; use std::sync::Arc; -use crate::config::{AppConfig, derive_encryption_key}; +use crate::config::{AppConfig, IdentityMode, derive_encryption_key}; use aionui_ai_agent::{ AcpSessionSyncService, AcpSkillManager, ActiveLeaseRegistry, AgentFactoryDeps, AgentRegistry, IWorkerTaskManager, RuntimeTokenService, WorkerTaskManagerImpl, build_agent_factory, @@ -48,6 +48,8 @@ pub struct AppServices { pub work_dir: PathBuf, /// When `true`, skip JWT authentication and use a fixed default user. pub local: bool, + pub identity_mode: IdentityMode, + pub bootstrap_secret: Option>, pub app_version: String, /// Resolved skill paths. Shared with the `ConversationService` for /// snapshot resolution at create time. @@ -92,7 +94,8 @@ impl AppServices { pub async fn from_config(database: Database, config: &AppConfig) -> anyhow::Result { let data_dir = config.data_dir.clone(); let work_dir = config.work_dir.clone(); - let local = config.local; + let identity_mode = config.effective_identity_mode(); + let local = identity_mode.is_local(); let dump_prompts = config.dump_prompts; let app_version = config.app_version.clone(); let user_repo: Arc = Arc::new(SqliteUserRepository::new(database.pool().clone())); @@ -242,6 +245,8 @@ impl AppServices { dump_prompts, work_dir, local, + identity_mode, + bootstrap_secret: config.bootstrap_secret.clone().map(Arc::::from), app_version, skill_paths, skill_repo, diff --git a/crates/aionui-app/tests/acp_e2e.rs b/crates/aionui-app/tests/acp_e2e.rs index dc0f1ec94..07b0b72dc 100644 --- a/crates/aionui-app/tests/acp_e2e.rs +++ b/crates/aionui-app/tests/acp_e2e.rs @@ -16,6 +16,16 @@ use aionui_db::{ use common::{body_json, build_app, get_with_token, json_with_token, setup_and_login}; +async fn user_id_for_username(services: &aionui_app::AppServices, username: &str) -> String { + services + .user_repo + .find_by_username(username) + .await + .unwrap() + .expect("test user should exist") + .id +} + // ── Global ACP routes ──────────────────────────────────────────── #[tokio::test] @@ -71,35 +81,39 @@ async fn test_custom_agent_nonexistent_command() { async fn management_list_includes_missing_custom_agents() { let (mut app, services) = build_app().await; let (token, _csrf) = setup_and_login(&mut app, &services, "user1", "pass123").await; + let user_id = user_id_for_username(&services, "user1").await; let repo: std::sync::Arc = std::sync::Arc::new(SqliteAgentMetadataRepository::new(services.database.pool().clone())); - repo.upsert(&UpsertAgentMetadataParams { - id: "custom-missing-agent", - icon: None, - name: "Missing Custom Agent", - name_i18n: None, - description: None, - description_i18n: None, - backend: Some("claude"), - agent_type: "acp", - agent_source: "custom", - agent_source_info: Some(r#"{"binary_name":"aionui-missing-agent-binary"}"#), - enabled: true, - command: Some("aionui-missing-agent-binary"), - args: Some("[]"), - env: Some("[]"), - native_skills_dirs: None, - behavior_policy: None, - yolo_id: None, - agent_capabilities: None, - auth_methods: None, - config_options: None, - available_modes: None, - available_models: None, - available_commands: None, - sort_order: 1500, - }) + repo.upsert_for_user( + &user_id, + &UpsertAgentMetadataParams { + id: "custom-missing-agent", + icon: None, + name: "Missing Custom Agent", + name_i18n: None, + description: None, + description_i18n: None, + backend: Some("claude"), + agent_type: "acp", + agent_source: "custom", + agent_source_info: Some(r#"{"binary_name":"aionui-missing-agent-binary"}"#), + enabled: true, + command: Some("aionui-missing-agent-binary"), + args: Some("[]"), + env: Some("[]"), + native_skills_dirs: None, + behavior_policy: None, + yolo_id: None, + agent_capabilities: None, + auth_methods: None, + config_options: None, + available_modes: None, + available_models: None, + available_commands: None, + sort_order: 1500, + }, + ) .await .unwrap(); services.agent_registry.hydrate().await.unwrap(); @@ -122,38 +136,43 @@ async fn management_list_includes_missing_custom_agents() { async fn management_list_marks_rows_with_unavailable_snapshot() { let (mut app, services) = build_app().await; let (token, _csrf) = setup_and_login(&mut app, &services, "user1", "pass123").await; + let user_id = user_id_for_username(&services, "user1").await; let repo: std::sync::Arc = std::sync::Arc::new(SqliteAgentMetadataRepository::new(services.database.pool().clone())); - repo.upsert(&UpsertAgentMetadataParams { - id: "custom-unavailable-agent", - icon: None, - name: "Unavailable Custom Agent", - name_i18n: None, - description: None, - description_i18n: None, - backend: Some("claude"), - agent_type: "acp", - agent_source: "custom", - agent_source_info: Some(r#"{"binary_name":"cargo"}"#), - enabled: true, - command: Some("cargo"), - args: Some("[]"), - env: Some("[]"), - native_skills_dirs: None, - behavior_policy: None, - yolo_id: None, - agent_capabilities: None, - auth_methods: None, - config_options: None, - available_modes: None, - available_models: None, - available_commands: None, - sort_order: 1500, - }) + repo.upsert_for_user( + &user_id, + &UpsertAgentMetadataParams { + id: "custom-unavailable-agent", + icon: None, + name: "Unavailable Custom Agent", + name_i18n: None, + description: None, + description_i18n: None, + backend: Some("claude"), + agent_type: "acp", + agent_source: "custom", + agent_source_info: Some(r#"{"binary_name":"cargo"}"#), + enabled: true, + command: Some("cargo"), + args: Some("[]"), + env: Some("[]"), + native_skills_dirs: None, + behavior_policy: None, + yolo_id: None, + agent_capabilities: None, + auth_methods: None, + config_options: None, + available_modes: None, + available_models: None, + available_commands: None, + sort_order: 1500, + }, + ) .await .unwrap(); - repo.update_availability_snapshot( + repo.update_availability_snapshot_for_user( + &user_id, "custom-unavailable-agent", &UpdateAgentAvailabilitySnapshotParams { last_check_status: Some("offline"), @@ -198,35 +217,39 @@ async fn legacy_agents_endpoint_is_not_found() { async fn health_check_by_id_returns_missing_status_for_uninstalled_agent() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "user1", "pass123").await; + let user_id = user_id_for_username(&services, "user1").await; let repo: std::sync::Arc = std::sync::Arc::new(SqliteAgentMetadataRepository::new(services.database.pool().clone())); - repo.upsert(&UpsertAgentMetadataParams { - id: "custom-missing-agent", - icon: None, - name: "Missing Custom Agent", - name_i18n: None, - description: None, - description_i18n: None, - backend: Some("claude"), - agent_type: "acp", - agent_source: "custom", - agent_source_info: Some(r#"{"binary_name":"aionui-missing-agent-binary"}"#), - enabled: true, - command: Some("aionui-missing-agent-binary"), - args: Some("[]"), - env: Some("[]"), - native_skills_dirs: None, - behavior_policy: None, - yolo_id: None, - agent_capabilities: None, - auth_methods: None, - config_options: None, - available_modes: None, - available_models: None, - available_commands: None, - sort_order: 1500, - }) + repo.upsert_for_user( + &user_id, + &UpsertAgentMetadataParams { + id: "custom-missing-agent", + icon: None, + name: "Missing Custom Agent", + name_i18n: None, + description: None, + description_i18n: None, + backend: Some("claude"), + agent_type: "acp", + agent_source: "custom", + agent_source_info: Some(r#"{"binary_name":"aionui-missing-agent-binary"}"#), + enabled: true, + command: Some("aionui-missing-agent-binary"), + args: Some("[]"), + env: Some("[]"), + native_skills_dirs: None, + behavior_policy: None, + yolo_id: None, + agent_capabilities: None, + auth_methods: None, + config_options: None, + available_modes: None, + available_models: None, + available_commands: None, + sort_order: 1500, + }, + ) .await .unwrap(); services.agent_registry.hydrate().await.unwrap(); diff --git a/crates/aionui-app/tests/active_lease_e2e.rs b/crates/aionui-app/tests/active_lease_e2e.rs index d789a7ab3..e96d62cc2 100644 --- a/crates/aionui-app/tests/active_lease_e2e.rs +++ b/crates/aionui-app/tests/active_lease_e2e.rs @@ -267,8 +267,8 @@ async fn team_active_lease_rejects_cross_user_without_renewing() { let req = empty_post_with_token("/api/teams/team-cross-user/active-lease", &other_token, &other_csrf); let resp = app.oneshot(req).await.unwrap(); - assert_eq!(resp.status(), StatusCode::FORBIDDEN); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); let body = body_json(resp).await; - assert_eq!(body["code"], "FORBIDDEN"); + assert_eq!(body["code"], "NOT_FOUND"); assert!(services.active_lease_registry.active_until("team-cross-conv").is_none()); } diff --git a/crates/aionui-app/tests/assistants_e2e.rs b/crates/aionui-app/tests/assistants_e2e.rs index ecdd20c47..3d19ac18c 100644 --- a/crates/aionui-app/tests/assistants_e2e.rs +++ b/crates/aionui-app/tests/assistants_e2e.rs @@ -37,6 +37,8 @@ use tower::ServiceExt; use common::{body_json, delete_with_token, get_with_token, json_with_token, setup_and_login}; +const DEFAULT_USER_ID: &str = "system_default_user"; + // --------------------------------------------------------------------------- // Fixture — router + temp dirs + services // --------------------------------------------------------------------------- @@ -65,7 +67,10 @@ struct TestAgentCatalog { #[async_trait::async_trait] impl AssistantAgentCatalogPort for TestAgentCatalog { - async fn list_management_agents(&self) -> Result, aionui_assistant::AssistantError> { + async fn list_management_agents( + &self, + _user_id: &str, + ) -> Result, aionui_assistant::AssistantError> { Ok(self.rows.clone()) } } @@ -289,6 +294,7 @@ async fn fixture() -> Fixture { // erroring out — mirroring a configured production setup. provider_repo .create(aionui_db::CreateProviderParams { + user_id: "system_default_user", id: None, platform: "openai", name: "Test OpenAI", @@ -939,6 +945,44 @@ async fn update_extension_registry_id_without_user_row_returns_404() { assert_eq!(resp.status(), StatusCode::NOT_FOUND); } +#[tokio::test] +async fn assistant_routes_hide_other_users_rows() { + let fx = fixture().await; + create_user(&fx, "private-a", "Private A").await; + + let mut app = fx.app.clone(); + let (token_b, csrf_b) = setup_and_login(&mut app, &fx.services, "other-user", "OtherP@ss1").await; + + let requests = [ + get_with_token("/api/assistants/private-a", &token_b), + json_with_token( + "PUT", + "/api/assistants/private-a", + json!({ "name": "hijacked" }), + &token_b, + &csrf_b, + ), + delete_with_token("/api/assistants/private-a", &token_b, &csrf_b), + ]; + + for request in requests { + let resp = fx.app.clone().oneshot(request).await.unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + let body = body_json(resp).await; + assert_eq!(body["success"], false); + assert_eq!(body["code"], "NOT_FOUND"); + } + + let resp = fx + .app + .clone() + .oneshot(get_with_token("/api/assistants/private-a", &fx.token)) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(body_json(resp).await["data"]["profile"]["name"], "Private A"); +} + // =========================================================================== // DELETE /api/assistants/{id} // =========================================================================== @@ -949,8 +993,8 @@ async fn delete_happy_path_removes_row_and_user_assets() { create_user(&fx, "u1", "A").await; // Drop a rule, skill, and avatar on disk so the fs-cleanup branch has // something to remove. - let rules_dir = fx.user_data_dir.join("assistant-rules"); - let skills_dir = fx.user_data_dir.join("assistant-skills"); + let rules_dir = fx.user_data_dir.join("assistant-rules").join(DEFAULT_USER_ID); + let skills_dir = fx.user_data_dir.join("assistant-skills").join(DEFAULT_USER_ID); let avatars_dir = fx.user_data_dir.join("assistant-avatars"); std::fs::create_dir_all(&rules_dir).unwrap(); std::fs::create_dir_all(&skills_dir).unwrap(); @@ -1318,6 +1362,92 @@ async fn read_rule_user_round_trip_through_write() { assert_eq!(json["data"], "my rule"); } +#[tokio::test] +async fn rule_and_skill_files_are_isolated_by_current_user() { + let fx = fixture().await; + let mut app = fx.app.clone(); + let (token_b, csrf_b) = setup_and_login(&mut app, &fx.services, "other-user", "OtherP@ss1").await; + let user_b = fx + .services + .user_repo + .find_by_username("other-user") + .await + .unwrap() + .unwrap(); + + for (token, csrf, endpoint, content) in [ + ( + fx.token.as_str(), + fx.csrf.as_str(), + "/api/skills/assistant-rule/write", + "rule-a", + ), + ( + token_b.as_str(), + csrf_b.as_str(), + "/api/skills/assistant-rule/write", + "rule-b", + ), + ( + fx.token.as_str(), + fx.csrf.as_str(), + "/api/skills/assistant-skill/write", + "skill-a", + ), + ( + token_b.as_str(), + csrf_b.as_str(), + "/api/skills/assistant-skill/write", + "skill-b", + ), + ] { + let req = json_with_token( + "POST", + endpoint, + json!({ "assistant_id": "shared-assistant", "content": content }), + token, + csrf, + ); + let resp = fx.app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + } + + for (token, csrf, expected) in [ + (fx.token.as_str(), fx.csrf.as_str(), "rule-a"), + (token_b.as_str(), csrf_b.as_str(), "rule-b"), + ] { + let req = json_with_token( + "POST", + "/api/skills/assistant-rule/read", + json!({ "assistant_id": "shared-assistant" }), + token, + csrf, + ); + let resp = fx.app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(body_json(resp).await["data"], expected); + } + + let rule_root = fx.user_data_dir.join("assistant-rules"); + let skill_root = fx.user_data_dir.join("assistant-skills"); + assert_eq!( + std::fs::read_to_string(rule_root.join(DEFAULT_USER_ID).join("shared-assistant.md")).unwrap(), + "rule-a" + ); + assert_eq!( + std::fs::read_to_string(rule_root.join(&user_b.id).join("shared-assistant.md")).unwrap(), + "rule-b" + ); + assert_eq!( + std::fs::read_to_string(skill_root.join(DEFAULT_USER_ID).join("shared-assistant.md")).unwrap(), + "skill-a" + ); + assert_eq!( + std::fs::read_to_string(skill_root.join(&user_b.id).join("shared-assistant.md")).unwrap(), + "skill-b" + ); +} + // =========================================================================== // POST /api/skills/assistant-rule/write // =========================================================================== @@ -1336,7 +1466,11 @@ async fn write_rule_user_happy_path() { let resp = fx.app.clone().oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); // File was actually written. - let file = fx.user_data_dir.join("assistant-rules/u1.md"); + let file = fx + .user_data_dir + .join("assistant-rules") + .join(DEFAULT_USER_ID) + .join("u1.md"); assert_eq!(std::fs::read_to_string(file).unwrap(), "rule body"); } @@ -1376,7 +1510,7 @@ async fn write_rule_extension_registry_id_behaves_like_user_id() { async fn delete_rule_user_removes_file() { let fx = fixture().await; create_user(&fx, "u1", "A").await; - let rules_dir = fx.user_data_dir.join("assistant-rules"); + let rules_dir = fx.user_data_dir.join("assistant-rules").join(DEFAULT_USER_ID); std::fs::create_dir_all(&rules_dir).unwrap(); std::fs::write(rules_dir.join("u1.md"), "body").unwrap(); @@ -1502,7 +1636,11 @@ async fn write_skill_user_happy_path() { ); let resp = fx.app.clone().oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); - let file = fx.user_data_dir.join("assistant-skills/u1.md"); + let file = fx + .user_data_dir + .join("assistant-skills") + .join(DEFAULT_USER_ID) + .join("u1.md"); assert_eq!(std::fs::read_to_string(file).unwrap(), "skill body"); } @@ -1542,7 +1680,7 @@ async fn write_skill_extension_registry_id_behaves_like_user_id() { async fn delete_skill_user_removes_file() { let fx = fixture().await; create_user(&fx, "u1", "A").await; - let skills_dir = fx.user_data_dir.join("assistant-skills"); + let skills_dir = fx.user_data_dir.join("assistant-skills").join(DEFAULT_USER_ID); std::fs::create_dir_all(&skills_dir).unwrap(); std::fs::write(skills_dir.join("u1.md"), "body").unwrap(); diff --git a/crates/aionui-app/tests/channel_e2e.rs b/crates/aionui-app/tests/channel_e2e.rs index f157ab5a0..efd730c32 100644 --- a/crates/aionui-app/tests/channel_e2e.rs +++ b/crates/aionui-app/tests/channel_e2e.rs @@ -14,6 +14,8 @@ use tower::ServiceExt; use common::{body_json, build_app, get_with_token, json_with_token, setup_and_login}; +const OWNER_ID: &str = "system_default_user"; + // =========================================================================== // §1 Plugin management // =========================================================================== @@ -415,15 +417,19 @@ async fn put_channel_assistant_setting_clears_active_sessions() { std::sync::Arc::new(SqliteChannelRepository::new(services.database.pool().clone())); let now = now_ms(); - repo.create_user(&AssistantUserRow { - id: "user-channel-assistant".to_owned(), - platform_user_id: "user-channel-assistant".to_owned(), - platform_type: "lark".to_owned(), - display_name: Some("Channel Assistant User".to_owned()), - authorized_at: now, - last_active: Some(now), - session_id: None, - }) + repo.create_user( + OWNER_ID, + &AssistantUserRow { + id: "user-channel-assistant".to_owned(), + owner_user_id: OWNER_ID.to_owned(), + platform_user_id: "user-channel-assistant".to_owned(), + platform_type: "lark".to_owned(), + display_name: Some("Channel Assistant User".to_owned()), + authorized_at: now, + last_active: Some(now), + session_id: None, + }, + ) .await .unwrap(); let new_session = AssistantSessionRow { @@ -436,10 +442,15 @@ async fn put_channel_assistant_setting_clears_active_sessions() { created_at: now, last_activity: now, }; - repo.get_or_create_session("user-channel-assistant", "chat-channel-assistant", &new_session) - .await - .unwrap(); - assert_eq!(repo.get_all_sessions().await.unwrap().len(), 1); + repo.get_or_create_session( + OWNER_ID, + "user-channel-assistant", + "chat-channel-assistant", + &new_session, + ) + .await + .unwrap(); + assert_eq!(repo.get_all_sessions(OWNER_ID).await.unwrap().len(), 1); let req = json_with_token( "PUT", @@ -453,7 +464,7 @@ async fn put_channel_assistant_setting_clears_active_sessions() { let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); - assert!(repo.get_all_sessions().await.unwrap().is_empty()); + assert!(repo.get_all_sessions(OWNER_ID).await.unwrap().is_empty()); } #[tokio::test] @@ -464,15 +475,19 @@ async fn put_channel_default_model_setting_clears_active_sessions() { std::sync::Arc::new(SqliteChannelRepository::new(services.database.pool().clone())); let now = now_ms(); - repo.create_user(&AssistantUserRow { - id: "user-channel-model".to_owned(), - platform_user_id: "user-channel-model".to_owned(), - platform_type: "lark".to_owned(), - display_name: Some("Channel Model User".to_owned()), - authorized_at: now, - last_active: Some(now), - session_id: None, - }) + repo.create_user( + OWNER_ID, + &AssistantUserRow { + id: "user-channel-model".to_owned(), + owner_user_id: OWNER_ID.to_owned(), + platform_user_id: "user-channel-model".to_owned(), + platform_type: "lark".to_owned(), + display_name: Some("Channel Model User".to_owned()), + authorized_at: now, + last_active: Some(now), + session_id: None, + }, + ) .await .unwrap(); let new_session = AssistantSessionRow { @@ -485,10 +500,10 @@ async fn put_channel_default_model_setting_clears_active_sessions() { created_at: now, last_activity: now, }; - repo.get_or_create_session("user-channel-model", "chat-channel-model", &new_session) + repo.get_or_create_session(OWNER_ID, "user-channel-model", "chat-channel-model", &new_session) .await .unwrap(); - assert_eq!(repo.get_all_sessions().await.unwrap().len(), 1); + assert_eq!(repo.get_all_sessions(OWNER_ID).await.unwrap().len(), 1); let req = json_with_token( "PUT", @@ -503,7 +518,7 @@ async fn put_channel_default_model_setting_clears_active_sessions() { let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); - assert!(repo.get_all_sessions().await.unwrap().is_empty()); + assert!(repo.get_all_sessions(OWNER_ID).await.unwrap().is_empty()); } // SS-1: Sync valid platform clears sessions @@ -573,7 +588,7 @@ async fn pairing_approve_creates_user() { let pairing_svc = aionui_channel::pairing::PairingService::new(repo.clone(), services.event_bus.clone()); let code = pairing_svc - .request_pairing("tg_user_42", "telegram", Some("Alice")) + .request_pairing(OWNER_ID, "tg_user_42", "telegram", Some("Alice")) .await .unwrap(); @@ -664,7 +679,7 @@ async fn pairing_reject_removes_from_pending() { let pairing_svc = aionui_channel::pairing::PairingService::new(repo.clone(), services.event_bus.clone()); let code = pairing_svc - .request_pairing("tg_user_99", "telegram", None) + .request_pairing(OWNER_ID, "tg_user_99", "telegram", None) .await .unwrap(); diff --git a/crates/aionui-app/tests/common/mod.rs b/crates/aionui-app/tests/common/mod.rs index fa69e2ab4..7c407e5ec 100644 --- a/crates/aionui-app/tests/common/mod.rs +++ b/crates/aionui-app/tests/common/mod.rs @@ -26,7 +26,8 @@ pub async fn build_app() -> (axum::Router, AppServices) { /// (E1 `/api/skills`, E3/E4 built-in reads, E5 `/api/skills/info`). /// Returns the router, services, and the `SkillPaths` so the test can seed /// fixtures at known locations. `/api/skills` reads the database only; tests -/// that seed skill directories should call `sync_skill_catalog_for_test`. +/// that seed skill directories should call `sync_skill_catalog_for_test` for +/// the logged-in user. #[allow(dead_code)] pub async fn build_app_with_skill_paths(root: &std::path::Path) -> (axum::Router, AppServices, SkillPaths) { let db = aionui_db::init_database_memory().await.unwrap(); @@ -78,9 +79,15 @@ pub async fn build_app_with_skill_paths(root: &std::path::Path) -> (axum::Router (router, services, paths) } -pub async fn sync_skill_catalog_for_test(services: &AppServices, paths: &SkillPaths) { +pub async fn sync_skill_catalog_for_test(services: &AppServices, paths: &SkillPaths, username: &str) { let repo = aionui_db::SqliteSkillRepository::new(services.database.pool().clone()); - aionui_extension::sync_skill_catalog_into_repo(paths, &repo) + let user = services + .user_repo + .find_by_username(username) + .await + .unwrap() + .unwrap_or_else(|| panic!("test user '{username}' should exist")); + aionui_extension::sync_skill_catalog_into_repo_for_user(paths, &repo, &user.id) .await .expect("sync skill catalog"); } diff --git a/crates/aionui-app/tests/conversation_e2e.rs b/crates/aionui-app/tests/conversation_e2e.rs index 69e643fcf..edb981178 100644 --- a/crates/aionui-app/tests/conversation_e2e.rs +++ b/crates/aionui-app/tests/conversation_e2e.rs @@ -257,8 +257,13 @@ async fn t1_3b_create_persists_available_locale_fallback_rule_in_assistant_snaps .any(|skill| skill == "override-skill") ); + let user_id = conversation_repo + .owner_user_id(data["id"].as_str().unwrap()) + .await + .unwrap() + .unwrap(); let snapshot = conversation_repo - .get_assistant_snapshot(data["id"].as_str().unwrap()) + .get_assistant_snapshot(&user_id, data["id"].as_str().unwrap()) .await .unwrap() .unwrap(); @@ -884,7 +889,8 @@ async fn t7_1_reset_conversation() { hidden: false, created_at: 1000, }; - aionui_db::IConversationRepository::insert_message(&repo, &msg) + let user_id = repo.owner_user_id(&id).await.unwrap().unwrap(); + aionui_db::IConversationRepository::insert_message(&repo, &user_id, &msg) .await .unwrap(); @@ -975,7 +981,8 @@ async fn team_owned_conversation_rejects_ordinary_send_but_allows_history_reads( hidden: false, created_at: 1000, }; - aionui_db::IConversationRepository::insert_message(&repo, &msg) + let user_id = repo.owner_user_id(&id).await.unwrap().unwrap(); + aionui_db::IConversationRepository::insert_message(&repo, &user_id, &msg) .await .unwrap(); diff --git a/crates/aionui-app/tests/cron_e2e.rs b/crates/aionui-app/tests/cron_e2e.rs index 62f152766..e69606b65 100644 --- a/crates/aionui-app/tests/cron_e2e.rs +++ b/crates/aionui-app/tests/cron_e2e.rs @@ -11,6 +11,7 @@ use axum::http::StatusCode; use serde_json::json; use tower::ServiceExt; +use aionui_db::models::ConversationRow; use aionui_db::{ CreateMcpServerParams, IConversationRepository, ICronRepository, IMcpServerRepository, SqliteConversationRepository, SqliteCronRepository, SqliteMcpServerRepository, @@ -86,8 +87,51 @@ async fn ensure_default_assistant(app: &mut axum::Router, token: &str, csrf: &st ); } -async fn create_job(app: &mut axum::Router, token: &str, csrf: &str, body: serde_json::Value) -> serde_json::Value { +async fn ensure_conversation(services: &aionui_app::AppServices, user_id: &str, conversation_id: &str, title: &str) { + if conversation_id.trim().is_empty() { + return; + } + let repo = SqliteConversationRepository::new(services.database.pool().clone()); + if repo.get(user_id, conversation_id).await.unwrap().is_some() { + return; + } + let now = aionui_common::now_ms(); + repo.create(&ConversationRow { + id: conversation_id.to_owned(), + user_id: user_id.to_owned(), + name: title.to_owned(), + r#type: "acp".to_owned(), + extra: "{}".to_owned(), + model: None, + status: Some("finished".to_owned()), + source: Some("aionui".to_owned()), + channel_chat_id: None, + pinned: false, + pinned_at: None, + created_at: now, + updated_at: now, + }) + .await + .unwrap(); +} + +async fn create_job( + app: &mut axum::Router, + services: &aionui_app::AppServices, + token: &str, + csrf: &str, + body: serde_json::Value, +) -> serde_json::Value { ensure_default_assistant(app, token, csrf).await; + let admin = services + .user_repo + .find_by_username("admin") + .await + .unwrap() + .expect("admin user should exist"); + let conversation_id = body["conversation_id"].as_str().unwrap_or_default(); + let title = body["conversation_title"].as_str().unwrap_or("Cron Test Conversation"); + ensure_conversation(services, &admin.id, conversation_id, title).await; let req = json_with_token("POST", "/api/cron/jobs", body, token, csrf); let resp = app.clone().oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::CREATED); @@ -169,7 +213,7 @@ async fn cj1_create_cron_job() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let data = create_job(&mut app, &token, &csrf, create_job_body("Daily Report")).await; + let data = create_job(&mut app, &services, &token, &csrf, create_job_body("Daily Report")).await; assert!(data["id"].as_str().unwrap().starts_with("cron_")); assert_eq!(data["name"], "Daily Report"); @@ -188,7 +232,7 @@ async fn cj1b_create_job_allows_missing_task_description() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let data = create_job(&mut app, &token, &csrf, create_job_body("No Description")).await; + let data = create_job(&mut app, &services, &token, &csrf, create_job_body("No Description")).await; assert!(data.get("description").is_none()); assert_eq!(data["schedule"]["description"], "every minute"); @@ -203,17 +247,25 @@ async fn cj2_create_three_schedule_types() { let now = aionui_common::now_ms(); - let at = create_job(&mut app, &token, &csrf, create_at_job_body("At Job", now + 3_600_000)).await; + let at = create_job( + &mut app, + &services, + &token, + &csrf, + create_at_job_body("At Job", now + 3_600_000), + ) + .await; assert_eq!(at["schedule"]["kind"], "at"); assert!(at["state"]["next_run_at_ms"].as_i64().unwrap() > now); - let every = create_job(&mut app, &token, &csrf, create_job_body("Every Job")).await; + let every = create_job(&mut app, &services, &token, &csrf, create_job_body("Every Job")).await; assert_eq!(every["schedule"]["kind"], "every"); let next = every["state"]["next_run_at_ms"].as_i64().unwrap(); assert!((next - now - 60000).abs() < 3000); let cron = create_job( &mut app, + &services, &token, &csrf, create_cron_job_body("Cron Job", "0 */5 * * * *"), @@ -333,7 +385,7 @@ async fn cj4_get_single_job() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let created = create_job(&mut app, &token, &csrf, create_job_body("Get Test")).await; + let created = create_job(&mut app, &services, &token, &csrf, create_job_body("Get Test")).await; let job_id = created["id"].as_str().unwrap(); let req = get_with_token(&format!("/api/cron/jobs/{job_id}"), &token); @@ -361,6 +413,7 @@ async fn cj5_get_nonexistent() { async fn cj5b_run_now_legacy_workspace_with_whitespace_succeeds() { let (mut app, services) = build_app_with_mock_agents().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; + let owner = services.user_repo.find_by_username("admin").await.unwrap().unwrap(); ensure_default_assistant(&mut app, &token, &csrf).await; let cron_repo = SqliteCronRepository::new(services.database.pool().clone()); let now = aionui_common::now_ms(); @@ -372,6 +425,7 @@ async fn cj5b_run_now_legacy_workspace_with_whitespace_succeeds() { cron_repo .insert(&aionui_db::models::CronJobRow { id: "cron_whitespace_workspace".into(), + user_id: owner.id, name: "Legacy Workspace".into(), enabled: true, schedule_kind: "every".into(), @@ -430,7 +484,7 @@ async fn cj6_list_all_jobs() { let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; for i in 0..3 { - create_job(&mut app, &token, &csrf, create_job_body(&format!("Job {i}"))).await; + create_job(&mut app, &services, &token, &csrf, create_job_body(&format!("Job {i}"))).await; } let req = get_with_token("/api/cron/jobs", &token); @@ -451,15 +505,15 @@ async fn cj7_list_by_conversation() { let mut body_a = create_job_body("Job A"); body_a["conversation_id"] = json!("conv_target"); - create_job(&mut app, &token, &csrf, body_a).await; + create_job(&mut app, &services, &token, &csrf, body_a).await; let mut body_b = create_job_body("Job B"); body_b["conversation_id"] = json!("conv_target"); - create_job(&mut app, &token, &csrf, body_b).await; + create_job(&mut app, &services, &token, &csrf, body_b).await; let mut body_c = create_job_body("Job C"); body_c["conversation_id"] = json!("conv_other"); - create_job(&mut app, &token, &csrf, body_c).await; + create_job(&mut app, &services, &token, &csrf, body_c).await; let req = get_with_token("/api/cron/jobs?conversation_id=conv_target", &token); let resp = app.oneshot(req).await.unwrap(); @@ -477,7 +531,7 @@ async fn cj8_update_job() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let created = create_job(&mut app, &token, &csrf, create_job_body("Original")).await; + let created = create_job(&mut app, &services, &token, &csrf, create_job_body("Original")).await; let job_id = created["id"].as_str().unwrap(); let update_body = json!({"name": "Updated Name", "enabled": false}); @@ -500,7 +554,7 @@ async fn cj9_update_schedule_type() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let created = create_job(&mut app, &token, &csrf, create_job_body("Schedule Change")).await; + let created = create_job(&mut app, &services, &token, &csrf, create_job_body("Schedule Change")).await; let job_id = created["id"].as_str().unwrap(); let update_body = json!({"schedule": {"kind": "cron", "expr": "0 */5 * * * *"}}); @@ -520,6 +574,7 @@ async fn cj9b_update_schedule_preserves_existing_timezone_when_omitted() { let created = create_job( &mut app, + &services, &token, &csrf, json!({ @@ -565,7 +620,7 @@ async fn cj11_delete_job() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let created = create_job(&mut app, &token, &csrf, create_job_body("To Delete")).await; + let created = create_job(&mut app, &services, &token, &csrf, create_job_body("To Delete")).await; let job_id = created["id"].as_str().unwrap(); let req = delete_with_token(&format!("/api/cron/jobs/{job_id}"), &token, &csrf); @@ -614,7 +669,7 @@ async fn rn1_run_now_returns_conversation_id_for_new_conversation_job() { let mut body = create_job_body("Run Now Job"); body["conversation_id"] = json!(conversation_id); - let created = create_job(&mut app, &token, &csrf, body).await; + let created = create_job(&mut app, &services, &token, &csrf, body).await; let job_id = created["id"].as_str().unwrap(); let req = json_with_token( @@ -654,7 +709,7 @@ async fn rn1b_run_now_returns_active_conversation_when_conversation_is_busy() { let mut body = create_job_body("Busy Run Now Job"); body["conversation_id"] = json!(conversation_id); - let created = create_job(&mut app, &token, &csrf, body).await; + let created = create_job(&mut app, &services, &token, &csrf, body).await; let job_id = created["id"].as_str().unwrap(); let claim = services @@ -686,6 +741,7 @@ async fn rn1c_run_now_new_conversation_preset_assistant_uses_fixed_assistant_mcp let mcp_repo = SqliteMcpServerRepository::new(services.database.pool().clone()); let fixed_mcp = mcp_repo .create(CreateMcpServerParams { + user_id: "system_default_user", name: "fixed-mcp", description: None, enabled: true, @@ -699,6 +755,7 @@ async fn rn1c_run_now_new_conversation_preset_assistant_uses_fixed_assistant_mcp .expect("create fixed mcp"); let extra_mcp = mcp_repo .create(CreateMcpServerParams { + user_id: "system_default_user", name: "extra-mcp", description: None, enabled: true, @@ -784,8 +841,13 @@ async fn rn1c_run_now_new_conversation_preset_assistant_uses_fixed_assistant_mcp .expect("run-now should return created conversation id"); let conversation_repo = SqliteConversationRepository::new(services.database.pool().clone()); + let user_id = conversation_repo + .owner_user_id(conversation_id) + .await + .expect("load conversation owner") + .expect("conversation should have an owner"); let conversation = conversation_repo - .get(conversation_id) + .get(&user_id, conversation_id) .await .expect("load conversation") .expect("conversation should exist"); @@ -805,7 +867,7 @@ async fn rn1c_run_now_new_conversation_preset_assistant_uses_fixed_assistant_mcp assert_ne!(fixed_mcp.id, extra_mcp.id, "fixture should seed two distinct MCP rows"); let snapshot = conversation_repo - .get_assistant_snapshot(conversation_id) + .get_assistant_snapshot(&user_id, conversation_id) .await .expect("load assistant snapshot") .expect("preset assistant cron conversation should persist snapshot"); @@ -830,7 +892,7 @@ async fn sk1_save_skill() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let created = create_job(&mut app, &token, &csrf, create_job_body("Skill Job")).await; + let created = create_job(&mut app, &services, &token, &csrf, create_job_body("Skill Job")).await; let job_id = created["id"].as_str().unwrap(); let skill_body = json!({"content": "---\nname: test\ndescription: test skill\n---\nDo something"}); @@ -852,7 +914,7 @@ async fn sk2_has_skill_true() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let created = create_job(&mut app, &token, &csrf, create_job_body("Skill Check")).await; + let created = create_job(&mut app, &services, &token, &csrf, create_job_body("Skill Check")).await; let job_id = created["id"].as_str().unwrap(); let skill_body = json!({"content": "---\nname: x\n---\nContent"}); @@ -880,7 +942,7 @@ async fn sk3_has_skill_false() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let created = create_job(&mut app, &token, &csrf, create_job_body("No Skill")).await; + let created = create_job(&mut app, &services, &token, &csrf, create_job_body("No Skill")).await; let job_id = created["id"].as_str().unwrap(); let req = get_with_token(&format!("/api/cron/jobs/{job_id}/skill"), &token); @@ -898,7 +960,7 @@ async fn sk4_save_empty_skill() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let created = create_job(&mut app, &token, &csrf, create_job_body("Empty Skill")).await; + let created = create_job(&mut app, &services, &token, &csrf, create_job_body("Empty Skill")).await; let job_id = created["id"].as_str().unwrap(); let skill_body = json!({"content": ""}); @@ -920,7 +982,7 @@ async fn sk5_save_placeholder_skill() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let created = create_job(&mut app, &token, &csrf, create_job_body("Placeholder Skill")).await; + let created = create_job(&mut app, &services, &token, &csrf, create_job_body("Placeholder Skill")).await; let job_id = created["id"].as_str().unwrap(); let skill_body = json!({"content": "TODO: fill in later"}); @@ -961,7 +1023,7 @@ async fn sk7_delete_skill() { let (mut app, services) = build_app().await; let (token, csrf) = setup_and_login(&mut app, &services, "admin", "StrongP@ss1").await; - let created = create_job(&mut app, &token, &csrf, create_job_body("Delete Skill Job")).await; + let created = create_job(&mut app, &services, &token, &csrf, create_job_body("Delete Skill Job")).await; let job_id = created["id"].as_str().unwrap(); let save_req = json_with_token( @@ -1028,7 +1090,7 @@ async fn sc6_cron_with_timezone() { "agent_config": default_assistant_agent_config("Shanghai Job") }); - let data = create_job(&mut app, &token, &csrf, body).await; + let data = create_job(&mut app, &services, &token, &csrf, body).await; let now = aionui_common::now_ms(); assert!(data["state"]["next_run_at_ms"].as_i64().unwrap() > now); } diff --git a/crates/aionui-app/tests/extension_e2e.rs b/crates/aionui-app/tests/extension_e2e.rs index ea2e7d30a..8b60bebab 100644 --- a/crates/aionui-app/tests/extension_e2e.rs +++ b/crates/aionui-app/tests/extension_e2e.rs @@ -1,6 +1,7 @@ mod common; -use axum::http::StatusCode; +use axum::body::Body; +use axum::http::{Request, StatusCode}; use serde_json::json; use tempfile::TempDir; use tower::ServiceExt; @@ -617,22 +618,27 @@ async fn eq19_channel_status_merges_extension_meta_for_persisted_row() { let ext_root = write_legacy_extension_fixture(&tmp); let (mut app, services) = build_app_with_extension_root(&ext_root).await; let repo = SqliteChannelRepository::new(services.database.pool().clone()); + let (token, _csrf) = setup_and_login(&mut app, &services, "user1", "pass1").await; + let owner_user_id = services.user_repo.find_by_username("user1").await.unwrap().unwrap().id; let now = now_ms(); - repo.upsert_plugin(&aionui_db::models::ChannelPluginRow { - id: "legacy-channel".to_string(), - r#type: "legacy-channel".to_string(), - name: "Legacy Channel Persisted".to_string(), - enabled: true, - config: "{\"token\":\"secret\"}".to_string(), - status: Some("running".to_string()), - last_connected: Some(now), - created_at: now, - updated_at: now, - }) + repo.upsert_plugin( + &owner_user_id, + &aionui_db::models::ChannelPluginRow { + id: "legacy-channel".to_string(), + owner_user_id: owner_user_id.clone(), + r#type: "legacy-channel".to_string(), + name: "Legacy Channel Persisted".to_string(), + enabled: true, + config: "{\"token\":\"secret\"}".to_string(), + status: Some("running".to_string()), + last_connected: Some(now), + created_at: now, + updated_at: now, + }, + ) .await .unwrap(); - let (token, _csrf) = setup_and_login(&mut app, &services, "user1", "pass1").await; let resp = app .oneshot(get_with_token("/api/channel/plugins", &token)) .await @@ -664,6 +670,7 @@ async fn eq20_enable_extension_channel_persists_config_and_exposes_status() { let (mut app, services) = build_app_with_extension_root(&ext_root).await; let repo = SqliteChannelRepository::new(services.database.pool().clone()); let (token, csrf) = setup_and_login(&mut app, &services, "user1", "pass1").await; + let owner_user_id = services.user_repo.find_by_username("user1").await.unwrap().unwrap().id; let enable_resp = app .clone() @@ -686,7 +693,11 @@ async fn eq20_enable_extension_channel_persists_config_and_exposes_status() { let enable_json = body_json(enable_resp).await; assert_eq!(enable_json["data"]["success"], true); - let row = repo.get_plugin("legacy-channel").await.unwrap().unwrap(); + let row = repo + .get_plugin(&owner_user_id, "legacy-channel") + .await + .unwrap() + .unwrap(); assert!(row.enabled); assert_eq!(row.r#type, "legacy-channel"); assert_eq!(row.status.as_deref(), Some("stopped")); @@ -811,6 +822,108 @@ async fn em4_disable_nonexistent_returns_not_found() { assert_eq!(resp.status(), StatusCode::NOT_FOUND); } +#[tokio::test] +async fn extension_enablement_and_contributions_are_isolated_by_user() { + let tmp = TempDir::new().unwrap(); + let ext_root = write_legacy_extension_fixture(&tmp); + let (mut app, services) = build_app_with_extension_root(&ext_root).await; + let (token_a, csrf_a) = setup_and_login(&mut app, &services, "extension-user-a", "pass-a").await; + let (token_b, csrf_b) = setup_and_login(&mut app, &services, "extension-user-b", "pass-b").await; + + let missing_csrf = Request::builder() + .method("POST") + .uri("/api/extensions/disable") + .header("content-type", "application/json") + .header("authorization", format!("Bearer {token_a}")) + .body(Body::from(r#"{"name":"legacy-suite"}"#)) + .unwrap(); + let response = app.clone().oneshot(missing_csrf).await.unwrap(); + assert_eq!(response.status(), StatusCode::FORBIDDEN); + assert_eq!(body_json(response).await["code"], "CSRF_INVALID"); + + let response = app + .clone() + .oneshot(json_with_token( + "POST", + "/api/extensions/disable", + json!({"name": "legacy-suite"}), + &token_a, + &csrf_a, + )) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let extensions_a = body_json( + app.clone() + .oneshot(get_with_token("/api/extensions", &token_a)) + .await + .unwrap(), + ) + .await; + let extensions_b = body_json( + app.clone() + .oneshot(get_with_token("/api/extensions", &token_b)) + .await + .unwrap(), + ) + .await; + assert_eq!(extensions_a["data"][0]["enabled"], false); + assert_eq!(extensions_b["data"][0]["enabled"], true); + + let skills_a = body_json( + app.clone() + .oneshot(get_with_token("/api/extensions/skills", &token_a)) + .await + .unwrap(), + ) + .await; + let skills_b = body_json( + app.clone() + .oneshot(get_with_token("/api/extensions/skills", &token_b)) + .await + .unwrap(), + ) + .await; + assert_eq!(skills_a["data"], json!([])); + assert_eq!(skills_b["data"].as_array().unwrap().len(), 1); + + let asset_a = app + .clone() + .oneshot(get_with_token( + "/api/extensions/legacy-suite/assets/assets/channel.png", + &token_a, + )) + .await + .unwrap(); + let asset_b = app + .clone() + .oneshot(get_with_token( + "/api/extensions/legacy-suite/assets/assets/channel.png", + &token_b, + )) + .await + .unwrap(); + assert_eq!(asset_a.status(), StatusCode::NOT_FOUND); + assert_eq!(asset_b.status(), StatusCode::OK); + + let response = app + .clone() + .oneshot(json_with_token( + "POST", + "/api/extensions/enable", + json!({"name": "legacy-suite"}), + &token_b, + &csrf_b, + )) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let extensions_a = body_json(app.oneshot(get_with_token("/api/extensions", &token_a)).await.unwrap()).await; + assert_eq!(extensions_a["data"][0]["enabled"], false); +} + // --------------------------------------------------------------------------- // HM — Hub marketplace // --------------------------------------------------------------------------- @@ -1303,8 +1416,33 @@ async fn skill_batch_import_reports_partial_failures_without_rolling_back_succes "limit_bytes": 50 * 1024 * 1024 }]) ); - assert!(paths.user_skills_dir.join("sample-alpha").join("SKILL.md").exists()); + // Imported skills for a real (non-default) user land in that user's + // scoped storage under `.users/`, never in the legacy shared root. + let user = services + .user_repo + .find_by_username("user1") + .await + .unwrap() + .expect("test user should exist"); + let skill_repo = aionui_db::SqliteSkillRepository::new(services.database.pool().clone()); + let alpha_row = aionui_db::ISkillRepository::find_by_name_for_user(&skill_repo, &user.id, "sample-alpha") + .await + .unwrap() + .expect("imported skill row should exist for the importing user"); + assert!( + alpha_row.path.contains(".users"), + "import must use user-scoped storage: {}", + alpha_row.path + ); + assert!(std::path::Path::new(&alpha_row.path).join("SKILL.md").exists()); + assert!(!paths.user_skills_dir.join("sample-alpha").exists()); assert!(!paths.user_skills_dir.join("sample-beta").exists()); + assert!( + aionui_db::ISkillRepository::find_by_name_for_user(&skill_repo, &user.id, "sample-beta") + .await + .unwrap() + .is_none() + ); let resp = app .oneshot(get_with_token("/api/skills/import-history", &token)) @@ -1341,12 +1479,28 @@ fn write_skill(dir: &std::path::Path, name: &str, description: &str) { async fn sl1_list_skills_tags_builtin_and_custom_with_source_field() { let tmp = TempDir::new().unwrap(); let (mut app, services, paths) = build_app_with_skill_paths(tmp.path()).await; - let (token, _csrf) = setup_and_login(&mut app, &services, "user1", "pass1").await; + let (token, csrf) = setup_and_login(&mut app, &services, "user1", "pass1").await; let builtin_dir = paths.builtin_skills_dir.clone(); write_skill(&builtin_dir, "review", "Built-in review skill"); - write_skill(&paths.user_skills_dir, "my-skill", "A user-imported skill"); - sync_skill_catalog_for_test(&services, &paths).await; + sync_skill_catalog_for_test(&services, &paths, "user1").await; + + // Real users get custom skills through the import API, which stores + // files under the user's scoped storage (never the legacy shared root). + let source_dir = tmp.path().join("import-src"); + write_skill(&source_dir, "my-skill", "A user-imported skill"); + let resp = app + .clone() + .oneshot(json_with_token( + "POST", + "/api/skills/import", + json!({ "skill_path": source_dir.join("my-skill").to_str().unwrap() }), + &token, + &csrf, + )) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); let resp = app.oneshot(get_with_token("/api/skills", &token)).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); @@ -1380,12 +1534,35 @@ async fn sl1_list_skills_tags_builtin_and_custom_with_source_field() { async fn sl2_list_skills_user_custom_overrides_builtin() { let tmp = TempDir::new().unwrap(); let (mut app, services, paths) = build_app_with_skill_paths(tmp.path()).await; - let (token, _csrf) = setup_and_login(&mut app, &services, "user1", "pass1").await; + let (token, csrf) = setup_and_login(&mut app, &services, "user1", "pass1").await; let builtin_dir = paths.builtin_skills_dir.clone(); write_skill(&builtin_dir, "review", "Built-in review"); - write_skill(&paths.user_skills_dir, "review", "Custom review override"); - sync_skill_catalog_for_test(&services, &paths).await; + sync_skill_catalog_for_test(&services, &paths, "user1").await; + + // A user-imported skill with the same name shadows the builtin row in + // the user's listing. The source dir name differs; the skill NAME in + // the frontmatter is what collides. + let source_dir = tmp.path().join("import-src"); + let override_dir = source_dir.join("review-override"); + std::fs::create_dir_all(&override_dir).unwrap(); + std::fs::write( + override_dir.join("SKILL.md"), + "---\nname: review\ndescription: Custom review override\n---\nBody", + ) + .unwrap(); + let resp = app + .clone() + .oneshot(json_with_token( + "POST", + "/api/skills/import", + json!({ "skill_path": override_dir.to_str().unwrap() }), + &token, + &csrf, + )) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); let resp = app.oneshot(get_with_token("/api/skills", &token)).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); @@ -1427,7 +1604,7 @@ async fn ba1_unified_skill_list_includes_auto_inject_builtin_entries() { write_skill(&auto_dir, "skill-creator", "Scaffold a new skill"); // A top-level builtin that must NOT appear in the auto list. write_skill(&builtin_dir, "review", "Top-level"); - sync_skill_catalog_for_test(&services, &paths).await; + sync_skill_catalog_for_test(&services, &paths, "user1").await; let resp = app.oneshot(get_with_token("/api/skills", &token)).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); diff --git a/crates/aionui-app/tests/local_mode.rs b/crates/aionui-app/tests/local_mode.rs index 85e7b0e7b..cbdf69f1d 100644 --- a/crates/aionui-app/tests/local_mode.rs +++ b/crates/aionui-app/tests/local_mode.rs @@ -52,3 +52,42 @@ async fn test_non_local_mode_requires_auth() { services.database.close().await; } + +#[tokio::test] +async fn test_local_mode_aionpro_requires_session_and_allows_bootstrap_provision_without_csrf() { + let db = aionui_db::init_database_memory().await.unwrap(); + let config = aionui_app::AppConfig { + identity_mode: aionui_app::IdentityMode::AionPro, + bootstrap_secret: Some("bootstrap-secret".to_string()), + ..Default::default() + }; + let services = aionui_app::AppServices::from_config(db, &config).await.unwrap(); + let router = aionui_app::create_router(&services).await.expect("build router"); + + let business_response = router + .clone() + .oneshot(Request::builder().uri("/api/settings").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(business_response.status(), StatusCode::UNAUTHORIZED); + + let provision_response = router + .oneshot( + Request::builder() + .method("PUT") + .uri("/api/auth/internal/external-users/pro-app") + .header("content-type", "application/json") + .header("x-aioncore-bootstrap-secret", "bootstrap-secret") + .body(Body::from(r#"{"user_type":"aionpro","username":"Pro App User"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(provision_response.status(), StatusCode::OK); + let body = to_bytes(provision_response.into_body(), usize::MAX).await.unwrap(); + let json: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(json["data"]["user_type"], "aionpro"); + assert_ne!(json["data"]["user_id"], "system_default_user"); + + services.database.close().await; +} diff --git a/crates/aionui-app/tests/message_e2e.rs b/crates/aionui-app/tests/message_e2e.rs index 9eb7e75e7..7f3d47384 100644 --- a/crates/aionui-app/tests/message_e2e.rs +++ b/crates/aionui-app/tests/message_e2e.rs @@ -35,6 +35,7 @@ async fn insert_message( created_at: i64, ) { let repo = aionui_db::SqliteConversationRepository::new(services.database.pool().clone()); + let user_id = repo.owner_user_id(conv_id).await.unwrap().unwrap(); let msg = aionui_db::models::MessageRow { id: msg_id.into(), conversation_id: conv_id.into(), @@ -46,15 +47,17 @@ async fn insert_message( hidden: false, created_at, }; - aionui_db::IConversationRepository::insert_message(&repo, &msg) + aionui_db::IConversationRepository::insert_message(&repo, &user_id, &msg) .await .unwrap(); } async fn update_conversation_workspace(services: &aionui_app::AppServices, conv_id: &str, workspace: &str) { let repo = aionui_db::SqliteConversationRepository::new(services.database.pool().clone()); + let user_id = repo.owner_user_id(conv_id).await.unwrap().unwrap(); IConversationRepository::update( &repo, + &user_id, conv_id, &ConversationRowUpdate { extra: Some(json!({ "workspace": workspace, "backend": "gemini" }).to_string()), @@ -73,6 +76,7 @@ async fn insert_acp_tool_message( created_at: i64, ) { let repo = aionui_db::SqliteConversationRepository::new(services.database.pool().clone()); + let user_id = repo.owner_user_id(conv_id).await.unwrap().unwrap(); let msg = aionui_db::models::MessageRow { id: msg_id.into(), conversation_id: conv_id.into(), @@ -99,14 +103,15 @@ async fn insert_acp_tool_message( hidden: false, created_at, }; - aionui_db::IConversationRepository::insert_message(&repo, &msg) + aionui_db::IConversationRepository::insert_message(&repo, &user_id, &msg) .await .unwrap(); } async fn upsert_artifact(services: &aionui_app::AppServices, artifact: aionui_db::ConversationArtifactRow) { let repo = aionui_db::SqliteConversationRepository::new(services.database.pool().clone()); - aionui_db::IConversationRepository::upsert_artifact(&repo, &artifact) + let user_id = repo.owner_user_id(&artifact.conversation_id).await.unwrap().unwrap(); + aionui_db::IConversationRepository::upsert_artifact(&repo, &user_id, &artifact) .await .unwrap(); } @@ -424,7 +429,8 @@ async fn t8_7_messages_exclude_legacy_cron_rows() { hidden: false, created_at: 2000, }; - aionui_db::IConversationRepository::insert_message(&repo, &msg) + let user_id = repo.owner_user_id(&conv_id).await.unwrap().unwrap(); + aionui_db::IConversationRepository::insert_message(&repo, &user_id, &msg) .await .unwrap(); } diff --git a/crates/aionui-app/tests/office_e2e.rs b/crates/aionui-app/tests/office_e2e.rs index 74a6b1570..e225601d9 100644 --- a/crates/aionui-app/tests/office_e2e.rs +++ b/crates/aionui-app/tests/office_e2e.rs @@ -19,7 +19,7 @@ use axum::http::StatusCode; use serde_json::json; use tower::ServiceExt; -use common::{body_json, get_request, json_with_token, setup_and_login}; +use common::{body_json, get_with_token, json_with_token, setup_and_login}; use aionui_app::{AppConfig, AppServices, build_module_states, create_router_with_states}; use aionui_office::{ConversionService, OfficeRouterState, OfficecliWatchManager, ProxyService, SnapshotService}; @@ -593,9 +593,12 @@ async fn dc9_invalid_conversion_target() { #[tokio::test] async fn rp2_ppt_proxy_ssrf_inactive_port() { - let (app, _services, _tmp) = build_office_app().await; + let (mut app, services, _tmp) = build_office_app().await; + let (token, _csrf) = setup_and_login(&mut app, &services, "user1", "pass123").await; - let req = get_request("/api/ppt-proxy/8080/index.html"); + // Proxy routes require auth (preview ports are scoped to the active + // Core user); the SSRF port check applies to authenticated requests. + let req = get_with_token("/api/ppt-proxy/8080/index.html", &token); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::FORBIDDEN); @@ -605,9 +608,10 @@ async fn rp2_ppt_proxy_ssrf_inactive_port() { #[tokio::test] async fn rp4_office_watch_proxy_ssrf_inactive_port() { - let (app, _services, _tmp) = build_office_app().await; + let (mut app, services, _tmp) = build_office_app().await; + let (token, _csrf) = setup_and_login(&mut app, &services, "user1", "pass123").await; - let req = get_request("/api/office-watch-proxy/9999/index.html"); + let req = get_with_token("/api/office-watch-proxy/9999/index.html", &token); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::FORBIDDEN); @@ -617,9 +621,10 @@ async fn rp4_office_watch_proxy_ssrf_inactive_port() { #[tokio::test] async fn ppt_proxy_root_path_returns_non_404() { - let (app, _services, _tmp) = build_office_app().await; + let (mut app, services, _tmp) = build_office_app().await; + let (token, _csrf) = setup_and_login(&mut app, &services, "user1", "pass123").await; - let req = get_request("/api/ppt-proxy/19999"); + let req = get_with_token("/api/ppt-proxy/19999", &token); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::FORBIDDEN); @@ -627,9 +632,10 @@ async fn ppt_proxy_root_path_returns_non_404() { #[tokio::test] async fn office_watch_proxy_root_path_returns_non_404() { - let (app, _services, _tmp) = build_office_app().await; + let (mut app, services, _tmp) = build_office_app().await; + let (token, _csrf) = setup_and_login(&mut app, &services, "user1", "pass123").await; - let req = get_request("/api/office-watch-proxy/19999"); + let req = get_with_token("/api/office-watch-proxy/19999", &token); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::FORBIDDEN); diff --git a/crates/aionui-app/tests/runtime_team_tools_e2e.rs b/crates/aionui-app/tests/runtime_team_tools_e2e.rs new file mode 100644 index 000000000..ddaa884f9 --- /dev/null +++ b/crates/aionui-app/tests/runtime_team_tools_e2e.rs @@ -0,0 +1,140 @@ +mod common; + +use aionui_ai_agent::{RuntimeTokenScope, TEAM_RUNTIME_TOKEN_SESSION_GENERATION}; +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use serde_json::{Value, json}; +use tower::ServiceExt; + +use common::{body_json, build_app}; + +const USER_ID: &str = "system_default_user"; +const CONVERSATION_ID: &str = "conv-runtime-team-tools"; + +fn context_request(user_id: &str, conversation_id: &str, token: Option<&str>) -> Request { + let mut builder = Request::builder() + .method("GET") + .uri("/api/runtime/team-tools/context") + .header("x-aionui-user-id", user_id) + .header("x-aionui-conversation-id", conversation_id); + if let Some(token) = token { + builder = builder.header("x-aionui-runtime-token", token); + } + builder.body(Body::empty()).unwrap() +} + +fn call_request(user_id: &str, conversation_id: &str, token: Option<&str>) -> Request { + let mut builder = Request::builder() + .method("POST") + .uri("/api/runtime/team-tools/call") + .header("content-type", "application/json") + .header("x-aionui-user-id", user_id) + .header("x-aionui-conversation-id", conversation_id); + if let Some(token) = token { + builder = builder.header("x-aionui-runtime-token", token); + } + builder + .body(Body::from( + serde_json::to_vec(&json!({ + "tool": "team_members", + "arguments": {} + })) + .unwrap(), + )) + .unwrap() +} + +fn assert_runtime_auth_failed(json: &Value) { + assert_eq!(json["success"], false); + assert_eq!(json["error"]["code"], "runtime_auth_failed"); +} + +#[tokio::test] +async fn runtime_team_tools_reject_missing_runtime_token() { + let (app, _services) = build_app().await; + + let resp = app + .oneshot(context_request(USER_ID, CONVERSATION_ID, None)) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = body_json(resp).await; + assert_runtime_auth_failed(&json); +} + +#[tokio::test] +async fn runtime_team_tools_reject_forged_user_header() { + let (app, services) = build_app().await; + let issue = services.runtime_token_service.issue( + USER_ID, + CONVERSATION_ID, + TEAM_RUNTIME_TOKEN_SESSION_GENERATION, + [RuntimeTokenScope::TeamContext], + ); + + let resp = app + .oneshot(context_request("other-user", CONVERSATION_ID, Some(&issue.token))) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = body_json(resp).await; + assert_runtime_auth_failed(&json); +} + +#[tokio::test] +async fn runtime_team_tools_reject_forged_conversation_header() { + let (app, services) = build_app().await; + let issue = services.runtime_token_service.issue( + USER_ID, + CONVERSATION_ID, + TEAM_RUNTIME_TOKEN_SESSION_GENERATION, + [RuntimeTokenScope::TeamContext], + ); + + let resp = app + .oneshot(context_request(USER_ID, "other-conversation", Some(&issue.token))) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = body_json(resp).await; + assert_runtime_auth_failed(&json); +} + +#[tokio::test] +async fn runtime_team_tools_reject_wrong_scope_for_call() { + let (app, services) = build_app().await; + let issue = services.runtime_token_service.issue( + USER_ID, + CONVERSATION_ID, + TEAM_RUNTIME_TOKEN_SESSION_GENERATION, + [RuntimeTokenScope::TeamContext], + ); + + let resp = app + .oneshot(call_request(USER_ID, CONVERSATION_ID, Some(&issue.token))) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = body_json(resp).await; + assert_runtime_auth_failed(&json); +} + +#[tokio::test] +async fn runtime_team_tools_accept_matching_runtime_token_before_service_lookup() { + let (app, services) = build_app().await; + let issue = services.runtime_token_service.issue( + USER_ID, + CONVERSATION_ID, + TEAM_RUNTIME_TOKEN_SESSION_GENERATION, + [RuntimeTokenScope::TeamContext], + ); + + let resp = app + .oneshot(context_request(USER_ID, CONVERSATION_ID, Some(&issue.token))) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + let json = body_json(resp).await; + assert_eq!(json["success"], false); + assert_eq!(json["error"]["code"], "conversation_not_found"); +} diff --git a/crates/aionui-app/tests/stt_stream_e2e.rs b/crates/aionui-app/tests/stt_stream_e2e.rs index fb7e65720..422a8e89e 100644 --- a/crates/aionui-app/tests/stt_stream_e2e.rs +++ b/crates/aionui-app/tests/stt_stream_e2e.rs @@ -66,7 +66,7 @@ async fn seed_stt_prefs(app: &TestApp, config: Value) { let service = aionui_system::ClientPrefService::new(repo); let mut req = aionui_api_types::UpdateClientPreferencesRequest::new(); req.insert("tools.speechToText".to_owned(), config); - service.update_preferences(req).await.unwrap(); + service.update_preferences("system_default_user", req).await.unwrap(); } type ClientWs = WebSocketStream>; diff --git a/crates/aionui-app/tests/team_e2e.rs b/crates/aionui-app/tests/team_e2e.rs index 423e87eda..809385a87 100644 --- a/crates/aionui-app/tests/team_e2e.rs +++ b/crates/aionui-app/tests/team_e2e.rs @@ -268,7 +268,8 @@ async fn tc3b_create_team_writes_legacy_extra_shape() { let conversation_id = data["assistants"][0]["conversation_id"].as_str().unwrap(); let repo = aionui_db::SqliteConversationRepository::new(services.database.pool().clone()); - let row = repo.get(conversation_id).await.unwrap().unwrap(); + let user_id = repo.owner_user_id(conversation_id).await.unwrap().unwrap(); + let row = repo.get(&user_id, conversation_id).await.unwrap().unwrap(); let extra: serde_json::Value = serde_json::from_str(&row.extra).unwrap(); assert_eq!(extra["teamId"], data["id"]); @@ -488,9 +489,9 @@ async fn team_api_rejects_cross_user_access() { let req = get_with_token(&format!("/api/teams/{team_id}"), &other_token); let resp = app.clone().oneshot(req).await.unwrap(); - assert_eq!(resp.status(), StatusCode::FORBIDDEN); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); - let forbidden_requests = [ + let hidden_requests = [ json_with_token( "PATCH", &format!("/api/teams/{team_id}/name"), @@ -535,9 +536,9 @@ async fn team_api_rejects_cross_user_access() { ), ]; - for req in forbidden_requests { + for req in hidden_requests { let resp = app.clone().oneshot(req).await.unwrap(); - assert_eq!(resp.status(), StatusCode::FORBIDDEN); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); } } @@ -621,7 +622,7 @@ async fn trs2_run_state_returns_active_run_payload() { assert!(body["data"]["active_run"]["starting_batch_count"].is_number()); assert!(body["data"]["active_run"]["running_batch_count"].is_number()); assert!(body["data"]["active_run"]["active_enqueue_lease_count"].is_number()); - assert!(body["data"]["active_run"]["slot_work"].as_array().unwrap().len() >= 1); + assert!(!body["data"]["active_run"]["slot_work"].as_array().unwrap().is_empty()); } #[tokio::test] @@ -657,9 +658,9 @@ async fn trs5_run_state_rejects_cross_user_access() { let req = get_with_token(&format!("/api/teams/{team_id}/run-state"), &other_token); let resp = app.oneshot(req).await.unwrap(); - assert_eq!(resp.status(), StatusCode::FORBIDDEN); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); let body = body_json(resp).await; - assert_eq!(body["code"], "FORBIDDEN"); + assert_eq!(body["code"], "NOT_FOUND"); } // TL-3: Each team contains full assistants info @@ -1087,9 +1088,15 @@ async fn es1b_team_mcp_list_assistants_matches_assistant_projection() { let ensure_resp = app.clone().oneshot(ensure_req).await.unwrap(); assert_eq!(ensure_resp.status(), StatusCode::OK); + let lead_user_id = services + .conversation_repo + .owner_user_id(lead_conversation_id) + .await + .unwrap() + .unwrap(); let lead_conversation = services .conversation_repo - .get(lead_conversation_id) + .get(&lead_user_id, lead_conversation_id) .await .unwrap() .unwrap(); @@ -1255,8 +1262,10 @@ async fn sm1b_team_send_persists_user_bubble_through_projection_adapter() { assert_eq!(resp.status(), StatusCode::OK); let repo = aionui_db::SqliteConversationRepository::new(services.database.pool().clone()); + let user_id = repo.owner_user_id(lead_conversation_id).await.unwrap().unwrap(); let messages = repo .list_messages_page( + &user_id, lead_conversation_id, &MessagePageParams { limit: 50, diff --git a/crates/aionui-app/tests/websocket_e2e.rs b/crates/aionui-app/tests/websocket_e2e.rs index 481190926..28ec2e0f9 100644 --- a/crates/aionui-app/tests/websocket_e2e.rs +++ b/crates/aionui-app/tests/websocket_e2e.rs @@ -9,7 +9,7 @@ use std::time::Duration; use aionui_api_types::WebSocketMessage; use aionui_app::{AppConfig, AppServices, create_router}; -use aionui_realtime::WebSocketManager; +use aionui_realtime::{EventBroadcaster, WebSocketManager}; use futures_util::{SinkExt, StreamExt}; use serde_json::{Value, json}; use tokio::net::TcpListener; @@ -39,8 +39,21 @@ async fn start_app() -> TestApp { TestApp { addr, services } } -/// Sign a valid JWT token for testing. -fn sign_token(app: &TestApp, user_id: &str) -> String { +/// Seed an active Core user and sign a valid JWT for it. +/// +/// The WS handshake resolves the token's user against the users table +/// (active row + matching session generation), so a bare signed token for +/// a nonexistent user is rejected with a realtime auth error. +async fn sign_token(app: &TestApp, user_id: &str) -> String { + sqlx::query( + "INSERT OR IGNORE INTO users (id, username, password_hash, created_at, updated_at) \ + VALUES (?, ?, 'hash', 0, 0)", + ) + .bind(user_id) + .bind(user_id) + .execute(app.services.database.pool()) + .await + .unwrap(); app.services.jwt_service.sign(user_id, "testuser").unwrap() } @@ -211,6 +224,27 @@ fn assert_invalid_message_error(msg: &Value) { ); } +/// Wait (bounded) until the manager reports exactly `expected` clients. +/// +/// Registration happens after the HTTP upgrade completes, so a fixed +/// post-connect sleep races under parallel test load; poll instead. +async fn wait_for_clients(app: &TestApp, expected: usize) { + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + loop { + if ws_manager(app).client_count() == expected { + return; + } + if tokio::time::Instant::now() >= deadline { + assert_eq!( + ws_manager(app).client_count(), + expected, + "timed out waiting for websocket client count" + ); + } + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + fn ws_manager(app: &TestApp) -> &Arc { &app.services.ws_manager } @@ -222,12 +256,10 @@ fn ws_manager(app: &TestApp) -> &Arc { #[tokio::test] async fn t1_1_valid_bearer_token_connects() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (_tx, _rx) = connect_bearer(app.addr, &token).await; - tokio::time::sleep(Duration::from_millis(50)).await; - - assert_eq!(ws_manager(&app).client_count(), 1); + wait_for_clients(&app, 1).await; } #[tokio::test] @@ -258,23 +290,19 @@ async fn t1_3_invalid_token_sends_auth_expired_then_closes() { #[tokio::test] async fn t1_4_token_from_cookie() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (_tx, _rx) = connect_cookie(app.addr, &token).await; - tokio::time::sleep(Duration::from_millis(50)).await; - - assert_eq!(ws_manager(&app).client_count(), 1); + wait_for_clients(&app, 1).await; } #[tokio::test] async fn t1_5_token_from_sec_websocket_protocol() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (_tx, _rx) = connect_protocol(app.addr, &token).await; - tokio::time::sleep(Duration::from_millis(50)).await; - - assert_eq!(ws_manager(&app).client_count(), 1); + wait_for_clients(&app, 1).await; } // =========================================================================== @@ -284,7 +312,7 @@ async fn t1_5_token_from_sec_websocket_protocol() { #[tokio::test] async fn t3_1_valid_json_with_no_registered_route_returns_unsupported_error() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (mut tx, mut rx) = connect_bearer(app.addr, &token).await; tokio::time::sleep(Duration::from_millis(50)).await; @@ -299,7 +327,7 @@ async fn t3_1_valid_json_with_no_registered_route_returns_unsupported_error() { #[tokio::test] async fn t3_2_invalid_json_returns_error() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (mut tx, mut rx) = connect_bearer(app.addr, &token).await; tokio::time::sleep(Duration::from_millis(50)).await; @@ -313,7 +341,7 @@ async fn t3_2_invalid_json_returns_error() { #[tokio::test] async fn t3_3_missing_fields_returns_error() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (mut tx, mut rx) = connect_bearer(app.addr, &token).await; tokio::time::sleep(Duration::from_millis(50)).await; @@ -331,13 +359,11 @@ async fn t3_3_missing_fields_returns_error() { #[tokio::test] async fn t4_1_broadcast_reaches_all_clients() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (_, mut rx1) = connect_bearer(app.addr, &token).await; let (_, mut rx2) = connect_bearer(app.addr, &token).await; - tokio::time::sleep(Duration::from_millis(50)).await; - - assert_eq!(ws_manager(&app).client_count(), 2); + wait_for_clients(&app, 2).await; let event = WebSocketMessage::new("test-broadcast", json!({"seq": 1})); ws_manager(&app).broadcast_all(event); @@ -349,18 +375,89 @@ async fn t4_1_broadcast_reaches_all_clients() { assert_eq!(msg2["name"], "test-broadcast"); } +#[tokio::test] +async fn t4_1_scoped_event_reaches_only_matching_user() { + let app = start_app().await; + let token_a = sign_token(&app, "user-a").await; + let token_b = sign_token(&app, "user-b").await; + + let (_, mut rx_a) = connect_bearer(app.addr, &token_a).await; + let (_, mut rx_b) = connect_bearer(app.addr, &token_b).await; + wait_for_clients(&app, 2).await; + assert_eq!(ws_manager(&app).client_count_for_user("user-a"), 1); + assert_eq!(ws_manager(&app).client_count_for_user("user-b"), 1); + + app.services.event_bus.broadcast(WebSocketMessage::new( + "scoped-broadcast", + json!({"user_id": "user-a", "seq": 1}), + )); + + let msg_a = read_text(&mut rx_a).await; + assert_eq!(msg_a["name"], "scoped-broadcast"); + assert_eq!(msg_a["data"]["user_id"], "user-a"); + + let timeout_result = tokio::time::timeout(Duration::from_millis(200), rx_b.next()).await; + assert!(timeout_result.is_err(), "user-b should not receive user-a event"); +} + +#[tokio::test] +async fn t4_1_unscoped_business_event_is_dropped_by_bridge() { + let app = start_app().await; + let token = sign_token(&app, "user-a").await; + + let (_, mut rx) = connect_bearer(app.addr, &token).await; + tokio::time::sleep(Duration::from_millis(50)).await; + + app.services.event_bus.broadcast(WebSocketMessage::new( + "conversation.listChanged", + json!({"conversation_id": "c1"}), + )); + + let timeout_result = tokio::time::timeout(Duration::from_millis(200), rx.next()).await; + assert!(timeout_result.is_err(), "unscoped business event should be dropped"); +} + +#[tokio::test] +async fn t4_1_whitelisted_global_event_reaches_all_users() { + let app = start_app().await; + let token_a = sign_token(&app, "user-a").await; + let token_b = sign_token(&app, "user-b").await; + + let (_, mut rx_a) = connect_bearer(app.addr, &token_a).await; + let (_, mut rx_b) = connect_bearer(app.addr, &token_b).await; + tokio::time::sleep(Duration::from_millis(50)).await; + + app.services.event_bus.broadcast(WebSocketMessage::new( + "runtime.statusChanged", + json!({"phase": "ready"}), + )); + + let msg_a = read_text(&mut rx_a).await; + let msg_b = read_text(&mut rx_b).await; + assert_eq!(msg_a["name"], "runtime.statusChanged"); + assert_eq!(msg_b["name"], "runtime.statusChanged"); + + app.services.event_bus.broadcast(WebSocketMessage::new( + "hub.state-changed", + json!({"name": "demo-extension", "status": "installed"}), + )); + + let msg_a = read_text(&mut rx_a).await; + let msg_b = read_text(&mut rx_b).await; + assert_eq!(msg_a["name"], "hub.state-changed"); + assert_eq!(msg_b["name"], "hub.state-changed"); +} + #[tokio::test] async fn t4_2_unicast_reaches_only_target() { use aionui_realtime::ConnectionId; let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (_, mut rx1) = connect_bearer(app.addr, &token).await; let (_, mut rx2) = connect_bearer(app.addr, &token).await; - tokio::time::sleep(Duration::from_millis(50)).await; - - assert_eq!(ws_manager(&app).client_count(), 2); + wait_for_clients(&app, 2).await; let first_conn_id = ConnectionId(1); let msg = WebSocketMessage::new("unicast-test", json!({"target": true})); @@ -376,19 +473,15 @@ async fn t4_2_unicast_reaches_only_target() { #[tokio::test] async fn t4_3_broadcast_after_disconnect_no_error() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (mut tx1, _rx1) = connect_bearer(app.addr, &token).await; let (_, mut rx2) = connect_bearer(app.addr, &token).await; - tokio::time::sleep(Duration::from_millis(50)).await; - - assert_eq!(ws_manager(&app).client_count(), 2); + wait_for_clients(&app, 2).await; // Disconnect client 1 tx1.send(tungstenite::Message::Close(None)).await.unwrap(); - tokio::time::sleep(Duration::from_millis(100)).await; - - assert_eq!(ws_manager(&app).client_count(), 1); + wait_for_clients(&app, 1).await; // Broadcast — should not error even though client 1 is gone let event = WebSocketMessage::new("after-disconnect", json!({})); @@ -405,7 +498,7 @@ async fn t4_3_broadcast_after_disconnect_no_error() { #[tokio::test] async fn t5_1_pong_does_not_generate_response() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (mut tx, mut rx) = connect_bearer(app.addr, &token).await; tokio::time::sleep(Duration::from_millis(50)).await; @@ -420,7 +513,7 @@ async fn t5_1_pong_does_not_generate_response() { #[tokio::test] async fn t5_2_subscribe_show_open_file_mode() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (mut tx, mut rx) = connect_bearer(app.addr, &token).await; tokio::time::sleep(Duration::from_millis(50)).await; @@ -441,7 +534,7 @@ async fn t5_2_subscribe_show_open_file_mode() { #[tokio::test] async fn t5_3_subscribe_show_open_directory_mode() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (mut tx, mut rx) = connect_bearer(app.addr, &token).await; tokio::time::sleep(Duration::from_millis(50)).await; @@ -461,7 +554,7 @@ async fn t5_3_subscribe_show_open_directory_mode() { #[tokio::test] async fn t5_4_subscribe_show_open_mixed_mode() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (mut tx, mut rx) = connect_bearer(app.addr, &token).await; tokio::time::sleep(Duration::from_millis(50)).await; @@ -485,16 +578,13 @@ async fn t5_4_subscribe_show_open_mixed_mode() { #[tokio::test] async fn t6_1_client_close_removes_from_manager() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let (mut tx, _rx) = connect_bearer(app.addr, &token).await; - tokio::time::sleep(Duration::from_millis(50)).await; - assert_eq!(ws_manager(&app).client_count(), 1); + wait_for_clients(&app, 1).await; tx.send(tungstenite::Message::Close(None)).await.unwrap(); - tokio::time::sleep(Duration::from_millis(100)).await; - - assert_eq!(ws_manager(&app).client_count(), 0); + wait_for_clients(&app, 0).await; } // =========================================================================== @@ -504,7 +594,7 @@ async fn t6_1_client_close_removes_from_manager() { #[tokio::test] async fn t7_1_multiple_concurrent_connections() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; let mut handles = Vec::new(); for _ in 0..10 { @@ -518,8 +608,7 @@ async fn t7_1_multiple_concurrent_connections() { connections.push(h.await.unwrap()); } - tokio::time::sleep(Duration::from_millis(100)).await; - assert_eq!(ws_manager(&app).client_count(), 10); + wait_for_clients(&app, 10).await; } // =========================================================================== @@ -529,7 +618,7 @@ async fn t7_1_multiple_concurrent_connections() { #[tokio::test] async fn t7_2_blacklisted_token_rejected() { let app = start_app().await; - let token = sign_token(&app, "user1"); + let token = sign_token(&app, "user1").await; // Blacklist the token app.services.jwt_service.blacklist_token(&token); diff --git a/crates/aionui-assistant/src/agent_catalog.rs b/crates/aionui-assistant/src/agent_catalog.rs index 7ba5bf5c3..33caa35d7 100644 --- a/crates/aionui-assistant/src/agent_catalog.rs +++ b/crates/aionui-assistant/src/agent_catalog.rs @@ -4,5 +4,5 @@ use crate::error::AssistantError; #[async_trait::async_trait] pub trait AssistantAgentCatalogPort: Send + Sync { - async fn list_management_agents(&self) -> Result, AssistantError>; + async fn list_management_agents(&self, user_id: &str) -> Result, AssistantError>; } diff --git a/crates/aionui-assistant/src/routes.rs b/crates/aionui-assistant/src/routes.rs index 28e2fdce8..e151ca446 100644 --- a/crates/aionui-assistant/src/routes.rs +++ b/crates/aionui-assistant/src/routes.rs @@ -5,7 +5,7 @@ use axum::Router; use axum::body::Body; use axum::extract::rejection::JsonRejection; -use axum::extract::{Json, Path, Query, State}; +use axum::extract::{Extension, Json, Path, Query, State}; use axum::http::{HeaderValue, StatusCode, header}; use axum::response::Response; use axum::routing::{get, patch, post}; @@ -14,6 +14,7 @@ use aionui_api_types::{ ApiResponse, AssistantDetailResponse, AssistantResponse, CreateAssistantRequest, ImportAssistantsRequest, ImportAssistantsResult, SetAssistantStateRequest, UpdateAssistantRequest, }; +use aionui_auth::CurrentUser; use aionui_common::ApiError; use crate::error::AssistantError; @@ -53,73 +54,87 @@ impl From for ApiError { async fn list( State(state): State, + Extension(current_user): Extension, ) -> Result>>, ApiError> { - let items = state.service.list().await?; + let items = state.service.list_for_user(¤t_user.id).await?; Ok(Json(ApiResponse::ok(items))) } async fn create( State(state): State, + Extension(current_user): Extension, body: Result, JsonRejection>, ) -> Result<(StatusCode, Json>), ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let created = state.service.create(req).await?; + let created = state.service.create_for_user(¤t_user.id, req).await?; Ok((StatusCode::CREATED, Json(ApiResponse::ok(created)))) } async fn get_one( State(state): State, + Extension(current_user): Extension, Path(id): Path, Query(query): Query, ) -> Result>, ApiError> { - let detail = state.service.get_detail(&id, query.locale.as_deref()).await?; + let detail = state + .service + .get_detail_for_user(¤t_user.id, &id, query.locale.as_deref()) + .await?; Ok(Json(ApiResponse::ok(detail))) } async fn update( State(state): State, + Extension(current_user): Extension, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let updated = state.service.update(&id, req).await?; + let updated = state.service.update_for_user(¤t_user.id, &id, req).await?; Ok(Json(ApiResponse::ok(updated))) } async fn delete_one( State(state): State, + Extension(current_user): Extension, Path(id): Path, ) -> Result>, ApiError> { - state.service.delete(&id).await?; + state.service.delete_for_user(¤t_user.id, &id).await?; Ok(Json(ApiResponse::success())) } async fn set_state( State(state): State, + Extension(current_user): Extension, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let resp = state.service.set_state(&id, req).await?; + let resp = state.service.set_state_for_user(¤t_user.id, &id, req).await?; Ok(Json(ApiResponse::ok(resp))) } async fn import( State(state): State, + Extension(current_user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let result = state.service.import(req).await?; + let result = state.service.import_for_user(¤t_user.id, req).await?; Ok(Json(ApiResponse::ok(result))) } /// Serve the raw avatar bytes for an assistant. Content-Type inferred from the /// file extension (png/jpg/svg default). Extensions return 404 — the frontend /// serves those via `aion-asset://`. -async fn get_avatar(State(state): State, Path(id): Path) -> Result { +async fn get_avatar( + State(state): State, + Extension(current_user): Extension, + Path(id): Path, +) -> Result { let asset = state .service - .avatar_asset(&id) + .avatar_asset_for_user(¤t_user.id, &id) .await .ok_or_else(|| ApiError::NotFound(format!("avatar '{id}' not found")))?; diff --git a/crates/aionui-assistant/src/service.rs b/crates/aionui-assistant/src/service.rs index ff1794af8..6fc294a89 100644 --- a/crates/aionui-assistant/src/service.rs +++ b/crates/aionui-assistant/src/service.rs @@ -19,7 +19,7 @@ use aionui_db::{ AssistantDefinitionRow, AssistantOverlayRow, AssistantRow, CreateAssistantParams, IAssistantDefinitionRepository, IAssistantOverlayRepository, IAssistantOverrideRepository, IAssistantPreferenceRepository, IAssistantRepository, IProviderRepository, SqlitePool, UpdateAssistantParams, UpsertAssistantDefinitionParams, - UpsertAssistantOverlayParams, UpsertAssistantPreferenceParams, resolve_agent_binding, + UpsertAssistantOverlayParams, UpsertAssistantPreferenceParams, resolve_agent_binding_for_user, }; use aionui_extension::{AssistantClassifier, AssistantRuleDispatcher, ExtensionError}; use serde_json; @@ -35,6 +35,7 @@ use crate::error::AssistantError; /// contended by a concurrent startup (Sentry 135525166 Option B). const BOOTSTRAP_RETRY_MAX_ATTEMPTS: u32 = 5; const BOOTSTRAP_RETRY_BACKOFF_MS: [u64; 4] = [50, 100, 200, 400]; +const DEFAULT_USER_ID: &str = "system_default_user"; /// Whether an assistant error is transient SQLite busy/locked contention. Repos /// convert `DbError` into `AssistantError::Internal(other.to_string())`, so the @@ -133,6 +134,14 @@ pub struct AssistantServiceDeps { pub agent_catalog: Option>, } +struct GeneratedAssistantReconcileContext<'a> { + definitions: &'a [AssistantDefinitionRow], + has_existing_generated: bool, + existing_min_sort_order: i32, + missing_generated_count: usize, + missing_index: &'a mut usize, +} + impl AssistantService { /// Construct an `AssistantService` pinned to the runtime data directory. /// @@ -228,14 +237,16 @@ impl AssistantService { .map_err(|e| AssistantError::Internal(format!("encode builtin disabled skills: {e}")))?; let (avatar_type, avatar_value) = serialize_avatar("builtin", builtin.avatar.as_deref()); let (definition_id, assistant_id) = self - .resolve_definition_identity("builtin", Some(&builtin.id), &builtin.id) + .resolve_global_definition_identity("builtin", Some(&builtin.id), &builtin.id) .await?; let existing_definition = self .definition_repo .get_by_id(&definition_id) .await .map_err(|e| AssistantError::Internal(format!("get builtin definition: {e}")))?; - let agent_id = self.resolve_agent_id_for_agent_ref(&builtin.agent_ref).await?; + let agent_id = self + .resolve_agent_id_for_agent_ref(DEFAULT_USER_ID, &builtin.agent_ref) + .await?; let default_model_mode = existing_definition .as_ref() .filter(|definition| definition.source == "builtin") @@ -265,7 +276,7 @@ impl AssistantService { .and_then(|definition| definition.default_thought_level_value.as_deref()); self.definition_repo - .upsert(&UpsertAssistantDefinitionParams { + .upsert_global(&UpsertAssistantDefinitionParams { id: &definition_id, assistant_id: &assistant_id, source: "builtin", @@ -311,7 +322,7 @@ impl AssistantService { for definition in self .definition_repo - .list() + .list_for_user(DEFAULT_USER_ID) .await .map_err(|e| AssistantError::Internal(format!("list assistant definitions: {e}")))? { @@ -341,8 +352,11 @@ impl AssistantService { } async fn sync_legacy_user_assistants_to_new_tables(&self) -> Result<(), AssistantError> { - for row in self.repo.list().await? { - if let Err(error) = self.sync_legacy_user_assistant_to_new_tables(&row).await { + for row in self.repo.list_for_user(DEFAULT_USER_ID).await? { + if let Err(error) = self + .sync_legacy_user_assistant_to_new_tables_for_user(DEFAULT_USER_ID, &row) + .await + { warn!( assistant_id = %row.id, error = %error, @@ -353,35 +367,44 @@ impl AssistantService { Ok(()) } - async fn sync_legacy_user_assistant_to_new_tables(&self, row: &AssistantRow) -> Result<(), AssistantError> { + async fn sync_legacy_user_assistant_to_new_tables_for_user( + &self, + user_id: &str, + row: &AssistantRow, + ) -> Result<(), AssistantError> { if self.builtin.has(&row.id) { return Ok(()); } if self .definition_repo - .get_by_source_ref_including_deleted("user", &row.id) + .get_by_source_ref_including_deleted_for_user(user_id, "user", &row.id) .await .map_err(|e| AssistantError::Internal(format!("get user definition by source_ref: {e}")))? .is_some() || self .definition_repo - .get_by_assistant_id_including_deleted(&row.id) + .get_by_assistant_id_including_deleted_for_user(user_id, &row.id) .await .map_err(|e| AssistantError::Internal(format!("get user definition by assistant_id: {e}")))? .is_some() { return Ok(()); } - self.upsert_definition_from_legacy_user_row(row, None).await?; + self.upsert_definition_from_legacy_user_row_for_user(user_id, row, None) + .await?; Ok(()) } async fn reconcile_user_avatar_assets(&self) -> Result<(), AssistantError> { - let definitions = self.definition_repo.list_including_deleted().await.map_err(|e| { - AssistantError::Internal(format!( - "list assistant definitions including deleted for avatar reconcile: {e}" - )) - })?; + let definitions = self + .definition_repo + .list_including_deleted_for_user(DEFAULT_USER_ID) + .await + .map_err(|e| { + AssistantError::Internal(format!( + "list assistant definitions including deleted for avatar reconcile: {e}" + )) + })?; for mut definition in definitions { if definition.avatar_type != "user_asset" { @@ -413,10 +436,10 @@ impl AssistantService { } async fn sync_legacy_overrides_to_new_states(&self) -> Result<(), AssistantError> { - for override_row in self.override_repo.get_all().await? { + for override_row in self.override_repo.get_all_for_user(DEFAULT_USER_ID).await? { let Some(definition) = self .definition_repo - .get_by_assistant_id(&override_row.assistant_id) + .get_by_assistant_id_for_user(DEFAULT_USER_ID, &override_row.assistant_id) .await? else { warn!( @@ -428,7 +451,7 @@ impl AssistantService { let existing_state = self .state_repo - .get(&definition.id) + .get_for_user(DEFAULT_USER_ID, &definition.id) .await .map_err(|e| AssistantError::Internal(format!("get assistant overlay: {e}")))?; @@ -442,13 +465,16 @@ impl AssistantService { } self.state_repo - .upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: &definition.id, - enabled: override_row.enabled, - sort_order: override_row.sort_order, - agent_id_override: None, - last_used_at: override_row.last_used_at, - }) + .upsert_for_user( + DEFAULT_USER_ID, + &UpsertAssistantOverlayParams { + assistant_definition_id: &definition.id, + enabled: override_row.enabled, + sort_order: override_row.sort_order, + agent_id_override: None, + last_used_at: override_row.last_used_at, + }, + ) .await .map_err(|e| AssistantError::Internal(format!("upsert assistant overlay: {e}")))?; } @@ -457,12 +483,19 @@ impl AssistantService { } async fn reconcile_generated_assistants(&self) -> Result, AssistantError> { + self.reconcile_generated_assistants_for_user(DEFAULT_USER_ID).await + } + + async fn reconcile_generated_assistants_for_user( + &self, + user_id: &str, + ) -> Result, AssistantError> { let Some(agent_catalog) = &self.agent_catalog else { return Ok(Vec::new()); }; - let rows = agent_catalog.list_management_agents().await?; - let definitions = self.definition_repo.list().await.map_err(|e| { + let rows = agent_catalog.list_management_agents(user_id).await?; + let definitions = self.definition_repo.list_for_user(user_id).await.map_err(|e| { AssistantError::Internal(format!("list assistant definitions for generated reconcile: {e}")) })?; let generated_source_refs: HashSet = definitions @@ -473,7 +506,7 @@ impl AssistantService { let has_existing_generated = !generated_source_refs.is_empty(); let existing_min_sort_order = self .state_repo - .list() + .list_for_user(user_id) .await .map_err(|e| AssistantError::Internal(format!("list assistant overlays for generated reconcile: {e}")))? .into_iter() @@ -502,12 +535,15 @@ impl AssistantService { for row in generated_rows { if let Err(error) = self .reconcile_generated_assistant( + user_id, row, - &definitions, - has_existing_generated, - existing_min_sort_order, - missing_generated_count, - &mut missing_index, + GeneratedAssistantReconcileContext { + definitions: &definitions, + has_existing_generated, + existing_min_sort_order, + missing_generated_count, + missing_index: &mut missing_index, + }, ) .await { @@ -524,14 +560,12 @@ impl AssistantService { async fn reconcile_generated_assistant( &self, + user_id: &str, row: &AgentManagementRow, - definitions: &[AssistantDefinitionRow], - has_existing_generated: bool, - existing_min_sort_order: i32, - missing_generated_count: usize, - missing_index: &mut usize, + context: GeneratedAssistantReconcileContext<'_>, ) -> Result<(), AssistantError> { - let existing_definition = definitions + let existing_definition = context + .definitions .iter() .find(|definition| { definition.source == "generated" && definition.source_ref.as_deref() == Some(row.id.as_str()) @@ -540,7 +574,7 @@ impl AssistantService { let is_missing = existing_definition.is_none(); let assistant_id = format!("bare:{}", row.id); let (definition_id, assistant_id) = self - .resolve_definition_identity("generated", Some(&row.id), &assistant_id) + .resolve_definition_identity_for_user(user_id, "generated", Some(&row.id), &assistant_id) .await?; let avatar_value = row.icon.as_deref().filter(|value| !value.trim().is_empty()); let (definition, should_upsert) = if let Some(mut definition) = existing_definition { @@ -612,7 +646,7 @@ impl AssistantService { if should_upsert { self.definition_repo - .upsert(&upsert_params_from_definition(&definition)) + .upsert_for_user(user_id, &upsert_params_from_definition(&definition)) .await .map_err(|e| AssistantError::Internal(format!("upsert generated assistant definition: {e}")))?; } @@ -623,26 +657,31 @@ impl AssistantService { if self .state_repo - .get(&definition_id) + .get_for_user(user_id, &definition_id) .await .map_err(|e| AssistantError::Internal(format!("get generated assistant overlay: {e}")))? .is_none() { - let current_missing_index = *missing_index; - *missing_index += 1; - let initial_generated_sort_order = if !has_existing_generated && missing_generated_count > 0 { - existing_min_sort_order as i64 - missing_generated_count as i64 + current_missing_index as i64 + let current_missing_index = *context.missing_index; + *context.missing_index += 1; + let initial_generated_sort_order = if !context.has_existing_generated && context.missing_generated_count > 0 + { + context.existing_min_sort_order as i64 - context.missing_generated_count as i64 + + current_missing_index as i64 } else { row.sort_order }; self.state_repo - .upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: &definition_id, - enabled: true, - sort_order: initial_generated_sort_order.clamp(i32::MIN as i64, i32::MAX as i64) as i32, - agent_id_override: None, - last_used_at: None, - }) + .upsert_for_user( + user_id, + &UpsertAssistantOverlayParams { + assistant_definition_id: &definition_id, + enabled: true, + sort_order: initial_generated_sort_order.clamp(i32::MIN as i64, i32::MAX as i64) as i32, + agent_id_override: None, + last_used_at: None, + }, + ) .await .map_err(|e| AssistantError::Internal(format!("upsert generated assistant overlay: {e}")))?; } @@ -650,8 +689,9 @@ impl AssistantService { Ok(()) } - async fn upsert_definition_from_legacy_user_row( + async fn upsert_definition_from_legacy_user_row_for_user( &self, + user_id: &str, row: &AssistantRow, requested_agent_id: Option<&str>, ) -> Result<(), AssistantError> { @@ -666,58 +706,67 @@ impl AssistantService { let custom_skill_names = normalize_json_array_string(row.custom_skill_names.as_deref(), "custom_skill_names")?; let default_disabled_builtin_skill_ids = normalize_json_array_string(row.disabled_builtin_skills.as_deref(), "disabled_builtin_skills")?; - let (definition_id, assistant_id) = self.resolve_definition_identity("user", Some(&row.id), &row.id).await?; + let (definition_id, assistant_id) = self + .resolve_definition_identity_for_user(user_id, "user", Some(&row.id), &row.id) + .await?; let (avatar_type, avatar_value) = self.normalize_legacy_user_avatar_input(&assistant_id, row.avatar.as_deref())?; - let existing_definition = self.definition_repo.get_by_assistant_id(&assistant_id).await?; + let existing_definition = self + .definition_repo + .get_by_assistant_id_for_user(user_id, &assistant_id) + .await?; let agent_id = match requested_agent_id { Some(agent_id) => agent_id.to_string(), None => match existing_definition { Some(definition) => definition.agent_id, - None => self.resolve_default_agent_id().await?, + None => self.resolve_default_agent_id_for_user(user_id).await?, }, }; - self.resolve_runtime_backend_for_agent_id(&agent_id).await?; + self.resolve_runtime_backend_for_agent_id(user_id, &agent_id).await?; self.definition_repo - .upsert(&UpsertAssistantDefinitionParams { - id: &definition_id, - assistant_id: &assistant_id, - source: "user", - owner_type: "user", - source_ref: Some(&row.id), - name: &row.name, - name_i18n: &name_i18n, - description: row.description.as_deref(), - description_i18n: &description_i18n, - avatar_type: &avatar_type, - avatar_value: avatar_value.as_deref(), - agent_id: &agent_id, - rule_resource_type: "user_file", - rule_resource_ref: Some(&row.id), - recommended_prompts: &recommended_prompts, - recommended_prompts_i18n: &recommended_prompts_i18n, - default_model_mode: "auto", - default_model_value: None, - default_permission_mode: "auto", - default_permission_value: None, - default_thought_level_mode: "auto", - default_thought_level_value: None, - default_skills_mode: "fixed", - default_skill_ids: &default_skill_ids, - custom_skill_names: &custom_skill_names, - default_disabled_builtin_skill_ids: &default_disabled_builtin_skill_ids, - default_mcps_mode: "auto", - default_mcp_ids: "[]", - }) + .upsert_for_user( + user_id, + &UpsertAssistantDefinitionParams { + id: &definition_id, + assistant_id: &assistant_id, + source: "user", + owner_type: "user", + source_ref: Some(&row.id), + name: &row.name, + name_i18n: &name_i18n, + description: row.description.as_deref(), + description_i18n: &description_i18n, + avatar_type: &avatar_type, + avatar_value: avatar_value.as_deref(), + agent_id: &agent_id, + rule_resource_type: "user_file", + rule_resource_ref: Some(&row.id), + recommended_prompts: &recommended_prompts, + recommended_prompts_i18n: &recommended_prompts_i18n, + default_model_mode: "auto", + default_model_value: None, + default_permission_mode: "auto", + default_permission_value: None, + default_thought_level_mode: "auto", + default_thought_level_value: None, + default_skills_mode: "fixed", + default_skill_ids: &default_skill_ids, + custom_skill_names: &custom_skill_names, + default_disabled_builtin_skill_ids: &default_disabled_builtin_skill_ids, + default_mcps_mode: "auto", + default_mcp_ids: "[]", + }, + ) .await .map_err(|e| AssistantError::Internal(format!("upsert user definition: {e}")))?; Ok(()) } - async fn apply_detail_overrides( + async fn apply_detail_overrides_for_user( &self, + user_id: &str, assistant_id: &str, overrides: SerializedDetailOverrides, reset_model_and_permission: bool, @@ -728,7 +777,7 @@ impl AssistantService { let Some(existing) = self .definition_repo - .get_by_assistant_id(assistant_id) + .get_by_assistant_id_for_user(user_id, assistant_id) .await .map_err(|e| AssistantError::Internal(format!("get assistant definition: {e}")))? else { @@ -736,10 +785,14 @@ impl AssistantService { }; let mut patched = existing.clone(); + if patched.source == "builtin" && patched.owner_type == "system" { + patched.id = generate_prefixed_id("asstdef"); + patched.owner_type = "user".to_string(); + } apply_detail_patch_to_definition(&mut patched, &overrides, reset_model_and_permission); self.definition_repo - .upsert(&upsert_params_from_definition(&patched)) + .upsert_for_user(user_id, &upsert_params_from_definition(&patched)) .await .map_err(|e| AssistantError::Internal(format!("upsert patched assistant definition: {e}")))?; @@ -752,10 +805,15 @@ impl AssistantService { /// Classify an assistant id into its source. pub async fn classify_source(&self, id: &str) -> AssistantSource { + self.classify_source_for_user(DEFAULT_USER_ID, id).await + } + + /// Classify an assistant id into its source for the current user. + pub async fn classify_source_for_user(&self, user_id: &str, id: &str) -> AssistantSource { if self.builtin.has(id) { return AssistantSource::Builtin; } - if let Ok(Some(definition)) = self.definition_repo.get_by_assistant_id(id).await { + if let Ok(Some(definition)) = self.definition_repo.get_by_assistant_id_for_user(user_id, id).await { return match definition.source.as_str() { "builtin" => AssistantSource::Builtin, "generated" => AssistantSource::Generated, @@ -773,15 +831,19 @@ impl AssistantService { /// application. Also performs opportunistic orphan cleanup on the /// overrides table. pub async fn list(&self) -> Result, AssistantError> { - let projections = self.reconcile_generated_assistants().await?; + self.list_for_user(DEFAULT_USER_ID).await + } + + pub async fn list_for_user(&self, user_id: &str) -> Result, AssistantError> { + let projections = self.reconcile_generated_assistants_for_user(user_id).await?; let definitions = self .definition_repo - .list() + .list_for_user(user_id) .await .map_err(|e| AssistantError::Internal(format!("list assistant definitions: {e}")))?; let states = self .state_repo - .list() + .list_for_user(user_id) .await .map_err(|e| AssistantError::Internal(format!("list assistant overlays: {e}")))?; let state_map: HashMap = states @@ -796,7 +858,7 @@ impl AssistantService { continue; } let projection = self - .project_definition(definition, state_map.get(&definition.id), &projections) + .project_definition(user_id, definition, state_map.get(&definition.id), &projections) .await?; result.push(self.definition_to_response(definition, state_map.get(&definition.id), &projection)?); } @@ -811,7 +873,7 @@ impl AssistantService { // Opportunistic orphan cleanup: any override row whose assistant_id no // longer appears in the merged list is stale. let valid_ids: Vec<&str> = result.iter().map(|a| a.id.as_str()).collect(); - if let Err(e) = self.override_repo.delete_orphans(&valid_ids).await { + if let Err(e) = self.override_repo.delete_orphans_for_user(user_id, &valid_ids).await { warn!("override orphan cleanup failed: {e}"); } @@ -819,14 +881,18 @@ impl AssistantService { } pub async fn get(&self, id: &str) -> Result { - let projections = self.reconcile_generated_assistants().await?; - if let Some(definition) = self.definition_repo.get_by_assistant_id(id).await? { + self.get_for_user(DEFAULT_USER_ID, id).await + } + + pub async fn get_for_user(&self, user_id: &str, id: &str) -> Result { + let projections = self.reconcile_generated_assistants_for_user(user_id).await?; + if let Some(definition) = self.definition_repo.get_by_assistant_id_for_user(user_id, id).await? { if generated_definition_is_uninstalled(&definition, &projections) { return Err(AssistantError::NotFound(format!("assistant '{id}' not found"))); } - let state = self.state_repo.get(&definition.id).await?; + let state = self.state_repo.get_for_user(user_id, &definition.id).await?; let projection = self - .project_definition(&definition, state.as_ref(), &projections) + .project_definition(user_id, &definition, state.as_ref(), &projections) .await?; return self.definition_to_response(&definition, state.as_ref(), &projection); } @@ -835,16 +901,25 @@ impl AssistantService { } pub async fn get_detail(&self, id: &str, locale: Option<&str>) -> Result { - let projections = self.reconcile_generated_assistants().await?; - if let Some(definition) = self.definition_repo.get_by_assistant_id(id).await? { + self.get_detail_for_user(DEFAULT_USER_ID, id, locale).await + } + + pub async fn get_detail_for_user( + &self, + user_id: &str, + id: &str, + locale: Option<&str>, + ) -> Result { + let projections = self.reconcile_generated_assistants_for_user(user_id).await?; + if let Some(definition) = self.definition_repo.get_by_assistant_id_for_user(user_id, id).await? { if generated_definition_is_uninstalled(&definition, &projections) { return Err(AssistantError::NotFound(format!("assistant '{id}' not found"))); } - let state = self.state_repo.get(&definition.id).await?; - let preference = self.preference_repo.get(&definition.id).await?; - let rules_content = self.read_rule(id, locale).await?; + let state = self.state_repo.get_for_user(user_id, &definition.id).await?; + let preference = self.preference_repo.get_for_user(user_id, &definition.id).await?; + let rules_content = self.read_rule_for_user(user_id, id, locale).await?; let projection = self - .project_definition(&definition, state.as_ref(), &projections) + .project_definition(user_id, &definition, state.as_ref(), &projections) .await?; return self.definition_to_detail_response( &definition, @@ -879,14 +954,18 @@ impl AssistantService { /// which on machines without the Gemini CLI 400'd within 1 ms /// with `Agent 'Gemini CLI' CLI not found in PATH`. pub async fn resolve_default_agent_id(&self) -> Result { + self.resolve_default_agent_id_for_user(DEFAULT_USER_ID).await + } + + pub async fn resolve_default_agent_id_for_user(&self, user_id: &str) -> Result { let providers = self .provider_repo - .list() + .list(user_id) .await .map_err(|e| AssistantError::Internal(format!("failed to list providers: {e}")))?; if providers.iter().any(|p| p.enabled) { - self.resolve_agent_id_for_agent_ref("aionrs").await + self.resolve_agent_id_for_agent_ref(user_id, "aionrs").await } else { Err(AssistantError::BadRequest( "Cannot create assistant: no providers configured. Add a provider before creating an assistant, \ @@ -896,12 +975,16 @@ impl AssistantService { } } - async fn resolve_runtime_backend_for_agent_id(&self, agent_id: &str) -> Result { + async fn resolve_runtime_backend_for_agent_id( + &self, + user_id: &str, + agent_id: &str, + ) -> Result { let trimmed = agent_id.trim(); if trimmed.is_empty() { return Err(AssistantError::BadRequest("agent_id is required".into())); } - let Some(binding) = resolve_agent_binding(&self.pool, trimmed) + let Some(binding) = resolve_agent_binding_for_user(&self.pool, user_id, trimmed) .await .map_err(|e| AssistantError::Internal(format!("resolve agent binding: {e}")))? else { @@ -910,9 +993,9 @@ impl AssistantService { Ok(binding.runtime_backend) } - async fn resolve_agent_id_for_agent_ref(&self, agent_ref: &str) -> Result { + async fn resolve_agent_id_for_agent_ref(&self, user_id: &str, agent_ref: &str) -> Result { let trimmed = agent_ref.trim(); - let Some(binding) = resolve_agent_binding(&self.pool, trimmed) + let Some(binding) = resolve_agent_binding_for_user(&self.pool, user_id, trimmed) .await .map_err(|e| AssistantError::Internal(format!("resolve agent binding: {e}")))? else { @@ -923,12 +1006,13 @@ impl AssistantService { async fn project_definition( &self, + user_id: &str, definition: &AssistantDefinitionRow, state: Option<&AssistantOverlayRow>, agent_rows: &[AgentManagementRow], ) -> Result { let effective_agent_id = effective_agent_id_for_definition(definition, state); - let runtime_backend = resolve_agent_binding(&self.pool, effective_agent_id) + let runtime_backend = resolve_agent_binding_for_user(&self.pool, user_id, effective_agent_id) .await .map_err(|e| AssistantError::Internal(format!("resolve agent binding: {e}")))? .map(|binding| binding.runtime_backend); @@ -945,6 +1029,14 @@ impl AssistantService { // ----------------------------------------------------------------------- pub async fn create(&self, req: CreateAssistantRequest) -> Result { + self.create_for_user(DEFAULT_USER_ID, req).await + } + + pub async fn create_for_user( + &self, + user_id: &str, + req: CreateAssistantRequest, + ) -> Result { let name = req.name.trim().to_string(); if name.is_empty() { return Err(AssistantError::BadRequest("name is required".into())); @@ -970,9 +1062,10 @@ impl AssistantService { // the Gemini CLI (ELECTRON-1J1, ELECTRON-1KV). let resolved_agent_id = match req.agent_id.as_deref() { Some(s) if !s.trim().is_empty() => s.trim().to_string(), - _ => self.resolve_default_agent_id().await?, + _ => self.resolve_default_agent_id_for_user(user_id).await?, }; - self.resolve_runtime_backend_for_agent_id(&resolved_agent_id).await?; + self.resolve_runtime_backend_for_agent_id(user_id, &resolved_agent_id) + .await?; let avatar = self.normalize_user_avatar_input(&id, req.avatar.as_deref())?; let params = CreateAssistantParams { id: &id, @@ -989,19 +1082,33 @@ impl AssistantService { prompts_i18n: serialized.prompts_i18n.as_deref(), }; - let row = self.repo.create(¶ms).await?; - self.upsert_definition_from_legacy_user_row(&row, Some(&resolved_agent_id)) + let row = self.repo.create_for_user(user_id, ¶ms).await?; + self.upsert_definition_from_legacy_user_row_for_user(user_id, &row, Some(&resolved_agent_id)) + .await?; + self.apply_detail_overrides_for_user(user_id, &row.id, detail_overrides, false) .await?; - self.apply_detail_overrides(&row.id, detail_overrides, false).await?; - if let Some(definition) = self.definition_repo.get_by_assistant_id(&row.id).await? { - self.sync_preferences_from_defaults_request(&definition, None, req.defaults.as_ref()) + if let Some(definition) = self + .definition_repo + .get_by_assistant_id_for_user(user_id, &row.id) + .await? + { + self.sync_preferences_from_defaults_request_for_user(user_id, &definition, None, req.defaults.as_ref()) .await?; } - self.get(&id).await + self.get_for_user(user_id, &id).await } pub async fn update(&self, id: &str, req: UpdateAssistantRequest) -> Result { - match self.classify_source(id).await { + self.update_for_user(DEFAULT_USER_ID, id, req).await + } + + pub async fn update_for_user( + &self, + user_id: &str, + id: &str, + req: UpdateAssistantRequest, + ) -> Result { + match self.classify_source_for_user(user_id, id).await { AssistantSource::Builtin => { let detail_overrides = SerializedDetailOverrides::from_update(&req)?; let builtin_defaults_forbidden = req @@ -1036,18 +1143,18 @@ impl AssistantService { let definition = self .definition_repo - .get_by_assistant_id(id) + .get_by_assistant_id_for_user(user_id, id) .await? .ok_or_else(|| AssistantError::NotFound(format!("assistant '{id}' not found")))?; - let existing = self.override_repo.get(id).await?; + let existing = self.override_repo.get_for_user(user_id, id).await?; let enabled = existing.as_ref().is_none_or(|o| o.enabled); let sort_order = existing.as_ref().map(|o| o.sort_order).unwrap_or(0); let last_used_at = existing.as_ref().and_then(|o| o.last_used_at); let requested_agent_id = req.agent_id.as_deref().map(|agent_id| agent_id.trim().to_string()); let current_agent_id = self .state_repo - .get(&definition.id) + .get_for_user(user_id, &definition.id) .await .map_err(|e| AssistantError::Internal(format!("get assistant overlay: {e}")))? .and_then(|row| row.agent_id_override) @@ -1056,28 +1163,39 @@ impl AssistantService { .as_deref() .is_some_and(|agent_id| agent_id != current_agent_id); if let Some(requested_agent_id) = requested_agent_id.as_deref() { - self.resolve_runtime_backend_for_agent_id(requested_agent_id).await?; - self.state_repo - .upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: &definition.id, - enabled, - sort_order, - agent_id_override: Some(requested_agent_id), - last_used_at, - }) - .await - .map_err(|e| AssistantError::Internal(format!("upsert assistant overlay: {e}")))?; + self.resolve_runtime_backend_for_agent_id(user_id, requested_agent_id) + .await?; } - self.apply_detail_overrides(id, detail_overrides, reset_model_and_permission) + self.apply_detail_overrides_for_user(user_id, id, detail_overrides, reset_model_and_permission) .await?; let definition = self .definition_repo - .get_by_assistant_id(id) + .get_by_assistant_id_for_user(user_id, id) .await? .ok_or_else(|| AssistantError::NotFound(format!("assistant '{id}' not found")))?; - self.sync_preferences_from_defaults_request(&definition, Some(&definition), req.defaults.as_ref()) - .await?; - return self.get(id).await; + if let Some(requested_agent_id) = requested_agent_id.as_deref() { + self.state_repo + .upsert_for_user( + user_id, + &UpsertAssistantOverlayParams { + assistant_definition_id: &definition.id, + enabled, + sort_order, + agent_id_override: Some(requested_agent_id), + last_used_at, + }, + ) + .await + .map_err(|e| AssistantError::Internal(format!("upsert assistant overlay: {e}")))?; + } + self.sync_preferences_from_defaults_request_for_user( + user_id, + &definition, + Some(&definition), + req.defaults.as_ref(), + ) + .await?; + return self.get_for_user(user_id, id).await; } AssistantSource::Generated => { if req.name.is_some() @@ -1100,7 +1218,7 @@ impl AssistantService { let detail_overrides = SerializedDetailOverrides::from_update(&req)?; let current_definition = self .definition_repo - .get_by_assistant_id(id) + .get_by_assistant_id_for_user(user_id, id) .await? .ok_or_else(|| AssistantError::NotFound(format!("assistant '{id}' not found")))?; let mut patched = current_definition.clone(); @@ -1128,12 +1246,17 @@ impl AssistantService { let patched = self .definition_repo - .upsert(&upsert_params_from_definition(&patched)) + .upsert_for_user(user_id, &upsert_params_from_definition(&patched)) .await .map_err(|e| AssistantError::Internal(format!("upsert generated assistant definition: {e}")))?; - self.sync_preferences_from_defaults_request(&patched, Some(¤t_definition), req.defaults.as_ref()) - .await?; - return self.get(id).await; + self.sync_preferences_from_defaults_request_for_user( + user_id, + &patched, + Some(¤t_definition), + req.defaults.as_ref(), + ) + .await?; + return self.get_for_user(user_id, id).await; } AssistantSource::User => {} } @@ -1142,7 +1265,7 @@ impl AssistantService { let detail_overrides = SerializedDetailOverrides::from_update(&req)?; let current_definition = self .definition_repo - .get_by_assistant_id(id) + .get_by_assistant_id_for_user(user_id, id) .await? .ok_or_else(|| AssistantError::NotFound(format!("assistant '{id}' not found")))?; let requested_agent_id = match req.agent_id.as_deref() { @@ -1151,7 +1274,7 @@ impl AssistantService { None => None, }; if let Some(agent_id) = requested_agent_id.as_deref() { - self.resolve_runtime_backend_for_agent_id(agent_id).await?; + self.resolve_runtime_backend_for_agent_id(user_id, agent_id).await?; } let reset_model_and_permission = requested_agent_id .as_deref() @@ -1177,22 +1300,28 @@ impl AssistantService { let row = self .repo - .update(id, ¶ms) + .update_for_user(user_id, id, ¶ms) .await? .ok_or_else(|| AssistantError::NotFound(format!("assistant '{id}' not found")))?; - self.upsert_definition_from_legacy_user_row(&row, requested_agent_id.as_deref()) + self.upsert_definition_from_legacy_user_row_for_user(user_id, &row, requested_agent_id.as_deref()) .await?; - self.apply_detail_overrides(id, detail_overrides, reset_model_and_permission) + self.apply_detail_overrides_for_user(user_id, id, detail_overrides, reset_model_and_permission) + .await?; + if let Some(definition) = self.definition_repo.get_by_assistant_id_for_user(user_id, id).await? { + self.sync_preferences_from_defaults_request_for_user( + user_id, + &definition, + Some(¤t_definition), + req.defaults.as_ref(), + ) .await?; - if let Some(definition) = self.definition_repo.get_by_assistant_id(id).await? { - self.sync_preferences_from_defaults_request(&definition, Some(¤t_definition), req.defaults.as_ref()) - .await?; } - self.get(id).await + self.get_for_user(user_id, id).await } - async fn sync_preferences_from_defaults_request( + async fn sync_preferences_from_defaults_request_for_user( &self, + user_id: &str, definition: &AssistantDefinitionRow, previous_definition: Option<&AssistantDefinitionRow>, defaults: Option<&AssistantDefaultsRequest>, @@ -1203,7 +1332,7 @@ impl AssistantService { let existing = self .preference_repo - .get(&definition.id) + .get_for_user(user_id, &definition.id) .await .map_err(|e| AssistantError::Internal(format!("get assistant preference: {e}")))?; @@ -1327,7 +1456,7 @@ impl AssistantService { { if existing.is_some() { self.preference_repo - .delete(&definition.id) + .delete_for_user(user_id, &definition.id) .await .map_err(|e| AssistantError::Internal(format!("delete assistant preference: {e}")))?; } @@ -1342,15 +1471,18 @@ impl AssistantService { .map_err(|e| AssistantError::Internal(format!("encode assistant mcp preference: {e}")))?; self.preference_repo - .upsert(&UpsertAssistantPreferenceParams { - assistant_definition_id: &definition.id, - last_model_id: last_model_id.as_deref(), - last_permission_value: last_permission_value.as_deref(), - last_thought_level_value: last_thought_level_value.as_deref(), - last_skill_ids: &last_skill_ids_json, - last_disabled_builtin_skill_ids: &last_disabled_builtin_skill_ids_json, - last_mcp_ids: &last_mcp_ids_json, - }) + .upsert_for_user( + user_id, + &UpsertAssistantPreferenceParams { + assistant_definition_id: &definition.id, + last_model_id: last_model_id.as_deref(), + last_permission_value: last_permission_value.as_deref(), + last_thought_level_value: last_thought_level_value.as_deref(), + last_skill_ids: &last_skill_ids_json, + last_disabled_builtin_skill_ids: &last_disabled_builtin_skill_ids_json, + last_mcp_ids: &last_mcp_ids_json, + }, + ) .await .map_err(|e| AssistantError::Internal(format!("upsert assistant preference: {e}")))?; @@ -1358,7 +1490,11 @@ impl AssistantService { } pub async fn delete(&self, id: &str) -> Result<(), AssistantError> { - match self.classify_source(id).await { + self.delete_for_user(DEFAULT_USER_ID, id).await + } + + pub async fn delete_for_user(&self, user_id: &str, id: &str) -> Result<(), AssistantError> { + match self.classify_source_for_user(user_id, id).await { AssistantSource::Builtin => { return Err(AssistantError::Forbidden("Cannot delete built-in assistant".into())); } @@ -1368,29 +1504,33 @@ impl AssistantService { AssistantSource::User => {} } - let removed = self.repo.delete(id).await?; + let removed = self.repo.delete_for_user(user_id, id).await?; if !removed { return Err(AssistantError::NotFound(format!("assistant '{id}' not found"))); } // Drop the override row (best-effort). - if let Err(e) = self.override_repo.delete(id).await { + if let Err(e) = self.override_repo.delete_for_user(user_id, id).await { warn!("failed to remove override for deleted assistant '{id}': {e}"); } - if let Some(definition) = self.definition_repo.get_by_assistant_id(id).await? { - if let Err(e) = self.state_repo.delete(&definition.id).await { + if let Some(definition) = self.definition_repo.get_by_assistant_id_for_user(user_id, id).await? { + if let Err(e) = self.state_repo.delete_for_user(user_id, &definition.id).await { warn!("failed to remove assistant overlay for deleted assistant '{id}': {e}"); } - if let Err(e) = self.preference_repo.delete(&definition.id).await { + if let Err(e) = self.preference_repo.delete_for_user(user_id, &definition.id).await { warn!("failed to remove assistant preferences for deleted assistant '{id}': {e}"); } - if let Err(e) = self.definition_repo.soft_delete(&definition.id, now_ms()).await { + if let Err(e) = self + .definition_repo + .soft_delete_for_user(user_id, &definition.id, now_ms()) + .await + { warn!("failed to soft-delete assistant definition for deleted assistant '{id}': {e}"); } } // Best-effort filesystem cleanup. - self.cleanup_user_assets(id); + self.cleanup_user_assets_for_user(user_id, id); Ok(()) } @@ -1400,11 +1540,20 @@ impl AssistantService { id: &str, req: SetAssistantStateRequest, ) -> Result { - match self.classify_source(id).await { + self.set_state_for_user(DEFAULT_USER_ID, id, req).await + } + + pub async fn set_state_for_user( + &self, + user_id: &str, + id: &str, + req: SetAssistantStateRequest, + ) -> Result { + match self.classify_source_for_user(user_id, id).await { AssistantSource::Builtin | AssistantSource::Generated => {} AssistantSource::User => { // Confirm the user row exists (otherwise 404). - if self.repo.get(id).await?.is_none() { + if self.repo.get_for_user(user_id, id).await?.is_none() { return Err(AssistantError::NotFound(format!("assistant '{id}' not found"))); } } @@ -1413,11 +1562,11 @@ impl AssistantService { // Merge with existing state/override to preserve fields not in this request. let definition = self .definition_repo - .get_by_assistant_id(id) + .get_by_assistant_id_for_user(user_id, id) .await? .ok_or_else(|| AssistantError::NotFound(format!("assistant '{id}' not found")))?; - let existing_state = self.state_repo.get(&definition.id).await?; - let existing = self.override_repo.get(id).await?; + let existing_state = self.state_repo.get_for_user(user_id, &definition.id).await?; + let existing = self.override_repo.get_for_user(user_id, id).await?; let enabled = req.enabled.unwrap_or_else(|| { existing_state .as_ref() @@ -1437,17 +1586,20 @@ impl AssistantService { .as_ref() .and_then(|state| state.agent_id_override.clone()); self.state_repo - .upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: &definition.id, - enabled, - sort_order, - agent_id_override: agent_id_override.as_deref(), - last_used_at, - }) + .upsert_for_user( + user_id, + &UpsertAssistantOverlayParams { + assistant_definition_id: &definition.id, + enabled, + sort_order, + agent_id_override: agent_id_override.as_deref(), + last_used_at, + }, + ) .await .map_err(|e| AssistantError::Internal(format!("upsert assistant overlay: {e}")))?; - self.get(id).await + self.get_for_user(user_id, id).await } // ----------------------------------------------------------------------- @@ -1458,6 +1610,14 @@ impl AssistantService { /// built-in id collision or already-imported user-id collision. /// Never overwrites an existing user row. pub async fn import(&self, req: ImportAssistantsRequest) -> Result { + self.import_for_user(DEFAULT_USER_ID, req).await + } + + pub async fn import_for_user( + &self, + user_id: &str, + req: ImportAssistantsRequest, + ) -> Result { let mut result = ImportAssistantsResult::default(); // Resolved-once cache for the inferred default agent id. We only @@ -1476,7 +1636,7 @@ impl AssistantService { result.skipped += 1; continue; } - match self.repo.get(&id).await { + match self.repo.get_for_user(user_id, &id).await { Ok(Some(_)) => { result.skipped += 1; continue; @@ -1528,7 +1688,7 @@ impl AssistantService { } _ => match cached_default_agent_id.as_deref() { Some(v) => v.to_string(), - None => match self.resolve_default_agent_id().await { + None => match self.resolve_default_agent_id_for_user(user_id).await { Ok(v) => { cached_default_agent_id = Some(v.clone()); v @@ -1544,7 +1704,10 @@ impl AssistantService { }, }, }; - if let Err(e) = self.resolve_runtime_backend_for_agent_id(&resolved_agent_id).await { + if let Err(e) = self + .resolve_runtime_backend_for_agent_id(user_id, &resolved_agent_id) + .await + { result.failed += 1; result.errors.push(ImportError { id, @@ -1580,9 +1743,9 @@ impl AssistantService { prompts_i18n: serialized.prompts_i18n.as_deref(), }; - match self.repo.create(¶ms).await { + match self.repo.create_for_user(user_id, ¶ms).await { Ok(row) => { - self.upsert_definition_from_legacy_user_row(&row, Some(&resolved_agent_id)) + self.upsert_definition_from_legacy_user_row_for_user(user_id, &row, Some(&resolved_agent_id)) .await?; result.imported += 1; } @@ -1610,9 +1773,21 @@ impl AssistantService { /// Read an assistant rule file, dispatching by source. pub async fn read_rule(&self, id: &str, locale: Option<&str>) -> Result { - match self.classify_source(id).await { + self.read_rule_for_user(DEFAULT_USER_ID, id, locale).await + } + + /// Read an assistant rule file for the current owner, dispatching by source. + pub async fn read_rule_for_user( + &self, + user_id: &str, + id: &str, + locale: Option<&str>, + ) -> Result { + match self.classify_source_for_user(user_id, id).await { AssistantSource::Builtin => Ok(self.read_builtin_rule_with_fallback(id, locale)), - AssistantSource::Generated | AssistantSource::User => Ok(self.read_user_rule_with_fallback(id, locale)), + AssistantSource::Generated | AssistantSource::User => { + Ok(self.read_user_rule_with_fallback(user_id, id, locale)) + } } } @@ -1645,8 +1820,8 @@ impl AssistantService { /// when the locale-specific `..md` is absent. Scheduled/cron runs /// create the conversation with `assistant: None`, so no UI locale reaches /// rule resolution and the localized file would otherwise be missed. - fn read_user_rule_with_fallback(&self, id: &str, locale: Option<&str>) -> String { - let rules_dir = self.user_rules_dir(); + fn read_user_rule_with_fallback(&self, user_id: &str, id: &str, locale: Option<&str>) -> String { + let rules_dir = self.user_rules_dir_for_user(user_id); let content = read_assistant_md_with_legacy(&rules_dir, id, locale); if !content.is_empty() { return content; @@ -1659,18 +1834,49 @@ impl AssistantService { } } - read_first_assistant_md(&rules_dir, id) + let content = read_first_assistant_md(&rules_dir, id); + if !content.is_empty() || user_id != DEFAULT_USER_ID { + return content; + } + + let legacy_rules_dir = self.user_rules_root_dir(); + let content = read_assistant_md_with_legacy(&legacy_rules_dir, id, locale); + if !content.is_empty() { + return content; + } + + // Mirror the scoped-dir fallback: a requested locale falls back to + // the locale-less file (migrating any legacy-named copy) before the + // glob-any last resort, which has no migration side effect. + if locale.is_some_and(|value| !value.is_empty()) { + let locale_less = read_assistant_md_with_legacy(&legacy_rules_dir, id, None); + if !locale_less.is_empty() { + return locale_less; + } + } + read_first_assistant_md(&legacy_rules_dir, id) } /// Write an assistant rule file. User-authored and generated assistants /// keep editable configuration in the local profile; built-ins reject. pub async fn write_rule(&self, id: &str, locale: Option<&str>, content: &str) -> Result<(), AssistantError> { - match self.classify_source(id).await { + self.write_rule_for_user(DEFAULT_USER_ID, id, locale, content).await + } + + /// Write an assistant rule file for the current owner. + pub async fn write_rule_for_user( + &self, + user_id: &str, + id: &str, + locale: Option<&str>, + content: &str, + ) -> Result<(), AssistantError> { + match self.classify_source_for_user(user_id, id).await { AssistantSource::Builtin => Err(AssistantError::BadRequest( "Cannot write rule for built-in assistant".into(), )), AssistantSource::Generated | AssistantSource::User => { - let path = self.user_rule_path(id, locale); + let path = self.user_rule_path_for_user(user_id, id, locale); if let Some(parent) = path.parent() { std::fs::create_dir_all(parent) .map_err(|e| AssistantError::Internal(format!("create dir failed: {e}")))?; @@ -1683,32 +1889,60 @@ impl AssistantService { /// Delete all locale versions of an assistant rule. pub async fn delete_rule(&self, id: &str) -> Result { - match self.classify_source(id).await { + self.delete_rule_for_user(DEFAULT_USER_ID, id).await + } + + /// Delete all locale versions of an assistant rule for the current owner. + pub async fn delete_rule_for_user(&self, user_id: &str, id: &str) -> Result { + match self.classify_source_for_user(user_id, id).await { AssistantSource::Builtin => Err(AssistantError::BadRequest( "Cannot delete rule for built-in assistant".into(), )), AssistantSource::Generated | AssistantSource::User => { - Ok(remove_assistant_md_files(&self.user_rules_dir(), id)) + Ok(remove_assistant_md_files(&self.user_rules_dir_for_user(user_id), id)) } } } pub async fn read_skill(&self, id: &str, locale: Option<&str>) -> Result { - match self.classify_source(id).await { + self.read_skill_for_user(DEFAULT_USER_ID, id, locale).await + } + + pub async fn read_skill_for_user( + &self, + user_id: &str, + id: &str, + locale: Option<&str>, + ) -> Result { + match self.classify_source_for_user(user_id, id).await { AssistantSource::Builtin => Ok(String::new()), AssistantSource::Generated | AssistantSource::User => { - Ok(read_assistant_md_with_legacy(&self.user_skills_dir(), id, locale)) + let content = read_assistant_md_with_legacy(&self.user_skills_dir_for_user(user_id), id, locale); + if !content.is_empty() || user_id != DEFAULT_USER_ID { + return Ok(content); + } + Ok(read_assistant_md_with_legacy(&self.user_skills_root_dir(), id, locale)) } } } pub async fn write_skill(&self, id: &str, locale: Option<&str>, content: &str) -> Result<(), AssistantError> { - match self.classify_source(id).await { + self.write_skill_for_user(DEFAULT_USER_ID, id, locale, content).await + } + + pub async fn write_skill_for_user( + &self, + user_id: &str, + id: &str, + locale: Option<&str>, + content: &str, + ) -> Result<(), AssistantError> { + match self.classify_source_for_user(user_id, id).await { AssistantSource::Builtin => Err(AssistantError::BadRequest( "Cannot write skill for built-in assistant".into(), )), AssistantSource::Generated | AssistantSource::User => { - let path = self.user_skill_path(id, locale); + let path = self.user_skill_path_for_user(user_id, id, locale); if let Some(parent) = path.parent() { std::fs::create_dir_all(parent) .map_err(|e| AssistantError::Internal(format!("create dir failed: {e}")))?; @@ -1720,12 +1954,16 @@ impl AssistantService { } pub async fn delete_skill(&self, id: &str) -> Result { - match self.classify_source(id).await { + self.delete_skill_for_user(DEFAULT_USER_ID, id).await + } + + pub async fn delete_skill_for_user(&self, user_id: &str, id: &str) -> Result { + match self.classify_source_for_user(user_id, id).await { AssistantSource::Builtin => Err(AssistantError::BadRequest( "Cannot delete skill for built-in assistant".into(), )), AssistantSource::Generated | AssistantSource::User => { - Ok(remove_assistant_md_files(&self.user_skills_dir(), id)) + Ok(remove_assistant_md_files(&self.user_skills_dir_for_user(user_id), id)) } } } @@ -1746,10 +1984,14 @@ impl AssistantService { /// has no on-disk file) also return `None`; clients fall back to the /// text avatar for those. pub async fn avatar_asset(&self, id: &str) -> Option { - match self.classify_source(id).await { + self.avatar_asset_for_user(DEFAULT_USER_ID, id).await + } + + pub async fn avatar_asset_for_user(&self, user_id: &str, id: &str) -> Option { + match self.classify_source_for_user(user_id, id).await { AssistantSource::Builtin => self.builtin.avatar_asset(id), AssistantSource::Generated | AssistantSource::User => { - if let Ok(Some(definition)) = self.definition_repo.get_by_assistant_id(id).await { + if let Ok(Some(definition)) = self.definition_repo.get_by_assistant_id_for_user(user_id, id).await { if definition.avatar_type != "user_asset" { return None; } @@ -1772,14 +2014,22 @@ impl AssistantService { // Internal helpers // ----------------------------------------------------------------------- - fn user_rules_dir(&self) -> PathBuf { + fn user_rules_root_dir(&self) -> PathBuf { self.user_data_dir.join("assistant-rules") } - fn user_skills_dir(&self) -> PathBuf { + fn user_rules_dir_for_user(&self, user_id: &str) -> PathBuf { + self.user_rules_root_dir().join(encode_filename_component(user_id)) + } + + fn user_skills_root_dir(&self) -> PathBuf { self.user_data_dir.join("assistant-skills") } + fn user_skills_dir_for_user(&self, user_id: &str) -> PathBuf { + self.user_skills_root_dir().join(encode_filename_component(user_id)) + } + fn user_avatars_dir(&self) -> PathBuf { self.user_data_dir.join("assistant-avatars") } @@ -2016,16 +2266,17 @@ impl AssistantService { self.read_user_avatar_asset_by_filename(value).is_some() } - fn user_rule_path(&self, id: &str, locale: Option<&str>) -> PathBuf { - assistant_md_path(&self.user_rules_dir(), id, locale) + fn user_rule_path_for_user(&self, user_id: &str, id: &str, locale: Option<&str>) -> PathBuf { + assistant_md_path(&self.user_rules_dir_for_user(user_id), id, locale) } - fn user_skill_path(&self, id: &str, locale: Option<&str>) -> PathBuf { - assistant_md_path(&self.user_skills_dir(), id, locale) + fn user_skill_path_for_user(&self, user_id: &str, id: &str, locale: Option<&str>) -> PathBuf { + assistant_md_path(&self.user_skills_dir_for_user(user_id), id, locale) } - async fn resolve_definition_identity( + async fn resolve_definition_identity_for_user( &self, + user_id: &str, source: &str, source_ref: Option<&str>, assistant_id: &str, @@ -2033,7 +2284,7 @@ impl AssistantService { if let Some(source_ref) = source_ref && let Some(existing) = self .definition_repo - .get_by_source_ref_including_deleted(source, source_ref) + .get_by_source_ref_including_deleted_for_user(user_id, source, source_ref) .await .map_err(|e| AssistantError::Internal(format!("get assistant definition by source_ref: {e}")))? { @@ -2042,7 +2293,7 @@ impl AssistantService { if let Some(existing) = self .definition_repo - .get_by_assistant_id_including_deleted(assistant_id) + .get_by_assistant_id_including_deleted_for_user(user_id, assistant_id) .await .map_err(|e| AssistantError::Internal(format!("get assistant definition by key: {e}")))? { @@ -2052,9 +2303,41 @@ impl AssistantService { Ok((generate_prefixed_id("asstdef"), assistant_id.to_string())) } - fn cleanup_user_assets(&self, id: &str) { - remove_assistant_md_files(&self.user_rules_dir(), id); - remove_assistant_md_files(&self.user_skills_dir(), id); + async fn resolve_global_definition_identity( + &self, + source: &str, + source_ref: Option<&str>, + assistant_id: &str, + ) -> Result<(String, String), AssistantError> { + if let Some(source_ref) = source_ref + && let Some(existing) = self + .definition_repo + .get_global_by_source_ref_including_deleted(source, source_ref) + .await + .map_err(|e| AssistantError::Internal(format!("get global assistant definition by source_ref: {e}")))? + { + return Ok((existing.id, existing.assistant_id)); + } + + if let Some(existing) = self + .definition_repo + .get_global_by_assistant_id_including_deleted(assistant_id) + .await + .map_err(|e| AssistantError::Internal(format!("get global assistant definition by key: {e}")))? + { + return Ok((existing.id, existing.assistant_id)); + } + + Ok((generate_prefixed_id("asstdef"), assistant_id.to_string())) + } + + fn cleanup_user_assets_for_user(&self, user_id: &str, id: &str) { + remove_assistant_md_files(&self.user_rules_dir_for_user(user_id), id); + remove_assistant_md_files(&self.user_skills_dir_for_user(user_id), id); + if user_id == DEFAULT_USER_ID { + remove_assistant_md_files(&self.user_rules_root_dir(), id); + remove_assistant_md_files(&self.user_skills_root_dir(), id); + } remove_assistant_avatar_files(&self.user_avatars_dir(), id); } } @@ -2068,38 +2351,50 @@ impl AssistantClassifier for AssistantService { #[async_trait::async_trait] impl AssistantRuleDispatcher for AssistantService { - async fn read_rule(&self, id: &str, locale: Option<&str>) -> Result { - AssistantService::read_rule(self, id, locale) + async fn read_rule(&self, user_id: &str, id: &str, locale: Option<&str>) -> Result { + AssistantService::read_rule_for_user(self, user_id, id, locale) .await .map_err(assistant_error_to_extension_error) } - async fn write_rule(&self, id: &str, locale: Option<&str>, content: &str) -> Result<(), ExtensionError> { - AssistantService::write_rule(self, id, locale, content) + async fn write_rule( + &self, + user_id: &str, + id: &str, + locale: Option<&str>, + content: &str, + ) -> Result<(), ExtensionError> { + AssistantService::write_rule_for_user(self, user_id, id, locale, content) .await .map_err(assistant_error_to_extension_error) } - async fn delete_rule(&self, id: &str) -> Result { - AssistantService::delete_rule(self, id) + async fn delete_rule(&self, user_id: &str, id: &str) -> Result { + AssistantService::delete_rule_for_user(self, user_id, id) .await .map_err(assistant_error_to_extension_error) } - async fn read_skill(&self, id: &str, locale: Option<&str>) -> Result { - AssistantService::read_skill(self, id, locale) + async fn read_skill(&self, user_id: &str, id: &str, locale: Option<&str>) -> Result { + AssistantService::read_skill_for_user(self, user_id, id, locale) .await .map_err(assistant_error_to_extension_error) } - async fn write_skill(&self, id: &str, locale: Option<&str>, content: &str) -> Result<(), ExtensionError> { - AssistantService::write_skill(self, id, locale, content) + async fn write_skill( + &self, + user_id: &str, + id: &str, + locale: Option<&str>, + content: &str, + ) -> Result<(), ExtensionError> { + AssistantService::write_skill_for_user(self, user_id, id, locale, content) .await .map_err(assistant_error_to_extension_error) } - async fn delete_skill(&self, id: &str) -> Result { - AssistantService::delete_skill(self, id) + async fn delete_skill(&self, user_id: &str, id: &str) -> Result { + AssistantService::delete_skill_for_user(self, user_id, id) .await .map_err(assistant_error_to_extension_error) } @@ -3053,7 +3348,7 @@ pub fn generate_user_id() -> String { mod tests { use super::*; use aionui_db::{ - CreateProviderParams, SqliteAssistantDefinitionRepository, SqliteAssistantOverlayRepository, + CreateProviderParams, IUserRepository, SqliteAssistantDefinitionRepository, SqliteAssistantOverlayRepository, SqliteAssistantOverrideRepository, SqliteAssistantPreferenceRepository, SqliteAssistantRepository, SqliteProviderRepository, UpsertOverrideParams, init_database_memory, }; @@ -3177,7 +3472,10 @@ mod tests { #[async_trait::async_trait] impl AssistantAgentCatalogPort for StubAgentCatalog { - async fn list_management_agents(&self) -> Result, AssistantError> { + async fn list_management_agents( + &self, + _user_id: &str, + ) -> Result, AssistantError> { Ok(self.rows.lock().expect("agent rows lock poisoned").clone()) } } @@ -3302,8 +3600,13 @@ mod tests { } async fn seed_provider(repo: &dyn IProviderRepository, platform: &str) { + seed_provider_for_user(repo, DEFAULT_USER_ID, platform).await; + } + + async fn seed_provider_for_user(repo: &dyn IProviderRepository, user_id: &str, platform: &str) { repo.create(CreateProviderParams { id: None, + user_id, platform, name: "Test Provider", base_url: "https://example.invalid", @@ -3323,6 +3626,13 @@ mod tests { .expect("seed provider"); } + async fn create_test_user(db: &aionui_db::Database, username: &str) -> String { + let repo = aionui_db::SqliteUserRepository::new(db.pool().clone()); + let user = repo.create_user(username, "hash").await.unwrap(); + assert_eq!(user.username.as_deref(), Some(username)); + user.id + } + fn mk_builtin(id: &str, name: &str) -> BuiltinAssistant { BuiltinAssistant { id: id.into(), @@ -3698,6 +4008,46 @@ mod tests { assert!(list.iter().any(|a| a.id == "u1")); } + #[tokio::test] + async fn list_for_user_isolates_user_authored_assistants() { + let fx = fixture_with_builtins(vec![mk_builtin("builtin-office", "Office")]).await; + let user_b = create_test_user(&fx._db, "assistant_user_b").await; + seed_provider_for_user(&*fx.provider_repo, &user_b, "openai").await; + + fx.service + .create_for_user( + DEFAULT_USER_ID, + CreateAssistantRequest { + id: Some("u-default".into()), + name: "Default User Assistant".into(), + ..req_default() + }, + ) + .await + .unwrap(); + fx.service + .create_for_user( + &user_b, + CreateAssistantRequest { + id: Some("u-b".into()), + name: "User B Assistant".into(), + ..req_default() + }, + ) + .await + .unwrap(); + + let default_list = fx.service.list_for_user(DEFAULT_USER_ID).await.unwrap(); + let user_b_list = fx.service.list_for_user(&user_b).await.unwrap(); + + assert!(default_list.iter().any(|assistant| assistant.id == "builtin-office")); + assert!(user_b_list.iter().any(|assistant| assistant.id == "builtin-office")); + assert!(default_list.iter().any(|assistant| assistant.id == "u-default")); + assert!(!default_list.iter().any(|assistant| assistant.id == "u-b")); + assert!(user_b_list.iter().any(|assistant| assistant.id == "u-b")); + assert!(!user_b_list.iter().any(|assistant| assistant.id == "u-default")); + } + #[tokio::test] async fn builtin_listing_uses_manifest_default_enabled_and_sort_order() { // A builtin with default_enabled=false + sort_order=50, and no user @@ -5823,6 +6173,57 @@ mod tests { assert_eq!(content, "rule body"); } + #[tokio::test] + async fn rule_and_skill_files_are_scoped_by_user() { + let fx = fixture().await; + + fx.service + .write_rule_for_user("user-a", "shared-assistant", Some("en-US"), "rule A") + .await + .unwrap(); + fx.service + .write_rule_for_user("user-b", "shared-assistant", Some("en-US"), "rule B") + .await + .unwrap(); + fx.service + .write_skill_for_user("user-a", "shared-assistant", None, "skill A") + .await + .unwrap(); + fx.service + .write_skill_for_user("user-b", "shared-assistant", None, "skill B") + .await + .unwrap(); + + assert_eq!( + fx.service + .read_rule_for_user("user-a", "shared-assistant", Some("en-US")) + .await + .unwrap(), + "rule A" + ); + assert_eq!( + fx.service + .read_rule_for_user("user-b", "shared-assistant", Some("en-US")) + .await + .unwrap(), + "rule B" + ); + assert_eq!( + fx.service + .read_skill_for_user("user-a", "shared-assistant", None) + .await + .unwrap(), + "skill A" + ); + assert_eq!( + fx.service + .read_skill_for_user("user-b", "shared-assistant", None) + .await + .unwrap(), + "skill B" + ); + } + #[tokio::test] async fn read_rule_user_falls_back_to_saved_locale_when_locale_missing() { // Scheduled/cron runs resolve rules without a locale (conversation is @@ -5900,7 +6301,7 @@ mod tests { assert!( fx._tmp .path() - .join("assistant-rules/bare%3Aagent-claude.en-US.md") + .join("assistant-rules/system_default_user/bare%3Aagent-claude.en-US.md") .is_file() ); assert!( @@ -6196,6 +6597,7 @@ mod tests { fx.provider_repo .create(CreateProviderParams { id: None, + user_id: DEFAULT_USER_ID, platform: "anthropic", name: "Disabled", base_url: "https://example.invalid", diff --git a/crates/aionui-auth/src/csrf.rs b/crates/aionui-auth/src/csrf.rs index 27e9f54a8..35b5504bf 100644 --- a/crates/aionui-auth/src/csrf.rs +++ b/crates/aionui-auth/src/csrf.rs @@ -18,7 +18,7 @@ use crate::extract::extract_cookie_value; /// /// Behavior: /// - Safe methods (GET, HEAD, OPTIONS) bypass validation. -/// - Exempt paths (`/login`, `/api/auth/qr-login`) bypass validation. +/// - Exempt paths (`/login`, `/api/auth/qr-login`, Bootstrap-secret internal auth) bypass validation. /// - All other requests must include an `x-csrf-token` header whose value /// matches the `aionui-csrf-token` cookie. /// - Sets the CSRF cookie on responses if the client does not have one. @@ -35,7 +35,10 @@ pub async fn csrf_middleware( // Validate CSRF for state-changing requests let needs_validation = matches!(method, Method::POST | Method::PUT | Method::DELETE | Method::PATCH); - let is_exempt = path == "/login" || path == "/api/auth/qr-login"; + let is_exempt = path == "/login" + || path == "/api/auth/qr-login" + || path.starts_with("/api/auth/internal/external-users/") + || path.starts_with("/api/auth/internal/external-sessions"); if needs_validation && !is_exempt { let header_token = request diff --git a/crates/aionui-auth/src/jwt.rs b/crates/aionui-auth/src/jwt.rs index ef2c19b8c..4ab821571 100644 --- a/crates/aionui-auth/src/jwt.rs +++ b/crates/aionui-auth/src/jwt.rs @@ -34,6 +34,9 @@ pub struct TokenPayload { pub iss: String, /// Audience (standard JWT claim). pub aud: String, + /// User session generation at token issuance time. + #[serde(default)] + pub session_generation: i64, } /// JWT service for signing, verification, and token blacklisting. @@ -59,6 +62,16 @@ impl JwtService { /// Sign a new JWT for the given user. The token expires after 24 hours. pub fn sign(&self, user_id: &str, username: &str) -> Result { + self.sign_with_session_generation(user_id, username, 0) + } + + /// Sign a new JWT bound to the user's current session generation. + pub fn sign_with_session_generation( + &self, + user_id: &str, + username: &str, + session_generation: i64, + ) -> Result { let now = now_secs()?; let exp = now + TOKEN_EXPIRY.as_secs(); @@ -69,6 +82,7 @@ impl JwtService { exp, iss: JWT_ISSUER.to_owned(), aud: JWT_AUDIENCE.to_owned(), + session_generation, }; let secret = self @@ -228,11 +242,21 @@ mod tests { let payload = service.verify(&token).unwrap(); assert_eq!(payload.user_id, "user_1"); assert_eq!(payload.username, "admin"); + assert_eq!(payload.session_generation, 0); assert_eq!(payload.iss, JWT_ISSUER); assert_eq!(payload.aud, JWT_AUDIENCE); assert!(payload.exp > payload.iat); } + #[test] + fn sign_with_session_generation_roundtrip() { + let service = test_service(); + let token = service.sign_with_session_generation("user_1", "admin", 7).unwrap(); + let payload = service.verify(&token).unwrap(); + assert_eq!(payload.user_id, "user_1"); + assert_eq!(payload.session_generation, 7); + } + #[test] fn verify_tampered_token_fails() { let service = test_service(); @@ -261,6 +285,7 @@ mod tests { exp: 1001, iss: JWT_ISSUER.into(), aud: JWT_AUDIENCE.into(), + session_generation: 0, }; let token = encode( &Header::default(), @@ -342,6 +367,7 @@ mod tests { exp: 1001, iss: JWT_ISSUER.into(), aud: JWT_AUDIENCE.into(), + session_generation: 0, }; let token = encode( &Header::default(), diff --git a/crates/aionui-auth/src/lib.rs b/crates/aionui-auth/src/lib.rs index fdefc280d..8d9c410c4 100644 --- a/crates/aionui-auth/src/lib.rs +++ b/crates/aionui-auth/src/lib.rs @@ -12,6 +12,7 @@ pub mod qr_token; mod rate_limit; mod routes; mod security; +mod service; mod validation; // Error type @@ -50,10 +51,12 @@ pub use security::security_headers_middleware; pub use csrf::csrf_middleware; // Auth middleware -pub use middleware::{AuthState, CurrentUser, auth_middleware, local_auth_middleware}; +pub use middleware::{AuthIdentityMode, AuthState, CurrentUser, auth_middleware, local_auth_middleware}; // QR token store pub use qr_token::QrTokenStore; // Routes -pub use routes::{AuthRouterState, auth_routes}; +pub use routes::{AuthRouterState, SessionRevokedHook, auth_routes}; + +pub use service::{AuthProvisionService, ProvisionError}; diff --git a/crates/aionui-auth/src/middleware.rs b/crates/aionui-auth/src/middleware.rs index d9b8520a5..05ccaa89e 100644 --- a/crates/aionui-auth/src/middleware.rs +++ b/crates/aionui-auth/src/middleware.rs @@ -3,15 +3,23 @@ use std::sync::Arc; use axum::extract::{Request, State}; +use axum::http::StatusCode; use axum::middleware::Next; use axum::response::Response; use aionui_common::ApiError; -use aionui_db::IUserRepository; +use aionui_db::{IUserRepository, UserStatus, UserType}; use crate::JwtService; use crate::extract::extract_token_from_headers; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AuthIdentityMode { + Local, + UserSession, + AionPro, +} + /// Authenticated user injected into request extensions by the auth middleware. /// /// Route handlers extract this from `request.extensions()` to identify @@ -22,6 +30,21 @@ pub struct CurrentUser { pub id: String, /// Username. pub username: String, + /// Internal identity source for the current user. + pub user_type: UserType, + /// Current account status. Authenticated requests only receive active users. + pub status: UserStatus, +} + +impl CurrentUser { + pub fn local_default() -> Self { + Self { + id: "system_default_user".to_string(), + username: "system_default_user".to_string(), + user_type: UserType::Local, + status: UserStatus::Active, + } + } } /// Shared state for the authentication middleware. @@ -29,8 +52,7 @@ pub struct CurrentUser { pub struct AuthState { pub jwt_service: Arc, pub user_repo: Arc, - /// When `true`, skip JWT verification and inject a fixed default user. - pub local: bool, + pub identity_mode: AuthIdentityMode, } /// Authentication middleware that verifies JWT tokens and injects `CurrentUser`. @@ -50,11 +72,8 @@ pub async fn auth_middleware( next: Next, ) -> Result { // In local mode, skip JWT verification and inject a fixed default user. - if state.local { - request.extensions_mut().insert(CurrentUser { - id: "system_default_user".to_string(), - username: "system_default_user".to_string(), - }); + if state.identity_mode == AuthIdentityMode::Local { + request.extensions_mut().insert(CurrentUser::local_default()); return Ok(next.run(request).await); } @@ -68,7 +87,7 @@ pub async fn auth_middleware( let user = state .user_repo - .find_by_id(&payload.user_id) + .find_active_by_id(&payload.user_id) .await .map_err(|e| { tracing::error!(error = %e, "auth middleware user lookup failed"); @@ -76,9 +95,24 @@ pub async fn auth_middleware( })? .ok_or_else(|| ApiError::Unauthorized("Invalid authentication subject".into()))?; + if state.identity_mode == AuthIdentityMode::AionPro && user.user_type != UserType::Aionpro { + return Err(ApiError::coded( + StatusCode::UNAUTHORIZED, + "USER_CONTEXT_REQUIRED", + "User context required.", + None, + )); + } + + if payload.session_generation != user.session_generation { + return Err(ApiError::Unauthorized("Invalid authentication session".into())); + } + request.extensions_mut().insert(CurrentUser { id: user.id, - username: user.username, + username: user.username.unwrap_or_else(|| "external_user".to_string()), + user_type: user.user_type, + status: user.status, }); Ok(next.run(request).await) @@ -89,10 +123,7 @@ pub async fn auth_middleware( /// Injects a fixed `CurrentUser` with id and username `system_default_user`. /// Used when the server runs as an embedded subprocess inside Electron. pub async fn local_auth_middleware(mut request: Request, next: Next) -> Response { - request.extensions_mut().insert(CurrentUser { - id: "system_default_user".to_string(), - username: "system_default_user".to_string(), - }); + request.extensions_mut().insert(CurrentUser::local_default()); next.run(request).await } diff --git a/crates/aionui-auth/src/routes.rs b/crates/aionui-auth/src/routes.rs index 670e5b545..57bc03c21 100644 --- a/crates/aionui-auth/src/routes.rs +++ b/crates/aionui-auth/src/routes.rs @@ -5,33 +5,40 @@ use std::time::Duration; use axum::extract::rejection::JsonRejection; use axum::extract::{Json, Path, State}; -use axum::http::{HeaderMap, header}; +use axum::http::{HeaderMap, StatusCode, header}; use axum::middleware::from_fn_with_state; use axum::response::{Html, IntoResponse, Response}; -use axum::routing::{get, post}; +use axum::routing::{get, post, put}; use axum::{Extension, Router}; -use serde::Deserialize; +use serde::{Deserialize, Serialize}; use aionui_api_types::{ - ApiResponse, AuthStatusResponse, ChangePasswordRequest, LoginRequest, LoginResponse, PublicUser, QrLoginRequest, - RefreshResponse, RefreshTokenRequest, UserInfoResponse, WebuiChangePasswordRequest, WebuiChangeUsernameRequest, - WebuiChangeUsernameResponse, WebuiGenerateQrTokenResponse, WebuiResetPasswordResponse, WsTokenResponse, + ApiResponse, AuthStatusResponse, ChangePasswordRequest, EnsureExternalSessionRequest, EnsureExternalUserRequest, + EnsureExternalUserResponse, LoginRequest, LoginResponse, PublicUser, QrLoginRequest, RefreshResponse, + RefreshTokenRequest, RevokeExternalSessionRequest, RevokeExternalSessionResponse, UserInfoResponse, + WebuiChangePasswordRequest, WebuiChangeUsernameRequest, WebuiChangeUsernameResponse, WebuiGenerateQrTokenResponse, + WebuiResetPasswordResponse, WsTokenResponse, }; use aionui_common::ApiError; use aionui_common::constants::COOKIE_MAX_AGE_DAYS; -use aionui_db::{DbError, IUserRepository, models::User}; +use aionui_db::{DbError, IUserRepository, UserStatus, UserType, models::User}; use crate::error::AuthError; use crate::extract::extract_token_from_headers; -use crate::middleware::{AuthState, CurrentUser, auth_middleware}; +use crate::middleware::{AuthIdentityMode, AuthState, CurrentUser, auth_middleware}; use crate::password::{dummy_password_hash, generate_password, hash_password, verify_password_timed}; use crate::qr_token::QrTokenStore; use crate::rate_limit::{ RateLimiter, api_rate_limit_middleware, auth_rate_limit_middleware, authenticated_action_rate_limit_middleware, }; +use crate::service::{AuthProvisionService, ProvisionError}; use crate::validation::{validate_password, validate_username}; use crate::{CookieConfig, JwtService}; +const BOOTSTRAP_SECRET_HEADER: &str = "x-aioncore-bootstrap-secret"; + +pub type SessionRevokedHook = dyn Fn(&str) + Send + Sync; + impl From for ApiError { fn from(err: AuthError) -> Self { match err { @@ -64,7 +71,11 @@ pub struct AuthRouterState { pub user_repo: Arc, pub cookie_config: Arc, pub qr_token_store: Arc, + pub identity_mode: AuthIdentityMode, + pub bootstrap_secret: Option>, + pub session_revoked_hook: Option>, pub local: bool, + pub aionpro_mode: bool, } #[derive(Debug, Deserialize)] @@ -94,6 +105,39 @@ struct UpdateJwtSecretRequest { jwt_secret: String, } +#[derive(Debug, Serialize)] +struct InternalUserResponse { + id: String, + user_type: UserType, + external_user_id: Option, + username: Option, + email: Option, + avatar_path: Option, + status: UserStatus, + session_generation: i64, + created_at: i64, + updated_at: i64, + last_login: Option, +} + +impl From for InternalUserResponse { + fn from(user: User) -> Self { + Self { + id: user.id, + user_type: user.user_type, + external_user_id: user.external_user_id, + username: user.username, + email: user.email, + avatar_path: user.avatar_path, + status: user.status, + session_generation: user.session_generation, + created_at: user.created_at, + updated_at: user.updated_at, + last_login: user.last_login, + } + } +} + fn ensure_local_mode(local: bool) -> Result<(), ApiError> { if local { return Ok(()); @@ -103,6 +147,76 @@ fn ensure_local_mode(local: bool) -> Result<(), ApiError> { )) } +fn require_bootstrap_secret(headers: &HeaderMap, expected: Option<&str>) -> Result<(), ApiError> { + let Some(expected) = expected else { + return Err(ApiError::coded( + StatusCode::UNAUTHORIZED, + "BOOTSTRAP_SECRET_REQUIRED", + "Bootstrap secret required.", + None, + )); + }; + let Some(actual) = headers.get(BOOTSTRAP_SECRET_HEADER).and_then(|v| v.to_str().ok()) else { + return Err(ApiError::coded( + StatusCode::UNAUTHORIZED, + "BOOTSTRAP_SECRET_REQUIRED", + "Bootstrap secret required.", + None, + )); + }; + if constant_time_eq(actual.as_bytes(), expected.as_bytes()) { + Ok(()) + } else { + Err(ApiError::coded( + StatusCode::UNAUTHORIZED, + "INVALID_BOOTSTRAP_SECRET", + "Invalid bootstrap secret.", + None, + )) + } +} + +fn constant_time_eq(left: &[u8], right: &[u8]) -> bool { + let max_len = left.len().max(right.len()); + let mut diff = left.len() ^ right.len(); + for idx in 0..max_len { + let l = left.get(idx).copied().unwrap_or(0); + let r = right.get(idx).copied().unwrap_or(0); + diff |= usize::from(l ^ r); + } + diff == 0 +} + +fn provision_error_to_api_error(err: ProvisionError) -> ApiError { + match err { + ProvisionError::UnsupportedUserType => ApiError::BadRequest("Unsupported external user type".into()), + ProvisionError::UserDisabled => ApiError::coded(StatusCode::FORBIDDEN, "USER_DISABLED", "User disabled.", None), + ProvisionError::UserNotProvisioned => ApiError::coded( + StatusCode::UNAUTHORIZED, + "USER_CONTEXT_REQUIRED", + "User context required.", + None, + ), + ProvisionError::Db(DbError::Conflict(_)) => ApiError::coded( + StatusCode::CONFLICT, + "EXTERNAL_USER_CONFLICT", + "External user conflict.", + None, + ), + ProvisionError::Db(e) => db_error_to_api_error(e), + ProvisionError::Token(e) => ApiError::Internal(format!("Token signing error: {e}")), + } +} + +fn user_context_required() -> ApiError { + ApiError::coded( + StatusCode::UNAUTHORIZED, + "USER_CONTEXT_REQUIRED", + "User context required.", + None, + ) +} + /// Build the auth router with all endpoints and middleware layers. /// /// Returns a `Router` with these endpoints: @@ -133,7 +247,11 @@ pub fn auth_routes(state: AuthRouterState) -> Router { let auth_state = AuthState { jwt_service: state.jwt_service.clone(), user_repo: state.user_repo.clone(), - local: false, + identity_mode: if state.aionpro_mode { + AuthIdentityMode::AionPro + } else { + AuthIdentityMode::UserSession + }, }; // Auth rate limited routes (login, qr-login) @@ -146,6 +264,18 @@ pub fn auth_routes(state: AuthRouterState) -> Router { // API rate limited public routes (no auth required) let api_public = Router::new() .route("/api/auth/status", get(status_handler)) + .route( + "/api/auth/internal/external-users/{external_user_id}", + put(ensure_external_user_handler), + ) + .route( + "/api/auth/internal/external-sessions", + post(create_external_session_handler), + ) + .route( + "/api/auth/internal/external-sessions/revoke", + post(revoke_external_session_handler), + ) .route( "/api/auth/internal/users", get(list_internal_users_handler).post(create_internal_user_handler), @@ -220,6 +350,82 @@ pub fn auth_routes(state: AuthRouterState) -> Router { .merge(static_routes) } +// --------------------------------------------------------------------------- +// PUT /api/auth/internal/external-users/{external_user_id} +// --------------------------------------------------------------------------- + +async fn ensure_external_user_handler( + State(state): State, + headers: HeaderMap, + Path(external_user_id): Path, + body: Result, JsonRejection>, +) -> Result>, ApiError> { + require_bootstrap_secret(&headers, state.bootstrap_secret.as_deref().map(AsRef::as_ref))?; + let Json(req) = body.map_err(ApiError::from)?; + let service = AuthProvisionService::new(state.user_repo, state.jwt_service); + let response = service + .ensure_external_user(&external_user_id, req) + .await + .map_err(provision_error_to_api_error)?; + tracing::info!( + user_id = %response.user_id, + user_type = ?response.user_type, + "external user provision succeeded" + ); + Ok(Json(ApiResponse::ok(response))) +} + +// --------------------------------------------------------------------------- +// POST /api/auth/internal/external-sessions +// --------------------------------------------------------------------------- + +async fn create_external_session_handler( + State(state): State, + headers: HeaderMap, + body: Result, JsonRejection>, +) -> Result { + require_bootstrap_secret(&headers, state.bootstrap_secret.as_deref().map(AsRef::as_ref))?; + let Json(req) = body.map_err(ApiError::from)?; + let service = AuthProvisionService::new(state.user_repo, state.jwt_service); + let exchange = service + .create_external_session(req) + .await + .map_err(provision_error_to_api_error)?; + tracing::info!( + user_id = %exchange.response.user.id, + "external core session exchange succeeded" + ); + let cookie = state.cookie_config.build_session_cookie(&exchange.token); + Ok(([(header::SET_COOKIE, cookie)], Json(ApiResponse::ok(exchange.response))).into_response()) +} + +// --------------------------------------------------------------------------- +// POST /api/auth/internal/external-sessions/revoke +// --------------------------------------------------------------------------- + +async fn revoke_external_session_handler( + State(state): State, + headers: HeaderMap, + body: Result, JsonRejection>, +) -> Result>, ApiError> { + require_bootstrap_secret(&headers, state.bootstrap_secret.as_deref().map(AsRef::as_ref))?; + let Json(req) = body.map_err(ApiError::from)?; + let service = AuthProvisionService::new(state.user_repo, state.jwt_service); + let response = service + .revoke_external_session(req) + .await + .map_err(provision_error_to_api_error)?; + tracing::info!( + user_id = %response.user_id, + session_generation = response.session_generation, + "external core session revoked" + ); + if let Some(hook) = &state.session_revoked_hook { + hook(&response.user_id); + } + Ok(Json(ApiResponse::ok(response))) +} + // --------------------------------------------------------------------------- // POST /login // --------------------------------------------------------------------------- @@ -228,6 +434,10 @@ async fn login_handler( State(state): State, body: Result, JsonRejection>, ) -> Result { + if state.aionpro_mode { + return Err(user_context_required()); + } + let Json(req) = body.map_err(ApiError::from)?; // Input length validation (per API spec) @@ -246,7 +456,7 @@ async fn login_handler( .map_err(|e| ApiError::Internal(format!("Database error: {e}")))?; let (found_user, password_valid) = match user { - Some(u) if u.password_hash.trim().is_empty() => { + Some(u) if u.password_hash.as_deref().unwrap_or_default().trim().is_empty() => { // Seeded user with no password yet (first-run local mode). // Treat as invalid credentials; run dummy verify for timing symmetry // and to avoid bcrypt error on empty hash leaking as a 500. @@ -254,7 +464,11 @@ async fn login_handler( (None, false) } Some(u) => { - let valid = verify_password_timed(&req.password, &u.password_hash).await?; + let Some(password_hash) = u.password_hash.as_deref() else { + let _ = verify_password_timed(&req.password, dummy_password_hash()).await; + return Err(ApiError::Unauthorized("Invalid username or password".into())); + }; + let valid = verify_password_timed(&req.password, password_hash).await?; (Some(u), valid) } None => { @@ -272,7 +486,11 @@ async fn login_handler( let token = state .jwt_service - .sign(&user.id, &user.username) + .sign_with_session_generation( + &user.id, + user.username.as_deref().unwrap_or("external_user"), + user.session_generation, + ) .map_err(|e| ApiError::Internal(format!("Token signing error: {e}")))?; // Update last login (best-effort) @@ -284,7 +502,7 @@ async fn login_handler( let resp = LoginResponse::new( PublicUser { id: user.id, - username: user.username, + username: user.username.unwrap_or_else(|| "external_user".to_string()), }, token, ); @@ -346,46 +564,48 @@ async fn status_handler( async fn list_internal_users_handler( State(state): State, -) -> Result>>, ApiError> { +) -> Result>>, ApiError> { ensure_local_mode(state.local)?; let users = state.user_repo.list_users().await.map_err(db_error_to_api_error)?; - Ok(Json(ApiResponse::ok(users))) + Ok(Json(ApiResponse::ok( + users.into_iter().map(InternalUserResponse::from).collect(), + ))) } async fn get_system_user_handler( State(state): State, -) -> Result>>, ApiError> { +) -> Result>>, ApiError> { ensure_local_mode(state.local)?; let user = state.user_repo.get_system_user().await.map_err(db_error_to_api_error)?; - Ok(Json(ApiResponse::ok(user))) + Ok(Json(ApiResponse::ok(user.map(InternalUserResponse::from)))) } async fn find_user_by_username_handler( State(state): State, Path(username): Path, -) -> Result>>, ApiError> { +) -> Result>>, ApiError> { ensure_local_mode(state.local)?; let user = state .user_repo .find_by_username(&username) .await .map_err(db_error_to_api_error)?; - Ok(Json(ApiResponse::ok(user))) + Ok(Json(ApiResponse::ok(user.map(InternalUserResponse::from)))) } async fn find_user_by_id_handler( State(state): State, Path(id): Path, -) -> Result>>, ApiError> { +) -> Result>>, ApiError> { ensure_local_mode(state.local)?; let user = state.user_repo.find_by_id(&id).await.map_err(db_error_to_api_error)?; - Ok(Json(ApiResponse::ok(user))) + Ok(Json(ApiResponse::ok(user.map(InternalUserResponse::from)))) } async fn create_internal_user_handler( State(state): State, body: Result, JsonRejection>, -) -> Result>, ApiError> { +) -> Result>, ApiError> { ensure_local_mode(state.local)?; let Json(req) = body.map_err(ApiError::from)?; let user = state @@ -393,7 +613,7 @@ async fn create_internal_user_handler( .create_user(&req.username, &req.password_hash) .await .map_err(db_error_to_api_error)?; - Ok(Json(ApiResponse::ok(user))) + Ok(Json(ApiResponse::ok(InternalUserResponse::from(user)))) } async fn set_system_user_credentials_handler( @@ -505,7 +725,10 @@ async fn change_password_handler( .ok_or_else(|| ApiError::NotFound("User not found".into()))?; // Verify current password - let valid = verify_password_timed(&req.current_password, &user.password_hash).await?; + let Some(password_hash) = user.password_hash.as_deref() else { + return Err(ApiError::Unauthorized("Current password is incorrect".into())); + }; + let valid = verify_password_timed(&req.current_password, password_hash).await?; if !valid { return Err(ApiError::Unauthorized("Current password is incorrect".into())); } @@ -554,9 +777,31 @@ async fn refresh_handler( .verify(&req.token) .map_err(|_| ApiError::Unauthorized("Invalid or expired token".into()))?; + let user = state + .user_repo + .find_active_by_id(&payload.user_id) + .await + .map_err(|e| { + tracing::error!(error = %e, "refresh token user lookup failed"); + ApiError::Internal("Authentication service unavailable".into()) + })? + .ok_or_else(|| ApiError::Unauthorized("Invalid authentication subject".into()))?; + + if state.aionpro_mode && user.user_type != aionui_db::UserType::Aionpro { + return Err(user_context_required()); + } + + if payload.session_generation != user.session_generation { + return Err(ApiError::Unauthorized("Invalid authentication session".into())); + } + let new_token = state .jwt_service - .sign(&payload.user_id, &payload.username) + .sign_with_session_generation( + &user.id, + user.username.as_deref().unwrap_or("external_user"), + user.session_generation, + ) .map_err(|e| ApiError::Internal(format!("Token signing error: {e}")))?; Ok(Json(RefreshResponse { @@ -603,6 +848,10 @@ async fn qr_login_handler( State(state): State, body: Result, JsonRejection>, ) -> Result { + if state.aionpro_mode { + return Err(user_context_required()); + } + let Json(req) = body.map_err(ApiError::from)?; // Validate and consume QR token (one-time use) @@ -618,7 +867,11 @@ async fn qr_login_handler( let token = state .jwt_service - .sign(&user.id, &user.username) + .sign_with_session_generation( + &user.id, + user.username.as_deref().unwrap_or("external_user"), + user.session_generation, + ) .map_err(|e| ApiError::Internal(format!("Token signing error: {e}")))?; // Update last login (best-effort) @@ -630,7 +883,7 @@ async fn qr_login_handler( let resp = LoginResponse::new( PublicUser { id: user.id, - username: user.username, + username: user.username.unwrap_or_else(|| "external_user".to_string()), }, token, ); @@ -763,7 +1016,7 @@ async fn webui_change_username_handler( let user = resolve_webui_admin(&*state.user_repo).await?; - if user.username != trimmed { + if user.username.as_deref() != Some(trimmed.as_str()) { state .user_repo .update_username(&user.id, &trimmed) diff --git a/crates/aionui-auth/src/service.rs b/crates/aionui-auth/src/service.rs new file mode 100644 index 000000000..a47904f76 --- /dev/null +++ b/crates/aionui-auth/src/service.rs @@ -0,0 +1,137 @@ +use std::sync::Arc; + +use aionui_api_types::{ + EnsureExternalSessionRequest, EnsureExternalSessionResponse, EnsureExternalUserRequest, EnsureExternalUserResponse, + ExternalUserType, PublicUser, RevokeExternalSessionRequest, RevokeExternalSessionResponse, +}; +use aionui_db::{ExternalUserProjection, IUserRepository, UserStatus, UserType, models::User}; + +use crate::JwtService; + +#[derive(Debug, thiserror::Error)] +pub enum ProvisionError { + #[error("unsupported external user type")] + UnsupportedUserType, + #[error("external user is disabled")] + UserDisabled, + #[error("external user is not provisioned")] + UserNotProvisioned, + #[error("database error: {0}")] + Db(#[from] aionui_db::DbError), + #[error("token signing failed: {0}")] + Token(#[from] crate::AuthError), +} + +#[derive(Clone)] +pub struct AuthProvisionService { + user_repo: Arc, + jwt_service: Arc, +} + +pub struct ExternalSessionExchange { + pub response: EnsureExternalSessionResponse, + pub token: String, +} + +impl AuthProvisionService { + pub fn new(user_repo: Arc, jwt_service: Arc) -> Self { + Self { user_repo, jwt_service } + } + + pub async fn ensure_external_user( + &self, + external_user_id: &str, + request: EnsureExternalUserRequest, + ) -> Result { + let user_type = map_external_user_type(request.user_type)?; + let user = self + .user_repo + .ensure_external_user( + user_type, + external_user_id, + ExternalUserProjection { + username: request.username, + email: request.email, + avatar_path: request.avatar_path, + }, + ) + .await?; + + if user.status == UserStatus::Disabled { + return Err(ProvisionError::UserDisabled); + } + + // Upgrade path (AionUi -> AionPro over the same database): while this + // is the machine's only external user, re-own the local default + // user's pre-multi-account data to it. Self-idempotent — moves + // nothing once the source set is empty or a second account exists. + let adopted = self.user_repo.adopt_system_default_data(&user.id).await?; + if adopted > 0 { + tracing::info!(user_id = %user.id, rows = adopted, "adopted system_default_user data into first external user"); + } + + Ok(external_user_response(user, request.user_type)) + } + + pub async fn create_external_session( + &self, + request: EnsureExternalSessionRequest, + ) -> Result { + let user_type = map_external_user_type(request.user_type)?; + let user = self + .user_repo + .find_by_external_user_id(user_type, &request.external_user_id) + .await? + .ok_or(ProvisionError::UserNotProvisioned)?; + + if user.status == UserStatus::Disabled { + return Err(ProvisionError::UserDisabled); + } + + let username = user.username.clone().unwrap_or_else(|| "external_user".to_string()); + let token = self + .jwt_service + .sign_with_session_generation(&user.id, &username, user.session_generation)?; + self.user_repo.update_last_login(&user.id).await?; + + Ok(ExternalSessionExchange { + response: EnsureExternalSessionResponse { + user: PublicUser { id: user.id, username }, + session_generation: user.session_generation, + }, + token, + }) + } + + pub async fn revoke_external_session( + &self, + request: RevokeExternalSessionRequest, + ) -> Result { + let user_type = map_external_user_type(request.user_type)?; + let user = self + .user_repo + .find_by_external_user_id(user_type, &request.external_user_id) + .await? + .ok_or(ProvisionError::UserNotProvisioned)?; + let session_generation = self.user_repo.increment_session_generation(&user.id).await?; + Ok(RevokeExternalSessionResponse { + user_id: user.id, + session_generation, + }) + } +} + +fn map_external_user_type(user_type: ExternalUserType) -> Result { + match user_type { + ExternalUserType::Aionpro => Ok(UserType::Aionpro), + } +} + +fn external_user_response(user: User, user_type: ExternalUserType) -> EnsureExternalUserResponse { + EnsureExternalUserResponse { + user_id: user.id, + user_type, + external_user_id: user.external_user_id.unwrap_or_default(), + session_generation: user.session_generation, + } +} diff --git a/crates/aionui-auth/tests/middleware_tests.rs b/crates/aionui-auth/tests/middleware_tests.rs index c4e83faf4..20bfab36a 100644 --- a/crates/aionui-auth/tests/middleware_tests.rs +++ b/crates/aionui-auth/tests/middleware_tests.rs @@ -10,11 +10,11 @@ use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation, decode, encode} use tower::ServiceExt; use aionui_auth::{ - AuthState, CookieConfig, CurrentUser, JwtService, RateLimiter, TokenPayload, api_rate_limit_middleware, - auth_middleware, auth_rate_limit_middleware, authenticated_action_rate_limit_middleware, csrf_middleware, - security_headers_middleware, + AuthIdentityMode, AuthState, CookieConfig, CurrentUser, JwtService, RateLimiter, TokenPayload, + api_rate_limit_middleware, auth_middleware, auth_rate_limit_middleware, authenticated_action_rate_limit_middleware, + csrf_middleware, security_headers_middleware, }; -use aionui_db::{IUserRepository, SqliteUserRepository, init_database_memory}; +use aionui_db::{IUserRepository, SqliteUserRepository, UserStatus, UserType, init_database_memory}; async fn json_body(resp: axum::response::Response) -> serde_json::Value { let body = axum::body::to_bytes(resp.into_body(), usize::MAX).await.unwrap(); @@ -131,10 +131,18 @@ async fn auth_app(jwt_service: Arc) -> Router { } fn protected_auth_app(jwt_service: Arc, user_repo: Arc) -> Router { + protected_auth_app_with_mode(jwt_service, user_repo, AuthIdentityMode::UserSession) +} + +fn protected_auth_app_with_mode( + jwt_service: Arc, + user_repo: Arc, + identity_mode: AuthIdentityMode, +) -> Router { let state = AuthState { jwt_service, user_repo, - local: false, + identity_mode, }; Router::new() @@ -239,6 +247,57 @@ async fn auth_middleware_missing_user_returns_unauthorized_code() { assert_eq!(json["code"], "UNAUTHORIZED"); } +#[tokio::test] +async fn auth_middleware_session_generation_mismatch_returns_unauthorized_code() { + let jwt_service = Arc::new(JwtService::new("middleware_test_secret".into())); + let token = jwt_service + .sign_with_session_generation("system_default_user", "system_default_user", 0) + .unwrap(); + let db = init_database_memory().await.unwrap(); + let repo = Arc::new(SqliteUserRepository::new(db.pool().clone())); + repo.increment_session_generation("system_default_user").await.unwrap(); + let app = protected_auth_app(jwt_service, repo as Arc); + + let resp = app + .oneshot( + Request::get("/protected") + .header(header::AUTHORIZATION, format!("Bearer {token}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = json_body(resp).await; + assert_eq!(json["code"], "UNAUTHORIZED"); +} + +#[tokio::test] +async fn auth_middleware_aionpro_rejects_local_user_token() { + let jwt_service = Arc::new(JwtService::new("middleware_test_secret".into())); + let token = jwt_service + .sign_with_session_generation("system_default_user", "system_default_user", 0) + .unwrap(); + let db = init_database_memory().await.unwrap(); + let repo = Arc::new(SqliteUserRepository::new(db.pool().clone())); + let app = protected_auth_app_with_mode(jwt_service, repo as Arc, AuthIdentityMode::AionPro); + + let resp = app + .oneshot( + Request::get("/protected") + .header(header::AUTHORIZATION, format!("Bearer {token}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = json_body(resp).await; + assert_eq!(json["code"], "USER_CONTEXT_REQUIRED"); +} + #[tokio::test] async fn auth_middleware_database_error_returns_internal_error_code() { let jwt_service = Arc::new(JwtService::new("middleware_test_secret".into())); @@ -410,6 +469,8 @@ async fn authenticated_action_limit_uses_user_id_key() { request.extensions_mut().insert(CurrentUser { id: "user_42".into(), username: "admin".into(), + user_type: UserType::Local, + status: UserStatus::Active, }); Ok::<_, std::convert::Infallible>(next.run(request).await) }, diff --git a/crates/aionui-auth/tests/route_tests.rs b/crates/aionui-auth/tests/route_tests.rs index 2f94e6225..9c2446287 100644 --- a/crates/aionui-auth/tests/route_tests.rs +++ b/crates/aionui-auth/tests/route_tests.rs @@ -4,7 +4,7 @@ //! T7 (current user), T8 (change password), T9 (refresh token), //! T10 (ws token), T11 (QR login). -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use axum::Router; use axum::body::Body; @@ -12,8 +12,11 @@ use axum::http::{Request, StatusCode, header}; use http_body_util::BodyExt; use tower::ServiceExt; -use aionui_auth::{AuthRouterState, CookieConfig, JwtService, QrTokenStore, auth_routes, hash_password}; -use aionui_db::{IUserRepository, SqliteUserRepository, init_database_memory}; +use aionui_auth::{ + AuthIdentityMode, AuthRouterState, CookieConfig, JwtService, QrTokenStore, SessionRevokedHook, auth_routes, + hash_password, +}; +use aionui_db::{IUserRepository, SqliteUserRepository, UserStatus, init_database_memory}; // --------------------------------------------------------------------------- // Test helpers @@ -25,6 +28,19 @@ async fn test_app() -> (Router, TestContext) { } async fn test_app_with_local(local: bool) -> (Router, TestContext) { + test_app_with_options(local, None).await +} + +async fn test_app_with_options(local: bool, bootstrap_secret: Option<&str>) -> (Router, TestContext) { + test_app_with_options_and_hook(local, bootstrap_secret, false, None).await +} + +async fn test_app_with_options_and_hook( + local: bool, + bootstrap_secret: Option<&str>, + aionpro_mode: bool, + session_revoked_hook: Option>, +) -> (Router, TestContext) { let db = init_database_memory().await.unwrap(); let user_repo = Arc::new(SqliteUserRepository::new(db.pool().clone())) as Arc; let jwt_service = Arc::new(JwtService::new("test_secret_for_routes".into())); @@ -39,7 +55,17 @@ async fn test_app_with_local(local: bool) -> (Router, TestContext) { user_repo: user_repo.clone(), cookie_config, qr_token_store: qr_token_store.clone(), + identity_mode: if local { + AuthIdentityMode::Local + } else if aionpro_mode { + AuthIdentityMode::AionPro + } else { + AuthIdentityMode::UserSession + }, + bootstrap_secret: bootstrap_secret.map(Arc::::from), + session_revoked_hook, local, + aionpro_mode, }; let app = auth_routes(state); @@ -88,6 +114,35 @@ fn json_post(uri: &str, body: &str) -> Request { .unwrap() } +fn json_post_with_bootstrap_secret(uri: &str, body: &str, secret: &str) -> Request { + Request::builder() + .method("POST") + .uri(uri) + .header("content-type", "application/json") + .header("x-aioncore-bootstrap-secret", secret) + .body(Body::from(body.to_owned())) + .unwrap() +} + +fn json_put_with_bootstrap_secret(uri: &str, body: &str, secret: &str) -> Request { + Request::builder() + .method("PUT") + .uri(uri) + .header("content-type", "application/json") + .header("x-aioncore-bootstrap-secret", secret) + .body(Body::from(body.to_owned())) + .unwrap() +} + +fn json_put_anonymous(uri: &str, body: &str) -> Request { + Request::builder() + .method("PUT") + .uri(uri) + .header("content-type", "application/json") + .body(Body::from(body.to_owned())) + .unwrap() +} + /// Helper: perform a JSON POST request with auth token. fn json_post_with_token(uri: &str, body: &str, token: &str) -> Request { Request::builder() @@ -120,6 +175,21 @@ async fn body_json(resp: axum::response::Response) -> serde_json::Value { serde_json::from_slice(&bytes).unwrap() } +fn extract_session_token(resp: &axum::response::Response) -> Option { + resp.headers() + .get_all(header::SET_COOKIE) + .iter() + .filter_map(|value| value.to_str().ok()) + .find_map(|cookie| { + cookie + .split(',') + .find(|part| part.trim_start().starts_with("aionui-session=")) + .and_then(|part| part.trim_start().strip_prefix("aionui-session=")) + .and_then(|value| value.split(';').next()) + .map(str::to_owned) + }) +} + /// Helper: login and return (token, user_id). async fn login(app: &mut Router, username: &str, password: &str) -> (String, String) { let req = json_post( @@ -215,6 +285,21 @@ async fn t4_4_login_missing_fields() { assert_eq!(resp.status(), StatusCode::BAD_REQUEST); } +#[tokio::test] +async fn login_rejects_aionpro_mode() { + let (app, ctx) = test_app_with_options_and_hook(false, Some("bootstrap-secret"), true, None).await; + create_test_user(&ctx, "admin", "StrongP@ss1").await; + + let resp = app + .oneshot(json_post("/login", r#"{"username":"admin","password":"StrongP@ss1"}"#)) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = body_json(resp).await; + assert_eq!(json["code"], "USER_CONTEXT_REQUIRED"); +} + #[tokio::test] async fn t4_5_login_empty_password_hash_returns_401() { // Regression: when the seeded system user has an empty password_hash @@ -579,6 +664,22 @@ async fn t9_3_refresh_missing_token() { assert_eq!(resp.status(), StatusCode::BAD_REQUEST); } +#[tokio::test] +async fn refresh_rejects_local_user_token_in_aionpro_mode() { + let (app, ctx) = test_app_with_options_and_hook(false, Some("bootstrap-secret"), true, None).await; + let token = ctx + .jwt_service + .sign_with_session_generation("system_default_user", "system_default_user", 0) + .unwrap(); + + let body = format!(r#"{{"token":"{token}"}}"#); + let resp = app.oneshot(json_post("/api/auth/refresh", &body)).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = body_json(resp).await; + assert_eq!(json["code"], "USER_CONTEXT_REQUIRED"); +} + // =========================================================================== // T10. WebSocket Token (GET /api/ws-token) // =========================================================================== @@ -697,6 +798,19 @@ async fn t11_5_qr_login_missing_token() { assert_eq!(resp.status(), StatusCode::BAD_REQUEST); } +#[tokio::test] +async fn qr_login_rejects_aionpro_mode() { + let (app, ctx) = test_app_with_options_and_hook(false, Some("bootstrap-secret"), true, None).await; + let qr_token = ctx.qr_token_store.generate(); + + let body = format!(r#"{{"qr_token":"{qr_token}"}}"#); + let resp = app.oneshot(json_post("/api/auth/qr-login", &body)).await.unwrap(); + + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = body_json(resp).await; + assert_eq!(json["code"], "USER_CONTEXT_REQUIRED"); +} + // =========================================================================== // QR Login Page (GET /qr-login) // =========================================================================== @@ -713,6 +827,224 @@ async fn qr_login_page_returns_html() { assert!(content_type.contains("text/html")); } +// =========================================================================== +// Internal external user provision and session exchange +// =========================================================================== + +#[tokio::test] +async fn external_user_provision_requires_bootstrap_secret() { + let (app, _ctx) = test_app_with_options(false, Some("bootstrap-secret")).await; + + let resp = app + .oneshot(json_put_anonymous( + "/api/auth/internal/external-users/pro-1", + r#"{"user_type":"aionpro","username":"Pro User"}"#, + )) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = body_json(resp).await; + assert_eq!(json["code"], "BOOTSTRAP_SECRET_REQUIRED"); +} + +#[tokio::test] +async fn external_user_provision_rejects_wrong_bootstrap_secret() { + let (app, _ctx) = test_app_with_options(false, Some("bootstrap-secret")).await; + + let resp = app + .oneshot(json_put_with_bootstrap_secret( + "/api/auth/internal/external-users/pro-1", + r#"{"user_type":"aionpro","username":"Pro User"}"#, + "wrong-secret", + )) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = body_json(resp).await; + assert_eq!(json["code"], "INVALID_BOOTSTRAP_SECRET"); +} + +#[tokio::test] +async fn external_user_provision_is_idempotent() { + let (app, _ctx) = test_app_with_options(false, Some("bootstrap-secret")).await; + let body = r#"{"user_type":"aionpro","username":"Pro User","email":"pro@example.com"}"#; + + let first = app + .clone() + .oneshot(json_put_with_bootstrap_secret( + "/api/auth/internal/external-users/pro-1", + body, + "bootstrap-secret", + )) + .await + .unwrap(); + assert_eq!(first.status(), StatusCode::OK); + let first_json = body_json(first).await; + let user_id = first_json["data"]["user_id"].as_str().unwrap().to_owned(); + assert_eq!(first_json["data"]["user_type"], "aionpro"); + assert_eq!(first_json["data"]["external_user_id"], "pro-1"); + + let second = app + .oneshot(json_put_with_bootstrap_secret( + "/api/auth/internal/external-users/pro-1", + r#"{"user_type":"aionpro","username":"Ignored"}"#, + "bootstrap-secret", + )) + .await + .unwrap(); + assert_eq!(second.status(), StatusCode::OK); + let second_json = body_json(second).await; + assert_eq!(second_json["data"]["user_id"], user_id); +} + +#[tokio::test] +async fn external_session_exchange_returns_cookie_without_json_token() { + let (app, _ctx) = test_app_with_options(false, Some("bootstrap-secret")).await; + + let provision = app + .clone() + .oneshot(json_put_with_bootstrap_secret( + "/api/auth/internal/external-users/pro-session", + r#"{"user_type":"aionpro","username":"Pro User"}"#, + "bootstrap-secret", + )) + .await + .unwrap(); + assert_eq!(provision.status(), StatusCode::OK); + + let session = app + .clone() + .oneshot(json_post_with_bootstrap_secret( + "/api/auth/internal/external-sessions", + r#"{"user_type":"aionpro","external_user_id":"pro-session"}"#, + "bootstrap-secret", + )) + .await + .unwrap(); + assert_eq!(session.status(), StatusCode::OK); + assert!(session.headers().get(header::SET_COOKIE).is_some()); + let token = extract_session_token(&session).unwrap(); + let json = body_json(session).await; + assert!(json["data"].get("token").is_none()); + assert_eq!(json["data"]["user"]["username"], "Pro User"); + + let user_resp = app.oneshot(get_with_token("/api/auth/user", &token)).await.unwrap(); + assert_eq!(user_resp.status(), StatusCode::OK); + let user_json = body_json(user_resp).await; + assert_eq!(user_json["user"]["username"], "Pro User"); +} + +#[tokio::test] +async fn external_session_exchange_rejects_unprovisioned_user() { + let (app, _ctx) = test_app_with_options(false, Some("bootstrap-secret")).await; + + let resp = app + .oneshot(json_post_with_bootstrap_secret( + "/api/auth/internal/external-sessions", + r#"{"user_type":"aionpro","external_user_id":"missing"}"#, + "bootstrap-secret", + )) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let json = body_json(resp).await; + assert_eq!(json["code"], "USER_CONTEXT_REQUIRED"); +} + +#[tokio::test] +async fn external_session_revoke_invalidates_existing_token() { + let revoked_users = Arc::new(Mutex::new(Vec::new())); + let hook_users = revoked_users.clone(); + let (app, _ctx) = test_app_with_options_and_hook( + false, + Some("bootstrap-secret"), + false, + Some(Arc::new(move |user_id: &str| { + hook_users.lock().unwrap().push(user_id.to_owned()); + })), + ) + .await; + + let provision = app + .clone() + .oneshot(json_put_with_bootstrap_secret( + "/api/auth/internal/external-users/pro-revoke", + r#"{"user_type":"aionpro","username":"Pro User"}"#, + "bootstrap-secret", + )) + .await + .unwrap(); + assert_eq!(provision.status(), StatusCode::OK); + + let session = app + .clone() + .oneshot(json_post_with_bootstrap_secret( + "/api/auth/internal/external-sessions", + r#"{"user_type":"aionpro","external_user_id":"pro-revoke"}"#, + "bootstrap-secret", + )) + .await + .unwrap(); + assert_eq!(session.status(), StatusCode::OK); + let token = extract_session_token(&session).unwrap(); + let session_json = body_json(session).await; + assert_eq!(session_json["data"]["session_generation"], 0); + assert!(session_json["data"].get("token").is_none()); + + let revoke = app + .clone() + .oneshot(json_post_with_bootstrap_secret( + "/api/auth/internal/external-sessions/revoke", + r#"{"user_type":"aionpro","external_user_id":"pro-revoke"}"#, + "bootstrap-secret", + )) + .await + .unwrap(); + assert_eq!(revoke.status(), StatusCode::OK); + let revoke_json = body_json(revoke).await; + assert_eq!(revoke_json["data"]["session_generation"], 1); + assert_eq!( + revoked_users.lock().unwrap().as_slice(), + &[revoke_json["data"]["user_id"].as_str().unwrap().to_owned()] + ); + + let user_resp = app.oneshot(get_with_token("/api/auth/user", &token)).await.unwrap(); + assert_eq!(user_resp.status(), StatusCode::UNAUTHORIZED); +} + +#[tokio::test] +async fn external_session_exchange_rejects_disabled_user() { + let (app, ctx) = test_app_with_options(false, Some("bootstrap-secret")).await; + let provision = app + .clone() + .oneshot(json_put_with_bootstrap_secret( + "/api/auth/internal/external-users/pro-disabled", + r#"{"user_type":"aionpro","username":"Disabled"}"#, + "bootstrap-secret", + )) + .await + .unwrap(); + let provision_json = body_json(provision).await; + let user_id = provision_json["data"]["user_id"].as_str().unwrap(); + ctx.user_repo.set_status(user_id, UserStatus::Disabled).await.unwrap(); + + let resp = app + .oneshot(json_post_with_bootstrap_secret( + "/api/auth/internal/external-sessions", + r#"{"user_type":"aionpro","external_user_id":"pro-disabled"}"#, + "bootstrap-secret", + )) + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + let json = body_json(resp).await; + assert_eq!(json["code"], "USER_DISABLED"); +} + // =========================================================================== // T12. Local-only internal user routes // =========================================================================== @@ -742,6 +1074,8 @@ async fn t12_2_internal_user_routes_work_in_local_mode() { assert_eq!(system_resp.status(), StatusCode::OK); let system_json = body_json(system_resp).await; assert_eq!(system_json["data"]["id"], "system_default_user"); + assert!(system_json["data"].get("password_hash").is_none()); + assert!(system_json["data"].get("jwt_secret").is_none()); let user_resp = app .clone() @@ -750,6 +1084,8 @@ async fn t12_2_internal_user_routes_work_in_local_mode() { .unwrap(); assert_eq!(user_resp.status(), StatusCode::OK); let user_json = body_json(user_resp).await; + assert!(user_json["data"].get("password_hash").is_none()); + assert!(user_json["data"].get("jwt_secret").is_none()); let user_id = user_json["data"]["id"].as_str().unwrap().to_owned(); let update_resp = app diff --git a/crates/aionui-channel/Cargo.toml b/crates/aionui-channel/Cargo.toml index a0de2d46d..a05104794 100644 --- a/crates/aionui-channel/Cargo.toml +++ b/crates/aionui-channel/Cargo.toml @@ -12,6 +12,7 @@ weixin = ["dep:reqwest", "dep:futures-util", "dep:base64", "dep:uuid"] [dependencies] aionui-common.workspace = true +aionui-auth.workspace = true aionui-db.workspace = true aionui-api-types.workspace = true aionui-realtime.workspace = true diff --git a/crates/aionui-channel/src/action.rs b/crates/aionui-channel/src/action.rs index 0d9d54c0c..88dd85b4f 100644 --- a/crates/aionui-channel/src/action.rs +++ b/crates/aionui-channel/src/action.rs @@ -21,6 +21,7 @@ pub enum MessageResult { /// Message was dispatched to the AI Agent. The caller should send /// a "thinking" placeholder and then relay stream events. Dispatched { + owner_user_id: String, session_id: String, conversation_id: Option, }, @@ -40,6 +41,7 @@ pub struct ActionExecutor { pairing: Arc, session_mgr: Arc, settings: Arc, + owner_user_id: Option, } impl ActionExecutor { @@ -47,14 +49,20 @@ impl ActionExecutor { pairing: Arc, session_mgr: Arc, settings: Arc, + owner_user_id: Option, ) -> Self { Self { pairing, session_mgr, settings, + owner_user_id, } } + pub fn owner_user_id(&self) -> Option<&str> { + self.owner_user_id.as_deref() + } + /// Main entry point: handle an incoming message from any platform. /// /// Flow: @@ -65,15 +73,23 @@ impl ActionExecutor { let platform_type = msg.platform.to_string(); let user_id = &msg.user.id; let chat_id = &msg.chat_id; + let owner_user_id = msg + .owner_user_id + .as_deref() + .or(self.owner_user_id.as_deref()) + .ok_or_else(|| ChannelError::InvalidConfig("channel owner user is required".to_owned()))?; // 1. Authorization check — resolve platform user → internal user ID - let internal_user_id = self.pairing.get_internal_user_id(user_id, &platform_type).await?; + let internal_user_id = self + .pairing + .get_internal_user_id(owner_user_id, user_id, &platform_type) + .await?; let internal_user_id = match internal_user_id { Some(id) => id, None => { let response = self - .handle_unauthorized(user_id, &platform_type, &msg.user.display_name) + .handle_unauthorized(owner_user_id, user_id, &platform_type, &msg.user.display_name) .await?; return Ok(MessageResult::Action(response)); } @@ -81,15 +97,21 @@ impl ActionExecutor { // 2. Button callback → action routing if let Some(action) = &msg.action { - let response = self.route_action(action, &internal_user_id).await?; + let response = self.route_action(owner_user_id, action, &internal_user_id).await?; return Ok(MessageResult::Action(response)); } // 3. Text message → session resolution → AI dispatch - let agent_config = self.settings.get_agent_config(msg.platform).await?; + let agent_config = self.settings.get_agent_config(owner_user_id, msg.platform).await?; let session = self .session_mgr - .get_or_create_session(&internal_user_id, chat_id, &agent_config.agent_type, None) + .get_or_create_session( + owner_user_id, + &internal_user_id, + chat_id, + &agent_config.agent_type, + None, + ) .await?; info!( @@ -101,6 +123,7 @@ impl ActionExecutor { ); Ok(MessageResult::Dispatched { + owner_user_id: owner_user_id.to_owned(), session_id: session.id, conversation_id: session.conversation_id, }) @@ -110,13 +133,14 @@ impl ActionExecutor { /// a response with instructions and action buttons. async fn handle_unauthorized( &self, + owner_user_id: &str, platform_user_id: &str, platform_type: &str, display_name: &str, ) -> Result { let code = self .pairing - .request_pairing(platform_user_id, platform_type, Some(display_name)) + .request_pairing(owner_user_id, platform_user_id, platform_type, Some(display_name)) .await?; debug!( @@ -131,31 +155,45 @@ impl ActionExecutor { /// Routes an action to the appropriate handler by category. async fn route_action( &self, + owner_user_id: &str, action: &UnifiedAction, internal_user_id: &str, ) -> Result { match action.category { - ActionCategory::Platform => self.handle_platform_action(action).await, - ActionCategory::System => self.handle_system_action(action, internal_user_id).await, + ActionCategory::Platform => self.handle_platform_action(owner_user_id, action).await, + ActionCategory::System => self.handle_system_action(owner_user_id, action, internal_user_id).await, ActionCategory::Chat => self.handle_chat_action(action).await, } } // ── Platform actions ──────────────────────────────────────────── - async fn handle_platform_action(&self, action: &UnifiedAction) -> Result { + async fn handle_platform_action( + &self, + owner_user_id: &str, + action: &UnifiedAction, + ) -> Result { match action.action.as_str() { "pairing.show" | "pairing.refresh" => { let code = self .pairing - .request_pairing(&action.context.user_id, &action.context.platform.to_string(), None) + .request_pairing( + owner_user_id, + &action.context.user_id, + &action.context.platform.to_string(), + None, + ) .await?; Ok(build_pairing_response(&code)) } "pairing.check" => { let authorized = self .pairing - .is_user_authorized(&action.context.user_id, &action.context.platform.to_string()) + .is_user_authorized( + owner_user_id, + &action.context.user_id, + &action.context.platform.to_string(), + ) .await?; if authorized { Ok(ActionResponse { @@ -217,6 +255,7 @@ impl ActionExecutor { async fn handle_system_action( &self, + owner_user_id: &str, action: &UnifiedAction, internal_user_id: &str, ) -> Result { @@ -224,10 +263,13 @@ impl ActionExecutor { "session.new" => { let user_id = internal_user_id; let chat_id = &action.context.chat_id; - let agent_config = self.settings.get_agent_config(action.context.platform).await?; + let agent_config = self + .settings + .get_agent_config(owner_user_id, action.context.platform) + .await?; let session = self .session_mgr - .reset_session(user_id, chat_id, &agent_config.agent_type, None) + .reset_session(owner_user_id, user_id, chat_id, &agent_config.agent_type, None) .await?; Ok(ActionResponse { @@ -251,10 +293,13 @@ impl ActionExecutor { "session.status" => { let user_id = internal_user_id; let chat_id = &action.context.chat_id; - let agent_config = self.settings.get_agent_config(action.context.platform).await?; + let agent_config = self + .settings + .get_agent_config(owner_user_id, action.context.platform) + .await?; let session = self .session_mgr - .get_or_create_session(user_id, chat_id, &agent_config.agent_type, None) + .get_or_create_session(owner_user_id, user_id, chat_id, &agent_config.agent_type, None) .await?; Ok(ActionResponse { @@ -513,6 +558,7 @@ mod tests { } // ── Mock IChannelRepository ──────────────────────────────────────── + const OWNER_ID: &str = "owner-test"; struct MockRepo { users: Mutex>, @@ -532,6 +578,7 @@ mod tests { fn add_authorized_user(&self, platform_user_id: &str, platform_type: &str) { let user = AssistantUserRow { id: format!("user_{platform_user_id}"), + owner_user_id: OWNER_ID.to_owned(), platform_user_id: platform_user_id.to_owned(), platform_type: platform_type.to_owned(), display_name: Some("Test User".into()), @@ -545,27 +592,33 @@ mod tests { #[async_trait::async_trait] impl IChannelRepository for MockRepo { - async fn get_all_plugins(&self) -> Result, DbError> { + async fn get_all_plugins(&self, _owner_user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn get_plugin(&self, _id: &str) -> Result, DbError> { + async fn get_plugin(&self, _owner_user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } - async fn upsert_plugin(&self, _row: &ChannelPluginRow) -> Result<(), DbError> { + async fn upsert_plugin(&self, _owner_user_id: &str, _row: &ChannelPluginRow) -> Result<(), DbError> { Ok(()) } - async fn update_plugin_status(&self, _id: &str, _params: &UpdatePluginStatusParams) -> Result<(), DbError> { + async fn update_plugin_status( + &self, + _owner_user_id: &str, + _id: &str, + _params: &UpdatePluginStatusParams, + ) -> Result<(), DbError> { Ok(()) } - async fn delete_plugin(&self, _id: &str) -> Result<(), DbError> { + async fn delete_plugin(&self, _owner_user_id: &str, _id: &str) -> Result<(), DbError> { Ok(()) } - async fn get_all_users(&self) -> Result, DbError> { + async fn get_all_users(&self, _owner_user_id: &str) -> Result, DbError> { Ok(self.users.lock().unwrap().clone()) } async fn get_user_by_platform( &self, + _owner_user_id: &str, platform_user_id: &str, platform_type: &str, ) -> Result, DbError> { @@ -575,26 +628,32 @@ mod tests { .find(|u| u.platform_user_id == platform_user_id && u.platform_type == platform_type) .cloned()) } - async fn create_user(&self, row: &AssistantUserRow) -> Result<(), DbError> { + async fn create_user(&self, _owner_user_id: &str, row: &AssistantUserRow) -> Result<(), DbError> { self.users.lock().unwrap().push(row.clone()); Ok(()) } - async fn update_user_last_active(&self, _id: &str, _last_active: TimestampMs) -> Result<(), DbError> { + async fn update_user_last_active( + &self, + _owner_user_id: &str, + _id: &str, + _last_active: TimestampMs, + ) -> Result<(), DbError> { Ok(()) } - async fn delete_user(&self, _id: &str) -> Result<(), DbError> { + async fn delete_user(&self, _owner_user_id: &str, _id: &str) -> Result<(), DbError> { Ok(()) } - async fn get_all_sessions(&self) -> Result, DbError> { + async fn get_all_sessions(&self, _owner_user_id: &str) -> Result, DbError> { Ok(self.sessions.lock().unwrap().clone()) } - async fn get_session(&self, id: &str) -> Result, DbError> { + async fn get_session(&self, _owner_user_id: &str, id: &str) -> Result, DbError> { let sessions = self.sessions.lock().unwrap(); Ok(sessions.iter().find(|s| s.id == id).cloned()) } async fn get_or_create_session( &self, + _owner_user_id: &str, user_id: &str, chat_id: &str, new_row: &AssistantSessionRow, @@ -610,10 +669,20 @@ mod tests { sessions.push(new_row.clone()); Ok(new_row.clone()) } - async fn update_session_activity(&self, _id: &str, _last_activity: TimestampMs) -> Result<(), DbError> { + async fn update_session_activity( + &self, + _owner_user_id: &str, + _id: &str, + _last_activity: TimestampMs, + ) -> Result<(), DbError> { Ok(()) } - async fn update_session_conversation(&self, id: &str, conversation_id: &str) -> Result<(), DbError> { + async fn update_session_conversation( + &self, + _owner_user_id: &str, + id: &str, + conversation_id: &str, + ) -> Result<(), DbError> { let mut sessions = self.sessions.lock().unwrap(); if let Some(s) = sessions.iter_mut().find(|s| s.id == id) { s.conversation_id = Some(conversation_id.to_owned()); @@ -622,7 +691,12 @@ mod tests { Err(DbError::NotFound(id.into())) } } - async fn update_session_agent_type(&self, id: &str, agent_type: &str) -> Result<(), DbError> { + async fn update_session_agent_type( + &self, + _owner_user_id: &str, + id: &str, + agent_type: &str, + ) -> Result<(), DbError> { let mut sessions = self.sessions.lock().unwrap(); if let Some(s) = sessions.iter_mut().find(|s| s.id == id) { s.agent_type = agent_type.to_owned(); @@ -631,29 +705,38 @@ mod tests { Err(DbError::NotFound(id.into())) } } - async fn delete_sessions_by_user(&self, user_id: &str) -> Result<(), DbError> { + async fn delete_sessions_by_user(&self, _owner_user_id: &str, user_id: &str) -> Result<(), DbError> { self.sessions.lock().unwrap().retain(|s| s.user_id != user_id); Ok(()) } - async fn delete_session_by_user_chat(&self, user_id: &str, chat_id: &str) -> Result<(), DbError> { + async fn delete_session_by_user_chat( + &self, + _owner_user_id: &str, + user_id: &str, + chat_id: &str, + ) -> Result<(), DbError> { let mut sessions = self.sessions.lock().unwrap(); sessions.retain(|s| !(s.user_id == user_id && s.chat_id.as_deref() == Some(chat_id))); Ok(()) } - async fn create_pairing(&self, row: &PairingCodeRow) -> Result<(), DbError> { + async fn create_pairing(&self, _owner_user_id: &str, row: &PairingCodeRow) -> Result<(), DbError> { self.pairings.lock().unwrap().push(row.clone()); Ok(()) } - async fn get_pending_pairings(&self) -> Result, DbError> { + async fn get_pending_pairings(&self, _owner_user_id: &str) -> Result, DbError> { let pairings = self.pairings.lock().unwrap(); Ok(pairings.iter().filter(|p| p.status == "pending").cloned().collect()) } - async fn get_pairing_by_code(&self, code: &str) -> Result, DbError> { + async fn get_pairing_by_code( + &self, + _owner_user_id: &str, + code: &str, + ) -> Result, DbError> { let pairings = self.pairings.lock().unwrap(); Ok(pairings.iter().find(|p| p.code == code).cloned()) } - async fn update_pairing_status(&self, code: &str, status: &str) -> Result<(), DbError> { + async fn update_pairing_status(&self, _owner_user_id: &str, code: &str, status: &str) -> Result<(), DbError> { let mut pairings = self.pairings.lock().unwrap(); if let Some(p) = pairings.iter_mut().find(|p| p.code == code) { p.status = status.to_owned(); @@ -662,7 +745,7 @@ mod tests { Err(DbError::NotFound(code.into())) } } - async fn cleanup_expired_pairings(&self, _now: TimestampMs) -> Result { + async fn cleanup_expired_pairings(&self, _owner_user_id: &str, _now: TimestampMs) -> Result { Ok(0) } } @@ -673,16 +756,16 @@ mod tests { #[async_trait::async_trait] impl IClientPreferenceRepository for MockPrefRepo { - async fn get_all(&self) -> Result, DbError> { + async fn get_all(&self, _user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn get_by_keys(&self, _keys: &[&str]) -> Result, DbError> { + async fn get_by_keys(&self, _user_id: &str, _keys: &[&str]) -> Result, DbError> { Ok(vec![]) } - async fn upsert_batch(&self, _entries: &[(&str, &str)]) -> Result<(), DbError> { + async fn upsert_batch(&self, _user_id: &str, _entries: &[(&str, &str)]) -> Result<(), DbError> { Ok(()) } - async fn delete_keys(&self, _keys: &[&str]) -> Result<(), DbError> { + async fn delete_keys(&self, _user_id: &str, _keys: &[&str]) -> Result<(), DbError> { Ok(()) } } @@ -696,12 +779,24 @@ mod tests { let session_mgr = Arc::new(SessionManager::new(repo.clone())); let pref_repo: Arc = Arc::new(MockPrefRepo); let settings = Arc::new(ChannelSettingsService::new(pref_repo)); - let executor = ActionExecutor::new(pairing, session_mgr, settings); + let executor = ActionExecutor::new(pairing, session_mgr, settings, Some(OWNER_ID.to_owned())); + (executor, repo) + } + + fn setup_without_owner() -> (ActionExecutor, Arc) { + let repo = Arc::new(MockRepo::new()); + let broadcaster = Arc::new(MockBroadcaster); + let pairing = Arc::new(PairingService::new(repo.clone(), broadcaster)); + let session_mgr = Arc::new(SessionManager::new(repo.clone())); + let pref_repo: Arc = Arc::new(MockPrefRepo); + let settings = Arc::new(ChannelSettingsService::new(pref_repo)); + let executor = ActionExecutor::new(pairing, session_mgr, settings, None); (executor, repo) } fn make_text_message(user_id: &str, chat_id: &str, text: &str, platform: PluginType) -> UnifiedIncomingMessage { UnifiedIncomingMessage { + owner_user_id: None, id: "msg_1".into(), platform, chat_id: chat_id.into(), @@ -732,6 +827,7 @@ mod tests { params: Option>, ) -> UnifiedIncomingMessage { UnifiedIncomingMessage { + owner_user_id: None, id: "msg_1".into(), platform, chat_id: chat_id.into(), @@ -783,6 +879,16 @@ mod tests { } } + #[tokio::test] + async fn missing_owner_user_id_is_rejected() { + let (executor, _repo) = setup_without_owner(); + let msg = make_text_message("tg_42", "chat_1", "Hello", PluginType::Telegram); + + let err = executor.handle_incoming_message(&msg).await.unwrap_err(); + assert!(matches!(err, ChannelError::InvalidConfig(_))); + assert!(err.to_string().contains("owner user")); + } + #[tokio::test] async fn authorized_user_text_dispatches_to_agent() { let (executor, repo) = setup(); diff --git a/crates/aionui-channel/src/channel_settings.rs b/crates/aionui-channel/src/channel_settings.rs index be223a69b..40cebfcb6 100644 --- a/crates/aionui-channel/src/channel_settings.rs +++ b/crates/aionui-channel/src/channel_settings.rs @@ -78,9 +78,13 @@ impl ChannelSettingsService { /// - **Legacy:** `{"backend":"claude","name":"Claude"}` (no agent_type field) /// /// Falls back to `agent_type=aionrs, backend=None` when no config exists. - pub async fn get_agent_config(&self, platform: PluginType) -> Result { + pub async fn get_agent_config( + &self, + user_id: &str, + platform: PluginType, + ) -> Result { let key = agent_key(platform); - let prefs = self.pref_repo.get_by_keys(&[&key]).await?; + let prefs = self.pref_repo.get_by_keys(user_id, &[&key]).await?; let Some(pref) = prefs.into_iter().next() else { return Ok(default_agent_config()); @@ -88,7 +92,7 @@ impl ChannelSettingsService { if let Some(setting) = parse_channel_assistant_setting(&pref.value) { if let Some(assistant_id) = setting.assistant_id.as_deref() { - if let Some(resolved) = self.resolve_assistant_agent_config(assistant_id).await? { + if let Some(resolved) = self.resolve_assistant_agent_config(user_id, assistant_id).await? { debug!( platform = %platform, assistant_id, @@ -137,9 +141,13 @@ impl ChannelSettingsService { /// Reads the model configuration for a platform from `client_preferences`. /// /// Returns `None` when no model is configured (common for ACP agents). - pub async fn get_model_config(&self, platform: PluginType) -> Result, ChannelError> { + pub async fn get_model_config( + &self, + user_id: &str, + platform: PluginType, + ) -> Result, ChannelError> { let key = model_key(platform); - let prefs = self.pref_repo.get_by_keys(&[&key]).await?; + let prefs = self.pref_repo.get_by_keys(user_id, &[&key]).await?; let Some(pref) = prefs.into_iter().next() else { return Ok(None); @@ -165,11 +173,12 @@ impl ChannelSettingsService { pub async fn get_platform_settings( &self, + user_id: &str, platform: PluginType, ) -> Result { let key_agent = agent_key(platform); let key_model = model_key(platform); - let prefs = self.pref_repo.get_by_keys(&[&key_agent, &key_model]).await?; + let prefs = self.pref_repo.get_by_keys(user_id, &[&key_agent, &key_model]).await?; let mut assistant = None; let mut default_model = None; @@ -177,7 +186,10 @@ impl ChannelSettingsService { for pref in prefs { if pref.key == key_agent { if let Some(parsed) = parse_channel_assistant_setting(&pref.value) { - assistant = Some(self.normalize_channel_assistant_setting_for_response(parsed).await?); + assistant = Some( + self.normalize_channel_assistant_setting_for_response(user_id, parsed) + .await?, + ); } } else if pref.key == key_model { default_model = parse_channel_model_setting(&pref.value); @@ -185,7 +197,7 @@ impl ChannelSettingsService { } if assistant.is_none() { - assistant = self.resolve_default_channel_assistant_setting().await?; + assistant = self.resolve_default_channel_assistant_setting(user_id).await?; } Ok(ChannelPlatformSettingsResponse { @@ -197,17 +209,21 @@ impl ChannelSettingsService { pub async fn get_assistant_setting( &self, + user_id: &str, platform: PluginType, ) -> Result, ChannelError> { let key = agent_key(platform); - let prefs = self.pref_repo.get_by_keys(&[&key]).await?; + let prefs = self.pref_repo.get_by_keys(user_id, &[&key]).await?; let Some(pref) = prefs.into_iter().next() else { - return self.resolve_default_channel_assistant_setting().await; + return self.resolve_default_channel_assistant_setting(user_id).await; }; let parsed = if let Some(assistant) = parse_channel_assistant_setting(&pref.value) { - Some(self.normalize_channel_assistant_setting_for_response(assistant).await?) + Some( + self.normalize_channel_assistant_setting_for_response(user_id, assistant) + .await?, + ) } else { None }; @@ -217,29 +233,36 @@ impl ChannelSettingsService { pub async fn set_assistant_setting( &self, + user_id: &str, platform: PluginType, assistant: &ChannelAssistantSettingRequest, ) -> Result<(), ChannelError> { let normalized = normalize_channel_assistant_setting_for_write(assistant); let payload = serde_json::to_string(&normalized).map_err(ChannelError::Json)?; let key = agent_key(platform); - self.pref_repo.upsert_batch(&[(&key, payload.as_str())]).await?; + self.pref_repo + .upsert_batch(user_id, &[(&key, payload.as_str())]) + .await?; Ok(()) } pub async fn set_model_setting( &self, + user_id: &str, platform: PluginType, model: &ChannelDefaultModelSetting, ) -> Result<(), ChannelError> { let payload = serde_json::to_string(model).map_err(ChannelError::Json)?; let key = model_key(platform); - self.pref_repo.upsert_batch(&[(&key, payload.as_str())]).await?; + self.pref_repo + .upsert_batch(user_id, &[(&key, payload.as_str())]) + .await?; Ok(()) } async fn resolve_assistant_agent_config( &self, + user_id: &str, assistant_id: &str, ) -> Result, ChannelError> { let (Some(definition_repo), Some(overlay_repo)) = @@ -248,16 +271,19 @@ impl ChannelSettingsService { return Ok(None); }; - let Some(definition) = definition_repo.get_by_assistant_id(assistant_id).await? else { + let Some(definition) = definition_repo + .get_by_assistant_id_for_user(user_id, assistant_id) + .await? + else { return Ok(None); }; let agent_id = overlay_repo - .get(&definition.id) + .get_for_user(user_id, &definition.id) .await? .and_then(|row| row.agent_id_override) .unwrap_or(definition.agent_id); - let agent_backend = self.runtime_backend_for_agent_id(&agent_id).await?; + let agent_backend = self.runtime_backend_for_agent_id(user_id, &agent_id).await?; let agent_type = backend_to_agent_type(&agent_backend); let backend = if agent_type == "acp" { Some(agent_backend) } else { None }; @@ -266,6 +292,7 @@ impl ChannelSettingsService { async fn resolve_assistant_identity_for_legacy_binding( &self, + user_id: &str, assistant: &ChannelAssistantSettingResponse, ) -> Result, ChannelError> { let (Some(definition_repo), Some(_overlay_repo)) = @@ -284,13 +311,13 @@ impl ChannelSettingsService { return Ok(None); }; - let definitions = definition_repo.list().await?; + let definitions = definition_repo.list_for_user(user_id).await?; for definition in definitions { if definition.source != "generated" { continue; } - let runtime_backend = self.runtime_backend_for_agent_id(&definition.agent_id).await?; + let runtime_backend = self.runtime_backend_for_agent_id(user_id, &definition.agent_id).await?; if runtime_backend == legacy_backend { return Ok(Some(definition.assistant_id)); } @@ -301,6 +328,7 @@ impl ChannelSettingsService { async fn normalize_channel_assistant_setting_for_response( &self, + user_id: &str, assistant: ChannelAssistantSettingResponse, ) -> Result { let assistant_id = assistant @@ -319,7 +347,8 @@ impl ChannelSettingsService { let canonical_assistant_id = if assistant_id.is_some() { assistant_id } else { - self.resolve_assistant_identity_for_legacy_binding(&assistant).await? + self.resolve_assistant_identity_for_legacy_binding(user_id, &assistant) + .await? }; if canonical_assistant_id.is_some() { @@ -337,8 +366,9 @@ impl ChannelSettingsService { async fn resolve_default_channel_assistant_setting( &self, + user_id: &str, ) -> Result, ChannelError> { - let Some(assistant_id) = self.resolve_default_assistant_identity().await? else { + let Some(assistant_id) = self.resolve_default_assistant_identity(user_id).await? else { return Ok(None); }; @@ -351,25 +381,25 @@ impl ChannelSettingsService { })) } - async fn resolve_default_assistant_identity(&self) -> Result, ChannelError> { + async fn resolve_default_assistant_identity(&self, user_id: &str) -> Result, ChannelError> { let (Some(definition_repo), Some(overlay_repo)) = (&self.assistant_definition_repo, &self.assistant_overlay_repo) else { return Ok(None); }; - let definitions = definition_repo.list().await?; - let overlays = overlay_repo.list().await?; + let definitions = definition_repo.list_for_user(user_id).await?; + let overlays = overlay_repo.list_for_user(user_id).await?; for definition in definitions.iter().filter(|definition| definition.source == "generated") { - if self.effective_assistant_backend(definition, &overlays).await? == DEFAULT_AGENT_TYPE { + if self.effective_assistant_backend(user_id, definition, &overlays).await? == DEFAULT_AGENT_TYPE { return Ok(Some(definition.assistant_id.clone())); } } let mut any_aionrs = None; for definition in &definitions { - if self.effective_assistant_backend(definition, &overlays).await? == DEFAULT_AGENT_TYPE { + if self.effective_assistant_backend(user_id, definition, &overlays).await? == DEFAULT_AGENT_TYPE { any_aionrs = Some(definition); break; } @@ -383,6 +413,7 @@ impl ChannelSettingsService { async fn effective_assistant_backend( &self, + user_id: &str, definition: &aionui_db::models::AssistantDefinitionRow, overlays: &[aionui_db::models::AssistantOverlayRow], ) -> Result { @@ -391,14 +422,14 @@ impl ChannelSettingsService { .find(|overlay| overlay.assistant_definition_id == definition.id) .and_then(|overlay| overlay.agent_id_override.as_deref()) .unwrap_or(definition.agent_id.as_str()); - self.runtime_backend_for_agent_id(agent_id).await + self.runtime_backend_for_agent_id(user_id, agent_id).await } - async fn runtime_backend_for_agent_id(&self, agent_id: &str) -> Result { + async fn runtime_backend_for_agent_id(&self, user_id: &str, agent_id: &str) -> Result { let Some(agent_metadata_repo) = self.agent_metadata_repo.as_ref() else { return Ok(agent_id.to_owned()); }; - let rows = agent_metadata_repo.list_all().await?; + let rows = agent_metadata_repo.list_all_for_user(user_id).await?; Ok(resolve_agent_binding_from_rows(&rows, agent_id) .map(|binding| binding.runtime_backend) .unwrap_or_else(|| agent_id.to_owned())) @@ -506,6 +537,8 @@ mod tests { use aionui_db::{IAssistantDefinitionRepository, IAssistantOverlayRepository}; use std::sync::Mutex; + const TEST_USER_ID: &str = "user-1"; + struct MockPrefRepo { data: Mutex>, } @@ -526,11 +559,12 @@ mod tests { #[async_trait::async_trait] impl IClientPreferenceRepository for MockPrefRepo { - async fn get_all(&self) -> Result, DbError> { + async fn get_all(&self, _user_id: &str) -> Result, DbError> { let data = self.data.lock().unwrap(); Ok(data .iter() .map(|(k, v)| ClientPreference { + user_id: TEST_USER_ID.to_owned(), key: k.clone(), value: v.clone(), updated_at: 0, @@ -538,12 +572,13 @@ mod tests { .collect()) } - async fn get_by_keys(&self, keys: &[&str]) -> Result, DbError> { + async fn get_by_keys(&self, _user_id: &str, keys: &[&str]) -> Result, DbError> { let data = self.data.lock().unwrap(); Ok(data .iter() .filter(|(k, _)| keys.contains(&k.as_str())) .map(|(k, v)| ClientPreference { + user_id: TEST_USER_ID.to_owned(), key: k.clone(), value: v.clone(), updated_at: 0, @@ -551,7 +586,7 @@ mod tests { .collect()) } - async fn upsert_batch(&self, entries: &[(&str, &str)]) -> Result<(), DbError> { + async fn upsert_batch(&self, _user_id: &str, entries: &[(&str, &str)]) -> Result<(), DbError> { let mut data = self.data.lock().unwrap(); for (key, value) in entries { if let Some(existing) = data.iter_mut().find(|(k, _)| k == key) { @@ -563,7 +598,7 @@ mod tests { Ok(()) } - async fn delete_keys(&self, keys: &[&str]) -> Result<(), DbError> { + async fn delete_keys(&self, _user_id: &str, keys: &[&str]) -> Result<(), DbError> { let mut data = self.data.lock().unwrap(); data.retain(|(k, _)| !keys.contains(&k.as_str())); Ok(()) @@ -580,14 +615,49 @@ mod tests { Ok(self.rows.clone()) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + + async fn list_including_deleted_for_user( + &self, + _user_id: &str, + ) -> Result, DbError> { + self.list().await + } + async fn get_by_assistant_id(&self, assistant_id: &str) -> Result, DbError> { Ok(self.rows.iter().find(|row| row.assistant_id == assistant_id).cloned()) } + async fn get_by_assistant_id_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + + async fn get_by_assistant_id_including_deleted_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + async fn get_by_id(&self, definition_id: &str) -> Result, DbError> { Ok(self.rows.iter().find(|row| row.id == definition_id).cloned()) } + async fn get_by_id_for_user( + &self, + _user_id: &str, + definition_id: &str, + ) -> Result, DbError> { + self.get_by_id(definition_id).await + } + async fn get_by_source_ref( &self, source: &str, @@ -600,6 +670,24 @@ mod tests { .cloned()) } + async fn get_by_source_ref_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + + async fn get_by_source_ref_including_deleted_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + async fn upsert( &self, _params: &UpsertAssistantDefinitionParams<'_>, @@ -607,9 +695,26 @@ mod tests { panic!("unused in channel settings tests") } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn soft_delete(&self, _definition_id: &str, _deleted_at: i64) -> Result { panic!("unused in channel settings tests") } + + async fn soft_delete_for_user( + &self, + _user_id: &str, + definition_id: &str, + deleted_at: i64, + ) -> Result { + self.soft_delete(definition_id, deleted_at).await + } } struct MockAssistantOverlayRepo { @@ -626,17 +731,41 @@ mod tests { .cloned()) } + async fn get_for_user( + &self, + _user_id: &str, + definition_id: &str, + ) -> Result, DbError> { + self.get(definition_id).await + } + async fn list(&self) -> Result, DbError> { Ok(self.rows.clone()) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + async fn upsert(&self, _params: &UpsertAssistantOverlayParams<'_>) -> Result { panic!("unused in channel settings tests") } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantOverlayParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn delete(&self, _definition_id: &str) -> Result { panic!("unused in channel settings tests") } + + async fn delete_for_user(&self, _user_id: &str, definition_id: &str) -> Result { + self.delete(definition_id).await + } } fn make_definition(assistant_id: &str, agent_id: &str) -> AssistantDefinitionRow { @@ -731,7 +860,7 @@ mod tests { let repo = Arc::new(MockPrefRepo::new()); let svc = ChannelSettingsService::new(repo); - let config = svc.get_agent_config(PluginType::Telegram).await.unwrap(); + let config = svc.get_agent_config(TEST_USER_ID, PluginType::Telegram).await.unwrap(); assert_eq!(config.agent_type, "aionrs"); assert!(config.backend.is_none()); } @@ -744,7 +873,7 @@ mod tests { )])); let svc = ChannelSettingsService::new(repo); - let config = svc.get_agent_config(PluginType::Telegram).await.unwrap(); + let config = svc.get_agent_config(TEST_USER_ID, PluginType::Telegram).await.unwrap(); assert_eq!(config.agent_type, "acp"); assert_eq!(config.backend.as_deref(), Some("codex")); } @@ -757,7 +886,7 @@ mod tests { )])); let svc = ChannelSettingsService::new(repo); - let config = svc.get_agent_config(PluginType::Lark).await.unwrap(); + let config = svc.get_agent_config(TEST_USER_ID, PluginType::Lark).await.unwrap(); assert_eq!(config.agent_type, "aionrs"); assert!(config.backend.is_none()); } @@ -772,7 +901,7 @@ mod tests { )])); let svc = ChannelSettingsService::new(repo); - let config = svc.get_agent_config(PluginType::Telegram).await.unwrap(); + let config = svc.get_agent_config(TEST_USER_ID, PluginType::Telegram).await.unwrap(); assert_eq!(config.agent_type, "acp"); assert_eq!(config.backend.as_deref(), Some("claude")); } @@ -785,7 +914,7 @@ mod tests { )])); let svc = ChannelSettingsService::new(repo); - let config = svc.get_agent_config(PluginType::Lark).await.unwrap(); + let config = svc.get_agent_config(TEST_USER_ID, PluginType::Lark).await.unwrap(); assert_eq!(config.agent_type, "aionrs"); assert!(config.backend.is_none()); } @@ -798,7 +927,7 @@ mod tests { )])); let svc = ChannelSettingsService::new(repo); - let config = svc.get_agent_config(PluginType::Weixin).await.unwrap(); + let config = svc.get_agent_config(TEST_USER_ID, PluginType::Weixin).await.unwrap(); assert_eq!(config.agent_type, "openclaw-gateway"); assert!(config.backend.is_none()); } @@ -815,7 +944,7 @@ mod tests { let overlay_repo: Arc = Arc::new(MockAssistantOverlayRepo { rows: vec![] }); let svc = ChannelSettingsService::new(repo).with_assistant_repos(definition_repo, overlay_repo); - let config = svc.get_agent_config(PluginType::Telegram).await.unwrap(); + let config = svc.get_agent_config(TEST_USER_ID, PluginType::Telegram).await.unwrap(); assert_eq!(config.agent_type, "acp"); assert_eq!(config.backend.as_deref(), Some("claude")); } @@ -835,7 +964,7 @@ mod tests { }); let svc = ChannelSettingsService::new(repo).with_assistant_repos(definition_repo, overlay_repo); - let config = svc.get_agent_config(PluginType::Telegram).await.unwrap(); + let config = svc.get_agent_config(TEST_USER_ID, PluginType::Telegram).await.unwrap(); assert_eq!(config.agent_type, "acp"); assert_eq!(config.backend.as_deref(), Some("codex")); } @@ -851,7 +980,10 @@ mod tests { let overlay_repo: Arc = Arc::new(MockAssistantOverlayRepo { rows: vec![] }); let svc = ChannelSettingsService::new(repo).with_assistant_repos(definition_repo, overlay_repo); - let err = svc.get_agent_config(PluginType::Telegram).await.unwrap_err(); + let err = svc + .get_agent_config(TEST_USER_ID, PluginType::Telegram) + .await + .unwrap_err(); assert!(matches!(err, ChannelError::InvalidConfig(_))); assert!( err.to_string().contains("missing-assistant"), @@ -866,7 +998,7 @@ mod tests { let repo = Arc::new(MockPrefRepo::new()); let svc = ChannelSettingsService::new(repo); - let config = svc.get_model_config(PluginType::Telegram).await.unwrap(); + let config = svc.get_model_config(TEST_USER_ID, PluginType::Telegram).await.unwrap(); assert!(config.is_none()); } @@ -878,7 +1010,11 @@ mod tests { )])); let svc = ChannelSettingsService::new(repo); - let config = svc.get_model_config(PluginType::Weixin).await.unwrap().unwrap(); + let config = svc + .get_model_config(TEST_USER_ID, PluginType::Weixin) + .await + .unwrap() + .unwrap(); assert_eq!(config.provider_id, "490fdb4e"); assert_eq!(config.use_model.as_deref(), Some("global.anthropic.claude-opus-4-6-v1")); } @@ -891,7 +1027,7 @@ mod tests { )])); let svc = ChannelSettingsService::new(repo); - let config = svc.get_model_config(PluginType::Telegram).await.unwrap(); + let config = svc.get_model_config(TEST_USER_ID, PluginType::Telegram).await.unwrap(); assert!(config.is_none()); } @@ -901,6 +1037,7 @@ mod tests { let svc = ChannelSettingsService::new(repo.clone()); svc.set_assistant_setting( + TEST_USER_ID, PluginType::Telegram, &ChannelAssistantSettingRequest { assistant_id: "assistant-1".into(), @@ -910,7 +1047,10 @@ mod tests { .await .unwrap(); - let stored = repo.get_by_keys(&["assistant.telegram.agent"]).await.unwrap(); + let stored = repo + .get_by_keys(TEST_USER_ID, &["assistant.telegram.agent"]) + .await + .unwrap(); let payload = serde_json::from_str::(&stored[0].value).unwrap(); assert_eq!(payload["assistant_id"], "assistant-1"); @@ -926,6 +1066,7 @@ mod tests { let svc = ChannelSettingsService::new(repo.clone()); svc.set_assistant_setting( + TEST_USER_ID, PluginType::Lark, &ChannelAssistantSettingRequest { assistant_id: " legacy-custom ".into(), @@ -935,7 +1076,7 @@ mod tests { .await .unwrap(); - let stored = repo.get_by_keys(&["assistant.lark.agent"]).await.unwrap(); + let stored = repo.get_by_keys(TEST_USER_ID, &["assistant.lark.agent"]).await.unwrap(); let payload = serde_json::from_str::(&stored[0].value).unwrap(); assert_eq!(payload["assistant_id"], "legacy-custom"); @@ -953,7 +1094,11 @@ mod tests { )])); let svc = ChannelSettingsService::new(repo); - let setting = svc.get_assistant_setting(PluginType::Telegram).await.unwrap().unwrap(); + let setting = svc + .get_assistant_setting(TEST_USER_ID, PluginType::Telegram) + .await + .unwrap() + .unwrap(); assert_eq!(setting.assistant_id.as_deref(), Some("legacy-custom")); assert!(setting.custom_agent_id.is_none()); @@ -974,7 +1119,11 @@ mod tests { let overlay_repo = Arc::new(MockAssistantOverlayRepo { rows: vec![] }); let svc = ChannelSettingsService::new(repo).with_assistant_repos(definition_repo, overlay_repo); - let setting = svc.get_assistant_setting(PluginType::Telegram).await.unwrap().unwrap(); + let setting = svc + .get_assistant_setting(TEST_USER_ID, PluginType::Telegram) + .await + .unwrap() + .unwrap(); assert_eq!(setting.assistant_id.as_deref(), Some("bare-aionrs")); assert!(setting.custom_agent_id.is_none()); @@ -991,7 +1140,11 @@ mod tests { )])); let svc = ChannelSettingsService::new(repo); - let setting = svc.get_assistant_setting(PluginType::Lark).await.unwrap().unwrap(); + let setting = svc + .get_assistant_setting(TEST_USER_ID, PluginType::Lark) + .await + .unwrap() + .unwrap(); assert!(setting.assistant_id.is_none()); assert!(setting.custom_agent_id.is_none()); @@ -1012,7 +1165,11 @@ mod tests { let overlay_repo = Arc::new(MockAssistantOverlayRepo { rows: vec![] }); let svc = ChannelSettingsService::new(repo).with_assistant_repos(definition_repo, overlay_repo); - let setting = svc.get_assistant_setting(PluginType::Lark).await.unwrap().unwrap(); + let setting = svc + .get_assistant_setting(TEST_USER_ID, PluginType::Lark) + .await + .unwrap() + .unwrap(); assert_eq!(setting.assistant_id.as_deref(), Some("bare-codex")); assert!(setting.custom_agent_id.is_none()); @@ -1029,7 +1186,10 @@ mod tests { )])); let svc = ChannelSettingsService::new(repo); - let settings = svc.get_platform_settings(PluginType::Telegram).await.unwrap(); + let settings = svc + .get_platform_settings(TEST_USER_ID, PluginType::Telegram) + .await + .unwrap(); let assistant = settings.assistant.expect("assistant settings"); assert_eq!(assistant.assistant_id.as_deref(), Some("legacy-custom")); @@ -1051,7 +1211,10 @@ mod tests { let overlay_repo = Arc::new(MockAssistantOverlayRepo { rows: vec![] }); let svc = ChannelSettingsService::new(repo).with_assistant_repos(definition_repo, overlay_repo); - let settings = svc.get_platform_settings(PluginType::Telegram).await.unwrap(); + let settings = svc + .get_platform_settings(TEST_USER_ID, PluginType::Telegram) + .await + .unwrap(); let assistant = settings.assistant.expect("assistant settings"); assert_eq!(assistant.assistant_id.as_deref(), Some("bare-aionrs")); @@ -1073,7 +1236,10 @@ mod tests { let overlay_repo = Arc::new(MockAssistantOverlayRepo { rows: vec![] }); let svc = ChannelSettingsService::new(repo).with_assistant_repos(definition_repo, overlay_repo); - let settings = svc.get_platform_settings(PluginType::Telegram).await.unwrap(); + let settings = svc + .get_platform_settings(TEST_USER_ID, PluginType::Telegram) + .await + .unwrap(); let assistant = settings.assistant.expect("assistant settings"); assert_eq!(assistant.assistant_id.as_deref(), Some("bare-codex")); diff --git a/crates/aionui-channel/src/manager.rs b/crates/aionui-channel/src/manager.rs index c673c76c1..3cc32de50 100644 --- a/crates/aionui-channel/src/manager.rs +++ b/crates/aionui-channel/src/manager.rs @@ -28,8 +28,8 @@ pub struct ChannelManager { repo: Arc, broadcaster: Arc, encryption_key: [u8; 32], - /// Active plugin instances keyed by plugin ID. - plugins: DashMap>, + /// Active plugin instances keyed by owner user ID + plugin ID. + plugins: DashMap>, /// Sender for incoming messages from all plugins. /// The `ActionExecutor` holds the receiving end. message_tx: mpsc::Sender, @@ -37,6 +37,21 @@ pub struct ChannelManager { confirm_tx: mpsc::Sender<(String, String)>, } +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct ChannelRuntimeKey { + owner_user_id: String, + plugin_id: String, +} + +impl ChannelRuntimeKey { + fn new(owner_user_id: &str, plugin_id: &str) -> Self { + Self { + owner_user_id: owner_user_id.to_owned(), + plugin_id: plugin_id.to_owned(), + } + } +} + /// Factory function type for creating plugin instances. /// /// Platform-specific implementations register their factory via @@ -77,12 +92,15 @@ impl ChannelManager { /// Returns the status of all registered plugins from the database. /// /// Merges DB state with live runtime status for active plugins. - pub async fn get_plugin_status(&self) -> Result, ChannelError> { - let rows = self.repo.get_all_plugins().await?; + pub async fn get_plugin_status(&self, owner_user_id: &str) -> Result, ChannelError> { + let rows = self.repo.get_all_plugins(owner_user_id).await?; let statuses: Vec = rows .into_iter() .map(|row| { - let live_status = self.plugins.get(&row.id).map(|p| p.status().to_string()); + let live_status = self + .plugins + .get(&Self::runtime_key(owner_user_id, &row.id)) + .map(|p| p.status().to_string()); self.row_to_status_response(&row, live_status) }) .collect(); @@ -102,6 +120,7 @@ impl ChannelManager { /// - `factory`: Function to create the platform-specific plugin instance pub async fn enable_plugin( &self, + owner_user_id: &str, plugin_id: &str, config_value: &serde_json::Value, factory: &PluginFactory, @@ -115,13 +134,10 @@ impl ChannelManager { // are supplied. let config: PluginConfig = match Self::config_with_credentials(config_value)? { Some(config) => config, - None => self.load_stored_config(plugin_id).await?, + None => self.load_stored_config(owner_user_id, plugin_id).await?, }; - // Stop existing plugin if running - if self.plugins.contains_key(plugin_id) { - self.stop_plugin(plugin_id).await; - } + self.stop_plugin(owner_user_id, plugin_id).await; // Encrypt config for storage let config_json = serde_json::to_string(&config)?; @@ -132,6 +148,7 @@ impl ChannelManager { let now = now_ms(); let row = ChannelPluginRow { id: plugin_id.to_owned(), + owner_user_id: owner_user_id.to_owned(), r#type: plugin_type.to_string(), name: self.default_plugin_name(plugin_type), enabled: true, @@ -141,26 +158,23 @@ impl ChannelManager { created_at: now, updated_at: now, }; - self.repo.upsert_plugin(&row).await?; + self.repo.upsert_plugin(owner_user_id, &row).await?; // Create and start plugin instance let mut plugin = factory(plugin_type) .ok_or_else(|| ChannelError::InvalidPluginType(format!("No implementation for {plugin_type}")))?; - let callbacks = PluginCallbacks { - message_tx: self.message_tx.clone(), - confirm_tx: self.confirm_tx.clone(), - }; + let callbacks = self.callbacks_for_owner(owner_user_id); if let Err(e) = plugin.initialize(config, callbacks).await { - self.update_plugin_error(plugin_id, &e.to_string()).await; - self.broadcast_status_change(plugin_id).await; + self.update_plugin_error(owner_user_id, plugin_id, &e.to_string()).await; + self.broadcast_status_change(owner_user_id, plugin_id).await; return Err(e); } if let Err(e) = plugin.start().await { - self.update_plugin_error(plugin_id, &e.to_string()).await; - self.broadcast_status_change(plugin_id).await; + self.update_plugin_error(owner_user_id, plugin_id, &e.to_string()).await; + self.broadcast_status_change(owner_user_id, plugin_id).await; return Err(e); } @@ -170,13 +184,15 @@ impl ChannelManager { last_connected: Some(now_ms()), enabled: None, }; - self.repo.update_plugin_status(plugin_id, ¶ms).await?; + self.repo + .update_plugin_status(owner_user_id, plugin_id, ¶ms) + .await?; // Store active instance - self.plugins.insert(plugin_id.to_owned(), plugin); + self.plugins.insert(Self::runtime_key(owner_user_id, plugin_id), plugin); - info!(plugin_id = %plugin_id, "plugin enabled and started"); - self.broadcast_status_change(plugin_id).await; + info!(owner_user_id = %owner_user_id, plugin_id = %plugin_id, "plugin enabled and started"); + self.broadcast_status_change(owner_user_id, plugin_id).await; Ok(()) } @@ -187,22 +203,22 @@ impl ChannelManager { /// can behave consistently and survive restarts. pub async fn enable_extension_plugin( &self, + owner_user_id: &str, plugin_id: &str, plugin_name: &str, config: &PluginConfig, ) -> Result<(), ChannelError> { - if self.plugins.contains_key(plugin_id) { - self.stop_plugin(plugin_id).await; - } + self.stop_plugin(owner_user_id, plugin_id).await; let config_json = serde_json::to_string(config)?; let encrypted_config = encrypt_string(&config_json, &self.encryption_key) .map_err(|e| ChannelError::EncryptionFailed(e.to_string()))?; let now = now_ms(); - let existing = self.repo.get_plugin(plugin_id).await?; + let existing = self.repo.get_plugin(owner_user_id, plugin_id).await?; let row = ChannelPluginRow { id: plugin_id.to_owned(), + owner_user_id: owner_user_id.to_owned(), r#type: plugin_id.to_owned(), name: plugin_name.to_owned(), enabled: true, @@ -212,10 +228,10 @@ impl ChannelManager { created_at: existing.as_ref().map(|row| row.created_at).unwrap_or(now), updated_at: now, }; - self.repo.upsert_plugin(&row).await?; + self.repo.upsert_plugin(owner_user_id, &row).await?; - info!(plugin_id = %plugin_id, "extension plugin enabled (metadata-only mode)"); - self.broadcast_status_change(plugin_id).await; + info!(owner_user_id = %owner_user_id, plugin_id = %plugin_id, "extension plugin enabled (metadata-only mode)"); + self.broadcast_status_change(owner_user_id, plugin_id).await; Ok(()) } @@ -223,9 +239,9 @@ impl ChannelManager { /// the active instance. /// /// Idempotent — disabling an already-disabled plugin is a no-op. - pub async fn disable_plugin(&self, plugin_id: &str) -> Result<(), ChannelError> { + pub async fn disable_plugin(&self, owner_user_id: &str, plugin_id: &str) -> Result<(), ChannelError> { // Stop running instance if any - self.stop_plugin(plugin_id).await; + self.stop_plugin(owner_user_id, plugin_id).await; // Update DB let params = UpdatePluginStatusParams { @@ -233,10 +249,12 @@ impl ChannelManager { last_connected: None, enabled: Some(false), }; - self.repo.update_plugin_status(plugin_id, ¶ms).await?; + self.repo + .update_plugin_status(owner_user_id, plugin_id, ¶ms) + .await?; - info!(plugin_id = %plugin_id, "plugin disabled"); - self.broadcast_status_change(plugin_id).await; + info!(owner_user_id = %owner_user_id, plugin_id = %plugin_id, "plugin disabled"); + self.broadcast_status_change(owner_user_id, plugin_id).await; Ok(()) } @@ -279,8 +297,8 @@ impl ChannelManager { /// Reads all enabled plugins from DB, decrypts their config, and /// starts them. Errors on individual plugins are logged but don't /// prevent other plugins from starting. - pub async fn restore_plugins(&self, factory: &PluginFactory) -> Result<(), ChannelError> { - let rows = self.repo.get_all_plugins().await?; + pub async fn restore_plugins(&self, owner_user_id: &str, factory: &PluginFactory) -> Result<(), ChannelError> { + let rows = self.repo.get_all_plugins(owner_user_id).await?; let enabled: Vec = rows.into_iter().filter(|r| r.enabled).collect(); if enabled.is_empty() { @@ -288,7 +306,7 @@ impl ChannelManager { return Ok(()); } - info!(count = enabled.len(), "restoring enabled plugins"); + info!(owner_user_id = %owner_user_id, count = enabled.len(), "restoring enabled plugins"); for row in enabled { if PluginType::from_str_opt(&row.r#type).is_none() { @@ -297,17 +315,18 @@ impl ChannelManager { plugin_type = %row.r#type, "skipping extension plugin runtime restore; metadata-only mode" ); - self.broadcast_status_change(&row.id).await; + self.broadcast_status_change(owner_user_id, &row.id).await; continue; } - if let Err(e) = self.restore_single_plugin(&row, factory).await { + if let Err(e) = self.restore_single_plugin(owner_user_id, &row, factory).await { warn!( + owner_user_id = %owner_user_id, plugin_id = %row.id, error = %e, "failed to restore plugin, marking as error" ); - self.update_plugin_error(&row.id, &e.to_string()).await; - self.broadcast_status_change(&row.id).await; + self.update_plugin_error(owner_user_id, &row.id, &e.to_string()).await; + self.broadcast_status_change(owner_user_id, &row.id).await; } } @@ -318,23 +337,45 @@ impl ChannelManager { /// /// Called during application shutdown. pub async fn shutdown(&self) { - let keys: Vec = self.plugins.iter().map(|entry| entry.key().clone()).collect(); + let keys: Vec = self.plugins.iter().map(|entry| entry.key().clone()).collect(); for key in keys { - self.stop_plugin(&key).await; + self.stop_plugin_by_key(&key).await; } info!("all plugins shut down"); } + /// Stops all active plugin connections for one owner user. + /// + /// Called when an external Core session is revoked so no channel runtime + /// continues handling messages for the old account. + pub async fn shutdown_for_user(&self, owner_user_id: &str) -> usize { + let keys: Vec = self + .plugins + .iter() + .filter(|entry| entry.key().owner_user_id == owner_user_id) + .map(|entry| entry.key().clone()) + .collect(); + let stopped = keys.len(); + + for key in keys { + self.stop_plugin_by_key(&key).await; + } + if stopped > 0 { + info!(owner_user_id = %owner_user_id, stopped, "channel plugins shut down for user"); + } + stopped + } + /// Returns the number of currently active (in-memory) plugins. pub fn active_plugin_count(&self) -> usize { self.plugins.len() } /// Checks whether a specific plugin is currently running. - pub fn is_plugin_running(&self, plugin_id: &str) -> bool { + pub fn is_plugin_running(&self, owner_user_id: &str, plugin_id: &str) -> bool { self.plugins - .get(plugin_id) + .get(&Self::runtime_key(owner_user_id, plugin_id)) .map(|p| p.status() == PluginStatus::Running) .unwrap_or(false) } @@ -345,13 +386,14 @@ impl ChannelManager { /// to the correct platform plugin. pub async fn send_message( &self, + owner_user_id: &str, plugin_id: &str, chat_id: &str, message: crate::types::UnifiedOutgoingMessage, ) -> Result { let plugin = self .plugins - .get(plugin_id) + .get(&Self::runtime_key(owner_user_id, plugin_id)) .ok_or_else(|| ChannelError::PluginNotFound(plugin_id.to_owned()))?; plugin.send_message(chat_id, message).await } @@ -359,6 +401,7 @@ impl ChannelManager { /// Edits an existing message through a specific plugin. pub async fn edit_message( &self, + owner_user_id: &str, plugin_id: &str, chat_id: &str, message_id: &str, @@ -366,7 +409,7 @@ impl ChannelManager { ) -> Result<(), ChannelError> { let plugin = self .plugins - .get(plugin_id) + .get(&Self::runtime_key(owner_user_id, plugin_id)) .ok_or_else(|| ChannelError::PluginNotFound(plugin_id.to_owned()))?; plugin.edit_message(chat_id, message_id, message).await } @@ -403,10 +446,10 @@ impl ChannelManager { /// Used when an enable request omits credentials and the stored /// configuration should be reused (Settings re-enable toggle). Returns /// `InvalidConfig` when there is no stored config to fall back to. - async fn load_stored_config(&self, plugin_id: &str) -> Result { + async fn load_stored_config(&self, owner_user_id: &str, plugin_id: &str) -> Result { let row = self .repo - .get_plugin(plugin_id) + .get_plugin(owner_user_id, plugin_id) .await? .filter(|row| !row.config.is_empty()) .ok_or_else(|| { @@ -422,8 +465,38 @@ impl ChannelManager { } /// Stops and removes an active plugin instance. - async fn stop_plugin(&self, plugin_id: &str) { - if let Some((_, mut plugin)) = self.plugins.remove(plugin_id) { + fn runtime_key(owner_user_id: &str, plugin_id: &str) -> ChannelRuntimeKey { + ChannelRuntimeKey::new(owner_user_id, plugin_id) + } + + fn callbacks_for_owner(&self, owner_user_id: &str) -> PluginCallbacks { + let (plugin_msg_tx, mut plugin_msg_rx) = mpsc::channel::(64); + let message_tx = self.message_tx.clone(); + let owner_user_id = owner_user_id.to_owned(); + tokio::spawn(async move { + while let Some(mut msg) = plugin_msg_rx.recv().await { + msg.owner_user_id = Some(owner_user_id.clone()); + if message_tx.send(msg).await.is_err() { + break; + } + } + }); + + PluginCallbacks { + message_tx: plugin_msg_tx, + confirm_tx: self.confirm_tx.clone(), + } + } + + async fn stop_plugin(&self, owner_user_id: &str, plugin_id: &str) { + let key = Self::runtime_key(owner_user_id, plugin_id); + self.stop_plugin_by_key(&key).await; + } + + /// Stops and removes an active plugin instance by runtime key. + async fn stop_plugin_by_key(&self, key: &ChannelRuntimeKey) { + if let Some((_, mut plugin)) = self.plugins.remove(key) { + let plugin_id = plugin.plugin_type().to_string(); if let Err(e) = plugin.stop().await { warn!( plugin_id = %plugin_id, @@ -436,7 +509,12 @@ impl ChannelManager { } /// Restores a single plugin from its DB row. - async fn restore_single_plugin(&self, row: &ChannelPluginRow, factory: &PluginFactory) -> Result<(), ChannelError> { + async fn restore_single_plugin( + &self, + owner_user_id: &str, + row: &ChannelPluginRow, + factory: &PluginFactory, + ) -> Result<(), ChannelError> { let plugin_type = PluginType::from_str_opt(&row.r#type).ok_or_else(|| ChannelError::InvalidPluginType(row.r#type.clone()))?; @@ -448,10 +526,7 @@ impl ChannelManager { let mut plugin = factory(plugin_type) .ok_or_else(|| ChannelError::InvalidPluginType(format!("No implementation for {plugin_type}")))?; - let callbacks = PluginCallbacks { - message_tx: self.message_tx.clone(), - confirm_tx: self.confirm_tx.clone(), - }; + let callbacks = self.callbacks_for_owner(owner_user_id); plugin.initialize(config, callbacks).await?; plugin.start().await?; @@ -462,23 +537,24 @@ impl ChannelManager { last_connected: Some(now_ms()), enabled: None, }; - self.repo.update_plugin_status(&row.id, ¶ms).await?; + self.repo.update_plugin_status(owner_user_id, &row.id, ¶ms).await?; - self.plugins.insert(row.id.clone(), plugin); - info!(plugin_id = %row.id, "plugin restored"); - self.broadcast_status_change(&row.id).await; + self.plugins.insert(Self::runtime_key(owner_user_id, &row.id), plugin); + info!(owner_user_id = %owner_user_id, plugin_id = %row.id, "plugin restored"); + self.broadcast_status_change(owner_user_id, &row.id).await; Ok(()) } /// Updates a plugin to error status in the DB. - async fn update_plugin_error(&self, plugin_id: &str, error_msg: &str) { + async fn update_plugin_error(&self, owner_user_id: &str, plugin_id: &str, error_msg: &str) { let params = UpdatePluginStatusParams { status: Some(PluginStatus::Error.to_string()), last_connected: None, enabled: None, }; - if let Err(e) = self.repo.update_plugin_status(plugin_id, ¶ms).await { + if let Err(e) = self.repo.update_plugin_status(owner_user_id, plugin_id, ¶ms).await { error!( + owner_user_id = %owner_user_id, plugin_id = %plugin_id, db_error = %e, original_error = %error_msg, @@ -488,15 +564,16 @@ impl ChannelManager { } /// Broadcasts a `channel.plugin-status-changed` event. - async fn broadcast_status_change(&self, plugin_id: &str) { - let row = match self.repo.get_plugin(plugin_id).await { + async fn broadcast_status_change(&self, owner_user_id: &str, plugin_id: &str) { + let row = match self.repo.get_plugin(owner_user_id, plugin_id).await { Ok(Some(row)) => row, Ok(None) => { - warn!(plugin_id = %plugin_id, "plugin not found for status broadcast"); + warn!(owner_user_id = %owner_user_id, plugin_id = %plugin_id, "plugin not found for status broadcast"); return; } Err(e) => { warn!( + owner_user_id = %owner_user_id, plugin_id = %plugin_id, error = %e, "failed to read plugin for status broadcast" @@ -505,10 +582,14 @@ impl ChannelManager { } }; - let live_status = self.plugins.get(plugin_id).map(|p| p.status().to_string()); + let live_status = self + .plugins + .get(&Self::runtime_key(owner_user_id, plugin_id)) + .map(|p| p.status().to_string()); let status_response = self.row_to_status_response(&row, live_status); let payload = PluginStatusChangedPayload { + user_id: owner_user_id.to_owned(), plugin_id: plugin_id.to_owned(), status: status_response, }; @@ -525,7 +606,9 @@ impl ChannelManager { /// Converts a DB row + optional live status to a `PluginStatusResponse`. fn row_to_status_response(&self, row: &ChannelPluginRow, live_status: Option) -> PluginStatusResponse { - let is_running = self.plugins.contains_key(&row.id); + let is_running = self + .plugins + .contains_key(&Self::runtime_key(&row.owner_user_id, &row.id)); let has_token = !row.config.is_empty(); PluginStatusResponse { plugin_id: row.id.clone(), @@ -560,21 +643,24 @@ impl ChannelManager { impl crate::stream_relay::ChannelSender for ChannelManager { async fn send_message( &self, + owner_user_id: &str, plugin_id: &str, chat_id: &str, message: crate::types::UnifiedOutgoingMessage, ) -> Result { - self.send_message(plugin_id, chat_id, message).await + self.send_message(owner_user_id, plugin_id, chat_id, message).await } async fn edit_message( &self, + owner_user_id: &str, plugin_id: &str, chat_id: &str, message_id: &str, message: crate::types::UnifiedOutgoingMessage, ) -> Result<(), crate::error::ChannelError> { - self.edit_message(plugin_id, chat_id, message_id, message).await + self.edit_message(owner_user_id, plugin_id, chat_id, message_id, message) + .await } } @@ -616,6 +702,7 @@ mod tests { } // ── Mock IChannelRepository ──────────────────────────────────────── + const OWNER_ID: &str = "owner-test"; struct MockRepo { plugins: Mutex>, @@ -635,16 +722,16 @@ mod tests { #[async_trait::async_trait] impl IChannelRepository for MockRepo { - async fn get_all_plugins(&self) -> Result, DbError> { + async fn get_all_plugins(&self, _owner_user_id: &str) -> Result, DbError> { Ok(self.plugins.lock().unwrap().clone()) } - async fn get_plugin(&self, id: &str) -> Result, DbError> { + async fn get_plugin(&self, _owner_user_id: &str, id: &str) -> Result, DbError> { let plugins = self.plugins.lock().unwrap(); Ok(plugins.iter().find(|p| p.id == id).cloned()) } - async fn upsert_plugin(&self, row: &ChannelPluginRow) -> Result<(), DbError> { + async fn upsert_plugin(&self, _owner_user_id: &str, row: &ChannelPluginRow) -> Result<(), DbError> { let mut plugins = self.plugins.lock().unwrap(); if let Some(existing) = plugins.iter_mut().find(|p| p.id == row.id) { *existing = row.clone(); @@ -654,7 +741,12 @@ mod tests { Ok(()) } - async fn update_plugin_status(&self, id: &str, params: &UpdatePluginStatusParams) -> Result<(), DbError> { + async fn update_plugin_status( + &self, + _owner_user_id: &str, + id: &str, + params: &UpdatePluginStatusParams, + ) -> Result<(), DbError> { let mut plugins = self.plugins.lock().unwrap(); if let Some(p) = plugins.iter_mut().find(|p| p.id == id) { if let Some(ref s) = params.status { @@ -673,7 +765,7 @@ mod tests { } } - async fn delete_plugin(&self, id: &str) -> Result<(), DbError> { + async fn delete_plugin(&self, _owner_user_id: &str, id: &str) -> Result<(), DbError> { let mut plugins = self.plugins.lock().unwrap(); let len_before = plugins.len(); plugins.retain(|p| p.id != id); @@ -685,67 +777,97 @@ mod tests { } // -- User CRUD (unused stubs) -- - async fn get_all_users(&self) -> Result, DbError> { + async fn get_all_users(&self, _owner_user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn get_user_by_platform(&self, _pid: &str, _pt: &str) -> Result, DbError> { + async fn get_user_by_platform( + &self, + _owner_user_id: &str, + _pid: &str, + _pt: &str, + ) -> Result, DbError> { Ok(None) } - async fn create_user(&self, _row: &AssistantUserRow) -> Result<(), DbError> { + async fn create_user(&self, _owner_user_id: &str, _row: &AssistantUserRow) -> Result<(), DbError> { Ok(()) } - async fn update_user_last_active(&self, _id: &str, _la: TimestampMs) -> Result<(), DbError> { + async fn update_user_last_active( + &self, + _owner_user_id: &str, + _id: &str, + _la: TimestampMs, + ) -> Result<(), DbError> { Ok(()) } - async fn delete_user(&self, _id: &str) -> Result<(), DbError> { + async fn delete_user(&self, _owner_user_id: &str, _id: &str) -> Result<(), DbError> { Ok(()) } // -- Session CRUD (unused stubs) -- - async fn get_all_sessions(&self) -> Result, DbError> { + async fn get_all_sessions(&self, _owner_user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn get_session(&self, _id: &str) -> Result, DbError> { + async fn get_session(&self, _owner_user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } async fn get_or_create_session( &self, + _owner_user_id: &str, _uid: &str, _cid: &str, new_row: &AssistantSessionRow, ) -> Result { Ok(new_row.clone()) } - async fn update_session_activity(&self, _id: &str, _la: TimestampMs) -> Result<(), DbError> { + async fn update_session_activity( + &self, + _owner_user_id: &str, + _id: &str, + _la: TimestampMs, + ) -> Result<(), DbError> { Ok(()) } - async fn update_session_conversation(&self, _id: &str, _cid: &str) -> Result<(), DbError> { + async fn update_session_conversation( + &self, + _owner_user_id: &str, + _id: &str, + _cid: &str, + ) -> Result<(), DbError> { Ok(()) } - async fn update_session_agent_type(&self, _id: &str, _at: &str) -> Result<(), DbError> { + async fn update_session_agent_type(&self, _owner_user_id: &str, _id: &str, _at: &str) -> Result<(), DbError> { Ok(()) } - async fn delete_sessions_by_user(&self, _uid: &str) -> Result<(), DbError> { + async fn delete_sessions_by_user(&self, _owner_user_id: &str, _uid: &str) -> Result<(), DbError> { Ok(()) } - async fn delete_session_by_user_chat(&self, _uid: &str, _cid: &str) -> Result<(), DbError> { + async fn delete_session_by_user_chat( + &self, + _owner_user_id: &str, + _uid: &str, + _cid: &str, + ) -> Result<(), DbError> { Ok(()) } // -- Pairing codes (unused stubs) -- - async fn create_pairing(&self, _row: &PairingCodeRow) -> Result<(), DbError> { + async fn create_pairing(&self, _owner_user_id: &str, _row: &PairingCodeRow) -> Result<(), DbError> { Ok(()) } - async fn get_pending_pairings(&self) -> Result, DbError> { + async fn get_pending_pairings(&self, _owner_user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn get_pairing_by_code(&self, _code: &str) -> Result, DbError> { + async fn get_pairing_by_code( + &self, + _owner_user_id: &str, + _code: &str, + ) -> Result, DbError> { Ok(None) } - async fn update_pairing_status(&self, _code: &str, _status: &str) -> Result<(), DbError> { + async fn update_pairing_status(&self, _owner_user_id: &str, _code: &str, _status: &str) -> Result<(), DbError> { Ok(()) } - async fn cleanup_expired_pairings(&self, _now: TimestampMs) -> Result { + async fn cleanup_expired_pairings(&self, _owner_user_id: &str, _now: TimestampMs) -> Result { Ok(0) } } @@ -934,7 +1056,7 @@ mod tests { #[tokio::test] async fn get_status_empty() { let (mgr, _repo, _bc) = make_manager(); - let statuses = mgr.get_plugin_status().await.unwrap(); + let statuses = mgr.get_plugin_status(OWNER_ID).await.unwrap(); assert!(statuses.is_empty()); } @@ -944,6 +1066,7 @@ mod tests { let now = now_ms(); repo.plugins.lock().unwrap().push(ChannelPluginRow { id: "telegram".into(), + owner_user_id: OWNER_ID.into(), r#type: "telegram".into(), name: "Telegram Bot".into(), enabled: true, @@ -954,7 +1077,7 @@ mod tests { updated_at: now, }); - let statuses = mgr.get_plugin_status().await.unwrap(); + let statuses = mgr.get_plugin_status(OWNER_ID).await.unwrap(); assert_eq!(statuses.len(), 1); assert_eq!(statuses[0].plugin_id, "telegram"); assert_eq!(statuses[0].plugin_type, "telegram"); @@ -968,7 +1091,7 @@ mod tests { let factory = make_factory(); // Enable the plugin (will set live status to Running) - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); @@ -980,7 +1103,7 @@ mod tests { } } - let statuses = mgr.get_plugin_status().await.unwrap(); + let statuses = mgr.get_plugin_status(OWNER_ID).await.unwrap(); assert_eq!(statuses.len(), 1); // Live status (running) should override DB status (stopped) assert_eq!(statuses[0].status.as_deref(), Some("running")); @@ -993,7 +1116,7 @@ mod tests { let (mgr, repo, _bc) = make_manager(); let factory = make_factory(); - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); @@ -1014,12 +1137,12 @@ mod tests { let (mgr, _repo, _bc) = make_manager(); let factory = make_factory(); - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); assert_eq!(mgr.active_plugin_count(), 1); - assert!(mgr.is_plugin_running("telegram")); + assert!(mgr.is_plugin_running(OWNER_ID, "telegram")); } #[tokio::test] @@ -1027,7 +1150,7 @@ mod tests { let (mgr, _repo, bc) = make_manager(); let factory = make_factory(); - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); @@ -1043,13 +1166,13 @@ mod tests { let (mgr, _repo, _bc) = make_manager(); let factory = make_factory(); - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); assert_eq!(mgr.active_plugin_count(), 1); // Re-enable should replace (stop old, start new) - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); assert_eq!(mgr.active_plugin_count(), 1); @@ -1061,7 +1184,7 @@ mod tests { let factory = make_factory(); let err = mgr - .enable_plugin("whatsapp", &make_test_config(), &factory) + .enable_plugin(OWNER_ID, "whatsapp", &make_test_config(), &factory) .await .unwrap_err(); assert!(matches!(err, ChannelError::InvalidPluginType(_))); @@ -1073,7 +1196,10 @@ mod tests { let factory = make_factory(); let bad_config = serde_json::json!({ "wrong": "shape" }); - let err = mgr.enable_plugin("telegram", &bad_config, &factory).await.unwrap_err(); + let err = mgr + .enable_plugin(OWNER_ID, "telegram", &bad_config, &factory) + .await + .unwrap_err(); assert!(matches!(err, ChannelError::InvalidConfig(_))); } @@ -1083,7 +1209,7 @@ mod tests { let factory = make_no_impl_factory(); let err = mgr - .enable_plugin("telegram", &make_test_config(), &factory) + .enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap_err(); assert!(matches!(err, ChannelError::InvalidPluginType(_))); @@ -1094,7 +1220,9 @@ mod tests { let (mgr, repo, _bc) = make_manager(); let factory = make_failing_init_factory(); - let err = mgr.enable_plugin("telegram", &make_test_config(), &factory).await; + let err = mgr + .enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) + .await; assert!(err.is_err()); // Plugin should not be in active map @@ -1110,7 +1238,9 @@ mod tests { let (mgr, repo, _bc) = make_manager(); let factory = make_failing_start_factory(); - let err = mgr.enable_plugin("telegram", &make_test_config(), &factory).await; + let err = mgr + .enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) + .await; assert!(err.is_err()); assert_eq!(mgr.active_plugin_count(), 0); @@ -1125,12 +1255,12 @@ mod tests { let (mgr, repo, _bc) = make_manager(); let factory = make_factory(); - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); - assert!(mgr.is_plugin_running("telegram")); + assert!(mgr.is_plugin_running(OWNER_ID, "telegram")); - mgr.disable_plugin("telegram").await.unwrap(); + mgr.disable_plugin(OWNER_ID, "telegram").await.unwrap(); assert_eq!(mgr.active_plugin_count(), 0); let plugins = repo.get_plugins(); @@ -1143,12 +1273,12 @@ mod tests { let (mgr, _repo, bc) = make_manager(); let factory = make_factory(); - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); bc.take_events(); // clear enable events - mgr.disable_plugin("telegram").await.unwrap(); + mgr.disable_plugin(OWNER_ID, "telegram").await.unwrap(); let events = bc.take_events(); assert!(!events.is_empty()); @@ -1161,6 +1291,7 @@ mod tests { // Manually insert a disabled plugin in DB repo.plugins.lock().unwrap().push(ChannelPluginRow { id: "telegram".into(), + owner_user_id: OWNER_ID.into(), r#type: "telegram".into(), name: "Telegram Bot".into(), enabled: false, @@ -1172,7 +1303,7 @@ mod tests { }); // Should not error - mgr.disable_plugin("telegram").await.unwrap(); + mgr.disable_plugin(OWNER_ID, "telegram").await.unwrap(); assert_eq!(mgr.active_plugin_count(), 0); } @@ -1237,6 +1368,7 @@ mod tests { repo.plugins.lock().unwrap().push(ChannelPluginRow { id: "telegram".into(), + owner_user_id: OWNER_ID.into(), r#type: "telegram".into(), name: "Telegram Bot".into(), enabled: false, @@ -1247,7 +1379,7 @@ mod tests { updated_at: now_ms(), }); - mgr.restore_plugins(&factory).await.unwrap(); + mgr.restore_plugins(OWNER_ID, &factory).await.unwrap(); assert_eq!(mgr.active_plugin_count(), 0); } @@ -1261,6 +1393,7 @@ mod tests { repo.plugins.lock().unwrap().push(ChannelPluginRow { id: "telegram".into(), + owner_user_id: OWNER_ID.into(), r#type: "telegram".into(), name: "Telegram Bot".into(), enabled: true, @@ -1271,9 +1404,9 @@ mod tests { updated_at: now_ms(), }); - mgr.restore_plugins(&factory).await.unwrap(); + mgr.restore_plugins(OWNER_ID, &factory).await.unwrap(); assert_eq!(mgr.active_plugin_count(), 1); - assert!(mgr.is_plugin_running("telegram")); + assert!(mgr.is_plugin_running(OWNER_ID, "telegram")); } #[tokio::test] @@ -1288,6 +1421,7 @@ mod tests { let mut plugins = repo.plugins.lock().unwrap(); plugins.push(ChannelPluginRow { id: "telegram".into(), + owner_user_id: OWNER_ID.into(), r#type: "telegram".into(), name: "Telegram Bot".into(), enabled: true, @@ -1299,6 +1433,7 @@ mod tests { }); plugins.push(ChannelPluginRow { id: "lark".into(), + owner_user_id: OWNER_ID.into(), r#type: "lark".into(), name: "Lark Bot".into(), enabled: true, @@ -1311,12 +1446,12 @@ mod tests { } let factory = make_factory(); - mgr.restore_plugins(&factory).await.unwrap(); + mgr.restore_plugins(OWNER_ID, &factory).await.unwrap(); // Telegram should have started, Lark should have failed assert_eq!(mgr.active_plugin_count(), 1); - assert!(mgr.is_plugin_running("telegram")); - assert!(!mgr.is_plugin_running("lark")); + assert!(mgr.is_plugin_running(OWNER_ID, "telegram")); + assert!(!mgr.is_plugin_running(OWNER_ID, "lark")); // Lark should have error status in DB let plugins = repo.get_plugins(); @@ -1328,7 +1463,7 @@ mod tests { async fn restore_empty_is_noop() { let (mgr, _repo, _bc) = make_manager(); let factory = make_factory(); - mgr.restore_plugins(&factory).await.unwrap(); + mgr.restore_plugins(OWNER_ID, &factory).await.unwrap(); assert_eq!(mgr.active_plugin_count(), 0); } @@ -1339,7 +1474,7 @@ mod tests { let (mgr, _repo, _bc) = make_manager(); let factory = make_factory(); - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); @@ -1349,7 +1484,9 @@ mod tests { "appSecret": "secret" } }); - mgr.enable_plugin("lark", &lark_config, &factory).await.unwrap(); + mgr.enable_plugin(OWNER_ID, "lark", &lark_config, &factory) + .await + .unwrap(); assert_eq!(mgr.active_plugin_count(), 2); @@ -1357,6 +1494,26 @@ mod tests { assert_eq!(mgr.active_plugin_count(), 0); } + #[tokio::test] + async fn shutdown_for_user_stops_only_matching_owner_plugins() { + let (mgr, _repo, _bc) = make_manager(); + let factory = make_factory(); + + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) + .await + .unwrap(); + mgr.enable_plugin("other-owner", "telegram", &make_test_config(), &factory) + .await + .unwrap(); + + assert_eq!(mgr.shutdown_for_user(OWNER_ID).await, 1); + assert_eq!(mgr.active_plugin_count(), 1); + assert!(!mgr.is_plugin_running(OWNER_ID, "telegram")); + assert!(mgr.is_plugin_running("other-owner", "telegram")); + + mgr.shutdown().await; + } + // ── send_message / edit_message ──────────────────────────────────── #[tokio::test] @@ -1364,12 +1521,12 @@ mod tests { let (mgr, _repo, _bc) = make_manager(); let factory = make_factory(); - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); let msg_id = mgr - .send_message("telegram", "chat_1", make_test_outgoing()) + .send_message(OWNER_ID, "telegram", "chat_1", make_test_outgoing()) .await .unwrap(); assert_eq!(msg_id, "mock_msg_id"); @@ -1379,7 +1536,7 @@ mod tests { async fn send_message_plugin_not_found() { let (mgr, _repo, _bc) = make_manager(); let err = mgr - .send_message("telegram", "chat_1", make_test_outgoing()) + .send_message(OWNER_ID, "telegram", "chat_1", make_test_outgoing()) .await .unwrap_err(); assert!(matches!(err, ChannelError::PluginNotFound(_))); @@ -1390,11 +1547,11 @@ mod tests { let (mgr, _repo, _bc) = make_manager(); let factory = make_factory(); - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); - mgr.edit_message("telegram", "chat_1", "msg_1", make_test_outgoing()) + mgr.edit_message(OWNER_ID, "telegram", "chat_1", "msg_1", make_test_outgoing()) .await .unwrap(); } @@ -1403,7 +1560,7 @@ mod tests { async fn edit_message_plugin_not_found() { let (mgr, _repo, _bc) = make_manager(); let err = mgr - .edit_message("telegram", "chat_1", "msg_1", make_test_outgoing()) + .edit_message(OWNER_ID, "telegram", "chat_1", "msg_1", make_test_outgoing()) .await .unwrap_err(); assert!(matches!(err, ChannelError::PluginNotFound(_))); @@ -1411,6 +1568,14 @@ mod tests { // ── helper methods ───────────────────────────────────────────────── + #[test] + fn runtime_key_keeps_owner_and_plugin_boundaries() { + let first = ChannelRuntimeKey::new("owner:a", "plugin"); + let second = ChannelRuntimeKey::new("owner", "a:plugin"); + + assert_ne!(first, second); + } + #[tokio::test] async fn active_plugin_count_tracks_correctly() { let (mgr, _repo, _bc) = make_manager(); @@ -1418,19 +1583,19 @@ mod tests { assert_eq!(mgr.active_plugin_count(), 0); - mgr.enable_plugin("telegram", &make_test_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_test_config(), &factory) .await .unwrap(); assert_eq!(mgr.active_plugin_count(), 1); - mgr.disable_plugin("telegram").await.unwrap(); + mgr.disable_plugin(OWNER_ID, "telegram").await.unwrap(); assert_eq!(mgr.active_plugin_count(), 0); } #[tokio::test] async fn is_plugin_running_false_for_missing() { let (mgr, _repo, _bc) = make_manager(); - assert!(!mgr.is_plugin_running("nonexistent")); + assert!(!mgr.is_plugin_running(OWNER_ID, "nonexistent")); } #[test] diff --git a/crates/aionui-channel/src/message_service.rs b/crates/aionui-channel/src/message_service.rs index 839cce9db..62c5a5f21 100644 --- a/crates/aionui-channel/src/message_service.rs +++ b/crates/aionui-channel/src/message_service.rs @@ -27,7 +27,6 @@ pub struct ChannelMessageService { conversation_svc: Arc, task_manager: Arc, settings: Arc, - owner_user_id: String, } impl ChannelMessageService { @@ -35,13 +34,11 @@ impl ChannelMessageService { conversation_svc: Arc, task_manager: Arc, settings: Arc, - owner_user_id: String, ) -> Self { Self { conversation_svc, task_manager, settings, - owner_user_id, } } @@ -56,6 +53,7 @@ impl ChannelMessageService { /// relaying them to the IM platform. pub async fn send_to_agent( &self, + owner_user_id: &str, session: &AssistantSessionRow, text: &str, platform: PluginType, @@ -63,7 +61,10 @@ impl ChannelMessageService { // Ensure conversation exists let conversation_id = match &session.conversation_id { Some(cid) => cid.clone(), - None => self.create_conversation_for_session(session, platform).await?, + None => { + self.create_conversation_for_session(owner_user_id, session, platform) + .await? + } }; // Send message through ConversationService. `msg_id` is now @@ -77,7 +78,7 @@ impl ChannelMessageService { hidden: false, }; - let user_id = &self.owner_user_id; + let user_id = owner_user_id; // Channel relays need a stream subscription before the agent starts // emitting. `ConversationService::send_message` returns immediately // and builds cold agents in the background, so warm the conversation @@ -121,16 +122,17 @@ impl ChannelMessageService { /// for per-chat isolation. async fn create_conversation_for_session( &self, + owner_user_id: &str, session: &AssistantSessionRow, platform: PluginType, ) -> Result { let source = platform_to_source(platform); let agent_config = self .settings - .get_agent_config(platform) + .get_agent_config(owner_user_id, platform) .await .map_err(|e| ChannelError::MessageSendFailed(e.to_string()))?; - let assistant_setting = self.settings.get_assistant_setting(platform).await?; + let assistant_setting = self.settings.get_assistant_setting(owner_user_id, platform).await?; let assistant_id = assistant_setting .as_ref() .and_then(|setting| setting.assistant_id.as_deref()) @@ -141,7 +143,7 @@ impl ChannelMessageService { .map(str::trim) .filter(|name| !name.is_empty()) .map(ToOwned::to_owned); - let model_config = self.settings.get_model_config(platform).await?; + let model_config = self.settings.get_model_config(owner_user_id, platform).await?; let agent_type = parse_agent_type(&agent_config.agent_type)?; let model = resolved_model_to_provider(model_config.as_ref()); let mut extra = Self::build_channel_extra(if assistant_id.is_some() { @@ -182,7 +184,7 @@ impl ChannelMessageService { let response = self .conversation_svc - .create(&self.owner_user_id, req) + .create(owner_user_id, req) .await .map_err(|e| ChannelError::MessageSendFailed(e.to_string()))?; diff --git a/crates/aionui-channel/src/orchestrator.rs b/crates/aionui-channel/src/orchestrator.rs index db64e6148..092e025e1 100644 --- a/crates/aionui-channel/src/orchestrator.rs +++ b/crates/aionui-channel/src/orchestrator.rs @@ -67,6 +67,10 @@ impl ChannelOrchestrator { let chat_id = msg.chat_id.clone(); let plugin_id = platform.to_string(); let text = msg.content.text.clone(); + let message_owner_user_id = msg + .owner_user_id + .clone() + .or_else(|| self.action_executor.owner_user_id().map(ToOwned::to_owned)); let executor = Arc::clone(&self.action_executor); let msg_svc = Arc::clone(&self.message_service); @@ -76,9 +80,14 @@ impl ChannelOrchestrator { tokio::spawn(async move { match executor.handle_incoming_message(&msg).await { Ok(MessageResult::Action(response)) => { - send_action_response(&sender, &plugin_id, &chat_id, &response).await; + if let Some(owner_user_id) = message_owner_user_id.as_deref() { + send_action_response(&sender, owner_user_id, &plugin_id, &chat_id, &response).await; + } else { + warn!("dropping channel action response without owner user"); + } } Ok(MessageResult::Dispatched { + owner_user_id, session_id, conversation_id, }) => { @@ -86,6 +95,7 @@ impl ChannelOrchestrator { &msg_svc, &session_mgr, &sender, + &owner_user_id, &session_id, conversation_id.as_deref(), &text, @@ -108,6 +118,7 @@ impl ChannelOrchestrator { async fn send_action_response( sender: &Arc, + owner_user_id: &str, plugin_id: &str, chat_id: &str, response: &crate::types::ActionResponse, @@ -130,11 +141,13 @@ async fn send_action_response( match response.behavior { ActionBehavior::Edit => { if let Some(ref edit_id) = response.edit_message_id { - let _ = sender.edit_message(plugin_id, chat_id, edit_id, outgoing).await; + let _ = sender + .edit_message(owner_user_id, plugin_id, chat_id, edit_id, outgoing) + .await; } } _ => { - let _ = sender.send_message(plugin_id, chat_id, outgoing).await; + let _ = sender.send_message(owner_user_id, plugin_id, chat_id, outgoing).await; } } } @@ -145,6 +158,7 @@ async fn handle_dispatched( msg_svc: &Arc, session_mgr: &Arc, sender: &Arc, + owner_user_id: &str, session_id: &str, conversation_id: Option<&str>, text: &str, @@ -152,7 +166,7 @@ async fn handle_dispatched( plugin_id: &str, chat_id: &str, ) { - let session = match session_mgr.get_session_by_id(session_id).await { + let session = match session_mgr.get_session_by_id(owner_user_id, session_id).await { Ok(Some(s)) => s, Ok(None) => { warn!(session_id = %session_id, "session not found after dispatch"); @@ -164,7 +178,7 @@ async fn handle_dispatched( } }; - let send_result = match msg_svc.send_to_agent(&session, text, platform).await { + let send_result = match msg_svc.send_to_agent(owner_user_id, &session, text, platform).await { Ok(r) => r, Err(e) => { error!(error = %e, "failed to send to agent"); @@ -181,7 +195,7 @@ async fn handle_dispatched( reply_to_message_id: None, silent: None, }; - let _ = sender.send_message(plugin_id, chat_id, err_msg).await; + let _ = sender.send_message(owner_user_id, plugin_id, chat_id, err_msg).await; return; } }; @@ -189,7 +203,7 @@ async fn handle_dispatched( // Bind conversation to session if newly created if conversation_id.is_none() && let Err(e) = session_mgr - .bind_conversation(session_id, &send_result.conversation_id) + .bind_conversation(owner_user_id, session_id, &send_result.conversation_id) .await { warn!(error = %e, "failed to bind conversation to session"); @@ -198,6 +212,7 @@ async fn handle_dispatched( // Spawn stream relay if we got a subscription if let Some(rx) = send_result.stream_rx { let relay_config = RelayConfig { + owner_user_id: owner_user_id.to_owned(), platform, plugin_id: plugin_id.to_owned(), chat_id: chat_id.to_owned(), diff --git a/crates/aionui-channel/src/pairing.rs b/crates/aionui-channel/src/pairing.rs index a82bc3b24..545f5ac59 100644 --- a/crates/aionui-channel/src/pairing.rs +++ b/crates/aionui-channel/src/pairing.rs @@ -50,12 +50,14 @@ impl PairingService { /// marked as expired before creating the new one. pub async fn request_pairing( &self, + owner_user_id: &str, platform_user_id: &str, platform_type: &str, display_name: Option<&str>, ) -> Result { // Expire any existing pending codes for this user - self.expire_user_pending_codes(platform_user_id, platform_type).await?; + self.expire_user_pending_codes(owner_user_id, platform_user_id, platform_type) + .await?; let code = generate_pairing_code()?; let now = now_ms(); @@ -63,6 +65,7 @@ impl PairingService { let row = PairingCodeRow { code: code.clone(), + owner_user_id: owner_user_id.to_owned(), platform_user_id: platform_user_id.to_owned(), platform_type: platform_type.to_owned(), display_name: display_name.map(String::from), @@ -71,10 +74,10 @@ impl PairingService { status: PairingStatus::Pending.to_string(), }; - self.repo.create_pairing(&row).await?; + self.repo.create_pairing(owner_user_id, &row).await?; info!( - code = %code, + owner_user_id = %owner_user_id, platform_user_id = %platform_user_id, platform_type = %platform_type, "pairing code created" @@ -82,6 +85,7 @@ impl PairingService { // Broadcast event let payload = PairingRequestedPayload { + user_id: owner_user_id.to_owned(), code: code.clone(), platform_user_id: platform_user_id.to_owned(), platform_type: platform_type.to_owned(), @@ -101,14 +105,15 @@ impl PairingService { /// - Creates an `assistant_users` record /// - Updates the pairing status to `approved` /// - Broadcasts a `channel.user-authorized` event - pub async fn approve_pairing(&self, code: &str) -> Result<(), ChannelError> { - let row = self.get_valid_pending_pairing(code).await?; + pub async fn approve_pairing(&self, owner_user_id: &str, code: &str) -> Result<(), ChannelError> { + let row = self.get_valid_pending_pairing(owner_user_id, code).await?; let now = now_ms(); // Create user record let user_id = generate_id(); let user_row = AssistantUserRow { id: user_id.clone(), + owner_user_id: owner_user_id.to_owned(), platform_user_id: row.platform_user_id.clone(), platform_type: row.platform_type.clone(), display_name: row.display_name.clone(), @@ -116,15 +121,15 @@ impl PairingService { last_active: None, session_id: None, }; - self.repo.create_user(&user_row).await?; + self.repo.create_user(owner_user_id, &user_row).await?; // Update pairing status self.repo - .update_pairing_status(code, &PairingStatus::Approved.to_string()) + .update_pairing_status(owner_user_id, code, &PairingStatus::Approved.to_string()) .await?; info!( - code = %code, + owner_user_id = %owner_user_id, user_id = %user_id, platform_user_id = %row.platform_user_id, "pairing approved, user created" @@ -132,6 +137,7 @@ impl PairingService { // Broadcast event let payload = UserAuthorizedPayload { + user_id: owner_user_id.to_owned(), id: user_id, platform_user_id: row.platform_user_id, platform_type: row.platform_type, @@ -148,20 +154,20 @@ impl PairingService { /// /// Validates the code exists and is still pending (not expired or /// already processed), then marks it as rejected. - pub async fn reject_pairing(&self, code: &str) -> Result<(), ChannelError> { - let _row = self.get_valid_pending_pairing(code).await?; + pub async fn reject_pairing(&self, owner_user_id: &str, code: &str) -> Result<(), ChannelError> { + let _row = self.get_valid_pending_pairing(owner_user_id, code).await?; self.repo - .update_pairing_status(code, &PairingStatus::Rejected.to_string()) + .update_pairing_status(owner_user_id, code, &PairingStatus::Rejected.to_string()) .await?; - info!(code = %code, "pairing rejected"); + info!(owner_user_id = %owner_user_id, "pairing rejected"); Ok(()) } /// Returns all pending (not expired) pairing requests. - pub async fn get_pending_pairings(&self) -> Result, ChannelError> { - let rows = self.repo.get_pending_pairings().await?; + pub async fn get_pending_pairings(&self, owner_user_id: &str) -> Result, ChannelError> { + let rows = self.repo.get_pending_pairings(owner_user_id).await?; let now = now_ms(); // Filter out expired ones that haven't been cleaned up yet let active: Vec = rows.into_iter().filter(|r| r.expires_at > now).collect(); @@ -169,8 +175,16 @@ impl PairingService { } /// Checks whether a platform user is already authorized. - pub async fn is_user_authorized(&self, platform_user_id: &str, platform_type: &str) -> Result { - let user = self.repo.get_user_by_platform(platform_user_id, platform_type).await?; + pub async fn is_user_authorized( + &self, + owner_user_id: &str, + platform_user_id: &str, + platform_type: &str, + ) -> Result { + let user = self + .repo + .get_user_by_platform(owner_user_id, platform_user_id, platform_type) + .await?; Ok(user.is_some()) } @@ -179,23 +193,27 @@ impl PairingService { /// Returns `None` if the user is not authorized. pub async fn get_internal_user_id( &self, + owner_user_id: &str, platform_user_id: &str, platform_type: &str, ) -> Result, ChannelError> { - let user = self.repo.get_user_by_platform(platform_user_id, platform_type).await?; + let user = self + .repo + .get_user_by_platform(owner_user_id, platform_user_id, platform_type) + .await?; Ok(user.map(|u| u.id)) } /// Starts a background task that periodically cleans up expired /// pairing codes. Returns a `JoinHandle` that can be used to cancel /// the task on shutdown. - pub fn start_cleanup_timer(repo: Arc) -> JoinHandle<()> { + pub fn start_cleanup_timer(owner_user_id: String, repo: Arc) -> JoinHandle<()> { tokio::spawn(async move { let mut interval = tokio::time::interval(PAIRING_CLEANUP_INTERVAL); loop { interval.tick().await; let now = now_ms(); - match repo.cleanup_expired_pairings(now).await { + match repo.cleanup_expired_pairings(&owner_user_id, now).await { Ok(count) if count > 0 => { debug!(count, "cleaned up expired pairing codes"); } @@ -209,10 +227,10 @@ impl PairingService { } /// Validates that a pairing code exists, is pending, and not expired. - async fn get_valid_pending_pairing(&self, code: &str) -> Result { + async fn get_valid_pending_pairing(&self, owner_user_id: &str, code: &str) -> Result { let row = self .repo - .get_pairing_by_code(code) + .get_pairing_by_code(owner_user_id, code) .await? .ok_or_else(|| ChannelError::PairingNotFound(code.to_owned()))?; @@ -225,7 +243,7 @@ impl PairingService { // Mark as expired for consistency let _ = self .repo - .update_pairing_status(code, &PairingStatus::Expired.to_string()) + .update_pairing_status(owner_user_id, code, &PairingStatus::Expired.to_string()) .await; return Err(ChannelError::PairingExpired(code.to_owned())); } @@ -237,15 +255,20 @@ impl PairingService { /// /// Called before creating a new code to ensure only one active code /// per user at a time. - async fn expire_user_pending_codes(&self, platform_user_id: &str, platform_type: &str) -> Result<(), ChannelError> { - let pending = self.repo.get_pending_pairings().await?; + async fn expire_user_pending_codes( + &self, + owner_user_id: &str, + platform_user_id: &str, + platform_type: &str, + ) -> Result<(), ChannelError> { + let pending = self.repo.get_pending_pairings(owner_user_id).await?; for row in pending { if row.platform_user_id == platform_user_id && row.platform_type == platform_type { self.repo - .update_pairing_status(&row.code, &PairingStatus::Expired.to_string()) + .update_pairing_status(owner_user_id, &row.code, &PairingStatus::Expired.to_string()) .await?; debug!( - code = %row.code, + owner_user_id = %owner_user_id, "expired old pending code for user" ); } @@ -314,30 +337,36 @@ mod tests { impl IChannelRepository for MockRepo { // -- Plugin CRUD (unused stubs) -- - async fn get_all_plugins(&self) -> Result, DbError> { + async fn get_all_plugins(&self, _owner_user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn get_plugin(&self, _id: &str) -> Result, DbError> { + async fn get_plugin(&self, _owner_user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } - async fn upsert_plugin(&self, _row: &ChannelPluginRow) -> Result<(), DbError> { + async fn upsert_plugin(&self, _owner_user_id: &str, _row: &ChannelPluginRow) -> Result<(), DbError> { Ok(()) } - async fn update_plugin_status(&self, _id: &str, _params: &UpdatePluginStatusParams) -> Result<(), DbError> { + async fn update_plugin_status( + &self, + _owner_user_id: &str, + _id: &str, + _params: &UpdatePluginStatusParams, + ) -> Result<(), DbError> { Ok(()) } - async fn delete_plugin(&self, _id: &str) -> Result<(), DbError> { + async fn delete_plugin(&self, _owner_user_id: &str, _id: &str) -> Result<(), DbError> { Ok(()) } // -- User CRUD -- - async fn get_all_users(&self) -> Result, DbError> { + async fn get_all_users(&self, _owner_user_id: &str) -> Result, DbError> { Ok(self.users.lock().unwrap().clone()) } async fn get_user_by_platform( &self, + _owner_user_id: &str, platform_user_id: &str, platform_type: &str, ) -> Result, DbError> { @@ -348,7 +377,7 @@ mod tests { .cloned()) } - async fn create_user(&self, row: &AssistantUserRow) -> Result<(), DbError> { + async fn create_user(&self, _owner_user_id: &str, row: &AssistantUserRow) -> Result<(), DbError> { let mut users = self.users.lock().unwrap(); if users .iter() @@ -360,7 +389,12 @@ mod tests { Ok(()) } - async fn update_user_last_active(&self, id: &str, last_active: TimestampMs) -> Result<(), DbError> { + async fn update_user_last_active( + &self, + _owner_user_id: &str, + id: &str, + last_active: TimestampMs, + ) -> Result<(), DbError> { let mut users = self.users.lock().unwrap(); if let Some(u) = users.iter_mut().find(|u| u.id == id) { u.last_active = Some(last_active); @@ -370,7 +404,7 @@ mod tests { } } - async fn delete_user(&self, id: &str) -> Result<(), DbError> { + async fn delete_user(&self, _owner_user_id: &str, id: &str) -> Result<(), DbError> { let mut users = self.users.lock().unwrap(); let len_before = users.len(); users.retain(|u| u.id != id); @@ -383,39 +417,60 @@ mod tests { // -- Session CRUD (unused stubs) -- - async fn get_all_sessions(&self) -> Result, DbError> { + async fn get_all_sessions(&self, _owner_user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn get_session(&self, _id: &str) -> Result, DbError> { + async fn get_session(&self, _owner_user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } async fn get_or_create_session( &self, + _owner_user_id: &str, _user_id: &str, _chat_id: &str, new_row: &AssistantSessionRow, ) -> Result { Ok(new_row.clone()) } - async fn update_session_activity(&self, _id: &str, _last_activity: TimestampMs) -> Result<(), DbError> { + async fn update_session_activity( + &self, + _owner_user_id: &str, + _id: &str, + _last_activity: TimestampMs, + ) -> Result<(), DbError> { Ok(()) } - async fn update_session_conversation(&self, _id: &str, _conversation_id: &str) -> Result<(), DbError> { + async fn update_session_conversation( + &self, + _owner_user_id: &str, + _id: &str, + _conversation_id: &str, + ) -> Result<(), DbError> { Ok(()) } - async fn update_session_agent_type(&self, _id: &str, _agent_type: &str) -> Result<(), DbError> { + async fn update_session_agent_type( + &self, + _owner_user_id: &str, + _id: &str, + _agent_type: &str, + ) -> Result<(), DbError> { Ok(()) } - async fn delete_sessions_by_user(&self, _user_id: &str) -> Result<(), DbError> { + async fn delete_sessions_by_user(&self, _owner_user_id: &str, _user_id: &str) -> Result<(), DbError> { Ok(()) } - async fn delete_session_by_user_chat(&self, _user_id: &str, _chat_id: &str) -> Result<(), DbError> { + async fn delete_session_by_user_chat( + &self, + _owner_user_id: &str, + _user_id: &str, + _chat_id: &str, + ) -> Result<(), DbError> { Ok(()) } // -- Pairing codes -- - async fn create_pairing(&self, row: &PairingCodeRow) -> Result<(), DbError> { + async fn create_pairing(&self, _owner_user_id: &str, row: &PairingCodeRow) -> Result<(), DbError> { let mut pairings = self.pairings.lock().unwrap(); if pairings.iter().any(|p| p.code == row.code) { return Err(DbError::Conflict("duplicate code".into())); @@ -424,17 +479,21 @@ mod tests { Ok(()) } - async fn get_pending_pairings(&self) -> Result, DbError> { + async fn get_pending_pairings(&self, _owner_user_id: &str) -> Result, DbError> { let pairings = self.pairings.lock().unwrap(); Ok(pairings.iter().filter(|p| p.status == "pending").cloned().collect()) } - async fn get_pairing_by_code(&self, code: &str) -> Result, DbError> { + async fn get_pairing_by_code( + &self, + _owner_user_id: &str, + code: &str, + ) -> Result, DbError> { let pairings = self.pairings.lock().unwrap(); Ok(pairings.iter().find(|p| p.code == code).cloned()) } - async fn update_pairing_status(&self, code: &str, status: &str) -> Result<(), DbError> { + async fn update_pairing_status(&self, _owner_user_id: &str, code: &str, status: &str) -> Result<(), DbError> { let mut pairings = self.pairings.lock().unwrap(); if let Some(p) = pairings.iter_mut().find(|p| p.code == code) { p.status = status.to_owned(); @@ -444,7 +503,7 @@ mod tests { } } - async fn cleanup_expired_pairings(&self, now: TimestampMs) -> Result { + async fn cleanup_expired_pairings(&self, _owner_user_id: &str, now: TimestampMs) -> Result { let mut pairings = self.pairings.lock().unwrap(); let mut count = 0u64; for p in pairings.iter_mut() { @@ -458,6 +517,7 @@ mod tests { } // ── Helpers ──────────────────────────────────────────────────────── + const OWNER_ID: &str = "owner-test"; fn make_service() -> (PairingService, Arc, Arc) { let repo = Arc::new(MockRepo::new()); @@ -503,7 +563,10 @@ mod tests { #[tokio::test] async fn request_pairing_creates_code() { let (svc, repo, _bc) = make_service(); - let code = svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); + let code = svc + .request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); assert_eq!(code.len(), PAIRING_CODE_LENGTH); let pairings = repo.get_pairings(); @@ -518,7 +581,9 @@ mod tests { #[tokio::test] async fn request_pairing_broadcasts_event() { let (svc, _repo, bc) = make_service(); - svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); + svc.request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); let events = bc.take_events(); assert_eq!(events.len(), 1); @@ -532,7 +597,7 @@ mod tests { async fn request_pairing_sets_correct_expiry() { let (svc, repo, _bc) = make_service(); let before = now_ms(); - svc.request_pairing("u1", "lark", None).await.unwrap(); + svc.request_pairing(OWNER_ID, "u1", "lark", None).await.unwrap(); let after = now_ms(); let p = &repo.get_pairings()[0]; @@ -545,8 +610,14 @@ mod tests { async fn request_pairing_expires_old_code() { let (svc, repo, _bc) = make_service(); - let code1 = svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); - let code2 = svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); + let code1 = svc + .request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); + let code2 = svc + .request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); assert_ne!(code1, code2); @@ -560,7 +631,7 @@ mod tests { #[tokio::test] async fn request_pairing_no_display_name() { let (svc, repo, _bc) = make_service(); - svc.request_pairing("u1", "dingtalk", None).await.unwrap(); + svc.request_pairing(OWNER_ID, "u1", "dingtalk", None).await.unwrap(); let pairings = repo.get_pairings(); assert!(pairings[0].display_name.is_none()); @@ -571,9 +642,12 @@ mod tests { #[tokio::test] async fn approve_creates_user_and_updates_status() { let (svc, repo, _bc) = make_service(); - let code = svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); + let code = svc + .request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); - svc.approve_pairing(&code).await.unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); // Check pairing status let pairings = repo.get_pairings(); @@ -591,10 +665,13 @@ mod tests { #[tokio::test] async fn approve_broadcasts_user_authorized() { let (svc, _repo, bc) = make_service(); - let code = svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); + let code = svc + .request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); bc.take_events(); // clear request event - svc.approve_pairing(&code).await.unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); let events = bc.take_events(); assert_eq!(events.len(), 1); @@ -608,17 +685,17 @@ mod tests { #[tokio::test] async fn approve_nonexistent_code_returns_not_found() { let (svc, _repo, _bc) = make_service(); - let err = svc.approve_pairing("000000").await.unwrap_err(); + let err = svc.approve_pairing(OWNER_ID, "000000").await.unwrap_err(); assert!(matches!(err, ChannelError::PairingNotFound(_))); } #[tokio::test] async fn approve_already_approved_returns_already_processed() { let (svc, _repo, _bc) = make_service(); - let code = svc.request_pairing("tg_42", "telegram", None).await.unwrap(); - svc.approve_pairing(&code).await.unwrap(); + let code = svc.request_pairing(OWNER_ID, "tg_42", "telegram", None).await.unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); - let err = svc.approve_pairing(&code).await.unwrap_err(); + let err = svc.approve_pairing(OWNER_ID, &code).await.unwrap_err(); assert!(matches!(err, ChannelError::PairingAlreadyProcessed(_))); } @@ -628,6 +705,7 @@ mod tests { // Manually insert an already-expired code let row = PairingCodeRow { code: "999999".into(), + owner_user_id: OWNER_ID.into(), platform_user_id: "u1".into(), platform_type: "telegram".into(), display_name: None, @@ -637,7 +715,7 @@ mod tests { }; repo.pairings.lock().unwrap().push(row); - let err = svc.approve_pairing("999999").await.unwrap_err(); + let err = svc.approve_pairing(OWNER_ID, "999999").await.unwrap_err(); assert!(matches!(err, ChannelError::PairingExpired(_))); } @@ -646,9 +724,9 @@ mod tests { #[tokio::test] async fn reject_updates_status() { let (svc, repo, _bc) = make_service(); - let code = svc.request_pairing("tg_42", "telegram", None).await.unwrap(); + let code = svc.request_pairing(OWNER_ID, "tg_42", "telegram", None).await.unwrap(); - svc.reject_pairing(&code).await.unwrap(); + svc.reject_pairing(OWNER_ID, &code).await.unwrap(); let pairings = repo.get_pairings(); let p = pairings.iter().find(|p| p.code == code).unwrap(); @@ -658,17 +736,17 @@ mod tests { #[tokio::test] async fn reject_nonexistent_code_returns_not_found() { let (svc, _repo, _bc) = make_service(); - let err = svc.reject_pairing("000000").await.unwrap_err(); + let err = svc.reject_pairing(OWNER_ID, "000000").await.unwrap_err(); assert!(matches!(err, ChannelError::PairingNotFound(_))); } #[tokio::test] async fn reject_already_approved_returns_already_processed() { let (svc, _repo, _bc) = make_service(); - let code = svc.request_pairing("tg_42", "telegram", None).await.unwrap(); - svc.approve_pairing(&code).await.unwrap(); + let code = svc.request_pairing(OWNER_ID, "tg_42", "telegram", None).await.unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); - let err = svc.reject_pairing(&code).await.unwrap_err(); + let err = svc.reject_pairing(OWNER_ID, &code).await.unwrap_err(); assert!(matches!(err, ChannelError::PairingAlreadyProcessed(_))); } @@ -679,11 +757,12 @@ mod tests { let (svc, repo, _bc) = make_service(); // Insert valid pending code - svc.request_pairing("u1", "telegram", None).await.unwrap(); + svc.request_pairing(OWNER_ID, "u1", "telegram", None).await.unwrap(); // Insert manually expired code let expired_row = PairingCodeRow { code: "000001".into(), + owner_user_id: OWNER_ID.into(), platform_user_id: "u2".into(), platform_type: "lark".into(), display_name: None, @@ -693,7 +772,7 @@ mod tests { }; repo.pairings.lock().unwrap().push(expired_row); - let pending = svc.get_pending_pairings().await.unwrap(); + let pending = svc.get_pending_pairings(OWNER_ID).await.unwrap(); assert_eq!(pending.len(), 1); assert_eq!(pending[0].platform_user_id, "u1"); } @@ -701,7 +780,7 @@ mod tests { #[tokio::test] async fn get_pending_empty_when_none() { let (svc, _repo, _bc) = make_service(); - let pending = svc.get_pending_pairings().await.unwrap(); + let pending = svc.get_pending_pairings(OWNER_ID).await.unwrap(); assert!(pending.is_empty()); } @@ -710,17 +789,17 @@ mod tests { #[tokio::test] async fn unauthorized_user_returns_false() { let (svc, _repo, _bc) = make_service(); - let authorized = svc.is_user_authorized("tg_42", "telegram").await.unwrap(); + let authorized = svc.is_user_authorized(OWNER_ID, "tg_42", "telegram").await.unwrap(); assert!(!authorized); } #[tokio::test] async fn authorized_user_returns_true_after_approval() { let (svc, _repo, _bc) = make_service(); - let code = svc.request_pairing("tg_42", "telegram", None).await.unwrap(); - svc.approve_pairing(&code).await.unwrap(); + let code = svc.request_pairing(OWNER_ID, "tg_42", "telegram", None).await.unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); - let authorized = svc.is_user_authorized("tg_42", "telegram").await.unwrap(); + let authorized = svc.is_user_authorized(OWNER_ID, "tg_42", "telegram").await.unwrap(); assert!(authorized); } @@ -733,6 +812,7 @@ mod tests { // Insert manually expired pending code let expired_row = PairingCodeRow { code: "111111".into(), + owner_user_id: OWNER_ID.into(), platform_user_id: "u1".into(), platform_type: "telegram".into(), display_name: None, @@ -743,9 +823,9 @@ mod tests { repo.pairings.lock().unwrap().push(expired_row); // Insert valid pending code - svc.request_pairing("u2", "lark", None).await.unwrap(); + svc.request_pairing(OWNER_ID, "u2", "lark", None).await.unwrap(); - let count = repo.cleanup_expired_pairings(now_ms()).await.unwrap(); + let count = repo.cleanup_expired_pairings(OWNER_ID, now_ms()).await.unwrap(); assert_eq!(count, 1); let pairings = repo.get_pairings(); diff --git a/crates/aionui-channel/src/plugins/dingtalk/plugin.rs b/crates/aionui-channel/src/plugins/dingtalk/plugin.rs index 4e880ac19..43814ea19 100644 --- a/crates/aionui-channel/src/plugins/dingtalk/plugin.rs +++ b/crates/aionui-channel/src/plugins/dingtalk/plugin.rs @@ -689,6 +689,7 @@ async fn handle_bot_message(data_str: &str, message_tx: &mpsc::Sender ApiError { match err { DbError::NotFound(msg) => ApiError::NotFound(msg), + DbError::Conflict(msg) if msg.starts_with("CROSS_ACCOUNT_REFERENCE:") => { + ApiError::coded(StatusCode::CONFLICT, "CROSS_ACCOUNT_REFERENCE", msg, None) + } DbError::Conflict(msg) => ApiError::Conflict(msg), DbError::Query(e) => ApiError::Internal(format!("Database error: {e}")), DbError::Migration(e) => ApiError::Internal(format!("Migration error: {e}")), @@ -132,9 +137,10 @@ pub fn weixin_login_route(state: ChannelRouterState) -> Router { /// `GET /api/channel/plugins` — get status of all registered plugins. async fn get_plugin_status( State(state): State, + Extension(user): Extension, ) -> Result>>, ApiError> { - let statuses = state.manager.get_plugin_status().await?; - let extension_plugins = state.extension_registry.get_channel_plugins().await; + let statuses = state.manager.get_plugin_status(&user.id).await?; + let extension_plugins = state.extension_registry.get_channel_plugins_for_user(&user.id).await; let extension_map: HashMap = extension_plugins .into_iter() @@ -307,15 +313,16 @@ impl From<&ResolvedChannelPlugin> for ChannelExtensionMetaView { /// `POST /api/channel/plugins/enable` — enable a plugin with config. async fn enable_plugin( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - if let Some(extension_plugin) = resolve_extension_channel_plugin(&state, &req.plugin_id).await { + if let Some(extension_plugin) = resolve_extension_channel_plugin(&state, &user.id, &req.plugin_id).await { let config = build_extension_config(&extension_plugin, &req.config)?; match state .manager - .enable_extension_plugin(&req.plugin_id, &extension_plugin.name, &config) + .enable_extension_plugin(&user.id, &req.plugin_id, &extension_plugin.name, &config) .await { Ok(()) => { @@ -338,7 +345,7 @@ async fn enable_plugin( match state .manager - .enable_plugin(&req.plugin_id, &req.config, state.plugin_factory.as_ref()) + .enable_plugin(&user.id, &req.plugin_id, &req.config, state.plugin_factory.as_ref()) .await { Ok(()) => Ok(Json(ApiResponse::ok(BridgeResponse { @@ -360,14 +367,17 @@ async fn enable_plugin( /// `POST /api/channel/plugins/disable` — disable a plugin. async fn disable_plugin( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - if resolve_extension_channel_plugin(&state, &req.plugin_id).await.is_some() + if resolve_extension_channel_plugin(&state, &user.id, &req.plugin_id) + .await + .is_some() && state .repo - .get_plugin(&req.plugin_id) + .get_plugin(&user.id, &req.plugin_id) .await .map_err(db_error_to_api_error)? .is_none() @@ -379,7 +389,7 @@ async fn disable_plugin( }))); } - match state.manager.disable_plugin(&req.plugin_id).await { + match state.manager.disable_plugin(&user.id, &req.plugin_id).await { Ok(()) => Ok(Json(ApiResponse::ok(BridgeResponse { success: true, message: Some("Plugin disabled".into()), @@ -399,11 +409,12 @@ async fn disable_plugin( /// `POST /api/channel/plugins/test` — test plugin credentials. async fn test_plugin( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - if let Some(extension_plugin) = resolve_extension_channel_plugin(&state, &req.plugin_id).await { + if let Some(extension_plugin) = resolve_extension_channel_plugin(&state, &user.id, &req.plugin_id).await { let _config = build_extension_test_config(&extension_plugin, &req)?; return Ok(Json(ApiResponse::ok(TestPluginResponse { success: true, @@ -439,8 +450,9 @@ async fn test_plugin( /// `GET /api/channel/pairings` — get all pending pairing requests. async fn get_pending_pairings( State(state): State, + Extension(user): Extension, ) -> Result>>, ApiError> { - let rows = state.pairing_service.get_pending_pairings().await?; + let rows = state.pairing_service.get_pending_pairings(&user.id).await?; let responses: Vec = rows .into_iter() .map(|r| PairingRequestResponse { @@ -458,11 +470,12 @@ async fn get_pending_pairings( /// `POST /api/channel/pairings/approve` — approve a pairing request. async fn approve_pairing( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - state.pairing_service.approve_pairing(&req.code).await?; + state.pairing_service.approve_pairing(&user.id, &req.code).await?; Ok(Json(ApiResponse::ok(BridgeResponse { success: true, @@ -474,11 +487,12 @@ async fn approve_pairing( /// `POST /api/channel/pairings/reject` — reject a pairing request. async fn reject_pairing( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - state.pairing_service.reject_pairing(&req.code).await?; + state.pairing_service.reject_pairing(&user.id, &req.code).await?; Ok(Json(ApiResponse::ok(BridgeResponse { success: true, @@ -494,8 +508,13 @@ async fn reject_pairing( /// `GET /api/channel/users` — get all authorized users. async fn get_authorized_users( State(state): State, + Extension(user): Extension, ) -> Result>>, ApiError> { - let rows = state.repo.get_all_users().await.map_err(db_error_to_api_error)?; + let rows = state + .repo + .get_all_users(&user.id) + .await + .map_err(db_error_to_api_error)?; let responses: Vec = rows .into_iter() .map(|r| ChannelUserResponse { @@ -515,17 +534,21 @@ async fn get_authorized_users( /// Also cleans up the user's sessions. async fn revoke_user( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; // Clean up sessions first - state.session_manager.cleanup_user_sessions(&req.user_id).await?; + state + .session_manager + .cleanup_user_sessions(&user.id, &req.user_id) + .await?; // Delete user record state .repo - .delete_user(&req.user_id) + .delete_user(&user.id, &req.user_id) .await .map_err(db_error_to_api_error)?; @@ -543,8 +566,9 @@ async fn revoke_user( /// `GET /api/channel/sessions` — get all active sessions. async fn get_active_sessions( State(state): State, + Extension(user): Extension, ) -> Result>>, ApiError> { - let rows = state.session_manager.get_active_sessions().await?; + let rows = state.session_manager.get_active_sessions(&user.id).await?; let responses: Vec = rows .into_iter() .map(|r| ChannelSessionResponse { @@ -568,18 +592,20 @@ async fn get_active_sessions( /// `GET /api/channel/settings/:platform` — return backend-owned channel settings. async fn get_channel_settings( State(state): State, + Extension(user): Extension, Path(platform): Path, ) -> Result>, ApiError> { let platform = PluginType::from_str_opt(&platform) .ok_or_else(|| ApiError::BadRequest(format!("Invalid platform: {}", platform)))?; - let settings = state.settings_service.get_platform_settings(platform).await?; + let settings = state.settings_service.get_platform_settings(&user.id, platform).await?; Ok(Json(ApiResponse::ok(settings))) } /// `PUT /api/channel/settings/:platform/assistant` — persist assistant binding for a platform. async fn set_channel_assistant_setting( State(state): State, + Extension(user): Extension, Path(platform): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { @@ -587,8 +613,11 @@ async fn set_channel_assistant_setting( .ok_or_else(|| ApiError::BadRequest(format!("Invalid platform: {}", platform)))?; let Json(req) = body.map_err(ApiError::from)?; - state.settings_service.set_assistant_setting(platform, &req).await?; - state.session_manager.clear_all_sessions().await?; + state + .settings_service + .set_assistant_setting(&user.id, platform, &req) + .await?; + state.session_manager.clear_all_sessions(&user.id).await?; Ok(Json(ApiResponse::ok(BridgeResponse { success: true, @@ -600,6 +629,7 @@ async fn set_channel_assistant_setting( /// `PUT /api/channel/settings/:platform/default-model` — persist default model for a platform. async fn set_channel_default_model_setting( State(state): State, + Extension(user): Extension, Path(platform): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { @@ -607,8 +637,11 @@ async fn set_channel_default_model_setting( .ok_or_else(|| ApiError::BadRequest(format!("Invalid platform: {}", platform)))?; let Json(req) = body.map_err(ApiError::from)?; - state.settings_service.set_model_setting(platform, &req).await?; - state.session_manager.clear_all_sessions().await?; + state + .settings_service + .set_model_setting(&user.id, platform, &req) + .await?; + state.session_manager.clear_all_sessions(&user.id).await?; Ok(Json(ApiResponse::ok(BridgeResponse { success: true, @@ -628,6 +661,7 @@ async fn set_channel_default_model_setting( /// Agent/model config is persisted separately via `PUT /api/settings/client`. async fn sync_channel_settings( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; @@ -635,7 +669,7 @@ async fn sync_channel_settings( let _platform = PluginType::from_str_opt(&req.platform) .ok_or_else(|| ApiError::BadRequest(format!("Invalid platform: {}", req.platform)))?; - state.session_manager.clear_all_sessions().await?; + state.session_manager.clear_all_sessions(&user.id).await?; Ok(Json(ApiResponse::ok(BridgeResponse { success: true, @@ -731,11 +765,12 @@ fn build_test_config(req: &TestPluginRequest) -> PluginConfig { async fn resolve_extension_channel_plugin( state: &ChannelRouterState, + user_id: &str, plugin_id: &str, ) -> Option { state .extension_registry - .get_channel_plugins() + .get_channel_plugins_for_user(user_id) .await .into_iter() .find(|plugin| plugin.id == plugin_id) @@ -855,6 +890,16 @@ mod tests { assert!(matches!(err, ApiError::Conflict(_))); } + #[test] + fn cross_account_db_conflict_maps_to_stable_code() { + let err = db_error_to_api_error(DbError::Conflict( + "CROSS_ACCOUNT_REFERENCE: channel session conversation belongs to another user".into(), + )); + + assert_eq!(err.status_code(), StatusCode::CONFLICT); + assert_eq!(err.error_code(), "CROSS_ACCOUNT_REFERENCE"); + } + #[test] fn invalid_config_maps_to_bad_request() { let err = ApiError::from(ChannelError::InvalidConfig("missing token".into())); diff --git a/crates/aionui-channel/src/session.rs b/crates/aionui-channel/src/session.rs index 2184849bf..1d895b80a 100644 --- a/crates/aionui-channel/src/session.rs +++ b/crates/aionui-channel/src/session.rs @@ -31,6 +31,7 @@ impl SessionManager { /// the `ChannelManager` when it knows the active workspace path. pub async fn get_or_create_session( &self, + owner_user_id: &str, user_id: &str, chat_id: &str, agent_type: &str, @@ -48,7 +49,10 @@ impl SessionManager { last_activity: now, }; - let session = self.repo.get_or_create_session(user_id, chat_id, &new_row).await?; + let session = self + .repo + .get_or_create_session(owner_user_id, user_id, chat_id, &new_row) + .await?; debug!( session_id = %session.id, @@ -61,8 +65,8 @@ impl SessionManager { } /// Returns all active sessions. - pub async fn get_active_sessions(&self) -> Result, ChannelError> { - let sessions = self.repo.get_all_sessions().await?; + pub async fn get_active_sessions(&self, owner_user_id: &str) -> Result, ChannelError> { + let sessions = self.repo.get_all_sessions(owner_user_id).await?; Ok(sessions) } @@ -72,13 +76,16 @@ impl SessionManager { /// Used by `session.new` to give the user a clean slate in a chat. pub async fn reset_session( &self, + owner_user_id: &str, user_id: &str, chat_id: &str, agent_type: &str, workspace: Option<&str>, ) -> Result { // Delete old session if it exists - self.repo.delete_session_by_user_chat(user_id, chat_id).await?; + self.repo + .delete_session_by_user_chat(owner_user_id, user_id, chat_id) + .await?; // Create a fresh session let now = now_ms(); @@ -93,7 +100,10 @@ impl SessionManager { last_activity: now, }; - let session = self.repo.get_or_create_session(user_id, chat_id, &new_row).await?; + let session = self + .repo + .get_or_create_session(owner_user_id, user_id, chat_id, &new_row) + .await?; info!( session_id = %session.id, @@ -106,8 +116,15 @@ impl SessionManager { } /// Updates the agent_type for an existing session. - pub async fn update_agent_type(&self, session_id: &str, agent_type: &str) -> Result<(), ChannelError> { - self.repo.update_session_agent_type(session_id, agent_type).await?; + pub async fn update_agent_type( + &self, + owner_user_id: &str, + session_id: &str, + agent_type: &str, + ) -> Result<(), ChannelError> { + self.repo + .update_session_agent_type(owner_user_id, session_id, agent_type) + .await?; debug!( session_id = %session_id, @@ -120,8 +137,8 @@ impl SessionManager { /// Removes all sessions belonging to a user. /// /// Called when a user is revoked to clean up their session state. - pub async fn cleanup_user_sessions(&self, user_id: &str) -> Result<(), ChannelError> { - self.repo.delete_sessions_by_user(user_id).await?; + pub async fn cleanup_user_sessions(&self, owner_user_id: &str, user_id: &str) -> Result<(), ChannelError> { + self.repo.delete_sessions_by_user(owner_user_id, user_id).await?; info!(user_id = %user_id, "cleaned up user sessions"); Ok(()) } @@ -130,12 +147,14 @@ impl SessionManager { /// /// Called after settings sync to force sessions to be recreated /// with updated agent/model configuration. - pub async fn clear_all_sessions(&self) -> Result<(), ChannelError> { - let sessions = self.repo.get_all_sessions().await?; + pub async fn clear_all_sessions(&self, owner_user_id: &str) -> Result<(), ChannelError> { + let sessions = self.repo.get_all_sessions(owner_user_id).await?; let mut cleared_users = std::collections::HashSet::new(); for session in &sessions { if cleared_users.insert(session.user_id.clone()) { - self.repo.delete_sessions_by_user(&session.user_id).await?; + self.repo + .delete_sessions_by_user(owner_user_id, &session.user_id) + .await?; } } info!(count = sessions.len(), "cleared all channel sessions"); @@ -143,17 +162,26 @@ impl SessionManager { } /// Looks up a session by its unique ID. - pub async fn get_session_by_id(&self, session_id: &str) -> Result, ChannelError> { - Ok(self.repo.get_session(session_id).await?) + pub async fn get_session_by_id( + &self, + owner_user_id: &str, + session_id: &str, + ) -> Result, ChannelError> { + Ok(self.repo.get_session(owner_user_id, session_id).await?) } /// Persists the conversation binding for a session. /// /// Called after a new conversation is created for this session, /// linking the session to its backing conversation in the database. - pub async fn bind_conversation(&self, session_id: &str, conversation_id: &str) -> Result<(), ChannelError> { + pub async fn bind_conversation( + &self, + owner_user_id: &str, + session_id: &str, + conversation_id: &str, + ) -> Result<(), ChannelError> { self.repo - .update_session_conversation(session_id, conversation_id) + .update_session_conversation(owner_user_id, session_id, conversation_id) .await?; debug!( @@ -174,6 +202,7 @@ mod tests { use std::sync::Mutex; // ── Mock IChannelRepository ──────────────────────────────────────── + const OWNER_ID: &str = "owner-test"; struct MockRepo { sessions: Mutex>, @@ -194,55 +223,67 @@ mod tests { #[async_trait::async_trait] impl IChannelRepository for MockRepo { // -- Plugin CRUD (unused stubs) -- - async fn get_all_plugins(&self) -> Result, DbError> { + async fn get_all_plugins(&self, _owner_user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn get_plugin(&self, _id: &str) -> Result, DbError> { + async fn get_plugin(&self, _owner_user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } - async fn upsert_plugin(&self, _row: &ChannelPluginRow) -> Result<(), DbError> { + async fn upsert_plugin(&self, _owner_user_id: &str, _row: &ChannelPluginRow) -> Result<(), DbError> { Ok(()) } - async fn update_plugin_status(&self, _id: &str, _params: &UpdatePluginStatusParams) -> Result<(), DbError> { + async fn update_plugin_status( + &self, + _owner_user_id: &str, + _id: &str, + _params: &UpdatePluginStatusParams, + ) -> Result<(), DbError> { Ok(()) } - async fn delete_plugin(&self, _id: &str) -> Result<(), DbError> { + async fn delete_plugin(&self, _owner_user_id: &str, _id: &str) -> Result<(), DbError> { Ok(()) } // -- User CRUD (unused stubs) -- - async fn get_all_users(&self) -> Result, DbError> { + async fn get_all_users(&self, _owner_user_id: &str) -> Result, DbError> { Ok(vec![]) } async fn get_user_by_platform( &self, + _owner_user_id: &str, _platform_user_id: &str, _platform_type: &str, ) -> Result, DbError> { Ok(None) } - async fn create_user(&self, _row: &AssistantUserRow) -> Result<(), DbError> { + async fn create_user(&self, _owner_user_id: &str, _row: &AssistantUserRow) -> Result<(), DbError> { Ok(()) } - async fn update_user_last_active(&self, _id: &str, _last_active: TimestampMs) -> Result<(), DbError> { + async fn update_user_last_active( + &self, + _owner_user_id: &str, + _id: &str, + _last_active: TimestampMs, + ) -> Result<(), DbError> { Ok(()) } - async fn delete_user(&self, _id: &str) -> Result<(), DbError> { + async fn delete_user(&self, _owner_user_id: &str, _id: &str) -> Result<(), DbError> { Ok(()) } // -- Session CRUD -- - async fn get_all_sessions(&self) -> Result, DbError> { + async fn get_all_sessions(&self, _owner_user_id: &str) -> Result, DbError> { Ok(self.sessions.lock().unwrap().clone()) } - async fn get_session(&self, id: &str) -> Result, DbError> { + async fn get_session(&self, _owner_user_id: &str, id: &str) -> Result, DbError> { let sessions = self.sessions.lock().unwrap(); Ok(sessions.iter().find(|s| s.id == id).cloned()) } async fn get_or_create_session( &self, + _owner_user_id: &str, user_id: &str, chat_id: &str, new_row: &AssistantSessionRow, @@ -261,7 +302,12 @@ mod tests { Ok(new_row.clone()) } - async fn update_session_activity(&self, id: &str, last_activity: TimestampMs) -> Result<(), DbError> { + async fn update_session_activity( + &self, + _owner_user_id: &str, + id: &str, + last_activity: TimestampMs, + ) -> Result<(), DbError> { let mut sessions = self.sessions.lock().unwrap(); if let Some(s) = sessions.iter_mut().find(|s| s.id == id) { s.last_activity = last_activity; @@ -271,7 +317,12 @@ mod tests { } } - async fn update_session_conversation(&self, id: &str, conversation_id: &str) -> Result<(), DbError> { + async fn update_session_conversation( + &self, + _owner_user_id: &str, + id: &str, + conversation_id: &str, + ) -> Result<(), DbError> { let mut sessions = self.sessions.lock().unwrap(); if let Some(s) = sessions.iter_mut().find(|s| s.id == id) { s.conversation_id = Some(conversation_id.to_owned()); @@ -282,7 +333,12 @@ mod tests { } } - async fn update_session_agent_type(&self, id: &str, agent_type: &str) -> Result<(), DbError> { + async fn update_session_agent_type( + &self, + _owner_user_id: &str, + id: &str, + agent_type: &str, + ) -> Result<(), DbError> { let mut sessions = self.sessions.lock().unwrap(); if let Some(s) = sessions.iter_mut().find(|s| s.id == id) { s.agent_type = agent_type.to_owned(); @@ -293,32 +349,41 @@ mod tests { } } - async fn delete_sessions_by_user(&self, user_id: &str) -> Result<(), DbError> { + async fn delete_sessions_by_user(&self, _owner_user_id: &str, user_id: &str) -> Result<(), DbError> { let mut sessions = self.sessions.lock().unwrap(); sessions.retain(|s| s.user_id != user_id); Ok(()) } - async fn delete_session_by_user_chat(&self, user_id: &str, chat_id: &str) -> Result<(), DbError> { + async fn delete_session_by_user_chat( + &self, + _owner_user_id: &str, + user_id: &str, + chat_id: &str, + ) -> Result<(), DbError> { let mut sessions = self.sessions.lock().unwrap(); sessions.retain(|s| !(s.user_id == user_id && s.chat_id.as_deref() == Some(chat_id))); Ok(()) } // -- Pairing codes (unused stubs) -- - async fn create_pairing(&self, _row: &PairingCodeRow) -> Result<(), DbError> { + async fn create_pairing(&self, _owner_user_id: &str, _row: &PairingCodeRow) -> Result<(), DbError> { Ok(()) } - async fn get_pending_pairings(&self) -> Result, DbError> { + async fn get_pending_pairings(&self, _owner_user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn get_pairing_by_code(&self, _code: &str) -> Result, DbError> { + async fn get_pairing_by_code( + &self, + _owner_user_id: &str, + _code: &str, + ) -> Result, DbError> { Ok(None) } - async fn update_pairing_status(&self, _code: &str, _status: &str) -> Result<(), DbError> { + async fn update_pairing_status(&self, _owner_user_id: &str, _code: &str, _status: &str) -> Result<(), DbError> { Ok(()) } - async fn cleanup_expired_pairings(&self, _now: TimestampMs) -> Result { + async fn cleanup_expired_pairings(&self, _owner_user_id: &str, _now: TimestampMs) -> Result { Ok(0) } } @@ -335,7 +400,7 @@ mod tests { async fn creates_new_session() { let (mgr, repo) = make_manager(); let session = mgr - .get_or_create_session("user1", "chat1", "gemini", None) + .get_or_create_session(OWNER_ID, "user1", "chat1", "gemini", None) .await .unwrap(); @@ -353,11 +418,11 @@ mod tests { let (mgr, repo) = make_manager(); let s1 = mgr - .get_or_create_session("user1", "chat1", "gemini", None) + .get_or_create_session(OWNER_ID, "user1", "chat1", "gemini", None) .await .unwrap(); let s2 = mgr - .get_or_create_session("user1", "chat1", "gemini", None) + .get_or_create_session(OWNER_ID, "user1", "chat1", "gemini", None) .await .unwrap(); @@ -369,8 +434,14 @@ mod tests { async fn different_chats_get_different_sessions() { let (mgr, repo) = make_manager(); - let s1 = mgr.get_or_create_session("user1", "chatA", "acp", None).await.unwrap(); - let s2 = mgr.get_or_create_session("user1", "chatB", "acp", None).await.unwrap(); + let s1 = mgr + .get_or_create_session(OWNER_ID, "user1", "chatA", "acp", None) + .await + .unwrap(); + let s2 = mgr + .get_or_create_session(OWNER_ID, "user1", "chatB", "acp", None) + .await + .unwrap(); assert_ne!(s1.id, s2.id); assert_eq!(repo.get_sessions().len(), 2); @@ -381,11 +452,11 @@ mod tests { let (mgr, repo) = make_manager(); let s1 = mgr - .get_or_create_session("user1", "chat1", "gemini", None) + .get_or_create_session(OWNER_ID, "user1", "chat1", "gemini", None) .await .unwrap(); let s2 = mgr - .get_or_create_session("user2", "chat1", "gemini", None) + .get_or_create_session(OWNER_ID, "user2", "chat1", "gemini", None) .await .unwrap(); @@ -397,7 +468,7 @@ mod tests { async fn session_with_workspace() { let (mgr, _repo) = make_manager(); let session = mgr - .get_or_create_session("u1", "c1", "acp", Some("/workspace")) + .get_or_create_session(OWNER_ID, "u1", "c1", "acp", Some("/workspace")) .await .unwrap(); @@ -409,17 +480,21 @@ mod tests { #[tokio::test] async fn get_active_sessions_empty() { let (mgr, _repo) = make_manager(); - let sessions = mgr.get_active_sessions().await.unwrap(); + let sessions = mgr.get_active_sessions(OWNER_ID).await.unwrap(); assert!(sessions.is_empty()); } #[tokio::test] async fn get_active_sessions_returns_all() { let (mgr, _repo) = make_manager(); - mgr.get_or_create_session("u1", "c1", "gemini", None).await.unwrap(); - mgr.get_or_create_session("u2", "c2", "acp", None).await.unwrap(); + mgr.get_or_create_session(OWNER_ID, "u1", "c1", "gemini", None) + .await + .unwrap(); + mgr.get_or_create_session(OWNER_ID, "u2", "c2", "acp", None) + .await + .unwrap(); - let sessions = mgr.get_active_sessions().await.unwrap(); + let sessions = mgr.get_active_sessions(OWNER_ID).await.unwrap(); assert_eq!(sessions.len(), 2); } @@ -428,11 +503,17 @@ mod tests { #[tokio::test] async fn cleanup_removes_user_sessions() { let (mgr, repo) = make_manager(); - mgr.get_or_create_session("u1", "c1", "gemini", None).await.unwrap(); - mgr.get_or_create_session("u1", "c2", "gemini", None).await.unwrap(); - mgr.get_or_create_session("u2", "c1", "acp", None).await.unwrap(); + mgr.get_or_create_session(OWNER_ID, "u1", "c1", "gemini", None) + .await + .unwrap(); + mgr.get_or_create_session(OWNER_ID, "u1", "c2", "gemini", None) + .await + .unwrap(); + mgr.get_or_create_session(OWNER_ID, "u2", "c1", "acp", None) + .await + .unwrap(); - mgr.cleanup_user_sessions("u1").await.unwrap(); + mgr.cleanup_user_sessions(OWNER_ID, "u1").await.unwrap(); let sessions = repo.get_sessions(); assert_eq!(sessions.len(), 1); @@ -442,9 +523,11 @@ mod tests { #[tokio::test] async fn cleanup_noop_for_unknown_user() { let (mgr, repo) = make_manager(); - mgr.get_or_create_session("u1", "c1", "gemini", None).await.unwrap(); + mgr.get_or_create_session(OWNER_ID, "u1", "c1", "gemini", None) + .await + .unwrap(); - mgr.cleanup_user_sessions("u999").await.unwrap(); + mgr.cleanup_user_sessions(OWNER_ID, "u999").await.unwrap(); assert_eq!(repo.get_sessions().len(), 1); } @@ -454,10 +537,13 @@ mod tests { #[tokio::test] async fn bind_conversation_persists_conversation_id() { let (mgr, repo) = make_manager(); - let session = mgr.get_or_create_session("u1", "c1", "acp", None).await.unwrap(); + let session = mgr + .get_or_create_session(OWNER_ID, "u1", "c1", "acp", None) + .await + .unwrap(); assert!(session.conversation_id.is_none()); - mgr.bind_conversation(&session.id, "conv_123").await.unwrap(); + mgr.bind_conversation(OWNER_ID, &session.id, "conv_123").await.unwrap(); let updated = repo.get_sessions().into_iter().find(|s| s.id == session.id).unwrap(); assert_eq!(updated.conversation_id.as_deref(), Some("conv_123")); @@ -466,7 +552,7 @@ mod tests { #[tokio::test] async fn bind_conversation_not_found() { let (mgr, _repo) = make_manager(); - let err = mgr.bind_conversation("nonexistent", "conv_123").await; + let err = mgr.bind_conversation(OWNER_ID, "nonexistent", "conv_123").await; assert!(err.is_err()); } @@ -475,9 +561,12 @@ mod tests { #[tokio::test] async fn reset_session_creates_fresh_session() { let (mgr, repo) = make_manager(); - let s1 = mgr.get_or_create_session("u1", "c1", "gemini", None).await.unwrap(); + let s1 = mgr + .get_or_create_session(OWNER_ID, "u1", "c1", "gemini", None) + .await + .unwrap(); - let s2 = mgr.reset_session("u1", "c1", "gemini", None).await.unwrap(); + let s2 = mgr.reset_session(OWNER_ID, "u1", "c1", "gemini", None).await.unwrap(); // New session should have a different ID assert_ne!(s1.id, s2.id); @@ -492,7 +581,7 @@ mod tests { #[tokio::test] async fn reset_session_noop_when_no_existing() { let (mgr, repo) = make_manager(); - let session = mgr.reset_session("u1", "c1", "acp", None).await.unwrap(); + let session = mgr.reset_session(OWNER_ID, "u1", "c1", "acp", None).await.unwrap(); assert_eq!(session.user_id, "u1"); assert_eq!(repo.get_sessions().len(), 1); @@ -503,10 +592,13 @@ mod tests { #[tokio::test] async fn update_agent_type_persists() { let (mgr, repo) = make_manager(); - let session = mgr.get_or_create_session("u1", "c1", "gemini", None).await.unwrap(); + let session = mgr + .get_or_create_session(OWNER_ID, "u1", "c1", "gemini", None) + .await + .unwrap(); assert_eq!(session.agent_type, "gemini"); - mgr.update_agent_type(&session.id, "acp").await.unwrap(); + mgr.update_agent_type(OWNER_ID, &session.id, "acp").await.unwrap(); let updated = repo.get_sessions().into_iter().find(|s| s.id == session.id).unwrap(); assert_eq!(updated.agent_type, "acp"); @@ -515,7 +607,7 @@ mod tests { #[tokio::test] async fn update_agent_type_not_found() { let (mgr, _repo) = make_manager(); - let err = mgr.update_agent_type("nonexistent", "acp").await; + let err = mgr.update_agent_type(OWNER_ID, "nonexistent", "acp").await; assert!(err.is_err()); } } diff --git a/crates/aionui-channel/src/stream_relay.rs b/crates/aionui-channel/src/stream_relay.rs index 892be0f82..aefaa2666 100644 --- a/crates/aionui-channel/src/stream_relay.rs +++ b/crates/aionui-channel/src/stream_relay.rs @@ -14,6 +14,7 @@ use crate::types::{OutgoingMessageType, PluginType, UnifiedOutgoingMessage}; /// Configuration for a stream relay session. #[derive(Debug, Clone)] pub struct RelayConfig { + pub owner_user_id: String, pub platform: PluginType, pub plugin_id: String, pub chat_id: String, @@ -27,6 +28,7 @@ pub struct RelayConfig { pub trait ChannelSender: Send + Sync { async fn send_message( &self, + owner_user_id: &str, plugin_id: &str, chat_id: &str, message: UnifiedOutgoingMessage, @@ -34,6 +36,7 @@ pub trait ChannelSender: Send + Sync { async fn edit_message( &self, + owner_user_id: &str, plugin_id: &str, chat_id: &str, message_id: &str, @@ -85,7 +88,12 @@ impl ChannelStreamRelay { let flush_msg = ChannelMessageService::build_streaming_message(&formatted); let _ = self .sender - .send_message(&self.config.plugin_id, &self.config.chat_id, flush_msg) + .send_message( + &self.config.owner_user_id, + &self.config.plugin_id, + &self.config.chat_id, + flush_msg, + ) .await; text_buffer.clear(); has_content = false; @@ -97,7 +105,12 @@ impl ChannelStreamRelay { let final_msg = ChannelMessageService::build_final_message(&formatted); let _ = self .sender - .send_message(&self.config.plugin_id, &self.config.chat_id, final_msg) + .send_message( + &self.config.owner_user_id, + &self.config.plugin_id, + &self.config.chat_id, + final_msg, + ) .await; } info!( @@ -124,7 +137,12 @@ impl ChannelStreamRelay { }; let _ = self .sender - .send_message(&self.config.plugin_id, &self.config.chat_id, error_msg) + .send_message( + &self.config.owner_user_id, + &self.config.plugin_id, + &self.config.chat_id, + error_msg, + ) .await; break; } @@ -136,7 +154,12 @@ impl ChannelStreamRelay { let final_msg = ChannelMessageService::build_final_message(&formatted); let _ = self .sender - .send_message(&self.config.plugin_id, &self.config.chat_id, final_msg) + .send_message( + &self.config.owner_user_id, + &self.config.plugin_id, + &self.config.chat_id, + final_msg, + ) .await; } break; @@ -161,7 +184,12 @@ impl ChannelStreamRelay { let thinking_msg = ChannelMessageService::build_thinking_message(); let thinking_msg_id = match self .sender - .send_message(&self.config.plugin_id, &self.config.chat_id, thinking_msg) + .send_message( + &self.config.owner_user_id, + &self.config.plugin_id, + &self.config.chat_id, + thinking_msg, + ) .await { Ok(id) => id, @@ -186,7 +214,13 @@ impl ChannelStreamRelay { let msg = ChannelMessageService::build_streaming_message(&formatted); let _ = self .sender - .edit_message(&self.config.plugin_id, &self.config.chat_id, &thinking_msg_id, msg) + .edit_message( + &self.config.owner_user_id, + &self.config.plugin_id, + &self.config.chat_id, + &thinking_msg_id, + msg, + ) .await; last_edit = Instant::now(); } @@ -196,7 +230,13 @@ impl ChannelStreamRelay { let msg = ChannelMessageService::build_streaming_message(&format!("\u{23f3} {name}...")); let _ = self .sender - .edit_message(&self.config.plugin_id, &self.config.chat_id, &thinking_msg_id, msg) + .edit_message( + &self.config.owner_user_id, + &self.config.plugin_id, + &self.config.chat_id, + &thinking_msg_id, + msg, + ) .await; } Some(StreamAction::Finish) => { @@ -226,6 +266,7 @@ impl ChannelStreamRelay { let _ = self .sender .edit_message( + &self.config.owner_user_id, &self.config.plugin_id, &self.config.chat_id, &thinking_msg_id, @@ -260,7 +301,13 @@ impl ChannelStreamRelay { let final_msg = ChannelMessageService::build_final_message(&formatted); let _ = self .sender - .edit_message(&self.config.plugin_id, &self.config.chat_id, msg_id, final_msg) + .edit_message( + &self.config.owner_user_id, + &self.config.plugin_id, + &self.config.chat_id, + msg_id, + final_msg, + ) .await; } } @@ -309,6 +356,7 @@ impl Default for MessageRecorder { impl ChannelSender for MessageRecorder { async fn send_message( &self, + _owner_user_id: &str, _plugin_id: &str, _chat_id: &str, message: UnifiedOutgoingMessage, @@ -319,6 +367,7 @@ impl ChannelSender for MessageRecorder { async fn edit_message( &self, + _owner_user_id: &str, _plugin_id: &str, _chat_id: &str, _message_id: &str, diff --git a/crates/aionui-channel/src/types.rs b/crates/aionui-channel/src/types.rs index 8c3d3e22c..58fca617e 100644 --- a/crates/aionui-channel/src/types.rs +++ b/crates/aionui-channel/src/types.rs @@ -273,6 +273,8 @@ pub struct BotInfo { /// Message received from an IM platform, normalized to a common format. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct UnifiedIncomingMessage { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub owner_user_id: Option, pub id: String, pub platform: PluginType, pub chat_id: String, @@ -716,6 +718,7 @@ mod tests { #[test] fn unified_incoming_message_text() { let msg = UnifiedIncomingMessage { + owner_user_id: None, id: "msg_1".into(), platform: PluginType::Telegram, chat_id: "chat_42".into(), @@ -969,6 +972,7 @@ mod tests { #[test] fn incoming_message_roundtrip() { let msg = UnifiedIncomingMessage { + owner_user_id: None, id: "m1".into(), platform: PluginType::Lark, chat_id: "c1".into(), diff --git a/crates/aionui-channel/tests/dingtalk_integration.rs b/crates/aionui-channel/tests/dingtalk_integration.rs index 436a010b4..e57ef820a 100644 --- a/crates/aionui-channel/tests/dingtalk_integration.rs +++ b/crates/aionui-channel/tests/dingtalk_integration.rs @@ -25,6 +25,8 @@ mod dingtalk_tests { use std::sync::Arc; use tokio::sync::mpsc; + const OWNER_USER_ID: &str = "system_default_user"; + // -- Test infrastructure ------------------------------------------------ struct MockBroadcaster { @@ -223,7 +225,9 @@ mod dingtalk_tests { let factory = dingtalk_factory(); let config = make_dingtalk_config_value(Some("key_123"), Some("secret")); - let result = manager.enable_plugin("nonexistent", &config, &factory).await; + let result = manager + .enable_plugin(OWNER_USER_ID, "nonexistent", &config, &factory) + .await; assert!(result.is_err()); } @@ -235,7 +239,9 @@ mod dingtalk_tests { let factory = dingtalk_factory(); let config = make_dingtalk_config_value(Some("bad_id"), Some("bad_secret")); - let result = manager.enable_plugin("dingtalk", &config, &factory).await; + let result = manager + .enable_plugin(OWNER_USER_ID, "dingtalk", &config, &factory) + .await; assert!(result.is_err()); } @@ -244,7 +250,7 @@ mod dingtalk_tests { #[tokio::test] async fn disable_without_db_row_returns_error() { let (manager, _repo, _bc) = setup().await; - let result = manager.disable_plugin("dingtalk").await; + let result = manager.disable_plugin(OWNER_USER_ID, "dingtalk").await; assert!(result.is_err()); } @@ -253,7 +259,7 @@ mod dingtalk_tests { #[tokio::test] async fn get_plugin_status_empty() { let (manager, _repo, _bc) = setup().await; - let statuses = manager.get_plugin_status().await.unwrap(); + let statuses = manager.get_plugin_status(OWNER_USER_ID).await.unwrap(); assert!(statuses.is_empty()); } @@ -263,7 +269,7 @@ mod dingtalk_tests { async fn restore_plugins_none_stored() { let (manager, _repo, _bc) = setup().await; let factory = dingtalk_factory(); - let result = manager.restore_plugins(&factory).await; + let result = manager.restore_plugins(OWNER_USER_ID, &factory).await; assert!(result.is_ok()); assert_eq!(manager.active_plugin_count(), 0); } @@ -273,6 +279,6 @@ mod dingtalk_tests { #[tokio::test] async fn is_plugin_running_false_when_not_enabled() { let (manager, _repo, _bc) = setup().await; - assert!(!manager.is_plugin_running("dingtalk")); + assert!(!manager.is_plugin_running(OWNER_USER_ID, "dingtalk")); } } diff --git a/crates/aionui-channel/tests/lark_integration.rs b/crates/aionui-channel/tests/lark_integration.rs index 1656ce0c0..8e3869c71 100644 --- a/crates/aionui-channel/tests/lark_integration.rs +++ b/crates/aionui-channel/tests/lark_integration.rs @@ -24,6 +24,8 @@ mod lark_tests { use std::sync::Arc; use tokio::sync::mpsc; + const OWNER_USER_ID: &str = "system_default_user"; + // -- Test infrastructure ------------------------------------------------ struct MockBroadcaster { @@ -222,7 +224,9 @@ mod lark_tests { let factory = lark_factory(); let config = make_lark_config_value(Some("cli_123"), Some("secret")); - let result = manager.enable_plugin("nonexistent", &config, &factory).await; + let result = manager + .enable_plugin(OWNER_USER_ID, "nonexistent", &config, &factory) + .await; assert!(result.is_err()); } @@ -234,7 +238,7 @@ mod lark_tests { let factory = lark_factory(); let config = make_lark_config_value(Some("bad_id"), Some("bad_secret")); - let result = manager.enable_plugin("lark", &config, &factory).await; + let result = manager.enable_plugin(OWNER_USER_ID, "lark", &config, &factory).await; assert!(result.is_err()); } @@ -243,7 +247,7 @@ mod lark_tests { #[tokio::test] async fn disable_without_db_row_returns_error() { let (manager, _repo, _bc) = setup().await; - let result = manager.disable_plugin("lark").await; + let result = manager.disable_plugin(OWNER_USER_ID, "lark").await; assert!(result.is_err()); } @@ -252,7 +256,7 @@ mod lark_tests { #[tokio::test] async fn get_plugin_status_empty() { let (manager, _repo, _bc) = setup().await; - let statuses = manager.get_plugin_status().await.unwrap(); + let statuses = manager.get_plugin_status(OWNER_USER_ID).await.unwrap(); assert!(statuses.is_empty()); } @@ -262,7 +266,7 @@ mod lark_tests { async fn restore_plugins_none_stored() { let (manager, _repo, _bc) = setup().await; let factory = lark_factory(); - let result = manager.restore_plugins(&factory).await; + let result = manager.restore_plugins(OWNER_USER_ID, &factory).await; assert!(result.is_ok()); assert_eq!(manager.active_plugin_count(), 0); } @@ -272,6 +276,6 @@ mod lark_tests { #[tokio::test] async fn is_plugin_running_false_when_not_enabled() { let (manager, _repo, _bc) = setup().await; - assert!(!manager.is_plugin_running("lark")); + assert!(!manager.is_plugin_running(OWNER_USER_ID, "lark")); } } diff --git a/crates/aionui-channel/tests/manager_integration.rs b/crates/aionui-channel/tests/manager_integration.rs index 5504949ae..ab207c62e 100644 --- a/crates/aionui-channel/tests/manager_integration.rs +++ b/crates/aionui-channel/tests/manager_integration.rs @@ -16,11 +16,14 @@ use aionui_channel::types::{ BotInfo, OutgoingMessageType, PluginConfig, PluginCredentials, PluginStatus, PluginType, UnifiedOutgoingMessage, }; use aionui_common::decrypt_string; -use aionui_db::{IChannelRepository, SqliteChannelRepository, init_database_memory}; +use aionui_db::{ + IChannelRepository, IUserRepository, SqliteChannelRepository, SqliteUserRepository, init_database_memory, +}; use aionui_realtime::EventBroadcaster; use tokio::sync::mpsc; // ── Test infrastructure ───────────────────────────────────────────── +const OWNER_ID: &str = "system_default_user"; struct MockBroadcaster { events: Mutex>>, @@ -145,7 +148,35 @@ fn test_key() -> [u8; 32] { } async fn setup() -> (ChannelManager, Arc, Arc) { + let (mgr, repo, bc, _) = setup_with_optional_second_owner(false).await; + (mgr, repo, bc) +} + +async fn setup_with_second_owner() -> ( + ChannelManager, + Arc, + Arc, + String, +) { + let (mgr, repo, bc, owner_b_id) = setup_with_optional_second_owner(true).await; + (mgr, repo, bc, owner_b_id.expect("second owner should be created")) +} + +async fn setup_with_optional_second_owner( + create_second_owner: bool, +) -> ( + ChannelManager, + Arc, + Arc, + Option, +) { let db = init_database_memory().await.unwrap(); + let owner_b_id = if create_second_owner { + let user_repo = SqliteUserRepository::new(db.pool().clone()); + Some(user_repo.create_user("channel_owner_b", "hash").await.unwrap().id) + } else { + None + }; let repo: Arc = Arc::new(SqliteChannelRepository::new(db.pool().clone())); let bc = Arc::new(MockBroadcaster::new()); let (msg_tx, _msg_rx) = mpsc::channel(16); @@ -153,7 +184,7 @@ async fn setup() -> (ChannelManager, Arc, Arc PluginFactory { @@ -234,7 +265,7 @@ fn make_test_outgoing() -> UnifiedOutgoingMessage { #[tokio::test] async fn ps1_get_status_empty() { let (mgr, _repo, _bc) = setup().await; - let statuses = mgr.get_plugin_status().await.unwrap(); + let statuses = mgr.get_plugin_status(OWNER_ID).await.unwrap(); assert!(statuses.is_empty()); } @@ -245,11 +276,11 @@ async fn ps2_get_status_with_plugins() { let (mgr, _repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); - let statuses = mgr.get_plugin_status().await.unwrap(); + let statuses = mgr.get_plugin_status(OWNER_ID).await.unwrap(); assert_eq!(statuses.len(), 1); assert_eq!(statuses[0].plugin_id, "telegram"); assert_eq!(statuses[0].plugin_type, "telegram"); @@ -265,19 +296,19 @@ async fn ep1_enable_telegram_plugin() { let (mgr, repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); // Plugin persisted in DB - let row = repo.get_plugin("telegram").await.unwrap().unwrap(); + let row = repo.get_plugin(OWNER_ID, "telegram").await.unwrap().unwrap(); assert!(row.enabled); assert_eq!(row.r#type, "telegram"); assert_eq!(row.name, "Telegram Bot"); assert!(row.last_connected.is_some()); // Plugin is running in memory - assert!(mgr.is_plugin_running("telegram")); + assert!(mgr.is_plugin_running(OWNER_ID, "telegram")); assert_eq!(mgr.active_plugin_count(), 1); } @@ -288,7 +319,7 @@ async fn ep2_re_enable_updates_config() { let (mgr, repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); @@ -297,13 +328,15 @@ async fn ep2_re_enable_updates_config() { "credentials": { "token": "bot:new_token_456" }, "config": { "mode": "webhook", "webhook_url": "https://example.com" } }); - mgr.enable_plugin("telegram", &new_config, &factory).await.unwrap(); + mgr.enable_plugin(OWNER_ID, "telegram", &new_config, &factory) + .await + .unwrap(); // Still only one plugin assert_eq!(mgr.active_plugin_count(), 1); // Config should be updated - let row = repo.get_plugin("telegram").await.unwrap().unwrap(); + let row = repo.get_plugin(OWNER_ID, "telegram").await.unwrap().unwrap(); let decrypted = decrypt_string(&row.config, &test_key()).unwrap(); let config: PluginConfig = serde_json::from_str(&decrypted).unwrap(); assert_eq!(config.credentials.token.as_deref(), Some("bot:new_token_456")); @@ -317,20 +350,20 @@ async fn ep6_re_enable_empty_config_reuses_stored_credentials() { let factory = make_factory(); // First enable persists the token, then the user disables the channel. - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); - mgr.disable_plugin("telegram").await.unwrap(); + mgr.disable_plugin(OWNER_ID, "telegram").await.unwrap(); // The Settings re-enable toggle sends an empty config and relies on the // previously stored credentials being reused instead of erroring out. - mgr.enable_plugin("telegram", &serde_json::json!({}), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &serde_json::json!({}), &factory) .await .unwrap(); - assert!(mgr.is_plugin_running("telegram")); + assert!(mgr.is_plugin_running(OWNER_ID, "telegram")); - let row = repo.get_plugin("telegram").await.unwrap().unwrap(); + let row = repo.get_plugin(OWNER_ID, "telegram").await.unwrap().unwrap(); assert!(row.enabled); let decrypted = decrypt_string(&row.config, &test_key()).unwrap(); let config: PluginConfig = serde_json::from_str(&decrypted).unwrap(); @@ -345,7 +378,7 @@ async fn ep5_invalid_plugin_id() { let factory = make_factory(); let err = mgr - .enable_plugin("nonexistent", &make_telegram_config(), &factory) + .enable_plugin(OWNER_ID, "nonexistent", &make_telegram_config(), &factory) .await .unwrap_err(); assert!(matches!(err, ChannelError::InvalidPluginType(_))); @@ -360,7 +393,10 @@ async fn ep3_ep4_invalid_config_structure() { // Missing credentials entirely let bad = serde_json::json!({ "wrong_key": "value" }); - let err = mgr.enable_plugin("telegram", &bad, &factory).await.unwrap_err(); + let err = mgr + .enable_plugin(OWNER_ID, "telegram", &bad, &factory) + .await + .unwrap_err(); assert!(matches!(err, ChannelError::InvalidConfig(_))); } @@ -371,15 +407,15 @@ async fn dp1_disable_enabled_plugin() { let (mgr, repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); - mgr.disable_plugin("telegram").await.unwrap(); + mgr.disable_plugin(OWNER_ID, "telegram").await.unwrap(); assert_eq!(mgr.active_plugin_count(), 0); - assert!(!mgr.is_plugin_running("telegram")); + assert!(!mgr.is_plugin_running(OWNER_ID, "telegram")); - let row = repo.get_plugin("telegram").await.unwrap().unwrap(); + let row = repo.get_plugin(OWNER_ID, "telegram").await.unwrap().unwrap(); assert!(!row.enabled); assert_eq!(row.status.as_deref(), Some("stopped")); } @@ -391,13 +427,13 @@ async fn dp2_disable_already_disabled() { let (mgr, _repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); - mgr.disable_plugin("telegram").await.unwrap(); + mgr.disable_plugin(OWNER_ID, "telegram").await.unwrap(); // Second disable should not error - mgr.disable_plugin("telegram").await.unwrap(); + mgr.disable_plugin(OWNER_ID, "telegram").await.unwrap(); assert_eq!(mgr.active_plugin_count(), 0); } @@ -469,7 +505,7 @@ async fn tp_test_does_not_persist() { .await .unwrap(); - let plugins = repo.get_all_plugins().await.unwrap(); + let plugins = repo.get_all_plugins(OWNER_ID).await.unwrap(); assert!(plugins.is_empty()); assert_eq!(mgr.active_plugin_count(), 0); } @@ -481,11 +517,11 @@ async fn cs1_credentials_stored_encrypted() { let (mgr, repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); - let row = repo.get_plugin("telegram").await.unwrap().unwrap(); + let row = repo.get_plugin(OWNER_ID, "telegram").await.unwrap().unwrap(); // Config should not contain plaintext token assert!(!row.config.contains("bot:valid123")); @@ -507,11 +543,11 @@ async fn cs2_status_does_not_leak_credentials() { let (mgr, _repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); - let statuses = mgr.get_plugin_status().await.unwrap(); + let statuses = mgr.get_plugin_status(OWNER_ID).await.unwrap(); let json = serde_json::to_string(&statuses).unwrap(); // No sensitive fields should appear @@ -530,7 +566,7 @@ async fn ws2_enable_broadcasts_status_change() { let (mgr, _repo, bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); @@ -548,12 +584,12 @@ async fn ws2_disable_broadcasts_status_change() { let (mgr, _repo, bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); bc.take_events(); // clear enable events - mgr.disable_plugin("telegram").await.unwrap(); + mgr.disable_plugin(OWNER_ID, "telegram").await.unwrap(); let events = bc.take_events(); let status_events: Vec<_> = events @@ -571,7 +607,7 @@ async fn restore_starts_enabled_plugins() { let factory = make_factory(); // First enable and persist a plugin - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); @@ -580,9 +616,9 @@ async fn restore_starts_enabled_plugins() { assert_eq!(mgr.active_plugin_count(), 0); // Restore should bring it back - mgr.restore_plugins(&factory).await.unwrap(); + mgr.restore_plugins(OWNER_ID, &factory).await.unwrap(); assert_eq!(mgr.active_plugin_count(), 1); - assert!(mgr.is_plugin_running("telegram")); + assert!(mgr.is_plugin_running(OWNER_ID, "telegram")); } // ── Restore: disabled plugins are skipped ───────────────────────── @@ -592,12 +628,12 @@ async fn restore_skips_disabled_plugins() { let (mgr, _repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); - mgr.disable_plugin("telegram").await.unwrap(); + mgr.disable_plugin(OWNER_ID, "telegram").await.unwrap(); - mgr.restore_plugins(&factory).await.unwrap(); + mgr.restore_plugins(OWNER_ID, &factory).await.unwrap(); assert_eq!(mgr.active_plugin_count(), 0); } @@ -608,19 +644,51 @@ async fn enable_multiple_plugins() { let (mgr, _repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) + .await + .unwrap(); + mgr.enable_plugin(OWNER_ID, "lark", &make_lark_config(), &factory) .await .unwrap(); - mgr.enable_plugin("lark", &make_lark_config(), &factory).await.unwrap(); assert_eq!(mgr.active_plugin_count(), 2); - assert!(mgr.is_plugin_running("telegram")); - assert!(mgr.is_plugin_running("lark")); + assert!(mgr.is_plugin_running(OWNER_ID, "telegram")); + assert!(mgr.is_plugin_running(OWNER_ID, "lark")); - let statuses = mgr.get_plugin_status().await.unwrap(); + let statuses = mgr.get_plugin_status(OWNER_ID).await.unwrap(); assert_eq!(statuses.len(), 2); } +#[tokio::test] +async fn same_plugin_id_is_runtime_isolated_by_owner() { + let (mgr, repo, _bc, owner_b_id) = setup_with_second_owner().await; + let factory = make_factory(); + + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) + .await + .unwrap(); + mgr.enable_plugin(&owner_b_id, "telegram", &make_telegram_config(), &factory) + .await + .unwrap(); + + assert_eq!(mgr.active_plugin_count(), 2); + assert!(mgr.is_plugin_running(OWNER_ID, "telegram")); + assert!(mgr.is_plugin_running(&owner_b_id, "telegram")); + + let owner_a_plugins = repo.get_all_plugins(OWNER_ID).await.unwrap(); + let owner_b_plugins = repo.get_all_plugins(&owner_b_id).await.unwrap(); + assert_eq!(owner_a_plugins.len(), 1); + assert_eq!(owner_b_plugins.len(), 1); + assert_eq!(owner_a_plugins[0].id, "telegram"); + assert_eq!(owner_b_plugins[0].id, "telegram"); + + mgr.disable_plugin(OWNER_ID, "telegram").await.unwrap(); + + assert_eq!(mgr.active_plugin_count(), 1); + assert!(!mgr.is_plugin_running(OWNER_ID, "telegram")); + assert!(mgr.is_plugin_running(&owner_b_id, "telegram")); +} + // ── Shutdown stops all ──────────────────────────────────────────── #[tokio::test] @@ -628,10 +696,12 @@ async fn shutdown_stops_all() { let (mgr, _repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) + .await + .unwrap(); + mgr.enable_plugin(OWNER_ID, "lark", &make_lark_config(), &factory) .await .unwrap(); - mgr.enable_plugin("lark", &make_lark_config(), &factory).await.unwrap(); mgr.shutdown().await; assert_eq!(mgr.active_plugin_count(), 0); @@ -644,12 +714,12 @@ async fn send_message_routes_to_plugin() { let (mgr, _repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); let msg_id = mgr - .send_message("telegram", "chat_1", make_test_outgoing()) + .send_message(OWNER_ID, "telegram", "chat_1", make_test_outgoing()) .await .unwrap(); assert_eq!(msg_id, "mock_msg_id"); @@ -660,7 +730,7 @@ async fn send_message_not_running_fails() { let (mgr, _repo, _bc) = setup().await; let err = mgr - .send_message("telegram", "chat_1", make_test_outgoing()) + .send_message(OWNER_ID, "telegram", "chat_1", make_test_outgoing()) .await .unwrap_err(); assert!(matches!(err, ChannelError::PluginNotFound(_))); @@ -671,11 +741,11 @@ async fn edit_message_routes_to_plugin() { let (mgr, _repo, _bc) = setup().await; let factory = make_factory(); - mgr.enable_plugin("telegram", &make_telegram_config(), &factory) + mgr.enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap(); - mgr.edit_message("telegram", "chat_1", "msg_1", make_test_outgoing()) + mgr.edit_message(OWNER_ID, "telegram", "chat_1", "msg_1", make_test_outgoing()) .await .unwrap(); } @@ -687,11 +757,13 @@ async fn enable_failure_sets_error_in_db() { let (mgr, repo, _bc) = setup().await; let factory = make_failing_factory(); - let err = mgr.enable_plugin("telegram", &make_telegram_config(), &factory).await; + let err = mgr + .enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) + .await; assert!(err.is_err()); // Plugin should exist in DB with error status - let row = repo.get_plugin("telegram").await.unwrap().unwrap(); + let row = repo.get_plugin(OWNER_ID, "telegram").await.unwrap().unwrap(); assert_eq!(row.status.as_deref(), Some("error")); assert_eq!(mgr.active_plugin_count(), 0); } @@ -704,7 +776,7 @@ async fn enable_no_implementation_fails() { let factory = make_no_impl_factory(); let err = mgr - .enable_plugin("telegram", &make_telegram_config(), &factory) + .enable_plugin(OWNER_ID, "telegram", &make_telegram_config(), &factory) .await .unwrap_err(); assert!(matches!(err, ChannelError::InvalidPluginType(_))); diff --git a/crates/aionui-channel/tests/message_service_integration.rs b/crates/aionui-channel/tests/message_service_integration.rs index 085cedc2d..6920735ed 100644 --- a/crates/aionui-channel/tests/message_service_integration.rs +++ b/crates/aionui-channel/tests/message_service_integration.rs @@ -24,6 +24,8 @@ use aionui_realtime::EventBroadcaster; use async_trait::async_trait; use tokio::sync::broadcast; +const TEST_OWNER_USER_ID: &str = "system_default_user"; + struct TestBroadcaster { events: Mutex>>, } @@ -237,12 +239,7 @@ async fn send_to_agent_warms_cold_task_before_returning_stream_subscription() { let settings = Arc::new(ChannelSettingsService::new(Arc::new( SqliteClientPreferenceRepository::new(pool), ))); - let message_svc = ChannelMessageService::new( - conversation_svc, - Arc::clone(&task_manager), - settings, - "system_default_user".to_owned(), - ); + let message_svc = ChannelMessageService::new(conversation_svc, Arc::clone(&task_manager), settings); let session = AssistantSessionRow { id: "session-1".to_owned(), @@ -261,7 +258,10 @@ async fn send_to_agent_warms_cold_task_before_returning_stream_subscription() { PluginType::Dingtalk, PluginType::Weixin, ] { - let result = message_svc.send_to_agent(&session, "hello", platform).await.unwrap(); + let result = message_svc + .send_to_agent(TEST_OWNER_USER_ID, &session, "hello", platform) + .await + .unwrap(); assert!( result.stream_rx.is_some(), @@ -306,20 +306,18 @@ async fn send_to_agent_persists_assistant_snapshot_for_channel_bound_assistant() .await .unwrap(); pref_repo - .upsert_batch(&[( - "assistant.telegram.agent", - r#"{"assistant_id":"bare-claude","name":"Claude"}"#, - )]) + .upsert_batch( + TEST_OWNER_USER_ID, + &[( + "assistant.telegram.agent", + r#"{"assistant_id":"bare-claude","name":"Claude"}"#, + )], + ) .await .unwrap(); let settings = Arc::new(ChannelSettingsService::new(pref_repo).with_assistant_repos(definition_repo, overlay_repo)); - let message_svc = ChannelMessageService::new( - conversation_svc, - Arc::clone(&task_manager), - settings, - "system_default_user".to_owned(), - ); + let message_svc = ChannelMessageService::new(conversation_svc, Arc::clone(&task_manager), settings); let session = AssistantSessionRow { id: "session-assisted".to_owned(), @@ -333,12 +331,12 @@ async fn send_to_agent_persists_assistant_snapshot_for_channel_bound_assistant() }; let result = message_svc - .send_to_agent(&session, "hello", PluginType::Telegram) + .send_to_agent(TEST_OWNER_USER_ID, &session, "hello", PluginType::Telegram) .await .unwrap(); let snapshot = conversation_repo - .get_assistant_snapshot(&result.conversation_id) + .get_assistant_snapshot(TEST_OWNER_USER_ID, &result.conversation_id) .await .unwrap(); assert!( @@ -346,10 +344,14 @@ async fn send_to_agent_persists_assistant_snapshot_for_channel_bound_assistant() "channel-created conversation should persist an assistant snapshot when the platform is bound to an assistant" ); let snapshot = snapshot.unwrap(); - let conversation = conversation_repo.get(&result.conversation_id).await.unwrap().unwrap(); + let conversation = conversation_repo + .get(TEST_OWNER_USER_ID, &result.conversation_id) + .await + .unwrap() + .unwrap(); assert_eq!(conversation.r#type, AgentType::Acp.serde_name()); let session_row = acp_session_repo - .get(&result.conversation_id) + .get_for_user(TEST_OWNER_USER_ID, &result.conversation_id) .await .unwrap() .expect("acp_session row should exist for ACP assistant conversations"); @@ -380,21 +382,19 @@ async fn send_to_agent_rejects_unresolvable_channel_assistant_binding() { let pref_repo = Arc::new(SqliteClientPreferenceRepository::new(pool.clone())); pref_repo - .upsert_batch(&[( - "assistant.telegram.agent", - r#"{"assistant_id":"missing-assistant","name":"Missing"}"#, - )]) + .upsert_batch( + TEST_OWNER_USER_ID, + &[( + "assistant.telegram.agent", + r#"{"assistant_id":"missing-assistant","name":"Missing"}"#, + )], + ) .await .unwrap(); let definition_repo = Arc::new(SqliteAssistantDefinitionRepository::new(pool.clone())); let overlay_repo = Arc::new(SqliteAssistantOverlayRepository::new(pool.clone())); let settings = Arc::new(ChannelSettingsService::new(pref_repo).with_assistant_repos(definition_repo, overlay_repo)); - let message_svc = ChannelMessageService::new( - conversation_svc, - Arc::clone(&task_manager), - settings, - "system_default_user".to_owned(), - ); + let message_svc = ChannelMessageService::new(conversation_svc, Arc::clone(&task_manager), settings); let session = AssistantSessionRow { id: "session-assisted-missing".to_owned(), @@ -408,7 +408,7 @@ async fn send_to_agent_rejects_unresolvable_channel_assistant_binding() { }; let err = message_svc - .send_to_agent(&session, "hello", PluginType::Telegram) + .send_to_agent(TEST_OWNER_USER_ID, &session, "hello", PluginType::Telegram) .await .unwrap_err(); assert!(matches!(err, ChannelError::MessageSendFailed(_))); @@ -454,12 +454,7 @@ async fn send_to_agent_without_saved_binding_defaults_to_bare_aionrs_assistant() .unwrap(); let settings = Arc::new(ChannelSettingsService::new(pref_repo).with_assistant_repos(definition_repo, overlay_repo)); - let message_svc = ChannelMessageService::new( - conversation_svc, - Arc::clone(&task_manager), - settings, - "system_default_user".to_owned(), - ); + let message_svc = ChannelMessageService::new(conversation_svc, Arc::clone(&task_manager), settings); let session = AssistantSessionRow { id: "session-assisted-default-aionrs".to_owned(), @@ -473,16 +468,20 @@ async fn send_to_agent_without_saved_binding_defaults_to_bare_aionrs_assistant() }; let result = message_svc - .send_to_agent(&session, "hello", PluginType::Telegram) + .send_to_agent(TEST_OWNER_USER_ID, &session, "hello", PluginType::Telegram) .await .unwrap(); let snapshot = conversation_repo - .get_assistant_snapshot(&result.conversation_id) + .get_assistant_snapshot(TEST_OWNER_USER_ID, &result.conversation_id) .await .unwrap() .expect("channel-created conversation should default to a bare assistant snapshot"); - let conversation = conversation_repo.get(&result.conversation_id).await.unwrap().unwrap(); + let conversation = conversation_repo + .get(TEST_OWNER_USER_ID, &result.conversation_id) + .await + .unwrap() + .unwrap(); assert_eq!(snapshot.assistant_id, "bare-aionrs"); assert_eq!(snapshot.agent_id, "632f31d2"); @@ -525,17 +524,15 @@ async fn send_to_agent_without_assistant_name_falls_back_to_legacy_channel_name( .await .unwrap(); pref_repo - .upsert_batch(&[("assistant.telegram.agent", r#"{"assistant_id":"bare-codex"}"#)]) + .upsert_batch( + TEST_OWNER_USER_ID, + &[("assistant.telegram.agent", r#"{"assistant_id":"bare-codex"}"#)], + ) .await .unwrap(); let settings = Arc::new(ChannelSettingsService::new(pref_repo).with_assistant_repos(definition_repo, overlay_repo)); - let message_svc = ChannelMessageService::new( - conversation_svc, - Arc::clone(&task_manager), - settings, - "system_default_user".to_owned(), - ); + let message_svc = ChannelMessageService::new(conversation_svc, Arc::clone(&task_manager), settings); let session = AssistantSessionRow { id: "session-assisted-fallback-name".to_owned(), @@ -549,10 +546,14 @@ async fn send_to_agent_without_assistant_name_falls_back_to_legacy_channel_name( }; let result = message_svc - .send_to_agent(&session, "hello", PluginType::Telegram) + .send_to_agent(TEST_OWNER_USER_ID, &session, "hello", PluginType::Telegram) .await .unwrap(); - let conversation = conversation_repo.get(&result.conversation_id).await.unwrap().unwrap(); + let conversation = conversation_repo + .get(TEST_OWNER_USER_ID, &result.conversation_id) + .await + .unwrap() + .unwrap(); assert_eq!(conversation.name, "tg-acp-codex-70880480"); } diff --git a/crates/aionui-channel/tests/orchestrator_test.rs b/crates/aionui-channel/tests/orchestrator_test.rs index 3b7c710aa..a1527ab34 100644 --- a/crates/aionui-channel/tests/orchestrator_test.rs +++ b/crates/aionui-channel/tests/orchestrator_test.rs @@ -8,8 +8,11 @@ use aionui_channel::types::{ MessageContentType, PluginType, UnifiedIncomingMessage, UnifiedMessageContent, UnifiedUser, }; +const OWNER_ID: &str = "system_default_user"; + fn make_text_message(user_id: &str, chat_id: &str, text: &str) -> UnifiedIncomingMessage { UnifiedIncomingMessage { + owner_user_id: None, id: "msg-1".into(), platform: PluginType::Telegram, chat_id: chat_id.into(), @@ -45,7 +48,12 @@ async fn unauthorized_user_gets_pairing_response() { let pairing = Arc::new(PairingService::new(repo.clone(), bus)); let session_mgr = Arc::new(SessionManager::new(repo)); - let executor = Arc::new(ActionExecutor::new(pairing, Arc::clone(&session_mgr), settings)); + let executor = Arc::new(ActionExecutor::new( + pairing, + Arc::clone(&session_mgr), + settings, + Some(OWNER_ID.to_owned()), + )); let msg = make_text_message("unknown_user", "chat_1", "hello"); let result = executor.handle_incoming_message(&msg).await.unwrap(); diff --git a/crates/aionui-channel/tests/pairing_integration.rs b/crates/aionui-channel/tests/pairing_integration.rs index 470239a32..e1193a3fe 100644 --- a/crates/aionui-channel/tests/pairing_integration.rs +++ b/crates/aionui-channel/tests/pairing_integration.rs @@ -17,6 +17,7 @@ use aionui_channel::error::ChannelError; use aionui_channel::pairing::PairingService; // ── Test infrastructure ───────────────────────────────────────────── +const OWNER_ID: &str = "system_default_user"; struct MockBroadcaster { events: Mutex>>, @@ -56,7 +57,10 @@ async fn setup() -> (PairingService, Arc, Arc= before + ttl); assert!(row.expires_at <= after + ttl); @@ -81,13 +85,19 @@ async fn pg2_code_expires_after_ten_minutes() { #[tokio::test] async fn pg3_same_user_re_request_expires_old_code() { let (svc, repo, _bc) = setup().await; - let code1 = svc.request_pairing("u1", "telegram", Some("Alice")).await.unwrap(); - let code2 = svc.request_pairing("u1", "telegram", Some("Alice")).await.unwrap(); + let code1 = svc + .request_pairing(OWNER_ID, "u1", "telegram", Some("Alice")) + .await + .unwrap(); + let code2 = svc + .request_pairing(OWNER_ID, "u1", "telegram", Some("Alice")) + .await + .unwrap(); assert_ne!(code1, code2); - let old = repo.get_pairing_by_code(&code1).await.unwrap().unwrap(); - let new = repo.get_pairing_by_code(&code2).await.unwrap().unwrap(); + let old = repo.get_pairing_by_code(OWNER_ID, &code1).await.unwrap().unwrap(); + let new = repo.get_pairing_by_code(OWNER_ID, &code2).await.unwrap().unwrap(); assert_eq!(old.status, "expired"); assert_eq!(new.status, "pending"); } @@ -97,7 +107,7 @@ async fn pg3_same_user_re_request_expires_old_code() { #[tokio::test] async fn pp1_no_pending_returns_empty() { let (svc, _repo, _bc) = setup().await; - let pending = svc.get_pending_pairings().await.unwrap(); + let pending = svc.get_pending_pairings(OWNER_ID).await.unwrap(); assert!(pending.is_empty()); } @@ -106,10 +116,12 @@ async fn pp1_no_pending_returns_empty() { #[tokio::test] async fn pp2_multiple_pending_returned() { let (svc, _repo, _bc) = setup().await; - svc.request_pairing("u1", "telegram", Some("Alice")).await.unwrap(); - svc.request_pairing("u2", "lark", Some("Bob")).await.unwrap(); + svc.request_pairing(OWNER_ID, "u1", "telegram", Some("Alice")) + .await + .unwrap(); + svc.request_pairing(OWNER_ID, "u2", "lark", Some("Bob")).await.unwrap(); - let pending = svc.get_pending_pairings().await.unwrap(); + let pending = svc.get_pending_pairings(OWNER_ID).await.unwrap(); assert_eq!(pending.len(), 2); } @@ -118,11 +130,12 @@ async fn pp2_multiple_pending_returned() { #[tokio::test] async fn pp3_expired_not_in_pending() { let (svc, repo, _bc) = setup().await; - svc.request_pairing("u1", "telegram", None).await.unwrap(); + svc.request_pairing(OWNER_ID, "u1", "telegram", None).await.unwrap(); // Insert already-expired code directly let expired_row = PairingCodeRow { code: "000001".into(), + owner_user_id: OWNER_ID.into(), platform_user_id: "u2".into(), platform_type: "lark".into(), display_name: None, @@ -130,9 +143,9 @@ async fn pp3_expired_not_in_pending() { expires_at: 1001, status: "pending".into(), }; - repo.create_pairing(&expired_row).await.unwrap(); + repo.create_pairing(OWNER_ID, &expired_row).await.unwrap(); - let pending = svc.get_pending_pairings().await.unwrap(); + let pending = svc.get_pending_pairings(OWNER_ID).await.unwrap(); assert_eq!(pending.len(), 1); assert_eq!(pending[0].platform_user_id, "u1"); } @@ -142,12 +155,15 @@ async fn pp3_expired_not_in_pending() { #[tokio::test] async fn ap1_approve_valid_pairing() { let (svc, repo, _bc) = setup().await; - let code = svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); + let code = svc + .request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); - svc.approve_pairing(&code).await.unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); // Status updated - let row = repo.get_pairing_by_code(&code).await.unwrap().unwrap(); + let row = repo.get_pairing_by_code(OWNER_ID, &code).await.unwrap().unwrap(); assert_eq!(row.status, "approved"); } @@ -156,10 +172,13 @@ async fn ap1_approve_valid_pairing() { #[tokio::test] async fn ap2_dc2_approved_user_in_authorized_list() { let (svc, repo, _bc) = setup().await; - let code = svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); - svc.approve_pairing(&code).await.unwrap(); + let code = svc + .request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); - let users = repo.get_all_users().await.unwrap(); + let users = repo.get_all_users(OWNER_ID).await.unwrap(); assert_eq!(users.len(), 1); assert_eq!(users[0].platform_user_id, "tg_42"); assert_eq!(users[0].platform_type, "telegram"); @@ -171,7 +190,7 @@ async fn ap2_dc2_approved_user_in_authorized_list() { #[tokio::test] async fn ap3_approve_nonexistent_code() { let (svc, _repo, _bc) = setup().await; - let err = svc.approve_pairing("000000").await.unwrap_err(); + let err = svc.approve_pairing(OWNER_ID, "000000").await.unwrap_err(); assert!(matches!(err, ChannelError::PairingNotFound(_))); } @@ -184,6 +203,7 @@ async fn ap4_approve_expired_code() { let expired_row = PairingCodeRow { code: "999999".into(), + owner_user_id: OWNER_ID.into(), platform_user_id: "u1".into(), platform_type: "telegram".into(), display_name: None, @@ -191,9 +211,9 @@ async fn ap4_approve_expired_code() { expires_at: 1001, status: "pending".into(), }; - repo.create_pairing(&expired_row).await.unwrap(); + repo.create_pairing(OWNER_ID, &expired_row).await.unwrap(); - let err = svc.approve_pairing("999999").await.unwrap_err(); + let err = svc.approve_pairing(OWNER_ID, "999999").await.unwrap_err(); assert!(matches!(err, ChannelError::PairingExpired(_))); } @@ -202,10 +222,10 @@ async fn ap4_approve_expired_code() { #[tokio::test] async fn ap5_double_approve_returns_already_processed() { let (svc, _repo, _bc) = setup().await; - let code = svc.request_pairing("u1", "telegram", None).await.unwrap(); - svc.approve_pairing(&code).await.unwrap(); + let code = svc.request_pairing(OWNER_ID, "u1", "telegram", None).await.unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); - let err = svc.approve_pairing(&code).await.unwrap_err(); + let err = svc.approve_pairing(OWNER_ID, &code).await.unwrap_err(); assert!(matches!(err, ChannelError::PairingAlreadyProcessed(_))); } @@ -214,7 +234,7 @@ async fn ap5_double_approve_returns_already_processed() { #[tokio::test] async fn ap6_empty_code_returns_not_found() { let (svc, _repo, _bc) = setup().await; - let err = svc.approve_pairing("").await.unwrap_err(); + let err = svc.approve_pairing(OWNER_ID, "").await.unwrap_err(); assert!(matches!(err, ChannelError::PairingNotFound(_))); } @@ -223,11 +243,11 @@ async fn ap6_empty_code_returns_not_found() { #[tokio::test] async fn rp1_reject_valid_pairing() { let (svc, repo, _bc) = setup().await; - let code = svc.request_pairing("u1", "telegram", None).await.unwrap(); + let code = svc.request_pairing(OWNER_ID, "u1", "telegram", None).await.unwrap(); - svc.reject_pairing(&code).await.unwrap(); + svc.reject_pairing(OWNER_ID, &code).await.unwrap(); - let row = repo.get_pairing_by_code(&code).await.unwrap().unwrap(); + let row = repo.get_pairing_by_code(OWNER_ID, &code).await.unwrap().unwrap(); assert_eq!(row.status, "rejected"); } @@ -236,10 +256,10 @@ async fn rp1_reject_valid_pairing() { #[tokio::test] async fn rp2_rejected_not_in_pending() { let (svc, _repo, _bc) = setup().await; - let code = svc.request_pairing("u1", "telegram", None).await.unwrap(); - svc.reject_pairing(&code).await.unwrap(); + let code = svc.request_pairing(OWNER_ID, "u1", "telegram", None).await.unwrap(); + svc.reject_pairing(OWNER_ID, &code).await.unwrap(); - let pending = svc.get_pending_pairings().await.unwrap(); + let pending = svc.get_pending_pairings(OWNER_ID).await.unwrap(); assert!(pending.is_empty()); } @@ -248,7 +268,7 @@ async fn rp2_rejected_not_in_pending() { #[tokio::test] async fn rp3_reject_nonexistent_code() { let (svc, _repo, _bc) = setup().await; - let err = svc.reject_pairing("000000").await.unwrap_err(); + let err = svc.reject_pairing(OWNER_ID, "000000").await.unwrap_err(); assert!(matches!(err, ChannelError::PairingNotFound(_))); } @@ -257,10 +277,10 @@ async fn rp3_reject_nonexistent_code() { #[tokio::test] async fn rp4_reject_already_approved() { let (svc, _repo, _bc) = setup().await; - let code = svc.request_pairing("u1", "telegram", None).await.unwrap(); - svc.approve_pairing(&code).await.unwrap(); + let code = svc.request_pairing(OWNER_ID, "u1", "telegram", None).await.unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); - let err = svc.reject_pairing(&code).await.unwrap_err(); + let err = svc.reject_pairing(OWNER_ID, &code).await.unwrap_err(); assert!(matches!(err, ChannelError::PairingAlreadyProcessed(_))); } @@ -273,6 +293,7 @@ async fn ec1_expired_codes_cleaned_up() { let expired_row = PairingCodeRow { code: "111111".into(), + owner_user_id: OWNER_ID.into(), platform_user_id: "u1".into(), platform_type: "telegram".into(), display_name: None, @@ -280,12 +301,12 @@ async fn ec1_expired_codes_cleaned_up() { expires_at: 2000, status: "pending".into(), }; - repo.create_pairing(&expired_row).await.unwrap(); + repo.create_pairing(OWNER_ID, &expired_row).await.unwrap(); - let count = repo.cleanup_expired_pairings(now_ms()).await.unwrap(); + let count = repo.cleanup_expired_pairings(OWNER_ID, now_ms()).await.unwrap(); assert_eq!(count, 1); - let row = repo.get_pairing_by_code("111111").await.unwrap().unwrap(); + let row = repo.get_pairing_by_code(OWNER_ID, "111111").await.unwrap().unwrap(); assert_eq!(row.status, "expired"); } @@ -294,12 +315,12 @@ async fn ec1_expired_codes_cleaned_up() { #[tokio::test] async fn ec2_non_expired_unaffected() { let (svc, repo, _bc) = setup().await; - let code = svc.request_pairing("u1", "telegram", None).await.unwrap(); + let code = svc.request_pairing(OWNER_ID, "u1", "telegram", None).await.unwrap(); - let count = repo.cleanup_expired_pairings(now_ms()).await.unwrap(); + let count = repo.cleanup_expired_pairings(OWNER_ID, now_ms()).await.unwrap(); assert_eq!(count, 0); - let row = repo.get_pairing_by_code(&code).await.unwrap().unwrap(); + let row = repo.get_pairing_by_code(OWNER_ID, &code).await.unwrap().unwrap(); assert_eq!(row.status, "pending"); } @@ -310,12 +331,18 @@ async fn dc3_same_platform_user_unique() { let (svc, _repo, _bc) = setup().await; // Approve first pairing - let code1 = svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); - svc.approve_pairing(&code1).await.unwrap(); + let code1 = svc + .request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); + svc.approve_pairing(OWNER_ID, &code1).await.unwrap(); // Second pairing for same user should fail on user creation (unique constraint) - let code2 = svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); - let result = svc.approve_pairing(&code2).await; + let code2 = svc + .request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); + let result = svc.approve_pairing(OWNER_ID, &code2).await; // DB should reject duplicate (platform_user_id, platform_type) assert!(result.is_err()); } @@ -325,7 +352,9 @@ async fn dc3_same_platform_user_unique() { #[tokio::test] async fn ws1_pairing_request_broadcasts_event() { let (svc, _repo, bc) = setup().await; - svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); + svc.request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); let events = bc.take_events(); assert_eq!(events.len(), 1); @@ -342,10 +371,13 @@ async fn ws1_pairing_request_broadcasts_event() { #[tokio::test] async fn ws3_approve_broadcasts_user_authorized() { let (svc, _repo, bc) = setup().await; - let code = svc.request_pairing("tg_42", "telegram", Some("Alice")).await.unwrap(); + let code = svc + .request_pairing(OWNER_ID, "tg_42", "telegram", Some("Alice")) + .await + .unwrap(); bc.take_events(); // clear request event - svc.approve_pairing(&code).await.unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); let events = bc.take_events(); assert_eq!(events.len(), 1); @@ -361,24 +393,24 @@ async fn ws3_approve_broadcasts_user_authorized() { #[tokio::test] async fn is_user_authorized_false_before_approval() { let (svc, _repo, _bc) = setup().await; - assert!(!svc.is_user_authorized("tg_42", "telegram").await.unwrap()); + assert!(!svc.is_user_authorized(OWNER_ID, "tg_42", "telegram").await.unwrap()); } #[tokio::test] async fn is_user_authorized_true_after_approval() { let (svc, _repo, _bc) = setup().await; - let code = svc.request_pairing("tg_42", "telegram", None).await.unwrap(); - svc.approve_pairing(&code).await.unwrap(); + let code = svc.request_pairing(OWNER_ID, "tg_42", "telegram", None).await.unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); - assert!(svc.is_user_authorized("tg_42", "telegram").await.unwrap()); + assert!(svc.is_user_authorized(OWNER_ID, "tg_42", "telegram").await.unwrap()); } #[tokio::test] async fn is_user_authorized_different_platform_false() { let (svc, _repo, _bc) = setup().await; - let code = svc.request_pairing("tg_42", "telegram", None).await.unwrap(); - svc.approve_pairing(&code).await.unwrap(); + let code = svc.request_pairing(OWNER_ID, "tg_42", "telegram", None).await.unwrap(); + svc.approve_pairing(OWNER_ID, &code).await.unwrap(); // Same user ID but different platform - assert!(!svc.is_user_authorized("tg_42", "lark").await.unwrap()); + assert!(!svc.is_user_authorized(OWNER_ID, "tg_42", "lark").await.unwrap()); } diff --git a/crates/aionui-channel/tests/session_action_integration.rs b/crates/aionui-channel/tests/session_action_integration.rs index 5700a4965..b84292cba 100644 --- a/crates/aionui-channel/tests/session_action_integration.rs +++ b/crates/aionui-channel/tests/session_action_integration.rs @@ -21,6 +21,7 @@ use aionui_channel::types::{ }; // ── Test infrastructure ───────────────────────────────────────────── +const OWNER_ID: &str = "system_default_user"; struct MockBroadcaster { events: Mutex>>, @@ -57,7 +58,7 @@ async fn setup() -> ( let pref_repo: Arc = Arc::new(aionui_db::SqliteClientPreferenceRepository::new(db.pool().clone())); let settings = Arc::new(ChannelSettingsService::new(pref_repo)); - let executor = ActionExecutor::new(pairing_arc, session_mgr_arc, settings); + let executor = ActionExecutor::new(pairing_arc, session_mgr_arc, settings, Some(OWNER_ID.to_owned())); // Keep db alive std::mem::forget(db); @@ -69,6 +70,7 @@ async fn create_user(repo: &Arc, platform_user_id: &str, let user_id = generate_id(); let row = AssistantUserRow { id: user_id.clone(), + owner_user_id: OWNER_ID.to_owned(), platform_user_id: platform_user_id.to_owned(), platform_type: platform_type.to_owned(), display_name: Some("Test User".into()), @@ -76,12 +78,13 @@ async fn create_user(repo: &Arc, platform_user_id: &str, last_active: None, session_id: None, }; - repo.create_user(&row).await.unwrap(); + repo.create_user(OWNER_ID, &row).await.unwrap(); user_id } fn make_text_message(user_id: &str, chat_id: &str, text: &str) -> UnifiedIncomingMessage { UnifiedIncomingMessage { + owner_user_id: None, id: format!("msg_{}", now_ms()), platform: PluginType::Telegram, chat_id: chat_id.into(), @@ -110,6 +113,7 @@ fn make_action_message( category: ActionCategory, ) -> UnifiedIncomingMessage { UnifiedIncomingMessage { + owner_user_id: None, id: format!("msg_{}", now_ms()), platform: PluginType::Telegram, chat_id: chat_id.into(), @@ -145,10 +149,10 @@ fn make_action_message( /// Helper: authorize a user via the pairing flow. async fn authorize_user(pairing: &PairingService, platform_user_id: &str, platform_type: &str) { let code = pairing - .request_pairing(platform_user_id, platform_type, Some("Test")) + .request_pairing(OWNER_ID, platform_user_id, platform_type, Some("Test")) .await .unwrap(); - pairing.approve_pairing(&code).await.unwrap(); + pairing.approve_pairing(OWNER_ID, &code).await.unwrap(); } // ── GS-1: No active sessions returns empty ───────────────────────── @@ -156,7 +160,7 @@ async fn authorize_user(pairing: &PairingService, platform_user_id: &str, platfo #[tokio::test] async fn gs1_no_sessions_returns_empty() { let (session_mgr, _, _, _) = setup().await; - let sessions = session_mgr.get_active_sessions().await.unwrap(); + let sessions = session_mgr.get_active_sessions(OWNER_ID).await.unwrap(); assert!(sessions.is_empty()); } @@ -171,15 +175,15 @@ async fn gs2_multiple_sessions_returned() { let uid2 = create_user(&repo, "p2", "telegram").await; session_mgr - .get_or_create_session(&uid1, "c1", "gemini", None) + .get_or_create_session(OWNER_ID, &uid1, "c1", "gemini", None) .await .unwrap(); session_mgr - .get_or_create_session(&uid2, "c2", "acp", None) + .get_or_create_session(OWNER_ID, &uid2, "c2", "acp", None) .await .unwrap(); - let sessions = session_mgr.get_active_sessions().await.unwrap(); + let sessions = session_mgr.get_active_sessions(OWNER_ID).await.unwrap(); assert_eq!(sessions.len(), 2); for s in &sessions { @@ -201,11 +205,11 @@ async fn pc1_same_user_different_chat() { let uid = create_user(&repo, "p1", "telegram").await; let s1 = session_mgr - .get_or_create_session(&uid, "chatA", "gemini", None) + .get_or_create_session(OWNER_ID, &uid, "chatA", "gemini", None) .await .unwrap(); let s2 = session_mgr - .get_or_create_session(&uid, "chatB", "gemini", None) + .get_or_create_session(OWNER_ID, &uid, "chatB", "gemini", None) .await .unwrap(); @@ -226,11 +230,11 @@ async fn pc2_different_users_same_chat() { let uid2 = create_user(&repo, "p2", "telegram").await; let s1 = session_mgr - .get_or_create_session(&uid1, "chatA", "gemini", None) + .get_or_create_session(OWNER_ID, &uid1, "chatA", "gemini", None) .await .unwrap(); let s2 = session_mgr - .get_or_create_session(&uid2, "chatA", "gemini", None) + .get_or_create_session(OWNER_ID, &uid2, "chatA", "gemini", None) .await .unwrap(); @@ -246,11 +250,11 @@ async fn pc3_same_user_same_chat_reuses() { let uid = create_user(&repo, "p1", "telegram").await; let s1 = session_mgr - .get_or_create_session(&uid, "chatA", "gemini", None) + .get_or_create_session(OWNER_ID, &uid, "chatA", "gemini", None) .await .unwrap(); let s2 = session_mgr - .get_or_create_session(&uid, "chatA", "gemini", None) + .get_or_create_session(OWNER_ID, &uid, "chatA", "gemini", None) .await .unwrap(); @@ -267,22 +271,22 @@ async fn ru3_revoke_clears_sessions() { let uid2 = create_user(&repo, "p2", "telegram").await; session_mgr - .get_or_create_session(&uid1, "c1", "gemini", None) + .get_or_create_session(OWNER_ID, &uid1, "c1", "gemini", None) .await .unwrap(); session_mgr - .get_or_create_session(&uid1, "c2", "acp", None) + .get_or_create_session(OWNER_ID, &uid1, "c2", "acp", None) .await .unwrap(); session_mgr - .get_or_create_session(&uid2, "c1", "gemini", None) + .get_or_create_session(OWNER_ID, &uid2, "c1", "gemini", None) .await .unwrap(); // Cleanup user1 sessions - session_mgr.cleanup_user_sessions(&uid1).await.unwrap(); + session_mgr.cleanup_user_sessions(OWNER_ID, &uid1).await.unwrap(); - let sessions = repo.get_all_sessions().await.unwrap(); + let sessions = repo.get_all_sessions(OWNER_ID).await.unwrap(); assert_eq!(sessions.len(), 1); assert_eq!(sessions[0].user_id, uid2); } @@ -409,7 +413,7 @@ async fn action_session_new_resets_existing() { assert_ne!(sid1, sid3); // Only 1 session should exist for this user+chat in the DB - let all = repo.get_all_sessions().await.unwrap(); + let all = repo.get_all_sessions(OWNER_ID).await.unwrap(); let user_sessions: Vec<_> = all.iter().filter(|s| s.chat_id.as_deref() == Some("chat1")).collect(); assert_eq!(user_sessions.len(), 1); } diff --git a/crates/aionui-channel/tests/stream_relay_test.rs b/crates/aionui-channel/tests/stream_relay_test.rs index 6691cec3b..42135555f 100644 --- a/crates/aionui-channel/tests/stream_relay_test.rs +++ b/crates/aionui-channel/tests/stream_relay_test.rs @@ -13,6 +13,7 @@ use tokio::sync::broadcast; #[test] fn relay_config_fields() { let config = RelayConfig { + owner_user_id: "system_default_user".to_owned(), platform: PluginType::Telegram, plugin_id: "telegram".into(), chat_id: "123".into(), @@ -30,6 +31,7 @@ async fn relay_sends_thinking_then_final_message() { let recorder = Arc::new(MessageRecorder::new()); let config = RelayConfig { + owner_user_id: "system_default_user".to_owned(), platform: PluginType::Telegram, plugin_id: "telegram".into(), chat_id: "chat_1".into(), @@ -71,6 +73,7 @@ async fn relay_handles_error_event() { let recorder = Arc::new(MessageRecorder::new()); let config = RelayConfig { + owner_user_id: "system_default_user".to_owned(), platform: PluginType::Telegram, plugin_id: "telegram".into(), chat_id: "chat_1".into(), @@ -102,6 +105,7 @@ async fn weixin_flushes_pending_text_before_tool_call() { let recorder = Arc::new(MessageRecorder::new()); let config = RelayConfig { + owner_user_id: "system_default_user".to_owned(), platform: PluginType::Weixin, plugin_id: "weixin".into(), chat_id: "chat_1".into(), @@ -154,6 +158,7 @@ async fn telegram_does_not_flush_text_before_tool_call() { let recorder = Arc::new(MessageRecorder::new()); let config = RelayConfig { + owner_user_id: "system_default_user".to_owned(), platform: PluginType::Telegram, plugin_id: "telegram".into(), chat_id: "chat_1".into(), @@ -196,6 +201,7 @@ async fn weixin_skips_flush_when_buffer_is_empty() { let recorder = Arc::new(MessageRecorder::new()); let config = RelayConfig { + owner_user_id: "system_default_user".to_owned(), platform: PluginType::Weixin, plugin_id: "weixin".into(), chat_id: "chat_1".into(), @@ -233,6 +239,7 @@ async fn relay_handles_channel_closed() { let recorder = Arc::new(MessageRecorder::new()); let config = RelayConfig { + owner_user_id: "system_default_user".to_owned(), platform: PluginType::Telegram, plugin_id: "telegram".into(), chat_id: "chat_1".into(), diff --git a/crates/aionui-channel/tests/telegram_integration.rs b/crates/aionui-channel/tests/telegram_integration.rs index c4491de22..e2067982b 100644 --- a/crates/aionui-channel/tests/telegram_integration.rs +++ b/crates/aionui-channel/tests/telegram_integration.rs @@ -25,6 +25,8 @@ mod telegram_tests { use std::sync::Arc; use tokio::sync::mpsc; + const OWNER_USER_ID: &str = "system_default_user"; + // -- Test infrastructure ------------------------------------------------ struct MockBroadcaster { @@ -203,7 +205,9 @@ mod telegram_tests { let factory = telegram_factory(); let config = make_config_value(Some("bot:123")); - let result = manager.enable_plugin("nonexistent", &config, &factory).await; + let result = manager + .enable_plugin(OWNER_USER_ID, "nonexistent", &config, &factory) + .await; assert!(result.is_err()); } @@ -215,7 +219,9 @@ mod telegram_tests { let factory = telegram_factory(); let config = make_config_value(Some("bad-token")); - let result = manager.enable_plugin("telegram", &config, &factory).await; + let result = manager + .enable_plugin(OWNER_USER_ID, "telegram", &config, &factory) + .await; assert!(result.is_err()); } @@ -225,7 +231,7 @@ mod telegram_tests { async fn disable_without_db_row_returns_error() { let (manager, _repo, _bc) = setup().await; // Plugin was never enabled (no DB row), so update_plugin_status fails - let result = manager.disable_plugin("telegram").await; + let result = manager.disable_plugin(OWNER_USER_ID, "telegram").await; assert!(result.is_err()); } @@ -234,7 +240,7 @@ mod telegram_tests { #[tokio::test] async fn get_plugin_status_empty() { let (manager, _repo, _bc) = setup().await; - let statuses = manager.get_plugin_status().await.unwrap(); + let statuses = manager.get_plugin_status(OWNER_USER_ID).await.unwrap(); assert!(statuses.is_empty()); } @@ -244,7 +250,7 @@ mod telegram_tests { async fn restore_plugins_none_stored() { let (manager, _repo, _bc) = setup().await; let factory = telegram_factory(); - let result = manager.restore_plugins(&factory).await; + let result = manager.restore_plugins(OWNER_USER_ID, &factory).await; assert!(result.is_ok()); assert_eq!(manager.active_plugin_count(), 0); } @@ -254,6 +260,6 @@ mod telegram_tests { #[tokio::test] async fn is_plugin_running_false_when_not_enabled() { let (manager, _repo, _bc) = setup().await; - assert!(!manager.is_plugin_running("telegram")); + assert!(!manager.is_plugin_running(OWNER_USER_ID, "telegram")); } } diff --git a/crates/aionui-channel/tests/weixin_integration.rs b/crates/aionui-channel/tests/weixin_integration.rs index c97f82f35..607c6a44b 100644 --- a/crates/aionui-channel/tests/weixin_integration.rs +++ b/crates/aionui-channel/tests/weixin_integration.rs @@ -25,6 +25,8 @@ mod weixin_tests { use std::sync::Arc; use tokio::sync::mpsc; + const OWNER_USER_ID: &str = "system_default_user"; + // -- Test infrastructure ------------------------------------------------ struct MockBroadcaster { @@ -218,7 +220,9 @@ mod weixin_tests { let factory = weixin_factory(); let config = make_config_value(Some("tok_1"), Some("acc_1")); - let result = manager.enable_plugin("nonexistent", &config, &factory).await; + let result = manager + .enable_plugin(OWNER_USER_ID, "nonexistent", &config, &factory) + .await; assert!(result.is_err()); } @@ -227,7 +231,7 @@ mod weixin_tests { #[tokio::test] async fn disable_without_db_row_returns_error() { let (manager, _repo, _bc) = setup().await; - let result = manager.disable_plugin("weixin").await; + let result = manager.disable_plugin(OWNER_USER_ID, "weixin").await; assert!(result.is_err()); } @@ -236,7 +240,7 @@ mod weixin_tests { #[tokio::test] async fn get_plugin_status_empty() { let (manager, _repo, _bc) = setup().await; - let statuses = manager.get_plugin_status().await.unwrap(); + let statuses = manager.get_plugin_status(OWNER_USER_ID).await.unwrap(); assert!(statuses.is_empty()); } @@ -246,7 +250,7 @@ mod weixin_tests { async fn restore_plugins_none_stored() { let (manager, _repo, _bc) = setup().await; let factory = weixin_factory(); - let result = manager.restore_plugins(&factory).await; + let result = manager.restore_plugins(OWNER_USER_ID, &factory).await; assert!(result.is_ok()); assert_eq!(manager.active_plugin_count(), 0); } @@ -256,7 +260,7 @@ mod weixin_tests { #[tokio::test] async fn is_plugin_running_false_when_not_enabled() { let (manager, _repo, _bc) = setup().await; - assert!(!manager.is_plugin_running("weixin")); + assert!(!manager.is_plugin_running(OWNER_USER_ID, "weixin")); } // -- Login event serialization ------------------------------------------ diff --git a/crates/aionui-common/src/constants.rs b/crates/aionui-common/src/constants.rs index 530535862..cf794ae83 100644 --- a/crates/aionui-common/src/constants.rs +++ b/crates/aionui-common/src/constants.rs @@ -98,6 +98,13 @@ fn bool_field(value: &serde_json::Value, key: &str) -> bool { value.get(key).and_then(serde_json::Value::as_bool) == Some(true) } +// --- Image processing --- + +pub const SUPPORTED_IMAGE_EXTENSIONS: &[&str] = &[".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".tiff", ".svg"]; +/// Remote image download size limit (5 MB). +pub const REMOTE_IMAGE_MAX_SIZE: usize = 5 * 1024 * 1024; +pub const REMOTE_IMAGE_MAX_REDIRECTS: u32 = 5; + #[cfg(test)] mod tests { use super::*; @@ -163,10 +170,3 @@ mod tests { assert!(!supports_team_cli_fallback(Some(&json!({"execution": {"cli": false}})))); } } - -// --- Image processing --- - -pub const SUPPORTED_IMAGE_EXTENSIONS: &[&str] = &[".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".tiff", ".svg"]; -/// Remote image download size limit (5 MB). -pub const REMOTE_IMAGE_MAX_SIZE: usize = 5 * 1024 * 1024; -pub const REMOTE_IMAGE_MAX_REDIRECTS: u32 = 5; diff --git a/crates/aionui-common/src/enums.rs b/crates/aionui-common/src/enums.rs index 9eb45dd7f..9c88aaa4b 100644 --- a/crates/aionui-common/src/enums.rs +++ b/crates/aionui-common/src/enums.rs @@ -249,6 +249,9 @@ pub enum AgentKillReason { /// The requested runtime capabilities changed, so the in-memory task must /// be rebuilt before handling the next turn. RuntimeCapabilityChanged, + /// The owning user's Core session was revoked, so foreground runtime state + /// and agent processes for that user must be torn down. + SessionRevoked, } /// Preview content type for document preview history. diff --git a/crates/aionui-common/src/hooks.rs b/crates/aionui-common/src/hooks.rs index 23558183d..610ebe1e3 100644 --- a/crates/aionui-common/src/hooks.rs +++ b/crates/aionui-common/src/hooks.rs @@ -10,10 +10,10 @@ use async_trait::async_trait; /// `ConversationService::delete`. /// /// Implementors are responsible for cleaning up their per-conversation state -/// (kill agent processes, drop cron jobs, etc.). Hooks run sequentially in -/// registration order; failures must be logged inside the hook and not +/// (kill agent processes, drop cron job state, etc.). Hooks run sequentially +/// in registration order; failures must be logged inside the hook and not /// propagated. #[async_trait] pub trait OnConversationDelete: Send + Sync { - async fn on_conversation_deleted(&self, conversation_id: &str); + async fn on_conversation_deleted(&self, user_id: &str, conversation_id: &str); } diff --git a/crates/aionui-conversation/Cargo.toml b/crates/aionui-conversation/Cargo.toml index 86f608dc8..08835b91c 100644 --- a/crates/aionui-conversation/Cargo.toml +++ b/crates/aionui-conversation/Cargo.toml @@ -28,6 +28,7 @@ thiserror.workspace = true tokio = { workspace = true, features = ["test-util"] } tempfile.workspace = true tracing-subscriber.workspace = true +sqlx.workspace = true # Enable the `AgentInstance::Mock` variant so tests can build fake agents # through the trait-object escape hatch without spawning real CLI processes. aionui-ai-agent = { workspace = true, features = ["test-support"] } diff --git a/crates/aionui-conversation/src/acp_error_recovery.rs b/crates/aionui-conversation/src/acp_error_recovery.rs index e03b68e55..5b30ca63b 100644 --- a/crates/aionui-conversation/src/acp_error_recovery.rs +++ b/crates/aionui-conversation/src/acp_error_recovery.rs @@ -15,6 +15,7 @@ use crate::runtime_persistence::RuntimeWriteKind; impl ConversationService { async fn clear_conversation_model_seed_after_model_not_found( &self, + user_id: &str, conversation_id: &str, error_code: Option, ) { @@ -28,7 +29,7 @@ impl ConversationService { return; } - let row = match self.conversation_repo().get(conversation_id).await { + let row = match self.conversation_repo().get(user_id, conversation_id).await { Ok(Some(row)) => row, Ok(None) => { warn!( @@ -106,7 +107,7 @@ impl ConversationService { updated_at: Some(now_ms()), ..Default::default() }; - if let Err(err) = self.conversation_repo().update(conversation_id, &update).await { + if let Err(err) = self.conversation_repo().update(user_id, conversation_id, &update).await { warn!( conversation_id, ?previous_model_id, @@ -122,7 +123,7 @@ impl ConversationService { .source .as_deref() .and_then(|value| string_to_enum::(value).ok()); - self.broadcast_list_changed(conversation_id, "updated", source.as_ref()); + self.broadcast_list_changed(user_id, conversation_id, "updated", source.as_ref()); info!( conversation_id, ?previous_model_id, @@ -134,6 +135,7 @@ impl ConversationService { async fn clear_persisted_acp_model_after_model_not_found( &self, + user_id: &str, conversation_id: &str, error_code: Option, ) { @@ -147,11 +149,16 @@ impl ConversationService { return; } - let previous_model_id = match self.acp_session_repo().load_runtime_state(conversation_id).await { + let previous_model_id = match self + .acp_session_repo() + .load_runtime_state_for_user(user_id, conversation_id) + .await + { Ok(Some(state)) => state.current_model_id, Ok(None) => None, Err(err) => { warn!( + user_id, conversation_id, error = %err, "Failed to load ACP persisted model before clearing after model_not_found" @@ -166,11 +173,12 @@ impl ConversationService { }; match self .acp_session_repo() - .save_runtime_state(conversation_id, ¶ms) + .save_runtime_state_for_user(user_id, conversation_id, ¶ms) .await { Ok(true) => { info!( + user_id, conversation_id, ?previous_model_id, error_code = ?error_code, @@ -180,6 +188,7 @@ impl ConversationService { } Ok(false) => { warn!( + user_id, conversation_id, ?previous_model_id, error_code = ?error_code, @@ -189,6 +198,7 @@ impl ConversationService { } Err(err) => { warn!( + user_id, conversation_id, ?previous_model_id, error = %err, @@ -202,6 +212,7 @@ impl ConversationService { pub(crate) async fn evict_acp_task_after_terminal_error( &self, + user_id: &str, conversation_id: &str, agent_type: AgentType, outcome: &RelayOutcome, @@ -225,9 +236,9 @@ impl ConversationService { task_manager .kill_and_wait(conversation_id, Some(AgentKillReason::AgentErrorRecovery)) .await; - self.clear_persisted_acp_model_after_model_not_found(conversation_id, error_code) + self.clear_persisted_acp_model_after_model_not_found(user_id, conversation_id, error_code) .await; - self.clear_conversation_model_seed_after_model_not_found(conversation_id, error_code) + self.clear_conversation_model_seed_after_model_not_found(user_id, conversation_id, error_code) .await; info!( conversation_id, diff --git a/crates/aionui-conversation/src/message_persistence.rs b/crates/aionui-conversation/src/message_persistence.rs index a69f8857c..268ff9335 100644 --- a/crates/aionui-conversation/src/message_persistence.rs +++ b/crates/aionui-conversation/src/message_persistence.rs @@ -9,6 +9,7 @@ use crate::service::ConversationService; impl ConversationService { pub(crate) async fn persist_send_failure_tip( &self, + user_id: &str, conversation_id: &str, err: &AgentSendError, top_level_code: Option<&'static str>, @@ -54,7 +55,7 @@ impl ConversationService { created_at: now_ms(), }; - if let Err(store_err) = self.conversation_repo().insert_message(&row).await { + if let Err(store_err) = self.conversation_repo().insert_message(user_id, &row).await { warn!( conversation_id, error = %ErrorChain(&store_err), diff --git a/crates/aionui-conversation/src/routes.rs b/crates/aionui-conversation/src/routes.rs index bbb06d2b3..2a7b38844 100644 --- a/crates/aionui-conversation/src/routes.rs +++ b/crates/aionui-conversation/src/routes.rs @@ -31,6 +31,9 @@ impl From for ApiError { } ConversationError::Archived { reason, .. } => ApiError::ConversationArchived(reason), ConversationError::BadRequest { reason } => ApiError::BadRequest(reason), + ConversationError::Busy { reason } if reason.starts_with("CROSS_ACCOUNT_REFERENCE:") => { + ApiError::coded(StatusCode::CONFLICT, "CROSS_ACCOUNT_REFERENCE", reason, None) + } ConversationError::Busy { reason } => ApiError::Conflict(reason), ConversationError::Forbidden { reason } => ApiError::Forbidden(reason), ConversationError::NotFoundReason { reason } => ApiError::NotFound(reason), @@ -419,9 +422,9 @@ async fn check_approval( async fn active_count( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, ) -> Result>, ApiError> { - let count = state.task_manager.active_count(); + let count = state.service.active_count_for_user(&user.id).await?; Ok(Json(ApiResponse::ok(ActiveCountResponse { count }))) } @@ -435,6 +438,16 @@ mod error_mapping_tests { assert!(matches!(app, ApiError::NotFound(message) if message == "Conversation conv_1 not found")); } + #[test] + fn cross_account_reference_maps_to_stable_conflict_code() { + let app = ApiError::from(ConversationError::Busy { + reason: "CROSS_ACCOUNT_REFERENCE: acp_session conversation 'conv-1' belongs to another user".into(), + }); + + assert_eq!(app.status_code(), StatusCode::CONFLICT); + assert_eq!(app.error_code(), "CROSS_ACCOUNT_REFERENCE"); + } + #[test] fn conversation_archived_maps_to_app_conversation_archived() { let app = ApiError::from(ConversationError::Archived { diff --git a/crates/aionui-conversation/src/routes_aux.rs b/crates/aionui-conversation/src/routes_aux.rs index 5a56e27dd..ec5d646d9 100644 --- a/crates/aionui-conversation/src/routes_aux.rs +++ b/crates/aionui-conversation/src/routes_aux.rs @@ -31,7 +31,7 @@ pub fn conversation_ops_routes(state: ConversationRouterState) -> Router { async fn set_config_option( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path((id, option_id)): Path<(String, String)>, body: Result, JsonRejection>, ) -> Result>, ApiError> { @@ -39,7 +39,7 @@ async fn set_config_option( Ok(Json(ApiResponse::ok( state .service - .set_config_option(&id, &option_id, req) + .set_config_option(&user.id, &id, &option_id, req) .await .map_err(ApiError::from)?, ))) @@ -47,24 +47,24 @@ async fn set_config_option( async fn get_usage( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>>, ApiError> { Ok(Json(ApiResponse::ok( - state.service.get_usage(&id).await.map_err(ApiError::from)?, + state.service.get_usage(&user.id, &id).await.map_err(ApiError::from)?, ))) } async fn side_question( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, Json(req): Json, ) -> Result>, ApiError> { Ok(Json(ApiResponse::ok( state .service - .handle_side_question(&id, req) + .handle_side_question(&user.id, &id, req) .await .map_err(ApiError::from)?, ))) @@ -72,24 +72,28 @@ async fn side_question( async fn get_slash_commands( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>>, ApiError> { Ok(Json(ApiResponse::ok( - state.service.get_slash_commands(&id).await.map_err(ApiError::from)?, + state + .service + .get_slash_commands(&user.id, &id) + .await + .map_err(ApiError::from)?, ))) } async fn browse_workspace( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, Query(query): Query, ) -> Result>>, ApiError> { Ok(Json(ApiResponse::ok( state .service - .browse_workspace(&id, query) + .browse_workspace(&user.id, &id, query) .await .map_err(ApiError::from)?, ))) diff --git a/crates/aionui-conversation/src/runtime_completion.rs b/crates/aionui-conversation/src/runtime_completion.rs index b7b394128..6b7161550 100644 --- a/crates/aionui-conversation/src/runtime_completion.rs +++ b/crates/aionui-conversation/src/runtime_completion.rs @@ -11,6 +11,7 @@ use crate::runtime_persistence::{RuntimePersistenceCoordinator, RuntimeWriteKind #[derive(Clone)] pub struct RuntimeCompletionPublisher { + user_id: String, repo: Arc, broadcaster: Arc, persistence: RuntimePersistenceCoordinator, @@ -18,11 +19,13 @@ pub struct RuntimeCompletionPublisher { impl RuntimeCompletionPublisher { pub fn new( + user_id: String, repo: Arc, broadcaster: Arc, persistence: RuntimePersistenceCoordinator, ) -> Self { Self { + user_id, repo, broadcaster, persistence, @@ -42,7 +45,7 @@ impl RuntimeCompletionPublisher { return; } - match self.repo.get(conversation_id).await { + match self.repo.get(&self.user_id, conversation_id).await { Ok(Some(_)) => {} Ok(None) => { debug!( @@ -67,7 +70,7 @@ impl RuntimeCompletionPublisher { updated_at: Some(now_ms()), ..Default::default() }; - if let Err(error) = self.repo.update(conversation_id, &update).await { + if let Err(error) = self.repo.update(&self.user_id, conversation_id, &update).await { error!( conversation_id, turn_id, @@ -78,6 +81,7 @@ impl RuntimeCompletionPublisher { } let payload = json!({ + "user_id": self.user_id, "conversation_id": conversation_id, "session_id": conversation_id, "turn_id": turn_id, diff --git a/crates/aionui-conversation/src/runtime_state.rs b/crates/aionui-conversation/src/runtime_state.rs index a0269916c..fca8532fd 100644 --- a/crates/aionui-conversation/src/runtime_state.rs +++ b/crates/aionui-conversation/src/runtime_state.rs @@ -199,6 +199,30 @@ impl ConversationRuntimeStateService { .unwrap_or(false) } + pub fn clear_conversation(&self, conversation_id: &str) { + match self.state.lock() { + Ok(mut state) => { + let had_active_turn = state.active_turns.remove(conversation_id).is_some(); + let had_deleting = state.deleting_conversations.remove(conversation_id); + let had_cancelling = state.cancelling_conversations.remove(conversation_id); + if had_active_turn || had_deleting || had_cancelling { + info!( + conversation_id, + had_active_turn, had_deleting, had_cancelling, "conversation runtime state cleared" + ); + drop(state); + self.release_notify.notify_waiters(); + } + } + Err(_) => { + warn!( + conversation_id, + "conversation runtime state lock poisoned while clearing conversation" + ); + } + } + } + pub fn mark_shutting_down(&self) -> usize { match self.state.lock() { Ok(mut state) => { @@ -530,6 +554,23 @@ mod tests { assert!(!state.is_cancelling("conv-1")); } + #[test] + fn clear_conversation_removes_active_turn_and_lifecycle_flags() { + let state = Arc::new(ConversationRuntimeStateService::default()); + let _claim = state + .try_claim_turn("conv-1", "turn-1") + .expect("claim should be created"); + + state.mark_deleting("conv-1"); + state.mark_cancelling("conv-1"); + state.clear_conversation("conv-1"); + + assert!(!state.is_claimed("conv-1")); + assert!(!state.is_deleting("conv-1")); + assert!(!state.is_cancelling("conv-1")); + assert!(state.active_turn_id_for("conv-1").is_none()); + } + #[test] fn summary_uses_claim_as_starting_state() { let state = Arc::new(ConversationRuntimeStateService::default()); diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index d8d1ea0af..79bffa983 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -464,8 +464,9 @@ impl ConversationService { /// Register a hook to be notified when a conversation is deleted. /// /// Hooks are dispatched sequentially in registration order before - /// `delete()` removes the conversation row. Used by `aionui-app` to wire up `WorkerTaskManagerImpl` - /// (kill the agent process) and `CronService` (cascade-delete cron jobs). + /// `delete()` removes the conversation row. Used by `aionui-app` to wire + /// up `WorkerTaskManagerImpl` (kill the agent process) and `CronService` + /// (clear deleted workspace references from cron jobs). pub fn with_delete_hook(&self, hook: Arc) { if let Ok(mut guard) = self.delete_hooks.write() { guard.push(hook); @@ -553,8 +554,9 @@ impl ConversationService { RuntimePersistenceCoordinator::new(self.runtime_state()) } - pub(crate) fn completion_publisher(&self) -> RuntimeCompletionPublisher { + pub(crate) fn completion_publisher(&self, user_id: &str) -> RuntimeCompletionPublisher { RuntimeCompletionPublisher::new( + user_id.to_owned(), self.conversation_repo.clone(), self.broadcaster.clone(), self.runtime_persistence(), @@ -583,6 +585,51 @@ impl ConversationService { .summary_from_parts(conversation_id, task_status, has_task, pending_confirmations) } + pub async fn active_count_for_user(&self, user_id: &str) -> Result { + let mut count = 0; + for conversation_id in self.task_manager.active_conversation_ids() { + let belongs_to_user = self + .conversation_repo + .get(user_id, &conversation_id) + .await + .map_err(|e| ConversationError::internal(format!("Failed to load conversation: {e}")))? + .is_some(); + if belongs_to_user { + count += 1; + } + } + Ok(count) + } + + pub async fn terminate_runtime_for_user(&self, user_id: &str) -> Result { + let mut terminated = 0; + for conversation_id in self.task_manager.active_conversation_ids() { + let belongs_to_user = self + .conversation_repo + .get(user_id, &conversation_id) + .await + .map_err(|e| ConversationError::internal(format!("Failed to load conversation: {e}")))? + .is_some(); + if !belongs_to_user { + continue; + } + + self.task_manager + .kill_and_wait(&conversation_id, Some(AgentKillReason::SessionRevoked)) + .await; + self.runtime_state.clear_conversation(&conversation_id); + terminated += 1; + } + if terminated > 0 { + tracing::info!( + user_id, + terminated, + "terminated conversation runtimes for revoked session" + ); + } + Ok(terminated) + } + async fn send_message_response( &self, conversation_id: &str, @@ -596,14 +643,20 @@ impl ConversationService { } } - pub async fn complete_turn(&self, conversation_id: &str, turn_id: &str) { + pub async fn complete_turn(&self, user_id: &str, conversation_id: &str, turn_id: &str) { let runtime = self.runtime_summary_for(conversation_id).await; - self.completion_publisher() + self.completion_publisher(user_id) .publish(conversation_id, turn_id, Some(runtime)) .await; } - pub(crate) async fn complete_released_turn(&self, conversation_id: &str, turn_id: &str, was_deleting: bool) { + pub(crate) async fn complete_released_turn( + &self, + user_id: &str, + conversation_id: &str, + turn_id: &str, + was_deleting: bool, + ) { if was_deleting { debug!( conversation_id, @@ -612,20 +665,28 @@ impl ConversationService { return; } - self.complete_turn(conversation_id, turn_id).await; + self.complete_turn(user_id, conversation_id, turn_id).await; } } // ── Conversation CRUD ─────────────────────────────────────────────── impl ConversationService { - async fn attach_assistant_identity(&self, response: &mut ConversationResponse) -> Result<(), ConversationError> { + async fn attach_assistant_identity( + &self, + user_id: &str, + response: &mut ConversationResponse, + ) -> Result<(), ConversationError> { if response.assistant.is_some() { return Ok(()); } - if let Some(snapshot) = self.conversation_repo.get_assistant_snapshot(&response.id).await? { - response.assistant = Some(self.assistant_identity_from_snapshot(&snapshot).await?); + if let Some(snapshot) = self + .conversation_repo + .get_assistant_snapshot(user_id, &response.id) + .await? + { + response.assistant = Some(self.assistant_identity_from_snapshot(user_id, &snapshot).await?); } Ok(()) @@ -633,14 +694,17 @@ impl ConversationService { async fn assistant_identity_from_snapshot( &self, + user_id: &str, snapshot: &ConversationAssistantSnapshotRow, ) -> Result { let runtime_backend = self - .resolve_assistant_agent_binding(&snapshot.agent_id) + .resolve_assistant_agent_binding(user_id, &snapshot.agent_id) .await? .map(|binding| binding.runtime_backend) .unwrap_or_else(|| snapshot.agent_id.clone()); - let current_definition = self.current_assistant_definition(&snapshot.assistant_id).await?; + let current_definition = self + .current_assistant_definition(user_id, &snapshot.assistant_id) + .await?; let (source, name, avatar) = match current_definition { Some(definition) => ( definition.source, @@ -671,13 +735,14 @@ impl ConversationService { async fn current_assistant_definition( &self, + user_id: &str, assistant_id: &str, ) -> Result, ConversationError> { let Some(definition_repo) = self.assistant_definition_repo() else { return Ok(None); }; definition_repo - .get_by_assistant_id(assistant_id) + .get_by_assistant_id_for_user(user_id, assistant_id) .await .map_err(|e| ConversationError::internal(format!("assistant definition lookup failed: {e}"))) } @@ -697,6 +762,7 @@ impl ConversationService { let source = req.source.unwrap_or(ConversationSource::Aionui); let mut extra = req.extra; + strip_request_owner_user_id(&mut extra); let assistant_id = req .assistant @@ -717,7 +783,7 @@ impl ConversationService { .unwrap_or_default(); let assistant_snapshot = match assistant_id.as_deref() { Some(id) => { - self.resolve_assistant_snapshot(id, assistant_locale.as_deref(), &assistant_overrides, &extra) + self.resolve_assistant_snapshot(user_id, id, assistant_locale.as_deref(), &assistant_overrides, &extra) .await? } None => None, @@ -953,6 +1019,7 @@ impl ConversationService { && !initial_skills.is_empty() && let Some(rel_dirs) = native_skills_dirs( &self.agent_metadata_repo, + user_id, &effective_type, effective_backend .as_ref() @@ -961,7 +1028,10 @@ impl ConversationService { ) .await { - let resolved = self.skill_resolver.resolve_skills(&initial_skills).await; + let resolved = self + .skill_resolver + .resolve_skills_for_user(user_id, &initial_skills) + .await; if !resolved.is_empty() { let rel_dirs_refs: Vec<&str> = rel_dirs.iter().map(String::as_str).collect(); let n = self @@ -1011,7 +1081,9 @@ impl ConversationService { None => None, }; - let mcp_support = self.resolve_mcp_support_policy(&effective_type, &extra).await?; + let mcp_support = self + .resolve_mcp_support_policy(user_id, &effective_type, &extra) + .await?; let mut selected_row_ids: Vec = Vec::new(); let mut selected_mcp_names: Vec = Vec::new(); let mut selected_mcp_statuses: Vec = Vec::new(); @@ -1025,11 +1097,11 @@ impl ConversationService { if let Some(repo) = repo { let rows = match selected_mcp_server_ids.as_ref() { Some(ids) => repo - .list_by_ids_any(ids) + .list_by_ids_any(user_id, ids) .await .map_err(|e| ConversationError::internal(format!("Failed to load selected MCP servers: {e}")))?, None => repo - .list() + .list(user_id) .await .map_err(|e| ConversationError::internal(format!("Failed to list MCP servers: {e}")))?, }; @@ -1130,25 +1202,28 @@ impl ConversationService { .map_err(|e| ConversationError::internal(format!("Failed to serialize assistant MCP snapshot: {e}")))?; self.conversation_repo - .upsert_assistant_snapshot(&UpsertConversationAssistantSnapshotParams { - conversation_id: &row.id, - assistant_definition_id: &snapshot.assistant_definition_id, - assistant_id: &snapshot.assistant_id, - assistant_source: &snapshot.assistant_source, - agent_id: &snapshot.agent_id, - rules_content: &snapshot.rules.content, - default_model_mode: &snapshot.default_modes.model, - resolved_model_id: snapshot.resolved_defaults.model.as_deref(), - default_permission_mode: &snapshot.default_modes.permission, - resolved_permission_value: snapshot.resolved_defaults.permission.as_deref(), - default_thought_level_mode: &snapshot.default_modes.thought_level, - resolved_thought_level_value: snapshot.resolved_defaults.thought_level.as_deref(), - default_skills_mode: &snapshot.default_modes.skills, - resolved_skill_ids: &resolved_skill_ids, - resolved_disabled_builtin_skill_ids: &resolved_disabled_builtin_skill_ids, - default_mcps_mode: &snapshot.default_modes.mcps, - resolved_mcp_ids: &resolved_mcp_ids, - }) + .upsert_assistant_snapshot( + user_id, + &UpsertConversationAssistantSnapshotParams { + conversation_id: &row.id, + assistant_definition_id: &snapshot.assistant_definition_id, + assistant_id: &snapshot.assistant_id, + assistant_source: &snapshot.assistant_source, + agent_id: &snapshot.agent_id, + rules_content: &snapshot.rules.content, + default_model_mode: &snapshot.default_modes.model, + resolved_model_id: snapshot.resolved_defaults.model.as_deref(), + default_permission_mode: &snapshot.default_modes.permission, + resolved_permission_value: snapshot.resolved_defaults.permission.as_deref(), + default_thought_level_mode: &snapshot.default_modes.thought_level, + resolved_thought_level_value: snapshot.resolved_defaults.thought_level.as_deref(), + default_skills_mode: &snapshot.default_modes.skills, + resolved_skill_ids: &resolved_skill_ids, + resolved_disabled_builtin_skill_ids: &resolved_disabled_builtin_skill_ids, + default_mcps_mode: &snapshot.default_modes.mcps, + resolved_mcp_ids: &resolved_mcp_ids, + }, + ) .await? .ok_or_else(|| ConversationError::internal("assistant snapshot upsert returned no row"))?; } @@ -1157,12 +1232,13 @@ impl ConversationService { // conversation_id). Other agent types have no session-level // state so we only create it for ACP. if effective_type == AgentType::Acp { - self.create_acp_session_row(&id, &extra, assistant_snapshot.as_ref()) + self.create_acp_session_row(user_id, &id, &extra, assistant_snapshot.as_ref()) .await?; } if let Some(snapshot) = assistant_snapshot.as_ref() { - self.persist_assistant_preferences_from_snapshot(snapshot).await?; + self.persist_assistant_preferences_from_snapshot(user_id, snapshot) + .await?; } let mut response = row_to_response(row, &self.workspace_root)?; @@ -1181,7 +1257,7 @@ impl ConversationService { }); } - self.broadcast_list_changed(&response.id, "created", response.source.as_ref()); + self.broadcast_list_changed(user_id, &response.id, "created", response.source.as_ref()); log_conversation_created(&response, &extra); @@ -1191,6 +1267,7 @@ impl ConversationService { #[tracing::instrument(skip_all, fields(conversation_id = %conversation_id))] async fn create_acp_session_row( &self, + user_id: &str, conversation_id: &str, extra: &serde_json::Value, assistant_snapshot: Option<&AssistantSnapshot>, @@ -1233,7 +1310,7 @@ impl ConversationService { Some(id) => id.to_owned(), None if !backend.is_empty() && agent_source == "builtin" => self .agent_metadata_repo - .find_builtin_by_backend(backend) + .find_builtin_by_backend_for_user(user_id, backend) .await .map_err(|e| ConversationError::internal(format!("agent_metadata lookup: {e}")))? .map(|row| row.id) @@ -1242,6 +1319,7 @@ impl ConversationService { }; let params = CreateAcpSessionParams { + user_id, conversation_id, agent_source, agent_id: &resolved_agent_id, @@ -1271,7 +1349,7 @@ impl ConversationService { context_usage_json: None, }; self.acp_session_repo - .save_runtime_state(conversation_id, ¶ms) + .save_runtime_state_for_user(user_id, conversation_id, ¶ms) .await .map_err(|e| ConversationError::internal(format!("Failed to seed acp_session runtime state: {e}")))?; } @@ -1280,11 +1358,12 @@ impl ConversationService { async fn resolve_assistant_agent_binding( &self, + user_id: &str, value: &str, ) -> Result, ConversationError> { let rows = self .agent_metadata_repo - .list_all() + .list_all_for_user(user_id) .await .map_err(|e| ConversationError::internal(format!("agent_metadata lookup failed: {e}")))?; Ok(resolve_agent_binding_from_rows(&rows, value)) @@ -1292,6 +1371,7 @@ impl ConversationService { async fn resolve_assistant_snapshot( &self, + user_id: &str, assistant_id: &str, locale: Option<&str>, overrides: &AssistantConversationOverrides, @@ -1306,7 +1386,7 @@ impl ConversationService { }; let Some(definition) = definition_repo - .get_by_assistant_id(assistant_id) + .get_by_assistant_id_for_user(user_id, assistant_id) .await .map_err(|e| ConversationError::internal(format!("assistant definition lookup failed: {e}")))? else { @@ -1314,11 +1394,11 @@ impl ConversationService { }; let state = state_repo - .get(&definition.id) + .get_for_user(user_id, &definition.id) .await .map_err(|e| ConversationError::internal(format!("assistant state lookup failed: {e}")))?; let preference = preference_repo - .get(&definition.id) + .get_for_user(user_id, &definition.id) .await .map_err(|e| ConversationError::internal(format!("assistant preference lookup failed: {e}")))?; @@ -1390,7 +1470,7 @@ impl ConversationService { let rules_content = if let Some(dispatcher) = self.assistant_dispatcher() { dispatcher - .read_rule(assistant_id, locale) + .read_rule(user_id, assistant_id, locale) .await .map_err(|e| ConversationError::internal(format!("assistant rule lookup failed: {e}")))? } else { @@ -1406,7 +1486,7 @@ impl ConversationService { .and_then(|row| row.agent_id_override.clone()) .unwrap_or_else(|| definition.agent_id.clone()); let agent_binding = self - .resolve_assistant_agent_binding(&effective_agent_id) + .resolve_assistant_agent_binding(user_id, &effective_agent_id) .await? .ok_or_else(|| ConversationError::BadRequest { reason: format!("assistant agent `{effective_agent_id}` is not registered in agent_metadata"), @@ -1452,6 +1532,7 @@ impl ConversationService { async fn persist_assistant_preferences_from_snapshot( &self, + user_id: &str, snapshot: &AssistantSnapshot, ) -> Result<(), ConversationError> { let Some(preference_repo) = self.assistant_preference_repo() else { @@ -1459,7 +1540,7 @@ impl ConversationService { }; let existing_preference = preference_repo - .get(&snapshot.assistant_definition_id) + .get_for_user(user_id, &snapshot.assistant_definition_id) .await .map_err(|e| ConversationError::internal(format!("assistant preference lookup failed: {e}")))?; let last_model_id = if snapshot.default_modes.model == "auto" { @@ -1510,15 +1591,18 @@ impl ConversationService { }; preference_repo - .upsert(&aionui_db::UpsertAssistantPreferenceParams { - assistant_definition_id: &snapshot.assistant_definition_id, - last_model_id: last_model_id.as_deref(), - last_permission_value: last_permission_value.as_deref(), - last_thought_level_value: last_thought_level_value.as_deref(), - last_skill_ids: &last_skill_ids, - last_disabled_builtin_skill_ids: &last_disabled_builtin_skill_ids, - last_mcp_ids: &last_mcp_ids, - }) + .upsert_for_user( + user_id, + &aionui_db::UpsertAssistantPreferenceParams { + assistant_definition_id: &snapshot.assistant_definition_id, + last_model_id: last_model_id.as_deref(), + last_permission_value: last_permission_value.as_deref(), + last_thought_level_value: last_thought_level_value.as_deref(), + last_skill_ids: &last_skill_ids, + last_disabled_builtin_skill_ids: &last_disabled_builtin_skill_ids, + last_mcp_ids: &last_mcp_ids, + }, + ) .await .map_err(|e| ConversationError::internal(format!("assistant preference upsert failed: {e}")))?; @@ -1527,12 +1611,13 @@ impl ConversationService { pub(crate) async fn persist_runtime_assistant_snapshot( &self, + user_id: &str, conversation_id: &str, updates: AssistantRuntimePreferenceUpdate<'_>, ) -> Result<(), ConversationError> { let Some(snapshot) = self .conversation_repo - .get_assistant_snapshot(conversation_id) + .get_assistant_snapshot(user_id, conversation_id) .await .map_err(|e| { ConversationError::internal(format!( @@ -1544,27 +1629,30 @@ impl ConversationService { }; self.conversation_repo - .upsert_assistant_snapshot(&UpsertConversationAssistantSnapshotParams { - conversation_id: &snapshot.conversation_id, - assistant_definition_id: &snapshot.assistant_definition_id, - assistant_id: &snapshot.assistant_id, - assistant_source: &snapshot.assistant_source, - agent_id: &snapshot.agent_id, - rules_content: &snapshot.rules_content, - default_model_mode: &snapshot.default_model_mode, - resolved_model_id: updates.model.or(snapshot.resolved_model_id.as_deref()), - default_permission_mode: &snapshot.default_permission_mode, - resolved_permission_value: updates.permission.or(snapshot.resolved_permission_value.as_deref()), - default_thought_level_mode: &snapshot.default_thought_level_mode, - resolved_thought_level_value: updates - .thought_level - .or(snapshot.resolved_thought_level_value.as_deref()), - default_skills_mode: &snapshot.default_skills_mode, - resolved_skill_ids: &snapshot.resolved_skill_ids, - resolved_disabled_builtin_skill_ids: &snapshot.resolved_disabled_builtin_skill_ids, - default_mcps_mode: &snapshot.default_mcps_mode, - resolved_mcp_ids: &snapshot.resolved_mcp_ids, - }) + .upsert_assistant_snapshot( + user_id, + &UpsertConversationAssistantSnapshotParams { + conversation_id: &snapshot.conversation_id, + assistant_definition_id: &snapshot.assistant_definition_id, + assistant_id: &snapshot.assistant_id, + assistant_source: &snapshot.assistant_source, + agent_id: &snapshot.agent_id, + rules_content: &snapshot.rules_content, + default_model_mode: &snapshot.default_model_mode, + resolved_model_id: updates.model.or(snapshot.resolved_model_id.as_deref()), + default_permission_mode: &snapshot.default_permission_mode, + resolved_permission_value: updates.permission.or(snapshot.resolved_permission_value.as_deref()), + default_thought_level_mode: &snapshot.default_thought_level_mode, + resolved_thought_level_value: updates + .thought_level + .or(snapshot.resolved_thought_level_value.as_deref()), + default_skills_mode: &snapshot.default_skills_mode, + resolved_skill_ids: &snapshot.resolved_skill_ids, + resolved_disabled_builtin_skill_ids: &snapshot.resolved_disabled_builtin_skill_ids, + default_mcps_mode: &snapshot.default_mcps_mode, + resolved_mcp_ids: &snapshot.resolved_mcp_ids, + }, + ) .await .map_err(|e| ConversationError::internal(format!("assistant snapshot upsert failed: {e}")))?; @@ -1573,6 +1661,7 @@ impl ConversationService { pub(crate) async fn persist_runtime_assistant_preferences( &self, + user_id: &str, conversation_id: &str, updates: AssistantRuntimePreferenceUpdate<'_>, ) -> Result<(), ConversationError> { @@ -1584,7 +1673,7 @@ impl ConversationService { let persisted_snapshot = self .conversation_repo - .get_assistant_snapshot(conversation_id) + .get_assistant_snapshot(user_id, conversation_id) .await .map_err(|e| { ConversationError::internal(format!( @@ -1593,11 +1682,15 @@ impl ConversationService { })?; let fallback = if persisted_snapshot.is_none() { - let Some(conversation) = self.conversation_repo.get(conversation_id).await.map_err(|e| { - ConversationError::internal(format!( - "Failed to load conversation for assistant preference sync: {e}" - )) - })? + let Some(conversation) = self + .conversation_repo + .get(user_id, conversation_id) + .await + .map_err(|e| { + ConversationError::internal(format!( + "Failed to load conversation for assistant preference sync: {e}" + )) + })? else { return Ok(()); }; @@ -1631,7 +1724,7 @@ impl ConversationService { return Ok(()); }; let Some(definition) = definition_repo - .get_by_assistant_id(&assistant_id) + .get_by_assistant_id_for_user(user_id, &assistant_id) .await .map_err(|e| ConversationError::internal(format!("assistant definition lookup failed: {e}")))? else { @@ -1669,7 +1762,7 @@ impl ConversationService { }; let existing_preference = preference_repo - .get(&definition_id) + .get_for_user(user_id, &definition_id) .await .map_err(|e| ConversationError::internal(format!("assistant preference lookup failed: {e}")))?; @@ -1705,24 +1798,27 @@ impl ConversationService { }; preference_repo - .upsert(&aionui_db::UpsertAssistantPreferenceParams { - assistant_definition_id: &definition_id, - last_model_id: last_model_id.as_deref(), - last_permission_value: last_permission_value.as_deref(), - last_thought_level_value: last_thought_level_value.as_deref(), - last_skill_ids: existing_preference - .as_ref() - .map(|row| row.last_skill_ids.as_str()) - .unwrap_or("[]"), - last_disabled_builtin_skill_ids: existing_preference - .as_ref() - .map(|row| row.last_disabled_builtin_skill_ids.as_str()) - .unwrap_or("[]"), - last_mcp_ids: existing_preference - .as_ref() - .map(|row| row.last_mcp_ids.as_str()) - .unwrap_or("[]"), - }) + .upsert_for_user( + user_id, + &aionui_db::UpsertAssistantPreferenceParams { + assistant_definition_id: &definition_id, + last_model_id: last_model_id.as_deref(), + last_permission_value: last_permission_value.as_deref(), + last_thought_level_value: last_thought_level_value.as_deref(), + last_skill_ids: existing_preference + .as_ref() + .map(|row| row.last_skill_ids.as_str()) + .unwrap_or("[]"), + last_disabled_builtin_skill_ids: existing_preference + .as_ref() + .map(|row| row.last_disabled_builtin_skill_ids.as_str()) + .unwrap_or("[]"), + last_mcp_ids: existing_preference + .as_ref() + .map(|row| row.last_mcp_ids.as_str()) + .unwrap_or("[]"), + }, + ) .await .map_err(|e| ConversationError::internal(format!("assistant runtime preference upsert failed: {e}")))?; @@ -1737,16 +1833,15 @@ impl ConversationService { pub async fn get(&self, user_id: &str, id: &str) -> Result { let row = self .conversation_repo - .get(id) + .get(user_id, id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: id.to_owned() })?; let mut extra: serde_json::Value = serde_json::from_str(&row.extra) .map_err(|e| ConversationError::internal(format!("Invalid extra JSON: {e}")))?; - self.backfill_extra_inplace(&row.id, &mut extra).await; + self.backfill_extra_inplace(user_id, &row.id, &mut extra).await; let mut response = row_to_response_with_extra(row, extra, &self.workspace_root)?; - self.attach_assistant_identity(&mut response).await?; + self.attach_assistant_identity(user_id, &mut response).await?; response.runtime = Some(self.runtime_summary_for(id).await); Ok(response) } @@ -1786,10 +1881,10 @@ impl ConversationService { continue; } }; - self.backfill_extra_inplace(&row_id, &mut extra).await; + self.backfill_extra_inplace(user_id, &row_id, &mut extra).await; match row_to_response_with_extra(row, extra, &self.workspace_root) { Ok(mut resp) => { - self.attach_assistant_identity(&mut resp).await?; + self.attach_assistant_identity(user_id, &mut resp).await?; items.push(resp); } Err(err) => warn!( @@ -1822,9 +1917,8 @@ impl ConversationService { ) -> Result { let existing = self .conversation_repo - .get(id) + .get(user_id, id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: id.to_owned() })?; let existing_type: AgentType = string_to_enum(&existing.r#type)?; @@ -1878,6 +1972,7 @@ impl ConversationService { let mut existing_extra: serde_json::Value = serde_json::from_str(&existing.extra).unwrap_or_else(|_| serde_json::json!({})); merge_json(&mut existing_extra, new_extra); + strip_request_owner_user_id(&mut existing_extra); if existing_type == AgentType::Aionrs && let Some(obj) = existing_extra.as_object_mut() && obj.remove("model").is_some() @@ -1923,11 +2018,12 @@ impl ConversationService { updated_at: Some(now), }; - self.conversation_repo.update(id, &updates).await?; + self.conversation_repo.update(user_id, id, &updates).await?; if let Some(model) = req.model.as_ref() { let selected_model = model.use_model.as_deref().unwrap_or(model.model.as_str()); self.persist_runtime_assistant_snapshot( + user_id, id, AssistantRuntimePreferenceUpdate { model: Some(selected_model), @@ -1936,6 +2032,7 @@ impl ConversationService { ) .await?; self.persist_runtime_assistant_preferences( + user_id, id, AssistantRuntimePreferenceUpdate { model: Some(selected_model), @@ -1958,14 +2055,14 @@ impl ConversationService { // Re-fetch to return the updated version let updated = self .conversation_repo - .get(id) + .get(user_id, id) .await? .ok_or_else(|| ConversationError::internal("Conversation vanished after update"))?; let response = row_to_response(updated, &self.workspace_root)?; info!("Conversation updated"); - self.broadcast_list_changed(id, "updated", response.source.as_ref()); + self.broadcast_list_changed(user_id, id, "updated", response.source.as_ref()); Ok(response) } @@ -1975,15 +2072,20 @@ impl ConversationService { /// (e.g. `TeamSessionService::ensure_session` writing /// `team_mcp_stdio_config`) where a full `update()` would kill the agent /// on a spurious model comparison. - #[tracing::instrument(skip_all, fields(conversation_id = %conversation_id))] - pub async fn update_extra(&self, conversation_id: &str, patch: serde_json::Value) -> Result<(), ConversationError> { - let existing = - self.conversation_repo - .get(conversation_id) - .await? - .ok_or_else(|| ConversationError::NotFound { - id: conversation_id.to_owned(), - })?; + #[tracing::instrument(skip_all, fields(user_id = %user_id, conversation_id = %conversation_id))] + pub async fn update_extra( + &self, + user_id: &str, + conversation_id: &str, + patch: serde_json::Value, + ) -> Result<(), ConversationError> { + let existing = self + .conversation_repo + .get(user_id, conversation_id) + .await? + .ok_or_else(|| ConversationError::NotFound { + id: conversation_id.to_owned(), + })?; let mut merged: serde_json::Value = serde_json::from_str(&existing.extra).unwrap_or_else(|_| serde_json::json!({})); @@ -2000,15 +2102,29 @@ impl ConversationService { updated_at: Some(now_ms()), ..Default::default() }; - self.conversation_repo.update(conversation_id, &updates).await?; + self.conversation_repo + .update(user_id, conversation_id, &updates) + .await?; debug!("Conversation extra merged"); Ok(()) } - pub async fn save_acp_runtime_mode(&self, conversation_id: &str, mode: &str) -> Result<(), ConversationError> { + pub async fn save_acp_runtime_mode( + &self, + user_id: &str, + conversation_id: &str, + mode: &str, + ) -> Result<(), ConversationError> { + self.conversation_repo + .get(user_id, conversation_id) + .await? + .ok_or_else(|| ConversationError::NotFound { + id: conversation_id.to_owned(), + })?; + let runtime_state = self .acp_session_repo - .load_runtime_state(conversation_id) + .load_runtime_state_for_user(user_id, conversation_id) .await .map_err(|e| ConversationError::internal(format!("Failed to load runtime mode state: {e}")))?; let mut config_selections = runtime_state @@ -2024,7 +2140,7 @@ impl ConversationService { ..Default::default() }; self.acp_session_repo - .save_runtime_state(conversation_id, ¶ms) + .save_runtime_state_for_user(user_id, conversation_id, ¶ms) .await .map_err(|e| ConversationError::internal(format!("Failed to persist runtime mode: {e}")))?; Ok(()) @@ -2038,9 +2154,8 @@ impl ConversationService { // Get existing to retrieve source for broadcast and verify ownership let existing = self .conversation_repo - .get(id) + .get(user_id, id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: id.to_owned() })?; let source: Option = existing @@ -2057,10 +2172,10 @@ impl ConversationService { let hooks: Vec> = self.delete_hooks.read().map(|guard| guard.clone()).unwrap_or_default(); for hook in hooks { - hook.on_conversation_deleted(id).await; + hook.on_conversation_deleted(user_id, id).await; } - if let Err(err) = self.conversation_repo.delete(id).await { + if let Err(err) = self.conversation_repo.delete(user_id, id).await { self.runtime_state.clear_deleting(id); return Err(err.into()); } @@ -2070,7 +2185,7 @@ impl ConversationService { // No FK / CASCADE on `acp_session`: clean it up here so non-ACP // conversations that used to be ACP (shouldn't happen but is // cheap to cover) still drop their orphaned session row. - if let Err(err) = self.acp_session_repo.delete(id).await { + if let Err(err) = self.acp_session_repo.delete_for_user(user_id, id).await { warn!( error = %ErrorChain(&err), "Failed to delete acp_session row on conversation delete" @@ -2103,7 +2218,7 @@ impl ConversationService { } info!("Conversation deleted"); - self.broadcast_list_changed(id, "deleted", source.as_ref()); + self.broadcast_list_changed(user_id, id, "deleted", source.as_ref()); Ok(()) } @@ -2129,14 +2244,17 @@ impl ConversationService { pub async fn reset(&self, user_id: &str, id: &str) -> Result<(), ConversationError> { // Verify existence and ownership self.conversation_repo - .get(id) + .get(user_id, id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: id.to_owned() })?; // Delete all messages - self.conversation_repo.delete_messages_by_conversation(id).await?; - self.conversation_repo.delete_artifacts_by_conversation(id).await?; + self.conversation_repo + .delete_messages_by_conversation(user_id, id) + .await?; + self.conversation_repo + .delete_artifacts_by_conversation(user_id, id) + .await?; // Reset status to pending let now = now_ms(); @@ -2145,7 +2263,7 @@ impl ConversationService { updated_at: Some(now), ..Default::default() }; - self.conversation_repo.update(id, &updates).await?; + self.conversation_repo.update(user_id, id, &updates).await?; info!("Conversation reset"); Ok(()) @@ -2158,16 +2276,15 @@ impl ConversationService { id: &str, ) -> Result, ConversationError> { self.conversation_repo - .get(id) + .get(user_id, id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: id.to_owned() })?; let rows = self.conversation_repo.list_associated(user_id, id).await?; let mut items = Vec::with_capacity(rows.len()); for row in rows { let mut response = row_to_response(row, &self.workspace_root)?; - self.attach_assistant_identity(&mut response).await?; + self.attach_assistant_identity(user_id, &mut response).await?; items.push(response); } Ok(items) @@ -2183,7 +2300,7 @@ impl ConversationService { let mut items = Vec::with_capacity(rows.len()); for row in rows { let mut response = row_to_response(row, &self.workspace_root)?; - self.attach_assistant_identity(&mut response).await?; + self.attach_assistant_identity(user_id, &mut response).await?; items.push(response); } Ok(items) @@ -2218,9 +2335,8 @@ impl ConversationService { ) -> Result { // Verify conversation exists and belongs to user self.conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: conversation_id.to_owned(), })?; @@ -2249,7 +2365,7 @@ impl ConversationService { let page = self .conversation_repo - .list_messages_page(conversation_id, &MessagePageParams { limit, direction }) + .list_messages_page(user_id, conversation_id, &MessagePageParams { limit, direction }) .await?; let mut compacted_count = 0usize; @@ -2319,16 +2435,15 @@ impl ConversationService { message_id: &str, ) -> Result { self.conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: conversation_id.to_owned(), })?; let row = self .conversation_repo - .get_message(conversation_id, message_id) + .get_message(user_id, conversation_id, message_id) .await? .ok_or_else(|| ConversationError::MessageNotFound { id: message_id.to_owned(), @@ -2356,16 +2471,15 @@ impl ConversationService { conversation_id: &str, ) -> Result { self.conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: conversation_id.to_owned(), })?; let mut items = self .conversation_repo - .list_artifacts(conversation_id) + .list_artifacts(user_id, conversation_id) .await? .into_iter() .map(row_to_artifact_response) @@ -2373,7 +2487,7 @@ impl ConversationService { let mut legacy_items = self .conversation_repo - .list_legacy_cron_trigger_messages(conversation_id) + .list_legacy_cron_trigger_messages(user_id, conversation_id) .await? .into_iter() .filter_map(|row| legacy_cron_trigger_to_artifact(row).ok()) @@ -2398,9 +2512,8 @@ impl ConversationService { req: UpdateConversationArtifactRequest, ) -> Result { self.conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: conversation_id.to_owned(), })?; @@ -2412,18 +2525,20 @@ impl ConversationService { let row = self .conversation_repo - .update_artifact_status(conversation_id, artifact_id, &status, now_ms()) + .update_artifact_status(user_id, conversation_id, artifact_id, &status, now_ms()) .await? .ok_or_else(|| ConversationError::ArtifactNotFound { id: artifact_id.to_owned(), })?; let response = row_to_artifact_response(row)?; - self.broadcaster.broadcast(WebSocketMessage::new( - "conversation.artifact", - serde_json::to_value(&response) - .map_err(|e| ConversationError::internal(format!("Failed to serialize artifact event: {e}")))?, - )); + let mut payload = serde_json::to_value(&response) + .map_err(|e| ConversationError::internal(format!("Failed to serialize artifact event: {e}")))?; + if let Some(obj) = payload.as_object_mut() { + obj.insert("user_id".to_owned(), serde_json::Value::String(user_id.to_owned())); + } + self.broadcaster + .broadcast(WebSocketMessage::new("conversation.artifact", payload)); Ok(response) } @@ -2473,9 +2588,8 @@ impl ConversationService { task_manager: &Arc, ) -> Result { self.conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: conversation_id.to_owned(), })?; @@ -2501,9 +2615,8 @@ impl ConversationService { task_manager: &Arc, ) -> Result<(), ConversationError> { self.conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: conversation_id.to_owned(), })?; @@ -2524,6 +2637,7 @@ impl ConversationService { if let Some(conf_id) = conf_id { let payload = serde_json::json!({ + "user_id": user_id, "conversation_id": conversation_id, "id": conf_id, }); @@ -2544,9 +2658,8 @@ impl ConversationService { task_manager: &Arc, ) -> Result { self.conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: conversation_id.to_owned(), })?; @@ -2587,9 +2700,8 @@ impl ConversationService { // Verify conversation exists and belongs to user let row = self .conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: conversation_id.to_owned(), })?; @@ -2634,11 +2746,11 @@ impl ConversationService { { let mut turn_claim = turn_claim; let was_deleting = turn_claim.release(); - self.complete_released_turn(conversation_id, &turn_id, was_deleting) + self.complete_released_turn(user_id, conversation_id, &turn_id, was_deleting) .await; return Ok(self.send_message_response(conversation_id, user_msg_id, turn_id).await); } - if let Err(e) = self.conversation_repo.insert_message(&user_msg).await { + if let Err(e) = self.conversation_repo.insert_message(user_id, &user_msg).await { warn!(msg_id = %user_msg_id, error = %ErrorChain(&e), "Failed to insert user message"); return Err(e.into()); } @@ -2648,6 +2760,7 @@ impl ConversationService { self.broadcaster.broadcast(WebSocketMessage::new( "message.userCreated", serde_json::json!({ + "user_id": user_id, "conversation_id": conversation_id, "msg_id": &user_msg_id, "content": &req.content, @@ -2670,6 +2783,7 @@ impl ConversationService { let top_level_code = err.error_code(); let send_error = AgentSendError::from_agent_error(err.to_agent_error()); self.persist_and_broadcast_send_failure_tip( + user_id, conversation_id, &turn_id, &send_error, @@ -2678,7 +2792,7 @@ impl ConversationService { .await; let mut turn_claim = turn_claim; let was_deleting = turn_claim.release(); - self.complete_released_turn(conversation_id, &turn_id, was_deleting) + self.complete_released_turn(user_id, conversation_id, &turn_id, was_deleting) .await; return Ok(self.send_message_response(conversation_id, user_msg_id, turn_id).await); } @@ -2727,9 +2841,8 @@ impl ConversationService { let row = self .conversation_repo - .get(&request.conversation_id) + .get(&request.user_id, &request.conversation_id) .await? - .filter(|r| r.user_id == request.user_id) .ok_or_else(|| ConversationError::NotFound { id: request.conversation_id.clone(), })?; @@ -2754,7 +2867,7 @@ impl ConversationService { if self .runtime_persistence() .allows(&request.conversation_id, RuntimeWriteKind::UserMessage) - && let Err(e) = self.conversation_repo.insert_message(&user_msg).await + && let Err(e) = self.conversation_repo.insert_message(&request.user_id, &user_msg).await { warn!( msg_id = %user_msg.id, @@ -2763,7 +2876,7 @@ impl ConversationService { ); let mut turn_claim = turn_claim; let was_deleting = turn_claim.release(); - self.complete_released_turn(&request.conversation_id, &turn_id, was_deleting) + self.complete_released_turn(&request.user_id, &request.conversation_id, &turn_id, was_deleting) .await; return Err(e.into()); } @@ -2782,6 +2895,7 @@ impl ConversationService { let top_level_code = err.error_code(); let send_error = AgentSendError::from_agent_error(err.to_agent_error()); self.persist_and_broadcast_send_failure_tip( + &request.user_id, &request.conversation_id, &turn_id, &send_error, @@ -2790,7 +2904,7 @@ impl ConversationService { .await; let mut turn_claim = turn_claim; let was_deleting = turn_claim.release(); - self.complete_released_turn(&request.conversation_id, &turn_id, was_deleting) + self.complete_released_turn(&request.user_id, &request.conversation_id, &turn_id, was_deleting) .await; return Ok(ConversationAgentTurnOutcome { conversation_id: request.conversation_id.clone(), @@ -2838,11 +2952,13 @@ impl ConversationService { pub async fn latest_conversation_error_message( &self, + user_id: &str, conversation_id: &str, ) -> Result, ConversationError> { let page = self .conversation_repo .list_messages_page( + user_id, conversation_id, &MessagePageParams { limit: 30, @@ -2856,13 +2972,14 @@ impl ConversationService { pub(crate) async fn persist_and_broadcast_send_failure_tip( &self, + user_id: &str, conversation_id: &str, turn_id: &str, err: &AgentSendError, top_level_code: Option<&'static str>, ) { let Some(row) = self - .persist_send_failure_tip(conversation_id, err, top_level_code) + .persist_send_failure_tip(user_id, conversation_id, err, top_level_code) .await else { return; @@ -2874,6 +2991,7 @@ impl ConversationService { self.broadcaster.broadcast(WebSocketMessage::new( "message.stream", serde_json::json!({ + "user_id": user_id, "conversation_id": row.conversation_id, "msg_id": msg_id, "turn_id": turn_id, @@ -2894,13 +3012,14 @@ impl ConversationService { /// Used by paths outside the normal user→agent turn (e.g. the team /// scheduler writing an incoming teammate message as a left bubble in the /// target agent's conversation so the UI shows who spoke). - pub async fn insert_raw_message(&self, row: &MessageRow) -> Result<(), ConversationError> { - self.conversation_repo.insert_message(row).await?; + pub async fn insert_raw_message(&self, user_id: &str, row: &MessageRow) -> Result<(), ConversationError> { + self.conversation_repo.insert_message(user_id, row).await?; let msg_id = row.msg_id.clone().unwrap_or_else(|| row.id.clone()); let content_value: serde_json::Value = serde_json::from_str(&row.content).unwrap_or_else(|_| serde_json::Value::String(row.content.clone())); let payload = serde_json::json!({ + "user_id": user_id, "conversation_id": row.conversation_id, "msg_id": msg_id, "type": row.r#type, @@ -2926,9 +3045,8 @@ impl ConversationService { ) -> Result { // Verify conversation exists and belongs to user self.conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: conversation_id.to_owned(), })?; @@ -2999,7 +3117,7 @@ impl ConversationService { conversation_id: &str, active_leases: &ActiveLeaseRegistry, ) -> Result<(), ConversationError> { - let row = match self.conversation_repo.get(conversation_id).await { + let row = match self.conversation_repo.get(user_id, conversation_id).await { Ok(row) => row, Err(error) => { warn!( @@ -3013,7 +3131,7 @@ impl ConversationService { } }; - let Some(row) = row.filter(|row| row.user_id == user_id) else { + let Some(row) = row else { debug!( kind = "conversation", conversation_id, user_id, "Conversation active lease renew rejected" @@ -3060,9 +3178,8 @@ impl ConversationService { ) -> Result { let row = self .conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: conversation_id.to_owned(), })?; @@ -3102,9 +3219,8 @@ impl ConversationService { ) -> Result<(AgentInstance, bool), ConversationError> { let row = self .conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? - .filter(|r| r.user_id == user_id) .ok_or_else(|| ConversationError::NotFound { id: conversation_id.to_owned(), })?; @@ -3146,7 +3262,7 @@ impl ConversationService { }; // Persist auto-resolved workspace if factory picked a different path. - self.maybe_persist_workspace(conversation_id, &stored_workspace, agent.workspace()) + self.maybe_persist_workspace(user_id, conversation_id, &stored_workspace, agent.workspace()) .await?; info!(conversation_id, phase, "Conversation runtime recovered"); @@ -3232,7 +3348,7 @@ impl ConversationService { } let snapshot = self .conversation_repo - .get_assistant_snapshot(&row.id) + .get_assistant_snapshot(&row.user_id, &row.id) .await .map_err(|e| { ConversationError::internal(format!( @@ -3334,6 +3450,7 @@ impl ConversationService { let Some(rel_dirs) = native_skills_dirs( &self.agent_metadata_repo, + &context.conversation.user_id, &context.conversation.agent_type, backend.as_ref(), ) @@ -3345,7 +3462,10 @@ impl ConversationService { return; } - let resolved = self.skill_resolver.resolve_skills(&skill_names).await; + let resolved = self + .skill_resolver + .resolve_skills_for_user(&context.conversation.user_id, &skill_names) + .await; if resolved.is_empty() { return; } @@ -3371,6 +3491,7 @@ impl ConversationService { /// path here so the frontend can display the workspace panel correctly. pub(crate) async fn maybe_persist_workspace( &self, + user_id: &str, conversation_id: &str, stored_workspace: &str, resolved_workspace: &str, @@ -3388,7 +3509,7 @@ impl ConversationService { // Fetch latest extra, merge the resolved workspace path in, and persist. let row = self .conversation_repo - .get(conversation_id) + .get(user_id, conversation_id) .await? .ok_or_else(|| ConversationError::internal("Conversation vanished during workspace sync"))?; @@ -3403,7 +3524,7 @@ impl ConversationService { updated_at: Some(now_ms()), ..Default::default() }; - self.conversation_repo.update(conversation_id, &update).await?; + self.conversation_repo.update(user_id, conversation_id, &update).await?; debug!( conversation_id, @@ -3416,11 +3537,13 @@ impl ConversationService { /// Broadcast a `conversation.listChanged` WebSocket event. pub(crate) fn broadcast_list_changed( &self, + user_id: &str, conversation_id: &str, action: &str, source: Option<&ConversationSource>, ) { let payload = serde_json::json!({ + "user_id": user_id, "conversation_id": conversation_id, "action": action, "source": source, @@ -3437,7 +3560,7 @@ impl ConversationService { /// Persists the mutation asynchronously; failures are logged and /// swallowed so a read path never 500s because of a backfill write /// failure. - async fn backfill_extra_inplace(&self, conversation_id: &str, extra: &mut serde_json::Value) { + async fn backfill_extra_inplace(&self, user_id: &str, conversation_id: &str, extra: &mut serde_json::Value) { let auto_inject = self.skill_resolver.auto_inject_names().await; let mut mutated = backfill_skills_if_missing(extra, &auto_inject); mutated |= backfill_cron_job_id_alias(extra); @@ -3459,7 +3582,7 @@ impl ConversationService { extra: Some(serialized), ..Default::default() }; - if let Err(e) = self.conversation_repo.update(conversation_id, &update).await { + if let Err(e) = self.conversation_repo.update(user_id, conversation_id, &update).await { warn!( conversation_id, error = %ErrorChain(&e), @@ -3519,6 +3642,12 @@ fn normalize_workspace_extra(extra: &mut serde_json::Value) -> Result<(), Conver Ok(()) } +fn strip_request_owner_user_id(extra: &mut serde_json::Value) { + if let Some(obj) = extra.as_object_mut() { + obj.remove("user_id"); + } +} + fn team_id_from_extra(extra: &str) -> Option { TeamSessionBinding::team_id_marker_from_extra_str(extra) } @@ -3723,6 +3852,7 @@ fn context_skill_names(context: &AgentSessionContext) -> Vec { /// rely on prompt injection instead. async fn native_skills_dirs( repo: &Arc, + user_id: &str, agent_type: &AgentType, backend: Option<&serde_json::Value>, ) -> Option> { @@ -3730,7 +3860,11 @@ async fn native_skills_dirs( && let Some(serde_json::Value::String(vendor)) = backend && !vendor.is_empty() { - let row = repo.find_builtin_by_backend(vendor).await.ok().flatten()?; + let row = repo + .find_builtin_by_backend_for_user(user_id, vendor) + .await + .ok() + .flatten()?; let raw = row.native_skills_dirs?; return serde_json::from_str::>(&raw).ok(); } @@ -3742,11 +3876,12 @@ async fn native_skills_dirs( impl ConversationService { async fn resolve_mcp_support_policy( &self, + user_id: &str, agent_type: &AgentType, extra: &serde_json::Value, ) -> Result { match agent_type { - AgentType::Acp => resolve_acp_mcp_support_policy(&self.agent_metadata_repo, extra).await, + AgentType::Acp => resolve_acp_mcp_support_policy(&self.agent_metadata_repo, user_id, extra).await, AgentType::Aionrs => Ok(McpSupportPolicy::AIONRS), _ => Ok(McpSupportPolicy::AIONRS), } @@ -3755,6 +3890,7 @@ impl ConversationService { async fn resolve_acp_mcp_support_policy( repo: &Arc, + user_id: &str, extra: &serde_json::Value, ) -> Result { let agent_id = extra @@ -3772,12 +3908,12 @@ async fn resolve_acp_mcp_support_policy( let row = match agent_id { Some(id) => repo - .get(id) + .get_for_user(user_id, id) .await .map_err(|e| ConversationError::internal(format!("agent_metadata lookup: {e}")))?, None if agent_source == "builtin" => match backend { Some(vendor) => repo - .find_builtin_by_backend(vendor) + .find_builtin_by_backend_for_user(user_id, vendor) .await .map_err(|e| ConversationError::internal(format!("agent_metadata lookup: {e}")))?, None => None, @@ -3953,6 +4089,7 @@ fn enum_to_db(val: &T) -> Result pub(crate) async fn persist_session_key( repo: &Arc, persistence: &RuntimePersistenceCoordinator, + user_id: &str, conversation_id: &str, session_key: &str, ) { @@ -3960,7 +4097,7 @@ pub(crate) async fn persist_session_key( return; } - let row = match repo.get(conversation_id).await { + let row = match repo.get(user_id, conversation_id).await { Ok(Some(r)) => r, _ => return, }; @@ -3986,7 +4123,7 @@ pub(crate) async fn persist_session_key( updated_at: Some(now_ms()), ..Default::default() }; - if let Err(e) = repo.update(conversation_id, &update).await { + if let Err(e) = repo.update(user_id, conversation_id, &update).await { warn!(conversation_id, error = %ErrorChain(&e), "Failed to persist session key"); } else { debug!(conversation_id, "Persisted session key to conversation.extra"); diff --git a/crates/aionui-conversation/src/service_ops.rs b/crates/aionui-conversation/src/service_ops.rs index a67797d3c..b1875cc3a 100644 --- a/crates/aionui-conversation/src/service_ops.rs +++ b/crates/aionui-conversation/src/service_ops.rs @@ -27,8 +27,10 @@ impl ConversationService { pub async fn get_config_options( &self, + user_id: &str, conversation_id: &str, ) -> Result { + self.ensure_owned_conversation(user_id, conversation_id).await?; self.task(conversation_id)? .get_config_options() .await @@ -37,6 +39,7 @@ impl ConversationService { pub async fn set_config_option( &self, + user_id: &str, conversation_id: &str, option_id: &str, req: SetConfigOptionRequest, @@ -51,6 +54,7 @@ impl ConversationService { reason: "value must not be empty".into(), }); } + self.ensure_owned_conversation(user_id, conversation_id).await?; let agent = self.task(conversation_id)?; let response = match agent.set_config_option(option_id, &req.value).await { Ok(response) => response, @@ -103,7 +107,10 @@ impl ConversationService { _ => None, }; if let Some(updates) = updates { - if let Err(err) = self.persist_runtime_assistant_snapshot(conversation_id, updates).await { + if let Err(err) = self + .persist_runtime_assistant_snapshot(user_id, conversation_id, updates) + .await + { warn!( conversation_id, option_id, @@ -112,7 +119,7 @@ impl ConversationService { ); } if let Err(err) = self - .persist_runtime_assistant_preferences(conversation_id, updates) + .persist_runtime_assistant_preferences(user_id, conversation_id, updates) .await { warn!( @@ -130,14 +137,24 @@ impl ConversationService { // ── Usage / Slash commands ────────────────────────────────────── - pub async fn get_usage(&self, conversation_id: &str) -> Result, ConversationError> { + pub async fn get_usage( + &self, + user_id: &str, + conversation_id: &str, + ) -> Result, ConversationError> { + self.ensure_owned_conversation(user_id, conversation_id).await?; self.task(conversation_id)? .get_usage() .await .map_err(ConversationError::from) } - pub async fn get_slash_commands(&self, conversation_id: &str) -> Result, ConversationError> { + pub async fn get_slash_commands( + &self, + user_id: &str, + conversation_id: &str, + ) -> Result, ConversationError> { + self.ensure_owned_conversation(user_id, conversation_id).await?; self.task(conversation_id)? .get_slash_commands() .await @@ -148,9 +165,11 @@ impl ConversationService { pub async fn handle_side_question( &self, + user_id: &str, conversation_id: &str, req: SideQuestionRequest, ) -> Result { + self.ensure_owned_conversation(user_id, conversation_id).await?; // `AgentInstance::handle_side_question` already validates that the // question is non-empty; no need to duplicate the check here. self.task(conversation_id)? @@ -159,6 +178,22 @@ impl ConversationService { .map_err(ConversationError::from) } + async fn ensure_owned_conversation(&self, user_id: &str, conversation_id: &str) -> Result<(), ConversationError> { + let exists = self + .conversation_repo() + .get(user_id, conversation_id) + .await + .map_err(|e| ConversationError::internal(format!("Failed to load conversation: {e}")))? + .is_some(); + if exists { + Ok(()) + } else { + Err(ConversationError::NotFound { + id: conversation_id.to_owned(), + }) + } + } + // ── Workspace browsing ────────────────────────────────────────── /// Enumerate entries under `query.path` inside the conversation's @@ -167,6 +202,7 @@ impl ConversationService { /// depth cap of [`MAX_DIR_DEPTH`]. pub async fn browse_workspace( &self, + user_id: &str, conversation_id: &str, query: WorkspaceBrowseQuery, ) -> Result, ConversationError> { @@ -178,7 +214,7 @@ impl ConversationService { let row = self .conversation_repo() - .get(conversation_id) + .get(user_id, conversation_id) .await .map_err(|e| ConversationError::internal(format!("Failed to load conversation: {e}")))? .ok_or_else(|| ConversationError::NotFound { diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index fff583766..4c8604a4f 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -75,27 +75,39 @@ struct StaticAssistantDispatcher { #[async_trait::async_trait] impl AssistantRuleDispatcher for StaticAssistantDispatcher { - async fn read_rule(&self, id: &str, _locale: Option<&str>) -> Result { + async fn read_rule(&self, _user_id: &str, id: &str, _locale: Option<&str>) -> Result { Ok(self.rules.get(id).cloned().unwrap_or_default()) } - async fn write_rule(&self, _id: &str, _locale: Option<&str>, _content: &str) -> Result<(), ExtensionError> { + async fn write_rule( + &self, + _user_id: &str, + _id: &str, + _locale: Option<&str>, + _content: &str, + ) -> Result<(), ExtensionError> { Ok(()) } - async fn delete_rule(&self, _id: &str) -> Result { + async fn delete_rule(&self, _user_id: &str, _id: &str) -> Result { Ok(true) } - async fn read_skill(&self, _id: &str, _locale: Option<&str>) -> Result { + async fn read_skill(&self, _user_id: &str, _id: &str, _locale: Option<&str>) -> Result { Ok(String::new()) } - async fn write_skill(&self, _id: &str, _locale: Option<&str>, _content: &str) -> Result<(), ExtensionError> { + async fn write_skill( + &self, + _user_id: &str, + _id: &str, + _locale: Option<&str>, + _content: &str, + ) -> Result<(), ExtensionError> { Ok(()) } - async fn delete_skill(&self, _id: &str) -> Result { + async fn delete_skill(&self, _user_id: &str, _id: &str) -> Result { Ok(true) } } @@ -174,6 +186,7 @@ impl EventBroadcaster for MockBroadcaster { #[derive(Debug, Clone, PartialEq, Eq)] struct RecordedAvailabilityFailure { + user_id: String, agent_id: String, code: String, message: String, @@ -181,19 +194,29 @@ struct RecordedAvailabilityFailure { #[derive(Default)] struct RecordingAvailabilityFeedback { - successes: Mutex>, + successes: Mutex>, failures: Mutex>, } #[async_trait::async_trait] impl AgentAvailabilityFeedbackPort for RecordingAvailabilityFeedback { - async fn record_session_success(&self, agent_id: &str) -> Result<(), AgentError> { - self.successes.lock().unwrap().push(agent_id.to_owned()); + async fn record_session_success(&self, user_id: &str, agent_id: &str) -> Result<(), AgentError> { + self.successes + .lock() + .unwrap() + .push((user_id.to_owned(), agent_id.to_owned())); Ok(()) } - async fn record_session_failure(&self, agent_id: &str, code: &str, message: &str) -> Result<(), AgentError> { + async fn record_session_failure( + &self, + user_id: &str, + agent_id: &str, + code: &str, + message: &str, + ) -> Result<(), AgentError> { self.failures.lock().unwrap().push(RecordedAvailabilityFailure { + user_id: user_id.to_owned(), agent_id: agent_id.to_owned(), code: code.to_owned(), message: message.to_owned(), @@ -232,6 +255,7 @@ fn message_is_after_cursor(message: &MessageRow, cursor: &MessagePageCursor) -> async fn repo_messages_asc(repo: &Arc, conv_id: &str, limit: u32) -> Vec { repo.list_messages_page( + "user_1", conv_id, &MessagePageParams { limit, @@ -245,9 +269,14 @@ async fn repo_messages_asc(repo: &Arc, conv_id: &str, limit: u32) -> V #[async_trait::async_trait] impl IConversationRepository for MockRepo { - async fn get(&self, id: &str) -> Result, aionui_db::DbError> { + async fn get(&self, user_id: &str, id: &str) -> Result, aionui_db::DbError> { + let rows = self.rows.lock().unwrap(); + Ok(rows.iter().find(|r| r.user_id == user_id && r.id == id).cloned()) + } + + async fn owner_user_id(&self, id: &str) -> Result, aionui_db::DbError> { let rows = self.rows.lock().unwrap(); - Ok(rows.iter().find(|r| r.id == id).cloned()) + Ok(rows.iter().find(|r| r.id == id).map(|row| row.user_id.clone())) } async fn create(&self, row: &ConversationRow) -> Result<(), aionui_db::DbError> { @@ -255,11 +284,11 @@ impl IConversationRepository for MockRepo { Ok(()) } - async fn update(&self, id: &str, updates: &ConversationRowUpdate) -> Result<(), aionui_db::DbError> { + async fn update(&self, user_id: &str, id: &str, updates: &ConversationRowUpdate) -> Result<(), aionui_db::DbError> { let mut rows = self.rows.lock().unwrap(); let row = rows .iter_mut() - .find(|r| r.id == id) + .find(|r| r.user_id == user_id && r.id == id) .ok_or_else(|| aionui_db::DbError::NotFound(format!("Conversation {id}")))?; if let Some(name) = &updates.name { @@ -286,10 +315,10 @@ impl IConversationRepository for MockRepo { Ok(()) } - async fn delete(&self, id: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, user_id: &str, id: &str) -> Result<(), aionui_db::DbError> { let mut rows = self.rows.lock().unwrap(); let len_before = rows.len(); - rows.retain(|r| r.id != id); + rows.retain(|r| !(r.user_id == user_id && r.id == id)); if rows.len() == len_before { return Err(aionui_db::DbError::NotFound(format!("Conversation {id}"))); } @@ -349,6 +378,7 @@ impl IConversationRepository for MockRepo { async fn get_assistant_snapshot( &self, + _user_id: &str, conversation_id: &str, ) -> Result, aionui_db::DbError> { let rows = self.assistant_snapshots.lock().unwrap(); @@ -357,6 +387,7 @@ impl IConversationRepository for MockRepo { async fn upsert_assistant_snapshot( &self, + _user_id: &str, params: &UpsertConversationAssistantSnapshotParams<'_>, ) -> Result, aionui_db::DbError> { let row = ConversationAssistantSnapshotRow { @@ -386,7 +417,11 @@ impl IConversationRepository for MockRepo { Ok(Some(row)) } - async fn delete_assistant_snapshot(&self, conversation_id: &str) -> Result { + async fn delete_assistant_snapshot( + &self, + _user_id: &str, + conversation_id: &str, + ) -> Result { let mut rows = self.assistant_snapshots.lock().unwrap(); let before = rows.len(); rows.retain(|row| row.conversation_id != conversation_id); @@ -395,6 +430,7 @@ impl IConversationRepository for MockRepo { async fn list_messages_page( &self, + _user_id: &str, conv_id: &str, params: &MessagePageParams, ) -> Result { @@ -474,17 +510,23 @@ impl IConversationRepository for MockRepo { }) } - async fn insert_message(&self, message: &MessageRow) -> Result<(), aionui_db::DbError> { + async fn insert_message(&self, _user_id: &str, message: &MessageRow) -> Result<(), aionui_db::DbError> { let mut messages = self.messages.lock().unwrap(); messages.push(message.clone()); Ok(()) } - async fn update_message(&self, id: &str, updates: &MessageRowUpdate) -> Result<(), aionui_db::DbError> { + async fn update_message( + &self, + _user_id: &str, + conversation_id: &str, + id: &str, + updates: &MessageRowUpdate, + ) -> Result<(), aionui_db::DbError> { let mut messages = self.messages.lock().unwrap(); let message = messages .iter_mut() - .find(|message| message.id == id) + .find(|message| message.conversation_id == conversation_id && message.id == id) .ok_or_else(|| aionui_db::DbError::NotFound(format!("Message {id}")))?; if let Some(content) = &updates.content { @@ -499,7 +541,7 @@ impl IConversationRepository for MockRepo { Ok(()) } - async fn delete_messages_by_conversation(&self, conv_id: &str) -> Result<(), aionui_db::DbError> { + async fn delete_messages_by_conversation(&self, _user_id: &str, conv_id: &str) -> Result<(), aionui_db::DbError> { self.messages .lock() .unwrap() @@ -509,6 +551,7 @@ impl IConversationRepository for MockRepo { async fn get_message_by_msg_id( &self, + _user_id: &str, conv_id: &str, msg_id: &str, msg_type: &str, @@ -524,8 +567,9 @@ impl IConversationRepository for MockRepo { .cloned()) } - async fn list_stale_runtime_messages(&self) -> Result, aionui_db::DbError> { + async fn list_stale_runtime_messages(&self) -> Result, aionui_db::DbError> { let messages = self.messages.lock().unwrap(); + let rows = self.rows.lock().unwrap(); Ok(messages .iter() .filter(|message| { @@ -533,7 +577,16 @@ impl IConversationRepository for MockRepo { && matches!(message.status.as_deref(), Some("work" | "pending")) && matches!(message.r#type.as_str(), "text" | "thinking") }) - .cloned() + .filter_map(|message| { + let user_id = rows + .iter() + .find(|row| row.id == message.conversation_id) + .map(|row| row.user_id.clone())?; + Some(aionui_db::StaleRuntimeMessageRow { + user_id, + message: message.clone(), + }) + }) .collect()) } @@ -551,7 +604,11 @@ impl IConversationRepository for MockRepo { }) } - async fn list_artifacts(&self, conversation_id: &str) -> Result, aionui_db::DbError> { + async fn list_artifacts( + &self, + _user_id: &str, + conversation_id: &str, + ) -> Result, aionui_db::DbError> { Ok(self .artifacts .lock() @@ -564,6 +621,7 @@ impl IConversationRepository for MockRepo { async fn get_artifact( &self, + _user_id: &str, conversation_id: &str, artifact_id: &str, ) -> Result, aionui_db::DbError> { @@ -578,6 +636,7 @@ impl IConversationRepository for MockRepo { async fn upsert_artifact( &self, + _user_id: &str, artifact: &ConversationArtifactRow, ) -> Result { let mut artifacts = self.artifacts.lock().unwrap(); @@ -591,6 +650,7 @@ impl IConversationRepository for MockRepo { async fn update_artifact_status( &self, + _user_id: &str, conversation_id: &str, artifact_id: &str, status: &str, @@ -610,6 +670,7 @@ impl IConversationRepository for MockRepo { async fn mark_skill_suggest_artifacts_saved( &self, + _user_id: &str, cron_job_id: &str, updated_at: TimestampMs, ) -> Result, aionui_db::DbError> { @@ -626,7 +687,11 @@ impl IConversationRepository for MockRepo { Ok(updated) } - async fn delete_artifacts_by_conversation(&self, conversation_id: &str) -> Result<(), aionui_db::DbError> { + async fn delete_artifacts_by_conversation( + &self, + _user_id: &str, + conversation_id: &str, + ) -> Result<(), aionui_db::DbError> { self.artifacts .lock() .unwrap() @@ -636,6 +701,7 @@ impl IConversationRepository for MockRepo { async fn list_legacy_cron_trigger_messages( &self, + _user_id: &str, conversation_id: &str, ) -> Result, aionui_db::DbError> { Ok(self @@ -669,6 +735,7 @@ fn stub_agent_metadata_rows() -> Vec { .into_iter() .map(|(id, backend, agent_type, name, sort_order)| AgentMetadataRow { id: id.to_owned(), + user_id: None, icon: None, name: name.to_owned(), name_i18n: None, @@ -714,9 +781,15 @@ impl IAgentMetadataRepository for StubAgentMetadataRepo { async fn list_all(&self) -> Result, DbError> { Ok(stub_agent_metadata_rows()) } + async fn list_all_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list_all().await + } async fn get(&self, id: &str) -> Result, DbError> { Ok(stub_agent_metadata_rows().into_iter().find(|row| row.id == id)) } + async fn get_for_user(&self, _user_id: &str, id: &str) -> Result, DbError> { + self.get(id).await + } async fn find_by_source_and_name( &self, _agent_source: &str, @@ -724,14 +797,36 @@ impl IAgentMetadataRepository for StubAgentMetadataRepo { ) -> Result, DbError> { Ok(None) } + async fn find_by_source_and_name_for_user( + &self, + _user_id: &str, + agent_source: &str, + name: &str, + ) -> Result, DbError> { + self.find_by_source_and_name(agent_source, name).await + } async fn find_builtin_by_backend(&self, backend: &str) -> Result, DbError> { Ok(stub_agent_metadata_rows() .into_iter() .find(|row| row.agent_source == "builtin" && row.backend.as_deref() == Some(backend))) } + async fn find_builtin_by_backend_for_user( + &self, + _user_id: &str, + backend: &str, + ) -> Result, DbError> { + self.find_builtin_by_backend(backend).await + } async fn upsert(&self, _params: &UpsertAgentMetadataParams<'_>) -> Result { Err(DbError::Init("stub".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAgentMetadataParams<'_>, + ) -> Result { + self.upsert(params).await + } async fn apply_handshake( &self, _id: &str, @@ -739,6 +834,14 @@ impl IAgentMetadataRepository for StubAgentMetadataRepo { ) -> Result, DbError> { Ok(None) } + async fn apply_handshake_for_user( + &self, + _user_id: &str, + id: &str, + params: &UpdateAgentHandshakeParams<'_>, + ) -> Result, DbError> { + self.apply_handshake(id, params).await + } async fn update_availability_snapshot( &self, _id: &str, @@ -746,6 +849,14 @@ impl IAgentMetadataRepository for StubAgentMetadataRepo { ) -> Result, DbError> { Ok(None) } + async fn update_availability_snapshot_for_user( + &self, + _user_id: &str, + id: &str, + params: &UpdateAgentAvailabilitySnapshotParams<'_>, + ) -> Result, DbError> { + self.update_availability_snapshot(id, params).await + } async fn update_agent_overrides( &self, _id: &str, @@ -754,12 +865,27 @@ impl IAgentMetadataRepository for StubAgentMetadataRepo { ) -> Result<(), DbError> { Ok(()) } + async fn update_agent_overrides_for_user( + &self, + _user_id: &str, + id: &str, + command_override: Option<&str>, + env_override: Option<&str>, + ) -> Result<(), DbError> { + self.update_agent_overrides(id, command_override, env_override).await + } async fn set_enabled(&self, _id: &str, _enabled: bool) -> Result { Ok(false) } + async fn set_enabled_for_user(&self, _user_id: &str, id: &str, enabled: bool) -> Result { + self.set_enabled(id, enabled).await + } async fn delete(&self, _id: &str) -> Result { Ok(false) } + async fn delete_for_user(&self, _user_id: &str, id: &str) -> Result { + self.delete(id).await + } } /// Metadata repo used by custom-workspace skill-link tests. It models the @@ -769,6 +895,7 @@ struct ClaudeNativeSkillMetadataRepo; fn claude_metadata_row() -> AgentMetadataRow { AgentMetadataRow { id: "agent-claude".into(), + user_id: None, icon: None, name: "Claude Code".into(), name_i18n: None, @@ -813,9 +940,15 @@ impl IAgentMetadataRepository for ClaudeNativeSkillMetadataRepo { async fn list_all(&self) -> Result, DbError> { Ok(vec![claude_metadata_row()]) } + async fn list_all_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list_all().await + } async fn get(&self, id: &str) -> Result, DbError> { Ok((id == "agent-claude").then(claude_metadata_row)) } + async fn get_for_user(&self, _user_id: &str, id: &str) -> Result, DbError> { + self.get(id).await + } async fn find_by_source_and_name( &self, agent_source: &str, @@ -823,12 +956,34 @@ impl IAgentMetadataRepository for ClaudeNativeSkillMetadataRepo { ) -> Result, DbError> { Ok((agent_source == "builtin" && name == "Claude Code").then(claude_metadata_row)) } + async fn find_by_source_and_name_for_user( + &self, + _user_id: &str, + agent_source: &str, + name: &str, + ) -> Result, DbError> { + self.find_by_source_and_name(agent_source, name).await + } async fn find_builtin_by_backend(&self, backend: &str) -> Result, DbError> { Ok((backend == "claude").then(claude_metadata_row)) } + async fn find_builtin_by_backend_for_user( + &self, + _user_id: &str, + backend: &str, + ) -> Result, DbError> { + self.find_builtin_by_backend(backend).await + } async fn upsert(&self, _params: &UpsertAgentMetadataParams<'_>) -> Result { Err(DbError::Init("stub".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAgentMetadataParams<'_>, + ) -> Result { + self.upsert(params).await + } async fn apply_handshake( &self, _id: &str, @@ -836,6 +991,14 @@ impl IAgentMetadataRepository for ClaudeNativeSkillMetadataRepo { ) -> Result, DbError> { Ok(None) } + async fn apply_handshake_for_user( + &self, + _user_id: &str, + id: &str, + params: &UpdateAgentHandshakeParams<'_>, + ) -> Result, DbError> { + self.apply_handshake(id, params).await + } async fn update_availability_snapshot( &self, _id: &str, @@ -843,6 +1006,14 @@ impl IAgentMetadataRepository for ClaudeNativeSkillMetadataRepo { ) -> Result, DbError> { Ok(None) } + async fn update_availability_snapshot_for_user( + &self, + _user_id: &str, + id: &str, + params: &UpdateAgentAvailabilitySnapshotParams<'_>, + ) -> Result, DbError> { + self.update_availability_snapshot(id, params).await + } async fn update_agent_overrides( &self, _id: &str, @@ -851,12 +1022,27 @@ impl IAgentMetadataRepository for ClaudeNativeSkillMetadataRepo { ) -> Result<(), DbError> { Ok(()) } + async fn update_agent_overrides_for_user( + &self, + _user_id: &str, + id: &str, + command_override: Option<&str>, + env_override: Option<&str>, + ) -> Result<(), DbError> { + self.update_agent_overrides(id, command_override, env_override).await + } async fn set_enabled(&self, _id: &str, _enabled: bool) -> Result { Ok(false) } + async fn set_enabled_for_user(&self, _user_id: &str, id: &str, enabled: bool) -> Result { + self.set_enabled(id, enabled).await + } async fn delete(&self, _id: &str) -> Result { Ok(false) } + async fn delete_for_user(&self, _user_id: &str, id: &str) -> Result { + self.delete(id).await + } } #[derive(Debug, Clone, PartialEq, Eq)] @@ -918,6 +1104,7 @@ impl StubAcpSessionRepo { #[derive(Debug, Clone, PartialEq, Eq)] struct CreateAcpSessionCall { + user_id: String, conversation_id: String, agent_source: String, agent_id: String, @@ -925,11 +1112,12 @@ struct CreateAcpSessionCall { #[async_trait::async_trait] impl IAcpSessionRepository for StubAcpSessionRepo { - async fn get(&self, conversation_id: &str) -> Result, DbError> { + async fn get_for_user(&self, _user_id: &str, conversation_id: &str) -> Result, DbError> { Ok(Some(self.row_for(conversation_id))) } async fn create(&self, params: &CreateAcpSessionParams<'_>) -> Result { self.create_calls.lock().unwrap().push(CreateAcpSessionCall { + user_id: params.user_id.to_owned(), conversation_id: params.conversation_id.to_owned(), agent_source: params.agent_source.to_owned(), agent_id: params.agent_id.to_owned(), @@ -947,18 +1135,27 @@ impl IAcpSessionRepository for StubAcpSessionRepo { suspended_at: None, }) } - async fn update_session_id(&self, _conversation_id: &str, session_id: &str) -> Result { + async fn update_session_id_for_user( + &self, + _user_id: &str, + _conversation_id: &str, + session_id: &str, + ) -> Result { *self.session_id.lock().unwrap() = Some(session_id.to_owned()); Ok(true) } - async fn clear_session_id(&self, _conversation_id: &str) -> Result { + async fn clear_session_id_for_user(&self, _user_id: &str, _conversation_id: &str) -> Result { *self.session_id.lock().unwrap() = None; Ok(true) } - async fn delete(&self, _conversation_id: &str) -> Result { + async fn delete_for_user(&self, _user_id: &str, _conversation_id: &str) -> Result { Ok(false) } - async fn load_runtime_state(&self, _conversation_id: &str) -> Result, DbError> { + async fn load_runtime_state_for_user( + &self, + _user_id: &str, + _conversation_id: &str, + ) -> Result, DbError> { Ok(Some(self.runtime_state.lock().unwrap().clone().unwrap_or( PersistedSessionState { current_model_id: Some("deepseek-v4-pro".to_owned()), @@ -966,8 +1163,9 @@ impl IAcpSessionRepository for StubAcpSessionRepo { }, ))) } - async fn save_runtime_state( + async fn save_runtime_state_for_user( &self, + _user_id: &str, conversation_id: &str, params: &SaveRuntimeStateParams<'_>, ) -> Result { @@ -1105,6 +1303,8 @@ async fn make_service_with_mock_task_manager_and_assistant_support( ) { let (svc, broadcaster, repo) = make_service_with_mock_task_manager(task_mgr); let db = init_database_memory().await.unwrap(); + seed_test_user(db.pool(), "user-1").await; + seed_test_user(db.pool(), "user_1").await; let definition_repo = Arc::new(SqliteAssistantDefinitionRepository::new(db.pool().clone())); let overlay_repo = Arc::new(SqliteAssistantOverlayRepository::new(db.pool().clone())); let preference_repo: Arc = @@ -1130,6 +1330,8 @@ async fn make_service_with_assistant_support( ) { let (svc, broadcaster, repo, _task_mgr) = make_service_with_resolver(skill_resolver); let db = init_database_memory().await.unwrap(); + seed_test_user(db.pool(), "user-1").await; + seed_test_user(db.pool(), "user_1").await; let definition_repo = Arc::new(SqliteAssistantDefinitionRepository::new(db.pool().clone())); let state_repo = Arc::new(SqliteAssistantOverlayRepository::new(db.pool().clone())); let preference_repo: Arc = @@ -1159,6 +1361,8 @@ async fn make_service_with_assistant_support_and_acp_session_repo( let (svc, broadcaster, repo, _task_mgr) = make_service_with_resolver_and_acp_session_repo(skill_resolver, acp_session_repo.clone()); let db = init_database_memory().await.unwrap(); + seed_test_user(db.pool(), "user-1").await; + seed_test_user(db.pool(), "user_1").await; let definition_repo = Arc::new(SqliteAssistantDefinitionRepository::new(db.pool().clone())); let state_repo = Arc::new(SqliteAssistantOverlayRepository::new(db.pool().clone())); let preference_repo: Arc = @@ -1180,6 +1384,22 @@ async fn make_service_with_assistant_support_and_acp_session_repo( ) } +async fn seed_test_user(pool: &sqlx::SqlitePool, user_id: &str) { + let now = aionui_common::now_ms(); + sqlx::query( + "INSERT OR IGNORE INTO users ( + id, user_type, username, password_hash, status, session_generation, created_at, updated_at + ) VALUES (?, 'local', ?, 'hash', 'active', 0, ?, ?)", + ) + .bind(user_id) + .bind(user_id) + .bind(now) + .bind(now) + .execute(pool) + .await + .unwrap(); +} + fn make_create_req() -> CreateConversationRequest { let workspace = ensure_test_workspace_path(); serde_json::from_value(json!({ @@ -1598,7 +1818,7 @@ async fn create_derives_aionrs_type_from_assistant_backend_when_type_is_missing( let resp = svc.create("user_1", req).await.unwrap(); assert_eq!(resp.r#type, AgentType::Aionrs); - assert!(repo.get_assistant_snapshot(&resp.id).await.unwrap().is_some()); + assert!(repo.get_assistant_snapshot("user_1", &resp.id).await.unwrap().is_some()); } #[tokio::test] @@ -1644,7 +1864,7 @@ async fn create_derives_acp_type_from_assistant_backend_when_type_is_missing() { let resp = svc.create("user_1", req).await.unwrap(); assert_eq!(resp.r#type, AgentType::Acp); - assert!(repo.get_assistant_snapshot(&resp.id).await.unwrap().is_some()); + assert!(repo.get_assistant_snapshot("user_1", &resp.id).await.unwrap().is_some()); } #[tokio::test] @@ -1692,7 +1912,7 @@ async fn create_derives_acp_type_from_openclaw_agent_metadata_when_type_is_missi assert_eq!(resp.r#type, AgentType::Acp); assert_eq!(resp.extra["agent_id"], json!("b7e8a9c4")); assert_eq!(resp.extra["backend"], json!("openclaw")); - assert!(repo.get_assistant_snapshot(&resp.id).await.unwrap().is_some()); + assert!(repo.get_assistant_snapshot("user_1", &resp.id).await.unwrap().is_some()); } #[tokio::test] @@ -1741,7 +1961,7 @@ async fn create_derives_acp_type_from_custom_agent_metadata_when_type_is_missing assert_eq!(resp.extra["agent_id"], json!("custom-acp-1")); assert_eq!(resp.extra["agent_source"], json!("custom")); assert_eq!(resp.extra["backend"], json!("acp")); - assert!(repo.get_assistant_snapshot(&resp.id).await.unwrap().is_some()); + assert!(repo.get_assistant_snapshot("user_1", &resp.id).await.unwrap().is_some()); } #[tokio::test] @@ -1862,6 +2082,7 @@ async fn get_reports_idle_runtime_when_only_persisted_status_is_running() { let (svc, _broadcaster, repo, _task_mgr) = make_service(); let created = svc.create("user_1", make_create_req()).await.unwrap(); repo.update( + "user_1", &created.id, &ConversationRowUpdate { status: Some("running".into()), @@ -2105,8 +2326,8 @@ async fn delete_invokes_registered_hook() { struct RecordingHook(Mutex>); #[async_trait::async_trait] impl OnConversationDelete for RecordingHook { - async fn on_conversation_deleted(&self, conversation_id: &str) { - self.0.lock().unwrap().push(conversation_id.to_owned()); + async fn on_conversation_deleted(&self, user_id: &str, conversation_id: &str) { + self.0.lock().unwrap().push(format!("{user_id}:{conversation_id}")); } } @@ -2118,7 +2339,7 @@ async fn delete_invokes_registered_hook() { svc.delete("user_1", &conv.id).await.unwrap(); let calls = hook.0.lock().unwrap(); - assert_eq!(calls.as_slice(), &[conv.id]); + assert_eq!(calls.as_slice(), &[format!("user_1:{}", conv.id)]); } #[tokio::test] @@ -2132,8 +2353,8 @@ async fn delete_invokes_registered_hook_before_row_delete() { #[async_trait::async_trait] impl OnConversationDelete for RowVisibleHook { - async fn on_conversation_deleted(&self, conversation_id: &str) { - let exists = self.repo.get(conversation_id).await.unwrap().is_some(); + async fn on_conversation_deleted(&self, user_id: &str, conversation_id: &str) { + let exists = self.repo.get(user_id, conversation_id).await.unwrap().is_some(); self.observations.lock().unwrap().push(exists); } } @@ -2152,7 +2373,7 @@ async fn delete_invokes_registered_hook_before_row_delete() { let observations = hook.observations.lock().unwrap(); assert_eq!(observations.as_slice(), &[true]); } - assert!(repo.get(&conv.id).await.unwrap().is_none()); + assert!(repo.get("user_1", &conv.id).await.unwrap().is_none()); } #[tokio::test] @@ -2351,22 +2572,25 @@ async fn reset_sets_status_to_pending() { async fn reset_clears_conversation_artifacts() { let (svc, _broadcaster, repo, _task_mgr) = make_service(); let conv = svc.create("user_1", make_create_req()).await.unwrap(); - repo.upsert_artifact(&ConversationArtifactRow { - id: format!("{}:skill_suggest:cron_1", conv.id), - conversation_id: conv.id.clone(), - cron_job_id: Some("cron_1".into()), - kind: "skill_suggest".into(), - status: "pending".into(), - payload: json!({ "cron_job_id": "cron_1", "name": "daily-report" }).to_string(), - created_at: 1000, - updated_at: 1000, - }) + repo.upsert_artifact( + "user_1", + &ConversationArtifactRow { + id: format!("{}:skill_suggest:cron_1", conv.id), + conversation_id: conv.id.clone(), + cron_job_id: Some("cron_1".into()), + kind: "skill_suggest".into(), + status: "pending".into(), + payload: json!({ "cron_job_id": "cron_1", "name": "daily-report" }).to_string(), + created_at: 1000, + updated_at: 1000, + }, + ) .await .unwrap(); svc.reset("user_1", &conv.id).await.unwrap(); - let artifacts = repo.list_artifacts(&conv.id).await.unwrap(); + let artifacts = repo.list_artifacts("user_1", &conv.id).await.unwrap(); assert!(artifacts.is_empty()); } @@ -2374,22 +2598,25 @@ async fn reset_clears_conversation_artifacts() { async fn list_artifacts_includes_legacy_cron_trigger_messages() { let (svc, _broadcaster, repo, _task_mgr) = make_service(); let conv = svc.create("user_1", make_create_req()).await.unwrap(); - repo.insert_message(&MessageRow { - id: "legacy-msg-1".into(), - conversation_id: conv.id.clone(), - msg_id: Some("legacy-trigger-1".into()), - r#type: "cron_trigger".into(), - content: json!({ - "cron_job_id": "cron_1", - "cron_job_name": "Daily Report", - "triggered_at": 1234 - }) - .to_string(), - position: Some("center".into()), - status: Some("finish".into()), - hidden: false, - created_at: 1234, - }) + repo.insert_message( + "user_1", + &MessageRow { + id: "legacy-msg-1".into(), + conversation_id: conv.id.clone(), + msg_id: Some("legacy-trigger-1".into()), + r#type: "cron_trigger".into(), + content: json!({ + "cron_job_id": "cron_1", + "cron_job_name": "Daily Report", + "triggered_at": 1234 + }) + .to_string(), + position: Some("center".into()), + status: Some("finish".into()), + hidden: false, + created_at: 1234, + }, + ) .await .unwrap(); @@ -3080,6 +3307,10 @@ impl IWorkerTaskManager for MockTaskManager { self.agents.lock().unwrap().len() } + fn active_conversation_ids(&self) -> Vec { + self.agents.lock().unwrap().keys().cloned().collect() + } + fn collect_idle(&self, _idle_threshold_ms: TimestampMs) -> Vec { vec![] } @@ -3202,6 +3433,10 @@ impl IWorkerTaskManager for MockTaskManagerWithWorkspace { self.agents.lock().unwrap().len() } + fn active_conversation_ids(&self) -> Vec { + self.agents.lock().unwrap().keys().cloned().collect() + } + fn collect_idle(&self, _idle_threshold_ms: TimestampMs) -> Vec { vec![] } @@ -3535,12 +3770,78 @@ async fn get_config_options_returns_active_agent_snapshot() { }]); task_mgr.insert_agent(&conv.id, AgentInstance::Mock(Arc::new(agent))); - let result = svc.get_config_options(&conv.id).await.unwrap(); + let result = svc.get_config_options("user_1", &conv.id).await.unwrap(); assert_eq!(result.config_options[0].id, "model"); assert_eq!(result.config_options[0].current_value.as_deref(), Some("gpt-5.5")); } +#[tokio::test] +async fn get_config_options_rejects_cross_user_active_task() { + let task_mgr = Arc::new(MockTaskManager::new()); + let (svc, _broadcaster, _repo) = make_service_with_mock_task_manager(task_mgr.clone()); + let conv = svc.create("user_1", make_create_req()).await.unwrap(); + let agent = MockAgent::new(&conv.id); + task_mgr.insert_agent(&conv.id, AgentInstance::Mock(Arc::new(agent))); + + let err = svc.get_config_options("user_2", &conv.id).await.unwrap_err(); + + assert!(matches!(err, ConversationError::NotFound { id } if id == conv.id)); +} + +#[tokio::test] +async fn active_count_for_user_counts_only_owned_active_tasks() { + let task_mgr = Arc::new(MockTaskManager::new()); + let (svc, _broadcaster, _repo) = make_service_with_mock_task_manager(task_mgr.clone()); + let user_1_conv = svc.create("user_1", make_create_req()).await.unwrap(); + let user_2_conv = svc.create("user_2", make_create_req()).await.unwrap(); + task_mgr.insert_agent( + &user_1_conv.id, + AgentInstance::Mock(Arc::new(MockAgent::new(&user_1_conv.id))), + ); + task_mgr.insert_agent( + &user_2_conv.id, + AgentInstance::Mock(Arc::new(MockAgent::new(&user_2_conv.id))), + ); + + assert_eq!(svc.active_count_for_user("user_1").await.unwrap(), 1); + assert_eq!(svc.active_count_for_user("user_2").await.unwrap(), 1); + assert_eq!(svc.active_count_for_user("user_3").await.unwrap(), 0); +} + +#[tokio::test] +async fn terminate_runtime_for_user_kills_only_owned_active_tasks() { + let task_mgr = Arc::new(MockTaskManager::new()); + let (svc, _broadcaster, _repo) = make_service_with_mock_task_manager(task_mgr.clone()); + let user_1_conv = svc.create("user_1", make_create_req()).await.unwrap(); + let user_2_conv = svc.create("user_2", make_create_req()).await.unwrap(); + task_mgr.insert_agent( + &user_1_conv.id, + AgentInstance::Mock(Arc::new(MockAgent::new(&user_1_conv.id))), + ); + task_mgr.insert_agent( + &user_2_conv.id, + AgentInstance::Mock(Arc::new(MockAgent::new(&user_2_conv.id))), + ); + let runtime_state = svc.runtime_state(); + let _claim = runtime_state + .try_claim_turn(&user_1_conv.id, "turn-1") + .expect("claim should be created"); + runtime_state.mark_cancelling(&user_1_conv.id); + + let terminated = svc.terminate_runtime_for_user("user_1").await.unwrap(); + + assert_eq!(terminated, 1); + assert_eq!( + task_mgr.kill_records(), + vec![(user_1_conv.id.clone(), Some(AgentKillReason::SessionRevoked))] + ); + assert!(!runtime_state.is_claimed(&user_1_conv.id)); + assert!(!runtime_state.is_cancelling(&user_1_conv.id)); + assert_eq!(task_mgr.active_count(), 1); + assert!(task_mgr.get_task(&user_2_conv.id).is_some()); +} + #[tokio::test] async fn ensure_runtime_recovers_missing_agent_and_returns_snapshot() { let task_mgr = Arc::new(MockTaskManager::new()); @@ -3638,6 +3939,7 @@ async fn set_config_option_returns_observed_confirmation() { let result = svc .set_config_option( + "user_1", &conv.id, "reasoning_effort", SetConfigOptionRequest { @@ -3669,11 +3971,15 @@ async fn save_acp_runtime_mode_updates_runtime_mode_config_selection() { ), context_usage_json: None, })); - let (svc, _, _, _) = make_service_with_resolver_and_acp_session_repo( + let (svc, _, repo, _) = make_service_with_resolver_and_acp_session_repo( Arc::new(FixedSkillResolver { names: vec![] }), acp_repo.clone(), ); - svc.save_acp_runtime_mode("conv_1", "full-access").await.unwrap(); + let conv = insert_conversation_with_type(&repo, "user_1", AgentType::Acp).await; + + svc.save_acp_runtime_mode("user_1", &conv.id, "full-access") + .await + .unwrap(); let saves = acp_repo.runtime_state_saves(); assert_eq!(saves.len(), 1); @@ -3685,6 +3991,27 @@ async fn save_acp_runtime_mode_updates_runtime_mode_config_selection() { assert_eq!(selections["reasoning_effort"], "low"); } +#[tokio::test] +async fn save_acp_runtime_mode_rejects_wrong_user() { + let acp_repo = Arc::new(StubAcpSessionRepo::with_runtime_state(PersistedSessionState { + current_mode_id: Some("read-only".to_owned()), + ..Default::default() + })); + let (svc, _, repo, _) = make_service_with_resolver_and_acp_session_repo( + Arc::new(FixedSkillResolver { names: vec![] }), + acp_repo.clone(), + ); + let conv = insert_conversation_with_type(&repo, "user_1", AgentType::Acp).await; + + let err = svc + .save_acp_runtime_mode("user_2", &conv.id, "full-access") + .await + .unwrap_err(); + + assert!(matches!(err, ConversationError::NotFound { .. })); + assert!(acp_repo.runtime_state_saves().is_empty()); +} + #[tokio::test] async fn run_agent_turn_applies_required_runtime_mode_after_stream_subscription() { let task_mgr = Arc::new(MockTaskManager::new()); @@ -3749,6 +4076,7 @@ async fn set_config_option_evicts_task_when_acp_protocol_is_not_connected() { let err = svc .set_config_option( + "user_1", &conv.id, "model", SetConfigOptionRequest { @@ -3783,6 +4111,7 @@ async fn command_ack_does_not_persist_assistant_preference_in_core_service() { let result = svc .set_config_option( + "user_1", &conv.id, "model", SetConfigOptionRequest { @@ -3823,15 +4152,18 @@ async fn set_config_option_persists_runtime_model_into_assistant_preference_when .await .unwrap(); preference_repo - .upsert(&UpsertAssistantPreferenceParams { - assistant_definition_id: "asstdef_acp_auto", - last_model_id: Some("legacy-acp-model"), - last_permission_value: Some("legacy-mode"), - last_thought_level_value: Some("legacy-low"), - last_skill_ids: "[]", - last_disabled_builtin_skill_ids: "[]", - last_mcp_ids: "[]", - }) + .upsert_for_user( + "user_1", + &UpsertAssistantPreferenceParams { + assistant_definition_id: "asstdef_acp_auto", + last_model_id: Some("legacy-acp-model"), + last_permission_value: Some("legacy-mode"), + last_thought_level_value: Some("legacy-low"), + last_skill_ids: "[]", + last_disabled_builtin_skill_ids: "[]", + last_mcp_ids: "[]", + }, + ) .await .unwrap(); @@ -3842,6 +4174,7 @@ async fn set_config_option_persists_runtime_model_into_assistant_preference_when let result = svc .set_config_option( + "user_1", &conv.id, "model", SetConfigOptionRequest { @@ -3852,13 +4185,18 @@ async fn set_config_option_persists_runtime_model_into_assistant_preference_when .unwrap(); assert_eq!(result.confirmation, ConfigOptionConfirmation::Observed); - let pref_after_model = preference_repo.get("asstdef_acp_auto").await.unwrap().unwrap(); + let pref_after_model = preference_repo + .get_for_user("user_1", "asstdef_acp_auto") + .await + .unwrap() + .unwrap(); assert_eq!(pref_after_model.last_model_id.as_deref(), Some("gpt-5.5")); assert_eq!(pref_after_model.last_permission_value.as_deref(), Some("legacy-mode")); - let snapshot_after_model = repo.get_assistant_snapshot(&conv.id).await.unwrap().unwrap(); + let snapshot_after_model = repo.get_assistant_snapshot("user_1", &conv.id).await.unwrap().unwrap(); assert_eq!(snapshot_after_model.resolved_model_id.as_deref(), Some("gpt-5.5")); svc.set_config_option( + "user_1", &conv.id, "mode", SetConfigOptionRequest { @@ -3867,15 +4205,20 @@ async fn set_config_option_persists_runtime_model_into_assistant_preference_when ) .await .unwrap(); - let pref_after_mode = preference_repo.get("asstdef_acp_auto").await.unwrap().unwrap(); + let pref_after_mode = preference_repo + .get_for_user("user_1", "asstdef_acp_auto") + .await + .unwrap() + .unwrap(); assert_eq!(pref_after_mode.last_model_id.as_deref(), Some("gpt-5.5")); assert_eq!(pref_after_mode.last_permission_value.as_deref(), Some("plan")); - let snapshot_after_mode = repo.get_assistant_snapshot(&conv.id).await.unwrap().unwrap(); + let snapshot_after_mode = repo.get_assistant_snapshot("user_1", &conv.id).await.unwrap().unwrap(); assert_eq!(snapshot_after_mode.resolved_permission_value.as_deref(), Some("plan")); // Thought-level config options follow the same auto-default writeback // contract as model and permission defaults. svc.set_config_option( + "user_1", &conv.id, "reasoning_effort", SetConfigOptionRequest { @@ -3884,11 +4227,15 @@ async fn set_config_option_persists_runtime_model_into_assistant_preference_when ) .await .unwrap(); - let pref_after_thought = preference_repo.get("asstdef_acp_auto").await.unwrap().unwrap(); + let pref_after_thought = preference_repo + .get_for_user("user_1", "asstdef_acp_auto") + .await + .unwrap() + .unwrap(); assert_eq!(pref_after_thought.last_model_id.as_deref(), Some("gpt-5.5")); assert_eq!(pref_after_thought.last_permission_value.as_deref(), Some("plan")); assert_eq!(pref_after_thought.last_thought_level_value.as_deref(), Some("high")); - let snapshot_after_thought = repo.get_assistant_snapshot(&conv.id).await.unwrap().unwrap(); + let snapshot_after_thought = repo.get_assistant_snapshot("user_1", &conv.id).await.unwrap().unwrap(); assert_eq!( snapshot_after_thought.resolved_thought_level_value.as_deref(), Some("high") @@ -3921,15 +4268,18 @@ async fn set_config_option_does_not_persist_preference_on_error() { .await .unwrap(); preference_repo - .upsert(&UpsertAssistantPreferenceParams { - assistant_definition_id: "asstdef_acp_auto", - last_model_id: Some("original-model"), - last_permission_value: Some("original-mode"), - last_thought_level_value: Some("original-low"), - last_skill_ids: "[]", - last_disabled_builtin_skill_ids: "[]", - last_mcp_ids: "[]", - }) + .upsert_for_user( + "user_1", + &UpsertAssistantPreferenceParams { + assistant_definition_id: "asstdef_acp_auto", + last_model_id: Some("original-model"), + last_permission_value: Some("original-mode"), + last_thought_level_value: Some("original-low"), + last_skill_ids: "[]", + last_disabled_builtin_skill_ids: "[]", + last_mcp_ids: "[]", + }, + ) .await .unwrap(); @@ -3947,6 +4297,7 @@ async fn set_config_option_does_not_persist_preference_on_error() { let result = svc .set_config_option( + "user_1", &conv.id, "model", SetConfigOptionRequest { @@ -3956,7 +4307,11 @@ async fn set_config_option_does_not_persist_preference_on_error() { .await; assert!(result.is_err(), "conflict from agent must surface as error"); - let pref_after = preference_repo.get("asstdef_acp_auto").await.unwrap().unwrap(); + let pref_after = preference_repo + .get_for_user("user_1", "asstdef_acp_auto") + .await + .unwrap() + .unwrap(); assert_eq!( pref_after.last_model_id.as_deref(), Some("original-model"), @@ -3991,15 +4346,18 @@ async fn set_config_option_skips_preference_write_back_when_default_mode_is_fixe .await .unwrap(); preference_repo - .upsert(&UpsertAssistantPreferenceParams { - assistant_definition_id: "asstdef_acp_fixed", - last_model_id: Some("legacy-fixed-model"), - last_permission_value: Some("legacy-fixed-mode"), - last_thought_level_value: None, - last_skill_ids: "[]", - last_disabled_builtin_skill_ids: "[]", - last_mcp_ids: "[]", - }) + .upsert_for_user( + "user_1", + &UpsertAssistantPreferenceParams { + assistant_definition_id: "asstdef_acp_fixed", + last_model_id: Some("legacy-fixed-model"), + last_permission_value: Some("legacy-fixed-mode"), + last_thought_level_value: None, + last_skill_ids: "[]", + last_disabled_builtin_skill_ids: "[]", + last_mcp_ids: "[]", + }, + ) .await .unwrap(); @@ -4008,6 +4366,7 @@ async fn set_config_option_skips_preference_write_back_when_default_mode_is_fixe task_mgr.insert_agent(&conv.id, AgentInstance::Mock(agent)); svc.set_config_option( + "user_1", &conv.id, "model", SetConfigOptionRequest { @@ -4017,6 +4376,7 @@ async fn set_config_option_skips_preference_write_back_when_default_mode_is_fixe .await .unwrap(); svc.set_config_option( + "user_1", &conv.id, "mode", SetConfigOptionRequest { @@ -4026,6 +4386,7 @@ async fn set_config_option_skips_preference_write_back_when_default_mode_is_fixe .await .unwrap(); svc.set_config_option( + "user_1", &conv.id, "reasoning_effort", SetConfigOptionRequest { @@ -4035,13 +4396,17 @@ async fn set_config_option_skips_preference_write_back_when_default_mode_is_fixe .await .unwrap(); - let pref = preference_repo.get("asstdef_acp_fixed").await.unwrap().unwrap(); + let pref = preference_repo + .get_for_user("user_1", "asstdef_acp_fixed") + .await + .unwrap() + .unwrap(); assert_eq!(pref.last_model_id.as_deref(), Some("legacy-fixed-model")); assert_eq!(pref.last_permission_value.as_deref(), Some("legacy-fixed-mode")); assert_eq!(pref.last_thought_level_value.as_deref(), None); // The snapshot still tracks the runtime override so the active session reflects it, // even though the persisted assistant preference must not change for fixed defaults. - let snapshot = repo.get_assistant_snapshot(&conv.id).await.unwrap().unwrap(); + let snapshot = repo.get_assistant_snapshot("user_1", &conv.id).await.unwrap().unwrap(); assert_eq!(snapshot.resolved_model_id.as_deref(), Some("transient-model")); assert_eq!(snapshot.resolved_permission_value.as_deref(), Some("transient-mode")); assert_eq!(snapshot.resolved_thought_level_value.as_deref(), Some("transient-high")); @@ -4073,15 +4438,18 @@ async fn set_config_option_command_ack_does_not_persist_assistant_preference() { .await .unwrap(); preference_repo - .upsert(&UpsertAssistantPreferenceParams { - assistant_definition_id: "asstdef_acp_ack", - last_model_id: Some("legacy-ack-model"), - last_permission_value: Some("legacy-ack-mode"), - last_thought_level_value: Some("legacy-ack-thought"), - last_skill_ids: "[]", - last_disabled_builtin_skill_ids: "[]", - last_mcp_ids: "[]", - }) + .upsert_for_user( + "user_1", + &UpsertAssistantPreferenceParams { + assistant_definition_id: "asstdef_acp_ack", + last_model_id: Some("legacy-ack-model"), + last_permission_value: Some("legacy-ack-mode"), + last_thought_level_value: Some("legacy-ack-thought"), + last_skill_ids: "[]", + last_disabled_builtin_skill_ids: "[]", + last_mcp_ids: "[]", + }, + ) .await .unwrap(); @@ -4095,6 +4463,7 @@ async fn set_config_option_command_ack_does_not_persist_assistant_preference() { task_mgr.insert_agent(&conv.id, AgentInstance::Mock(agent)); svc.set_config_option( + "user_1", &conv.id, "model", SetConfigOptionRequest { @@ -4104,10 +4473,14 @@ async fn set_config_option_command_ack_does_not_persist_assistant_preference() { .await .unwrap(); - let pref = preference_repo.get("asstdef_acp_ack").await.unwrap().unwrap(); + let pref = preference_repo + .get_for_user("user_1", "asstdef_acp_ack") + .await + .unwrap() + .unwrap(); assert_eq!(pref.last_model_id.as_deref(), Some("legacy-ack-model")); assert_eq!(pref.last_permission_value.as_deref(), Some("legacy-ack-mode")); - let snapshot = repo.get_assistant_snapshot(&conv.id).await.unwrap().unwrap(); + let snapshot = repo.get_assistant_snapshot("user_1", &conv.id).await.unwrap().unwrap(); assert_eq!(snapshot.resolved_model_id.as_deref(), Some("legacy-ack-model")); } @@ -4137,15 +4510,18 @@ async fn update_aionrs_model_updates_assistant_preference_only_when_snapshot_mod .await .unwrap(); preference_repo - .upsert(&UpsertAssistantPreferenceParams { - assistant_definition_id: "asstdef_aionrs_auto", - last_model_id: Some("legacy-aionrs-model"), - last_permission_value: None, - last_thought_level_value: None, - last_skill_ids: "[]", - last_disabled_builtin_skill_ids: "[]", - last_mcp_ids: "[]", - }) + .upsert_for_user( + "user_1", + &UpsertAssistantPreferenceParams { + assistant_definition_id: "asstdef_aionrs_auto", + last_model_id: Some("legacy-aionrs-model"), + last_permission_value: None, + last_thought_level_value: None, + last_skill_ids: "[]", + last_disabled_builtin_skill_ids: "[]", + last_mcp_ids: "[]", + }, + ) .await .unwrap(); @@ -4174,9 +4550,17 @@ async fn update_aionrs_model_updates_assistant_preference_only_when_snapshot_mod updated.model.as_ref().and_then(|model| model.use_model.as_deref()), Some("model-z") ); - let auto_pref = preference_repo.get("asstdef_aionrs_auto").await.unwrap().unwrap(); + let auto_pref = preference_repo + .get_for_user("user_1", "asstdef_aionrs_auto") + .await + .unwrap() + .unwrap(); assert_eq!(auto_pref.last_model_id.as_deref(), Some("model-z")); - let auto_snapshot = repo.get_assistant_snapshot(&auto_conv.id).await.unwrap().unwrap(); + let auto_snapshot = repo + .get_assistant_snapshot("user_1", &auto_conv.id) + .await + .unwrap() + .unwrap(); assert_eq!(auto_snapshot.resolved_model_id.as_deref(), Some("model-z")); upsert_test_assistant_definition( @@ -4199,15 +4583,18 @@ async fn update_aionrs_model_updates_assistant_preference_only_when_snapshot_mod .await .unwrap(); preference_repo - .upsert(&UpsertAssistantPreferenceParams { - assistant_definition_id: "asstdef_aionrs_fixed", - last_model_id: Some("legacy-aionrs-fixed-model"), - last_permission_value: None, - last_thought_level_value: None, - last_skill_ids: "[]", - last_disabled_builtin_skill_ids: "[]", - last_mcp_ids: "[]", - }) + .upsert_for_user( + "user_1", + &UpsertAssistantPreferenceParams { + assistant_definition_id: "asstdef_aionrs_fixed", + last_model_id: Some("legacy-aionrs-fixed-model"), + last_permission_value: None, + last_thought_level_value: None, + last_skill_ids: "[]", + last_disabled_builtin_skill_ids: "[]", + last_mcp_ids: "[]", + }, + ) .await .unwrap(); @@ -4232,13 +4619,21 @@ async fn update_aionrs_model_updates_assistant_preference_only_when_snapshot_mod .await .unwrap(); - let fixed_pref = preference_repo.get("asstdef_aionrs_fixed").await.unwrap().unwrap(); + let fixed_pref = preference_repo + .get_for_user("user_1", "asstdef_aionrs_fixed") + .await + .unwrap() + .unwrap(); assert_eq!(fixed_pref.last_model_id.as_deref(), Some("legacy-aionrs-fixed-model")); - let fixed_snapshot = repo.get_assistant_snapshot(&fixed_conv.id).await.unwrap().unwrap(); + let fixed_snapshot = repo + .get_assistant_snapshot("user_1", &fixed_conv.id) + .await + .unwrap() + .unwrap(); assert_eq!(fixed_snapshot.resolved_model_id.as_deref(), Some("model-y")); // Ensure the update path used the repository row, not a no-op. - let updated_row = repo.get(&fixed_conv.id).await.unwrap().unwrap(); + let updated_row = repo.get("user_1", &fixed_conv.id).await.unwrap().unwrap(); assert!( updated_row .model @@ -4256,6 +4651,7 @@ async fn send_message_missing_workspace_persists_message_and_failure_tip() { broadcaster.take_events(); let legacy_workspace = format!("/tmp/does-not-exist-{}", aionui_common::generate_short_id()); repo.update( + "user_1", &conv.id, &ConversationRowUpdate { extra: Some(json!({ "workspace": legacy_workspace }).to_string()), @@ -4278,6 +4674,7 @@ async fn send_message_missing_workspace_persists_message_and_failure_tip() { loop { let messages = repo .list_messages_page( + "user_1", &conv.id, &MessagePageParams { limit: 20, @@ -4383,7 +4780,7 @@ async fn send_message_returns_before_cold_agent_build_completes() { "cold agent build should continue in the background after send_message returns" ); - let updated = repo.get(&conv.id).await.unwrap().unwrap(); + let updated = repo.get("user_1", &conv.id).await.unwrap().unwrap(); assert_ne!(updated.status.as_deref(), Some("running")); assert!( svc.runtime_state().is_claimed(&conv.id), @@ -4407,6 +4804,7 @@ async fn send_message_persists_hidden_user_message_when_requested() { let messages = repo .list_messages_page( + "user_1", &conv.id, &MessagePageParams { limit: 20, @@ -4447,6 +4845,7 @@ async fn send_message_persists_error_tip_when_agent_build_fails() { loop { let messages = repo .list_messages_page( + "user_1", &conv.id, &MessagePageParams { limit: 20, @@ -4487,7 +4886,7 @@ async fn send_message_persists_error_tip_when_agent_build_fails() { "The upstream Agent failed while handling the request" ); - let updated = repo.get(&conv.id).await.unwrap().unwrap(); + let updated = repo.get("user_1", &conv.id).await.unwrap().unwrap(); assert_eq!(updated.status.as_deref(), Some("finished")); assert!( !svc.runtime_state().is_claimed(&conv.id), @@ -4547,32 +4946,35 @@ async fn run_agent_turn_returns_error_message_when_agent_build_fails() { async fn latest_conversation_error_message_prefers_error_detail() { let (svc, _broadcaster, repo, _task_mgr) = make_service(); let conv = svc.create("user_1", make_create_req()).await.unwrap(); - repo.insert_message(&MessageRow { - id: "msg_error".into(), - conversation_id: conv.id.clone(), - msg_id: None, - r#type: "tips".into(), - content: serde_json::json!({ - "content": "Current agent failed", - "type": "error", - "details": { - "detail": "Cannot resume Claude session with Codex assistant" - }, - "error": { - "message": "Current agent failed", - "detail": "Codex cannot resume a conversation created by Claude" - } - }) - .to_string(), - position: Some("center".into()), - status: Some("error".into()), - hidden: false, - created_at: 10, - }) + repo.insert_message( + "user_1", + &MessageRow { + id: "msg_error".into(), + conversation_id: conv.id.clone(), + msg_id: None, + r#type: "tips".into(), + content: serde_json::json!({ + "content": "Current agent failed", + "type": "error", + "details": { + "detail": "Cannot resume Claude session with Codex assistant" + }, + "error": { + "message": "Current agent failed", + "detail": "Codex cannot resume a conversation created by Claude" + } + }) + .to_string(), + position: Some("center".into()), + status: Some("error".into()), + hidden: false, + created_at: 10, + }, + ) .await .unwrap(); - let message = svc.latest_conversation_error_message(&conv.id).await.unwrap(); + let message = svc.latest_conversation_error_message("user_1", &conv.id).await.unwrap(); assert_eq!( message.as_deref(), @@ -4616,6 +5018,7 @@ async fn send_message_persists_openclaw_gateway_unreachable_tip_when_turn_build_ let messages = repo .list_messages_page( + "user_1", &conv.id, &MessagePageParams { limit: 20, @@ -4725,7 +5128,7 @@ async fn send_message_allows_stale_db_running_without_runtime_claim() { status: Some("running".into()), ..Default::default() }; - repo.update(&conv.id, &update).await.unwrap(); + repo.update("user_1", &conv.id, &update).await.unwrap(); let result = svc.send_message("user_1", &conv.id, make_send_req(), &task_mgr).await; assert!(result.is_ok(), "stale DB running must not block sending"); @@ -4765,6 +5168,7 @@ async fn send_message_rejects_when_runtime_is_shutting_down() { let messages = repo .list_messages_page( + "user_1", &conv.id, &MessagePageParams { limit: 20, @@ -4802,6 +5206,7 @@ async fn send_message_build_failure_while_deleting_skips_failure_tip_and_complet let messages = repo .list_messages_page( + "user_1", &conv.id, &MessagePageParams { limit: 20, @@ -4846,7 +5251,7 @@ async fn send_message_persists_factory_resolved_workspace() { extra: Some(r#"{"workspace":""}"#.to_owned()), ..Default::default() }; - repo.update(&conv.id, &empty_ws_update).await.unwrap(); + repo.update("user_1", &conv.id, &empty_ws_update).await.unwrap(); svc.send_message("user_1", &conv.id, make_send_req(), &task_mgr) .await @@ -4872,30 +5277,36 @@ async fn startup_recovery_closes_stale_runtime_messages_without_failure_tip() { let (svc, _broadcaster, repo, _task_mgr) = make_service(); let conv = svc.create("user_1", make_create_req()).await.unwrap(); - repo.insert_message(&MessageRow { - id: "visible-stale".into(), - conversation_id: conv.id.clone(), - msg_id: Some("visible-stale".into()), - r#type: "text".into(), - content: json!({ "content": "partial output" }).to_string(), - position: Some("left".into()), - status: Some("work".into()), - hidden: false, - created_at: 1, - }) + repo.insert_message( + "user_1", + &MessageRow { + id: "visible-stale".into(), + conversation_id: conv.id.clone(), + msg_id: Some("visible-stale".into()), + r#type: "text".into(), + content: json!({ "content": "partial output" }).to_string(), + position: Some("left".into()), + status: Some("work".into()), + hidden: false, + created_at: 1, + }, + ) .await .unwrap(); - repo.insert_message(&MessageRow { - id: "empty-stale".into(), - conversation_id: conv.id.clone(), - msg_id: Some("empty-stale".into()), - r#type: "thinking".into(), - content: json!({ "content": "" }).to_string(), - position: Some("left".into()), - status: Some("pending".into()), - hidden: false, - created_at: 2, - }) + repo.insert_message( + "user_1", + &MessageRow { + id: "empty-stale".into(), + conversation_id: conv.id.clone(), + msg_id: Some("empty-stale".into()), + r#type: "thinking".into(), + content: json!({ "content": "" }).to_string(), + position: Some("left".into()), + status: Some("pending".into()), + hidden: false, + created_at: 2, + }, + ) .await .unwrap(); @@ -4903,6 +5314,7 @@ async fn startup_recovery_closes_stale_runtime_messages_without_failure_tip() { let messages = repo .list_messages_page( + "user_1", &conv.id, &MessagePageParams { limit: 20, @@ -5004,6 +5416,7 @@ async fn send_message_does_not_inject_send_error_when_runtime_terminal_exists() let messages = repo .list_messages_page( + "user_1", &conv.id, &MessagePageParams { limit: 20, @@ -5042,6 +5455,7 @@ async fn send_message_injects_send_error_when_runtime_terminal_missing() { let messages = repo .list_messages_page( + "user_1", &conv.id, &MessagePageParams { limit: 20, @@ -5097,11 +5511,13 @@ async fn send_message_records_agent_availability_feedback_on_send_failure() { failures, vec![ RecordedAvailabilityFailure { + user_id: "user_1".into(), agent_id: "agent-feedback-1".into(), code: "session_send_failed".into(), message: "provider returned 401 invalid api key".into(), }, RecordedAvailabilityFailure { + user_id: "user_1".into(), agent_id: "agent-feedback-1".into(), code: "auth_required".into(), message: "provider returned 401 invalid api key".into(), @@ -5142,7 +5558,10 @@ async fn send_message_records_agent_availability_feedback_on_send_success() { wait_for_turn_released(&svc, &conv.id).await; let successes = feedback.successes.lock().unwrap().clone(); - assert_eq!(successes, vec!["agent-feedback-success".to_owned()]); + assert_eq!( + successes, + vec![("user_1".to_owned(), "agent-feedback-success".to_owned())] + ); assert!(feedback.failures.lock().unwrap().is_empty()); } @@ -5811,6 +6230,7 @@ async fn warmup_rejects_legacy_workspace_with_runtime_error_code() { let conv = svc.create("user_1", make_create_req()).await.unwrap(); let legacy_workspace = format!("/tmp/does-not-exist-{}", aionui_common::generate_short_id()); repo.update( + "user_1", &conv.id, &ConversationRowUpdate { extra: Some(json!({ "workspace": legacy_workspace }).to_string()), @@ -6264,25 +6684,31 @@ async fn create_resolves_assistant_snapshot_and_updates_preferences() { .await .unwrap(); state_repo - .upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: "asstdef_preset_1", - enabled: true, - sort_order: 0, - agent_id_override: Some("codex"), - last_used_at: None, - }) + .upsert_for_user( + "user-1", + &UpsertAssistantOverlayParams { + assistant_definition_id: "asstdef_preset_1", + enabled: true, + sort_order: 0, + agent_id_override: Some("codex"), + last_used_at: None, + }, + ) .await .unwrap(); preference_repo - .upsert(&UpsertAssistantPreferenceParams { - assistant_definition_id: "asstdef_preset_1", - last_model_id: Some("old-model"), - last_permission_value: Some("workspace-write"), - last_thought_level_value: Some("high"), - last_skill_ids: r#"["legacy-skill"]"#, - last_disabled_builtin_skill_ids: r#"["legacy-disabled"]"#, - last_mcp_ids: r#"["legacy-mcp"]"#, - }) + .upsert_for_user( + "user-1", + &UpsertAssistantPreferenceParams { + assistant_definition_id: "asstdef_preset_1", + last_model_id: Some("old-model"), + last_permission_value: Some("workspace-write"), + last_thought_level_value: Some("high"), + last_skill_ids: r#"["legacy-skill"]"#, + last_disabled_builtin_skill_ids: r#"["legacy-disabled"]"#, + last_mcp_ids: r#"["legacy-mcp"]"#, + }, + ) .await .unwrap(); @@ -6328,7 +6754,7 @@ async fn create_resolves_assistant_snapshot_and_updates_preferences() { assert_eq!(resp.extra["skills"], json!(["cron", "pdf"])); assert!(resp.extra.get("assistant_snapshot").is_none()); - let snapshot = repo.get_assistant_snapshot(&resp.id).await.unwrap().unwrap(); + let snapshot = repo.get_assistant_snapshot("user_1", &resp.id).await.unwrap().unwrap(); assert_eq!(snapshot.assistant_definition_id, "asstdef_preset_1"); assert_eq!(snapshot.assistant_id, "preset-1"); assert_eq!(snapshot.agent_id, "8e1acf31"); @@ -6338,7 +6764,11 @@ async fn create_resolves_assistant_snapshot_and_updates_preferences() { assert_eq!(snapshot.default_skills_mode, "auto"); assert_eq!(snapshot.resolved_skill_ids, r#"["pdf"]"#); - let updated_pref = preference_repo.get("asstdef_preset_1").await.unwrap().unwrap(); + let updated_pref = preference_repo + .get_for_user("user-1", "asstdef_preset_1") + .await + .unwrap() + .unwrap(); assert_eq!(updated_pref.last_model_id.as_deref(), Some("new-model")); assert_eq!(updated_pref.last_skill_ids, r#"["pdf"]"#); assert_eq!(updated_pref.last_disabled_builtin_skill_ids, r#"["todo-tracker"]"#); @@ -6391,46 +6821,52 @@ async fn existing_conversation_reads_current_assistant_identity() { let workspace = ensure_test_workspace_path(); definition_repo - .upsert(&UpsertAssistantDefinitionParams { - id: "asstdef_live_identity", - assistant_id: "live-identity", - source: "user", - owner_type: "user", - source_ref: Some("live-identity"), - name: "Old Name", - name_i18n: "{}", - description: None, - description_i18n: "{}", - avatar_type: "emoji", - avatar_value: Some("🤖"), - agent_id: "claude", - rule_resource_type: "none", - rule_resource_ref: None, - recommended_prompts: "[]", - recommended_prompts_i18n: "{}", - default_model_mode: "auto", - default_model_value: None, - default_permission_mode: "auto", - default_permission_value: None, - default_thought_level_mode: "auto", - default_thought_level_value: None, - default_skills_mode: "auto", - default_skill_ids: "[]", - custom_skill_names: "[]", - default_disabled_builtin_skill_ids: "[]", - default_mcps_mode: "auto", - default_mcp_ids: "[]", - }) + .upsert_for_user( + "user-1", + &UpsertAssistantDefinitionParams { + id: "asstdef_live_identity", + assistant_id: "live-identity", + source: "user", + owner_type: "user", + source_ref: Some("live-identity"), + name: "Old Name", + name_i18n: "{}", + description: None, + description_i18n: "{}", + avatar_type: "emoji", + avatar_value: Some("🤖"), + agent_id: "claude", + rule_resource_type: "none", + rule_resource_ref: None, + recommended_prompts: "[]", + recommended_prompts_i18n: "{}", + default_model_mode: "auto", + default_model_value: None, + default_permission_mode: "auto", + default_permission_value: None, + default_thought_level_mode: "auto", + default_thought_level_value: None, + default_skills_mode: "auto", + default_skill_ids: "[]", + custom_skill_names: "[]", + default_disabled_builtin_skill_ids: "[]", + default_mcps_mode: "auto", + default_mcp_ids: "[]", + }, + ) .await .unwrap(); state_repo - .upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: "asstdef_live_identity", - enabled: true, - sort_order: 0, - agent_id_override: None, - last_used_at: None, - }) + .upsert_for_user( + "user-1", + &UpsertAssistantOverlayParams { + assistant_definition_id: "asstdef_live_identity", + enabled: true, + sort_order: 0, + agent_id_override: None, + last_used_at: None, + }, + ) .await .unwrap(); @@ -6447,36 +6883,39 @@ async fn existing_conversation_reads_current_assistant_identity() { let created = svc.create("user-1", req).await.unwrap(); definition_repo - .upsert(&UpsertAssistantDefinitionParams { - id: "asstdef_live_identity", - assistant_id: "live-identity", - source: "user", - owner_type: "user", - source_ref: Some("live-identity"), - name: "New Name", - name_i18n: "{}", - description: None, - description_i18n: "{}", - avatar_type: "emoji", - avatar_value: Some("🧪"), - agent_id: "claude", - rule_resource_type: "none", - rule_resource_ref: None, - recommended_prompts: "[]", - recommended_prompts_i18n: "{}", - default_model_mode: "auto", - default_model_value: None, - default_permission_mode: "auto", - default_permission_value: None, - default_thought_level_mode: "auto", - default_thought_level_value: None, - default_skills_mode: "auto", - default_skill_ids: "[]", - custom_skill_names: "[]", - default_disabled_builtin_skill_ids: "[]", - default_mcps_mode: "auto", - default_mcp_ids: "[]", - }) + .upsert_for_user( + "user-1", + &UpsertAssistantDefinitionParams { + id: "asstdef_live_identity", + assistant_id: "live-identity", + source: "user", + owner_type: "user", + source_ref: Some("live-identity"), + name: "New Name", + name_i18n: "{}", + description: None, + description_i18n: "{}", + avatar_type: "emoji", + avatar_value: Some("🧪"), + agent_id: "claude", + rule_resource_type: "none", + rule_resource_ref: None, + recommended_prompts: "[]", + recommended_prompts_i18n: "{}", + default_model_mode: "auto", + default_model_value: None, + default_permission_mode: "auto", + default_permission_value: None, + default_thought_level_mode: "auto", + default_thought_level_value: None, + default_skills_mode: "auto", + default_skill_ids: "[]", + custom_skill_names: "[]", + default_disabled_builtin_skill_ids: "[]", + default_mcps_mode: "auto", + default_mcp_ids: "[]", + }, + ) .await .unwrap(); @@ -6518,46 +6957,52 @@ async fn create_routes_asset_avatar_in_assistant_identity_through_backend() { let workspace = ensure_test_workspace_path(); definition_repo - .upsert(&UpsertAssistantDefinitionParams { - id: "asstdef_data_avatar", - assistant_id: "custom-data-avatar", - source: "user", - owner_type: "user", - source_ref: Some("custom-data-avatar"), - name: "Data Avatar", - name_i18n: "{}", - description: None, - description_i18n: "{}", - avatar_type: "user_asset", - avatar_value: Some("data:image/svg+xml;base64,PHN2Zy8+"), - agent_id: "claude", - rule_resource_type: "none", - rule_resource_ref: None, - recommended_prompts: "[]", - recommended_prompts_i18n: "{}", - default_model_mode: "auto", - default_model_value: None, - default_permission_mode: "auto", - default_permission_value: None, - default_thought_level_mode: "auto", - default_thought_level_value: None, - default_skills_mode: "auto", - default_skill_ids: "[]", - custom_skill_names: "[]", - default_disabled_builtin_skill_ids: "[]", - default_mcps_mode: "auto", - default_mcp_ids: "[]", - }) + .upsert_for_user( + "user-1", + &UpsertAssistantDefinitionParams { + id: "asstdef_data_avatar", + assistant_id: "custom-data-avatar", + source: "user", + owner_type: "user", + source_ref: Some("custom-data-avatar"), + name: "Data Avatar", + name_i18n: "{}", + description: None, + description_i18n: "{}", + avatar_type: "user_asset", + avatar_value: Some("data:image/svg+xml;base64,PHN2Zy8+"), + agent_id: "claude", + rule_resource_type: "none", + rule_resource_ref: None, + recommended_prompts: "[]", + recommended_prompts_i18n: "{}", + default_model_mode: "auto", + default_model_value: None, + default_permission_mode: "auto", + default_permission_value: None, + default_thought_level_mode: "auto", + default_thought_level_value: None, + default_skills_mode: "auto", + default_skill_ids: "[]", + custom_skill_names: "[]", + default_disabled_builtin_skill_ids: "[]", + default_mcps_mode: "auto", + default_mcp_ids: "[]", + }, + ) .await .unwrap(); state_repo - .upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: "asstdef_data_avatar", - enabled: true, - sort_order: 0, - agent_id_override: None, - last_used_at: None, - }) + .upsert_for_user( + "user-1", + &UpsertAssistantOverlayParams { + assistant_definition_id: "asstdef_data_avatar", + enabled: true, + sort_order: 0, + agent_id_override: None, + last_used_at: None, + }, + ) .await .unwrap(); @@ -6610,7 +7055,7 @@ async fn assistant_backed_acp_build_options_include_snapshot_rule_as_preset_cont .unwrap(); let conv = create_assistant_backed_conversation(&svc, "user_1", Some("acp"), "codex", "preset-acp-rule").await; - let row = repo.get(&conv.id).await.unwrap().unwrap(); + let row = repo.get("user_1", &conv.id).await.unwrap().unwrap(); let options = svc.build_task_options(&row).await.unwrap(); match options.context.kind { @@ -6652,7 +7097,7 @@ async fn assistant_backed_aionrs_build_options_include_snapshot_rule_as_preset_r let conv = create_assistant_backed_conversation(&svc, "user_1", Some("aionrs"), "aionrs", "preset-aionrs-rule").await; - let row = repo.get(&conv.id).await.unwrap().unwrap(); + let row = repo.get("user_1", &conv.id).await.unwrap().unwrap(); let options = svc.build_task_options(&row).await.unwrap(); match options.context.kind { @@ -6709,25 +7154,31 @@ async fn create_prefers_assistant_snapshot_over_legacy_runtime_seed_fields() { .await .unwrap(); state_repo - .upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: "asstdef_preset_legacy_seed", - enabled: true, - sort_order: 0, - agent_id_override: Some("codex"), - last_used_at: None, - }) + .upsert_for_user( + "user-1", + &UpsertAssistantOverlayParams { + assistant_definition_id: "asstdef_preset_legacy_seed", + enabled: true, + sort_order: 0, + agent_id_override: Some("codex"), + last_used_at: None, + }, + ) .await .unwrap(); preference_repo - .upsert(&UpsertAssistantPreferenceParams { - assistant_definition_id: "asstdef_preset_legacy_seed", - last_model_id: Some("preferred-model"), - last_permission_value: Some("workspace-write"), - last_thought_level_value: Some("medium"), - last_skill_ids: r#"["legacy-skill"]"#, - last_disabled_builtin_skill_ids: r#"["legacy-disabled"]"#, - last_mcp_ids: r#"["legacy-mcp"]"#, - }) + .upsert_for_user( + "user-1", + &UpsertAssistantPreferenceParams { + assistant_definition_id: "asstdef_preset_legacy_seed", + last_model_id: Some("preferred-model"), + last_permission_value: Some("workspace-write"), + last_thought_level_value: Some("medium"), + last_skill_ids: r#"["legacy-skill"]"#, + last_disabled_builtin_skill_ids: r#"["legacy-disabled"]"#, + last_mcp_ids: r#"["legacy-mcp"]"#, + }, + ) .await .unwrap(); @@ -6756,7 +7207,7 @@ async fn create_prefers_assistant_snapshot_over_legacy_runtime_seed_fields() { assert_eq!(resp.extra["session_mode"], json!("workspace-write")); assert_eq!(resp.extra["current_mode_id"], json!("workspace-write")); - let snapshot = repo.get_assistant_snapshot(&resp.id).await.unwrap().unwrap(); + let snapshot = repo.get_assistant_snapshot("user_1", &resp.id).await.unwrap().unwrap(); assert_eq!(snapshot.agent_id, "8e1acf31"); assert_eq!(snapshot.resolved_model_id.as_deref(), Some("override-model")); assert_eq!(snapshot.resolved_permission_value.as_deref(), Some("workspace-write")); @@ -6782,13 +7233,16 @@ async fn create_prefers_snapshot_runtime_identity_over_legacy_extra_identity() { ) .await; overlay_repo - .upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: "asstdef_snapshot_identity", - enabled: true, - sort_order: 0, - agent_id_override: None, - last_used_at: None, - }) + .upsert_for_user( + "user-1", + &UpsertAssistantOverlayParams { + assistant_definition_id: "asstdef_snapshot_identity", + enabled: true, + sort_order: 0, + agent_id_override: None, + last_used_at: None, + }, + ) .await .unwrap(); @@ -6819,6 +7273,7 @@ async fn create_prefers_snapshot_runtime_identity_over_legacy_extra_identity() { let create_calls = acp_repo.create_calls(); assert_eq!(create_calls.len(), 1); + assert_eq!(create_calls[0].user_id, "user-1"); assert_eq!(create_calls[0].agent_id, "8e1acf31"); assert_eq!(create_calls[0].agent_source, "builtin"); assert_eq!(create_calls[0].agent_id, "8e1acf31"); @@ -6870,25 +7325,31 @@ async fn create_does_not_overwrite_preferences_for_fixed_skills_and_mcps() { .await .unwrap(); state_repo - .upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: "asstdef_preset_fixed", - enabled: true, - sort_order: 0, - agent_id_override: Some("codex"), - last_used_at: None, - }) + .upsert_for_user( + "user-1", + &UpsertAssistantOverlayParams { + assistant_definition_id: "asstdef_preset_fixed", + enabled: true, + sort_order: 0, + agent_id_override: Some("codex"), + last_used_at: None, + }, + ) .await .unwrap(); preference_repo - .upsert(&UpsertAssistantPreferenceParams { - assistant_definition_id: "asstdef_preset_fixed", - last_model_id: Some("legacy-model"), - last_permission_value: Some("workspace-write"), - last_thought_level_value: Some("legacy-thought"), - last_skill_ids: r#"["legacy-skill"]"#, - last_disabled_builtin_skill_ids: r#"["legacy-disabled"]"#, - last_mcp_ids: r#"["legacy-mcp"]"#, - }) + .upsert_for_user( + "user-1", + &UpsertAssistantPreferenceParams { + assistant_definition_id: "asstdef_preset_fixed", + last_model_id: Some("legacy-model"), + last_permission_value: Some("workspace-write"), + last_thought_level_value: Some("legacy-thought"), + last_skill_ids: r#"["legacy-skill"]"#, + last_disabled_builtin_skill_ids: r#"["legacy-disabled"]"#, + last_mcp_ids: r#"["legacy-mcp"]"#, + }, + ) .await .unwrap(); @@ -6914,7 +7375,11 @@ async fn create_does_not_overwrite_preferences_for_fixed_skills_and_mcps() { .unwrap(); let _resp = svc.create("user-1", req).await.unwrap(); - let updated_pref = preference_repo.get("asstdef_preset_fixed").await.unwrap().unwrap(); + let updated_pref = preference_repo + .get_for_user("user-1", "asstdef_preset_fixed") + .await + .unwrap() + .unwrap(); assert_eq!(updated_pref.last_model_id.as_deref(), Some("new-model")); assert_eq!(updated_pref.last_permission_value.as_deref(), Some("workspace-read")); assert_eq!(updated_pref.last_skill_ids, r#"["legacy-skill"]"#); @@ -6968,25 +7433,31 @@ async fn create_with_auto_builtin_defaults_without_preferences_keeps_snapshot_va .await .unwrap(); state_repo - .upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: "asstdef_preset_auto", - enabled: true, - sort_order: 0, - agent_id_override: Some("codex"), - last_used_at: None, - }) + .upsert_for_user( + "user-1", + &UpsertAssistantOverlayParams { + assistant_definition_id: "asstdef_preset_auto", + enabled: true, + sort_order: 0, + agent_id_override: Some("codex"), + last_used_at: None, + }, + ) .await .unwrap(); preference_repo - .upsert(&UpsertAssistantPreferenceParams { - assistant_definition_id: "asstdef_preset_auto", - last_model_id: None, - last_permission_value: None, - last_thought_level_value: None, - last_skill_ids: "[]", - last_disabled_builtin_skill_ids: "[]", - last_mcp_ids: "[]", - }) + .upsert_for_user( + "user-1", + &UpsertAssistantPreferenceParams { + assistant_definition_id: "asstdef_preset_auto", + last_model_id: None, + last_permission_value: None, + last_thought_level_value: None, + last_skill_ids: "[]", + last_disabled_builtin_skill_ids: "[]", + last_mcp_ids: "[]", + }, + ) .await .unwrap(); @@ -7009,7 +7480,11 @@ async fn create_with_auto_builtin_defaults_without_preferences_keeps_snapshot_va assert!(resp.extra.get("permission_mode").is_none()); assert!(resp.extra.get("assistant_snapshot").is_none()); - let updated_pref = preference_repo.get("asstdef_preset_auto").await.unwrap().unwrap(); + let updated_pref = preference_repo + .get_for_user("user-1", "asstdef_preset_auto") + .await + .unwrap() + .unwrap(); assert_eq!(updated_pref.last_model_id, None); assert_eq!(updated_pref.last_permission_value, None); assert_eq!(updated_pref.last_mcp_ids, "[]"); @@ -7237,7 +7712,7 @@ async fn get_backfills_legacy_row_and_persists() { assert_eq!(resp2.extra["skills"], json!(["cron", "pdf"])); // Verify the row on disk was persisted with the new shape. - let persisted = repo.get("legacy-1").await.unwrap().unwrap(); + let persisted = repo.get("user-1", "legacy-1").await.unwrap().unwrap(); let persisted_extra: serde_json::Value = serde_json::from_str(&persisted.extra).unwrap(); assert_eq!(persisted_extra["skills"], json!(["cron", "pdf"])); assert!(persisted_extra.get("enabled_skills").is_none()); @@ -7360,7 +7835,7 @@ async fn insert_raw_message_persists_row_and_broadcasts_stream() { created_at: 1234, }; - svc.insert_raw_message(&row).await.unwrap(); + svc.insert_raw_message("user_1", &row).await.unwrap(); let stored = repo.messages.lock().unwrap().clone(); assert_eq!(stored.len(), 1, "row must be persisted via repo.insert_message"); @@ -7410,25 +7885,28 @@ async fn seed_aionrs_conversation_with_snapshot( updated_at: 1, }; repo.create(&row).await.unwrap(); - repo.upsert_assistant_snapshot(&UpsertConversationAssistantSnapshotParams { - conversation_id: &row.id, - assistant_definition_id: "asstdef-seed", - assistant_id: "assistant-seed", - assistant_source: "builtin", - agent_id: "agent-seed", - rules_content: "", - default_model_mode: "auto", - resolved_model_id: None, - default_permission_mode, - resolved_permission_value, - default_thought_level_mode: "auto", - resolved_thought_level_value: None, - default_skills_mode: "auto", - resolved_skill_ids: "[]", - resolved_disabled_builtin_skill_ids: "[]", - default_mcps_mode: "auto", - resolved_mcp_ids: "[]", - }) + repo.upsert_assistant_snapshot( + &row.user_id, + &UpsertConversationAssistantSnapshotParams { + conversation_id: &row.id, + assistant_definition_id: "asstdef-seed", + assistant_id: "assistant-seed", + assistant_source: "builtin", + agent_id: "agent-seed", + rules_content: "", + default_model_mode: "auto", + resolved_model_id: None, + default_permission_mode, + resolved_permission_value, + default_thought_level_mode: "auto", + resolved_thought_level_value: None, + default_skills_mode: "auto", + resolved_skill_ids: "[]", + resolved_disabled_builtin_skill_ids: "[]", + default_mcps_mode: "auto", + resolved_mcp_ids: "[]", + }, + ) .await .unwrap(); row diff --git a/crates/aionui-conversation/src/service_test/acp_error_recovery_test.rs b/crates/aionui-conversation/src/service_test/acp_error_recovery_test.rs index 9c01cf755..b3bf29037 100644 --- a/crates/aionui-conversation/src/service_test/acp_error_recovery_test.rs +++ b/crates/aionui-conversation/src/service_test/acp_error_recovery_test.rs @@ -50,6 +50,7 @@ async fn send_message_clears_persisted_acp_model_after_model_not_found() { let conv = svc.create("user_1", make_create_req()).await.unwrap(); let workspace = ensure_test_workspace_path(); repo.update( + "user_1", &conv.id, &ConversationRowUpdate { extra: Some( @@ -106,7 +107,7 @@ async fn send_message_clears_persisted_acp_model_after_model_not_found() { }] ); - let row = repo.get(&conv.id).await.unwrap().unwrap(); + let row = repo.get("user_1", &conv.id).await.unwrap().unwrap(); let extra: serde_json::Value = serde_json::from_str(&row.extra).unwrap(); assert!(extra.get("workspace").is_some()); assert!( diff --git a/crates/aionui-conversation/src/session_context.rs b/crates/aionui-conversation/src/session_context.rs index d11e7a23a..6af20e2eb 100644 --- a/crates/aionui-conversation/src/session_context.rs +++ b/crates/aionui-conversation/src/session_context.rs @@ -205,7 +205,7 @@ impl<'a> SessionContextBuilder<'a> { serde_json::from_value(extra.clone()).map_err(|e| ConversationError::BadRequest { reason: format!("Invalid ACP build options: {e}"), })?; - config.user_id.get_or_insert_with(|| row.user_id.clone()); + config.user_id = Some(row.user_id.clone()); apply_team_seed_to_acp_config(&team, &mut config); normalize_cron_alias(row, &extra, &mut config.cron_job_id); @@ -225,7 +225,7 @@ impl<'a> SessionContextBuilder<'a> { let belongs_to_team = team.is_some(); let session_row = self .acp_session_repo - .get(&row.id) + .get_for_user(&row.user_id, &row.id) .await .map_err(|e| ConversationError::internal(format!("Failed to load acp_session row: {e}")))?; self.resolve_acp_identity(row, &mut config, &extra, session_row.as_ref()) @@ -259,7 +259,7 @@ impl<'a> SessionContextBuilder<'a> { if let Some(session_row) = session_row.filter(|row| !row.agent_id.is_empty()) { let metadata = self .agent_metadata_repo - .get(&session_row.agent_id) + .get_for_user(&row.user_id, &session_row.agent_id) .await .map_err(|e| ConversationError::internal(format!("agent_metadata lookup: {e}")))?; debug!( @@ -292,7 +292,7 @@ impl<'a> SessionContextBuilder<'a> { let Some(row_meta) = self .agent_metadata_repo - .find_builtin_by_backend(backend) + .find_builtin_by_backend_for_user(&row.user_id, backend) .await .map_err(|e| ConversationError::internal(format!("agent_metadata lookup: {e}")))? else { @@ -329,7 +329,7 @@ impl<'a> SessionContextBuilder<'a> { let db_state = self .acp_session_repo - .load_runtime_state(&row.id) + .load_runtime_state_for_user(&row.user_id, &row.id) .await .map_err(|e| ConversationError::internal(format!("Failed to load acp_session runtime state: {e}")))?; let snapshot = db_state.map(decode_persisted_session_state); @@ -384,7 +384,7 @@ fn build_aionrs_context( AionrsBuildExtra::default() } }; - config.user_id.get_or_insert_with(|| row.user_id.clone()); + config.user_id = Some(row.user_id.clone()); apply_team_seed_to_aionrs_config(&team, &mut config); let belongs_to_team = team.is_some(); // Team-bound sessions keep the team seed / create-time value; runtime @@ -605,12 +605,50 @@ mod tests { workspace_root: PathBuf, metadata_repo: Arc, acp_session_repo: Arc, + pool: sqlx::SqlitePool, } impl TestRepos { fn builder(&self) -> SessionContextBuilder<'_> { SessionContextBuilder::new(&self.workspace_root, &self.metadata_repo, &self.acp_session_repo) } + + async fn insert_conversation(&self, row: &ConversationRow) { + sqlx::query( + "INSERT OR IGNORE INTO users \ + (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, ?, ?)", + ) + .bind(&row.user_id) + .bind(&row.user_id) + .bind(row.created_at) + .bind(row.updated_at) + .execute(&self.pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO conversations \ + (id, user_id, name, type, extra, model, status, source, \ + channel_chat_id, pinned, pinned_at, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(&row.id) + .bind(&row.user_id) + .bind(&row.name) + .bind(&row.r#type) + .bind(&row.extra) + .bind(&row.model) + .bind(row.status.as_deref().unwrap_or("pending")) + .bind(&row.source) + .bind(&row.channel_chat_id) + .bind(row.pinned) + .bind(row.pinned_at) + .bind(row.created_at) + .bind(row.updated_at) + .execute(&self.pool) + .await + .unwrap(); + } } #[derive(Clone)] @@ -657,6 +695,7 @@ mod tests { workspace_root, metadata_repo, acp_session_repo, + pool: db.pool().clone(), } } @@ -744,6 +783,23 @@ mod tests { assert_eq!(acp.config.backend.as_deref(), Some("claude")); } + #[tokio::test] + async fn acp_extra_user_id_is_overridden_by_conversation_owner() { + let repos = setup().await; + let row = row( + "acp", + serde_json::json!({ + "backend": "claude", + "user_id": "other-user" + }), + None, + ); + + let context = repos.builder().build(&row).await.unwrap(); + let acp = acp_context(context); + assert_eq!(acp.config.user_id.as_deref(), Some("user-1")); + } + #[tokio::test] async fn acp_builtin_backend_fallback_resolves_agent_id() { let repos = setup().await; @@ -785,9 +841,20 @@ mod tests { async fn acp_persisted_runtime_is_loaded_before_legacy_seed() { let repos = setup().await; upsert_builtin(&repos, "builtin-claude-test", "claude").await; + let row = row( + "acp", + serde_json::json!({ + "backend": "claude", + "current_mode_id": "legacy-mode", + "current_model_id": "legacy-model" + }), + None, + ); + repos.insert_conversation(&row).await; repos .acp_session_repo .create(&CreateAcpSessionParams { + user_id: "user-1", conversation_id: "conv-1", agent_source: "builtin", agent_id: "builtin-claude-test", @@ -796,12 +863,13 @@ mod tests { .unwrap(); repos .acp_session_repo - .update_session_id("conv-1", "sess-1") + .update_session_id_for_user("user-1", "conv-1", "sess-1") .await .unwrap(); repos .acp_session_repo - .save_runtime_state( + .save_runtime_state_for_user( + "user-1", "conv-1", &SaveRuntimeStateParams { current_mode_id: Some(Some("persisted-mode")), @@ -812,15 +880,6 @@ mod tests { ) .await .unwrap(); - let row = row( - "acp", - serde_json::json!({ - "backend": "claude", - "current_mode_id": "legacy-mode", - "current_model_id": "legacy-model" - }), - None, - ); let context = repos.builder().build(&row).await.unwrap(); let acp = acp_context(context); @@ -835,9 +894,20 @@ mod tests { async fn acp_unassigned_session_runtime_is_startup_seed_not_resume_snapshot() { let repos = setup().await; upsert_builtin(&repos, "builtin-codex-test", "codex").await; + let row = row( + "acp", + serde_json::json!({ + "backend": "codex", + "current_mode_id": "full-access", + "current_model_id": "gpt-5.5" + }), + None, + ); + repos.insert_conversation(&row).await; repos .acp_session_repo .create(&CreateAcpSessionParams { + user_id: "user-1", conversation_id: "conv-1", agent_source: "builtin", agent_id: "builtin-codex-test", @@ -846,7 +916,8 @@ mod tests { .unwrap(); repos .acp_session_repo - .save_runtime_state( + .save_runtime_state_for_user( + "user-1", "conv-1", &SaveRuntimeStateParams { current_mode_id: Some(Some("full-access")), @@ -857,15 +928,6 @@ mod tests { ) .await .unwrap(); - let row = row( - "acp", - serde_json::json!({ - "backend": "codex", - "current_mode_id": "full-access", - "current_model_id": "gpt-5.5" - }), - None, - ); let context = repos.builder().build(&row).await.unwrap(); let acp = acp_context(context); @@ -879,16 +941,18 @@ mod tests { let repos = setup().await; upsert_builtin(&repos, "builtin-claude-test", "claude").await; upsert_builtin(&repos, "builtin-codex-test", "codex").await; + let row = row("acp", serde_json::json!({ "backend": "claude" }), None); + repos.insert_conversation(&row).await; repos .acp_session_repo .create(&CreateAcpSessionParams { + user_id: "user-1", conversation_id: "conv-1", agent_source: "builtin", agent_id: "builtin-codex-test", }) .await .unwrap(); - let row = row("acp", serde_json::json!({ "backend": "claude" }), None); let context = repos.builder().build(&row).await.unwrap(); let acp = acp_context(context); @@ -1009,6 +1073,16 @@ mod tests { assert_eq!(aionrs.config.team_mcp_stdio_config.unwrap().port, 5252); } + #[tokio::test] + async fn aionrs_extra_user_id_is_overridden_by_conversation_owner() { + let repos = setup().await; + let row = row("aionrs", serde_json::json!({ "user_id": "other-user" }), None); + + let context = repos.builder().build(&row).await.unwrap(); + let aionrs = aionrs_context(context); + assert_eq!(aionrs.config.user_id.as_deref(), Some("user-1")); + } + #[tokio::test] async fn aionrs_uses_conversation_model_and_ignores_legacy_extra_model() { let repos = setup().await; diff --git a/crates/aionui-conversation/src/skill_resolver.rs b/crates/aionui-conversation/src/skill_resolver.rs index 8b47f4e4e..30a731249 100644 --- a/crates/aionui-conversation/src/skill_resolver.rs +++ b/crates/aionui-conversation/src/skill_resolver.rs @@ -26,6 +26,11 @@ pub trait SkillResolver: Send + Sync { /// same search order as `materialize_skills_for_agent`. async fn resolve_skills(&self, names: &[String]) -> Vec; + /// Resolve each skill name for a specific Core user. + async fn resolve_skills_for_user(&self, _user_id: &str, names: &[String]) -> Vec { + self.resolve_skills(names).await + } + /// Load full skill bodies for prompt-protocol agents that request /// `[LOAD_SKILL: name]` in their response. async fn load_skill_bodies(&self, names: &[String]) -> Vec { @@ -33,6 +38,12 @@ pub trait SkillResolver: Send + Sync { load_resolved_skill_bodies(&resolved).await } + /// Load full skill bodies for prompt-protocol agents under one Core user. + async fn load_skill_bodies_for_user(&self, user_id: &str, names: &[String]) -> Vec { + let resolved = self.resolve_skills_for_user(user_id, names).await; + load_resolved_skill_bodies(&resolved).await + } + /// Create symlinks pointing at each resolved skill inside the given /// workspace's per-backend native skills directories. `rel_dirs` is /// the list of relative paths (e.g. `.claude/skills`) to populate. @@ -119,14 +130,19 @@ impl SkillResolver for ExtensionSkillResolver { } async fn resolve_skills(&self, names: &[String]) -> Vec { + self.resolve_skills_for_user("system_default_user", names).await + } + + async fn resolve_skills_for_user(&self, user_id: &str, names: &[String]) -> Vec { if names.is_empty() { return Vec::new(); } // Conversation_id is validated upstream; we don't use a real one here // because this resolver is purely a path-resolution helper. - match aionui_extension::materialize_skills_for_agent_with_repo( + match aionui_extension::materialize_skills_for_agent_with_repo_for_user( &self.paths, self.skill_repo.as_ref(), + user_id, "workspace-link", names, ) @@ -190,7 +206,7 @@ impl SkillResolver for FixedSkillResolver { #[cfg(test)] mod tests { use super::*; - use aionui_db::SqliteSkillRepository; + use aionui_db::{SqliteSkillRepository, UpsertSkillParams}; fn write_skill(dir: &Path, name: &str, description: &str) { let skill_dir = dir.join(name); @@ -238,4 +254,65 @@ mod tests { assert_eq!(resolver.auto_inject_names().await, vec!["auto-cron".to_string()]); } + + #[tokio::test] + async fn extension_resolver_resolves_user_scoped_skill_rows() { + let tmp = tempfile::TempDir::new().unwrap(); + let paths = Arc::new(aionui_extension::SkillPaths { + data_dir: tmp.path().to_path_buf(), + user_skills_dir: tmp.path().join("skills"), + cron_skills_dir: tmp.path().join("cron").join("skills"), + builtin_skills_dir: tmp.path().join("builtin-skills"), + builtin_rules_dir: tmp.path().join("builtin-rules"), + assistant_rules_dir: tmp.path().join("assistant-rules"), + assistant_skills_dir: tmp.path().join("assistant-skills"), + }); + let user_a_skill = tmp.path().join("user-a-skill"); + let user_b_skill = tmp.path().join("user-b-skill"); + write_skill(&user_a_skill, "shared", "User A skill"); + write_skill(&user_b_skill, "shared", "User B skill"); + + let db = aionui_db::init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES ('user_b', 'local', 'user_b', 'hash', 'active', 0, 1, 1)", + ) + .execute(db.pool()) + .await + .unwrap(); + + let repo = Arc::new(SqliteSkillRepository::new(db.pool().clone())); + let user_a_skill_path = user_a_skill.join("shared").to_string_lossy().into_owned(); + let user_b_skill_path = user_b_skill.join("shared").to_string_lossy().into_owned(); + repo.upsert_for_user( + "system_default_user", + UpsertSkillParams { + name: "shared", + description: Some("User A skill"), + path: &user_a_skill_path, + source: "user", + enabled: true, + }, + ) + .await + .unwrap(); + repo.upsert_for_user( + "user_b", + UpsertSkillParams { + name: "shared", + description: Some("User B skill"), + path: &user_b_skill_path, + source: "user", + enabled: true, + }, + ) + .await + .unwrap(); + + let resolver = ExtensionSkillResolver::new(paths, repo); + let resolved = resolver.resolve_skills_for_user("user_b", &["shared".to_owned()]).await; + + assert_eq!(resolved.len(), 1); + assert_eq!(resolved[0].source_path, user_b_skill.join("shared")); + } } diff --git a/crates/aionui-conversation/src/startup_recovery.rs b/crates/aionui-conversation/src/startup_recovery.rs index 6193cdc0a..ecf94f349 100644 --- a/crates/aionui-conversation/src/startup_recovery.rs +++ b/crates/aionui-conversation/src/startup_recovery.rs @@ -26,7 +26,8 @@ impl ConversationService { }; let mut recovered = 0usize; - for row in rows { + for stale in rows { + let row = stale.message; if !self .runtime_persistence() .allows(&row.conversation_id, RuntimeWriteKind::StartupRecovery) @@ -41,7 +42,11 @@ impl ConversationService { hidden: Some(matches!(action, StartupRecoveryAction::FinishEmptyPlaceholder)), }; - match self.conversation_repo().update_message(&row.id, &update).await { + match self + .conversation_repo() + .update_message(&stale.user_id, &row.conversation_id, &row.id, &update) + .await + { Ok(()) => { recovered += 1; info!( diff --git a/crates/aionui-conversation/src/stream_persistence.rs b/crates/aionui-conversation/src/stream_persistence.rs index dfb381729..eae727ae3 100644 --- a/crates/aionui-conversation/src/stream_persistence.rs +++ b/crates/aionui-conversation/src/stream_persistence.rs @@ -66,6 +66,7 @@ pub(crate) struct FinalTextOverride { #[derive(Clone)] pub(crate) struct StreamPersistenceAdapter { + user_id: String, conversation_id: String, msg_id: String, repo: Arc, @@ -74,12 +75,14 @@ pub(crate) struct StreamPersistenceAdapter { impl StreamPersistenceAdapter { pub fn new( + user_id: String, conversation_id: String, msg_id: String, repo: Arc, persistence: Option, ) -> Self { Self { + user_id, conversation_id, msg_id, repo, @@ -100,9 +103,14 @@ impl StreamPersistenceAdapter { runtime: Option, ) { if let Some(persistence) = &self.persistence { - RuntimeCompletionPublisher::new(self.repo.clone(), broadcaster.clone(), persistence.clone()) - .publish(&self.conversation_id, turn_id, runtime) - .await; + RuntimeCompletionPublisher::new( + self.user_id.clone(), + self.repo.clone(), + broadcaster.clone(), + persistence.clone(), + ) + .publish(&self.conversation_id, turn_id, runtime) + .await; return; } @@ -111,11 +119,12 @@ impl StreamPersistenceAdapter { updated_at: Some(now_ms()), ..Default::default() }; - if let Err(e) = self.repo.update(&self.conversation_id, &update).await { + if let Err(e) = self.repo.update(&self.user_id, &self.conversation_id, &update).await { log_persist_error(&e, "Failed to update conversation status"); } let payload = json!({ + "user_id": self.user_id, "conversation_id": self.conversation_id, "session_id": self.conversation_id, "turn_id": turn_id, @@ -151,7 +160,11 @@ impl StreamPersistenceAdapter { status: Some(Some("work".into())), hidden: None, }; - if let Err(e) = self.repo.update_message(&segment.id, &update).await { + if let Err(e) = self + .repo + .update_message(&self.user_id, &self.conversation_id, &segment.id, &update) + .await + { log_persist_error(&e, "Failed to update streaming text segment"); } } else { @@ -166,7 +179,7 @@ impl StreamPersistenceAdapter { hidden: false, created_at: segment.created_at, }; - if let Err(e) = self.repo.insert_message(&row).await { + if let Err(e) = self.repo.insert_message(&self.user_id, &row).await { log_persist_error(&e, "Failed to create streaming text segment"); } segment.record_created = true; @@ -189,7 +202,11 @@ impl StreamPersistenceAdapter { status: Some(Some(status.to_owned())), hidden: Some(false), }; - if let Err(e) = self.repo.update_message(&segment.id, &update).await { + if let Err(e) = self + .repo + .update_message(&self.user_id, &self.conversation_id, &segment.id, &update) + .await + { log_persist_error(&e, "Failed to finalize text segment"); return None; } @@ -205,7 +222,7 @@ impl StreamPersistenceAdapter { hidden: false, created_at: segment.created_at, }; - if let Err(e) = self.repo.insert_message(&row).await { + if let Err(e) = self.repo.insert_message(&self.user_id, &row).await { log_persist_error(&e, "Failed to create finalized text segment"); return None; } @@ -236,7 +253,11 @@ impl StreamPersistenceAdapter { status: Some(Some(status.to_owned())), hidden: Some(hidden), }; - if let Err(e) = self.repo.update_message(&primary_segment.id, &update).await { + if let Err(e) = self + .repo + .update_message(&self.user_id, &self.conversation_id, &primary_segment.id, &update) + .await + { log_persist_error(&e, "Failed to rewrite finalized text segment"); } overrides.push(FinalTextOverride { @@ -251,7 +272,11 @@ impl StreamPersistenceAdapter { status: Some(Some(status.to_owned())), hidden: Some(true), }; - if let Err(e) = self.repo.update_message(&segment.id, &hide_update).await { + if let Err(e) = self + .repo + .update_message(&self.user_id, &self.conversation_id, &segment.id, &hide_update) + .await + { log_persist_error(&e, "Failed to hide superseded text segment"); } overrides.push(FinalTextOverride { @@ -267,7 +292,11 @@ impl StreamPersistenceAdapter { status: Some(Some(status.to_owned())), hidden: Some(false), }; - if let Err(e) = self.repo.update_message(&segment.id, &status_update).await { + if let Err(e) = self + .repo + .update_message(&self.user_id, &self.conversation_id, &segment.id, &status_update) + .await + { log_persist_error(&e, "Failed to finalize text segment status"); } } @@ -284,7 +313,7 @@ impl StreamPersistenceAdapter { hidden: false, created_at: now_ms(), }; - if let Err(e) = self.repo.insert_message(&row).await { + if let Err(e) = self.repo.insert_message(&self.user_id, &row).await { log_persist_error(&e, "Failed to create final fallback message"); } } @@ -310,7 +339,7 @@ impl StreamPersistenceAdapter { hidden: false, created_at: now_ms(), }; - if let Err(e) = self.repo.insert_message(&row).await { + if let Err(e) = self.repo.insert_message(&self.user_id, &row).await { log_persist_error(&e, "Failed to store error message"); } } @@ -343,7 +372,7 @@ impl StreamPersistenceAdapter { hidden: false, created_at: now_ms(), }; - if let Err(e) = self.repo.insert_message(&row).await { + if let Err(e) = self.repo.insert_message(&self.user_id, &row).await { log_persist_error(&e, "Failed to store tip message"); } } @@ -373,7 +402,7 @@ impl StreamPersistenceAdapter { hidden: false, created_at: segment.started_at, }; - if let Err(e) = self.repo.insert_message(&row).await { + if let Err(e) = self.repo.insert_message(&self.user_id, &row).await { log_persist_error(&e, "Failed to persist thinking message"); } } @@ -414,7 +443,7 @@ impl StreamPersistenceAdapter { hidden: false, created_at: now_ms(), }; - if let Err(e) = self.repo.upsert_message(&row).await { + if let Err(e) = self.repo.upsert_message(&self.user_id, &row).await { error!( call_id = %data.call_id, tool = %data.name, @@ -464,7 +493,7 @@ impl StreamPersistenceAdapter { hidden: false, created_at: now_ms(), }; - if let Err(e) = self.repo.upsert_message(&row).await { + if let Err(e) = self.repo.upsert_message(&self.user_id, &row).await { error!(error = %ErrorChain(&e), "Failed to upsert acp_tool_call message"); } } @@ -491,7 +520,7 @@ impl StreamPersistenceAdapter { let existing = self .repo - .get_message_by_msg_id(&self.conversation_id, &group_id, "tool_group") + .get_message_by_msg_id(&self.user_id, &self.conversation_id, &group_id, "tool_group") .await .unwrap_or(None); @@ -501,7 +530,11 @@ impl StreamPersistenceAdapter { status: Some(Some(status.to_owned())), hidden: None, }; - if let Err(e) = self.repo.update_message(&group_id, &update).await { + if let Err(e) = self + .repo + .update_message(&self.user_id, &self.conversation_id, &group_id, &update) + .await + { error!(error = %ErrorChain(&e), "Failed to update tool_group message"); } } else { @@ -516,7 +549,7 @@ impl StreamPersistenceAdapter { hidden: false, created_at: now_ms(), }; - if let Err(e) = self.repo.insert_message(&row).await { + if let Err(e) = self.repo.insert_message(&self.user_id, &row).await { error!(error = %ErrorChain(&e), "Failed to persist tool_group message"); } } diff --git a/crates/aionui-conversation/src/stream_relay.rs b/crates/aionui-conversation/src/stream_relay.rs index ab77b7188..243551a53 100644 --- a/crates/aionui-conversation/src/stream_relay.rs +++ b/crates/aionui-conversation/src/stream_relay.rs @@ -123,7 +123,8 @@ impl StreamRelay { repo: Arc, broadcaster: Arc, ) -> Self { - let adapter = StreamPersistenceAdapter::new(conversation_id.clone(), msg_id.clone(), repo, None); + let adapter = + StreamPersistenceAdapter::new(user_id.clone(), conversation_id.clone(), msg_id.clone(), repo, None); Self { conversation_id, msg_id, @@ -743,6 +744,7 @@ impl StreamRelay { let middleware = MessageMiddleware::new_with_skill_loader(self.skill_resolver.as_ref().map(|resolver| { Box::new(SharedSkillResolver { resolver: Arc::clone(resolver), + user_id: self.user_id.clone(), allowed_skill_names: self.allowed_skill_names.clone(), }) as Box })); @@ -777,6 +779,8 @@ impl StreamRelay { fn broadcast_stream_payload(&self, mut payload: serde_json::Value) { if let Some(obj) = payload.as_object_mut() { + obj.entry("user_id") + .or_insert_with(|| serde_json::Value::String(self.user_id.clone())); obj.entry("turn_id") .or_insert_with(|| serde_json::Value::String(self.turn_id.clone())); } @@ -787,6 +791,7 @@ impl StreamRelay { struct SharedSkillResolver { resolver: Arc, + user_id: String, allowed_skill_names: Vec, } @@ -801,7 +806,7 @@ impl ISkillLoadService for SharedSkillResolver { .filter(|name| self.allowed_skill_names.iter().any(|allowed| allowed == *name)) .cloned() .collect(); - self.resolver.load_skill_bodies(&filtered).await + self.resolver.load_skill_bodies_for_user(&self.user_id, &filtered).await } } @@ -854,6 +859,7 @@ mod tests { #[derive(Default)] struct RecordingSkillResolverForRelay { requested: Mutex>, + requested_user_ids: Mutex>, } #[async_trait::async_trait] @@ -866,7 +872,8 @@ mod tests { Vec::new() } - async fn load_skill_bodies(&self, names: &[String]) -> Vec { + async fn load_skill_bodies_for_user(&self, user_id: &str, names: &[String]) -> Vec { + self.requested_user_ids.lock().unwrap().push(user_id.to_owned()); self.requested.lock().unwrap().extend(names.iter().cloned()); names .iter() @@ -893,6 +900,7 @@ mod tests { let resolver: Arc = concrete.clone(); let loader = SharedSkillResolver { resolver, + user_id: "system_default_user".into(), allowed_skill_names: vec!["cron".into()], }; @@ -901,6 +909,10 @@ mod tests { assert_eq!(loaded.len(), 1); assert_eq!(loaded[0].name, "cron"); assert_eq!(concrete.requested.lock().unwrap().as_slice(), ["cron"]); + assert_eq!( + concrete.requested_user_ids.lock().unwrap().as_slice(), + ["system_default_user"] + ); } #[tokio::test] @@ -2124,7 +2136,8 @@ mod tests { repo.set_not_found(true); let repo: Arc = repo; let bus: Arc = Arc::new(aionui_realtime::BroadcastEventBus::new(64)); - let adapter = StreamPersistenceAdapter::new("deleted-conv".into(), "msg-1".into(), repo, None); + let adapter = + StreamPersistenceAdapter::new("user-test".into(), "deleted-conv".into(), "msg-1".into(), repo, None); adapter.complete_conversation(&bus, "turn-1", None).await; } @@ -2328,19 +2341,27 @@ mod tests { #[async_trait::async_trait] impl IConversationRepository for RecordingRepo { - async fn get(&self, _id: &str) -> Result, DbError> { + async fn get(&self, _user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } + async fn owner_user_id(&self, _id: &str) -> Result, DbError> { + Ok(Some("user-1".into())) + } async fn create(&self, _row: &aionui_db::models::ConversationRow) -> Result<(), DbError> { Ok(()) } - async fn update(&self, _id: &str, _updates: &aionui_db::ConversationRowUpdate) -> Result<(), DbError> { + async fn update( + &self, + _user_id: &str, + _id: &str, + _updates: &aionui_db::ConversationRowUpdate, + ) -> Result<(), DbError> { if self.not_found.load(Ordering::Acquire) { return Err(DbError::NotFound("Conversation deleted-conv not found".into())); } Ok(()) } - async fn delete(&self, _id: &str) -> Result<(), DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), DbError> { Ok(()) } async fn list_paginated( @@ -2379,6 +2400,7 @@ mod tests { } async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &aionui_db::MessagePageParams, ) -> Result { @@ -2388,7 +2410,7 @@ mod tests { has_more_after: false, }) } - async fn insert_message(&self, row: &MessageRow) -> Result<(), DbError> { + async fn insert_message(&self, _user_id: &str, row: &MessageRow) -> Result<(), DbError> { if self.not_found.load(Ordering::Acquire) { return Err(DbError::NotFound(format!("Message '{}'", row.id))); } @@ -2398,7 +2420,7 @@ mod tests { self.inserts.lock().unwrap().push(row.clone()); Ok(()) } - async fn upsert_message(&self, row: &MessageRow) -> Result<(), DbError> { + async fn upsert_message(&self, _user_id: &str, row: &MessageRow) -> Result<(), DbError> { if self.not_found.load(Ordering::Acquire) { return Err(DbError::NotFound(format!("Message '{}'", row.id))); } @@ -2422,18 +2444,25 @@ mod tests { } Ok(()) } - async fn update_message(&self, id: &str, updates: &aionui_db::MessageRowUpdate) -> Result<(), DbError> { + async fn update_message( + &self, + _user_id: &str, + _conversation_id: &str, + id: &str, + updates: &aionui_db::MessageRowUpdate, + ) -> Result<(), DbError> { if self.not_found.load(Ordering::Acquire) { return Err(DbError::NotFound(format!("Message '{id}' not found"))); } self.updates.lock().unwrap().push((id.to_owned(), updates.clone())); Ok(()) } - async fn delete_messages_by_conversation(&self, _conv_id: &str) -> Result<(), DbError> { + async fn delete_messages_by_conversation(&self, _user_id: &str, _conv_id: &str) -> Result<(), DbError> { Ok(()) } async fn get_message_by_msg_id( &self, + _user_id: &str, _conv_id: &str, msg_id: &str, msg_type: &str, diff --git a/crates/aionui-conversation/src/turn_orchestrator.rs b/crates/aionui-conversation/src/turn_orchestrator.rs index db2078da4..85d435b75 100644 --- a/crates/aionui-conversation/src/turn_orchestrator.rs +++ b/crates/aionui-conversation/src/turn_orchestrator.rs @@ -129,6 +129,7 @@ impl ConversationTurnOrchestrator { let failure_message = send_error_display_message(&send_error); record_agent_session_failure( &self.service, + &input.user_id, availability_agent_id.as_deref(), "session_build_failed", &failure_message, @@ -136,6 +137,7 @@ impl ConversationTurnOrchestrator { .await; self.service .persist_and_broadcast_send_failure_tip( + &input.user_id, &input.conv_id, &input.turn_id, &send_error, @@ -151,7 +153,12 @@ impl ConversationTurnOrchestrator { if let Err(err) = self .service - .maybe_persist_workspace(&input.conv_id, &input.stored_workspace, agent.workspace()) + .maybe_persist_workspace( + &input.user_id, + &input.conv_id, + &input.stored_workspace, + agent.workspace(), + ) .await { let top_level_code = err.error_code(); @@ -166,6 +173,7 @@ impl ConversationTurnOrchestrator { ); self.service .persist_and_broadcast_send_failure_tip( + &input.user_id, &input.conv_id, &input.turn_id, &send_error, @@ -244,6 +252,7 @@ impl ConversationTurnOrchestrator { ); self.service .persist_and_broadcast_send_failure_tip( + &input.user_id, &input.conv_id, &input.turn_id, &send_error, @@ -260,6 +269,7 @@ impl ConversationTurnOrchestrator { let send_agent = agent.clone(); let conv_id_send = input.conv_id.clone(); let turn_id_for_send = input.turn_id.clone(); + let feedback_user_id = input.user_id.clone(); let feedback_service = self.service.clone(); let feedback_agent_id = availability_agent_id.clone(); let (send_error_tx, send_error_rx) = oneshot::channel(); @@ -269,6 +279,7 @@ impl ConversationTurnOrchestrator { let failure_message = send_error_display_message(&e); record_agent_session_failure( &feedback_service, + &feedback_user_id, feedback_agent_id.as_deref(), "session_send_failed", &failure_message, @@ -311,6 +322,7 @@ impl ConversationTurnOrchestrator { persist_session_key( self.service.conversation_repo(), &persistence, + &input.user_id, &input.conv_id, &session_key, ) @@ -451,6 +463,7 @@ impl ConversationTurnOrchestrator { ); self.service .evict_acp_task_after_terminal_error( + &input.user_id, &conv_id, attempt_result.agent_type, &attempt_result.outcome, @@ -466,7 +479,13 @@ impl ConversationTurnOrchestrator { { let send_error = AgentSendError::from_stream_error_data(data); self.service - .persist_and_broadcast_send_failure_tip(&conv_id, &turn_id, &send_error, None) + .persist_and_broadcast_send_failure_tip( + &input.user_id, + &conv_id, + &turn_id, + &send_error, + None, + ) .await; } @@ -475,6 +494,7 @@ impl ConversationTurnOrchestrator { AgentHealthAction::EvictAcpTask { .. } => { self.service .evict_acp_task_after_terminal_error( + &input.user_id, &conv_id, attempt_result.agent_type, &attempt_result.outcome, @@ -494,6 +514,7 @@ impl ConversationTurnOrchestrator { // availability so the list stops showing it as plainly usable. record_agent_session_failure( &self.service, + &input.user_id, availability_agent_id(&input.build_options).as_deref(), "auth_required", final_error_message @@ -502,12 +523,17 @@ impl ConversationTurnOrchestrator { ) .await; } else if !final_failed { - record_agent_session_success(&self.service, availability_agent_id(&input.build_options).as_deref()).await; + record_agent_session_success( + &self.service, + &input.user_id, + availability_agent_id(&input.build_options).as_deref(), + ) + .await; } let was_deleting = turn_claim.release_for_turn(&turn_id); self.service - .complete_released_turn(&conv_id, &turn_id, was_deleting) + .complete_released_turn(&input.user_id, &conv_id, &turn_id, was_deleting) .await; ConversationTurnResult { @@ -581,6 +607,7 @@ fn turn_attempt_error_message(summary: &TurnAttemptSummary) -> Option { async fn record_agent_session_failure( service: &ConversationService, + user_id: &str, agent_id: Option<&str>, code: &str, message: &str, @@ -591,8 +618,9 @@ async fn record_agent_session_failure( let Some(feedback) = service.agent_availability_feedback() else { return; }; - if let Err(error) = feedback.record_session_failure(agent_id, code, message).await { + if let Err(error) = feedback.record_session_failure(user_id, agent_id, code, message).await { warn!( + user_id, agent_id, code, error = %ErrorChain(&error), @@ -601,15 +629,16 @@ async fn record_agent_session_failure( } } -async fn record_agent_session_success(service: &ConversationService, agent_id: Option<&str>) { +async fn record_agent_session_success(service: &ConversationService, user_id: &str, agent_id: Option<&str>) { let Some(agent_id) = agent_id else { return; }; let Some(feedback) = service.agent_availability_feedback() else { return; }; - if let Err(error) = feedback.record_session_success(agent_id).await { + if let Err(error) = feedback.record_session_success(user_id, agent_id).await { warn!( + user_id, agent_id, error = %ErrorChain(&error), "Failed to record agent availability session success" diff --git a/crates/aionui-conversation/tests/acp_tool_call_persistence.rs b/crates/aionui-conversation/tests/acp_tool_call_persistence.rs index 45bcefde3..eadba1fcd 100644 --- a/crates/aionui-conversation/tests/acp_tool_call_persistence.rs +++ b/crates/aionui-conversation/tests/acp_tool_call_persistence.rs @@ -19,10 +19,11 @@ async fn run_acp_tool_call_update_without_insert_creates_placeholder() { let db = init_database_memory().await.unwrap(); let user_repo = SqliteUserRepository::new(db.pool().clone()); let user = user_repo.create_user("user-1", "hash").await.unwrap(); + let user_id = user.id.clone(); let repo = Arc::new(SqliteConversationRepository::new(db.pool().clone())); repo.create(&ConversationRow { id: "conv-1".into(), - user_id: user.id, + user_id: user_id.clone(), name: "test".into(), r#type: "acp".into(), extra: "{}".into(), @@ -44,7 +45,7 @@ async fn run_acp_tool_call_update_without_insert_creates_placeholder() { "conv-1".into(), "asst-1".into(), "turn-1".into(), - "user-1".into(), + user_id.clone(), repo.clone(), bus, ); @@ -72,6 +73,7 @@ async fn run_acp_tool_call_update_without_insert_creates_placeholder() { let messages = repo .list_messages_page( + &user_id, "conv-1", &MessagePageParams { limit: 20, @@ -93,10 +95,11 @@ async fn run_acp_tool_call_late_initial_event_merges_with_update_placeholder() { let db = init_database_memory().await.unwrap(); let user_repo = SqliteUserRepository::new(db.pool().clone()); let user = user_repo.create_user("user-1", "hash").await.unwrap(); + let user_id = user.id.clone(); let repo = Arc::new(SqliteConversationRepository::new(db.pool().clone())); repo.create(&ConversationRow { id: "conv-1".into(), - user_id: user.id, + user_id: user_id.clone(), name: "test".into(), r#type: "acp".into(), extra: "{}".into(), @@ -118,7 +121,7 @@ async fn run_acp_tool_call_late_initial_event_merges_with_update_placeholder() { "conv-1".into(), "asst-1".into(), "turn-1".into(), - "user-1".into(), + user_id.clone(), repo.clone(), bus, ); @@ -162,6 +165,7 @@ async fn run_acp_tool_call_late_initial_event_merges_with_update_placeholder() { let messages = repo .list_messages_page( + &user_id, "conv-1", &MessagePageParams { limit: 20, diff --git a/crates/aionui-conversation/tests/conversation_crud.rs b/crates/aionui-conversation/tests/conversation_crud.rs index d72c2b8ba..d0a9288f9 100644 --- a/crates/aionui-conversation/tests/conversation_crud.rs +++ b/crates/aionui-conversation/tests/conversation_crud.rs @@ -228,6 +228,25 @@ async fn t1_3_create_with_optional_fields() { assert_eq!(resp.channel_chat_id.as_deref(), Some("user:123")); } +#[tokio::test] +async fn create_ignores_extra_user_id() { + let (svc, _, _task_mgr) = setup().await; + let workspace = ensure_test_workspace_path(); + + let req: CreateConversationRequest = serde_json::from_value(json!({ + "type": "acp", + "extra": { "workspace": workspace, "user_id": "forged-other-user" } + })) + .unwrap(); + let resp = svc.create(USER_ID, req).await.unwrap(); + + assert!( + !resp.extra.as_object().unwrap().contains_key("user_id"), + "create must not persist client-supplied extra.user_id; got {:?}", + resp.extra + ); +} + // ── T2: List conversations ───────────────────────────────────────── #[tokio::test] @@ -441,6 +460,25 @@ async fn t4_4_extra_merge_preserves_existing_keys() { assert_eq!(updated.extra["contextFileName"], "ctx.md"); } +#[tokio::test] +async fn update_ignores_extra_user_id() { + let (svc, _, task_mgr) = setup().await; + let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); + + let req: UpdateConversationRequest = serde_json::from_value(json!({ + "extra": { "user_id": "forged-other-user", "display_density": "compact" } + })) + .unwrap(); + let updated = svc.update(USER_ID, &conv.id, req, &task_mgr).await.unwrap(); + + assert!( + !updated.extra.as_object().unwrap().contains_key("user_id"), + "update must not persist client-supplied extra.user_id; got {:?}", + updated.extra + ); + assert_eq!(updated.extra["display_density"], "compact"); +} + #[tokio::test] async fn t4_5_update_model() { let (svc, _, task_mgr) = setup().await; @@ -790,9 +828,9 @@ async fn complete_turn_skips_status_update_when_conversation_is_deleting() { let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); svc.runtime_state().mark_deleting(&conv.id); - svc.complete_turn(&conv.id, "turn-1").await; + svc.complete_turn(USER_ID, &conv.id, "turn-1").await; - let row = svc.conversation_repo().get(&conv.id).await.unwrap().unwrap(); + let row = svc.conversation_repo().get(USER_ID, &conv.id).await.unwrap().unwrap(); assert_eq!(row.status.as_deref(), Some("pending")); } @@ -881,7 +919,7 @@ async fn create_acp_seeds_acp_session_runtime_from_extra() { let conv = svc.create(USER_ID, req).await.unwrap(); let runtime = acp_session_repo - .load_runtime_state(&conv.id) + .load_runtime_state_for_user(USER_ID, &conv.id) .await .unwrap() .expect("acp_session runtime state should exist after create"); @@ -927,7 +965,10 @@ async fn create_acp_skips_seed_when_extra_has_empty_runtime_fields() { .unwrap(); let conv = svc.create(USER_ID, req).await.unwrap(); - let runtime = acp_session_repo.load_runtime_state(&conv.id).await.unwrap(); + let runtime = acp_session_repo + .load_runtime_state_for_user(USER_ID, &conv.id) + .await + .unwrap(); // Either `None` (no runtime key yet) or Some(default) — both mean "nothing seeded". assert!( runtime diff --git a/crates/aionui-conversation/tests/conversation_extended.rs b/crates/aionui-conversation/tests/conversation_extended.rs index de6f9503f..ffbfb39a5 100644 --- a/crates/aionui-conversation/tests/conversation_extended.rs +++ b/crates/aionui-conversation/tests/conversation_extended.rs @@ -207,7 +207,7 @@ async fn t7_1_reset_clears_messages_and_status() { let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); // Insert messages for i in 0..3 { - repo.insert_message(&make_message(&conv.id, &format!("msg {i}"), i)) + repo.insert_message(USER_ID, &make_message(&conv.id, &format!("msg {i}"), i)) .await .unwrap(); } @@ -259,7 +259,7 @@ async fn t8_2_cursor_pagination() { let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); for i in 0..10 { - repo.insert_message(&make_message(&conv.id, &format!("msg {i}"), i * 100)) + repo.insert_message(USER_ID, &make_message(&conv.id, &format!("msg {i}"), i * 100)) .await .unwrap(); } @@ -299,7 +299,7 @@ async fn t8_3_asc_order_default() { let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); for i in 0..3 { - repo.insert_message(&make_message(&conv.id, &format!("msg {i}"), i * 1000)) + repo.insert_message(USER_ID, &make_message(&conv.id, &format!("msg {i}"), i * 1000)) .await .unwrap(); } @@ -319,7 +319,7 @@ async fn t8_4_asc_order() { let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); for i in 0..3 { - repo.insert_message(&make_message(&conv.id, &format!("msg {i}"), i * 1000)) + repo.insert_message(USER_ID, &make_message(&conv.id, &format!("msg {i}"), i * 1000)) .await .unwrap(); } @@ -419,7 +419,7 @@ async fn t8_9_anchor_returns_window_containing_target() { if i == 3 { target_id = msg.id.clone(); } - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(USER_ID, &msg).await.unwrap(); } let page = svc @@ -449,7 +449,7 @@ async fn t8_6_compact_mode_truncates_large_tool_content_only_for_list_response() let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); let large_output = "match line\n".repeat(10_000); - repo.insert_message(&make_acp_tool_message(&conv.id, "tool-big", &large_output, 0)) + repo.insert_message(USER_ID, &make_acp_tool_message(&conv.id, "tool-big", &large_output, 0)) .await .unwrap(); @@ -492,9 +492,12 @@ async fn t8_7_get_message_returns_full_tool_content_after_compact_list() { let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); let large_output = "wide rg output\n".repeat(10_000); - repo.insert_message(&make_acp_tool_message(&conv.id, "tool-detail", &large_output, 0)) - .await - .unwrap(); + repo.insert_message( + USER_ID, + &make_acp_tool_message(&conv.id, "tool-detail", &large_output, 0), + ) + .await + .unwrap(); let _ = svc .list_messages( @@ -537,10 +540,10 @@ async fn t9_1_keyword_match() { let workspace = ensure_test_workspace_path(); let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); - repo.insert_message(&make_message(&conv.id, "Rust review report", 0)) + repo.insert_message(USER_ID, &make_message(&conv.id, "Rust review report", 0)) .await .unwrap(); - repo.insert_message(&make_message(&conv.id, "Python test", 100)) + repo.insert_message(USER_ID, &make_message(&conv.id, "Python test", 100)) .await .unwrap(); @@ -567,7 +570,7 @@ async fn t9_1_keyword_match() { async fn t9_2_no_match() { let (svc, repo, _b) = setup().await; let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); - repo.insert_message(&make_message(&conv.id, "hello world", 0)) + repo.insert_message(USER_ID, &make_message(&conv.id, "hello world", 0)) .await .unwrap(); @@ -587,9 +590,12 @@ async fn t9_3_search_pagination() { let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); for i in 0..5 { - repo.insert_message(&make_message(&conv.id, &format!("match keyword item {i}"), i * 100)) - .await - .unwrap(); + repo.insert_message( + USER_ID, + &make_message(&conv.id, &format!("match keyword item {i}"), i * 100), + ) + .await + .unwrap(); } let query = SearchMessagesQuery { @@ -632,7 +638,7 @@ async fn t9_5_preview_text_extracts_from_json_content() { hidden: false, created_at: now_ms(), }; - repo.insert_message(&complex_msg).await.unwrap(); + repo.insert_message(USER_ID, &complex_msg).await.unwrap(); let query = SearchMessagesQuery { keyword: "search".into(), @@ -664,7 +670,7 @@ async fn t9_6_search_result_includes_conversation_model() { .unwrap(); let conv = svc.create(USER_ID, aionrs_req).await.unwrap(); - repo.insert_message(&make_message(&conv.id, "model test keyword", 0)) + repo.insert_message(USER_ID, &make_message(&conv.id, "model test keyword", 0)) .await .unwrap(); @@ -687,7 +693,7 @@ async fn t9_7_search_does_not_leak_other_users_messages() { let (svc, repo, _b) = setup().await; let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); - repo.insert_message(&make_message(&conv.id, "secret keyword data", 0)) + repo.insert_message(USER_ID, &make_message(&conv.id, "secret keyword data", 0)) .await .unwrap(); @@ -768,7 +774,7 @@ async fn t10_3_associated_not_found() { async fn t12_4_search_sql_injection() { let (svc, repo, _b) = setup().await; let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); - repo.insert_message(&make_message(&conv.id, "safe content", 0)) + repo.insert_message(USER_ID, &make_message(&conv.id, "safe content", 0)) .await .unwrap(); @@ -788,7 +794,9 @@ async fn t12_4_search_sql_injection() { async fn messages_wrong_user_returns_not_found() { let (svc, repo, _b) = setup().await; let conv = svc.create(USER_ID, make_create_req()).await.unwrap(); - repo.insert_message(&make_message(&conv.id, "hello", 0)).await.unwrap(); + repo.insert_message(USER_ID, &make_message(&conv.id, "hello", 0)) + .await + .unwrap(); let err = svc .list_messages("other_user", &conv.id, ListMessagesQuery::default()) diff --git a/crates/aionui-conversation/tests/stream_relay_tips.rs b/crates/aionui-conversation/tests/stream_relay_tips.rs index 926f61527..5cf3bab8a 100644 --- a/crates/aionui-conversation/tests/stream_relay_tips.rs +++ b/crates/aionui-conversation/tests/stream_relay_tips.rs @@ -68,6 +68,7 @@ async fn persist_info_tip_preserves_code_and_params() { let messages = repo .list_messages_page( + "system_default_user", "conv-1", &MessagePageParams { limit: 100, diff --git a/crates/aionui-conversation/tests/stream_relay_tool_call.rs b/crates/aionui-conversation/tests/stream_relay_tool_call.rs index 57da66b4d..1e0c2391b 100644 --- a/crates/aionui-conversation/tests/stream_relay_tool_call.rs +++ b/crates/aionui-conversation/tests/stream_relay_tool_call.rs @@ -100,6 +100,7 @@ async fn run_tool_call_with_empty_call_id_is_not_persisted() { let messages = repo .list_messages_page( + "system_default_user", "conv-1", &MessagePageParams { limit: 100, @@ -157,6 +158,7 @@ async fn run_tool_call_late_running_event_does_not_regress_completed_message() { let messages = repo .list_messages_page( + "system_default_user", "conv-1", &MessagePageParams { limit: 100, @@ -223,6 +225,7 @@ async fn run_tool_call_canceled_status_persists_as_terminal_finish() { let messages = repo .list_messages_page( + "system_default_user", "conv-1", &MessagePageParams { limit: 100, @@ -398,6 +401,7 @@ async fn run_agent_turn_with_empty_call_id_tool_call_is_not_persisted() { let messages = repo .list_messages_page( + "system_default_user", "conv-1", &MessagePageParams { limit: 100, diff --git a/crates/aionui-cron/Cargo.toml b/crates/aionui-cron/Cargo.toml index 1b5b47ee7..73255f3ea 100644 --- a/crates/aionui-cron/Cargo.toml +++ b/crates/aionui-cron/Cargo.toml @@ -29,6 +29,7 @@ async-trait.workspace = true tokio = { workspace = true, features = ["test-util"] } tempfile.workspace = true tower = { workspace = true, features = ["util"] } +sqlx.workspace = true # Enable the `AgentInstance::Mock` variant so tests can build fake agents # through the trait-object escape hatch without spawning real CLI processes. aionui-ai-agent = { workspace = true, features = ["test-support"] } diff --git a/crates/aionui-cron/src/artifacts.rs b/crates/aionui-cron/src/artifacts.rs index f0097f016..e688cf1d7 100644 --- a/crates/aionui-cron/src/artifacts.rs +++ b/crates/aionui-cron/src/artifacts.rs @@ -80,10 +80,12 @@ pub(crate) fn artifact_response_from_row( pub(crate) fn broadcast_artifact( broadcaster: &Arc, + user_id: &str, row: &ConversationArtifactRow, ) -> Result<(), CronError> { - let payload = serde_json::to_value(artifact_response_from_row(row)?) + let mut payload = serde_json::to_value(artifact_response_from_row(row)?) .map_err(|e| CronError::Scheduler(format!("failed to serialize artifact event: {e}")))?; + payload["user_id"] = serde_json::Value::String(user_id.to_owned()); broadcaster.broadcast(WebSocketMessage::new("conversation.artifact", payload)); Ok(()) } @@ -101,6 +103,7 @@ mod tests { fn sample_job() -> CronJob { CronJob { id: "cron_1".into(), + user_id: "user1".into(), name: "Daily Report".into(), enabled: true, schedule: CronSchedule::Every { diff --git a/crates/aionui-cron/src/error.rs b/crates/aionui-cron/src/error.rs index db62029a8..9b8907d69 100644 --- a/crates/aionui-cron/src/error.rs +++ b/crates/aionui-cron/src/error.rs @@ -29,6 +29,9 @@ pub enum CronError { #[error("Invalid agent config: {0}")] InvalidAgentConfig(String), + #[error("Cross-account reference: {0}")] + CrossAccountReference(String), + #[error("Scheduler error: {0}")] Scheduler(String), diff --git a/crates/aionui-cron/src/events.rs b/crates/aionui-cron/src/events.rs index 48d9701d4..3c78a665a 100644 --- a/crates/aionui-cron/src/events.rs +++ b/crates/aionui-cron/src/events.rs @@ -16,16 +16,17 @@ impl CronEventEmitter { Self { broadcaster } } - pub fn emit_job_created(&self, job: &CronJobResponse) { - self.broadcast("cron.job-created", job); + pub fn emit_job_created(&self, user_id: &str, job: &CronJobResponse) { + self.broadcast(user_id, "cron.job-created", job); } - pub fn emit_job_updated(&self, job: &CronJobResponse) { - self.broadcast("cron.job-updated", job); + pub fn emit_job_updated(&self, user_id: &str, job: &CronJobResponse) { + self.broadcast(user_id, "cron.job-updated", job); } - pub fn emit_job_removed(&self, job_id: &str) { + pub fn emit_job_removed(&self, user_id: &str, job_id: &str) { self.broadcast( + user_id, "cron.job-removed", &CronJobRemovedPayload { job_id: job_id.to_owned(), @@ -33,8 +34,9 @@ impl CronEventEmitter { ); } - pub fn emit_job_executed(&self, job_id: &str, status: &str, err: Option<&str>) { + pub fn emit_job_executed(&self, user_id: &str, job_id: &str, status: &str, err: Option<&str>) { self.broadcast( + user_id, "cron.job-executed", &CronJobExecutedEvent { job_id: job_id.to_owned(), @@ -44,8 +46,9 @@ impl CronEventEmitter { ); } - pub fn emit_conversation_tips(&self, conversation_id: &str, content: &str, tip_type: &str) { + pub fn emit_conversation_tips(&self, user_id: &str, conversation_id: &str, content: &str, tip_type: &str) { let payload = json!({ + "user_id": user_id, "conversation_id": conversation_id, "msg_id": ConversationService::mint_msg_id(), "type": "tips", @@ -59,14 +62,15 @@ impl CronEventEmitter { .broadcast(WebSocketMessage::new("message.stream", payload)); } - fn broadcast(&self, event_name: &str, payload: &T) { - let value = match serde_json::to_value(payload) { + fn broadcast(&self, user_id: &str, event_name: &str, payload: &T) { + let mut value = match serde_json::to_value(payload) { Ok(v) => v, Err(e) => { error!(event_name, error = %e, "Failed to serialize event payload"); return; } }; + value["user_id"] = serde_json::Value::String(user_id.to_owned()); self.broadcaster.broadcast(WebSocketMessage::new(event_name, value)); } } @@ -147,11 +151,12 @@ mod tests { fn job_created_event_shape() { let (emitter, bc) = make_emitter(); let resp = sample_response(); - emitter.emit_job_created(&resp); + emitter.emit_job_created("user-1", &resp); let events = bc.events(); assert_eq!(events.len(), 1); assert_eq!(events[0].name, "cron.job-created"); + assert_eq!(events[0].data["user_id"], "user-1"); let parsed: CronJobResponse = serde_json::from_value(events[0].data.clone()).unwrap(); assert_eq!(parsed.id, "cron_123"); @@ -162,11 +167,12 @@ mod tests { fn job_updated_event_shape() { let (emitter, bc) = make_emitter(); let resp = sample_response(); - emitter.emit_job_updated(&resp); + emitter.emit_job_updated("user-1", &resp); let events = bc.events(); assert_eq!(events.len(), 1); assert_eq!(events[0].name, "cron.job-updated"); + assert_eq!(events[0].data["user_id"], "user-1"); let parsed: CronJobResponse = serde_json::from_value(events[0].data.clone()).unwrap(); assert_eq!(parsed.id, "cron_123"); @@ -175,11 +181,12 @@ mod tests { #[test] fn job_removed_event_shape() { let (emitter, bc) = make_emitter(); - emitter.emit_job_removed("cron_456"); + emitter.emit_job_removed("user-1", "cron_456"); let events = bc.events(); assert_eq!(events.len(), 1); assert_eq!(events[0].name, "cron.job-removed"); + assert_eq!(events[0].data["user_id"], "user-1"); let parsed: CronJobRemovedPayload = serde_json::from_value(events[0].data.clone()).unwrap(); assert_eq!(parsed.job_id, "cron_456"); @@ -188,11 +195,12 @@ mod tests { #[test] fn job_executed_success_event() { let (emitter, bc) = make_emitter(); - emitter.emit_job_executed("cron_789", "ok", None); + emitter.emit_job_executed("user-1", "cron_789", "ok", None); let events = bc.events(); assert_eq!(events.len(), 1); assert_eq!(events[0].name, "cron.job-executed"); + assert_eq!(events[0].data["user_id"], "user-1"); let parsed: CronJobExecutedEvent = serde_json::from_value(events[0].data.clone()).unwrap(); assert_eq!(parsed.job_id, "cron_789"); @@ -203,7 +211,7 @@ mod tests { #[test] fn job_executed_error_event() { let (emitter, bc) = make_emitter(); - emitter.emit_job_executed("cron_789", "error", Some("timeout")); + emitter.emit_job_executed("user-1", "cron_789", "error", Some("timeout")); let events = bc.events(); assert_eq!(events.len(), 1); @@ -216,7 +224,7 @@ mod tests { #[test] fn job_executed_skipped_event() { let (emitter, bc) = make_emitter(); - emitter.emit_job_executed("cron_789", "skipped", None); + emitter.emit_job_executed("user-1", "cron_789", "skipped", None); let events = bc.events(); let parsed: CronJobExecutedEvent = serde_json::from_value(events[0].data.clone()).unwrap(); @@ -227,10 +235,10 @@ mod tests { fn multiple_events_accumulate() { let (emitter, bc) = make_emitter(); let resp = sample_response(); - emitter.emit_job_created(&resp); - emitter.emit_job_updated(&resp); - emitter.emit_job_removed("cron_123"); - emitter.emit_job_executed("cron_123", "ok", None); + emitter.emit_job_created("user-1", &resp); + emitter.emit_job_updated("user-1", &resp); + emitter.emit_job_removed("user-1", "cron_123"); + emitter.emit_job_executed("user-1", "cron_123", "ok", None); let events = bc.events(); assert_eq!(events.len(), 4); diff --git a/crates/aionui-cron/src/executor.rs b/crates/aionui-cron/src/executor.rs index 5dfc92d60..4152576ff 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -28,7 +28,6 @@ use crate::types::{CronJob, ExecutionMode}; pub const RETRY_INTERVAL_MS: u64 = 30_000; pub const MAX_RETRIES_DEFAULT: i64 = 3; -const SYSTEM_DEFAULT_USER_ID: &str = "system_default_user"; const DEPRECATED_AGENT_TYPE_MESSAGE: &str = "This agent type is no longer supported for new conversations."; #[derive(Debug, Clone, PartialEq, Eq)] pub enum ExecutionResult { @@ -50,6 +49,7 @@ pub(crate) enum PreparedRunNow { } struct SkillSuggestContext { + user_id: String, conversation_id: String, job_id: String, workspace: String, @@ -212,7 +212,7 @@ impl JobExecutor { pub async fn conversation_exists(&self, conversation_id: &str) -> Result { let row = self .conversation_repo - .get(conversation_id) + .owner_user_id(conversation_id) .await .map_err(CronError::Database)?; Ok(row.is_some()) @@ -226,17 +226,37 @@ impl JobExecutor { &self, conversation_id: &str, ) -> Result, CronError> { + let Some(user_id) = self + .conversation_repo + .owner_user_id(conversation_id) + .await + .map_err(CronError::Database)? + else { + return Ok(None); + }; self.conversation_repo - .get(conversation_id) + .get(&user_id, conversation_id) + .await + .map_err(CronError::Database) + } + + pub async fn get_conversation_row_for_user( + &self, + user_id: &str, + conversation_id: &str, + ) -> Result, CronError> { + self.conversation_repo + .get(user_id, conversation_id) .await .map_err(CronError::Database) } pub async fn auto_workspace_to_delete_for_conversation( &self, + user_id: &str, conversation_id: &str, ) -> Result, CronError> { - let Some(row) = self.get_conversation_row(conversation_id).await? else { + let Some(row) = self.get_conversation_row_for_user(user_id, conversation_id).await? else { return Ok(None); }; Ok(self @@ -248,8 +268,16 @@ impl JobExecutor { &self, conversation_id: &str, ) -> Result, CronError> { + let Some(user_id) = self + .conversation_repo + .owner_user_id(conversation_id) + .await + .map_err(CronError::Database)? + else { + return Ok(None); + }; self.conversation_repo - .get_assistant_snapshot(conversation_id) + .get_assistant_snapshot(&user_id, conversation_id) .await .map_err(CronError::Database) } @@ -298,8 +326,9 @@ impl JobExecutor { created_at: aionui_common::now_ms(), }; + let user_id = self.resolve_target_conversation_user_id(conversation_id).await?; self.conversation_repo - .insert_message(&row) + .insert_message(&user_id, &row) .await .map_err(CronError::Database) } @@ -348,14 +377,14 @@ impl JobExecutor { ..Default::default() }; self.conversation_repo - .update(conversation_id, &update) + .update(&row.user_id, conversation_id, &update) .await .map_err(CronError::Database)?; } if is_acp_conversation && let Some(mode) = session_mode { self.conversation_service - .save_acp_runtime_mode(conversation_id, mode) + .save_acp_runtime_mode(&row.user_id, conversation_id, mode) .await?; } @@ -395,7 +424,7 @@ impl JobExecutor { ..Default::default() }; self.conversation_repo - .update(conversation_id, &update) + .update(&row.user_id, conversation_id, &update) .await .map_err(CronError::Database) } @@ -428,7 +457,7 @@ impl JobExecutor { }; return self .conversation_repo - .update(conversation_id, &update) + .update(&row.user_id, conversation_id, &update) .await .map_err(CronError::Database); }; @@ -454,7 +483,7 @@ impl JobExecutor { ..Default::default() }; self.conversation_repo - .update(conversation_id, &update) + .update(&row.user_id, conversation_id, &update) .await .map_err(CronError::Database) } @@ -508,7 +537,12 @@ impl JobExecutor { .create_new_conversation(job, saved_skill, ConversationPurpose::ExistingReplacement) .await; } - if !self.conversation_exists(conversation_id).await? { + let Some(owner_user_id) = self + .conversation_repo + .owner_user_id(conversation_id) + .await + .map_err(CronError::Database)? + else { warn!( job_id = %job.id, conversation_id, @@ -517,6 +551,19 @@ impl JobExecutor { return self .create_new_conversation(job, saved_skill, ConversationPurpose::ExistingReplacement) .await; + }; + if !job.user_id.trim().is_empty() && owner_user_id != job.user_id { + warn!( + job_id = %job.id, + job_user_id = %job.user_id, + conversation_id, + conversation_user_id = %owner_user_id, + "Cron existing-mode conversation owner mismatch; refusing to dispatch" + ); + return Err(CronError::Scheduler(format!( + "cron job {} targets conversation {} owned by another user", + job.id, conversation_id + ))); } Ok(job.conversation_id.clone()) } @@ -542,7 +589,7 @@ impl JobExecutor { saved_skill: Option<&SavedSkillContext>, purpose: ConversationPurpose, ) -> Result { - let agent_type = parse_agent_type(&self.agent_registry, &job.agent_type).await?; + let agent_type = parse_agent_type(&self.agent_registry, &job.user_id, &job.agent_type).await?; let model = resolve_model(job); let user_id = self.resolve_conversation_owner_user_id(job).await?; @@ -594,14 +641,32 @@ impl JobExecutor { } async fn resolve_conversation_owner_user_id(&self, job: &CronJob) -> Result { + if !job.user_id.trim().is_empty() { + return Ok(job.user_id.clone()); + } + + if job.conversation_id.trim().is_empty() { + return Err(CronError::Scheduler(format!( + "cron job {} has no conversation owner", + job.id + ))); + } + if !job.conversation_id.trim().is_empty() - && let Some(row) = self.get_conversation_row(&job.conversation_id).await? - && !row.user_id.trim().is_empty() + && let Some(user_id) = self + .conversation_repo + .owner_user_id(&job.conversation_id) + .await + .map_err(CronError::Database)? + && !user_id.trim().is_empty() { - return Ok(row.user_id); + return Ok(user_id); } - Ok(SYSTEM_DEFAULT_USER_ID.to_owned()) + Err(CronError::Scheduler(format!( + "cron job {} conversation {} has no resolvable owner", + job.id, job.conversation_id + ))) } async fn execute_inner_with_busy_retry( @@ -653,18 +718,20 @@ impl JobExecutor { let on_started = { let job = job.clone(); + let user_id = user_id.clone(); let conversation_repo = Arc::clone(&self.conversation_repo); let broadcaster = Arc::clone(&self.broadcaster); Some(Arc::new(move |started: ConversationAgentTurnStarted| { let job = job.clone(); + let user_id = user_id.clone(); let conversation_repo = Arc::clone(&conversation_repo); let broadcaster = Arc::clone(&broadcaster); let fut: std::pin::Pin + Send>> = Box::pin(async move { let created_at = now_ms(); let row = build_cron_trigger_artifact(&started.conversation_id, &job, created_at); - match conversation_repo.upsert_artifact(&row).await { + match conversation_repo.upsert_artifact(&user_id, &row).await { Ok(row) => { - if let Err(e) = broadcast_artifact(&broadcaster, &row) { + if let Err(e) = broadcast_artifact(&broadcaster, &user_id, &row) { warn!( job_id = %job.id, conversation_id = %started.conversation_id, @@ -688,7 +755,7 @@ impl JobExecutor { }; let turn_req = ConversationAgentTurnRequest { - user_id, + user_id: user_id.clone(), conversation_id: conversation_id.to_owned(), content: prompt, files: vec![], @@ -703,7 +770,9 @@ impl JobExecutor { Ok(outcome) => { if outcome.status == ConversationAgentTurnStatus::Failed { return ExecutionResult::Error { - message: self.conversation_turn_failed_message(conversation_id, outcome).await, + message: self + .conversation_turn_failed_message(&user_id, conversation_id, outcome) + .await, }; } if saved_skill.is_none() && matches!(job.execution_mode, ExecutionMode::NewConversation) { @@ -720,6 +789,7 @@ impl JobExecutor { String::new() }); self.schedule_skill_suggest_check(SkillSuggestContext { + user_id: user_id.clone(), conversation_id: conversation_id.to_owned(), job_id: job.id.clone(), workspace, @@ -757,6 +827,7 @@ impl JobExecutor { async fn conversation_turn_failed_message( &self, + user_id: &str, conversation_id: &str, outcome: ConversationAgentTurnOutcome, ) -> String { @@ -771,7 +842,7 @@ impl JobExecutor { match self .conversation_service - .latest_conversation_error_message(conversation_id) + .latest_conversation_error_message(user_id, conversation_id) .await { Ok(Some(message)) => message, @@ -802,21 +873,23 @@ impl JobExecutor { }; if row.user_id.trim().is_empty() { - return Ok(SYSTEM_DEFAULT_USER_ID.to_owned()); + return Err(CronError::Scheduler(format!( + "conversation {conversation_id} has no owner" + ))); } Ok(row.user_id) } - pub async fn mark_skill_suggest_artifacts_saved(&self, job_id: &str) -> Result<(), CronError> { + pub async fn mark_skill_suggest_artifacts_saved(&self, user_id: &str, job_id: &str) -> Result<(), CronError> { let rows = self .conversation_repo - .mark_skill_suggest_artifacts_saved(job_id, now_ms()) + .mark_skill_suggest_artifacts_saved(user_id, job_id, now_ms()) .await .map_err(CronError::Database)?; for row in rows { - broadcast_artifact(&self.broadcaster, &row)?; + broadcast_artifact(&self.broadcaster, user_id, &row)?; } Ok(()) @@ -849,7 +922,7 @@ impl JobExecutor { fn schedule_skill_suggest_check(&self, ctx: SkillSuggestContext) { self.skill_suggest_detector - .schedule_check(ctx.conversation_id, ctx.job_id, ctx.workspace); + .schedule_check(ctx.user_id, ctx.conversation_id, ctx.job_id, ctx.workspace); } async fn prepare_saved_skill(&self, job: &CronJob) -> Result, CronError> { @@ -907,7 +980,10 @@ impl JobExecutor { async fn load_conversation_skill_names(&self, conversation_id: &str) -> Result, CronError> { let Some(row) = self .conversation_repo - .get(conversation_id) + .get( + &self.resolve_target_conversation_user_id(conversation_id).await?, + conversation_id, + ) .await .map_err(CronError::Database)? else { @@ -941,7 +1017,11 @@ impl JobExecutor { /// 2. ACP vendor lookup via the registry — any builtin ACP row's /// `backend` aliases to [`AgentType::Acp`]. /// 3. Fallback to [`AgentType::Acp`] to preserve the prior default. -async fn parse_agent_type(registry: &AgentRegistry, agent_type_str: &str) -> Result { +async fn parse_agent_type( + registry: &AgentRegistry, + user_id: &str, + agent_type_str: &str, +) -> Result { if let Ok(agent_type) = serde_json::from_value::(serde_json::Value::String(agent_type_str.to_owned())) { if agent_type.is_deprecated_runtime() { return Err(CronError::InvalidAgentConfig(DEPRECATED_AGENT_TYPE_MESSAGE.into())); @@ -949,7 +1029,11 @@ async fn parse_agent_type(registry: &AgentRegistry, agent_type_str: &str) -> Res return Ok(agent_type); } - if registry.find_builtin_by_backend(agent_type_str).await.is_some() { + if registry + .find_builtin_by_backend_for_user(user_id, agent_type_str) + .await + .is_some() + { return Ok(AgentType::Acp); } @@ -984,7 +1068,10 @@ async fn inject_agent_identity( return; } - if let Some(meta) = registry.find_builtin_by_backend(lookup_label).await { + if let Some(meta) = registry + .find_builtin_by_backend_for_user(&job.user_id, lookup_label) + .await + { extra.insert("agent_id".to_owned(), serde_json::Value::String(meta.id.clone())); if let Some(backend) = meta.backend { extra.insert("backend".to_owned(), serde_json::Value::String(backend)); @@ -1240,6 +1327,7 @@ mod tests { fn sample_job() -> CronJob { CronJob { id: "cron_test1".into(), + user_id: "user1".into(), name: "Test Job".into(), enabled: true, schedule: CronSchedule::Every { @@ -1403,6 +1491,28 @@ mod tests { drop(claim); } + #[tokio::test] + async fn prepare_scheduled_rejects_cross_user_existing_conversation() { + let agent = Arc::new(RecordingAgent::new("conv_1", "default", true)); + let task_manager = Arc::new(RecordingTaskManager::new(AgentInstance::Mock(agent.clone()))); + let repo = Arc::new(MissingWorkspaceConversationRepo::new( + "conv_1", + serde_json::json!({ + "workspace": ensure_named_workspace_path("aionui-cron-cross-user-conversation-workspace") + }), + )); + let executor = make_executor_with_task_manager_and_repo(task_manager, repo); + let result = executor.prepare_scheduled(&sample_job()).await.unwrap_err(); + + match result { + ExecutionResult::Error { message } => { + assert!(message.contains("owned by another user"), "unexpected error: {message}"); + } + other => panic!("expected owner mismatch error, got {other:?}"), + } + assert_eq!(agent.send_calls(), 0, "owner mismatch must not dispatch"); + } + #[tokio::test] async fn prepare_run_now_returns_active_conversation_when_runtime_state_is_already_claimed() { let agent = Arc::new(RecordingAgent::new("conv_1", "default", true)); @@ -1508,15 +1618,21 @@ mod tests { #[tokio::test] async fn parse_agent_type_known_types() { let registry = hydrated_registry().await; - assert_eq!(parse_agent_type(®istry, "acp").await.unwrap(), AgentType::Acp); - assert_eq!(parse_agent_type(®istry, "aionrs").await.unwrap(), AgentType::Aionrs); + assert_eq!( + parse_agent_type(®istry, "user1", "acp").await.unwrap(), + AgentType::Acp + ); + assert_eq!( + parse_agent_type(®istry, "user1", "aionrs").await.unwrap(), + AgentType::Aionrs + ); } #[tokio::test] async fn parse_agent_type_rejects_deprecated_runtime_types() { let registry = hydrated_registry().await; for agent_type in ["openclaw-gateway", "nanobot", "remote", "gemini", "codex"] { - let err = parse_agent_type(®istry, agent_type).await.unwrap_err(); + let err = parse_agent_type(®istry, "user1", agent_type).await.unwrap_err(); assert!(matches!(err, CronError::InvalidAgentConfig(_))); assert!( err.to_string() @@ -1529,15 +1645,21 @@ mod tests { #[tokio::test] async fn parse_agent_type_acp_backend_aliases_to_acp() { let registry = hydrated_registry().await; - assert_eq!(parse_agent_type(®istry, "claude").await.unwrap(), AgentType::Acp); - assert_eq!(parse_agent_type(®istry, "qwen").await.unwrap(), AgentType::Acp); + assert_eq!( + parse_agent_type(®istry, "user1", "claude").await.unwrap(), + AgentType::Acp + ); + assert_eq!( + parse_agent_type(®istry, "user1", "qwen").await.unwrap(), + AgentType::Acp + ); } #[tokio::test] async fn parse_agent_type_unknown_defaults_to_acp() { let registry = hydrated_registry().await; assert_eq!( - parse_agent_type(®istry, "unknown_type").await.unwrap(), + parse_agent_type(®istry, "user1", "unknown_type").await.unwrap(), AgentType::Acp ); } @@ -2360,16 +2482,28 @@ mod tests { #[async_trait::async_trait] impl IConversationRepository for StubConvRepo { - async fn get(&self, _id: &str) -> Result, aionui_db::DbError> { + async fn get( + &self, + _user_id: &str, + _id: &str, + ) -> Result, aionui_db::DbError> { + Ok(None) + } + async fn owner_user_id(&self, _id: &str) -> Result, aionui_db::DbError> { Ok(None) } async fn create(&self, _row: &aionui_db::models::ConversationRow) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn update(&self, _id: &str, _updates: &ConversationRowUpdate) -> Result<(), aionui_db::DbError> { + async fn update( + &self, + _user_id: &str, + _id: &str, + _updates: &ConversationRowUpdate, + ) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn delete(&self, _id: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), aionui_db::DbError> { Ok(()) } async fn list_paginated( @@ -2408,6 +2542,7 @@ mod tests { } async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &MessagePageParams, ) -> Result { @@ -2417,17 +2552,32 @@ mod tests { has_more_after: false, }) } - async fn insert_message(&self, _message: &aionui_db::models::MessageRow) -> Result<(), aionui_db::DbError> { + async fn insert_message( + &self, + _user_id: &str, + _message: &aionui_db::models::MessageRow, + ) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn update_message(&self, _id: &str, _updates: &MessageRowUpdate) -> Result<(), aionui_db::DbError> { + async fn update_message( + &self, + _user_id: &str, + _conversation_id: &str, + _id: &str, + _updates: &MessageRowUpdate, + ) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn delete_messages_by_conversation(&self, _conv_id: &str) -> Result<(), aionui_db::DbError> { + async fn delete_messages_by_conversation( + &self, + _user_id: &str, + _conv_id: &str, + ) -> Result<(), aionui_db::DbError> { Ok(()) } async fn get_message_by_msg_id( &self, + _user_id: &str, _conv_id: &str, _msg_id: &str, _msg_type: &str, @@ -2797,10 +2947,14 @@ mod tests { #[async_trait::async_trait] impl IConversationRepository for ExistingConversationRepo { - async fn get(&self, id: &str) -> Result, aionui_db::DbError> { + async fn get( + &self, + _user_id: &str, + id: &str, + ) -> Result, aionui_db::DbError> { Ok(Some(aionui_db::models::ConversationRow { id: id.to_owned(), - user_id: "cron".into(), + user_id: "user1".into(), name: "Cron Conversation".into(), r#type: "acp".into(), extra: serde_json::json!({ @@ -2818,15 +2972,24 @@ mod tests { })) } + async fn owner_user_id(&self, _id: &str) -> Result, aionui_db::DbError> { + Ok(Some("user1".into())) + } + async fn create(&self, _row: &aionui_db::models::ConversationRow) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn update(&self, _id: &str, _updates: &ConversationRowUpdate) -> Result<(), aionui_db::DbError> { + async fn update( + &self, + _user_id: &str, + _id: &str, + _updates: &ConversationRowUpdate, + ) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn delete(&self, _id: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), aionui_db::DbError> { Ok(()) } @@ -2870,6 +3033,7 @@ mod tests { async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &MessagePageParams, ) -> Result { @@ -2880,20 +3044,35 @@ mod tests { }) } - async fn insert_message(&self, _message: &aionui_db::models::MessageRow) -> Result<(), aionui_db::DbError> { + async fn insert_message( + &self, + _user_id: &str, + _message: &aionui_db::models::MessageRow, + ) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn update_message(&self, _id: &str, _updates: &MessageRowUpdate) -> Result<(), aionui_db::DbError> { + async fn update_message( + &self, + _user_id: &str, + _conversation_id: &str, + _id: &str, + _updates: &MessageRowUpdate, + ) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn delete_messages_by_conversation(&self, _conv_id: &str) -> Result<(), aionui_db::DbError> { + async fn delete_messages_by_conversation( + &self, + _user_id: &str, + _conv_id: &str, + ) -> Result<(), aionui_db::DbError> { Ok(()) } async fn get_message_by_msg_id( &self, + _user_id: &str, _conv_id: &str, _msg_id: &str, _msg_type: &str, @@ -3001,20 +3180,33 @@ mod tests { #[async_trait::async_trait] impl IConversationRepository for MissingWorkspaceConversationRepo { - async fn get(&self, _id: &str) -> Result, aionui_db::DbError> { + async fn get( + &self, + _user_id: &str, + _id: &str, + ) -> Result, aionui_db::DbError> { Ok(Some(self.row.clone())) } + async fn owner_user_id(&self, _id: &str) -> Result, aionui_db::DbError> { + Ok(Some(self.row.user_id.clone())) + } + async fn create(&self, _row: &aionui_db::models::ConversationRow) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn update(&self, _id: &str, updates: &ConversationRowUpdate) -> Result<(), aionui_db::DbError> { + async fn update( + &self, + _user_id: &str, + _id: &str, + updates: &ConversationRowUpdate, + ) -> Result<(), aionui_db::DbError> { self.updates.lock().unwrap().push(updates.clone()); Ok(()) } - async fn delete(&self, _id: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), aionui_db::DbError> { Ok(()) } @@ -3058,6 +3250,7 @@ mod tests { async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &MessagePageParams, ) -> Result { @@ -3068,7 +3261,11 @@ mod tests { }) } - async fn insert_message(&self, message: &aionui_db::models::MessageRow) -> Result<(), aionui_db::DbError> { + async fn insert_message( + &self, + _user_id: &str, + message: &aionui_db::models::MessageRow, + ) -> Result<(), aionui_db::DbError> { self.operations.lock().unwrap().push(format!( "insert_message:{}:{}", message.position.as_deref().unwrap_or("none"), @@ -3078,16 +3275,27 @@ mod tests { Ok(()) } - async fn update_message(&self, _id: &str, _updates: &MessageRowUpdate) -> Result<(), aionui_db::DbError> { + async fn update_message( + &self, + _user_id: &str, + _conversation_id: &str, + _id: &str, + _updates: &MessageRowUpdate, + ) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn delete_messages_by_conversation(&self, _conv_id: &str) -> Result<(), aionui_db::DbError> { + async fn delete_messages_by_conversation( + &self, + _user_id: &str, + _conv_id: &str, + ) -> Result<(), aionui_db::DbError> { Ok(()) } async fn get_message_by_msg_id( &self, + _user_id: &str, _conv_id: &str, _msg_id: &str, _msg_type: &str, @@ -3111,6 +3319,7 @@ mod tests { async fn upsert_artifact( &self, + _user_id: &str, artifact: &ConversationArtifactRow, ) -> Result { self.operations @@ -3202,8 +3411,9 @@ mod tests { #[async_trait::async_trait] impl aionui_db::IAcpSessionRepository for StubAcpSessionRepo { - async fn get( + async fn get_for_user( &self, + _user_id: &str, _conversation_id: &str, ) -> Result, aionui_db::DbError> { Ok(None) @@ -3214,27 +3424,34 @@ mod tests { ) -> Result { Err(aionui_db::DbError::Init("stub".into())) } - async fn update_session_id( + async fn update_session_id_for_user( &self, + _user_id: &str, _conversation_id: &str, _session_id: &str, ) -> Result { Ok(false) } - async fn clear_session_id(&self, _conversation_id: &str) -> Result { + async fn clear_session_id_for_user( + &self, + _user_id: &str, + _conversation_id: &str, + ) -> Result { Ok(false) } - async fn delete(&self, _conversation_id: &str) -> Result { + async fn delete_for_user(&self, _user_id: &str, _conversation_id: &str) -> Result { Ok(false) } - async fn load_runtime_state( + async fn load_runtime_state_for_user( &self, + _user_id: &str, _conversation_id: &str, ) -> Result, aionui_db::DbError> { Ok(None) } - async fn save_runtime_state( + async fn save_runtime_state_for_user( &self, + _user_id: &str, _conversation_id: &str, _params: &aionui_db::SaveRuntimeStateParams<'_>, ) -> Result { @@ -3249,9 +3466,22 @@ mod tests { async fn list_all(&self) -> Result, aionui_db::DbError> { Ok(Vec::new()) } + async fn list_all_for_user( + &self, + _user_id: &str, + ) -> Result, aionui_db::DbError> { + self.list_all().await + } async fn get(&self, _id: &str) -> Result, aionui_db::DbError> { Ok(None) } + async fn get_for_user( + &self, + _user_id: &str, + id: &str, + ) -> Result, aionui_db::DbError> { + self.get(id).await + } async fn find_by_source_and_name( &self, _agent_source: &str, @@ -3259,18 +3489,40 @@ mod tests { ) -> Result, aionui_db::DbError> { Ok(None) } + async fn find_by_source_and_name_for_user( + &self, + _user_id: &str, + agent_source: &str, + name: &str, + ) -> Result, aionui_db::DbError> { + self.find_by_source_and_name(agent_source, name).await + } async fn find_builtin_by_backend( &self, _backend: &str, ) -> Result, aionui_db::DbError> { Ok(None) } + async fn find_builtin_by_backend_for_user( + &self, + _user_id: &str, + backend: &str, + ) -> Result, aionui_db::DbError> { + self.find_builtin_by_backend(backend).await + } async fn upsert( &self, _params: &aionui_db::models::UpsertAgentMetadataParams<'_>, ) -> Result { Err(aionui_db::DbError::Init("stub".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &aionui_db::models::UpsertAgentMetadataParams<'_>, + ) -> Result { + self.upsert(params).await + } async fn apply_handshake( &self, _id: &str, @@ -3278,6 +3530,14 @@ mod tests { ) -> Result, aionui_db::DbError> { Ok(None) } + async fn apply_handshake_for_user( + &self, + _user_id: &str, + id: &str, + params: &aionui_db::models::UpdateAgentHandshakeParams<'_>, + ) -> Result, aionui_db::DbError> { + self.apply_handshake(id, params).await + } async fn update_availability_snapshot( &self, _id: &str, @@ -3285,6 +3545,14 @@ mod tests { ) -> Result, aionui_db::DbError> { Ok(None) } + async fn update_availability_snapshot_for_user( + &self, + _user_id: &str, + id: &str, + params: &aionui_db::models::UpdateAgentAvailabilitySnapshotParams<'_>, + ) -> Result, aionui_db::DbError> { + self.update_availability_snapshot(id, params).await + } async fn update_agent_overrides( &self, _id: &str, @@ -3293,11 +3561,31 @@ mod tests { ) -> Result<(), aionui_db::DbError> { Ok(()) } + async fn update_agent_overrides_for_user( + &self, + _user_id: &str, + id: &str, + command_override: Option<&str>, + env_override: Option<&str>, + ) -> Result<(), aionui_db::DbError> { + self.update_agent_overrides(id, command_override, env_override).await + } async fn set_enabled(&self, _id: &str, _enabled: bool) -> Result { Ok(false) } + async fn set_enabled_for_user( + &self, + _user_id: &str, + id: &str, + enabled: bool, + ) -> Result { + self.set_enabled(id, enabled).await + } async fn delete(&self, _id: &str) -> Result { Ok(false) } + async fn delete_for_user(&self, _user_id: &str, id: &str) -> Result { + self.delete(id).await + } } } diff --git a/crates/aionui-cron/src/routes.rs b/crates/aionui-cron/src/routes.rs index 5b7ff0f9f..6e2ff2d7f 100644 --- a/crates/aionui-cron/src/routes.rs +++ b/crates/aionui-cron/src/routes.rs @@ -41,6 +41,9 @@ impl From for ApiError { CronError::InvalidTimezone(msg) => ApiError::BadRequest(msg), CronError::InvalidSkillContent(msg) => ApiError::BadRequest(msg), CronError::InvalidAgentConfig(msg) => ApiError::BadRequest(msg), + CronError::CrossAccountReference(msg) => { + ApiError::coded(StatusCode::CONFLICT, "CROSS_ACCOUNT_REFERENCE", msg, None) + } CronError::Scheduler(msg) => ApiError::Internal(msg), CronError::WorkspacePathUnavailable(path) => ApiError::WorkspacePathUnavailable(path), CronError::WorkspacePathRuntimeUnavailable(path) => ApiError::WorkspacePathRuntimeUnavailable(path), @@ -73,60 +76,60 @@ pub fn cron_routes(state: CronRouterState) -> Router { async fn create_job( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result<(StatusCode, Json>), ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let job = state.cron_service.add_job(req).await?; + let job = state.cron_service.add_job(&user.id, req).await?; let resp = CronService::to_response(&job); Ok((StatusCode::CREATED, Json(ApiResponse::ok(resp)))) } async fn list_jobs( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Query(query): Query, ) -> Result>>, ApiError> { - let jobs = state.cron_service.list_jobs(&query).await?; + let jobs = state.cron_service.list_jobs(&user.id, &query).await?; let items: Vec = jobs.iter().map(CronService::to_response).collect(); Ok(Json(ApiResponse::ok(items))) } async fn get_job( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - let job = state.cron_service.get_job(&id).await?; + let job = state.cron_service.get_job(&user.id, &id).await?; Ok(Json(ApiResponse::ok(CronService::to_response(&job)))) } async fn update_job( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let job = state.cron_service.update_job(&id, req).await?; + let job = state.cron_service.update_job(&user.id, &id, req).await?; Ok(Json(ApiResponse::ok(CronService::to_response(&job)))) } async fn delete_job( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - state.cron_service.remove_job(&id).await?; + state.cron_service.remove_job(&user.id, &id).await?; Ok(Json(ApiResponse::success())) } async fn run_now( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - let resp = state.cron_service.run_now(&id).await?; + let resp = state.cron_service.run_now(&user.id, &id).await?; Ok(Json(ApiResponse::ok(resp))) } @@ -145,10 +148,11 @@ async fn system_resume( async fn create_conversation_cron( State(state): State, + Extension(current_user): Extension, headers: HeaderMap, body: Result, JsonRejection>, ) -> Result>, ApiError> { - let user_id = header_value(&headers, "x-aionui-user-id")?; + let user_id = trusted_header_user_id(&headers, ¤t_user)?; let conversation_id = header_value(&headers, "x-aionui-conversation-id")?; let Json(req) = body.map_err(ApiError::from)?; let resp = state @@ -160,9 +164,10 @@ async fn create_conversation_cron( async fn list_conversation_cron( State(state): State, + Extension(current_user): Extension, headers: HeaderMap, ) -> Result>>, ApiError> { - let user_id = header_value(&headers, "x-aionui-user-id")?; + let user_id = trusted_header_user_id(&headers, ¤t_user)?; let conversation_id = header_value(&headers, "x-aionui-conversation-id")?; let jobs = state .cron_service @@ -174,11 +179,12 @@ async fn list_conversation_cron( async fn update_conversation_cron( State(state): State, + Extension(current_user): Extension, headers: HeaderMap, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { - let user_id = header_value(&headers, "x-aionui-user-id")?; + let user_id = trusted_header_user_id(&headers, ¤t_user)?; let conversation_id = header_value(&headers, "x-aionui-conversation-id")?; let Json(req) = body.map_err(ApiError::from)?; let job = state @@ -198,14 +204,27 @@ fn header_value(headers: &HeaderMap, name: &'static str) -> Result Result { + let header_user_id = header_value(headers, "x-aionui-user-id")?; + if header_user_id != current_user.id { + return Err(ApiError::coded( + StatusCode::FORBIDDEN, + "CRON_HELPER_USER_MISMATCH", + "Authenticated user does not match cron helper user header.", + None, + )); + } + Ok(header_user_id) +} + async fn save_skill( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - state.cron_service.save_skill(&id, req).await?; + state.cron_service.save_skill(&user.id, &id, req).await?; Ok(Json(ApiResponse::success())) } @@ -220,19 +239,19 @@ async fn list_conversations_by_cron_job( async fn has_skill( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - let resp = state.cron_service.has_skill(&id).await?; + let resp = state.cron_service.has_skill(&user.id, &id).await?; Ok(Json(ApiResponse::ok(resp))) } async fn delete_skill( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - state.cron_service.delete_skill(&id).await?; + state.cron_service.delete_skill(&user.id, &id).await?; Ok(Json(ApiResponse::success())) } @@ -265,6 +284,40 @@ mod tests { assert_eq!(header_value(&headers, "x-aionui-user-id").unwrap(), "user_1"); } + #[test] + fn trusted_header_user_id_rejects_user_mismatch() { + let mut headers = HeaderMap::new(); + headers.insert("x-aionui-user-id", "user_b".parse().unwrap()); + let current_user = CurrentUser { + id: "user_a".to_owned(), + username: "alice".to_owned(), + user_type: aionui_db::UserType::Aionpro, + status: aionui_db::UserStatus::Active, + }; + + assert!(matches!( + trusted_header_user_id(&headers, ¤t_user), + Err(ApiError::Coded { + code: "CRON_HELPER_USER_MISMATCH", + .. + }) + )); + } + + #[test] + fn trusted_header_user_id_accepts_authenticated_user() { + let mut headers = HeaderMap::new(); + headers.insert("x-aionui-user-id", "user_a".parse().unwrap()); + let current_user = CurrentUser { + id: "user_a".to_owned(), + username: "alice".to_owned(), + user_type: aionui_db::UserType::Aionpro, + status: aionui_db::UserStatus::Active, + }; + + assert_eq!(trusted_header_user_id(&headers, ¤t_user).unwrap(), "user_a"); + } + #[test] fn job_not_found_maps_to_not_found() { let err: ApiError = CronError::JobNotFound("cron_abc".into()).into(); @@ -319,6 +372,13 @@ mod tests { assert!(matches!(err, ApiError::BadRequest(_))); } + #[test] + fn cross_account_reference_maps_to_stable_conflict_code() { + let err: ApiError = CronError::CrossAccountReference("cross account".into()).into(); + assert_eq!(err.status_code(), StatusCode::CONFLICT); + assert_eq!(err.error_code(), "CROSS_ACCOUNT_REFERENCE"); + } + #[test] fn scheduler_error_maps_to_internal() { let err: ApiError = CronError::Scheduler("timer failed".into()).into(); diff --git a/crates/aionui-cron/src/scheduler.rs b/crates/aionui-cron/src/scheduler.rs index f0e409df1..17996c042 100644 --- a/crates/aionui-cron/src/scheduler.rs +++ b/crates/aionui-cron/src/scheduler.rs @@ -691,6 +691,7 @@ mod tests { use crate::types::{CreatedBy, ExecutionMode}; CronJob { id: id.to_owned(), + user_id: "user1".into(), name: "Test".into(), enabled, schedule: CronSchedule::Every { diff --git a/crates/aionui-cron/src/service.rs b/crates/aionui-cron/src/service.rs index 69292314c..4d04257c5 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -12,9 +12,10 @@ use aionui_common::{ validate_workspace_path_availability, }; use aionui_db::{ - ClaimCronRunParams, CronRunClaimResult, FinishCronRunParams, IAgentMetadataRepository, - IAssistantDefinitionRepository, IAssistantOverlayRepository, ICronRepository, UpdateCronJobParams, - models::AgentMetadataRow, resolve_agent_binding_from_rows, runtime_backend_for_agent, + ClaimCronRunParams, CronRunClaimResult, DbError, FinishCronRunParams, IAgentMetadataRepository, + IAssistantDefinitionRepository, IAssistantOverlayRepository, ICronRepository, ISkillRepository, + UpdateCronJobParams, UpsertSkillParams, models::AgentMetadataRow, resolve_agent_binding_from_rows, + runtime_backend_for_agent, }; use tracing::{debug, error, info, warn}; @@ -25,7 +26,10 @@ use crate::executor::{ExecutionResult, JobExecutor, PreparedRunNow, RETRY_INTERV use crate::scheduler::{ CronScheduler, compute_next_run, compute_next_run_after_occurrence, validate_schedule, validate_timezone, }; -use crate::skill_file::{delete_skill_file, has_skill_file, write_raw_skill_file, write_skill_file}; +use crate::skill_file::{ + cron_skill_dir, cron_skill_name, delete_skill_file, has_skill_file, parse_skill_content, read_skill_content, + write_raw_skill_file, write_skill_file, +}; use crate::types::{ CreatedBy, CronAgentConfig, CronJob, CronSchedule, ExecutionMode, cron_job_from_row, cron_job_to_response, cron_job_to_row, schedule_from_dto, @@ -52,6 +56,7 @@ pub struct CronService { agent_metadata_repo: Arc, assistant_definition_repo: Arc, assistant_overlay_repo: Arc, + skill_repo: Arc, scheduler: Arc, executor: Arc, emitter: CronEventEmitter, @@ -64,6 +69,7 @@ pub struct CronServiceDeps { pub agent_metadata_repo: Arc, pub assistant_definition_repo: Arc, pub assistant_overlay_repo: Arc, + pub skill_repo: Arc, pub scheduler: Arc, pub executor: Arc, pub emitter: CronEventEmitter, @@ -77,6 +83,7 @@ impl CronService { agent_metadata_repo: deps.agent_metadata_repo, assistant_definition_repo: deps.assistant_definition_repo, assistant_overlay_repo: deps.assistant_overlay_repo, + skill_repo: deps.skill_repo, scheduler: deps.scheduler, executor: deps.executor, emitter: deps.emitter, @@ -89,8 +96,8 @@ impl CronService { // CRUD // ----------------------------------------------------------------------- - pub async fn add_job(&self, req: CreateCronJobRequest) -> Result { - self.add_job_internal(req, None, None).await + pub async fn add_job(&self, user_id: &str, req: CreateCronJobRequest) -> Result { + self.add_job_internal(user_id, req, None, None).await } pub async fn create_for_conversation_helper( @@ -123,7 +130,7 @@ impl CronService { }; let job = self - .add_job_internal(create_req, Some(agent_type), assistant_backend_override) + .add_job_internal(user_id, create_req, Some(agent_type), assistant_backend_override) .await?; if let Err(err) = self .executor @@ -134,7 +141,7 @@ impl CronService { ) .await { - if let Err(cleanup_err) = self.remove_job(&job.id).await { + if let Err(cleanup_err) = self.remove_job(user_id, &job.id).await { warn!( conversation_id, job_id = %job.id, @@ -164,9 +171,12 @@ impl CronService { ) -> Result, CronError> { self.verify_conversation_helper_context(user_id, conversation_id) .await?; - self.list_jobs(&ListCronJobsQuery { - conversation_id: Some(conversation_id.to_owned()), - }) + self.list_jobs( + user_id, + &ListCronJobsQuery { + conversation_id: Some(conversation_id.to_owned()), + }, + ) .await } @@ -180,13 +190,14 @@ impl CronService { self.verify_conversation_helper_context(user_id, conversation_id) .await?; - let existing = self.get_job(job_id).await?; + let existing = self.get_job(user_id, job_id).await?; if existing.conversation_id != conversation_id { return Err(CronError::JobNotFound(job_id.to_owned())); } let job = self .update_job( + user_id, job_id, UpdateCronJobRequest { name: Some(req.name), @@ -230,9 +241,8 @@ impl CronService { } self.executor - .get_conversation_row(conversation_id) + .get_conversation_row_for_user(user_id, conversation_id) .await? - .filter(|row| row.user_id == user_id) .ok_or_else(|| { CronError::Conversation(aionui_conversation::ConversationError::NotFound { id: conversation_id.to_owned(), @@ -242,6 +252,7 @@ impl CronService { async fn add_job_internal( &self, + user_id: &str, req: CreateCronJobRequest, runtime_agent_type: Option, assistant_backend_override: Option, @@ -250,17 +261,26 @@ impl CronService { validate_schedule(&schedule)?; let resolved_agent_type = match runtime_agent_type { Some(agent_type) => agent_type, - None => self.resolve_new_job_agent_type(req.agent_config.as_ref()).await?, + None => { + self.resolve_new_job_agent_type(user_id, req.agent_config.as_ref()) + .await? + } }; validate_aionrs_agent_config(&resolved_agent_type, req.agent_config.as_ref())?; let execution_mode = parse_execution_mode(req.execution_mode.as_deref())?; let created_by = CreatedBy::from_str(&req.created_by)?; let message = req.message.or(req.prompt).unwrap_or_default(); + let conversation_id = req.conversation_id.trim(); + if matches!(execution_mode, ExecutionMode::Existing) { + self.require_existing_conversation_scope(user_id, conversation_id) + .await?; + } let agent_config = match req.agent_config { Some(config) => Some( self.build_cron_agent_config( + user_id, &resolved_agent_type, sanitize_agent_config_dto(config), assistant_backend_override.as_deref(), @@ -275,13 +295,14 @@ impl CronService { let job = CronJob { id: generate_prefixed_id("cron"), + user_id: user_id.to_owned(), name: req.name, enabled: true, schedule, message, execution_mode, agent_config, - conversation_id: req.conversation_id, + conversation_id: conversation_id.to_owned(), conversation_title: req.conversation_title, agent_type: resolved_agent_type, created_by, @@ -305,16 +326,21 @@ impl CronService { self.repo.insert(&row).await?; self.bind_existing_conversation_if_needed(&job).await; self.scheduler.schedule_job(&job); - self.emitter.emit_job_created(&cron_job_to_response(&job)); + self.emitter.emit_job_created(user_id, &cron_job_to_response(&job)); info!(job_id = %job.id, name = %job.name, "Cron job created"); Ok(job) } - pub async fn update_job(&self, job_id: &str, req: UpdateCronJobRequest) -> Result { + pub async fn update_job( + &self, + user_id: &str, + job_id: &str, + req: UpdateCronJobRequest, + ) -> Result { let existing_row = self .repo - .get_by_id(job_id) + .get_by_id_for_user(user_id, job_id) .await? .ok_or_else(|| CronError::JobNotFound(job_id.to_owned()))?; let mut job = cron_job_from_row(existing_row)?; @@ -350,9 +376,7 @@ impl CronService { } job.execution_mode = requested_mode; if requested_mode != original_execution_mode && matches!(requested_mode, ExecutionMode::NewConversation) { - clear_conversation_binding = !job.conversation_id.trim().is_empty(); - job.conversation_id.clear(); - job.conversation_title = None; + clear_conversation_binding = true; } } if req.agent_config.is_some() @@ -365,9 +389,12 @@ impl CronService { } if let Some(config_dto) = &req.agent_config { let config_dto = sanitize_agent_config_dto(config_dto.clone()); - job.agent_type = self.resolve_new_job_agent_type(Some(&config_dto)).await?; + job.agent_type = self.resolve_new_job_agent_type(&job.user_id, Some(&config_dto)).await?; validate_aionrs_agent_config(&job.agent_type, Some(&config_dto))?; - job.agent_config = Some(self.build_cron_agent_config(&job.agent_type, config_dto, None).await?); + job.agent_config = Some( + self.build_cron_agent_config(&job.user_id, &job.agent_type, config_dto, None) + .await?, + ); } if let Some(title) = &req.conversation_title && !clear_conversation_binding @@ -397,13 +424,12 @@ impl CronService { let mut params = build_update_params(&job, &req); if clear_conversation_binding { - params.conversation_id = Some(String::new()); params.conversation_title = Some(None); } if agent_config_changed { params.agent_config = Some(job.agent_config.as_ref().map(serde_json::to_string).transpose()?); } - self.repo.update(job_id, ¶ms).await?; + self.repo.update_for_user(user_id, job_id, ¶ms).await?; if clear_conversation_binding && let Err(err) = self @@ -421,27 +447,28 @@ impl CronService { self.bind_existing_conversation_if_needed(&job).await; self.scheduler.reschedule_job(&job); - self.emitter.emit_job_updated(&cron_job_to_response(&job)); + self.emitter.emit_job_updated(user_id, &cron_job_to_response(&job)); info!(job_id = %job.id, "Cron job updated"); Ok(job) } - pub async fn remove_job(&self, job_id: &str) -> Result<(), CronError> { + pub async fn remove_job(&self, user_id: &str, job_id: &str) -> Result<(), CronError> { self.scheduler.cancel_job(job_id); if let Err(err) = delete_skill_file(&self.data_dir, job_id).await { warn!(job_id, error = %err, "Failed to delete cron skill file during job removal"); } - self.repo.delete(job_id).await?; - self.emitter.emit_job_removed(job_id); + self.delete_skill_row_for_job(user_id, job_id).await; + self.repo.delete_for_user(user_id, job_id).await?; + self.emitter.emit_job_removed(user_id, job_id); info!(job_id, "Cron job removed"); Ok(()) } - pub async fn get_job(&self, job_id: &str) -> Result { + pub async fn get_job(&self, user_id: &str, job_id: &str) -> Result { let row = self .repo - .get_by_id(job_id) + .get_by_id_for_user(user_id, job_id) .await? .ok_or_else(|| CronError::JobNotFound(job_id.to_owned()))?; let mut job = cron_job_from_row(row)?; @@ -449,11 +476,11 @@ impl CronService { Ok(job) } - pub async fn list_jobs(&self, query: &ListCronJobsQuery) -> Result, CronError> { + pub async fn list_jobs(&self, user_id: &str, query: &ListCronJobsQuery) -> Result, CronError> { let rows = if let Some(conv_id) = &query.conversation_id { - self.repo.list_by_conversation(conv_id).await? + self.repo.list_by_conversation_for_user(user_id, conv_id).await? } else { - self.repo.list_all().await? + self.repo.list_all_for_user(user_id).await? }; let mut jobs = Vec::with_capacity(rows.len()); @@ -475,7 +502,9 @@ impl CronService { warn!(error = %error, "Failed to clean up old cron run records"); } - let rows = match self.repo.list_enabled().await { + self.reconcile_skill_rows().await; + + let rows = match self.repo.list_enabled_system().await { Ok(rows) => rows, Err(e) => { error!(error = %e, "Failed to load enabled cron jobs"); @@ -516,8 +545,78 @@ impl CronService { info!(scheduled, "Cron service initialized"); } + /// Reconcile owner-scoped skill rows for saved cron skills at startup. + /// + /// For each job whose skill file exists on disk, ensure an owner-scoped + /// row named `cron-{job_id}` exists. Legacy rows produced by the old + /// machine-level catalog scan were named after the SKILL.md frontmatter + /// name instead; those duplicates are soft-deleted so listings do not + /// double up while historical conversations can still materialize them. + async fn reconcile_skill_rows(&self) { + let rows = match self.repo.list_all_system().await { + Ok(rows) => rows, + Err(e) => { + warn!(error = %e, "Failed to list cron jobs for skill row reconciliation"); + return; + } + }; + + let mut reconciled = 0u32; + let mut cleaned = 0u32; + for row in rows { + let owner_user_id = row.user_id.clone(); + let job = match cron_job_from_row(row) { + Ok(job) => job, + Err(_) => continue, + }; + match has_skill_file(&self.data_dir, &job.id).await { + Ok(true) => {} + Ok(false) => continue, + Err(e) => { + warn!(job_id = %job.id, error = %e, "Failed to check cron skill file during reconciliation"); + continue; + } + } + + if let Ok(Some(content)) = read_skill_content(&self.data_dir, &job.id).await + && let Ok(parsed) = parse_skill_content(&content) + && let Ok(dir_name) = cron_skill_name(&job.id) + && parsed.name != dir_name + && let Ok(Some(legacy)) = self + .skill_repo + .find_by_name_for_user(&owner_user_id, &parsed.name) + .await + && legacy.source == "cron" + && legacy.user_id.is_some() + { + match self + .skill_repo + .delete_by_name_for_user(&owner_user_id, &parsed.name) + .await + { + Ok(_) => cleaned += 1, + Err(DbError::NotFound(_)) => {} + Err(e) => { + warn!(job_id = %job.id, error = %e, "Failed to clean legacy cron skill row"); + } + } + } + + match self.upsert_skill_row_for_job(&owner_user_id, &job).await { + Ok(()) => reconciled += 1, + Err(e) => { + warn!(job_id = %job.id, error = %e, "Failed to reconcile cron skill row"); + } + } + } + + if reconciled > 0 || cleaned > 0 { + info!(reconciled, cleaned, "Cron skill rows reconciled"); + } + } + pub async fn tick(&self, job_id: &str, scheduled_at: i64) { - let row = match self.repo.get_by_id(job_id).await { + let row = match self.repo.get_by_id_system(job_id).await { Ok(Some(r)) => r, Ok(None) => { warn!(job_id, "Tick: job not found, cancelling timer"); @@ -581,7 +680,9 @@ impl CronService { Ok(CronRunClaimResult::QueueBusy) => { self.record_queue_busy_skip(&job).await; self.reschedule_after_execution(&job, scheduled_at).await; - self.emitter.emit_job_executed(job_id, "skipped", None); + if let Some(user_id) = self.owner_user_id_for_job(&job).await { + self.emitter.emit_job_executed(&user_id, job_id, "skipped", None); + } return; } Err(error) => { @@ -629,7 +730,7 @@ impl CronService { } pub async fn handle_system_resume(&self) { - let rows = match self.repo.list_enabled().await { + let rows = match self.repo.list_enabled_system().await { Ok(r) => r, Err(e) => { error!(error = %e, "Resume: failed to load enabled jobs"); @@ -674,7 +775,9 @@ impl CronService { self.record_missed_execution(&job).await; self.insert_missed_job_tips(&job).await; self.reschedule_after_missed(&job).await; - self.emitter.emit_job_executed(&job.id, "missed", None); + if let Some(user_id) = self.owner_user_id_for_job(&job).await { + self.emitter.emit_job_executed(&user_id, &job.id, "missed", None); + } continue; } @@ -684,10 +787,10 @@ impl CronService { info!("System resume: all cron timers rescheduled"); } - pub async fn run_now(&self, job_id: &str) -> Result { + pub async fn run_now(&self, user_id: &str, job_id: &str) -> Result { let row = self .repo - .get_by_id(job_id) + .get_by_id_for_user(user_id, job_id) .await? .ok_or_else(|| CronError::JobNotFound(job_id.to_owned()))?; let mut job = cron_job_from_row(row)?; @@ -703,10 +806,11 @@ impl CronService { let conversation_id = prepared.conversation_id.clone(); let service = self.clone(); let job_id = job.id.clone(); + let user_id = user_id.to_owned(); tokio::spawn(async move { let result = service.executor.execute_prepared(&job, prepared).await; - service.handle_run_now_result(&job_id, result).await; + service.handle_run_now_result(&user_id, &job_id, result).await; }); Ok(RunNowResponse { conversation_id }) @@ -716,10 +820,10 @@ impl CronService { // Skill management // ----------------------------------------------------------------------- - pub async fn save_skill(&self, job_id: &str, req: SaveCronSkillRequest) -> Result<(), CronError> { + pub async fn save_skill(&self, user_id: &str, job_id: &str, req: SaveCronSkillRequest) -> Result<(), CronError> { let row = self .repo - .get_by_id(job_id) + .get_by_id_for_user(user_id, job_id) .await? .ok_or_else(|| CronError::JobNotFound(job_id.to_owned()))?; @@ -731,17 +835,69 @@ impl CronService { skill_content: Some(Some(req.content)), ..Default::default() }; - self.repo.update(job_id, ¶ms).await?; - self.executor.mark_skill_suggest_artifacts_saved(job_id).await?; + self.repo.update_for_user(user_id, job_id, ¶ms).await?; + self.upsert_skill_row_for_job(user_id, &job).await?; + self.executor + .mark_skill_suggest_artifacts_saved(user_id, job_id) + .await?; info!(job_id, "Skill content saved"); Ok(()) } - pub async fn has_skill(&self, job_id: &str) -> Result { + /// Upsert the owner-scoped skill row for a job's saved cron skill. + /// + /// Cron skill content is user data (the job prompt), so the catalog row + /// is owned by the job owner and named after the on-disk skill dir + /// (`cron-{job_id}`) — the stable name the executor materializes by. + async fn upsert_skill_row_for_job(&self, user_id: &str, job: &CronJob) -> Result<(), CronError> { + let name = cron_skill_name(&job.id)?; + let path = cron_skill_dir(&self.data_dir, &job.id)?; + let description = match read_skill_content(&self.data_dir, &job.id).await? { + Some(content) => parse_skill_content(&content) + .map(|parsed| parsed.description) + .ok() + .filter(|desc| !desc.trim().is_empty()), + None => None, + }; + let description = description + .or_else(|| job.description.clone()) + .unwrap_or_else(|| format!("Saved cron skill for {}", job.name)); + + self.skill_repo + .upsert_for_user( + user_id, + UpsertSkillParams { + name: &name, + description: Some(&description), + path: &path.to_string_lossy(), + source: "cron", + enabled: true, + }, + ) + .await + .map_err(|e| CronError::Scheduler(format!("Failed to upsert cron skill row: {e}")))?; + Ok(()) + } + + /// Soft-delete the owner-scoped skill row for a job. Missing rows are + /// fine — jobs without a saved skill never had one. + async fn delete_skill_row_for_job(&self, user_id: &str, job_id: &str) { + let Ok(name) = cron_skill_name(job_id) else { + return; + }; + match self.skill_repo.delete_by_name_for_user(user_id, &name).await { + Ok(_) | Err(DbError::NotFound(_)) => {} + Err(e) => { + warn!(job_id, error = %e, "Failed to delete cron skill row"); + } + } + } + + pub async fn has_skill(&self, user_id: &str, job_id: &str) -> Result { let row = self .repo - .get_by_id(job_id) + .get_by_id_for_user(user_id, job_id) .await? .ok_or_else(|| CronError::JobNotFound(job_id.to_owned()))?; @@ -751,9 +907,9 @@ impl CronService { Ok(HasSkillResponse { has_skill }) } - pub async fn delete_skill(&self, job_id: &str) -> Result<(), CronError> { + pub async fn delete_skill(&self, user_id: &str, job_id: &str) -> Result<(), CronError> { self.repo - .get_by_id(job_id) + .get_by_id_for_user(user_id, job_id) .await? .ok_or_else(|| CronError::JobNotFound(job_id.to_owned()))?; @@ -763,7 +919,8 @@ impl CronService { skill_content: Some(None), ..Default::default() }; - self.repo.update(job_id, ¶ms).await?; + self.repo.update_for_user(user_id, job_id, ¶ms).await?; + self.delete_skill_row_for_job(user_id, job_id).await; info!(job_id, "Skill content deleted"); Ok(()) @@ -779,6 +936,7 @@ impl CronService { async fn resolve_new_job_agent_type( &self, + user_id: &str, agent_config: Option<&aionui_api_types::CronAgentConfigWriteDto>, ) -> Result { let Some(assistant_id) = agent_config.and_then(|config| config.assistant_id.as_deref()) else { @@ -787,7 +945,7 @@ impl CronService { )); }; - self.resolve_agent_type_for_assistant_id(assistant_id).await + self.resolve_agent_type_for_assistant_id(user_id, assistant_id).await } async fn resolve_job_agent_type(&self, job: &CronJob) -> Result { @@ -806,7 +964,75 @@ impl CronService { )); }; - self.resolve_agent_type_for_assistant_id(assistant_id).await + self.resolve_agent_type_for_assistant_id(&job.user_id, assistant_id) + .await + } + + async fn owner_user_id_for_job(&self, job: &CronJob) -> Option { + if !job.user_id.trim().is_empty() { + return Some(job.user_id.clone()); + } + + let conversation_id = job.conversation_id.trim(); + if conversation_id.is_empty() { + warn!(job_id = %job.id, "Cron event skipped because job has no conversation owner"); + return None; + } + + match self.executor.get_conversation_row(conversation_id).await { + Ok(Some(row)) if !row.user_id.trim().is_empty() => Some(row.user_id), + Ok(Some(_)) => { + warn!( + job_id = %job.id, + conversation_id, + "Cron event skipped because conversation owner is empty" + ); + None + } + Ok(None) => { + warn!( + job_id = %job.id, + conversation_id, + "Cron event skipped because conversation owner could not be resolved" + ); + None + } + Err(err) => { + warn!( + job_id = %job.id, + conversation_id, + error = %err, + "Cron event skipped because conversation owner lookup failed" + ); + None + } + } + } + + async fn require_existing_conversation_scope(&self, user_id: &str, conversation_id: &str) -> Result<(), CronError> { + if conversation_id.is_empty() { + return Err(CronError::Conversation( + aionui_conversation::ConversationError::NotFound { + id: conversation_id.to_owned(), + }, + )); + } + + let Some(row) = self.executor.get_conversation_row(conversation_id).await? else { + return Err(CronError::Conversation( + aionui_conversation::ConversationError::NotFound { + id: conversation_id.to_owned(), + }, + )); + }; + + if row.user_id != user_id { + return Err(CronError::CrossAccountReference( + "Referenced conversation belongs to another user.".into(), + )); + } + + Ok(()) } async fn is_team_conversation_job(&self, job: &CronJob) -> Result { @@ -826,18 +1052,25 @@ impl CronService { .is_some_and(|value| !value.trim().is_empty())) } - async fn resolve_agent_type_for_assistant_id(&self, assistant_id: &str) -> Result { + async fn resolve_agent_type_for_assistant_id( + &self, + user_id: &str, + assistant_id: &str, + ) -> Result { let definition = self .assistant_definition_repo - .get_by_assistant_id(assistant_id) + .get_by_assistant_id_for_user(user_id, assistant_id) .await? .ok_or_else(|| CronError::InvalidAgentConfig(format!("assistant '{assistant_id}' not found")))?; - let overlay = self.assistant_overlay_repo.get(&definition.id).await?; + let overlay = self + .assistant_overlay_repo + .get_for_user(user_id, &definition.id) + .await?; let effective_agent_id = overlay .as_ref() .and_then(|item| item.agent_id_override.as_deref()) .unwrap_or(definition.agent_id.as_str()); - let effective_backend = self.runtime_backend_for_agent_id(effective_agent_id).await?; + let effective_backend = self.runtime_backend_for_agent_id(user_id, effective_agent_id).await?; Ok(runtime_agent_type_for_backend(&effective_backend).to_owned()) } @@ -882,7 +1115,7 @@ impl CronService { conversation_id: Some(conversation_id.to_owned()), ..Default::default() }; - self.repo.update(&job.id, ¶ms).await?; + self.repo.update_for_user(&job.user_id, &job.id, ¶ms).await?; self.executor .bind_cron_job_to_conversation( conversation_id, @@ -917,7 +1150,7 @@ impl CronService { }; let workspace_to_clear = match self .executor - .auto_workspace_to_delete_for_conversation(conversation_id) + .auto_workspace_to_delete_for_conversation(&job.user_id, conversation_id) .await { Ok(Some(path)) => path, @@ -943,6 +1176,7 @@ impl CronService { async fn handle_execution_result(&self, job: CronJob, scheduled_at: i64, result: ExecutionResult) { let job_id = &job.id; + let owner_user_id = self.owner_user_id_for_job(&job).await; match result { ExecutionResult::Success { conversation_id } => { @@ -952,9 +1186,12 @@ impl CronService { { return; } - self.update_job_after_success(job_id, &conversation_id).await; + self.update_job_after_success(&job.user_id, job_id, &conversation_id) + .await; self.reschedule_after_execution(&job, scheduled_at).await; - self.emitter.emit_job_executed(job_id, "ok", None); + if let Some(user_id) = owner_user_id.as_deref() { + self.emitter.emit_job_executed(user_id, job_id, "ok", None); + } } ExecutionResult::Retrying { attempt } => { if !self.schedule_retry(job_id, scheduled_at, attempt).await { @@ -964,7 +1201,7 @@ impl CronService { retry_count: Some(attempt), ..Default::default() }; - if let Err(e) = self.repo.update(job_id, ¶ms).await { + if let Err(e) = self.repo.update_for_user(&job.user_id, job_id, ¶ms).await { error!(job_id, error = %e, "Failed to update retry count"); } } @@ -980,11 +1217,13 @@ impl CronService { retry_count: Some(0), ..Default::default() }; - if let Err(e) = self.repo.update(job_id, ¶ms).await { + if let Err(e) = self.repo.update_for_user(&job.user_id, job_id, ¶ms).await { error!(job_id, error = %e, "Failed to update skipped status"); } self.reschedule_after_execution(&job, scheduled_at).await; - self.emitter.emit_job_executed(job_id, "skipped", None); + if let Some(user_id) = owner_user_id.as_deref() { + self.emitter.emit_job_executed(user_id, job_id, "skipped", None); + } } ExecutionResult::Error { message } => { if !self @@ -993,29 +1232,31 @@ impl CronService { { return; } - self.update_job_after_error(job_id, &message).await; + self.update_job_after_error(&job.user_id, job_id, &message).await; self.reschedule_after_execution(&job, scheduled_at).await; - self.emitter.emit_job_executed(job_id, "error", Some(&message)); + if let Some(user_id) = owner_user_id.as_deref() { + self.emitter.emit_job_executed(user_id, job_id, "error", Some(&message)); + } } } } - async fn handle_run_now_result(&self, job_id: &str, result: ExecutionResult) { + async fn handle_run_now_result(&self, user_id: &str, job_id: &str, result: ExecutionResult) { match result { ExecutionResult::Success { conversation_id } => { - self.update_job_after_success(job_id, &conversation_id).await; - self.emitter.emit_job_executed(job_id, "ok", None); + self.update_job_after_success(user_id, job_id, &conversation_id).await; + self.emitter.emit_job_executed(user_id, job_id, "ok", None); } ExecutionResult::Error { message } => { - self.update_job_after_error(job_id, &message).await; - self.emitter.emit_job_executed(job_id, "error", Some(&message)); + self.update_job_after_error(user_id, job_id, &message).await; + self.emitter.emit_job_executed(user_id, job_id, "error", Some(&message)); } ExecutionResult::Retrying { attempt } => { let params = UpdateCronJobParams { retry_count: Some(attempt), ..Default::default() }; - if let Err(err) = self.repo.update(job_id, ¶ms).await { + if let Err(err) = self.repo.update_for_user(user_id, job_id, ¶ms).await { error!( job_id, error = %err, @@ -1029,20 +1270,20 @@ impl CronService { retry_count: Some(0), ..Default::default() }; - if let Err(err) = self.repo.update(job_id, ¶ms).await { + if let Err(err) = self.repo.update_for_user(user_id, job_id, ¶ms).await { error!( job_id, error = %err, "Failed to update run-now skipped status" ); } - self.emitter.emit_job_executed(job_id, "skipped", None); + self.emitter.emit_job_executed(user_id, job_id, "skipped", None); } } } - async fn update_job_after_success(&self, job_id: &str, conversation_id: &str) { - let existing_row = match self.repo.get_by_id(job_id).await { + async fn update_job_after_success(&self, user_id: &str, job_id: &str, conversation_id: &str) { + let existing_row = match self.repo.get_by_id_for_user(user_id, job_id).await { Ok(Some(r)) => r, Ok(None) => return, Err(e) => { @@ -1068,7 +1309,7 @@ impl CronService { conversation_id: needs_conversation_bind.then(|| conversation_id.to_owned()), ..Default::default() }; - if let Err(e) = self.repo.update(job_id, ¶ms).await { + if let Err(e) = self.repo.update_for_user(user_id, job_id, ¶ms).await { error!(job_id, error = %e, "Failed to update job after success"); return; } @@ -1088,8 +1329,8 @@ impl CronService { } } - async fn update_job_after_error(&self, job_id: &str, message: &str) { - let run_count = match self.repo.get_by_id(job_id).await { + async fn update_job_after_error(&self, user_id: &str, job_id: &str, message: &str) { + let run_count = match self.repo.get_by_id_for_user(user_id, job_id).await { Ok(Some(r)) => r.run_count, Ok(None) => return, Err(e) => { @@ -1106,7 +1347,7 @@ impl CronService { run_count: Some(run_count + 1), ..Default::default() }; - if let Err(e) = self.repo.update(job_id, ¶ms).await { + if let Err(e) = self.repo.update_for_user(user_id, job_id, ¶ms).await { error!(job_id, error = %e, "Failed to update job after error"); } } @@ -1130,7 +1371,7 @@ impl CronService { next_run_at: Some(None), ..Default::default() }; - if let Err(e) = self.repo.update(&job.id, ¶ms).await { + if let Err(e) = self.repo.update_for_user(&job.user_id, &job.id, ¶ms).await { error!(job_id = %job.id, error = %e, "Failed to disable at-type job"); } self.scheduler.cancel_job(&job.id); @@ -1140,7 +1381,10 @@ impl CronService { next_run_at: None, ..job.clone() }; - self.emitter.emit_job_updated(&cron_job_to_response(&disabled)); + if let Some(user_id) = self.owner_user_id_for_job(job).await { + self.emitter + .emit_job_updated(&user_id, &cron_job_to_response(&disabled)); + } info!(job_id = %job.id, "At-type job executed, auto-disabled"); return; @@ -1155,7 +1399,7 @@ impl CronService { next_run_at: Some(next), ..Default::default() }; - if let Err(e) = self.repo.update(&job.id, ¶ms).await { + if let Err(e) = self.repo.update_for_user(&job.user_id, &job.id, ¶ms).await { error!(job_id = %job.id, error = %e, "Failed to update next_run_at"); } self.scheduler.reschedule_job(&updated); @@ -1168,7 +1412,7 @@ impl CronService { retry_count: Some(0), ..Default::default() }; - if let Err(err) = self.repo.update(&job.id, ¶ms).await { + if let Err(err) = self.repo.update_for_user(&job.user_id, &job.id, ¶ms).await { error!( job_id = %job.id, error = %err, @@ -1192,9 +1436,12 @@ impl CronService { .insert_tips_message(&job.conversation_id, &content, "warning") .await { - Ok(()) => self - .emitter - .emit_conversation_tips(&job.conversation_id, &content, "warning"), + Ok(()) => { + if let Some(user_id) = self.owner_user_id_for_job(job).await { + self.emitter + .emit_conversation_tips(&user_id, &job.conversation_id, &content, "warning"); + } + } Err(err) => { warn!( job_id = %job.id, @@ -1214,7 +1461,7 @@ impl CronService { next_run_at: Some(None), ..Default::default() }; - if let Err(err) = self.repo.update(&job.id, ¶ms).await { + if let Err(err) = self.repo.update_for_user(&job.user_id, &job.id, ¶ms).await { error!( job_id = %job.id, error = %err, @@ -1230,7 +1477,7 @@ impl CronService { next_run_at: Some(next), ..Default::default() }; - if let Err(err) = self.repo.update(&job.id, ¶ms).await { + if let Err(err) = self.repo.update_for_user(&job.user_id, &job.id, ¶ms).await { error!( job_id = %job.id, error = %err, @@ -1332,15 +1579,15 @@ impl CronService { retry_count: Some(0), ..Default::default() }; - if let Err(error) = self.repo.update(&job.id, ¶ms).await { + if let Err(error) = self.repo.update_for_user(&job.user_id, &job.id, ¶ms).await { error!(job_id = %job.id, error = %error, "Failed to record queue-busy cron skip"); } } - pub async fn delete_jobs_by_conversation(&self, conversation_id: &str) { + pub async fn delete_jobs_by_conversation(&self, user_id: &str, conversation_id: &str) { let workspace_to_clear = match self .executor - .auto_workspace_to_delete_for_conversation(conversation_id) + .auto_workspace_to_delete_for_conversation(user_id, conversation_id) .await { Ok(value) => value, @@ -1359,16 +1606,17 @@ impl CronService { return; }; - self.clear_deleted_workspace_from_jobs(conversation_id, &workspace_to_clear) + self.clear_deleted_workspace_from_jobs(user_id, conversation_id, &workspace_to_clear) .await; debug!(conversation_id, "Conversation deleted; cron jobs are preserved"); } - async fn clear_deleted_workspace_from_jobs(&self, conversation_id: &str, workspace_to_clear: &Path) { - let jobs = match self.repo.list_by_conversation(conversation_id).await { + async fn clear_deleted_workspace_from_jobs(&self, user_id: &str, conversation_id: &str, workspace_to_clear: &Path) { + let jobs = match self.repo.list_by_conversation_for_user(user_id, conversation_id).await { Ok(rows) => rows, Err(err) => { error!( + user_id, conversation_id, error = %err, "Failed to list cron jobs for deleted workspace cleanup" @@ -1419,7 +1667,7 @@ impl CronService { agent_config: Some(Some(agent_config_json)), ..Default::default() }; - match self.repo.update(&row.id, ¶ms).await { + match self.repo.update_for_user(&row.user_id, &row.id, ¶ms).await { Ok(()) => cleared += 1, Err(err) => { error!( @@ -1489,7 +1737,7 @@ impl CronService { extra_assistant_id.as_ref(), legacy_agent_label, ) { - (None, None, Some(label)) => self.resolve_assistant_id_for_agent_label(&label).await, + (None, None, Some(label)) => self.resolve_assistant_id_for_agent_label(&row.user_id, &label).await, _ => None, }; let fallback_assistant_id = match ( @@ -1497,7 +1745,7 @@ impl CronService { extra_assistant_id.as_ref(), legacy_assistant_id.as_ref(), ) { - (None, None, None) => self.resolve_default_assistant_id().await, + (None, None, None) => self.resolve_default_assistant_id(&row.user_id).await, _ => None, }; let uses_default_assistant_fallback = fallback_assistant_id.is_some(); @@ -1506,7 +1754,7 @@ impl CronService { .or(legacy_assistant_id) .or(fallback_assistant_id); let assistant_name = match assistant_id.as_deref() { - Some(assistant_id) => match self.resolve_assistant_name(Some(assistant_id)).await { + Some(assistant_id) => match self.resolve_assistant_name(&row.user_id, Some(assistant_id)).await { Ok(value) => value, Err(err) => { warn!( @@ -1521,7 +1769,10 @@ impl CronService { None => None, }; let snapshot_backend = match assistant_snapshot.as_ref() { - Some(snapshot) => match self.runtime_backend_for_agent_id(snapshot.agent_id.trim()).await { + Some(snapshot) => match self + .runtime_backend_for_agent_id(&row.user_id, snapshot.agent_id.trim()) + .await + { Ok(value) => Some(value).filter(|value| !value.is_empty()), Err(err) => { warn!( @@ -1538,7 +1789,7 @@ impl CronService { None } else { snapshot_backend.clone().or(self - .resolve_assistant_backend(assistant_id.as_deref()) + .resolve_assistant_backend(&row.user_id, assistant_id.as_deref()) .await .unwrap_or(None)) }; @@ -1569,6 +1820,7 @@ impl CronService { }; let full_auto_mode = match self .resolve_cron_full_auto_mode( + &row.user_id, &row.r#type, assistant_id_for_mode, assistant_snapshot.as_ref().map(|snapshot| snapshot.agent_id.as_str()), @@ -1625,6 +1877,7 @@ impl CronService { async fn build_cron_agent_config( &self, + user_id: &str, runtime_agent_type: &str, config: aionui_api_types::CronAgentConfigWriteDto, _assistant_backend_override: Option<&str>, @@ -1636,7 +1889,7 @@ impl CronService { }; let assistant_backend = self - .resolve_assistant_backend(Some(assistant_id)) + .resolve_assistant_backend(user_id, Some(assistant_id)) .await? .ok_or_else(|| { CronError::InvalidAgentConfig(format!( @@ -1645,6 +1898,7 @@ impl CronService { })?; let full_auto_mode = self .resolve_cron_full_auto_mode( + user_id, runtime_agent_type, Some(assistant_id), None, @@ -1666,41 +1920,58 @@ impl CronService { }) } - async fn resolve_assistant_backend(&self, assistant_id: Option<&str>) -> Result, CronError> { + async fn resolve_assistant_backend( + &self, + user_id: &str, + assistant_id: Option<&str>, + ) -> Result, CronError> { let Some(assistant_id) = assistant_id.filter(|value| !value.is_empty()) else { return Ok(None); }; - let Some(definition) = self.assistant_definition_repo.get_by_assistant_id(assistant_id).await? else { + let Some(definition) = self + .assistant_definition_repo + .get_by_assistant_id_for_user(user_id, assistant_id) + .await? + else { return Ok(None); }; - let overlay = self.assistant_overlay_repo.get(&definition.id).await?; + let overlay = self + .assistant_overlay_repo + .get_for_user(user_id, &definition.id) + .await?; let effective_agent_id = overlay .as_ref() .and_then(|item| item.agent_id_override.as_deref()) .unwrap_or(definition.agent_id.as_str()); - Ok(Some(self.runtime_backend_for_agent_id(effective_agent_id).await?)) + Ok(Some( + self.runtime_backend_for_agent_id(user_id, effective_agent_id).await?, + )) } - async fn resolve_assistant_name(&self, assistant_id: Option<&str>) -> Result, CronError> { + async fn resolve_assistant_name( + &self, + user_id: &str, + assistant_id: Option<&str>, + ) -> Result, CronError> { let Some(assistant_id) = assistant_id.filter(|value| !value.is_empty()) else { return Ok(None); }; Ok(self .assistant_definition_repo - .get_by_assistant_id(assistant_id) + .get_by_assistant_id_for_user(user_id, assistant_id) .await? .map(|definition| definition.name.trim().to_owned()) .filter(|value| !value.is_empty())) } - async fn resolve_assistant_id_for_agent_label(&self, agent_label: &str) -> Option { - let rows = self.agent_metadata_repo.list_all().await.ok()?; + async fn resolve_assistant_id_for_agent_label(&self, user_id: &str, agent_label: &str) -> Option { + let rows = self.agent_metadata_repo.list_all_for_user(user_id).await.ok()?; let binding = resolve_agent_binding_from_rows(&rows, agent_label)?; self.assistant_definition_repo - .list() + .list_for_user(user_id) .await .ok()? .into_iter() @@ -1717,9 +1988,9 @@ impl CronService { .map(|definition| definition.assistant_id) } - async fn resolve_default_assistant_id(&self) -> Option { + async fn resolve_default_assistant_id(&self, user_id: &str) -> Option { self.assistant_definition_repo - .list() + .list_for_user(user_id) .await .ok()? .into_iter() @@ -1736,8 +2007,8 @@ impl CronService { .map(|definition| definition.assistant_id) } - async fn runtime_backend_for_agent_id(&self, agent_id: &str) -> Result { - let rows = self.agent_metadata_repo.list_all().await?; + async fn runtime_backend_for_agent_id(&self, user_id: &str, agent_id: &str) -> Result { + let rows = self.agent_metadata_repo.list_all_for_user(user_id).await?; Ok(resolve_agent_binding_from_rows(&rows, agent_id) .map(|binding| binding.runtime_backend) .unwrap_or_else(|| agent_id.to_owned())) @@ -1745,20 +2016,21 @@ impl CronService { async fn resolve_cron_full_auto_mode( &self, + user_id: &str, runtime_agent_type: &str, assistant_id: Option<&str>, agent_id_hint: Option<&str>, backend_hint: Option<&str>, ) -> Result { - if let Some(row) = self.resolve_agent_metadata_for_assistant(assistant_id).await? { + if let Some(row) = self.resolve_agent_metadata_for_assistant(user_id, assistant_id).await? { return Ok(full_auto_mode_from_metadata(&row, runtime_agent_type)); } - if let Some(row) = self.resolve_agent_metadata_for_value(agent_id_hint).await? { + if let Some(row) = self.resolve_agent_metadata_for_value(user_id, agent_id_hint).await? { return Ok(full_auto_mode_from_metadata(&row, runtime_agent_type)); } - if let Some(row) = self.resolve_agent_metadata_for_value(backend_hint).await? { + if let Some(row) = self.resolve_agent_metadata_for_value(user_id, backend_hint).await? { return Ok(full_auto_mode_from_metadata(&row, runtime_agent_type)); } @@ -1767,33 +2039,43 @@ impl CronService { async fn resolve_agent_metadata_for_assistant( &self, + user_id: &str, assistant_id: Option<&str>, ) -> Result, CronError> { let Some(assistant_id) = assistant_id.map(str::trim).filter(|value| !value.is_empty()) else { return Ok(None); }; - let Some(definition) = self.assistant_definition_repo.get_by_assistant_id(assistant_id).await? else { + let Some(definition) = self + .assistant_definition_repo + .get_by_assistant_id_for_user(user_id, assistant_id) + .await? + else { return Ok(None); }; - let overlay = self.assistant_overlay_repo.get(&definition.id).await?; + let overlay = self + .assistant_overlay_repo + .get_for_user(user_id, &definition.id) + .await?; let effective_agent_id = overlay .as_ref() .and_then(|item| item.agent_id_override.as_deref()) .unwrap_or(definition.agent_id.as_str()); - self.resolve_agent_metadata_for_value(Some(effective_agent_id)).await + self.resolve_agent_metadata_for_value(user_id, Some(effective_agent_id)) + .await } async fn resolve_agent_metadata_for_value( &self, + user_id: &str, value: Option<&str>, ) -> Result, CronError> { let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else { return Ok(None); }; - let rows = self.agent_metadata_repo.list_all().await?; + let rows = self.agent_metadata_repo.list_all_for_user(user_id).await?; let Some(binding) = resolve_agent_binding_from_rows(&rows, value) else { return Ok(None); }; @@ -1826,8 +2108,8 @@ fn fallback_full_auto_mode(runtime_agent_type: &str, backend_hint: Option<&str>) #[async_trait::async_trait] impl aionui_common::OnConversationDelete for CronService { - async fn on_conversation_deleted(&self, conversation_id: &str) { - self.delete_jobs_by_conversation(conversation_id).await; + async fn on_conversation_deleted(&self, user_id: &str, conversation_id: &str) { + self.delete_jobs_by_conversation(user_id, conversation_id).await; } } @@ -2313,6 +2595,7 @@ mod tests { fn sample_job() -> CronJob { CronJob { id: "cron_test".into(), + user_id: "user1".into(), name: "Test".into(), enabled: true, schedule: CronSchedule::Every { diff --git a/crates/aionui-cron/src/skill_suggest.rs b/crates/aionui-cron/src/skill_suggest.rs index 029f59bd5..0c3d479a7 100644 --- a/crates/aionui-cron/src/skill_suggest.rs +++ b/crates/aionui-cron/src/skill_suggest.rs @@ -39,17 +39,19 @@ impl SkillSuggestDetector { } } - pub fn schedule_check(&self, conversation_id: String, job_id: String, workspace: String) { + pub fn schedule_check(&self, user_id: String, conversation_id: String, job_id: String, workspace: String) { let detector = self.clone(); tokio::spawn(async move { - detector.check_with_retry(&conversation_id, &job_id, &workspace).await; + detector + .check_with_retry(&user_id, &conversation_id, &job_id, &workspace) + .await; }); } - async fn check_with_retry(&self, conversation_id: &str, job_id: &str, workspace: &str) { + async fn check_with_retry(&self, user_id: &str, conversation_id: &str, job_id: &str, workspace: &str) { for delay_ms in RETRY_DELAYS_MS { sleep(Duration::from_millis(delay_ms)).await; - match self.check_and_emit(conversation_id, job_id, workspace).await { + match self.check_and_emit(user_id, conversation_id, job_id, workspace).await { Ok(true) => return, Ok(false) => continue, Err(err) => { @@ -64,7 +66,13 @@ impl SkillSuggestDetector { } } - async fn check_and_emit(&self, conversation_id: &str, job_id: &str, workspace: &str) -> Result { + async fn check_and_emit( + &self, + user_id: &str, + conversation_id: &str, + job_id: &str, + workspace: &str, + ) -> Result { if workspace.trim().is_empty() { return Ok(false); } @@ -99,6 +107,7 @@ impl SkillSuggestDetector { self.set_last_hash(conversation_id, job_id, hash); self.emit( + user_id, conversation_id, job_id, &validated.name, @@ -109,13 +118,22 @@ impl SkillSuggestDetector { Ok(true) } - async fn emit(&self, conversation_id: &str, job_id: &str, name: &str, description: &str, skill_content: &str) { - self.persist_and_broadcast(conversation_id, job_id, name, description, skill_content) + async fn emit( + &self, + user_id: &str, + conversation_id: &str, + job_id: &str, + name: &str, + description: &str, + skill_content: &str, + ) { + self.persist_and_broadcast(user_id, conversation_id, job_id, name, description, skill_content) .await; } async fn persist_and_broadcast( &self, + user_id: &str, conversation_id: &str, job_id: &str, name: &str, @@ -124,7 +142,7 @@ impl SkillSuggestDetector { ) { let row = build_skill_suggest_artifact(conversation_id, job_id, name, description, skill_content, now_ms()); - let row = match self.conversation_repo.upsert_artifact(&row).await { + let row = match self.conversation_repo.upsert_artifact(user_id, &row).await { Ok(row) => row, Err(err) => { warn!( @@ -137,7 +155,7 @@ impl SkillSuggestDetector { } }; - if let Err(err) = broadcast_artifact(&self.broadcaster, &row) { + if let Err(err) = broadcast_artifact(&self.broadcaster, user_id, &row) { warn!( conversation_id, job_id, @@ -222,7 +240,7 @@ mod tests { let mut rx = bus.subscribe(); let emitted = detector - .check_and_emit("conv-1", "cron-1", &workspace.to_string_lossy()) + .check_and_emit("system_default_user", "conv-1", "cron-1", &workspace.to_string_lossy()) .await .unwrap(); @@ -234,7 +252,7 @@ mod tests { assert_eq!(msg.data["payload"]["cron_job_id"], "cron-1"); assert_eq!(msg.data["payload"]["name"], "daily-report"); - let rows = repo.list_artifacts("conv-1").await.unwrap(); + let rows = repo.list_artifacts("system_default_user", "conv-1").await.unwrap(); assert_eq!(rows.len(), 1); assert_eq!(rows[0].kind, "skill_suggest"); assert_eq!(rows[0].status, "pending"); @@ -264,7 +282,7 @@ mod tests { assert!( detector - .check_and_emit("conv-1", "cron-1", &workspace.to_string_lossy()) + .check_and_emit("system_default_user", "conv-1", "cron-1", &workspace.to_string_lossy(),) .await .unwrap() ); @@ -272,7 +290,7 @@ mod tests { assert!( detector - .check_and_emit("conv-2", "cron-1", &workspace.to_string_lossy()) + .check_and_emit("system_default_user", "conv-2", "cron-1", &workspace.to_string_lossy(),) .await .unwrap() ); @@ -304,7 +322,7 @@ mod tests { assert!( detector - .check_and_emit("conv-1", "cron-1", &workspace.to_string_lossy()) + .check_and_emit("system_default_user", "conv-1", "cron-1", &workspace.to_string_lossy(),) .await .unwrap() ); @@ -312,7 +330,7 @@ mod tests { assert!( detector - .check_and_emit("conv-1", "cron-1", &workspace.to_string_lossy()) + .check_and_emit("system_default_user", "conv-1", "cron-1", &workspace.to_string_lossy(),) .await .unwrap() ); @@ -348,7 +366,7 @@ mod tests { let mut rx = bus.subscribe(); let emitted = detector - .check_and_emit("conv-1", "cron-1", &workspace.to_string_lossy()) + .check_and_emit("system_default_user", "conv-1", "cron-1", &workspace.to_string_lossy()) .await .unwrap(); diff --git a/crates/aionui-cron/src/types.rs b/crates/aionui-cron/src/types.rs index b3f742442..bc4c63d9f 100644 --- a/crates/aionui-cron/src/types.rs +++ b/crates/aionui-cron/src/types.rs @@ -153,6 +153,7 @@ pub struct CronAgentConfig { #[derive(Debug, Clone, PartialEq)] pub struct CronJob { pub id: String, + pub user_id: String, pub name: String, pub enabled: bool, pub schedule: CronSchedule, @@ -202,6 +203,7 @@ pub fn cron_job_from_row(row: CronJobRow) -> Result { Ok(CronJob { id: row.id, + user_id: row.user_id, name: row.name, enabled: row.enabled, schedule, @@ -272,6 +274,7 @@ pub fn cron_job_to_row(job: &CronJob) -> Result { Ok(CronJobRow { id: job.id.clone(), + user_id: job.user_id.clone(), name: job.name.clone(), enabled: job.enabled, schedule_kind, @@ -556,6 +559,7 @@ mod tests { fn sample_row() -> CronJobRow { CronJobRow { id: "cron_test1".into(), + user_id: "user1".into(), name: "Test Job".into(), enabled: true, schedule_kind: "every".into(), @@ -586,6 +590,7 @@ mod tests { fn sample_job() -> CronJob { CronJob { id: "cron_test1".into(), + user_id: "user1".into(), name: "Test Job".into(), enabled: true, schedule: CronSchedule::Every { diff --git a/crates/aionui-cron/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index 61cb7b928..c63aad457 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -18,6 +18,7 @@ use aionui_api_types::{ CronScheduleDto, ListCronJobsQuery, SaveCronSkillRequest, UpdateConversationCronRequest, UpdateCronJobRequest, WebSocketMessage, }; +use aionui_auth::CurrentUser; use aionui_common::{PaginatedResult, ProviderWithModel, TimestampMs, now_ms}; use aionui_conversation::ConversationService; use aionui_cron::{CronRouterState, cron_routes}; @@ -26,12 +27,13 @@ use aionui_db::{ IAssistantDefinitionRepository, IAssistantOverlayRepository, IAssistantPreferenceRepository, IConversationRepository, ICronRepository, MessagePageParams, MessagePageResult, MessageRowUpdate, MessageSearchRow, SqliteAcpSessionRepository, SqliteAgentMetadataRepository, SqliteAssistantDefinitionRepository, - SqliteAssistantOverlayRepository, SqliteAssistantPreferenceRepository, SqliteCronRepository, - UpsertAgentMetadataParams, UpsertAssistantDefinitionParams, UpsertAssistantOverlayParams, - UpsertConversationAssistantSnapshotParams, init_database_memory, + SqliteAssistantOverlayRepository, SqliteAssistantPreferenceRepository, SqliteConversationRepository, + SqliteCronRepository, SqliteSkillRepository, UpsertAgentMetadataParams, UpsertAssistantDefinitionParams, + UpsertAssistantOverlayParams, UpsertConversationAssistantSnapshotParams, init_database_memory, models::{ConversationAssistantSnapshotRow, CronJobRow, MessageRow}, }; use aionui_realtime::EventBroadcaster; +use axum::Extension; use axum::body::{Body, to_bytes}; use axum::http::{Method, Request, StatusCode}; @@ -45,6 +47,46 @@ use tower::ServiceExt; // ── Test infrastructure ──────────────────────────────────────────── +fn current_user(id: &str) -> CurrentUser { + CurrentUser { + id: id.to_owned(), + username: id.to_owned(), + user_type: aionui_db::UserType::Local, + status: aionui_db::UserStatus::Active, + } +} + +async fn seed_sqlite_conversations(db: &aionui_db::Database, conversation_ids: &[&str]) { + sqlx::query( + "INSERT OR IGNORE INTO users (id, username, password_hash, created_at, updated_at) \ + VALUES ('u1', 'u1', 'hash', 0, 0)", + ) + .execute(db.pool()) + .await + .unwrap(); + + let repo = SqliteConversationRepository::new(db.pool().clone()); + for conversation_id in conversation_ids { + repo.create(&aionui_db::models::ConversationRow { + id: (*conversation_id).to_owned(), + user_id: "u1".to_owned(), + name: (*conversation_id).to_owned(), + r#type: "normal".to_owned(), + model: None, + status: Some("pending".to_owned()), + source: None, + channel_chat_id: None, + extra: "{}".to_owned(), + pinned: false, + pinned_at: None, + created_at: 0, + updated_at: 0, + }) + .await + .unwrap(); + } +} + fn ensure_named_workspace_path(name: &str) -> String { let workspace = std::env::temp_dir().join(name); std::fs::create_dir_all(&workspace).unwrap(); @@ -113,17 +155,31 @@ struct StubConvRepo { rows: Mutex>, assistant_snapshots: Mutex>, update_failures: Mutex>, + sqlite_pool: aionui_db::SqlitePool, } impl StubConvRepo { - fn new() -> Self { + fn new(sqlite_pool: aionui_db::SqlitePool) -> Self { Self { messages: Mutex::new(Vec::new()), artifacts: Mutex::new(Vec::new()), rows: Mutex::new(HashMap::new()), assistant_snapshots: Mutex::new(HashMap::new()), update_failures: Mutex::new(Vec::new()), + sqlite_pool, + } + } + + async fn seed_sqlite_row(&self, row: &aionui_db::models::ConversationRow) -> Result<(), aionui_db::DbError> { + let repo = SqliteConversationRepository::new(self.sqlite_pool.clone()); + if repo.get(&row.user_id, &row.id).await?.is_some() { + return Ok(()); + } + let mut sqlite_row = row.clone(); + if sqlite_row.status.as_deref() == Some("active") { + sqlite_row.status = Some("pending".to_owned()); } + repo.create(&sqlite_row).await } fn take_messages(&self) -> Vec { @@ -173,11 +229,17 @@ impl StubConvRepo { #[async_trait::async_trait] impl IConversationRepository for StubConvRepo { - async fn get(&self, id: &str) -> Result, aionui_db::DbError> { - let mut rows = self.rows.lock().unwrap(); - - if let Some(existing) = rows.get(id) { - return Ok(Some(existing.clone())); + async fn get( + &self, + user_id: &str, + id: &str, + ) -> Result, aionui_db::DbError> { + if let Some(existing) = { self.rows.lock().unwrap().get(id).cloned() } { + if existing.user_id != user_id { + return Ok(None); + } + self.seed_sqlite_row(&existing).await?; + return Ok(Some(existing)); } if id.starts_with("missing") { return Ok(None); @@ -477,12 +539,27 @@ impl IConversationRepository for StubConvRepo { } }; - rows.insert(id.to_owned(), row.clone()); + if row.user_id != user_id { + return Ok(None); + } + self.rows.lock().unwrap().insert(id.to_owned(), row.clone()); + self.seed_sqlite_row(&row).await?; Ok(Some(row)) } + async fn owner_user_id(&self, id: &str) -> Result, aionui_db::DbError> { + if let Some(existing) = self.rows.lock().unwrap().get(id) { + return Ok(Some(existing.user_id.clone())); + } + if id.starts_with("missing") { + return Ok(None); + } + Ok(Some("u1".into())) + } + async fn get_assistant_snapshot( &self, + _user_id: &str, conversation_id: &str, ) -> Result, aionui_db::DbError> { Ok(self.assistant_snapshots.lock().unwrap().get(conversation_id).cloned()) @@ -490,6 +567,7 @@ impl IConversationRepository for StubConvRepo { async fn upsert_assistant_snapshot( &self, + _user_id: &str, params: &UpsertConversationAssistantSnapshotParams<'_>, ) -> Result, aionui_db::DbError> { let now = now_ms(); @@ -525,7 +603,12 @@ impl IConversationRepository for StubConvRepo { self.rows.lock().unwrap().insert(row.id.clone(), row.clone()); Ok(()) } - async fn update(&self, id: &str, updates: &ConversationRowUpdate) -> Result<(), aionui_db::DbError> { + async fn update( + &self, + _user_id: &str, + id: &str, + updates: &ConversationRowUpdate, + ) -> Result<(), aionui_db::DbError> { if self.update_failures.lock().unwrap().iter().any(|item| item == id) { return Err(aionui_db::DbError::Init(format!("forced update failure for {id}"))); } @@ -556,7 +639,7 @@ impl IConversationRepository for StubConvRepo { } Ok(()) } - async fn delete(&self, _id: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), aionui_db::DbError> { Ok(()) } async fn list_paginated( @@ -609,6 +692,7 @@ impl IConversationRepository for StubConvRepo { } async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &MessagePageParams, ) -> Result { @@ -618,18 +702,29 @@ impl IConversationRepository for StubConvRepo { has_more_after: false, }) } - async fn insert_message(&self, message: &aionui_db::models::MessageRow) -> Result<(), aionui_db::DbError> { + async fn insert_message( + &self, + _user_id: &str, + message: &aionui_db::models::MessageRow, + ) -> Result<(), aionui_db::DbError> { self.messages.lock().unwrap().push(message.clone()); Ok(()) } - async fn update_message(&self, _id: &str, _updates: &MessageRowUpdate) -> Result<(), aionui_db::DbError> { + async fn update_message( + &self, + _user_id: &str, + _conversation_id: &str, + _id: &str, + _updates: &MessageRowUpdate, + ) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn delete_messages_by_conversation(&self, _conv_id: &str) -> Result<(), aionui_db::DbError> { + async fn delete_messages_by_conversation(&self, _user_id: &str, _conv_id: &str) -> Result<(), aionui_db::DbError> { Ok(()) } async fn get_message_by_msg_id( &self, + _user_id: &str, _conv_id: &str, _msg_id: &str, _msg_type: &str, @@ -651,6 +746,7 @@ impl IConversationRepository for StubConvRepo { } async fn list_artifacts( &self, + _user_id: &str, conversation_id: &str, ) -> Result, aionui_db::DbError> { Ok(self @@ -664,6 +760,7 @@ impl IConversationRepository for StubConvRepo { } async fn get_artifact( &self, + _user_id: &str, conversation_id: &str, artifact_id: &str, ) -> Result, aionui_db::DbError> { @@ -677,6 +774,7 @@ impl IConversationRepository for StubConvRepo { } async fn upsert_artifact( &self, + _user_id: &str, artifact: &aionui_db::ConversationArtifactRow, ) -> Result { self.upsert_artifact_row(artifact.clone()); @@ -684,6 +782,7 @@ impl IConversationRepository for StubConvRepo { } async fn update_artifact_status( &self, + _user_id: &str, conversation_id: &str, artifact_id: &str, status: &str, @@ -702,6 +801,7 @@ impl IConversationRepository for StubConvRepo { } async fn mark_skill_suggest_artifacts_saved( &self, + _user_id: &str, cron_job_id: &str, updated_at: TimestampMs, ) -> Result, aionui_db::DbError> { @@ -740,7 +840,7 @@ async fn setup_with_conv_runtime() -> ( Arc, Arc, ) { - let (svc, cron_repo, bc, stub_conv_repo, conv_service, _) = setup_with_conv_runtime_and_agent_metadata().await; + let (svc, cron_repo, bc, stub_conv_repo, conv_service, _, _) = setup_with_conv_runtime_and_agent_metadata().await; (svc, cron_repo, bc, stub_conv_repo, conv_service) } @@ -751,9 +851,11 @@ async fn setup_with_conv_runtime_and_agent_metadata() -> ( Arc, Arc, Arc, + Arc, ) { let db = init_database_memory().await.unwrap(); let pool = db.pool().clone(); + seed_sqlite_conversations(&db, &["conv_1"]).await; let cron_repo: Arc = Arc::new(SqliteCronRepository::new(pool.clone())); let agent_metadata_repo: Arc = Arc::new(SqliteAgentMetadataRepository::new(pool.clone())); @@ -763,7 +865,7 @@ async fn setup_with_conv_runtime_and_agent_metadata() -> ( Arc::new(SqliteAssistantOverlayRepository::new(pool.clone())); let assistant_preference_repo: Arc = Arc::new(SqliteAssistantPreferenceRepository::new(pool.clone())); - let acp_session_repo: Arc = Arc::new(SqliteAcpSessionRepository::new(pool)); + let acp_session_repo: Arc = Arc::new(SqliteAcpSessionRepository::new(pool.clone())); let bc = Arc::new(MockBroadcaster::new()); let data_dir = std::env::temp_dir().join(format!("aionui-cron-test-{}", now_ms())); std::fs::create_dir_all(&data_dir).unwrap(); @@ -792,7 +894,7 @@ async fn setup_with_conv_runtime_and_agent_metadata() -> ( } } - let stub_conv_repo = Arc::new(StubConvRepo::new()); + let stub_conv_repo = Arc::new(StubConvRepo::new(pool.clone())); let stub_conv_repo_trait: Arc = stub_conv_repo.clone(); let task_manager: Arc = Arc::new(StubTaskManager); let conv_service = Arc::new(ConversationService::new( @@ -822,11 +924,13 @@ async fn setup_with_conv_runtime_and_agent_metadata() -> ( let scheduler = Arc::new(CronScheduler::new(Arc::new(|_| {}))); let emitter = CronEventEmitter::new(bc.clone() as Arc); + let skill_repo: Arc = Arc::new(SqliteSkillRepository::new(pool.clone())); let svc = CronService::new(CronServiceDeps { repo: cron_repo.clone(), agent_metadata_repo: agent_metadata_repo.clone(), assistant_definition_repo: assistant_definition_repo.clone(), assistant_overlay_repo: assistant_overlay_repo.clone(), + skill_repo: skill_repo.clone(), scheduler, executor, emitter, @@ -843,7 +947,20 @@ async fn setup_with_conv_runtime_and_agent_metadata() -> ( seed_bare_assistant_definitions(&assistant_definition_repo).await; std::mem::forget(db); - (svc, cron_repo, bc, stub_conv_repo, conv_service, agent_metadata_repo) + ( + svc, + cron_repo, + bc, + stub_conv_repo, + conv_service, + agent_metadata_repo, + skill_repo, + ) +} + +async fn setup_with_skill_repo() -> (CronService, Arc) { + let (svc, _, _, _, _, _, skill_repo) = setup_with_conv_runtime_and_agent_metadata().await; + (svc, skill_repo) } async fn setup_with_assistant_repos() -> ( @@ -856,6 +973,7 @@ async fn setup_with_assistant_repos() -> ( ) { let db = init_database_memory().await.unwrap(); let pool = db.pool().clone(); + seed_sqlite_conversations(&db, &["conv_1"]).await; let cron_repo: Arc = Arc::new(SqliteCronRepository::new(pool.clone())); let agent_metadata_repo: Arc = Arc::new(SqliteAgentMetadataRepository::new(pool.clone())); @@ -865,7 +983,7 @@ async fn setup_with_assistant_repos() -> ( Arc::new(SqliteAssistantOverlayRepository::new(pool.clone())); let assistant_preference_repo: Arc = Arc::new(SqliteAssistantPreferenceRepository::new(pool.clone())); - let acp_session_repo: Arc = Arc::new(SqliteAcpSessionRepository::new(pool)); + let acp_session_repo: Arc = Arc::new(SqliteAcpSessionRepository::new(pool.clone())); let bc = Arc::new(MockBroadcaster::new()); let data_dir = std::env::temp_dir().join(format!("aionui-cron-test-{}", now_ms())); std::fs::create_dir_all(&data_dir).unwrap(); @@ -894,7 +1012,7 @@ async fn setup_with_assistant_repos() -> ( } } - let stub_conv_repo = Arc::new(StubConvRepo::new()); + let stub_conv_repo = Arc::new(StubConvRepo::new(pool.clone())); let stub_conv_repo_trait: Arc = stub_conv_repo.clone(); let task_manager: Arc = Arc::new(StubTaskManager); let conv_service = Arc::new(ConversationService::new( @@ -928,6 +1046,7 @@ async fn setup_with_assistant_repos() -> ( agent_metadata_repo, assistant_definition_repo: assistant_definition_repo.clone(), assistant_overlay_repo: assistant_overlay_repo.clone(), + skill_repo: Arc::new(SqliteSkillRepository::new(pool.clone())), scheduler, executor, emitter, @@ -979,6 +1098,52 @@ fn make_create_req(name: &str, schedule: CronScheduleDto) -> CreateCronJobReques } } +fn make_cron_row_with_workspace(id: &str, user_id: &str, conversation_id: &str, workspace: &str) -> CronJobRow { + let now = now_ms(); + CronJobRow { + id: id.into(), + user_id: user_id.into(), + name: id.into(), + enabled: true, + schedule_kind: "every".into(), + schedule_value: "60000".into(), + schedule_tz: None, + schedule_description: Some("every minute".into()), + payload_message: "test message".into(), + execution_mode: "existing".into(), + agent_config: Some( + serde_json::to_string(&CronAgentConfig { + name: "Default Assistant".into(), + cli_path: None, + is_preset: None, + assistant_id: Some("assistant-default".into()), + custom_agent_id: None, + mode: Some("default".into()), + model_id: Some("claude-sonnet-4".into()), + model: None, + config_options: None, + workspace: Some(workspace.into()), + }) + .unwrap(), + ), + conversation_id: conversation_id.into(), + conversation_title: Some("Test Conv".into()), + created_by: "user".into(), + skill_content: None, + description: None, + created_at: now, + updated_at: now, + next_run_at: Some(now + 60_000), + last_run_at: None, + last_status: None, + last_error: None, + run_count: 0, + retry_count: 0, + max_retries: 3, + queue_enabled: false, + } +} + async fn seed_assistant_definition( repo: &Arc, definition_id: &str, @@ -986,36 +1151,39 @@ async fn seed_assistant_definition( agent_backend: &str, ) { let agent_id = seeded_agent_id(agent_backend); - repo.upsert(&UpsertAssistantDefinitionParams { - id: definition_id, - assistant_id, - source: "user", - owner_type: "user", - source_ref: Some(assistant_id), - name: assistant_id, - name_i18n: "{}", - description: Some("test assistant"), - description_i18n: "{}", - avatar_type: "emoji", - avatar_value: Some("🤖"), - agent_id, - rule_resource_type: "user_file", - rule_resource_ref: None, - recommended_prompts: "[]", - recommended_prompts_i18n: "{}", - default_model_mode: "auto", - default_model_value: None, - default_permission_mode: "auto", - default_permission_value: None, - default_thought_level_mode: "auto", - default_thought_level_value: None, - default_skills_mode: "auto", - default_skill_ids: "[]", - custom_skill_names: "[]", - default_disabled_builtin_skill_ids: "[]", - default_mcps_mode: "auto", - default_mcp_ids: "[]", - }) + repo.upsert_for_user( + "u1", + &UpsertAssistantDefinitionParams { + id: definition_id, + assistant_id, + source: "user", + owner_type: "user", + source_ref: Some(assistant_id), + name: assistant_id, + name_i18n: "{}", + description: Some("test assistant"), + description_i18n: "{}", + avatar_type: "emoji", + avatar_value: Some("🤖"), + agent_id, + rule_resource_type: "user_file", + rule_resource_ref: None, + recommended_prompts: "[]", + recommended_prompts_i18n: "{}", + default_model_mode: "auto", + default_model_value: None, + default_permission_mode: "auto", + default_permission_value: None, + default_thought_level_mode: "auto", + default_thought_level_value: None, + default_skills_mode: "auto", + default_skill_ids: "[]", + custom_skill_names: "[]", + default_disabled_builtin_skill_ids: "[]", + default_mcps_mode: "auto", + default_mcp_ids: "[]", + }, + ) .await .unwrap(); } @@ -1036,13 +1204,16 @@ async fn seed_assistant_overlay( agent_backend_override: Option<&str>, ) { let agent_id_override = agent_backend_override.map(seeded_agent_id); - repo.upsert(&UpsertAssistantOverlayParams { - assistant_definition_id: definition_id, - enabled: true, - sort_order: 0, - agent_id_override, - last_used_at: None, - }) + repo.upsert_for_user( + "u1", + &UpsertAssistantOverlayParams { + assistant_definition_id: definition_id, + enabled: true, + sort_order: 0, + agent_id_override, + last_used_at: None, + }, + ) .await .unwrap(); } @@ -1086,7 +1257,7 @@ async fn cj1_create_cron_job() { let (svc, _, bc) = setup().await; let req = make_create_req("Daily Report", every_60s()); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); assert!(job.id.starts_with("cron_")); assert_eq!(job.name, "Daily Report"); @@ -1105,10 +1276,10 @@ async fn create_job_allows_missing_task_description() { let mut req = make_create_req("No Description", every_60s()); req.description = None; - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); assert_eq!(job.description, None); - let row = cron_repo.get_by_id(&job.id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&job.id).await.unwrap().unwrap(); assert_eq!(row.description, None); } @@ -1128,7 +1299,7 @@ async fn create_job_strips_legacy_agent_ids_when_assistant_id_present() { workspace: None, }); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); let config = job.agent_config.expect("agent config"); assert_eq!(config.assistant_id.as_deref(), Some("assistant-1")); @@ -1141,7 +1312,7 @@ async fn create_job_uses_assistant_full_auto_mode_instead_of_requested_mode() { let (svc, _, _) = setup().await; let req = make_create_req("Full Auto Mode", every_60s()); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); let config = job.agent_config.expect("agent config"); assert_eq!(config.assistant_id.as_deref(), Some("assistant-default")); @@ -1164,7 +1335,7 @@ async fn create_job_derives_assistant_runtime_without_backend_hint() { workspace: None, }); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); assert_eq!(job.agent_type, "acp"); } @@ -1175,7 +1346,7 @@ async fn create_job_requires_assistant_id_for_new_jobs() { let mut req = make_create_req("Missing Runtime Type", every_60s()); req.agent_config = None; - let err = svc.add_job(req).await.unwrap_err(); + let err = svc.add_job("u1", req).await.unwrap_err(); assert!(matches!(err, aionui_cron::error::CronError::InvalidAgentConfig(_))); assert!(err.to_string().contains("assistant_id is required for new cron jobs")); } @@ -1201,7 +1372,7 @@ async fn create_job_derives_runtime_type_from_aionrs_assistant() { workspace: None, }); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); assert_eq!(job.agent_type, "aionrs"); } @@ -1234,7 +1405,7 @@ async fn create_job_derives_runtime_type_from_assistant_overlay_override() { workspace: None, }); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); assert_eq!(job.agent_type, "aionrs"); } @@ -1256,7 +1427,7 @@ async fn create_job_allows_assistant_backed_acp_jobs_without_backend_hint() { workspace: None, }); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); let config = job.agent_config.expect("agent config"); assert_eq!(job.agent_type, "acp"); @@ -1280,7 +1451,7 @@ async fn create_job_rejects_backend_fallback_when_assistant_id_cannot_resolve() }); let err = svc - .add_job(req) + .add_job("u1", req) .await .expect_err("missing assistant must not fall back to backend"); @@ -1296,17 +1467,20 @@ async fn cj2_create_three_schedule_types() { let now = now_ms(); let at_job = svc - .add_job(make_create_req("At Job", at_future(3600000))) + .add_job("u1", make_create_req("At Job", at_future(3600000))) .await .unwrap(); assert!(at_job.next_run_at.unwrap() > now); - let every_job = svc.add_job(make_create_req("Every Job", every_60s())).await.unwrap(); + let every_job = svc + .add_job("u1", make_create_req("Every Job", every_60s())) + .await + .unwrap(); let next = every_job.next_run_at.unwrap(); assert!((next - now - 60000).abs() < 2000); let cron_job = svc - .add_job(make_create_req("Cron Job", cron_every_5min())) + .add_job("u1", make_create_req("Cron Job", cron_every_5min())) .await .unwrap(); assert!(cron_job.next_run_at.unwrap() > now); @@ -1317,9 +1491,12 @@ async fn cj2_create_three_schedule_types() { #[tokio::test] async fn cj4_get_single_job() { let (svc, _, _) = setup().await; - let created = svc.add_job(make_create_req("Get Test", every_60s())).await.unwrap(); + let created = svc + .add_job("u1", make_create_req("Get Test", every_60s())) + .await + .unwrap(); - let fetched = svc.get_job(&created.id).await.unwrap(); + let fetched = svc.get_job("u1", &created.id).await.unwrap(); assert_eq!(fetched.id, created.id); assert_eq!(fetched.name, "Get Test"); } @@ -1329,7 +1506,7 @@ async fn cj4_get_single_job() { #[tokio::test] async fn cj5_get_nonexistent_job() { let (svc, _, _) = setup().await; - let err = svc.get_job("cron_nonexistent").await.unwrap_err(); + let err = svc.get_job("u1", "cron_nonexistent").await.unwrap_err(); assert!(matches!(err, aionui_cron::error::CronError::JobNotFound(_))); } @@ -1339,12 +1516,12 @@ async fn cj5_get_nonexistent_job() { async fn cj6_list_all_jobs() { let (svc, _, _) = setup().await; for i in 0..3 { - svc.add_job(make_create_req(&format!("Job {i}"), every_60s())) + svc.add_job("u1", make_create_req(&format!("Job {i}"), every_60s())) .await .unwrap(); } - let jobs = svc.list_jobs(&ListCronJobsQuery::default()).await.unwrap(); + let jobs = svc.list_jobs("u1", &ListCronJobsQuery::default()).await.unwrap(); assert!(jobs.len() >= 3); } @@ -1354,6 +1531,7 @@ async fn list_jobs_allows_legacy_custom_agent_id_without_assistant_id() { cron_repo .insert(&CronJobRow { id: "cron_legacy_custom_agent".into(), + user_id: "u1".into(), name: "Legacy custom agent job".into(), enabled: true, schedule_kind: "every".into(), @@ -1370,7 +1548,7 @@ async fn list_jobs_allows_legacy_custom_agent_id_without_assistant_id() { }) .to_string(), ), - conversation_id: "conv_legacy".into(), + conversation_id: "conv_1".into(), conversation_title: None, created_by: "user".into(), skill_content: None, @@ -1389,7 +1567,7 @@ async fn list_jobs_allows_legacy_custom_agent_id_without_assistant_id() { .await .unwrap(); - let jobs = svc.list_jobs(&ListCronJobsQuery::default()).await.unwrap(); + let jobs = svc.list_jobs("u1", &ListCronJobsQuery::default()).await.unwrap(); let legacy = jobs .iter() @@ -1406,20 +1584,20 @@ async fn cj7_list_by_conversation() { let mut req1 = make_create_req("Job A", every_60s()); req1.conversation_id = "conv_target".into(); - svc.add_job(req1).await.unwrap(); + svc.add_job("u1", req1).await.unwrap(); let mut req2 = make_create_req("Job B", every_60s()); req2.conversation_id = "conv_target".into(); - svc.add_job(req2).await.unwrap(); + svc.add_job("u1", req2).await.unwrap(); let mut req3 = make_create_req("Job C", every_60s()); req3.conversation_id = "conv_other".into(); - svc.add_job(req3).await.unwrap(); + svc.add_job("u1", req3).await.unwrap(); let query = ListCronJobsQuery { conversation_id: Some("conv_target".into()), }; - let jobs = svc.list_jobs(&query).await.unwrap(); + let jobs = svc.list_jobs("u1", &query).await.unwrap(); assert_eq!(jobs.len(), 2); } @@ -1430,9 +1608,9 @@ async fn cj7b_add_job_binds_existing_conversation_to_job() { let mut req = make_create_req("Bound Existing Conversation", every_60s()); req.conversation_id = "conv_existing_bind".into(); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); - let bound = conv_repo.get("conv_existing_bind").await.unwrap().unwrap(); + let bound = conv_repo.get("u1", "conv_existing_bind").await.unwrap().unwrap(); let extra: serde_json::Value = serde_json::from_str(&bound.extra).unwrap(); assert_eq!(extra["cron_job_id"], job.id); assert_eq!(extra["cronJobId"], job.id); @@ -1447,7 +1625,10 @@ async fn cj7b_add_job_binds_existing_conversation_to_job() { #[tokio::test] async fn cj8_update_job() { let (svc, _, bc) = setup().await; - let created = svc.add_job(make_create_req("Original", every_60s())).await.unwrap(); + let created = svc + .add_job("u1", make_create_req("Original", every_60s())) + .await + .unwrap(); bc.take_events(); let req = UpdateCronJobRequest { @@ -1463,7 +1644,7 @@ async fn cj8_update_job() { queue_enabled: None, }; - let updated = svc.update_job(&created.id, req).await.unwrap(); + let updated = svc.update_job("u1", &created.id, req).await.unwrap(); assert_eq!(updated.name, "Updated Name"); assert_eq!(updated.description.as_deref(), Some("Updated description")); assert!(!updated.enabled); @@ -1478,7 +1659,10 @@ async fn cj8_update_job() { async fn update_existing_conversation_job_rejects_agent_config_changes() { let (svc, _, _) = setup().await; let created = svc - .add_job(make_create_req("Existing Conversation Assistant Lock", every_60s())) + .add_job( + "u1", + make_create_req("Existing Conversation Assistant Lock", every_60s()), + ) .await .unwrap(); @@ -1504,7 +1688,7 @@ async fn update_existing_conversation_job_rejects_agent_config_changes() { queue_enabled: None, }; - let err = svc.update_job(&created.id, req).await.unwrap_err(); + let err = svc.update_job("u1", &created.id, req).await.unwrap_err(); assert!( matches!(err, aionui_cron::error::CronError::InvalidAgentConfig(message) if message.contains("ongoing conversation")) ); @@ -1514,10 +1698,10 @@ async fn update_existing_conversation_job_rejects_agent_config_changes() { async fn update_existing_conversation_job_rejects_agent_config_even_when_switching_to_new_conversation() { let (svc, _, _) = setup().await; let created = svc - .add_job(make_create_req( - "Existing Conversation Assistant Lock Mode Switch", - every_60s(), - )) + .add_job( + "u1", + make_create_req("Existing Conversation Assistant Lock Mode Switch", every_60s()), + ) .await .unwrap(); @@ -1543,29 +1727,28 @@ async fn update_existing_conversation_job_rejects_agent_config_even_when_switchi queue_enabled: None, }; - let err = svc.update_job(&created.id, req).await.unwrap_err(); + let err = svc.update_job("u1", &created.id, req).await.unwrap_err(); assert!( matches!(err, aionui_cron::error::CronError::InvalidAgentConfig(message) if message.contains("ongoing conversation")) ); } #[tokio::test] -async fn update_existing_job_to_new_conversation_removes_previous_conversation_binding() { - use aionui_common::OnConversationDelete; - +async fn update_existing_job_to_new_conversation_keeps_owner_anchor() { let (svc, cron_repo, _, conv_repo) = setup_with_conv_repo().await; let mut create_req = make_create_req("Mode Switch Clears Binding", every_60s()); create_req.conversation_id = "conv_mode_switch".into(); create_req.execution_mode = Some("existing".into()); - let created = svc.add_job(create_req).await.unwrap(); + let created = svc.add_job("u1", create_req).await.unwrap(); - let bound_before = conv_repo.get("conv_mode_switch").await.unwrap().unwrap(); + let bound_before = conv_repo.get("u1", "conv_mode_switch").await.unwrap().unwrap(); let extra_before: serde_json::Value = serde_json::from_str(&bound_before.extra).unwrap(); assert_eq!(extra_before["cron_job_id"], created.id); assert_eq!(extra_before["cronJobId"], created.id); let updated = svc .update_job( + "u1", &created.id, UpdateCronJobRequest { name: None, @@ -1584,21 +1767,15 @@ async fn update_existing_job_to_new_conversation_removes_previous_conversation_b .unwrap(); assert_eq!(updated.execution_mode.as_str(), "new_conversation"); - let row = cron_repo.get_by_id(&created.id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&created.id).await.unwrap().unwrap(); assert_eq!(row.execution_mode, "new_conversation"); - assert_eq!(row.conversation_id, ""); + assert_eq!(row.conversation_id, "conv_mode_switch"); assert!(row.conversation_title.is_none()); - let bound_after = conv_repo.get("conv_mode_switch").await.unwrap().unwrap(); + let bound_after = conv_repo.get("u1", "conv_mode_switch").await.unwrap().unwrap(); let extra_after: serde_json::Value = serde_json::from_str(&bound_after.extra).unwrap(); assert!(extra_after.get("cron_job_id").is_none()); assert!(extra_after.get("cronJobId").is_none()); - - svc.on_conversation_deleted("conv_mode_switch").await; - assert!( - svc.get_job(&created.id).await.is_ok(), - "deleting the previous existing-mode conversation must not delete the switched new-conversation job" - ); } #[tokio::test] @@ -1621,9 +1798,9 @@ async fn update_existing_job_to_new_conversation_clears_previous_auto_workspace( create_req.conversation_id = conversation_id.clone(); create_req.execution_mode = Some("existing".into()); create_req.agent_config.as_mut().unwrap().workspace = Some(auto_workspace); - let created = svc.add_job(create_req).await.unwrap(); + let created = svc.add_job("u1", create_req).await.unwrap(); - let bound_conversation = conv_repo.get(&conversation_id).await.unwrap().unwrap(); + let bound_conversation = conv_repo.get("u1", &conversation_id).await.unwrap().unwrap(); assert!( conv_service .auto_workspace_to_delete_for_row(&bound_conversation, &conversation_id) @@ -1632,6 +1809,7 @@ async fn update_existing_job_to_new_conversation_clears_previous_auto_workspace( ); svc.update_job( + "u1", &created.id, UpdateCronJobRequest { name: None, @@ -1649,7 +1827,7 @@ async fn update_existing_job_to_new_conversation_clears_previous_auto_workspace( .await .unwrap(); - let row = cron_repo.get_by_id(&created.id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&created.id).await.unwrap().unwrap(); let config: CronAgentConfig = serde_json::from_str(row.agent_config.as_deref().unwrap()).unwrap(); assert!( config.workspace.is_none(), @@ -1673,9 +1851,10 @@ async fn update_existing_job_to_new_conversation_preserves_custom_workspace() { create_req.conversation_id = conversation_id.clone(); create_req.execution_mode = Some("existing".into()); create_req.agent_config.as_mut().unwrap().workspace = Some(custom_workspace.clone()); - let created = svc.add_job(create_req).await.unwrap(); + let created = svc.add_job("u1", create_req).await.unwrap(); svc.update_job( + "u1", &created.id, UpdateCronJobRequest { name: None, @@ -1693,7 +1872,7 @@ async fn update_existing_job_to_new_conversation_preserves_custom_workspace() { .await .unwrap(); - let row = cron_repo.get_by_id(&created.id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&created.id).await.unwrap().unwrap(); let config: CronAgentConfig = serde_json::from_str(row.agent_config.as_deref().unwrap()).unwrap(); assert_eq!(config.workspace.as_deref(), Some(custom_workspace.as_str())); } @@ -1710,7 +1889,7 @@ async fn update_team_conversation_job_rejects_execution_mode_change() { "workspace": ensure_named_workspace_path("aionui-cron-service-team-workspace") }), ); - let created = svc.add_job(create_req).await.unwrap(); + let created = svc.add_job("u1", create_req).await.unwrap(); let req = UpdateCronJobRequest { name: None, @@ -1725,7 +1904,7 @@ async fn update_team_conversation_job_rejects_execution_mode_change() { queue_enabled: None, }; - let err = svc.update_job(&created.id, req).await.unwrap_err(); + let err = svc.update_job("u1", &created.id, req).await.unwrap_err(); assert!(matches!(err, aionui_cron::error::CronError::InvalidExecutionMode(message) if message.contains("Team"))); } @@ -1735,7 +1914,7 @@ async fn update_job_strips_legacy_agent_ids_when_assistant_id_present() { seed_assistant_definition(&definition_repo, "asstdef_update_assistant_1", "assistant-1", "claude").await; let mut create_req = make_create_req("Assistant Only Update", every_60s()); create_req.execution_mode = Some("new_conversation".into()); - let created = svc.add_job(create_req).await.unwrap(); + let created = svc.add_job("u1", create_req).await.unwrap(); let req = UpdateCronJobRequest { name: None, @@ -1759,7 +1938,7 @@ async fn update_job_strips_legacy_agent_ids_when_assistant_id_present() { queue_enabled: None, }; - let updated = svc.update_job(&created.id, req).await.unwrap(); + let updated = svc.update_job("u1", &created.id, req).await.unwrap(); let config = updated.agent_config.expect("agent config"); assert_eq!(config.assistant_id.as_deref(), Some("assistant-1")); @@ -1772,7 +1951,7 @@ async fn update_job_rejects_when_assistant_id_cannot_resolve() { let (svc, _, _) = setup().await; let mut create_req = make_create_req("Assistant Missing Update", every_60s()); create_req.execution_mode = Some("new_conversation".into()); - let created = svc.add_job(create_req).await.unwrap(); + let created = svc.add_job("u1", create_req).await.unwrap(); let req = UpdateCronJobRequest { name: None, @@ -1797,7 +1976,7 @@ async fn update_job_rejects_when_assistant_id_cannot_resolve() { }; let err = svc - .update_job(&created.id, req) + .update_job("u1", &created.id, req) .await .expect_err("missing assistant must not fall back to backend"); @@ -1814,7 +1993,7 @@ async fn update_job_rejects_when_assistant_id_cannot_resolve() { async fn cj9_update_schedule_type() { let (svc, _, _) = setup().await; let created = svc - .add_job(make_create_req("Schedule Change", every_60s())) + .add_job("u1", make_create_req("Schedule Change", every_60s())) .await .unwrap(); @@ -1831,7 +2010,7 @@ async fn cj9_update_schedule_type() { queue_enabled: None, }; - let updated = svc.update_job(&created.id, req).await.unwrap(); + let updated = svc.update_job("u1", &created.id, req).await.unwrap(); assert!(matches!( updated.schedule, aionui_cron::types::CronSchedule::Cron { .. } @@ -1856,7 +2035,7 @@ async fn cj10_update_nonexistent() { max_retries: None, queue_enabled: None, }; - let err = svc.update_job("cron_nonexistent", req).await.unwrap_err(); + let err = svc.update_job("u1", "cron_nonexistent", req).await.unwrap_err(); assert!(matches!(err, aionui_cron::error::CronError::JobNotFound(_))); } @@ -1865,12 +2044,15 @@ async fn cj10_update_nonexistent() { #[tokio::test] async fn cj11_delete_job() { let (svc, _, bc) = setup().await; - let created = svc.add_job(make_create_req("To Delete", every_60s())).await.unwrap(); + let created = svc + .add_job("u1", make_create_req("To Delete", every_60s())) + .await + .unwrap(); bc.take_events(); - svc.remove_job(&created.id).await.unwrap(); + svc.remove_job("u1", &created.id).await.unwrap(); - let err = svc.get_job(&created.id).await.unwrap_err(); + let err = svc.get_job("u1", &created.id).await.unwrap_err(); assert!(matches!(err, aionui_cron::error::CronError::JobNotFound(_))); let events = bc.take_events(); @@ -1883,7 +2065,7 @@ async fn cj11_delete_job() { #[tokio::test] async fn cj12_delete_nonexistent() { let (svc, _, _) = setup().await; - let err = svc.remove_job("cron_nonexistent").await.unwrap_err(); + let err = svc.remove_job("u1", "cron_nonexistent").await.unwrap_err(); assert!(matches!( err, aionui_cron::error::CronError::Database(aionui_db::DbError::NotFound(_)) @@ -1895,19 +2077,22 @@ async fn cj12_delete_nonexistent() { #[tokio::test] async fn sk1_save_skill() { let (svc, _, _) = setup().await; - let job = svc.add_job(make_create_req("Skill Job", every_60s())).await.unwrap(); + let job = svc + .add_job("u1", make_create_req("Skill Job", every_60s())) + .await + .unwrap(); let req = SaveCronSkillRequest { content: "---\nname: test\ndescription: test skill\n---\nDo something".into(), }; - svc.save_skill(&job.id, req).await.unwrap(); + svc.save_skill("u1", &job.id, req).await.unwrap(); } #[tokio::test] async fn sk1_1_save_skill_marks_related_skill_suggest_artifacts_saved() { let (svc, _, bc, conv_repo) = setup_with_conv_repo().await; let job = svc - .add_job(make_create_req("Skill Artifact Job", every_60s())) + .add_job("u1", make_create_req("Skill Artifact Job", every_60s())) .await .unwrap(); @@ -1929,6 +2114,7 @@ async fn sk1_1_save_skill_marks_related_skill_suggest_artifacts_saved() { }); svc.save_skill( + "u1", &job.id, SaveCronSkillRequest { content: "---\nname: daily-report\ndescription: Daily report\n---\nUse it.".into(), @@ -1958,9 +2144,13 @@ async fn sk1_1_save_skill_marks_related_skill_suggest_artifacts_saved() { #[tokio::test] async fn sk2_has_skill_true() { let (svc, _, _) = setup().await; - let job = svc.add_job(make_create_req("Skill Check", every_60s())).await.unwrap(); + let job = svc + .add_job("u1", make_create_req("Skill Check", every_60s())) + .await + .unwrap(); svc.save_skill( + "u1", &job.id, SaveCronSkillRequest { content: "---\nname: x\n---\nContent".into(), @@ -1969,7 +2159,7 @@ async fn sk2_has_skill_true() { .await .unwrap(); - let resp = svc.has_skill(&job.id).await.unwrap(); + let resp = svc.has_skill("u1", &job.id).await.unwrap(); assert!(resp.has_skill); } @@ -1978,9 +2168,12 @@ async fn sk2_has_skill_true() { #[tokio::test] async fn sk3_has_skill_false() { let (svc, _, _) = setup().await; - let job = svc.add_job(make_create_req("No Skill", every_60s())).await.unwrap(); + let job = svc + .add_job("u1", make_create_req("No Skill", every_60s())) + .await + .unwrap(); - let resp = svc.has_skill(&job.id).await.unwrap(); + let resp = svc.has_skill("u1", &job.id).await.unwrap(); assert!(!resp.has_skill); } @@ -1989,10 +2182,13 @@ async fn sk3_has_skill_false() { #[tokio::test] async fn sk4_save_empty_skill() { let (svc, _, _) = setup().await; - let job = svc.add_job(make_create_req("Empty Skill", every_60s())).await.unwrap(); + let job = svc + .add_job("u1", make_create_req("Empty Skill", every_60s())) + .await + .unwrap(); let err = svc - .save_skill(&job.id, SaveCronSkillRequest { content: "".into() }) + .save_skill("u1", &job.id, SaveCronSkillRequest { content: "".into() }) .await .unwrap_err(); assert!(matches!(err, aionui_cron::error::CronError::InvalidSkillContent(_))); @@ -2004,12 +2200,13 @@ async fn sk4_save_empty_skill() { async fn sk5_save_placeholder_skill() { let (svc, _, _) = setup().await; let job = svc - .add_job(make_create_req("Placeholder Skill", every_60s())) + .add_job("u1", make_create_req("Placeholder Skill", every_60s())) .await .unwrap(); let err = svc .save_skill( + "u1", &job.id, SaveCronSkillRequest { content: "TODO: fill in later".into(), @@ -2027,6 +2224,7 @@ async fn sk6_save_skill_nonexistent() { let (svc, _, _) = setup().await; let err = svc .save_skill( + "u1", "cron_nonexistent", SaveCronSkillRequest { content: "---\nname: x\n---\nOk".into(), @@ -2043,10 +2241,11 @@ async fn sk6_save_skill_nonexistent() { async fn sk7_delete_cleans_skill() { let (svc, _, _) = setup().await; let job = svc - .add_job(make_create_req("Skill Cleanup", every_60s())) + .add_job("u1", make_create_req("Skill Cleanup", every_60s())) .await .unwrap(); svc.save_skill( + "u1", &job.id, SaveCronSkillRequest { content: "---\nname: x\n---\nContent".into(), @@ -2055,19 +2254,172 @@ async fn sk7_delete_cleans_skill() { .await .unwrap(); - svc.remove_job(&job.id).await.unwrap(); + svc.remove_job("u1", &job.id).await.unwrap(); - let err = svc.has_skill(&job.id).await.unwrap_err(); + let err = svc.has_skill("u1", &job.id).await.unwrap_err(); assert!(matches!(err, aionui_cron::error::CronError::JobNotFound(_))); } +// ── SK-8..11: Owner-scoped cron skill rows ──────────────────────── + +#[tokio::test] +async fn sk8_save_skill_upserts_owner_scoped_row() { + let (svc, skill_repo) = setup_with_skill_repo().await; + let job = svc + .add_job("u1", make_create_req("Skill Row Job", every_60s())) + .await + .unwrap(); + + svc.save_skill( + "u1", + &job.id, + SaveCronSkillRequest { + content: "---\nname: analyze\ndescription: Analyze numbers\n---\nDo the analysis".into(), + }, + ) + .await + .unwrap(); + + let row_name = format!("cron-{}", job.id); + let row = skill_repo + .find_by_name_for_user("u1", &row_name) + .await + .unwrap() + .expect("owner-scoped cron skill row must exist after save"); + assert_eq!(row.user_id.as_deref(), Some("u1")); + assert_eq!(row.source, "cron"); + assert_eq!(row.description.as_deref(), Some("Analyze numbers")); + assert!( + row.path.ends_with(&row_name), + "row path should be the cron skill dir: {}", + row.path + ); +} + +#[tokio::test] +async fn sk9_delete_skill_removes_owner_row() { + let (svc, skill_repo) = setup_with_skill_repo().await; + let job = svc + .add_job("u1", make_create_req("Skill Row Delete", every_60s())) + .await + .unwrap(); + svc.save_skill( + "u1", + &job.id, + SaveCronSkillRequest { + content: "---\nname: x\ndescription: d\n---\nContent".into(), + }, + ) + .await + .unwrap(); + + svc.delete_skill("u1", &job.id).await.unwrap(); + + let row_name = format!("cron-{}", job.id); + assert!( + skill_repo + .find_by_name_for_user("u1", &row_name) + .await + .unwrap() + .is_none(), + "active row must be gone after delete_skill" + ); +} + +#[tokio::test] +async fn sk10_remove_job_removes_owner_row() { + let (svc, skill_repo) = setup_with_skill_repo().await; + let job = svc + .add_job("u1", make_create_req("Skill Row Remove Job", every_60s())) + .await + .unwrap(); + svc.save_skill( + "u1", + &job.id, + SaveCronSkillRequest { + content: "---\nname: x\ndescription: d\n---\nContent".into(), + }, + ) + .await + .unwrap(); + + svc.remove_job("u1", &job.id).await.unwrap(); + + let row_name = format!("cron-{}", job.id); + assert!( + skill_repo + .find_by_name_for_user("u1", &row_name) + .await + .unwrap() + .is_none(), + "active row must be gone after remove_job" + ); +} + +#[tokio::test] +async fn sk11_init_reconciles_owner_rows_and_cleans_legacy_names() { + let (svc, skill_repo) = setup_with_skill_repo().await; + let job = svc + .add_job("u1", make_create_req("Skill Row Reconcile", every_60s())) + .await + .unwrap(); + svc.save_skill( + "u1", + &job.id, + SaveCronSkillRequest { + content: "---\nname: analyze\ndescription: Analyze numbers\n---\nDo the analysis".into(), + }, + ) + .await + .unwrap(); + + let row_name = format!("cron-{}", job.id); + // Simulate a fresh database: the file exists but the row is missing. + skill_repo.delete_by_name_for_user("u1", &row_name).await.unwrap(); + // Simulate a legacy machine-level scan row named by SKILL.md frontmatter. + skill_repo + .upsert_for_user( + "u1", + aionui_db::UpsertSkillParams { + name: "analyze", + description: Some("Analyze numbers"), + path: "/legacy/path", + source: "cron", + enabled: true, + }, + ) + .await + .unwrap(); + + svc.init().await; + + let row = skill_repo + .find_by_name_for_user("u1", &row_name) + .await + .unwrap() + .expect("reconcile must restore the owner-scoped row"); + assert_eq!(row.user_id.as_deref(), Some("u1")); + assert_eq!(row.source, "cron"); + assert!( + skill_repo + .find_by_name_for_user("u1", "analyze") + .await + .unwrap() + .is_none(), + "legacy frontmatter-named cron row must be soft-deleted by reconcile" + ); +} + // ── SC-3: Every type next_run ───────────────────────────────────── #[tokio::test] async fn sc3_every_type_next_run() { let (svc, _, _) = setup().await; let now = now_ms(); - let job = svc.add_job(make_create_req("Every 60s", every_60s())).await.unwrap(); + let job = svc + .add_job("u1", make_create_req("Every 60s", every_60s())) + .await + .unwrap(); let next = job.next_run_at.unwrap(); let diff = (next - now - 60000).abs(); @@ -2087,7 +2439,7 @@ async fn sc5_invalid_cron_expression() { description: None, }, ); - let err = svc.add_job(req).await.unwrap_err(); + let err = svc.add_job("u1", req).await.unwrap_err(); assert!(matches!(err, aionui_cron::error::CronError::InvalidCronExpression(_))); } @@ -2105,7 +2457,7 @@ async fn sc6_cron_with_timezone() { description: None, }, ); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); assert!(job.next_run_at.unwrap() > now); } @@ -2121,7 +2473,7 @@ async fn sc7_every_zero_interval() { description: None, }, ); - let err = svc.add_job(req).await.unwrap_err(); + let err = svc.add_job("u1", req).await.unwrap_err(); assert!(matches!(err, aionui_cron::error::CronError::InvalidSchedule(_))); } @@ -2137,111 +2489,101 @@ async fn sc8_every_negative_interval() { description: None, }, ); - let err = svc.add_job(req).await.unwrap_err(); + let err = svc.add_job("u1", req).await.unwrap_err(); assert!(matches!(err, aionui_cron::error::CronError::InvalidSchedule(_))); } -// ── OC-1: Init preserves lazy-bind "existing" jobs with empty conversation_id ───── +// ── OC-1: Creation rejects jobs without a scoped parent conversation ───── #[tokio::test] -async fn oc1_init_preserves_lazy_existing_jobs() { - // "existing + empty conversation_id" is a legitimate lazy-binding job: - // the frontend creates a cron from the standalone cron page before any - // conversation exists, and the first execution materializes it. Those - // jobs must survive init, not be cleaned up as orphans. +async fn oc1_rejects_lazy_existing_jobs() { let (svc, _repo, _) = setup().await; let mut req = make_create_req("Lazy Existing", every_60s()); req.conversation_id = "".into(); req.execution_mode = Some("existing".into()); - let lazy = svc.add_job(req).await.unwrap(); - - let normal_req = make_create_req("Normal", every_60s()); - let normal = svc.add_job(normal_req).await.unwrap(); - - svc.init().await; - - let found_lazy = svc.get_job(&lazy.id).await; - assert!(found_lazy.is_ok(), "lazy-bind existing job should survive init"); - - let found = svc.get_job(&normal.id).await; - assert!(found.is_ok()); + let err = svc.add_job("u1", req).await.unwrap_err(); + assert!(matches!(err, aionui_cron::error::CronError::Conversation(_))); } -// NewConversation jobs don't depend on any existing conversation — they -// create one on every run. They must never be cleaned up as orphans. #[tokio::test] -async fn oc1b_init_preserves_new_conversation_jobs() { +async fn oc1b_allows_new_conversation_jobs_without_owner_anchor() { let (svc, _repo, _) = setup().await; let mut empty_req = make_create_req("New-conv empty", every_60s()); empty_req.conversation_id = "".into(); empty_req.execution_mode = Some("new_conversation".into()); - let empty = svc.add_job(empty_req).await.unwrap(); + let empty_job = svc.add_job("u1", empty_req).await.unwrap(); + assert_eq!(empty_job.user_id, "u1"); + assert_eq!(empty_job.conversation_id, ""); let mut stale_req = make_create_req("New-conv with stale id", every_60s()); - stale_req.conversation_id = "conv-that-no-longer-exists".into(); + stale_req.conversation_id = "missing-conv-that-no-longer-exists".into(); stale_req.execution_mode = Some("new_conversation".into()); - let stale = svc.add_job(stale_req).await.unwrap(); - - svc.init().await; - - assert!( - svc.get_job(&empty.id).await.is_ok(), - "empty new_conversation job must survive" - ); - assert!( - svc.get_job(&stale.id).await.is_ok(), - "new_conversation job with stale id must survive" - ); + let stale_job = svc.add_job("u1", stale_req).await.unwrap(); + assert_eq!(stale_job.user_id, "u1"); + assert_eq!(stale_job.conversation_id, "missing-conv-that-no-longer-exists"); } #[tokio::test] -async fn oc2_init_preserves_existing_jobs_with_missing_conversation() { +async fn oc2_rejects_existing_jobs_with_missing_conversation() { let (svc, _repo, _) = setup().await; let mut missing_req = make_create_req("Missing Conversation", every_60s()); missing_req.conversation_id = "missing-conv-1".into(); - let missing = svc.add_job(missing_req).await.unwrap(); - - let mut normal_req = make_create_req("Existing Conversation", every_60s()); - normal_req.conversation_id = "conv-existing".into(); - let normal = svc.add_job(normal_req).await.unwrap(); + let err = svc.add_job("u1", missing_req).await.unwrap_err(); + assert!(matches!(err, aionui_cron::error::CronError::Conversation(_))); +} - svc.init().await; +#[tokio::test] +async fn oc2b_rejects_existing_jobs_with_cross_user_conversation_code() { + let (svc, _repo, _bc, conv_repo) = setup_with_conv_repo().await; + sqlx::query( + "INSERT OR IGNORE INTO users (id, username, password_hash, created_at, updated_at) \ + VALUES ('u2', 'u2', 'hash', 0, 0)", + ) + .execute(&conv_repo.sqlite_pool) + .await + .unwrap(); + conv_repo + .create(&aionui_db::models::ConversationRow { + id: "conv_user_b".to_owned(), + user_id: "u2".to_owned(), + name: "User B Conversation".to_owned(), + r#type: "normal".to_owned(), + model: None, + status: Some("pending".to_owned()), + source: None, + channel_chat_id: None, + extra: "{}".to_owned(), + pinned: false, + pinned_at: None, + created_at: 0, + updated_at: 0, + }) + .await + .unwrap(); - let missing_found = svc.get_job(&missing.id).await; - assert!( - missing_found.is_ok(), - "existing job with deleted conversation should survive init and recover on next execution" - ); + let mut req = make_create_req("Cross User Conversation", every_60s()); + req.conversation_id = "conv_user_b".into(); - let found = svc.get_job(&normal.id).await; - assert!(found.is_ok()); + let err = svc.add_job("u1", req).await.unwrap_err(); + assert!(matches!( + err, + aionui_cron::error::CronError::CrossAccountReference(message) + if message == "Referenced conversation belongs to another user." + )); } #[tokio::test] -async fn existing_job_with_missing_conversation_run_now_creates_replacement_conversation() { - let (svc, _repo, _bc, conv_repo) = setup_with_conv_repo().await; +async fn existing_job_with_missing_conversation_is_rejected() { + let (svc, _repo, _bc, _conv_repo) = setup_with_conv_repo().await; let mut req = make_create_req("Missing Existing RunNow", every_60s()); req.conversation_id = "missing-conv-run-now".into(); req.execution_mode = Some("existing".into()); - let job = svc.add_job(req).await.unwrap(); - - let response = svc.run_now(&job.id).await.unwrap(); - - assert_ne!(response.conversation_id, "missing-conv-run-now"); - assert!( - conv_repo.get(&response.conversation_id).await.unwrap().is_some(), - "run-now should create a replacement conversation for an existing job whose previous conversation was deleted" - ); - - let rebound = svc.get_job(&job.id).await.unwrap(); - assert_eq!( - rebound.conversation_id, response.conversation_id, - "existing-mode replacement conversations must bind before the async turn finishes" - ); + let err = svc.add_job("u1", req).await.unwrap_err(); + assert!(matches!(err, aionui_cron::error::CronError::Conversation(_))); } #[tokio::test] @@ -2251,7 +2593,7 @@ async fn run_now_on_running_existing_conversation_returns_active_conversation_wi let mut req = make_create_req("Running Existing RunNow", every_60s()); req.conversation_id = "conv-running-run-now".into(); req.execution_mode = Some("existing".into()); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); bc.take_events(); let claim = conv_service @@ -2259,14 +2601,14 @@ async fn run_now_on_running_existing_conversation_returns_active_conversation_wi .try_claim_turn(&job.conversation_id, "turn-active") .expect("runtime claim should succeed"); - let response = svc.run_now(&job.id).await.unwrap(); + let response = svc.run_now("u1", &job.id).await.unwrap(); assert_eq!(response.conversation_id, job.conversation_id); for _ in 0..20 { tokio::task::yield_now().await; } - let row = cron_repo.get_by_id(&job.id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&job.id).await.unwrap().unwrap(); assert_eq!(row.run_count, 0); assert!(row.last_status.is_none()); assert!( @@ -2282,9 +2624,13 @@ async fn run_now_on_running_existing_conversation_returns_active_conversation_wi #[tokio::test] async fn delete_skill_clears_content() { let (svc, _, _) = setup().await; - let job = svc.add_job(make_create_req("Del Skill", every_60s())).await.unwrap(); + let job = svc + .add_job("u1", make_create_req("Del Skill", every_60s())) + .await + .unwrap(); svc.save_skill( + "u1", &job.id, SaveCronSkillRequest { content: "---\nname: x\n---\nOk".into(), @@ -2292,10 +2638,10 @@ async fn delete_skill_clears_content() { ) .await .unwrap(); - assert!(svc.has_skill(&job.id).await.unwrap().has_skill); + assert!(svc.has_skill("u1", &job.id).await.unwrap().has_skill); - svc.delete_skill(&job.id).await.unwrap(); - assert!(!svc.has_skill(&job.id).await.unwrap().has_skill); + svc.delete_skill("u1", &job.id).await.unwrap(); + assert!(!svc.has_skill("u1", &job.id).await.unwrap().has_skill); } fn conversation_cron_request(message: &str) -> CreateConversationCronRequest { @@ -2323,12 +2669,12 @@ async fn create_for_conversation_helper_creates_claimed_conversation_job_with_mu assert!(response.job_id.starts_with("cron_")); assert!(response.message.contains("Agent Helper Job")); - let row = cron_repo.get_by_id(&response.job_id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&response.job_id).await.unwrap().unwrap(); assert_eq!(row.payload_message, "first\nsecond\nthird"); assert_eq!(row.conversation_id, "conv_1"); assert_eq!(row.created_by, "agent"); - let bound = conv_repo.get("conv_1").await.unwrap().unwrap(); + let bound = conv_repo.get("u1", "conv_1").await.unwrap().unwrap(); let extra: serde_json::Value = serde_json::from_str(&bound.extra).unwrap(); assert_eq!(extra["cron_job_id"], response.job_id); assert_eq!(extra["cronJobId"], response.job_id); @@ -2355,11 +2701,11 @@ async fn create_for_conversation_helper_keeps_conversation_extra_mode_unchanged( .await .unwrap(); - let row = cron_repo.get_by_id(&response.job_id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&response.job_id).await.unwrap().unwrap(); let config: CronAgentConfig = serde_json::from_str(row.agent_config.as_deref().unwrap()).unwrap(); assert_eq!(config.mode.as_deref(), Some("yolo")); - let bound = conv_repo.get("conv_mode_default").await.unwrap().unwrap(); + let bound = conv_repo.get("u1", "conv_mode_default").await.unwrap().unwrap(); let extra: serde_json::Value = serde_json::from_str(&bound.extra).unwrap(); assert_eq!(extra["cron_job_id"], response.job_id); assert_eq!(extra["cronJobId"], response.job_id); @@ -2384,14 +2730,15 @@ async fn create_for_conversation_helper_uses_assistant_metadata_full_auto_mode() .await .unwrap(); - let row = cron_repo.get_by_id(&response.job_id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&response.job_id).await.unwrap().unwrap(); let config: CronAgentConfig = serde_json::from_str(row.agent_config.as_deref().unwrap()).unwrap(); assert_eq!(config.mode.as_deref(), Some("agent-full-access")); } #[tokio::test] async fn create_for_conversation_helper_uses_codex_canonical_full_auto_mode_from_fallback() { - let (svc, cron_repo, _, _, conv_service, agent_metadata_repo) = setup_with_conv_runtime_and_agent_metadata().await; + let (svc, cron_repo, _, _, conv_service, agent_metadata_repo, _) = + setup_with_conv_runtime_and_agent_metadata().await; let codex = agent_metadata_repo .find_builtin_by_backend("codex") .await @@ -2441,7 +2788,7 @@ async fn create_for_conversation_helper_uses_codex_canonical_full_auto_mode_from .await .unwrap(); - let row = cron_repo.get_by_id(&response.job_id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&response.job_id).await.unwrap().unwrap(); let config: CronAgentConfig = serde_json::from_str(row.agent_config.as_deref().unwrap()).unwrap(); assert_eq!(config.mode.as_deref(), Some("agent-full-access")); } @@ -2462,7 +2809,7 @@ async fn create_for_conversation_helper_fails_when_conversation_binding_fails() assert!(matches!(err, aionui_cron::error::CronError::Database(_))); - let rows = cron_repo.list_by_conversation("conv_1").await.unwrap(); + let rows = cron_repo.list_by_conversation_system("conv_1").await.unwrap(); assert!(rows.is_empty()); } @@ -2477,7 +2824,8 @@ async fn conversation_cron_routes_create_list_and_update_claimed_job() { let app = cron_routes(CronRouterState { cron_service: Arc::new(svc), conversation_service: (*conv_service).clone(), - }); + }) + .layer(Extension(current_user("u1"))); let response = app .clone() @@ -2501,7 +2849,7 @@ async fn conversation_cron_routes_create_list_and_update_claimed_job() { let envelope: ApiResponse = serde_json::from_slice(&body).unwrap(); let payload = envelope.data.expect("response should contain created job id"); - let row = cron_repo.get_by_id(&payload.job_id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&payload.job_id).await.unwrap().unwrap(); assert_eq!(row.payload_message, "first\nsecond\nthird"); assert_eq!(row.conversation_id, "conv_1"); assert_eq!(row.created_by, "agent"); @@ -2549,7 +2897,7 @@ async fn conversation_cron_routes_create_list_and_update_claimed_job() { .unwrap(); assert_eq!(update_response.status(), StatusCode::OK); - let row = cron_repo.get_by_id(&payload.job_id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&payload.job_id).await.unwrap().unwrap(); assert_eq!(row.name, "Updated Route Job"); assert_eq!(row.payload_message, "updated\nsecond\nthird"); assert_eq!(row.schedule_value, "0 */20 * * * *"); @@ -2561,7 +2909,8 @@ async fn conversation_cron_routes_reject_missing_headers_unclaimed_and_wrong_use let app = cron_routes(CronRouterState { cron_service: Arc::new(svc), conversation_service: (*conv_service).clone(), - }); + }) + .layer(Extension(current_user("u1"))); let missing_header = app .clone() @@ -2622,9 +2971,9 @@ async fn conversation_cron_routes_reject_missing_headers_unclaimed_and_wrong_use ) .await .unwrap(); - assert_eq!(wrong_user.status(), StatusCode::NOT_FOUND); + assert_eq!(wrong_user.status(), StatusCode::FORBIDDEN); let body = String::from_utf8(to_bytes(wrong_user.into_body(), usize::MAX).await.unwrap().to_vec()).unwrap(); - assert!(body.contains("conv_1")); + assert!(body.contains("CRON_HELPER_USER_MISMATCH")); } #[tokio::test] @@ -2707,12 +3056,13 @@ async fn update_for_conversation_helper_updates_claimed_conversation_job() { .unwrap(); assert_eq!(updated.name, "Updated Helper Job"); - let row = cron_repo.get_by_id(&created.job_id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&created.job_id).await.unwrap().unwrap(); assert_eq!(row.payload_message, "new message\nsecond line"); assert_eq!(row.schedule_value, "0 */20 * * * *"); conv_repo .update( + "u1", "conv_1", &ConversationRowUpdate { extra: Some("{}".into()), @@ -2737,7 +3087,7 @@ async fn update_for_conversation_helper_updates_claimed_conversation_job() { .await .unwrap(); - let bound = conv_repo.get("conv_1").await.unwrap().unwrap(); + let bound = conv_repo.get("u1", "conv_1").await.unwrap().unwrap(); let extra: serde_json::Value = serde_json::from_str(&bound.extra).unwrap(); assert_eq!(extra["cron_job_id"], created.job_id); assert_eq!(extra["cronJobId"], created.job_id); @@ -2758,6 +3108,7 @@ async fn update_for_conversation_helper_fails_when_conversation_binding_fails() conv_repo .update( + "u1", "conv_1", &ConversationRowUpdate { extra: Some("{}".into()), @@ -2828,7 +3179,10 @@ async fn update_for_conversation_helper_rejects_job_from_other_conversation() { #[tokio::test] async fn update_max_retries() { let (svc, _, _) = setup().await; - let job = svc.add_job(make_create_req("Retries", every_60s())).await.unwrap(); + let job = svc + .add_job("u1", make_create_req("Retries", every_60s())) + .await + .unwrap(); assert_eq!(job.max_retries, 3); let req = UpdateCronJobRequest { @@ -2843,7 +3197,7 @@ async fn update_max_retries() { max_retries: Some(5), queue_enabled: None, }; - let updated = svc.update_job(&job.id, req).await.unwrap(); + let updated = svc.update_job("u1", &job.id, req).await.unwrap(); assert_eq!(updated.max_retries, 5); } @@ -2852,12 +3206,13 @@ async fn create_and_update_queue_enabled() { let (svc, _, _) = setup().await; let mut create = make_create_req("Queued", every_60s()); create.queue_enabled = true; - let job = svc.add_job(create).await.unwrap(); + let job = svc.add_job("u1", create).await.unwrap(); assert!(job.queue_enabled); assert!(CronService::to_response(&job).state.queue_enabled); let updated = svc .update_job( + "u1", &job.id, UpdateCronJobRequest { queue_enabled: Some(false), @@ -2893,7 +3248,7 @@ async fn sc1_at_type_future_timestamp() { description: Some("once in 1h".into()), }, ); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); assert_eq!(job.next_run_at, Some(target_ms)); } @@ -2910,7 +3265,7 @@ async fn sc2_at_type_past_timestamp() { description: Some("once in the past".into()), }, ); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); assert_eq!(job.next_run_at, Some(target_ms)); } @@ -2921,7 +3276,7 @@ async fn sr1_system_resume_missed_job() { let (svc, repo, bc, conv_repo) = setup_with_conv_repo().await; let req = make_create_req("Resume Job", every_60s()); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); bc.take_events(); let past_ms = now_ms() - 10_000; @@ -2929,11 +3284,11 @@ async fn sr1_system_resume_missed_job() { next_run_at: Some(Some(past_ms)), ..Default::default() }; - repo.update(&job.id, ¶ms).await.unwrap(); + repo.update_system(&job.id, ¶ms).await.unwrap(); svc.handle_system_resume().await; - let updated = svc.get_job(&job.id).await.unwrap(); + let updated = svc.get_job("u1", &job.id).await.unwrap(); assert!( updated.last_run_at.is_none(), "missed job should not be auto-executed on resume" @@ -2977,24 +3332,24 @@ async fn cd1_delete_by_conversation_preserves_jobs() { let mut req_a = make_create_req("Cascade A", every_60s()); req_a.conversation_id = "conv_cascade".into(); - let job_a = svc.add_job(req_a).await.unwrap(); + let job_a = svc.add_job("u1", req_a).await.unwrap(); let mut req_b = make_create_req("Cascade B", every_60s()); req_b.conversation_id = "conv_cascade".into(); - let job_b = svc.add_job(req_b).await.unwrap(); + let job_b = svc.add_job("u1", req_b).await.unwrap(); let mut req_c = make_create_req("Unrelated", every_60s()); req_c.conversation_id = "conv_other".into(); - let _job_c = svc.add_job(req_c).await.unwrap(); + let _job_c = svc.add_job("u1", req_c).await.unwrap(); bc.take_events(); - svc.delete_jobs_by_conversation("conv_cascade").await; + svc.delete_jobs_by_conversation("u1", "conv_cascade").await; - assert!(svc.get_job(&job_a.id).await.is_ok()); - assert!(svc.get_job(&job_b.id).await.is_ok()); + assert!(svc.get_job("u1", &job_a.id).await.is_ok()); + assert!(svc.get_job("u1", &job_b.id).await.is_ok()); - let remaining = svc.list_jobs(&ListCronJobsQuery::default()).await.unwrap(); + let remaining = svc.list_jobs("u1", &ListCronJobsQuery::default()).await.unwrap(); assert_eq!(remaining.len(), 3, "all cron jobs should remain"); let events = bc.take_events(); @@ -3011,15 +3366,17 @@ async fn cd1_delete_by_conversation_preserves_jobs() { async fn cd2_delete_by_conversation_no_matching_jobs() { let (svc, _repo, bc) = setup().await; - svc.add_job(make_create_req("Existing", every_60s())).await.unwrap(); + svc.add_job("u1", make_create_req("Existing", every_60s())) + .await + .unwrap(); bc.take_events(); - svc.delete_jobs_by_conversation("conv_nonexistent").await; + svc.delete_jobs_by_conversation("u1", "conv_nonexistent").await; let events = bc.take_events(); assert!(events.is_empty(), "no events should be emitted when no jobs match"); - let all = svc.list_jobs(&ListCronJobsQuery::default()).await.unwrap(); + let all = svc.list_jobs("u1", &ListCronJobsQuery::default()).await.unwrap(); assert_eq!(all.len(), 1, "existing job should remain untouched"); } @@ -3033,12 +3390,12 @@ async fn cd3_on_conversation_delete_trait_preserves_jobs() { let mut req = make_create_req("Trait Cascade", every_60s()); req.conversation_id = "conv_trait_del".into(); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); bc.take_events(); - svc.on_conversation_deleted("conv_trait_del").await; + svc.on_conversation_deleted("u1", "conv_trait_del").await; - assert!(svc.get_job(&job.id).await.is_ok()); + assert!(svc.get_job("u1", &job.id).await.is_ok()); let events = bc.take_events(); assert!(events.is_empty()); @@ -3065,27 +3422,27 @@ async fn cd3b_on_conversation_delete_clears_deleted_workspace_from_jobs() { let mut req = make_create_req("Clears Deleted Workspace", every_60s()); req.conversation_id = conversation_id.clone(); req.agent_config.as_mut().unwrap().workspace = Some(deleted_workspace); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); bc.take_events(); - let bound_conversation = conv_repo.get(&conversation_id).await.unwrap().unwrap(); + let bound_conversation = conv_repo.get("u1", &conversation_id).await.unwrap().unwrap(); assert!( conv_service .auto_workspace_to_delete_for_row(&bound_conversation, &conversation_id) .is_some(), "test setup should use a workspace ConversationService will delete" ); - let row_before = cron_repo.get_by_id(&job.id).await.unwrap().unwrap(); + let row_before = cron_repo.get_by_id_system(&job.id).await.unwrap().unwrap(); let config_before: CronAgentConfig = serde_json::from_str(row_before.agent_config.as_deref().unwrap()).unwrap(); assert_eq!( config_before.workspace.as_deref(), Some(deleted_workspace_path.to_str().unwrap()) ); - svc.on_conversation_deleted(&conversation_id).await; + svc.on_conversation_deleted("u1", &conversation_id).await; - assert!(svc.get_job(&job.id).await.is_ok()); - let row = cron_repo.get_by_id(&job.id).await.unwrap().unwrap(); + assert!(svc.get_job("u1", &job.id).await.is_ok()); + let row = cron_repo.get_by_id_system(&job.id).await.unwrap().unwrap(); let config: CronAgentConfig = serde_json::from_str(row.agent_config.as_deref().unwrap()).unwrap(); assert!( config.workspace.is_none(), @@ -3096,6 +3453,71 @@ async fn cd3b_on_conversation_delete_clears_deleted_workspace_from_jobs() { assert!(events.is_empty()); } +#[tokio::test] +async fn cd3b_on_conversation_delete_only_clears_current_user_jobs() { + use aionui_common::OnConversationDelete; + + let (svc, cron_repo, _bc, conv_repo, conv_service) = setup_with_conv_runtime().await; + let conversation_id = format!("conv_workspace_scope_{}", now_ms()); + let deleted_workspace_path = std::env::temp_dir() + .join("conversations") + .join(format!("acp-temp-{conversation_id}")); + std::fs::create_dir_all(&deleted_workspace_path).unwrap(); + let deleted_workspace = deleted_workspace_path.to_string_lossy().to_string(); + conv_repo.set_conversation_extra( + &conversation_id, + serde_json::json!({ + "workspace": deleted_workspace, + }), + ); + + let bound_conversation = conv_repo.get("u1", &conversation_id).await.unwrap().unwrap(); + assert!( + conv_service + .auto_workspace_to_delete_for_row(&bound_conversation, &conversation_id) + .is_some(), + "test setup should use a workspace ConversationService will delete" + ); + + // New-conversation jobs may keep a stale conversation anchor (unanchored + // until run), so rows for both users can reference the same conversation + // id; existing-mode rows would be rejected at insert for a conversation + // the owner does not hold. + let mut u1_row = + make_cron_row_with_workspace("cron_workspace_scope_u1", "u1", &conversation_id, &deleted_workspace); + u1_row.execution_mode = "new_conversation".into(); + cron_repo.insert(&u1_row).await.unwrap(); + let mut u2_row = + make_cron_row_with_workspace("cron_workspace_scope_u2", "u2", &conversation_id, &deleted_workspace); + u2_row.execution_mode = "new_conversation".into(); + cron_repo.insert(&u2_row).await.unwrap(); + + svc.on_conversation_deleted("u1", &conversation_id).await; + + let u1_row = cron_repo + .get_by_id_system("cron_workspace_scope_u1") + .await + .unwrap() + .unwrap(); + let u1_config: CronAgentConfig = serde_json::from_str(u1_row.agent_config.as_deref().unwrap()).unwrap(); + assert!( + u1_config.workspace.is_none(), + "deleted conversation owner job should drop cached workspace" + ); + + let u2_row = cron_repo + .get_by_id_system("cron_workspace_scope_u2") + .await + .unwrap() + .unwrap(); + let u2_config: CronAgentConfig = serde_json::from_str(u2_row.agent_config.as_deref().unwrap()).unwrap(); + assert_eq!( + u2_config.workspace.as_deref(), + Some(deleted_workspace.as_str()), + "other user job must not be changed by the delete hook" + ); +} + #[tokio::test] async fn cd3c_on_conversation_delete_preserves_custom_workspace_on_jobs() { use aionui_common::OnConversationDelete; @@ -3113,12 +3535,12 @@ async fn cd3c_on_conversation_delete_preserves_custom_workspace_on_jobs() { let mut req = make_create_req("Preserves Custom Workspace", every_60s()); req.conversation_id = conversation_id.clone(); req.agent_config.as_mut().unwrap().workspace = Some(custom_workspace.clone()); - let job = svc.add_job(req).await.unwrap(); + let job = svc.add_job("u1", req).await.unwrap(); bc.take_events(); - svc.on_conversation_deleted(&conversation_id).await; + svc.on_conversation_deleted("u1", &conversation_id).await; - let row = cron_repo.get_by_id(&job.id).await.unwrap().unwrap(); + let row = cron_repo.get_by_id_system(&job.id).await.unwrap().unwrap(); let config: CronAgentConfig = serde_json::from_str(row.agent_config.as_deref().unwrap()).unwrap(); assert_eq!(config.workspace.as_deref(), Some(custom_workspace.as_str())); @@ -3135,23 +3557,23 @@ async fn cd4_on_conversation_delete_preserves_all_cron_jobs() { let mut new_conversation_req = make_create_req("Generated Run History", every_60s()); new_conversation_req.conversation_id = "conv_generated_run".into(); new_conversation_req.execution_mode = Some("new_conversation".into()); - let new_conversation_job = svc.add_job(new_conversation_req).await.unwrap(); + let new_conversation_job = svc.add_job("u1", new_conversation_req).await.unwrap(); let mut existing_req = make_create_req("Existing Bound Job", every_60s()); existing_req.conversation_id = "conv_generated_run".into(); existing_req.execution_mode = Some("existing".into()); - let existing_job = svc.add_job(existing_req).await.unwrap(); + let existing_job = svc.add_job("u1", existing_req).await.unwrap(); bc.take_events(); - svc.on_conversation_deleted("conv_generated_run").await; + svc.on_conversation_deleted("u1", "conv_generated_run").await; assert!( - svc.get_job(&new_conversation_job.id).await.is_ok(), + svc.get_job("u1", &new_conversation_job.id).await.is_ok(), "deleting a generated run conversation must not delete its new-conversation cron job" ); assert!( - svc.get_job(&existing_job.id).await.is_ok(), + svc.get_job("u1", &existing_job.id).await.is_ok(), "existing-mode jobs should also survive conversation deletion" ); diff --git a/crates/aionui-db/migrations/029_user_scope.sql b/crates/aionui-db/migrations/029_user_scope.sql new file mode 100644 index 000000000..a4d75f596 --- /dev/null +++ b/crates/aionui-db/migrations/029_user_scope.sql @@ -0,0 +1,984 @@ +-- Migration 028: add user scope for local, aggregate, and configuration data. + +PRAGMA foreign_keys = OFF; + +CREATE TABLE users_new ( + id TEXT PRIMARY KEY NOT NULL, + user_type TEXT NOT NULL DEFAULT 'local' + CHECK(user_type IN ('local', 'aionpro')), + external_user_id TEXT, + username TEXT, + email TEXT, + password_hash TEXT, + avatar_path TEXT, + jwt_secret TEXT, + status TEXT NOT NULL DEFAULT 'active' + CHECK(status IN ('active', 'disabled')), + session_generation INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + last_login INTEGER, + CHECK ( + (user_type = 'local' AND password_hash IS NOT NULL) + OR + (user_type = 'aionpro') + ), + CHECK ( + (external_user_id IS NULL) + OR + (length(external_user_id) > 0) + ) +); + +INSERT INTO users_new ( + id, + user_type, + external_user_id, + username, + email, + password_hash, + avatar_path, + jwt_secret, + status, + session_generation, + created_at, + updated_at, + last_login +) +SELECT + id, + 'local', + NULL, + username, + email, + password_hash, + avatar_path, + jwt_secret, + 'active', + 0, + created_at, + updated_at, + last_login +FROM users; + +CREATE TEMP TABLE user_scope_migration_checks ( + ok INTEGER NOT NULL CHECK (ok = 1) +); + +INSERT INTO user_scope_migration_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM users_new) = (SELECT COUNT(*) FROM users) + THEN 1 + ELSE 0 +END; + +INSERT INTO user_scope_migration_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM messages m + LEFT JOIN conversations c ON c.id = m.conversation_id + WHERE c.id IS NULL + ) + THEN 1 + ELSE 0 +END; + +INSERT INTO user_scope_migration_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM conversation_artifacts a + LEFT JOIN conversations c ON c.id = a.conversation_id + WHERE c.id IS NULL + ) + THEN 1 + ELSE 0 +END; + +INSERT INTO user_scope_migration_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM conversation_assistant_snapshots s + LEFT JOIN conversations c ON c.id = s.conversation_id + WHERE c.id IS NULL + ) + THEN 1 + ELSE 0 +END; + +INSERT INTO user_scope_migration_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM acp_session s + LEFT JOIN conversations c ON c.id = s.conversation_id + WHERE c.id IS NULL + ) + THEN 1 + ELSE 0 +END; + +INSERT INTO user_scope_migration_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM cron_jobs j + LEFT JOIN conversations c ON c.id = j.conversation_id + WHERE COALESCE(j.conversation_id, '') <> '' + AND c.id IS NULL + ) + THEN 1 + ELSE 0 +END; + +INSERT INTO user_scope_migration_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM mailbox m + LEFT JOIN teams t ON t.id = m.team_id + WHERE t.id IS NULL + ) + THEN 1 + ELSE 0 +END; + +INSERT INTO user_scope_migration_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM team_tasks tt + LEFT JOIN teams t ON t.id = tt.team_id + WHERE t.id IS NULL + ) + THEN 1 + ELSE 0 +END; + +DROP TABLE user_scope_migration_checks; + +DROP TABLE users; +ALTER TABLE users_new RENAME TO users; + +CREATE UNIQUE INDEX idx_users_local_username + ON users(username) + WHERE user_type = 'local' AND username IS NOT NULL; +CREATE UNIQUE INDEX idx_users_email + ON users(email) + WHERE email IS NOT NULL; +CREATE UNIQUE INDEX idx_users_external_user + ON users(user_type, external_user_id) + WHERE external_user_id IS NOT NULL; +CREATE INDEX idx_users_username ON users(username); +CREATE INDEX idx_users_status ON users(status); + +CREATE TRIGGER IF NOT EXISTS trg_mailbox_team_parent_insert +BEFORE INSERT ON mailbox +FOR EACH ROW +WHEN NOT EXISTS (SELECT 1 FROM teams WHERE id = NEW.team_id) +BEGIN + SELECT RAISE(ABORT, 'mailbox.team_id must reference teams.id'); +END; + +CREATE TRIGGER IF NOT EXISTS trg_mailbox_team_parent_update +BEFORE UPDATE OF team_id ON mailbox +FOR EACH ROW +WHEN NOT EXISTS (SELECT 1 FROM teams WHERE id = NEW.team_id) +BEGIN + SELECT RAISE(ABORT, 'mailbox.team_id must reference teams.id'); +END; + +CREATE TRIGGER IF NOT EXISTS trg_team_tasks_team_parent_insert +BEFORE INSERT ON team_tasks +FOR EACH ROW +WHEN NOT EXISTS (SELECT 1 FROM teams WHERE id = NEW.team_id) +BEGIN + SELECT RAISE(ABORT, 'team_tasks.team_id must reference teams.id'); +END; + +CREATE TRIGGER IF NOT EXISTS trg_team_tasks_team_parent_update +BEFORE UPDATE OF team_id ON team_tasks +FOR EACH ROW +WHEN NOT EXISTS (SELECT 1 FROM teams WHERE id = NEW.team_id) +BEGIN + SELECT RAISE(ABORT, 'team_tasks.team_id must reference teams.id'); +END; + +ALTER TABLE cron_jobs + ADD COLUMN user_id TEXT NOT NULL DEFAULT 'system_default_user'; +CREATE INDEX IF NOT EXISTS idx_cron_jobs_user_next_run ON cron_jobs(user_id, next_run_at); + +ALTER TABLE providers + ADD COLUMN user_id TEXT NOT NULL DEFAULT 'system_default_user'; +CREATE INDEX IF NOT EXISTS idx_providers_user_platform ON providers(user_id, platform); + +ALTER TABLE remote_agents + ADD COLUMN user_id TEXT NOT NULL DEFAULT 'system_default_user'; +CREATE INDEX IF NOT EXISTS idx_remote_agents_user_status ON remote_agents(user_id, status); + +CREATE TEMP TABLE user_scope_rebuild_checks ( + ok INTEGER NOT NULL CHECK (ok = 1) +); + +CREATE TABLE mcp_servers_new ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + name TEXT NOT NULL, + description TEXT, + enabled INTEGER NOT NULL DEFAULT 0, + transport_type TEXT NOT NULL, + transport_config TEXT NOT NULL, + tools TEXT, + last_test_status TEXT NOT NULL DEFAULT 'disconnected', + last_connected INTEGER, + original_json TEXT, + builtin INTEGER NOT NULL DEFAULT 0, + deleted_at INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (user_id, name) +); + +INSERT INTO mcp_servers_new ( + id, user_id, name, description, enabled, transport_type, transport_config, + tools, last_test_status, last_connected, original_json, builtin, deleted_at, + created_at, updated_at +) +SELECT + id, 'system_default_user', name, description, enabled, transport_type, + transport_config, tools, last_test_status, last_connected, original_json, + builtin, deleted_at, created_at, updated_at +FROM mcp_servers; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM mcp_servers_new) = (SELECT COUNT(*) FROM mcp_servers) + THEN 1 + ELSE 0 +END; + +DROP TABLE mcp_servers; +ALTER TABLE mcp_servers_new RENAME TO mcp_servers; +CREATE INDEX IF NOT EXISTS idx_mcp_servers_user_name ON mcp_servers(user_id, name); +CREATE INDEX IF NOT EXISTS idx_mcp_servers_user_enabled ON mcp_servers(user_id, enabled); +CREATE INDEX IF NOT EXISTS idx_mcp_servers_deleted_at ON mcp_servers(deleted_at); + +CREATE TABLE oauth_tokens_new ( + user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + server_url TEXT NOT NULL, + access_token TEXT NOT NULL, + refresh_token TEXT, + token_type TEXT NOT NULL DEFAULT 'bearer', + expires_at INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + PRIMARY KEY (user_id, server_url) +); + +INSERT INTO oauth_tokens_new ( + user_id, server_url, access_token, refresh_token, token_type, expires_at, + created_at, updated_at +) +SELECT + 'system_default_user', server_url, access_token, refresh_token, token_type, + expires_at, created_at, updated_at +FROM oauth_tokens; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM oauth_tokens_new) = (SELECT COUNT(*) FROM oauth_tokens) + THEN 1 + ELSE 0 +END; + +DROP TABLE oauth_tokens; +ALTER TABLE oauth_tokens_new RENAME TO oauth_tokens; + +CREATE TABLE system_settings_new ( + user_id TEXT PRIMARY KEY NOT NULL REFERENCES users(id), + language TEXT NOT NULL DEFAULT 'en-US', + notification_enabled INTEGER NOT NULL DEFAULT 1, + cron_notification_enabled INTEGER NOT NULL DEFAULT 0, + command_queue_enabled INTEGER NOT NULL DEFAULT 0, + save_upload_to_workspace INTEGER NOT NULL DEFAULT 0, + updated_at INTEGER NOT NULL +); + +INSERT INTO system_settings_new ( + user_id, language, notification_enabled, cron_notification_enabled, + command_queue_enabled, save_upload_to_workspace, updated_at +) +SELECT + 'system_default_user', language, notification_enabled, + cron_notification_enabled, command_queue_enabled, save_upload_to_workspace, + updated_at +FROM system_settings +WHERE id = 1; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM system_settings_new) = (SELECT COUNT(*) FROM system_settings WHERE id = 1) + THEN 1 + ELSE 0 +END; + +DROP TABLE system_settings; +ALTER TABLE system_settings_new RENAME TO system_settings; + +CREATE TABLE client_preferences_new ( + user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + key TEXT NOT NULL, + value TEXT NOT NULL, + updated_at INTEGER NOT NULL, + PRIMARY KEY (user_id, key) +); + +INSERT INTO client_preferences_new (user_id, key, value, updated_at) +SELECT 'system_default_user', key, value, updated_at +FROM client_preferences; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM client_preferences_new) = (SELECT COUNT(*) FROM client_preferences) + THEN 1 + ELSE 0 +END; + +DROP TABLE client_preferences; +ALTER TABLE client_preferences_new RENAME TO client_preferences; + +CREATE TABLE agent_metadata_new ( + id TEXT NOT NULL, + user_id TEXT REFERENCES users(id), + icon TEXT, + name TEXT NOT NULL, + name_i18n TEXT, + description TEXT, + description_i18n TEXT, + backend TEXT, + agent_type TEXT NOT NULL, + agent_source TEXT NOT NULL, + agent_source_info TEXT, + enabled INTEGER NOT NULL DEFAULT 1, + command TEXT, + args TEXT, + env TEXT, + native_skills_dirs TEXT, + behavior_policy TEXT, + yolo_id TEXT, + agent_capabilities TEXT, + auth_methods TEXT, + config_options TEXT, + available_modes TEXT, + available_models TEXT, + available_commands TEXT, + sort_order INTEGER NOT NULL DEFAULT 1000, + last_check_status TEXT, + last_check_kind TEXT, + last_check_error_code TEXT, + last_check_error_message TEXT, + last_check_guidance TEXT, + last_check_latency_ms INTEGER, + last_check_at INTEGER, + last_success_at INTEGER, + last_failure_at INTEGER, + command_override TEXT, + env_override TEXT, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); + +INSERT INTO agent_metadata_new ( + id, user_id, icon, name, name_i18n, description, description_i18n, + backend, agent_type, agent_source, agent_source_info, + enabled, command, args, env, native_skills_dirs, + behavior_policy, yolo_id, + agent_capabilities, auth_methods, config_options, + available_modes, available_models, available_commands, + sort_order, + last_check_status, last_check_kind, last_check_error_code, last_check_error_message, + last_check_guidance, last_check_latency_ms, last_check_at, last_success_at, last_failure_at, + command_override, env_override, created_at, updated_at +) +SELECT + id, + CASE WHEN agent_source IN ('builtin', 'internal') THEN NULL ELSE 'system_default_user' END, + icon, name, name_i18n, description, description_i18n, + backend, agent_type, agent_source, agent_source_info, + enabled, command, args, env, native_skills_dirs, + behavior_policy, yolo_id, + agent_capabilities, auth_methods, config_options, + available_modes, available_models, available_commands, + sort_order, + last_check_status, last_check_kind, last_check_error_code, last_check_error_message, + last_check_guidance, last_check_latency_ms, last_check_at, last_success_at, last_failure_at, + command_override, env_override, created_at, updated_at +FROM agent_metadata; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM agent_metadata_new) = (SELECT COUNT(*) FROM agent_metadata) + THEN 1 + ELSE 0 +END; + +DROP TABLE agent_metadata; +ALTER TABLE agent_metadata_new RENAME TO agent_metadata; +CREATE UNIQUE INDEX IF NOT EXISTS idx_agent_metadata_global_id + ON agent_metadata(id) + WHERE user_id IS NULL; +CREATE UNIQUE INDEX IF NOT EXISTS idx_agent_metadata_user_id + ON agent_metadata(user_id, id) + WHERE user_id IS NOT NULL; +CREATE INDEX IF NOT EXISTS idx_agent_metadata_backend ON agent_metadata(backend); +CREATE INDEX IF NOT EXISTS idx_agent_metadata_agent_type ON agent_metadata(agent_type); +CREATE INDEX IF NOT EXISTS idx_agent_metadata_sort_order ON agent_metadata(sort_order); +CREATE INDEX IF NOT EXISTS idx_agent_metadata_user_sort + ON agent_metadata(user_id, sort_order, name); + +ALTER TABLE assistants + ADD COLUMN user_id TEXT REFERENCES users(id); +UPDATE assistants +SET user_id = 'system_default_user' +WHERE user_id IS NULL; +CREATE INDEX IF NOT EXISTS idx_assistants_user_updated_at + ON assistants(user_id, updated_at DESC); + +CREATE TABLE assistant_overrides_new ( + user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + assistant_id TEXT NOT NULL, + enabled INTEGER NOT NULL DEFAULT 1, + sort_order INTEGER NOT NULL DEFAULT 0, + last_used_at INTEGER, + updated_at INTEGER NOT NULL, + PRIMARY KEY (user_id, assistant_id) +); + +INSERT INTO assistant_overrides_new ( + user_id, assistant_id, enabled, sort_order, last_used_at, updated_at +) +SELECT + 'system_default_user', assistant_id, enabled, sort_order, last_used_at, updated_at +FROM assistant_overrides; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM assistant_overrides_new) = (SELECT COUNT(*) FROM assistant_overrides) + THEN 1 + ELSE 0 +END; + +DROP TABLE assistant_overrides; +ALTER TABLE assistant_overrides_new RENAME TO assistant_overrides; +CREATE INDEX IF NOT EXISTS idx_assistant_overrides_user_sort + ON assistant_overrides(user_id, sort_order); + +DROP INDEX IF EXISTS idx_assistant_definitions_source_ref; +DROP INDEX IF EXISTS idx_assistant_definitions_assistant_id; + +CREATE TABLE assistant_definitions_new ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT REFERENCES users(id), + assistant_id TEXT NOT NULL, + source TEXT NOT NULL + CHECK (source IN ('builtin', 'user', 'generated')), + owner_type TEXT NOT NULL + CHECK (owner_type IN ('system', 'user')), + source_ref TEXT, + name TEXT NOT NULL, + name_i18n TEXT NOT NULL DEFAULT '{}', + description TEXT, + description_i18n TEXT NOT NULL DEFAULT '{}', + avatar_type TEXT NOT NULL DEFAULT 'none' + CHECK (avatar_type IN ('none', 'emoji', 'builtin_asset', 'user_asset')), + avatar_value TEXT, + agent_id TEXT NOT NULL, + rule_resource_type TEXT NOT NULL + CHECK (rule_resource_type IN ('none', 'builtin_asset', 'user_file', 'extension')), + rule_resource_ref TEXT, + recommended_prompts TEXT NOT NULL DEFAULT '[]', + recommended_prompts_i18n TEXT NOT NULL DEFAULT '{}', + default_model_mode TEXT NOT NULL + CHECK (default_model_mode IN ('auto', 'fixed')), + default_model_value TEXT, + default_permission_mode TEXT NOT NULL + CHECK (default_permission_mode IN ('auto', 'fixed')), + default_permission_value TEXT, + default_thought_level_mode TEXT NOT NULL DEFAULT 'auto' + CHECK (default_thought_level_mode IN ('auto', 'fixed')), + default_thought_level_value TEXT, + default_skills_mode TEXT NOT NULL + CHECK (default_skills_mode IN ('auto', 'fixed')), + default_skill_ids TEXT NOT NULL DEFAULT '[]', + custom_skill_names TEXT NOT NULL DEFAULT '[]', + default_disabled_builtin_skill_ids TEXT NOT NULL DEFAULT '[]', + default_mcps_mode TEXT NOT NULL + CHECK (default_mcps_mode IN ('auto', 'fixed')), + default_mcp_ids TEXT NOT NULL DEFAULT '[]', + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + deleted_at INTEGER +); + +INSERT INTO assistant_definitions_new ( + id, user_id, assistant_id, source, owner_type, source_ref, + name, name_i18n, description, description_i18n, avatar_type, avatar_value, + agent_id, rule_resource_type, rule_resource_ref, + recommended_prompts, recommended_prompts_i18n, + default_model_mode, default_model_value, + default_permission_mode, default_permission_value, + default_thought_level_mode, default_thought_level_value, + default_skills_mode, default_skill_ids, custom_skill_names, default_disabled_builtin_skill_ids, + default_mcps_mode, default_mcp_ids, + created_at, updated_at, deleted_at +) +SELECT + id, + CASE + WHEN source = 'builtin' AND owner_type = 'system' THEN NULL + ELSE 'system_default_user' + END, + assistant_id, source, owner_type, source_ref, + name, name_i18n, description, description_i18n, avatar_type, avatar_value, + agent_id, rule_resource_type, rule_resource_ref, + recommended_prompts, recommended_prompts_i18n, + default_model_mode, default_model_value, + default_permission_mode, default_permission_value, + default_thought_level_mode, default_thought_level_value, + default_skills_mode, default_skill_ids, custom_skill_names, default_disabled_builtin_skill_ids, + default_mcps_mode, default_mcp_ids, + created_at, updated_at, deleted_at +FROM assistant_definitions; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM assistant_definitions_new) = (SELECT COUNT(*) FROM assistant_definitions) + THEN 1 + ELSE 0 +END; + +DROP TABLE assistant_definitions; +ALTER TABLE assistant_definitions_new RENAME TO assistant_definitions; + +CREATE UNIQUE INDEX IF NOT EXISTS idx_assistant_definitions_global_source_ref + ON assistant_definitions(source, source_ref) + WHERE user_id IS NULL AND source_ref IS NOT NULL; +CREATE UNIQUE INDEX IF NOT EXISTS idx_assistant_definitions_user_source_ref + ON assistant_definitions(user_id, source, source_ref) + WHERE user_id IS NOT NULL AND source_ref IS NOT NULL; +CREATE UNIQUE INDEX IF NOT EXISTS idx_assistant_definitions_global_assistant_id + ON assistant_definitions(assistant_id) + WHERE user_id IS NULL; +CREATE UNIQUE INDEX IF NOT EXISTS idx_assistant_definitions_user_assistant_id + ON assistant_definitions(user_id, assistant_id) + WHERE user_id IS NOT NULL; +CREATE INDEX IF NOT EXISTS idx_assistant_definitions_source + ON assistant_definitions(source); +CREATE INDEX IF NOT EXISTS idx_assistant_definitions_agent_id + ON assistant_definitions(agent_id); +CREATE INDEX IF NOT EXISTS idx_assistant_definitions_user_updated_at + ON assistant_definitions(user_id, updated_at DESC); + +CREATE TABLE assistant_overlays_new ( + user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + assistant_definition_id TEXT NOT NULL, + enabled INTEGER NOT NULL DEFAULT 1, + sort_order INTEGER NOT NULL DEFAULT 0, + agent_id_override TEXT, + last_used_at INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + PRIMARY KEY (user_id, assistant_definition_id), + FOREIGN KEY (assistant_definition_id) REFERENCES assistant_definitions(id) ON DELETE CASCADE +); + +INSERT INTO assistant_overlays_new ( + user_id, assistant_definition_id, enabled, sort_order, agent_id_override, + last_used_at, created_at, updated_at +) +SELECT + 'system_default_user', assistant_definition_id, enabled, sort_order, + agent_id_override, last_used_at, created_at, updated_at +FROM assistant_overlays; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM assistant_overlays_new) = (SELECT COUNT(*) FROM assistant_overlays) + THEN 1 + ELSE 0 +END; + +DROP TABLE assistant_overlays; +ALTER TABLE assistant_overlays_new RENAME TO assistant_overlays; +CREATE INDEX IF NOT EXISTS idx_assistant_overlays_user_enabled + ON assistant_overlays(user_id, enabled); +CREATE INDEX IF NOT EXISTS idx_assistant_overlays_user_sort_order + ON assistant_overlays(user_id, sort_order); +CREATE INDEX IF NOT EXISTS idx_assistant_overlays_agent_id_override + ON assistant_overlays(agent_id_override) + WHERE agent_id_override IS NOT NULL; + +CREATE TABLE assistant_preferences_new ( + user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + assistant_definition_id TEXT NOT NULL, + last_model_id TEXT, + last_permission_value TEXT, + last_skill_ids TEXT NOT NULL DEFAULT '[]', + last_disabled_builtin_skill_ids TEXT NOT NULL DEFAULT '[]', + last_mcp_ids TEXT NOT NULL DEFAULT '[]', + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + last_thought_level_value TEXT, + PRIMARY KEY (user_id, assistant_definition_id), + FOREIGN KEY (assistant_definition_id) REFERENCES assistant_definitions(id) ON DELETE CASCADE +); + +INSERT INTO assistant_preferences_new ( + user_id, assistant_definition_id, last_model_id, last_permission_value, + last_skill_ids, last_disabled_builtin_skill_ids, last_mcp_ids, + created_at, updated_at, last_thought_level_value +) +SELECT + 'system_default_user', assistant_definition_id, last_model_id, last_permission_value, + last_skill_ids, last_disabled_builtin_skill_ids, last_mcp_ids, + created_at, updated_at, last_thought_level_value +FROM assistant_preferences; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM assistant_preferences_new) = (SELECT COUNT(*) FROM assistant_preferences) + THEN 1 + ELSE 0 +END; + +DROP TABLE assistant_preferences; +ALTER TABLE assistant_preferences_new RENAME TO assistant_preferences; + +CREATE TABLE skills_new ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT REFERENCES users(id), + name TEXT NOT NULL, + description TEXT, + path TEXT NOT NULL, + source TEXT NOT NULL DEFAULT 'user' + CHECK (source IN ('user', 'builtin', 'extension', 'cron')), + enabled INTEGER NOT NULL DEFAULT 1, + deleted_at INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); + +INSERT INTO skills_new ( + id, user_id, name, description, path, source, enabled, deleted_at, created_at, updated_at +) +SELECT + id, + CASE WHEN source = 'builtin' THEN NULL ELSE 'system_default_user' END, + name, description, path, source, enabled, deleted_at, created_at, updated_at +FROM skills; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM skills_new) = (SELECT COUNT(*) FROM skills) + THEN 1 + ELSE 0 +END; + +DROP TABLE skills; +ALTER TABLE skills_new RENAME TO skills; +CREATE UNIQUE INDEX IF NOT EXISTS idx_skills_global_name + ON skills(name) + WHERE user_id IS NULL; +CREATE UNIQUE INDEX IF NOT EXISTS idx_skills_user_name + ON skills(user_id, name) + WHERE user_id IS NOT NULL; +CREATE INDEX IF NOT EXISTS idx_skills_user_deleted_at ON skills(user_id, deleted_at); +CREATE INDEX IF NOT EXISTS idx_skills_source ON skills(source); +CREATE INDEX IF NOT EXISTS idx_skills_updated_at ON skills(updated_at DESC); + +ALTER TABLE skill_import_records + ADD COLUMN user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id); +CREATE INDEX IF NOT EXISTS idx_skill_import_records_user_created_at + ON skill_import_records(user_id, created_at DESC); + +CREATE TABLE assistant_plugins_new ( + id TEXT NOT NULL, + owner_user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + type TEXT NOT NULL, + name TEXT NOT NULL, + enabled INTEGER NOT NULL DEFAULT 0, + config TEXT NOT NULL, + status TEXT, + last_connected INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + PRIMARY KEY (owner_user_id, id) +); + +INSERT INTO assistant_plugins_new ( + id, owner_user_id, type, name, enabled, config, status, last_connected, + created_at, updated_at +) +SELECT + id, 'system_default_user', type, name, enabled, config, status, + last_connected, created_at, updated_at +FROM assistant_plugins; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM assistant_plugins_new) = (SELECT COUNT(*) FROM assistant_plugins) + THEN 1 + ELSE 0 +END; + +DROP TABLE assistant_plugins; +ALTER TABLE assistant_plugins_new RENAME TO assistant_plugins; +CREATE INDEX IF NOT EXISTS idx_assistant_plugins_owner_created_at + ON assistant_plugins(owner_user_id, created_at ASC); + +CREATE TABLE assistant_users_new ( + id TEXT PRIMARY KEY NOT NULL, + owner_user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + platform_user_id TEXT NOT NULL, + platform_type TEXT NOT NULL, + display_name TEXT, + authorized_at INTEGER NOT NULL, + last_active INTEGER, + session_id TEXT, + UNIQUE (owner_user_id, platform_user_id, platform_type) +); + +INSERT INTO assistant_users_new ( + id, owner_user_id, platform_user_id, platform_type, display_name, + authorized_at, last_active, session_id +) +SELECT + id, 'system_default_user', platform_user_id, platform_type, display_name, + authorized_at, last_active, session_id +FROM assistant_users; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM assistant_users_new) = (SELECT COUNT(*) FROM assistant_users) + THEN 1 + ELSE 0 +END; + +DROP TABLE assistant_users; +ALTER TABLE assistant_users_new RENAME TO assistant_users; +CREATE INDEX IF NOT EXISTS idx_assistant_users_owner_authorized_at + ON assistant_users(owner_user_id, authorized_at DESC); + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM assistant_sessions s + LEFT JOIN assistant_users u ON u.id = s.user_id + WHERE u.id IS NULL + ) + THEN 1 + ELSE 0 +END; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM assistant_sessions s + LEFT JOIN conversations c ON c.id = s.conversation_id + WHERE s.conversation_id IS NOT NULL + AND c.id IS NULL + ) + THEN 1 + ELSE 0 +END; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM assistant_sessions s + JOIN assistant_users u ON u.id = s.user_id + JOIN conversations c ON c.id = s.conversation_id + WHERE s.conversation_id IS NOT NULL + AND c.user_id != u.owner_user_id + ) + THEN 1 + ELSE 0 +END; + +CREATE TABLE assistant_pairing_codes_new ( + code TEXT NOT NULL, + owner_user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + platform_user_id TEXT NOT NULL, + platform_type TEXT NOT NULL, + display_name TEXT, + requested_at INTEGER NOT NULL, + expires_at INTEGER NOT NULL, + status TEXT NOT NULL DEFAULT 'pending' + CHECK (status IN ('pending', 'approved', 'rejected', 'expired')), + PRIMARY KEY (owner_user_id, code) +); + +INSERT INTO assistant_pairing_codes_new ( + code, owner_user_id, platform_user_id, platform_type, display_name, + requested_at, expires_at, status +) +SELECT + code, 'system_default_user', platform_user_id, platform_type, display_name, + requested_at, expires_at, status +FROM assistant_pairing_codes; + +INSERT INTO user_scope_rebuild_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM assistant_pairing_codes_new) = (SELECT COUNT(*) FROM assistant_pairing_codes) + THEN 1 + ELSE 0 +END; + +DROP TABLE assistant_pairing_codes; +ALTER TABLE assistant_pairing_codes_new RENAME TO assistant_pairing_codes; +CREATE INDEX IF NOT EXISTS idx_pairing_codes_owner_status + ON assistant_pairing_codes(owner_user_id, status); + +-- --------------------------------------------------------------------------- +-- Project-bind tables (created by 028_project_bind, which precedes this +-- migration). `projects` holds user data (user-named projects, user-selected +-- workspace roots; temp projects follow user-owned conversations), so it is +-- an independent user-scoped root. `folders` stays global by design: a +-- canonical filesystem path is a machine fact, deduped across users. +-- `project_explorer` inherits its user through the project parent chain but +-- carries a denormalized owner so the one-workspace-per-folder uniqueness +-- can be expressed per user (SQLite cannot index across tables) — same +-- pattern as the temporary Channel owner fields. +-- +-- Both tables are guaranteed empty here (the feature is un-wired in 028's +-- phase 1), so the rebuilds carry no data. The DEFAULT keeps phase-1 store +-- code (which does not write owner columns) working unchanged; phase 2 must +-- write owners explicitly and resolve workspace entries per user. +-- --------------------------------------------------------------------------- + +CREATE TEMP TABLE user_scope_project_checks ( + ok INTEGER NOT NULL CHECK (ok = 1) +); + +CREATE TABLE projects_new ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + project_id TEXT NOT NULL, + user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + name TEXT NOT NULL, + kind TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); + +INSERT INTO projects_new (id, project_id, user_id, name, kind, created_at, updated_at) +SELECT id, project_id, 'system_default_user', name, kind, created_at, updated_at +FROM projects; + +INSERT INTO user_scope_project_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM projects_new) = (SELECT COUNT(*) FROM projects) + THEN 1 + ELSE 0 +END; + +DROP TABLE projects; +ALTER TABLE projects_new RENAME TO projects; + +CREATE UNIQUE INDEX IF NOT EXISTS idx_projects_project_id_unique ON projects(project_id); +CREATE INDEX IF NOT EXISTS idx_projects_user_updated ON projects(user_id, updated_at DESC); +CREATE INDEX IF NOT EXISTS idx_projects_user_kind_updated ON projects(user_id, kind, updated_at DESC); + +CREATE TABLE project_explorer_new ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + pe_id TEXT NOT NULL, + project_id TEXT NOT NULL, + owner_user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + folder_id TEXT NOT NULL, + role TEXT NOT NULL, + display_name TEXT, + order_index INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); + +INSERT INTO project_explorer_new ( + id, pe_id, project_id, owner_user_id, folder_id, role, display_name, order_index, created_at, updated_at +) +SELECT + id, pe_id, project_id, 'system_default_user', folder_id, role, display_name, order_index, created_at, updated_at +FROM project_explorer; + +INSERT INTO user_scope_project_checks (ok) +SELECT CASE + WHEN (SELECT COUNT(*) FROM project_explorer_new) = (SELECT COUNT(*) FROM project_explorer) + THEN 1 + ELSE 0 +END; + +DROP TABLE project_explorer; +ALTER TABLE project_explorer_new RENAME TO project_explorer; + +CREATE UNIQUE INDEX IF NOT EXISTS idx_project_explorer_pe_id_unique ON project_explorer(pe_id); +CREATE UNIQUE INDEX IF NOT EXISTS idx_project_explorer_project_folder_unique + ON project_explorer(project_id, folder_id); +CREATE UNIQUE INDEX IF NOT EXISTS idx_project_explorer_one_workspace + ON project_explorer(project_id) WHERE role = 'workspace'; +-- Per user: a folder can be the workspace root of at most one project PER +-- OWNER. Replaces 028's global (folder_id) uniqueness, which would let one +-- user permanently reserve a directory against every other user. +CREATE UNIQUE INDEX IF NOT EXISTS idx_project_explorer_one_workspace_folder + ON project_explorer(owner_user_id, folder_id) WHERE role = 'workspace'; +CREATE INDEX IF NOT EXISTS idx_project_explorer_project_order ON project_explorer(project_id, order_index); +CREATE INDEX IF NOT EXISTS idx_project_explorer_folder_id ON project_explorer(folder_id); + +-- Cross-scope integrity: explorer owner must match its project's owner, and +-- user-scoped roots referencing a project must share its owner. All tables +-- involved are empty at migration time; these guard future upgrade paths. +INSERT INTO user_scope_project_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM project_explorer pe + JOIN projects p ON p.project_id = pe.project_id + WHERE pe.owner_user_id != p.user_id + ) + THEN 1 + ELSE 0 +END; + +INSERT INTO user_scope_project_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM conversations c + JOIN projects p ON p.project_id = c.project_id + WHERE c.project_id IS NOT NULL + AND p.user_id != c.user_id + ) + THEN 1 + ELSE 0 +END; + +INSERT INTO user_scope_project_checks (ok) +SELECT CASE + WHEN NOT EXISTS ( + SELECT 1 + FROM teams t + JOIN projects p ON p.project_id = t.project_id + WHERE t.project_id IS NOT NULL + AND p.user_id != t.user_id + ) + THEN 1 + ELSE 0 +END; + +DROP TABLE user_scope_project_checks; + +DROP TABLE user_scope_rebuild_checks; + +PRAGMA foreign_keys = ON; diff --git a/crates/aionui-db/src/agent_binding.rs b/crates/aionui-db/src/agent_binding.rs index dd8175eb8..9b1a7f4d7 100644 --- a/crates/aionui-db/src/agent_binding.rs +++ b/crates/aionui-db/src/agent_binding.rs @@ -52,12 +52,20 @@ pub fn resolve_agent_binding_from_rows(rows: &[AgentMetadataRow], value: &str) - .map(binding_resolution_for_agent) } -pub async fn resolve_agent_binding(pool: &SqlitePool, value: &str) -> Result, DbError> { +pub async fn resolve_agent_binding_for_user( + pool: &SqlitePool, + user_id: &str, + value: &str, +) -> Result, DbError> { let repo = SqliteAgentMetadataRepository::new(pool.clone()); - let rows = repo.list_all().await?; + let rows = repo.list_all_for_user(user_id).await?; Ok(resolve_agent_binding_from_rows(&rows, value)) } +pub async fn resolve_agent_binding(pool: &SqlitePool, value: &str) -> Result, DbError> { + resolve_agent_binding_for_user(pool, "system_default_user", value).await +} + fn agent_match_rank(row: &AgentMetadataRow) -> (i32, i64, &str) { let source_rank = match row.agent_source.as_str() { "builtin" => 0, diff --git a/crates/aionui-db/src/lib.rs b/crates/aionui-db/src/lib.rs index 29cd497a4..f56665261 100644 --- a/crates/aionui-db/src/lib.rs +++ b/crates/aionui-db/src/lib.rs @@ -10,8 +10,8 @@ pub mod models; mod repository; pub use agent_binding::{ - AgentBindingResolution, binding_resolution_for_agent, resolve_agent_binding, resolve_agent_binding_from_rows, - runtime_backend_for_agent, + AgentBindingResolution, binding_resolution_for_agent, resolve_agent_binding, resolve_agent_binding_for_user, + resolve_agent_binding_from_rows, runtime_backend_for_agent, }; pub use database::{ Database, DatabaseInitError, DatabaseInitOptions, init_database, init_database_memory, init_database_staged, @@ -24,16 +24,17 @@ pub use error::{ pub use instance_lock::{DataDirInstanceGuard, instance_lock_path}; pub use models::{ AgentMetadataRow, AssistantDefinitionRow, AssistantOverlayRow, AssistantOverrideRow, AssistantPreferenceRow, - AssistantRow, ConversationArtifactRow, ConversationAssistantSnapshotRow, CreateAssistantParams, FolderRow, - ProjectExplorerRow, ProjectKind, ProjectRow, Role, SkillImportRecordRow, SkillRow, - UpdateAgentAvailabilitySnapshotParams, UpdateAgentHandshakeParams, UpdateAssistantParams, + AssistantRow, ConversationArtifactRow, ConversationAssistantSnapshotRow, CreateAssistantParams, + ExternalUserProjection, FolderRow, ProjectExplorerRow, ProjectKind, ProjectRow, Role, SkillImportRecordRow, + SkillRow, UpdateAgentAvailabilitySnapshotParams, UpdateAgentHandshakeParams, UpdateAssistantParams, UpsertAgentMetadataParams, UpsertAssistantDefinitionParams, UpsertAssistantOverlayParams, - UpsertAssistantPreferenceParams, UpsertConversationAssistantSnapshotParams, UpsertOverrideParams, + UpsertAssistantPreferenceParams, UpsertConversationAssistantSnapshotParams, UpsertOverrideParams, UserStatus, + UserType, }; pub use repository::channel::UpdatePluginStatusParams; pub use repository::conversation::{ ConversationFilters, ConversationRowUpdate, MessagePageCursor, MessagePageDirection, MessagePageParams, - MessagePageResult, MessageRowUpdate, MessageSearchRow, + MessagePageResult, MessageRowUpdate, MessageSearchRow, StaleRuntimeMessageRow, }; pub use repository::cron::{ ClaimCronRunParams, CronRunClaimResult, FinishCronRunParams, RecoverableCronRun, UpdateCronJobParams, diff --git a/crates/aionui-db/src/models/agent_metadata.rs b/crates/aionui-db/src/models/agent_metadata.rs index ec1b19653..bfbcf3b55 100644 --- a/crates/aionui-db/src/models/agent_metadata.rs +++ b/crates/aionui-db/src/models/agent_metadata.rs @@ -12,6 +12,7 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct AgentMetadataRow { pub id: String, + pub user_id: Option, pub icon: Option, pub name: String, pub name_i18n: Option, diff --git a/crates/aionui-db/src/models/channel.rs b/crates/aionui-db/src/models/channel.rs index 51ee8a320..cbadca203 100644 --- a/crates/aionui-db/src/models/channel.rs +++ b/crates/aionui-db/src/models/channel.rs @@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct ChannelPluginRow { pub id: String, + pub owner_user_id: String, /// Platform type (telegram, lark, dingtalk, weixin, slack, discord). #[sqlx(rename = "type")] pub r#type: String, @@ -24,10 +25,11 @@ pub struct ChannelPluginRow { /// Row mapping for the `assistant_users` table. /// /// Represents an IM user authorized to chat with the assistant. -/// UNIQUE constraint on (platform_user_id, platform_type). +/// UNIQUE constraint on (owner_user_id, platform_user_id, platform_type). #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct AssistantUserRow { pub id: String, + pub owner_user_id: String, pub platform_user_id: String, pub platform_type: String, pub display_name: Option, @@ -60,6 +62,7 @@ pub struct AssistantSessionRow { #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct PairingCodeRow { pub code: String, + pub owner_user_id: String, pub platform_user_id: String, pub platform_type: String, pub display_name: Option, diff --git a/crates/aionui-db/src/models/client_preference.rs b/crates/aionui-db/src/models/client_preference.rs index 69d980df8..c75858e2b 100644 --- a/crates/aionui-db/src/models/client_preference.rs +++ b/crates/aionui-db/src/models/client_preference.rs @@ -6,6 +6,7 @@ use serde::{Deserialize, Serialize}; /// Generic key-value store. Values are stored as JSON-serialized TEXT. #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct ClientPreference { + pub user_id: String, pub key: String, pub value: String, pub updated_at: TimestampMs, diff --git a/crates/aionui-db/src/models/cron_job.rs b/crates/aionui-db/src/models/cron_job.rs index 9f9151785..4ad53333f 100644 --- a/crates/aionui-db/src/models/cron_job.rs +++ b/crates/aionui-db/src/models/cron_job.rs @@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct CronJobRow { pub id: String, + pub user_id: String, pub name: String, pub enabled: bool, pub schedule_kind: String, @@ -39,6 +40,7 @@ mod tests { fn cron_job_row_serialization_roundtrip() { let row = CronJobRow { id: "cron_abc123".into(), + user_id: "user1".into(), name: "Daily report".into(), enabled: true, schedule_kind: "cron".into(), @@ -77,6 +79,7 @@ mod tests { fn cron_job_row_optional_fields_default_to_none() { let row = CronJobRow { id: "cron_min".into(), + user_id: "user1".into(), name: "Minimal".into(), enabled: true, schedule_kind: "every".into(), diff --git a/crates/aionui-db/src/models/mcp_server.rs b/crates/aionui-db/src/models/mcp_server.rs index 556c1018d..a705c2ecd 100644 --- a/crates/aionui-db/src/models/mcp_server.rs +++ b/crates/aionui-db/src/models/mcp_server.rs @@ -11,6 +11,7 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct McpServerRow { pub id: String, + pub user_id: String, /// Unique server name (used as identifier when syncing to Agent CLIs). pub name: String, pub description: Option, diff --git a/crates/aionui-db/src/models/mod.rs b/crates/aionui-db/src/models/mod.rs index 4b2dfdb5d..2c5ac5953 100644 --- a/crates/aionui-db/src/models/mod.rs +++ b/crates/aionui-db/src/models/mod.rs @@ -40,4 +40,4 @@ pub use remote_agent::RemoteAgentRow; pub use skill::{SkillImportRecordRow, SkillRow}; pub use system_settings::SystemSettings; pub use team::{MailboxMessageRow, TeamRow, TeamTaskRow}; -pub use user::User; +pub use user::{ExternalUserProjection, User, UserStatus, UserType}; diff --git a/crates/aionui-db/src/models/oauth_token.rs b/crates/aionui-db/src/models/oauth_token.rs index 71c9ba3cf..415cb29a1 100644 --- a/crates/aionui-db/src/models/oauth_token.rs +++ b/crates/aionui-db/src/models/oauth_token.rs @@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize}; /// encrypted; callers handle encryption/decryption. #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct OAuthTokenRow { + pub user_id: String, /// MCP server URL (primary key). pub server_url: String, /// Encrypted OAuth access token. diff --git a/crates/aionui-db/src/models/provider.rs b/crates/aionui-db/src/models/provider.rs index fbe256828..37ae62fff 100644 --- a/crates/aionui-db/src/models/provider.rs +++ b/crates/aionui-db/src/models/provider.rs @@ -9,6 +9,7 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct Provider { pub id: String, + pub user_id: String, pub platform: String, pub name: String, pub base_url: String, diff --git a/crates/aionui-db/src/models/remote_agent.rs b/crates/aionui-db/src/models/remote_agent.rs index 0dac823d8..c0e7df701 100644 --- a/crates/aionui-db/src/models/remote_agent.rs +++ b/crates/aionui-db/src/models/remote_agent.rs @@ -12,6 +12,7 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct RemoteAgentRow { pub id: String, + pub user_id: String, pub name: String, /// One of: "openClaw", "zeroClaw", "acp". pub protocol: String, diff --git a/crates/aionui-db/src/models/skill.rs b/crates/aionui-db/src/models/skill.rs index 70e37295a..e92363d89 100644 --- a/crates/aionui-db/src/models/skill.rs +++ b/crates/aionui-db/src/models/skill.rs @@ -5,6 +5,7 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct SkillRow { pub id: String, + pub user_id: Option, pub name: String, pub description: Option, pub path: String, diff --git a/crates/aionui-db/src/models/system_settings.rs b/crates/aionui-db/src/models/system_settings.rs index f52a276f2..90e431c6b 100644 --- a/crates/aionui-db/src/models/system_settings.rs +++ b/crates/aionui-db/src/models/system_settings.rs @@ -3,11 +3,11 @@ use serde::{Deserialize, Serialize}; /// Row mapping for the `system_settings` table. /// -/// Single-row table (id is always 1). Boolean fields are stored as INTEGER +/// Per-user settings table. Boolean fields are stored as INTEGER /// in SQLite (0/1) and mapped to `bool` via sqlx. #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct SystemSettings { - pub id: i64, + pub user_id: String, pub language: String, pub notification_enabled: bool, pub cron_notification_enabled: bool, diff --git a/crates/aionui-db/src/models/user.rs b/crates/aionui-db/src/models/user.rs index 915a0b639..2cfa9dcbd 100644 --- a/crates/aionui-db/src/models/user.rs +++ b/crates/aionui-db/src/models/user.rs @@ -1,6 +1,40 @@ use aionui_common::TimestampMs; use serde::{Deserialize, Serialize}; +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, sqlx::Type)] +#[sqlx(type_name = "TEXT", rename_all = "lowercase")] +#[serde(rename_all = "lowercase")] +pub enum UserType { + Local, + Aionpro, +} + +impl UserType { + pub fn as_str(self) -> &'static str { + match self { + Self::Local => "local", + Self::Aionpro => "aionpro", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, sqlx::Type)] +#[sqlx(type_name = "TEXT", rename_all = "lowercase")] +#[serde(rename_all = "lowercase")] +pub enum UserStatus { + Active, + Disabled, +} + +impl UserStatus { + pub fn as_str(self) -> &'static str { + match self { + Self::Active => "active", + Self::Disabled => "disabled", + } + } +} + /// Row mapping for the `users` table. /// /// All fields match the SQLite column names and types exactly. @@ -8,12 +42,23 @@ use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)] pub struct User { pub id: String, - pub username: String, + pub user_type: UserType, + pub external_user_id: Option, + pub username: Option, pub email: Option, - pub password_hash: String, + pub password_hash: Option, pub avatar_path: Option, pub jwt_secret: Option, + pub status: UserStatus, + pub session_generation: i64, pub created_at: TimestampMs, pub updated_at: TimestampMs, pub last_login: Option, } + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ExternalUserProjection { + pub username: Option, + pub email: Option, + pub avatar_path: Option, +} diff --git a/crates/aionui-db/src/repository/acp_session.rs b/crates/aionui-db/src/repository/acp_session.rs index e0169dab8..affe7a985 100644 --- a/crates/aionui-db/src/repository/acp_session.rs +++ b/crates/aionui-db/src/repository/acp_session.rs @@ -8,9 +8,9 @@ //! session identity. Under the `"runtime"` key it holds the user's last //! per-session choices: current mode, current model, config selections, //! context usage. `AcpAgentService` updates those fields through -//! [`IAcpSessionRepository::save_runtime_state`] and +//! [`IAcpSessionRepository::save_runtime_state_for_user`] and //! `AcpAgentManager` preloads them on resume through -//! [`IAcpSessionRepository::load_runtime_state`]. +//! [`IAcpSessionRepository::load_runtime_state_for_user`]. use crate::error::DbError; use crate::models::AcpSessionRow; @@ -19,9 +19,10 @@ use crate::models::AcpSessionRow; /// /// `session_id` stays `None` until the CLI returns one (first /// `session/new` or `session/load`), at which point the caller flips -/// it through [`IAcpSessionRepository::update_session_id`]. +/// it through [`IAcpSessionRepository::update_session_id_for_user`]. #[derive(Debug, Clone)] pub struct CreateAcpSessionParams<'a> { + pub user_id: &'a str, pub conversation_id: &'a str, pub agent_source: &'a str, pub agent_id: &'a str, @@ -43,7 +44,7 @@ pub struct PersistedSessionState { pub context_usage_json: Option, } -/// Partial update for [`IAcpSessionRepository::save_runtime_state`]. +/// Partial update for [`IAcpSessionRepository::save_runtime_state_for_user`]. /// /// `Option>` lets callers distinguish "leave untouched" /// (outer `None`) from "clear to null" (inner `None`). @@ -66,8 +67,8 @@ impl SaveRuntimeStateParams<'_> { #[async_trait::async_trait] pub trait IAcpSessionRepository: Send + Sync { - /// Fetch the full row by conversation id. - async fn get(&self, conversation_id: &str) -> Result, DbError>; + /// Fetch the full row only when the owning conversation belongs to `user_id`. + async fn get_for_user(&self, user_id: &str, conversation_id: &str) -> Result, DbError>; /// Insert a fresh `acp_session` row. Called by `ConversationService` /// when an ACP-type conversation is created; primary-key conflict @@ -75,34 +76,49 @@ pub trait IAcpSessionRepository: Send + Sync { async fn create(&self, params: &CreateAcpSessionParams<'_>) -> Result; /// Record the CLI-assigned `session_id` after `session/new` or - /// `session/load` succeeds. Returns `true` when the row existed. - async fn update_session_id(&self, conversation_id: &str, session_id: &str) -> Result; + /// `session/load` succeeds. Returns `true` when the row existed and + /// belongs to `user_id`. + async fn update_session_id_for_user( + &self, + user_id: &str, + conversation_id: &str, + session_id: &str, + ) -> Result; /// Null the stored `session_id`, dropping the resume anchor while keeping the /// row (config/runtime state) intact. Called on an unrecoverable resume error /// ("No conversation found" / `error_during_execution`) so the NEXT turn opens /// Fresh instead of re-resuming a dead backend session forever. Distinct from - /// [`delete`](Self::delete), which drops the whole row. Returns `true` when the - /// row existed. This is the direct-CLI equivalent of the clean-slate - /// `Orchestrator` emitting `BackendBound{None}` and the legacy ACP + /// [`delete_for_user`](Self::delete_for_user), which drops the whole row. + /// Returns `true` when the row existed under `user_id`'s conversation. This + /// is the direct-CLI equivalent of the clean-slate `Orchestrator` emitting + /// `BackendBound{None}` and the legacy ACP /// `rebuild_after_session_not_found` → `clear_session_id` self-heal. - async fn clear_session_id(&self, conversation_id: &str) -> Result; + async fn clear_session_id_for_user(&self, user_id: &str, conversation_id: &str) -> Result; - /// Delete the row. Called by the conversation delete hook — no DB - /// foreign key, so this must be invoked explicitly. - async fn delete(&self, conversation_id: &str) -> Result; + /// Delete the row when the owning conversation belongs to `user_id`. + /// Called by the conversation delete hook — no DB foreign key, so this + /// must be invoked explicitly. + async fn delete_for_user(&self, user_id: &str, conversation_id: &str) -> Result; /// Decode and return the `session_config.runtime` sub-object. /// Returns `None` when the row does not exist or the JSON lacks a /// `runtime` key; returns `Some(Default::default())` when the key - /// is present but empty. - async fn load_runtime_state(&self, conversation_id: &str) -> Result, DbError>; + /// is present but empty. Only returns a row when the owning conversation + /// belongs to `user_id`. + async fn load_runtime_state_for_user( + &self, + user_id: &str, + conversation_id: &str, + ) -> Result, DbError>; - /// Merge a partial runtime update into `session_config.runtime`. - /// Assumes the row exists (created alongside the conversation); - /// returns `Ok(false)` when it does not. - async fn save_runtime_state( + /// Merge a partial runtime update into `session_config.runtime` when the + /// owning conversation belongs to `user_id`. Assumes the row exists + /// (created alongside the conversation); returns `Ok(false)` when it does + /// not. + async fn save_runtime_state_for_user( &self, + user_id: &str, conversation_id: &str, params: &SaveRuntimeStateParams<'_>, ) -> Result; diff --git a/crates/aionui-db/src/repository/agent_metadata.rs b/crates/aionui-db/src/repository/agent_metadata.rs index 33df906ad..621c78e47 100644 --- a/crates/aionui-db/src/repository/agent_metadata.rs +++ b/crates/aionui-db/src/repository/agent_metadata.rs @@ -17,9 +17,13 @@ pub trait IAgentMetadataRepository: Send + Sync { /// Return every row, in insertion order. async fn list_all(&self) -> Result, DbError>; + async fn list_all_for_user(&self, user_id: &str) -> Result, DbError>; + /// Look up by primary key. async fn get(&self, id: &str) -> Result, DbError>; + async fn get_for_user(&self, user_id: &str, id: &str) -> Result, DbError>; + /// Look up by the unique `(agent_source, name)` pair. async fn find_by_source_and_name( &self, @@ -27,14 +31,37 @@ pub trait IAgentMetadataRepository: Send + Sync { name: &str, ) -> Result, DbError>; + async fn find_by_source_and_name_for_user( + &self, + user_id: &str, + agent_source: &str, + name: &str, + ) -> Result, DbError>; + /// Look up the first `builtin` row whose vendor label matches. /// Useful when the caller only has the legacy `backend` string and /// not a full agent id. async fn find_builtin_by_backend(&self, backend: &str) -> Result, DbError>; + async fn find_builtin_by_backend_for_user( + &self, + user_id: &str, + backend: &str, + ) -> Result, DbError>; + /// Insert or replace a row. Returns the row as stored. async fn upsert(&self, params: &UpsertAgentMetadataParams<'_>) -> Result; + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAgentMetadataParams<'_>, + ) -> Result; + + async fn upsert_global(&self, params: &UpsertAgentMetadataParams<'_>) -> Result { + self.upsert(params).await + } + /// Apply handshake-derived fields on top of an existing row. /// Returns `Ok(None)` if no row matches `id`. async fn apply_handshake( @@ -43,6 +70,13 @@ pub trait IAgentMetadataRepository: Send + Sync { params: &UpdateAgentHandshakeParams<'_>, ) -> Result, DbError>; + async fn apply_handshake_for_user( + &self, + user_id: &str, + id: &str, + params: &UpdateAgentHandshakeParams<'_>, + ) -> Result, DbError>; + /// Persist the latest availability snapshot for an existing row. /// Returns `Ok(None)` if no row matches `id`. async fn update_availability_snapshot( @@ -51,6 +85,13 @@ pub trait IAgentMetadataRepository: Send + Sync { params: &UpdateAgentAvailabilitySnapshotParams<'_>, ) -> Result, DbError>; + async fn update_availability_snapshot_for_user( + &self, + user_id: &str, + id: &str, + params: &UpdateAgentAvailabilitySnapshotParams<'_>, + ) -> Result, DbError>; + /// Write only the self-repair override columns for an agent, leaving all /// other columns (seed truth + availability snapshot) untouched. Kept /// separate from the full-row upsert so startup reconcile never clobbers @@ -62,9 +103,21 @@ pub trait IAgentMetadataRepository: Send + Sync { env_override: Option<&str>, ) -> Result<(), DbError>; + async fn update_agent_overrides_for_user( + &self, + user_id: &str, + id: &str, + command_override: Option<&str>, + env_override: Option<&str>, + ) -> Result<(), DbError>; + /// Toggle the `enabled` flag. Returns `true` if a row was updated. async fn set_enabled(&self, id: &str, enabled: bool) -> Result; + async fn set_enabled_for_user(&self, user_id: &str, id: &str, enabled: bool) -> Result; + /// Delete a row. Returns `true` if a row was removed. async fn delete(&self, id: &str) -> Result; + + async fn delete_for_user(&self, user_id: &str, id: &str) -> Result; } diff --git a/crates/aionui-db/src/repository/assistant.rs b/crates/aionui-db/src/repository/assistant.rs index 5f6085d5e..d53f969f1 100644 --- a/crates/aionui-db/src/repository/assistant.rs +++ b/crates/aionui-db/src/repository/assistant.rs @@ -15,24 +15,43 @@ pub trait IAssistantRepository: Send + Sync { /// Return all user-authored assistants, ordered by `updated_at` descending. async fn list(&self) -> Result, DbError>; + async fn list_for_user(&self, user_id: &str) -> Result, DbError>; + /// Look up a single assistant by id. async fn get(&self, id: &str) -> Result, DbError>; + async fn get_for_user(&self, user_id: &str, id: &str) -> Result, DbError>; + /// Insert a new assistant row. Primary-key conflict surfaces as /// `DbError::Conflict`. async fn create(&self, params: &CreateAssistantParams<'_>) -> Result; + async fn create_for_user(&self, user_id: &str, params: &CreateAssistantParams<'_>) + -> Result; + /// Partial update of an existing assistant row. Returns `Ok(None)` if /// no row matches. async fn update(&self, id: &str, params: &UpdateAssistantParams<'_>) -> Result, DbError>; + async fn update_for_user( + &self, + user_id: &str, + id: &str, + params: &UpdateAssistantParams<'_>, + ) -> Result, DbError>; + /// Delete an assistant row by id. Returns `true` if a row was removed. async fn delete(&self, id: &str) -> Result; + async fn delete_for_user(&self, user_id: &str, id: &str) -> Result; + /// Insert or replace by id. Exists for callers outside of the /// migration/import path; the import endpoint must use `create` and /// skip on conflict per spec §6.3. async fn upsert(&self, params: &CreateAssistantParams<'_>) -> Result; + + async fn upsert_for_user(&self, user_id: &str, params: &CreateAssistantParams<'_>) + -> Result; } /// Per-assistant user state (enabled flag, sort order, last-used timestamp). @@ -41,41 +60,74 @@ pub trait IAssistantOverrideRepository: Send + Sync { /// Fetch the override row for a given assistant id, if any. async fn get(&self, assistant_id: &str) -> Result, DbError>; + async fn get_for_user(&self, user_id: &str, assistant_id: &str) -> Result, DbError>; + /// Fetch all override rows. async fn get_all(&self) -> Result, DbError>; + async fn get_all_for_user(&self, user_id: &str) -> Result, DbError>; + /// Insert or update the override row for an assistant. async fn upsert(&self, params: &UpsertOverrideParams<'_>) -> Result; + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertOverrideParams<'_>, + ) -> Result; + /// Delete the override row for an assistant. Returns `true` if a row was /// removed. async fn delete(&self, assistant_id: &str) -> Result; + async fn delete_for_user(&self, user_id: &str, assistant_id: &str) -> Result; + /// Remove override rows whose `assistant_id` is not in `valid_ids`. /// Returns the number of rows deleted. async fn delete_orphans(&self, valid_ids: &[&str]) -> Result; + + async fn delete_orphans_for_user(&self, user_id: &str, valid_ids: &[&str]) -> Result; } /// Runtime assistant definitions across builtin / user / generated / extension sources. #[async_trait::async_trait] pub trait IAssistantDefinitionRepository: Send + Sync { async fn list(&self) -> Result, DbError>; + async fn list_for_user(&self, user_id: &str) -> Result, DbError>; async fn list_including_deleted(&self) -> Result, DbError> { self.list().await } + async fn list_including_deleted_for_user(&self, user_id: &str) -> Result, DbError>; async fn get_by_assistant_id(&self, assistant_id: &str) -> Result, DbError>; + async fn get_by_assistant_id_for_user( + &self, + user_id: &str, + assistant_id: &str, + ) -> Result, DbError>; async fn get_by_assistant_id_including_deleted( &self, assistant_id: &str, ) -> Result, DbError> { self.get_by_assistant_id(assistant_id).await } + async fn get_by_assistant_id_including_deleted_for_user( + &self, + user_id: &str, + assistant_id: &str, + ) -> Result, DbError>; async fn get_by_id(&self, id: &str) -> Result, DbError>; + async fn get_by_id_for_user(&self, user_id: &str, id: &str) -> Result, DbError>; async fn get_by_source_ref( &self, source: &str, source_ref: &str, ) -> Result, DbError>; + async fn get_by_source_ref_for_user( + &self, + user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError>; async fn get_by_source_ref_including_deleted( &self, source: &str, @@ -83,7 +135,37 @@ pub trait IAssistantDefinitionRepository: Send + Sync { ) -> Result, DbError> { self.get_by_source_ref(source, source_ref).await } + async fn get_by_source_ref_including_deleted_for_user( + &self, + user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError>; + async fn get_global_by_source_ref_including_deleted( + &self, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref_including_deleted(source, source_ref).await + } + async fn get_global_by_assistant_id_including_deleted( + &self, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id_including_deleted(assistant_id).await + } async fn upsert(&self, params: &UpsertAssistantDefinitionParams<'_>) -> Result; + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result; + async fn upsert_global( + &self, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert(params).await + } async fn update_avatar_fields_preserving_deleted( &self, id: &str, @@ -96,21 +178,45 @@ pub trait IAssistantDefinitionRepository: Send + Sync { )) } async fn soft_delete(&self, id: &str, deleted_at: i64) -> Result; + async fn soft_delete_for_user(&self, user_id: &str, id: &str, deleted_at: i64) -> Result; } /// Runtime per-user assistant overlay used by the current app version. #[async_trait::async_trait] pub trait IAssistantOverlayRepository: Send + Sync { async fn get(&self, assistant_definition_id: &str) -> Result, DbError>; + async fn get_for_user( + &self, + user_id: &str, + assistant_definition_id: &str, + ) -> Result, DbError>; async fn list(&self) -> Result, DbError>; + async fn list_for_user(&self, user_id: &str) -> Result, DbError>; async fn upsert(&self, params: &UpsertAssistantOverlayParams<'_>) -> Result; + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAssistantOverlayParams<'_>, + ) -> Result; async fn delete(&self, assistant_definition_id: &str) -> Result; + async fn delete_for_user(&self, user_id: &str, assistant_definition_id: &str) -> Result; } /// Assistant-scoped "auto remember last" preferences. #[async_trait::async_trait] pub trait IAssistantPreferenceRepository: Send + Sync { async fn get(&self, assistant_definition_id: &str) -> Result, DbError>; + async fn get_for_user( + &self, + user_id: &str, + assistant_definition_id: &str, + ) -> Result, DbError>; async fn upsert(&self, params: &UpsertAssistantPreferenceParams<'_>) -> Result; + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAssistantPreferenceParams<'_>, + ) -> Result; async fn delete(&self, assistant_definition_id: &str) -> Result; + async fn delete_for_user(&self, user_id: &str, assistant_definition_id: &str) -> Result; } diff --git a/crates/aionui-db/src/repository/channel.rs b/crates/aionui-db/src/repository/channel.rs index 8edb8d19f..b53b7d2de 100644 --- a/crates/aionui-db/src/repository/channel.rs +++ b/crates/aionui-db/src/repository/channel.rs @@ -13,94 +13,121 @@ use crate::models::{AssistantSessionRow, AssistantUserRow, ChannelPluginRow, Pai pub trait IChannelRepository: Send + Sync { // ── Plugin CRUD ────────────────────────────────────────────────── - /// Returns all registered plugins. - async fn get_all_plugins(&self) -> Result, DbError>; + /// Returns all registered plugins for an owner. + async fn get_all_plugins(&self, owner_user_id: &str) -> Result, DbError>; /// Returns a single plugin by id, or `None` if not found. - async fn get_plugin(&self, id: &str) -> Result, DbError>; + async fn get_plugin(&self, owner_user_id: &str, id: &str) -> Result, DbError>; /// Inserts a new plugin or updates an existing one (by id). - async fn upsert_plugin(&self, row: &ChannelPluginRow) -> Result<(), DbError>; + async fn upsert_plugin(&self, owner_user_id: &str, row: &ChannelPluginRow) -> Result<(), DbError>; /// Updates only the `status` and `last_connected` of a plugin. - async fn update_plugin_status(&self, id: &str, params: &UpdatePluginStatusParams) -> Result<(), DbError>; + async fn update_plugin_status( + &self, + owner_user_id: &str, + id: &str, + params: &UpdatePluginStatusParams, + ) -> Result<(), DbError>; /// Deletes a plugin by id. Returns `DbError::NotFound` if absent. - async fn delete_plugin(&self, id: &str) -> Result<(), DbError>; + async fn delete_plugin(&self, owner_user_id: &str, id: &str) -> Result<(), DbError>; // ── User CRUD ──────────────────────────────────────────────────── - /// Returns all authorized users. - async fn get_all_users(&self) -> Result, DbError>; + /// Returns all authorized users for an owner. + async fn get_all_users(&self, owner_user_id: &str) -> Result, DbError>; /// Finds a user by platform identity. Returns `None` if not found. async fn get_user_by_platform( &self, + owner_user_id: &str, platform_user_id: &str, platform_type: &str, ) -> Result, DbError>; /// Creates a new authorized user record. - async fn create_user(&self, row: &AssistantUserRow) -> Result<(), DbError>; + async fn create_user(&self, owner_user_id: &str, row: &AssistantUserRow) -> Result<(), DbError>; /// Updates `last_active` timestamp for a user. - async fn update_user_last_active(&self, id: &str, last_active: TimestampMs) -> Result<(), DbError>; + async fn update_user_last_active( + &self, + owner_user_id: &str, + id: &str, + last_active: TimestampMs, + ) -> Result<(), DbError>; /// Deletes a user by id. Returns `DbError::NotFound` if absent. /// Associated sessions are cascade-deleted by the database. - async fn delete_user(&self, id: &str) -> Result<(), DbError>; + async fn delete_user(&self, owner_user_id: &str, id: &str) -> Result<(), DbError>; // ── Session CRUD ───────────────────────────────────────────────── - /// Returns all sessions. - async fn get_all_sessions(&self) -> Result, DbError>; + /// Returns all sessions for an owner. + async fn get_all_sessions(&self, owner_user_id: &str) -> Result, DbError>; /// Returns a single session by id. - async fn get_session(&self, id: &str) -> Result, DbError>; + async fn get_session(&self, owner_user_id: &str, id: &str) -> Result, DbError>; /// Finds an existing session by user + chat, or creates a new one. /// If found, updates `last_activity` and returns the existing row. /// If not found, inserts `new_row` and returns it. async fn get_or_create_session( &self, - user_id: &str, + owner_user_id: &str, + channel_user_id: &str, chat_id: &str, new_row: &AssistantSessionRow, ) -> Result; /// Updates `last_activity` timestamp for a session. - async fn update_session_activity(&self, id: &str, last_activity: TimestampMs) -> Result<(), DbError>; + async fn update_session_activity( + &self, + owner_user_id: &str, + id: &str, + last_activity: TimestampMs, + ) -> Result<(), DbError>; /// Updates the `conversation_id` of a session. - async fn update_session_conversation(&self, id: &str, conversation_id: &str) -> Result<(), DbError>; + async fn update_session_conversation( + &self, + owner_user_id: &str, + id: &str, + conversation_id: &str, + ) -> Result<(), DbError>; /// Updates the `agent_type` of a session. - async fn update_session_agent_type(&self, id: &str, agent_type: &str) -> Result<(), DbError>; + async fn update_session_agent_type(&self, owner_user_id: &str, id: &str, agent_type: &str) -> Result<(), DbError>; /// Deletes all sessions belonging to a user. - async fn delete_sessions_by_user(&self, user_id: &str) -> Result<(), DbError>; + async fn delete_sessions_by_user(&self, owner_user_id: &str, channel_user_id: &str) -> Result<(), DbError>; /// Deletes the session for a specific user + chat pair. - async fn delete_session_by_user_chat(&self, user_id: &str, chat_id: &str) -> Result<(), DbError>; + async fn delete_session_by_user_chat( + &self, + owner_user_id: &str, + channel_user_id: &str, + chat_id: &str, + ) -> Result<(), DbError>; // ── Pairing Codes ──────────────────────────────────────────────── /// Creates a new pairing code record. - async fn create_pairing(&self, row: &PairingCodeRow) -> Result<(), DbError>; + async fn create_pairing(&self, owner_user_id: &str, row: &PairingCodeRow) -> Result<(), DbError>; /// Returns all pairing codes with status = 'pending'. - async fn get_pending_pairings(&self) -> Result, DbError>; + async fn get_pending_pairings(&self, owner_user_id: &str) -> Result, DbError>; /// Retrieves a single pairing code, or `None` if not found. - async fn get_pairing_by_code(&self, code: &str) -> Result, DbError>; + async fn get_pairing_by_code(&self, owner_user_id: &str, code: &str) -> Result, DbError>; /// Updates the status of a pairing code. /// Returns `DbError::NotFound` if the code doesn't exist. - async fn update_pairing_status(&self, code: &str, status: &str) -> Result<(), DbError>; + async fn update_pairing_status(&self, owner_user_id: &str, code: &str, status: &str) -> Result<(), DbError>; /// Marks all expired-but-still-pending pairing codes as 'expired'. /// `now` is the current timestamp in milliseconds. - async fn cleanup_expired_pairings(&self, now: TimestampMs) -> Result; + async fn cleanup_expired_pairings(&self, owner_user_id: &str, now: TimestampMs) -> Result; } /// Parameters for updating plugin runtime status. diff --git a/crates/aionui-db/src/repository/client_preference.rs b/crates/aionui-db/src/repository/client_preference.rs index cf9edf388..06d39da05 100644 --- a/crates/aionui-db/src/repository/client_preference.rs +++ b/crates/aionui-db/src/repository/client_preference.rs @@ -7,15 +7,15 @@ use crate::models::ClientPreference; #[async_trait::async_trait] pub trait IClientPreferenceRepository: Send + Sync { /// Returns all client preferences. - async fn get_all(&self) -> Result, DbError>; + async fn get_all(&self, user_id: &str) -> Result, DbError>; /// Returns preferences for the given keys only. /// Keys that don't exist are simply omitted from the result. - async fn get_by_keys(&self, keys: &[&str]) -> Result, DbError>; + async fn get_by_keys(&self, user_id: &str, keys: &[&str]) -> Result, DbError>; /// Inserts or updates a batch of key-value pairs. - async fn upsert_batch(&self, entries: &[(&str, &str)]) -> Result<(), DbError>; + async fn upsert_batch(&self, user_id: &str, entries: &[(&str, &str)]) -> Result<(), DbError>; /// Deletes the given keys. - async fn delete_keys(&self, keys: &[&str]) -> Result<(), DbError>; + async fn delete_keys(&self, user_id: &str, keys: &[&str]) -> Result<(), DbError>; } diff --git a/crates/aionui-db/src/repository/conversation.rs b/crates/aionui-db/src/repository/conversation.rs index 469dbe537..b0c9f8a6a 100644 --- a/crates/aionui-db/src/repository/conversation.rs +++ b/crates/aionui-db/src/repository/conversation.rs @@ -18,18 +18,21 @@ use crate::models::{ pub trait IConversationRepository: Send + Sync { // ── Conversation CRUD ─────────────────────────────────────────── - /// Returns a conversation by ID, or `None` if not found. - async fn get(&self, id: &str) -> Result, DbError>; + /// Returns a conversation by user and ID, or `None` if not found. + async fn get(&self, user_id: &str, id: &str) -> Result, DbError>; + + /// Returns the owner user ID for a conversation, or `None` if the conversation does not exist. + async fn owner_user_id(&self, id: &str) -> Result, DbError>; /// Inserts a new conversation row. async fn create(&self, row: &ConversationRow) -> Result<(), DbError>; - /// Partially updates a conversation. Returns `DbError::NotFound` if ID is missing. - async fn update(&self, id: &str, updates: &ConversationRowUpdate) -> Result<(), DbError>; + /// Partially updates a conversation. Returns `DbError::NotFound` if ID is missing for the user. + async fn update(&self, user_id: &str, id: &str, updates: &ConversationRowUpdate) -> Result<(), DbError>; /// Deletes a conversation (messages cascade via FK). - /// Returns `DbError::NotFound` if ID is missing. - async fn delete(&self, id: &str) -> Result<(), DbError>; + /// Returns `DbError::NotFound` if ID is missing for the user. + async fn delete(&self, user_id: &str, id: &str) -> Result<(), DbError>; /// Lists conversations with cursor-based pagination and optional filters. async fn list_paginated( @@ -59,6 +62,7 @@ pub trait IConversationRepository: Send + Sync { /// Returns the persisted assistant snapshot for a conversation, if any. async fn get_assistant_snapshot( &self, + _user_id: &str, _conversation_id: &str, ) -> Result, DbError> { Ok(None) @@ -67,36 +71,48 @@ pub trait IConversationRepository: Send + Sync { /// Inserts or updates a persisted assistant snapshot for a conversation. async fn upsert_assistant_snapshot( &self, + _user_id: &str, _params: &UpsertConversationAssistantSnapshotParams<'_>, ) -> Result, DbError> { Ok(None) } /// Deletes the assistant snapshot bound to a conversation. - async fn delete_assistant_snapshot(&self, _conversation_id: &str) -> Result { + async fn delete_assistant_snapshot(&self, _user_id: &str, _conversation_id: &str) -> Result { Ok(false) } // ── Message operations ────────────────────────────────────────── /// Returns cursor-paginated messages for a conversation in ascending display order. - async fn list_messages_page(&self, conv_id: &str, params: &MessagePageParams) - -> Result; + async fn list_messages_page( + &self, + user_id: &str, + conv_id: &str, + params: &MessagePageParams, + ) -> Result; /// Returns a single message scoped to a conversation. - async fn get_message(&self, _conv_id: &str, _message_id: &str) -> Result, DbError> { + async fn get_message( + &self, + _user_id: &str, + _conv_id: &str, + _message_id: &str, + ) -> Result, DbError> { Ok(None) } /// Inserts a new message row. - async fn insert_message(&self, message: &MessageRow) -> Result<(), DbError>; + async fn insert_message(&self, user_id: &str, message: &MessageRow) -> Result<(), DbError>; /// Inserts a message row, or merges mutable fields into the existing row with the same ID. - async fn upsert_message(&self, message: &MessageRow) -> Result<(), DbError> { - match self.insert_message(message).await { + async fn upsert_message(&self, user_id: &str, message: &MessageRow) -> Result<(), DbError> { + match self.insert_message(user_id, message).await { Ok(()) => Ok(()), Err(DbError::Conflict(_)) => { self.update_message( + user_id, + &message.conversation_id, &message.id, &MessageRowUpdate { content: Some(message.content.clone()), @@ -111,14 +127,21 @@ pub trait IConversationRepository: Send + Sync { } /// Partially updates a message. Returns `DbError::NotFound` if ID is missing. - async fn update_message(&self, id: &str, updates: &MessageRowUpdate) -> Result<(), DbError>; + async fn update_message( + &self, + user_id: &str, + conversation_id: &str, + id: &str, + updates: &MessageRowUpdate, + ) -> Result<(), DbError>; /// Deletes all messages belonging to a conversation. - async fn delete_messages_by_conversation(&self, conv_id: &str) -> Result<(), DbError>; + async fn delete_messages_by_conversation(&self, user_id: &str, conv_id: &str) -> Result<(), DbError>; /// Finds a message by (conversation_id, msg_id, type) triple. async fn get_message_by_msg_id( &self, + user_id: &str, conv_id: &str, msg_id: &str, msg_type: &str, @@ -126,7 +149,7 @@ pub trait IConversationRepository: Send + Sync { /// Lists stale assistant-side runtime messages that were left in a /// non-terminal state by a previous process. - async fn list_stale_runtime_messages(&self) -> Result, DbError> { + async fn list_stale_runtime_messages(&self) -> Result, DbError> { Ok(Vec::new()) } @@ -140,13 +163,18 @@ pub trait IConversationRepository: Send + Sync { ) -> Result, DbError>; /// Returns persisted conversation artifacts ordered by `created_at`. - async fn list_artifacts(&self, _conversation_id: &str) -> Result, DbError> { + async fn list_artifacts( + &self, + _user_id: &str, + _conversation_id: &str, + ) -> Result, DbError> { Ok(Vec::new()) } /// Returns a conversation artifact by ID scoped to a conversation. async fn get_artifact( &self, + _user_id: &str, _conversation_id: &str, _artifact_id: &str, ) -> Result, DbError> { @@ -154,13 +182,18 @@ pub trait IConversationRepository: Send + Sync { } /// Inserts or updates a conversation artifact by primary key. - async fn upsert_artifact(&self, artifact: &ConversationArtifactRow) -> Result { + async fn upsert_artifact( + &self, + _user_id: &str, + artifact: &ConversationArtifactRow, + ) -> Result { Ok(artifact.clone()) } /// Updates artifact status and returns the updated row if found. async fn update_artifact_status( &self, + _user_id: &str, _conversation_id: &str, _artifact_id: &str, _status: &str, @@ -172,6 +205,7 @@ pub trait IConversationRepository: Send + Sync { /// Marks all skill suggestion artifacts for a cron job as saved. async fn mark_skill_suggest_artifacts_saved( &self, + _user_id: &str, _cron_job_id: &str, _updated_at: TimestampMs, ) -> Result, DbError> { @@ -179,13 +213,17 @@ pub trait IConversationRepository: Send + Sync { } /// Deletes all artifacts belonging to a conversation. - async fn delete_artifacts_by_conversation(&self, _conversation_id: &str) -> Result<(), DbError> { + async fn delete_artifacts_by_conversation(&self, _user_id: &str, _conversation_id: &str) -> Result<(), DbError> { Ok(()) } /// Returns legacy persisted cron trigger rows so callers can synthesize /// artifact cards for historical conversations created before artifact migration. - async fn list_legacy_cron_trigger_messages(&self, _conversation_id: &str) -> Result, DbError> { + async fn list_legacy_cron_trigger_messages( + &self, + _user_id: &str, + _conversation_id: &str, + ) -> Result, DbError> { Ok(Vec::new()) } } @@ -229,6 +267,12 @@ pub struct MessagePageResult { pub has_more_after: bool, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StaleRuntimeMessageRow { + pub user_id: String, + pub message: MessageRow, +} + /// Filters for paginated conversation listing. #[derive(Debug, Clone, Default)] pub struct ConversationFilters { diff --git a/crates/aionui-db/src/repository/cron.rs b/crates/aionui-db/src/repository/cron.rs index d2c27359c..b135cd910 100644 --- a/crates/aionui-db/src/repository/cron.rs +++ b/crates/aionui-db/src/repository/cron.rs @@ -70,28 +70,51 @@ pub trait ICronRepository: Send + Sync { /// Inserts a new cron job row. async fn insert(&self, row: &CronJobRow) -> Result<(), DbError>; - /// Updates a cron job by ID with the provided fields. - /// Returns `DbError::NotFound` if absent. - async fn update(&self, id: &str, params: &UpdateCronJobParams) -> Result<(), DbError>; + /// System-only update path for scheduler/recovery code that already + /// operates on persisted job identity. User-facing code must use + /// `update_for_user`. + async fn update_system(&self, id: &str, params: &UpdateCronJobParams) -> Result<(), DbError>; - /// Deletes a cron job by ID. Returns `DbError::NotFound` if absent. - async fn delete(&self, id: &str) -> Result<(), DbError>; + /// Updates a cron job whose conversation is owned by `user_id`. + async fn update_for_user(&self, user_id: &str, id: &str, params: &UpdateCronJobParams) -> Result<(), DbError>; - /// Returns a single cron job by ID, or `None` if not found. - async fn get_by_id(&self, id: &str) -> Result, DbError>; + /// System-only delete path for scheduler/recovery cleanup. User-facing + /// code must use `delete_for_user`. + async fn delete_system(&self, id: &str) -> Result<(), DbError>; - /// Returns all cron jobs ordered by creation time ascending. - async fn list_all(&self) -> Result, DbError>; + /// Deletes a cron job whose conversation is owned by `user_id`. + async fn delete_for_user(&self, user_id: &str, id: &str) -> Result<(), DbError>; - /// Returns all enabled cron jobs. - async fn list_enabled(&self) -> Result, DbError>; + /// System-only lookup path for scheduler/recovery code. User-facing code + /// must use `get_by_id_for_user`. + async fn get_by_id_system(&self, id: &str) -> Result, DbError>; - /// Returns all cron jobs for a given conversation. - async fn list_by_conversation(&self, conversation_id: &str) -> Result, DbError>; + /// Returns a single cron job whose conversation is owned by `user_id`. + async fn get_by_id_for_user(&self, user_id: &str, id: &str) -> Result, DbError>; - /// Deletes all cron jobs associated with a conversation. - /// Returns the number of deleted rows. - async fn delete_by_conversation(&self, conversation_id: &str) -> Result; + /// System-only full scan. User-facing code must use `list_all_for_user`. + async fn list_all_system(&self) -> Result, DbError>; + + /// Returns all cron jobs whose conversations are owned by `user_id`. + async fn list_all_for_user(&self, user_id: &str) -> Result, DbError>; + + /// System-only enabled-job scan used by the scheduler. + async fn list_enabled_system(&self) -> Result, DbError>; + + /// System-only conversation lookup. User-facing code must use + /// `list_by_conversation_for_user`. + async fn list_by_conversation_system(&self, conversation_id: &str) -> Result, DbError>; + + /// Returns all cron jobs for a given conversation owned by `user_id`. + async fn list_by_conversation_for_user( + &self, + user_id: &str, + conversation_id: &str, + ) -> Result, DbError>; + + /// System-only conversation cleanup. User-facing code must verify the + /// conversation owner first and use scoped paths. + async fn delete_by_conversation_system(&self, conversation_id: &str) -> Result; /// Atomically claims one scheduled occurrence across all backend processes. async fn claim_run(&self, params: &ClaimCronRunParams<'_>) -> Result; diff --git a/crates/aionui-db/src/repository/mcp_server.rs b/crates/aionui-db/src/repository/mcp_server.rs index cad4d6b1a..56fe57bb9 100644 --- a/crates/aionui-db/src/repository/mcp_server.rs +++ b/crates/aionui-db/src/repository/mcp_server.rs @@ -12,29 +12,29 @@ use crate::models::McpServerRow; #[async_trait::async_trait] pub trait IMcpServerRepository: Send + Sync { /// Returns all MCP servers, ordered by creation time ascending. - async fn list(&self) -> Result, DbError>; + async fn list(&self, user_id: &str) -> Result, DbError>; /// Finds an MCP server by ID, or `None` if not found. - async fn find_by_id(&self, id: &str) -> Result, DbError>; + async fn find_by_id(&self, user_id: &str, id: &str) -> Result, DbError>; /// Finds an MCP server by name, or `None` if not found. - async fn find_by_name(&self, name: &str) -> Result, DbError>; + async fn find_by_name(&self, user_id: &str, name: &str) -> Result, DbError>; /// Finds an MCP server by ID, including soft-deleted rows. - async fn find_by_id_any(&self, id: &str) -> Result, DbError> { - self.find_by_id(id).await + async fn find_by_id_any(&self, user_id: &str, id: &str) -> Result, DbError> { + self.find_by_id(user_id, id).await } /// Finds an MCP server by name, including soft-deleted rows. - async fn find_by_name_any(&self, name: &str) -> Result, DbError> { - self.find_by_name(name).await + async fn find_by_name_any(&self, user_id: &str, name: &str) -> Result, DbError> { + self.find_by_name(user_id, name).await } /// Finds a set of MCP servers by ID, including soft-deleted rows. - async fn list_by_ids_any(&self, ids: &[String]) -> Result, DbError> { + async fn list_by_ids_any(&self, user_id: &str, ids: &[String]) -> Result, DbError> { let mut rows = Vec::with_capacity(ids.len()); for id in ids { - if let Some(row) = self.find_by_id_any(id).await? { + if let Some(row) = self.find_by_id_any(user_id, id).await? { rows.push(row); } } @@ -47,21 +47,27 @@ pub trait IMcpServerRepository: Send + Sync { /// Updates an existing MCP server. Returns `DbError::NotFound` if the ID /// doesn't exist, `DbError::Conflict` if the new name collides with another. - async fn update(&self, id: &str, params: UpdateMcpServerParams<'_>) -> Result; + async fn update(&self, user_id: &str, id: &str, params: UpdateMcpServerParams<'_>) + -> Result; /// Soft-deletes an MCP server by ID. Returns `DbError::NotFound` if the ID /// doesn't exist. - async fn delete(&self, id: &str) -> Result<(), DbError>; + async fn delete(&self, user_id: &str, id: &str) -> Result<(), DbError>; /// Upserts multiple servers by name: existing names are updated, /// new names are inserted. Returns the count of affected rows. - async fn batch_upsert(&self, servers: &[CreateMcpServerParams<'_>]) -> Result, DbError>; + async fn batch_upsert( + &self, + user_id: &str, + servers: &[CreateMcpServerParams<'_>], + ) -> Result, DbError>; /// Updates only the latest connection-test result status /// (and optionally `last_connected`). /// Returns `DbError::NotFound` if the ID doesn't exist. async fn update_status( &self, + user_id: &str, id: &str, status: &str, last_connected: Option, @@ -69,12 +75,13 @@ pub trait IMcpServerRepository: Send + Sync { /// Updates only the tools JSON for a server. /// Returns `DbError::NotFound` if the ID doesn't exist. - async fn update_tools(&self, id: &str, tools: Option<&str>) -> Result<(), DbError>; + async fn update_tools(&self, user_id: &str, id: &str, tools: Option<&str>) -> Result<(), DbError>; } /// Parameters for creating a new MCP server. #[derive(Debug, Clone)] pub struct CreateMcpServerParams<'a> { + pub user_id: &'a str, pub name: &'a str, pub description: Option<&'a str>, pub enabled: bool, diff --git a/crates/aionui-db/src/repository/oauth_token.rs b/crates/aionui-db/src/repository/oauth_token.rs index 2df0e15db..0c4387c56 100644 --- a/crates/aionui-db/src/repository/oauth_token.rs +++ b/crates/aionui-db/src/repository/oauth_token.rs @@ -10,22 +10,23 @@ use crate::models::OAuthTokenRow; #[async_trait::async_trait] pub trait IOAuthTokenRepository: Send + Sync { /// Gets a token by server URL, or `None` if not found. - async fn get_by_url(&self, server_url: &str) -> Result, DbError>; + async fn get_by_url(&self, user_id: &str, server_url: &str) -> Result, DbError>; /// Inserts or updates a token for the given server URL. async fn upsert(&self, params: UpsertOAuthTokenParams<'_>) -> Result; /// Deletes a token by server URL. Returns `DbError::NotFound` if the URL /// doesn't exist. - async fn delete(&self, server_url: &str) -> Result<(), DbError>; + async fn delete(&self, user_id: &str, server_url: &str) -> Result<(), DbError>; /// Returns the list of server URLs that have stored tokens. - async fn list_authenticated_urls(&self) -> Result, DbError>; + async fn list_authenticated_urls(&self, user_id: &str) -> Result, DbError>; } /// Parameters for inserting or updating an OAuth token. #[derive(Debug)] pub struct UpsertOAuthTokenParams<'a> { + pub user_id: &'a str, pub server_url: &'a str, pub access_token: &'a str, pub refresh_token: Option<&'a str>, diff --git a/crates/aionui-db/src/repository/provider.rs b/crates/aionui-db/src/repository/provider.rs index 688520c1c..5688e2417 100644 --- a/crates/aionui-db/src/repository/provider.rs +++ b/crates/aionui-db/src/repository/provider.rs @@ -8,19 +8,19 @@ use crate::models::Provider; #[async_trait::async_trait] pub trait IProviderRepository: Send + Sync { /// Returns all providers, ordered by creation time ascending. - async fn list(&self) -> Result, DbError>; + async fn list(&self, user_id: &str) -> Result, DbError>; /// Finds a provider by ID, or `None` if not found. - async fn find_by_id(&self, id: &str) -> Result, DbError>; + async fn find_by_id(&self, user_id: &str, id: &str) -> Result, DbError>; /// Creates a new provider and returns the inserted row. async fn create(&self, params: CreateProviderParams<'_>) -> Result; /// Updates an existing provider. Returns `DbError::NotFound` if the ID doesn't exist. - async fn update(&self, id: &str, params: UpdateProviderParams<'_>) -> Result; + async fn update(&self, user_id: &str, id: &str, params: UpdateProviderParams<'_>) -> Result; /// Deletes a provider by ID. Returns `DbError::NotFound` if the ID doesn't exist. - async fn delete(&self, id: &str) -> Result<(), DbError>; + async fn delete(&self, user_id: &str, id: &str) -> Result<(), DbError>; } /// Parameters for creating a new provider. @@ -28,6 +28,7 @@ pub trait IProviderRepository: Send + Sync { pub struct CreateProviderParams<'a> { /// Optional caller-supplied id. When `None`, the repository generates one. pub id: Option<&'a str>, + pub user_id: &'a str, pub platform: &'a str, pub name: &'a str, pub base_url: &'a str, diff --git a/crates/aionui-db/src/repository/remote_agent.rs b/crates/aionui-db/src/repository/remote_agent.rs index 1ec5c771a..62367041a 100644 --- a/crates/aionui-db/src/repository/remote_agent.rs +++ b/crates/aionui-db/src/repository/remote_agent.rs @@ -11,24 +11,30 @@ use crate::models::RemoteAgentRow; #[async_trait::async_trait] pub trait IRemoteAgentRepository: Send + Sync { /// Returns all remote agents, ordered by creation time ascending. - async fn list(&self) -> Result, DbError>; + async fn list(&self, user_id: &str) -> Result, DbError>; /// Finds a remote agent by ID, or `None` if not found. - async fn find_by_id(&self, id: &str) -> Result, DbError>; + async fn find_by_id(&self, user_id: &str, id: &str) -> Result, DbError>; /// Creates a new remote agent and returns the inserted row. async fn create(&self, params: CreateRemoteAgentParams<'_>) -> Result; /// Updates an existing remote agent. Returns `DbError::NotFound` if the ID doesn't exist. - async fn update(&self, id: &str, params: UpdateRemoteAgentParams<'_>) -> Result; + async fn update( + &self, + user_id: &str, + id: &str, + params: UpdateRemoteAgentParams<'_>, + ) -> Result; /// Deletes a remote agent by ID. Returns `DbError::NotFound` if the ID doesn't exist. - async fn delete(&self, id: &str) -> Result<(), DbError>; + async fn delete(&self, user_id: &str, id: &str) -> Result<(), DbError>; /// Updates only the connection status (and optionally last_connected_at). /// Returns `DbError::NotFound` if the ID doesn't exist. async fn update_status( &self, + user_id: &str, id: &str, status: &str, last_connected_at: Option, @@ -38,6 +44,7 @@ pub trait IRemoteAgentRepository: Send + Sync { /// Parameters for creating a new remote agent. #[derive(Debug)] pub struct CreateRemoteAgentParams<'a> { + pub user_id: &'a str, pub name: &'a str, pub protocol: &'a str, pub url: &'a str, diff --git a/crates/aionui-db/src/repository/settings.rs b/crates/aionui-db/src/repository/settings.rs index ba5a79497..bb926d495 100644 --- a/crates/aionui-db/src/repository/settings.rs +++ b/crates/aionui-db/src/repository/settings.rs @@ -3,17 +3,18 @@ use crate::models::SystemSettings; /// System settings data access abstraction. /// -/// The `system_settings` table holds a single row (id=1). +/// The `system_settings` table holds one row per user. /// `get_settings` returns `None` if no row exists yet (caller uses defaults). -/// `upsert_settings` inserts or replaces the single row. +/// `upsert_settings` inserts or replaces the current user's row. #[async_trait::async_trait] pub trait ISettingsRepository: Send + Sync { /// Returns the settings row, or `None` if no settings have been persisted. - async fn get_settings(&self) -> Result, DbError>; + async fn get_settings(&self, user_id: &str) -> Result, DbError>; /// Inserts or replaces the single settings row. async fn upsert_settings( &self, + user_id: &str, language: &str, notification_enabled: bool, cron_notification_enabled: bool, diff --git a/crates/aionui-db/src/repository/skill.rs b/crates/aionui-db/src/repository/skill.rs index 1a2c80931..3f4d09202 100644 --- a/crates/aionui-db/src/repository/skill.rs +++ b/crates/aionui-db/src/repository/skill.rs @@ -7,26 +7,60 @@ pub trait ISkillRepository: Send + Sync { /// Returns active skills ordered by most recent update first. async fn list(&self) -> Result, DbError>; + /// Returns active global and user-owned skills for `user_id`. + async fn list_for_user(&self, user_id: &str) -> Result, DbError>; + /// Finds an active skill by name. async fn find_by_name(&self, name: &str) -> Result, DbError>; + /// Finds an active global or user-owned skill by name. + async fn find_by_name_for_user(&self, user_id: &str, name: &str) -> Result, DbError>; + /// Finds a skill by name, including soft-deleted rows. async fn find_by_name_any(&self, name: &str) -> Result, DbError>; + /// Finds a global or user-owned skill by name, including soft-deleted rows. + async fn find_by_name_any_for_user(&self, user_id: &str, name: &str) -> Result, DbError>; + /// Creates or updates a user skill by name and clears soft-delete state. async fn upsert(&self, params: UpsertSkillParams<'_>) -> Result; + /// Creates or updates a skill owned by `user_id` and clears soft-delete state. + async fn upsert_for_user(&self, user_id: &str, params: UpsertSkillParams<'_>) -> Result; + + /// Creates or updates a global skill row (`user_id IS NULL`). + async fn upsert_global(&self, params: UpsertSkillParams<'_>) -> Result { + self.upsert(params).await + } + /// Soft-deletes an active skill by name. async fn delete_by_name(&self, name: &str) -> Result; + /// Soft-deletes an active skill owned by `user_id`. + async fn delete_by_name_for_user(&self, user_id: &str, name: &str) -> Result; + /// Appends one import record. async fn create_import_record( &self, params: CreateSkillImportRecordParams<'_>, ) -> Result; + /// Appends one import record owned by `user_id`. + async fn create_import_record_for_user( + &self, + user_id: &str, + params: CreateSkillImportRecordParams<'_>, + ) -> Result; + /// Lists recent import records ordered by creation time descending. async fn list_import_records(&self, limit: i64) -> Result, DbError>; + + /// Lists recent import records owned by `user_id`. + async fn list_import_records_for_user( + &self, + user_id: &str, + limit: i64, + ) -> Result, DbError>; } /// Parameters for creating or updating a skill row. diff --git a/crates/aionui-db/src/repository/sqlite_acp_session.rs b/crates/aionui-db/src/repository/sqlite_acp_session.rs index f373ea2a8..47a71eb6b 100644 --- a/crates/aionui-db/src/repository/sqlite_acp_session.rs +++ b/crates/aionui-db/src/repository/sqlite_acp_session.rs @@ -19,15 +19,9 @@ impl SqliteAcpSessionRepository { pub fn new(pool: SqlitePool) -> Self { Self { pool } } -} - -fn is_unique_violation(err: &dyn sqlx::error::DatabaseError) -> bool { - err.code().is_some_and(|c| c == "2067" || c == "1555") -} -#[async_trait::async_trait] -impl IAcpSessionRepository for SqliteAcpSessionRepository { - async fn get(&self, conversation_id: &str) -> Result, DbError> { + #[cfg(test)] + async fn get_unscoped(&self, conversation_id: &str) -> Result, DbError> { let row = sqlx::query_as::<_, AcpSessionRow>("SELECT * FROM acp_session WHERE conversation_id = ?") .bind(conversation_id) .fetch_optional(&self.pool) @@ -35,37 +29,8 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { Ok(row) } - async fn create(&self, params: &CreateAcpSessionParams<'_>) -> Result { - let now = now_ms(); - sqlx::query( - "INSERT INTO acp_session \ - (conversation_id, agent_source, agent_id, \ - session_id, session_status, session_config, last_active_at) \ - VALUES (?, ?, ?, NULL, 'idle', '{}', ?)", - ) - .bind(params.conversation_id) - .bind(params.agent_source) - .bind(params.agent_id) - .bind(now) - .execute(&self.pool) - .await - .map_err(|e| match &e { - sqlx::Error::Database(db_err) if is_unique_violation(db_err.as_ref()) => DbError::Conflict(format!( - "acp_session row for conversation '{}' already exists", - params.conversation_id - )), - _ => DbError::Query(e), - })?; - - self.get(params.conversation_id).await?.ok_or_else(|| { - DbError::Init(format!( - "create did not produce acp_session row for '{}'", - params.conversation_id - )) - }) - } - - async fn update_session_id(&self, conversation_id: &str, session_id: &str) -> Result { + #[cfg(test)] + async fn update_session_id_unscoped(&self, conversation_id: &str, session_id: &str) -> Result { let now = now_ms(); let result = sqlx::query("UPDATE acp_session SET session_id = ?, last_active_at = ? WHERE conversation_id = ?") .bind(session_id) @@ -76,18 +41,8 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { Ok(result.rows_affected() > 0) } - async fn clear_session_id(&self, conversation_id: &str) -> Result { - let now = now_ms(); - let result = - sqlx::query("UPDATE acp_session SET session_id = NULL, last_active_at = ? WHERE conversation_id = ?") - .bind(now) - .bind(conversation_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - async fn delete(&self, conversation_id: &str) -> Result { + #[cfg(test)] + async fn delete_unscoped(&self, conversation_id: &str) -> Result { let result = sqlx::query("DELETE FROM acp_session WHERE conversation_id = ?") .bind(conversation_id) .execute(&self.pool) @@ -95,7 +50,11 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { Ok(result.rows_affected() > 0) } - async fn load_runtime_state(&self, conversation_id: &str) -> Result, DbError> { + #[cfg(test)] + async fn load_runtime_state_unscoped( + &self, + conversation_id: &str, + ) -> Result, DbError> { let raw: Option = sqlx::query_scalar("SELECT session_config FROM acp_session WHERE conversation_id = ?") .bind(conversation_id) @@ -106,24 +65,11 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { return Ok(None); }; - let parsed: Value = - serde_json::from_str(&raw).map_err(|e| DbError::Init(format!("invalid session_config JSON: {e}")))?; - let runtime = parsed.get("runtime"); - - let mut state = PersistedSessionState::default(); - if let Some(rt) = runtime { - state.current_mode_id = rt.get("current_mode_id").and_then(Value::as_str).map(ToOwned::to_owned); - state.current_model_id = rt - .get("current_model_id") - .and_then(Value::as_str) - .map(ToOwned::to_owned); - state.config_selections_json = rt.get("config_selections").map(serde_json::Value::to_string); - state.context_usage_json = rt.get("context_usage").map(serde_json::Value::to_string); - } - Ok(Some(state)) + Ok(Some(decode_runtime_state(&raw)?)) } - async fn save_runtime_state( + #[cfg(test)] + async fn save_runtime_state_unscoped( &self, conversation_id: &str, params: &SaveRuntimeStateParams<'_>, @@ -145,71 +91,280 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { return Ok(false); }; - let mut parsed: Value = serde_json::from_str(&raw).unwrap_or_else(|_| Value::Object(Default::default())); - let runtime = parsed - .as_object_mut() - .ok_or_else(|| DbError::Init("session_config is not a JSON object".into()))? - .entry("runtime") - .or_insert_with(|| Value::Object(Default::default())); - let runtime = runtime - .as_object_mut() - .ok_or_else(|| DbError::Init("session_config.runtime is not a JSON object".into()))?; - - if let Some(outer) = params.current_mode_id { - match outer { - Some(v) => { - runtime.insert("current_mode_id".into(), Value::String(v.to_owned())); - } - None => { - runtime.remove("current_mode_id"); - } + let new_config = merge_runtime_state(&raw, params)?; + let now = now_ms(); + let result = sqlx::query( + "UPDATE acp_session SET session_config = ?, last_active_at = ? \ + WHERE conversation_id = ?", + ) + .bind(new_config) + .bind(now) + .bind(conversation_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } +} + +fn is_unique_violation(err: &dyn sqlx::error::DatabaseError) -> bool { + err.code().is_some_and(|c| c == "2067" || c == "1555") +} + +fn decode_runtime_state(raw: &str) -> Result { + let parsed: Value = + serde_json::from_str(raw).map_err(|e| DbError::Init(format!("invalid session_config JSON: {e}")))?; + let runtime = parsed.get("runtime"); + + let mut state = PersistedSessionState::default(); + if let Some(rt) = runtime { + state.current_mode_id = rt.get("current_mode_id").and_then(Value::as_str).map(ToOwned::to_owned); + state.current_model_id = rt + .get("current_model_id") + .and_then(Value::as_str) + .map(ToOwned::to_owned); + state.config_selections_json = rt.get("config_selections").map(serde_json::Value::to_string); + state.context_usage_json = rt.get("context_usage").map(serde_json::Value::to_string); + } + Ok(state) +} + +fn merge_runtime_state(raw: &str, params: &SaveRuntimeStateParams<'_>) -> Result { + let mut parsed: Value = serde_json::from_str(raw).unwrap_or_else(|_| Value::Object(Default::default())); + let runtime = parsed + .as_object_mut() + .ok_or_else(|| DbError::Init("session_config is not a JSON object".into()))? + .entry("runtime") + .or_insert_with(|| Value::Object(Default::default())); + let runtime = runtime + .as_object_mut() + .ok_or_else(|| DbError::Init("session_config.runtime is not a JSON object".into()))?; + + if let Some(outer) = params.current_mode_id { + match outer { + Some(v) => { + runtime.insert("current_mode_id".into(), Value::String(v.to_owned())); + } + None => { + runtime.remove("current_mode_id"); } } - if let Some(outer) = params.current_model_id { - match outer { - Some(v) => { - runtime.insert("current_model_id".into(), Value::String(v.to_owned())); - } - None => { - runtime.remove("current_model_id"); - } + } + if let Some(outer) = params.current_model_id { + match outer { + Some(v) => { + runtime.insert("current_model_id".into(), Value::String(v.to_owned())); + } + None => { + runtime.remove("current_model_id"); + } + } + } + if let Some(outer) = params.config_selections_json { + match outer { + Some(json) => { + let v: Value = serde_json::from_str(json) + .map_err(|e| DbError::Init(format!("invalid config_selections JSON: {e}")))?; + runtime.insert("config_selections".into(), v); + } + None => { + runtime.remove("config_selections"); } } - if let Some(outer) = params.config_selections_json { - match outer { - Some(json) => { - let v: Value = serde_json::from_str(json) - .map_err(|e| DbError::Init(format!("invalid config_selections JSON: {e}")))?; - runtime.insert("config_selections".into(), v); - } - None => { - runtime.remove("config_selections"); - } + } + if let Some(outer) = params.context_usage_json { + match outer { + Some(json) => { + let v: Value = serde_json::from_str(json) + .map_err(|e| DbError::Init(format!("invalid context_usage JSON: {e}")))?; + runtime.insert("context_usage".into(), v); + } + None => { + runtime.remove("context_usage"); } } - if let Some(outer) = params.context_usage_json { - match outer { - Some(json) => { - let v: Value = serde_json::from_str(json) - .map_err(|e| DbError::Init(format!("invalid context_usage JSON: {e}")))?; - runtime.insert("context_usage".into(), v); - } - None => { - runtime.remove("context_usage"); - } + } + + serde_json::to_string(&parsed).map_err(|e| DbError::Init(format!("encode session_config: {e}"))) +} + +#[async_trait::async_trait] +impl IAcpSessionRepository for SqliteAcpSessionRepository { + async fn get_for_user(&self, user_id: &str, conversation_id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, AcpSessionRow>( + "SELECT a.* FROM acp_session a \ + WHERE a.conversation_id = ? \ + AND EXISTS (SELECT 1 FROM conversations c WHERE c.id = a.conversation_id AND c.user_id = ?)", + ) + .bind(conversation_id) + .bind(user_id) + .fetch_optional(&self.pool) + .await?; + Ok(row) + } + + async fn create(&self, params: &CreateAcpSessionParams<'_>) -> Result { + let now = now_ms(); + let result = sqlx::query( + "INSERT INTO acp_session \ + (conversation_id, agent_source, agent_id, \ + session_id, session_status, session_config, last_active_at) \ + SELECT c.id, ?, ?, NULL, 'idle', '{}', ? \ + FROM conversations c \ + WHERE c.id = ? AND c.user_id = ?", + ) + .bind(params.agent_source) + .bind(params.agent_id) + .bind(now) + .bind(params.conversation_id) + .bind(params.user_id) + .execute(&self.pool) + .await + .map_err(|e| match &e { + sqlx::Error::Database(db_err) if is_unique_violation(db_err.as_ref()) => DbError::Conflict(format!( + "acp_session row for conversation '{}' already exists", + params.conversation_id + )), + _ => DbError::Query(e), + })?; + + if result.rows_affected() == 0 { + let conversation_owner = sqlx::query_scalar::<_, String>("SELECT user_id FROM conversations WHERE id = ?") + .bind(params.conversation_id) + .fetch_optional(&self.pool) + .await?; + + if conversation_owner + .as_deref() + .is_some_and(|user_id| user_id != params.user_id) + { + return Err(DbError::Conflict(format!( + "CROSS_ACCOUNT_REFERENCE: acp_session conversation '{}' belongs to another user", + params.conversation_id + ))); } + + return Err(DbError::NotFound(format!( + "conversation '{}' for user '{}'", + params.conversation_id, params.user_id + ))); } - let new_config = - serde_json::to_string(&parsed).map_err(|e| DbError::Init(format!("encode session_config: {e}")))?; + self.get_for_user(params.user_id, params.conversation_id) + .await? + .ok_or_else(|| { + DbError::Init(format!( + "create did not produce acp_session row for '{}'", + params.conversation_id + )) + }) + } + + async fn update_session_id_for_user( + &self, + user_id: &str, + conversation_id: &str, + session_id: &str, + ) -> Result { + let now = now_ms(); + let result = sqlx::query( + "UPDATE acp_session SET session_id = ?, last_active_at = ? \ + WHERE conversation_id = ? \ + AND EXISTS (SELECT 1 FROM conversations c WHERE c.id = acp_session.conversation_id AND c.user_id = ?)", + ) + .bind(session_id) + .bind(now) + .bind(conversation_id) + .bind(user_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + async fn clear_session_id_for_user(&self, user_id: &str, conversation_id: &str) -> Result { + let now = now_ms(); + let result = sqlx::query( + "UPDATE acp_session SET session_id = NULL, last_active_at = ? \ + WHERE conversation_id = ? \ + AND EXISTS (SELECT 1 FROM conversations c WHERE c.id = acp_session.conversation_id AND c.user_id = ?)", + ) + .bind(now) + .bind(conversation_id) + .bind(user_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + async fn delete_for_user(&self, user_id: &str, conversation_id: &str) -> Result { + let result = sqlx::query( + "DELETE FROM acp_session \ + WHERE conversation_id = ? \ + AND EXISTS (SELECT 1 FROM conversations c WHERE c.id = acp_session.conversation_id AND c.user_id = ?)", + ) + .bind(conversation_id) + .bind(user_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + async fn load_runtime_state_for_user( + &self, + user_id: &str, + conversation_id: &str, + ) -> Result, DbError> { + let raw: Option = sqlx::query_scalar( + "SELECT a.session_config FROM acp_session a \ + WHERE a.conversation_id = ? \ + AND EXISTS (SELECT 1 FROM conversations c WHERE c.id = a.conversation_id AND c.user_id = ?)", + ) + .bind(conversation_id) + .bind(user_id) + .fetch_optional(&self.pool) + .await?; + + let Some(raw) = raw else { + return Ok(None); + }; + + Ok(Some(decode_runtime_state(&raw)?)) + } + + async fn save_runtime_state_for_user( + &self, + user_id: &str, + conversation_id: &str, + params: &SaveRuntimeStateParams<'_>, + ) -> Result { + if params.is_empty() { + return Ok(true); + } + + let raw: Option = sqlx::query_scalar( + "SELECT a.session_config FROM acp_session a \ + WHERE a.conversation_id = ? \ + AND EXISTS (SELECT 1 FROM conversations c WHERE c.id = a.conversation_id AND c.user_id = ?)", + ) + .bind(conversation_id) + .bind(user_id) + .fetch_optional(&self.pool) + .await?; + + let Some(raw) = raw else { + return Ok(false); + }; + + let new_config = merge_runtime_state(&raw, params)?; let now = now_ms(); let result = sqlx::query( "UPDATE acp_session SET session_config = ?, last_active_at = ? \ - WHERE conversation_id = ?", + WHERE conversation_id = ? \ + AND EXISTS (SELECT 1 FROM conversations c WHERE c.id = acp_session.conversation_id AND c.user_id = ?)", ) .bind(new_config) .bind(now) .bind(conversation_id) + .bind(user_id) .execute(&self.pool) .await?; Ok(result.rows_affected() > 0) @@ -224,17 +379,48 @@ mod tests { async fn setup() -> (SqliteAcpSessionRepository, crate::Database) { let db = init_database_memory().await.unwrap(); let repo = SqliteAcpSessionRepository::new(db.pool().clone()); + insert_conversation(&repo, "user-1", "conv-1").await; (repo, db) } fn create_params<'a>(conversation_id: &'a str) -> CreateAcpSessionParams<'a> { CreateAcpSessionParams { + user_id: "user-1", conversation_id, agent_source: "builtin", agent_id: "2d23ff1c", } } + async fn insert_conversation(repo: &SqliteAcpSessionRepository, user_id: &str, conversation_id: &str) { + let now = now_ms(); + sqlx::query( + "INSERT OR IGNORE INTO users \ + (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, ?, ?)", + ) + .bind(user_id) + .bind(user_id) + .bind(now) + .bind(now) + .execute(&repo.pool) + .await + .unwrap(); + sqlx::query( + "INSERT OR IGNORE INTO conversations \ + (id, user_id, name, type, extra, model, status, source, \ + channel_chat_id, pinned, pinned_at, created_at, updated_at) \ + VALUES (?, ?, 'Test', 'acp', '{}', NULL, 'pending', NULL, NULL, 0, NULL, ?, ?)", + ) + .bind(conversation_id) + .bind(user_id) + .bind(now) + .bind(now) + .execute(&repo.pool) + .await + .unwrap(); + } + #[tokio::test] async fn create_then_get_roundtrips() { let (repo, _db) = setup().await; @@ -245,7 +431,7 @@ mod tests { assert_eq!(row.session_status, "idle"); assert_eq!(row.session_config, "{}"); - let fetched = repo.get("conv-1").await.unwrap().unwrap(); + let fetched = repo.get_unscoped("conv-1").await.unwrap().unwrap(); assert_eq!(fetched.conversation_id, "conv-1"); } @@ -257,13 +443,33 @@ mod tests { assert!(matches!(err, DbError::Conflict(_))); } + #[tokio::test] + async fn create_rejects_cross_user_parent() { + let (repo, _db) = setup().await; + let err = repo + .create(&CreateAcpSessionParams { + user_id: "user-2", + conversation_id: "conv-1", + agent_source: "builtin", + agent_id: "2d23ff1c", + }) + .await + .unwrap_err(); + + assert!(matches!( + err, + DbError::Conflict(msg) if msg.starts_with("CROSS_ACCOUNT_REFERENCE:") + )); + assert!(repo.get_for_user("user-1", "conv-1").await.unwrap().is_none()); + } + #[tokio::test] async fn update_session_id_flips_field() { let (repo, _db) = setup().await; repo.create(&create_params("conv-1")).await.unwrap(); - assert!(repo.update_session_id("conv-1", "sess-abc").await.unwrap()); + assert!(repo.update_session_id_unscoped("conv-1", "sess-abc").await.unwrap()); - let fetched = repo.get("conv-1").await.unwrap().unwrap(); + let fetched = repo.get_unscoped("conv-1").await.unwrap().unwrap(); assert_eq!(fetched.session_id.as_deref(), Some("sess-abc")); assert!(fetched.last_active_at.is_some()); } @@ -271,17 +477,40 @@ mod tests { #[tokio::test] async fn update_session_id_missing_row_returns_false() { let (repo, _db) = setup().await; - assert!(!repo.update_session_id("nope", "sid").await.unwrap()); + assert!(!repo.update_session_id_unscoped("nope", "sid").await.unwrap()); + } + + #[tokio::test] + async fn update_session_id_for_user_rejects_other_owner() { + let (repo, _db) = setup().await; + insert_conversation(&repo, "user-1", "conv-1").await; + repo.create(&create_params("conv-1")).await.unwrap(); + + assert!( + !repo + .update_session_id_for_user("user-2", "conv-1", "sess-other") + .await + .unwrap() + ); + assert!( + repo.update_session_id_for_user("user-1", "conv-1", "sess-owner") + .await + .unwrap() + ); + + let fetched = repo.get_unscoped("conv-1").await.unwrap().unwrap(); + assert_eq!(fetched.session_id.as_deref(), Some("sess-owner")); } #[tokio::test] async fn clear_session_id_nulls_field_but_keeps_row() { let (repo, _db) = setup().await; + insert_conversation(&repo, "user-1", "conv-1").await; repo.create(&create_params("conv-1")).await.unwrap(); - repo.update_session_id("conv-1", "sess-abc").await.unwrap(); + repo.update_session_id_unscoped("conv-1", "sess-abc").await.unwrap(); - assert!(repo.clear_session_id("conv-1").await.unwrap()); - let fetched = repo.get("conv-1").await.unwrap().unwrap(); + assert!(repo.clear_session_id_for_user("user-1", "conv-1").await.unwrap()); + let fetched = repo.get_unscoped("conv-1").await.unwrap().unwrap(); assert_eq!(fetched.session_id, None, "the resume anchor must be nulled"); assert_eq!( fetched.session_status, "idle", @@ -289,32 +518,62 @@ mod tests { ); } + #[tokio::test] + async fn clear_session_id_rejects_other_owner() { + let (repo, _db) = setup().await; + insert_conversation(&repo, "user-1", "conv-1").await; + repo.create(&create_params("conv-1")).await.unwrap(); + repo.update_session_id_unscoped("conv-1", "sess-abc").await.unwrap(); + + assert!(!repo.clear_session_id_for_user("intruder", "conv-1").await.unwrap()); + let fetched = repo.get_unscoped("conv-1").await.unwrap().unwrap(); + assert_eq!( + fetched.session_id.as_deref(), + Some("sess-abc"), + "a foreign user must not clear another owner's resume anchor" + ); + } + #[tokio::test] async fn clear_session_id_missing_row_returns_false() { let (repo, _db) = setup().await; - assert!(!repo.clear_session_id("nope").await.unwrap()); + assert!(!repo.clear_session_id_for_user("user-1", "nope").await.unwrap()); } #[tokio::test] async fn delete_removes_row() { let (repo, _db) = setup().await; repo.create(&create_params("conv-1")).await.unwrap(); - assert!(repo.delete("conv-1").await.unwrap()); - assert!(repo.get("conv-1").await.unwrap().is_none()); - assert!(!repo.delete("conv-1").await.unwrap()); + assert!(repo.delete_unscoped("conv-1").await.unwrap()); + assert!(repo.get_unscoped("conv-1").await.unwrap().is_none()); + assert!(!repo.delete_unscoped("conv-1").await.unwrap()); + } + + #[tokio::test] + async fn get_and_delete_for_user_are_owner_scoped() { + let (repo, _db) = setup().await; + insert_conversation(&repo, "user-1", "conv-1").await; + repo.create(&create_params("conv-1")).await.unwrap(); + + assert!(repo.get_for_user("user-2", "conv-1").await.unwrap().is_none()); + assert!(repo.get_for_user("user-1", "conv-1").await.unwrap().is_some()); + assert!(!repo.delete_for_user("user-2", "conv-1").await.unwrap()); + assert!(repo.get_unscoped("conv-1").await.unwrap().is_some()); + assert!(repo.delete_for_user("user-1", "conv-1").await.unwrap()); + assert!(repo.get_unscoped("conv-1").await.unwrap().is_none()); } #[tokio::test] async fn load_runtime_state_missing_row() { let (repo, _db) = setup().await; - assert!(repo.load_runtime_state("nope").await.unwrap().is_none()); + assert!(repo.load_runtime_state_unscoped("nope").await.unwrap().is_none()); } #[tokio::test] async fn load_runtime_state_empty_config_returns_defaults() { let (repo, _db) = setup().await; repo.create(&create_params("conv-1")).await.unwrap(); - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo.load_runtime_state_unscoped("conv-1").await.unwrap().unwrap(); assert_eq!(state, PersistedSessionState::default()); } @@ -324,7 +583,7 @@ mod tests { repo.create(&create_params("conv-1")).await.unwrap(); assert!( - repo.save_runtime_state( + repo.save_runtime_state_unscoped( "conv-1", &SaveRuntimeStateParams { current_mode_id: Some(Some("code")), @@ -337,7 +596,7 @@ mod tests { .unwrap() ); - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo.load_runtime_state_unscoped("conv-1").await.unwrap().unwrap(); assert_eq!(state.current_mode_id.as_deref(), Some("code")); assert_eq!(state.current_model_id.as_deref(), Some("claude-sonnet-4")); // The stored JSON should parse back to the same payload @@ -350,12 +609,58 @@ mod tests { assert_eq!(usage["total"], 100); } + #[tokio::test] + async fn runtime_state_for_user_is_owner_scoped() { + let (repo, _db) = setup().await; + insert_conversation(&repo, "user-1", "conv-1").await; + repo.create(&create_params("conv-1")).await.unwrap(); + + assert!( + !repo + .save_runtime_state_for_user( + "user-2", + "conv-1", + &SaveRuntimeStateParams { + current_mode_id: Some(Some("other-mode")), + ..Default::default() + }, + ) + .await + .unwrap() + ); + assert!( + repo.save_runtime_state_for_user( + "user-1", + "conv-1", + &SaveRuntimeStateParams { + current_mode_id: Some(Some("owner-mode")), + ..Default::default() + }, + ) + .await + .unwrap() + ); + + assert!( + repo.load_runtime_state_for_user("user-2", "conv-1") + .await + .unwrap() + .is_none() + ); + let state = repo + .load_runtime_state_for_user("user-1", "conv-1") + .await + .unwrap() + .unwrap(); + assert_eq!(state.current_mode_id.as_deref(), Some("owner-mode")); + } + #[tokio::test] async fn save_runtime_state_partial_preserves_siblings() { let (repo, _db) = setup().await; repo.create(&create_params("conv-1")).await.unwrap(); - repo.save_runtime_state( + repo.save_runtime_state_unscoped( "conv-1", &SaveRuntimeStateParams { current_mode_id: Some(Some("code")), @@ -367,7 +672,7 @@ mod tests { .unwrap(); // Later write only touches current_model_id. - repo.save_runtime_state( + repo.save_runtime_state_unscoped( "conv-1", &SaveRuntimeStateParams { current_model_id: Some(Some("opus-4")), @@ -377,7 +682,7 @@ mod tests { .await .unwrap(); - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo.load_runtime_state_unscoped("conv-1").await.unwrap().unwrap(); assert_eq!( state.current_mode_id.as_deref(), Some("code"), @@ -391,7 +696,7 @@ mod tests { let (repo, _db) = setup().await; repo.create(&create_params("conv-1")).await.unwrap(); - repo.save_runtime_state( + repo.save_runtime_state_unscoped( "conv-1", &SaveRuntimeStateParams { current_mode_id: Some(Some("code")), @@ -400,7 +705,7 @@ mod tests { ) .await .unwrap(); - repo.save_runtime_state( + repo.save_runtime_state_unscoped( "conv-1", &SaveRuntimeStateParams { current_mode_id: Some(None), @@ -410,7 +715,7 @@ mod tests { .await .unwrap(); - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo.load_runtime_state_unscoped("conv-1").await.unwrap().unwrap(); assert!(state.current_mode_id.is_none()); } @@ -419,11 +724,11 @@ mod tests { let (repo, _db) = setup().await; repo.create(&create_params("conv-1")).await.unwrap(); assert!( - repo.save_runtime_state("conv-1", &SaveRuntimeStateParams::default()) + repo.save_runtime_state_unscoped("conv-1", &SaveRuntimeStateParams::default()) .await .unwrap() ); - let state = repo.load_runtime_state("conv-1").await.unwrap().unwrap(); + let state = repo.load_runtime_state_unscoped("conv-1").await.unwrap().unwrap(); assert_eq!(state, PersistedSessionState::default()); } @@ -431,7 +736,7 @@ mod tests { async fn save_runtime_state_missing_row_returns_false() { let (repo, _db) = setup().await; let ok = repo - .save_runtime_state( + .save_runtime_state_unscoped( "nope", &SaveRuntimeStateParams { current_mode_id: Some(Some("x")), diff --git a/crates/aionui-db/src/repository/sqlite_agent_metadata.rs b/crates/aionui-db/src/repository/sqlite_agent_metadata.rs index e3eb7bf03..e09baea02 100644 --- a/crates/aionui-db/src/repository/sqlite_agent_metadata.rs +++ b/crates/aionui-db/src/repository/sqlite_agent_metadata.rs @@ -15,8 +15,10 @@ pub struct SqliteAgentMetadataRepository { pool: SqlitePool, } +const DEFAULT_USER_ID: &str = "system_default_user"; + const AGENT_METADATA_SAFE_COLUMNS: &str = "\ - id, icon, name, name_i18n, description, description_i18n, \ + id, user_id, icon, name, name_i18n, description, description_i18n, \ backend, agent_type, agent_source, agent_source_info, \ enabled, command, args, env, native_skills_dirs, \ behavior_policy, yolo_id, \ @@ -57,6 +59,7 @@ impl AgentMetadataCacheField { #[derive(Debug)] struct AgentMetadataSafeRow { id: String, + user_id: Option, icon: Option, name: String, name_i18n: Option, @@ -99,6 +102,7 @@ impl AgentMetadataSafeRow { fn from_sqlite_row(row: SqliteRow) -> Result { Ok(Self { id: row.try_get("id")?, + user_id: row.try_get("user_id")?, icon: row.try_get("icon")?, name: row.try_get("name")?, name_i18n: row.try_get("name_i18n")?, @@ -174,6 +178,7 @@ impl AgentMetadataSafeRow { ( AgentMetadataRow { id: self.id, + user_id: self.user_id, icon: self.icon, name: self.name, name_i18n: self.name_i18n, @@ -236,15 +241,6 @@ impl SqliteAgentMetadataRepository { Self { pool } } - async fn fetch_all_safe(&self, sql: &str) -> Result, DbError> { - let rows = sqlx::query(sql).fetch_all(&self.pool).await?; - let mut decoded = Vec::with_capacity(rows.len()); - for row in rows { - decoded.push(self.decode_and_repair(row).await?); - } - Ok(decoded) - } - async fn fetch_optional_safe(&self, sql: &str, bind: &str) -> Result, DbError> { let row = sqlx::query(sql).bind(bind).fetch_optional(&self.pool).await?; match row { @@ -270,20 +266,98 @@ impl SqliteAgentMetadataRepository { } } + async fn fetch_optional_safe_three_binds( + &self, + sql: &str, + first: &str, + second: &str, + third: &str, + ) -> Result, DbError> { + let row = sqlx::query(sql) + .bind(first) + .bind(second) + .bind(third) + .fetch_optional(&self.pool) + .await?; + match row { + Some(row) => Ok(Some(self.decode_and_repair(row).await?)), + None => Ok(None), + } + } + + async fn ensure_user_row_for_write(&self, user_id: &str, id: &str) -> Result { + let result = sqlx::query( + "INSERT INTO agent_metadata ( + id, user_id, icon, name, name_i18n, description, description_i18n, + backend, agent_type, agent_source, agent_source_info, + enabled, command, args, env, native_skills_dirs, + behavior_policy, yolo_id, + agent_capabilities, auth_methods, config_options, + available_modes, available_models, available_commands, + sort_order, + last_check_status, last_check_kind, last_check_error_code, last_check_error_message, + last_check_guidance, last_check_latency_ms, last_check_at, last_success_at, last_failure_at, + command_override, env_override, created_at, updated_at + ) + SELECT + id, ?, icon, name, name_i18n, description, description_i18n, + backend, agent_type, agent_source, agent_source_info, + enabled, command, args, env, native_skills_dirs, + behavior_policy, yolo_id, + agent_capabilities, auth_methods, config_options, + available_modes, available_models, available_commands, + sort_order, + last_check_status, last_check_kind, last_check_error_code, last_check_error_message, + last_check_guidance, last_check_latency_ms, last_check_at, last_success_at, last_failure_at, + command_override, env_override, created_at, ? + FROM agent_metadata + WHERE id = ? AND (user_id = ? OR user_id IS NULL) + ORDER BY user_id IS NULL ASC + LIMIT 1 + ON CONFLICT(user_id, id) WHERE user_id IS NOT NULL DO NOTHING", + ) + .bind(user_id) + .bind(now_ms()) + .bind(id) + .bind(user_id) + .execute(&self.pool) + .await?; + + Ok(result.rows_affected() > 0 || self.user_row_exists(user_id, id).await?) + } + + async fn user_row_exists(&self, user_id: &str, id: &str) -> Result { + let exists = + sqlx::query_scalar::<_, i64>("SELECT EXISTS(SELECT 1 FROM agent_metadata WHERE user_id = ? AND id = ?)") + .bind(user_id) + .bind(id) + .fetch_one(&self.pool) + .await?; + Ok(exists != 0) + } + async fn decode_and_repair(&self, row: SqliteRow) -> Result { let safe = AgentMetadataSafeRow::from_sqlite_row(row)?; + let user_id = safe.user_id.clone(); let (model, invalid_fields) = safe.into_model(); if !invalid_fields.is_empty() { - self.clear_invalid_utf8_cache_fields(&model.id, &invalid_fields).await; + self.clear_invalid_utf8_cache_fields(&model.id, user_id.as_deref(), &invalid_fields) + .await; } Ok(model) } - async fn clear_invalid_utf8_cache_fields(&self, id: &str, fields: &[AgentMetadataCacheField]) { + async fn clear_invalid_utf8_cache_fields( + &self, + id: &str, + user_id: Option<&str>, + fields: &[AgentMetadataCacheField], + ) { for field in fields { warn!( table = "agent_metadata", row_id = %id, + user_id = user_id, field = field.column_name(), action = "clear_invalid_utf8", "Clearing invalid UTF-8 from rebuildable agent metadata cache field" @@ -292,14 +366,21 @@ impl SqliteAgentMetadataRepository { for field in fields { let sql = format!( - "UPDATE agent_metadata SET {} = NULL, updated_at = ? WHERE id = ?", + "UPDATE agent_metadata SET {} = NULL, updated_at = ? WHERE id = ? AND user_id IS ?", field.column_name() ); - if let Err(err) = sqlx::query(&sql).bind(now_ms()).bind(id).execute(&self.pool).await { + if let Err(err) = sqlx::query(&sql) + .bind(now_ms()) + .bind(id) + .bind(user_id) + .execute(&self.pool) + .await + { warn!( table = "agent_metadata", row_id = %id, + user_id = user_id, field = field.column_name(), action = "clear_invalid_utf8_failed", error = %err, @@ -313,16 +394,51 @@ impl SqliteAgentMetadataRepository { #[async_trait::async_trait] impl IAgentMetadataRepository for SqliteAgentMetadataRepository { async fn list_all(&self) -> Result, DbError> { - self.fetch_all_safe(&format!( - "SELECT {AGENT_METADATA_SAFE_COLUMNS} FROM agent_metadata ORDER BY sort_order ASC, name ASC" - )) - .await + self.list_all_for_user(DEFAULT_USER_ID).await + } + + async fn list_all_for_user(&self, user_id: &str) -> Result, DbError> { + let sql = format!( + "SELECT {AGENT_METADATA_SAFE_COLUMNS} + FROM agent_metadata am + WHERE am.user_id = ? + UNION ALL + SELECT {AGENT_METADATA_SAFE_COLUMNS} + FROM agent_metadata am + WHERE am.user_id IS NULL + AND NOT EXISTS ( + SELECT 1 FROM agent_metadata u + WHERE u.user_id = ? AND u.id = am.id + ) + ORDER BY sort_order ASC, name ASC" + ); + let rows = sqlx::query(&sql) + .bind(user_id) + .bind(user_id) + .fetch_all(&self.pool) + .await?; + let mut decoded = Vec::with_capacity(rows.len()); + for row in rows { + decoded.push(self.decode_and_repair(row).await?); + } + Ok(decoded) } async fn get(&self, id: &str) -> Result, DbError> { - self.fetch_optional_safe( - &format!("SELECT {AGENT_METADATA_SAFE_COLUMNS} FROM agent_metadata WHERE id = ?"), + self.get_for_user(DEFAULT_USER_ID, id).await + } + + async fn get_for_user(&self, user_id: &str, id: &str) -> Result, DbError> { + self.fetch_optional_safe_two_binds( + &format!( + "SELECT {AGENT_METADATA_SAFE_COLUMNS} + FROM agent_metadata + WHERE id = ? AND (user_id = ? OR user_id IS NULL) + ORDER BY user_id IS NULL ASC + LIMIT 1" + ), id, + user_id, ) .await } @@ -332,107 +448,155 @@ impl IAgentMetadataRepository for SqliteAgentMetadataRepository { agent_source: &str, name: &str, ) -> Result, DbError> { - self.fetch_optional_safe_two_binds( - &format!("SELECT {AGENT_METADATA_SAFE_COLUMNS} FROM agent_metadata WHERE agent_source = ? AND name = ?"), + self.find_by_source_and_name_for_user(DEFAULT_USER_ID, agent_source, name) + .await + } + + async fn find_by_source_and_name_for_user( + &self, + user_id: &str, + agent_source: &str, + name: &str, + ) -> Result, DbError> { + self.fetch_optional_safe_three_binds( + &format!( + "SELECT {AGENT_METADATA_SAFE_COLUMNS} + FROM agent_metadata + WHERE agent_source = ? AND name = ? AND (user_id = ? OR user_id IS NULL) + ORDER BY user_id IS NULL ASC, sort_order ASC, name ASC + LIMIT 1" + ), agent_source, name, + user_id, ) .await } async fn find_builtin_by_backend(&self, backend: &str) -> Result, DbError> { - self.fetch_optional_safe( + self.find_builtin_by_backend_for_user(DEFAULT_USER_ID, backend).await + } + + async fn find_builtin_by_backend_for_user( + &self, + user_id: &str, + backend: &str, + ) -> Result, DbError> { + self.fetch_optional_safe_two_binds( &format!( - "SELECT {AGENT_METADATA_SAFE_COLUMNS} FROM agent_metadata \ - WHERE agent_source = 'builtin' AND backend = ? \ - ORDER BY sort_order ASC, name ASC LIMIT 1" + "SELECT {AGENT_METADATA_SAFE_COLUMNS} + FROM agent_metadata + WHERE agent_source = 'builtin' + AND backend = ? + AND (user_id = ? OR user_id IS NULL) + ORDER BY user_id IS NULL ASC, sort_order ASC, name ASC + LIMIT 1" ), backend, + user_id, ) .await } async fn upsert(&self, params: &UpsertAgentMetadataParams<'_>) -> Result { + if matches!(params.agent_source, "builtin" | "internal") { + self.upsert_global(params).await + } else { + self.upsert_for_user(DEFAULT_USER_ID, params).await + } + } + + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAgentMetadataParams<'_>, + ) -> Result { + self.upsert_with_user_id(Some(user_id), params).await + } + + async fn upsert_global(&self, params: &UpsertAgentMetadataParams<'_>) -> Result { + self.upsert_with_user_id(None, params).await + } + + async fn apply_handshake( + &self, + id: &str, + params: &UpdateAgentHandshakeParams<'_>, + ) -> Result, DbError> { + if self.user_row_exists(DEFAULT_USER_ID, id).await? { + return self.apply_handshake_for_user(DEFAULT_USER_ID, id, params).await; + } + + let Some(existing) = self + .fetch_optional_safe( + &format!("SELECT {AGENT_METADATA_SAFE_COLUMNS} FROM agent_metadata WHERE user_id IS NULL AND id = ?"), + id, + ) + .await? + else { + return Ok(None); + }; + let now = now_ms(); + let agent_capabilities = params + .agent_capabilities + .map_or(existing.agent_capabilities, |v| v.map(String::from)); + let auth_methods = params + .auth_methods + .map_or(existing.auth_methods, |v| v.map(String::from)); + let config_options = params + .config_options + .map_or(existing.config_options, |v| v.map(String::from)); + let available_modes = params + .available_modes + .map_or(existing.available_modes, |v| v.map(String::from)); + let available_models = params + .available_models + .map_or(existing.available_models, |v| v.map(String::from)); + let available_commands = params + .available_commands + .map_or(existing.available_commands, |v| v.map(String::from)); sqlx::query( - "INSERT INTO agent_metadata \ - (id, icon, name, name_i18n, description, description_i18n, \ - backend, agent_type, agent_source, agent_source_info, \ - enabled, command, args, env, native_skills_dirs, \ - behavior_policy, yolo_id, \ - agent_capabilities, auth_methods, config_options, \ - available_modes, available_models, available_commands, \ - sort_order, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \ - ON CONFLICT(id) DO UPDATE SET \ - icon = excluded.icon, \ - name = excluded.name, \ - name_i18n = excluded.name_i18n, \ - description = excluded.description, \ - description_i18n = excluded.description_i18n, \ - backend = excluded.backend, \ - agent_type = excluded.agent_type, \ - agent_source = excluded.agent_source, \ - agent_source_info = excluded.agent_source_info, \ - enabled = excluded.enabled, \ - command = excluded.command, \ - args = excluded.args, \ - env = excluded.env, \ - native_skills_dirs = excluded.native_skills_dirs, \ - behavior_policy = excluded.behavior_policy, \ - yolo_id = excluded.yolo_id, \ - agent_capabilities = excluded.agent_capabilities, \ - auth_methods = excluded.auth_methods, \ - config_options = excluded.config_options, \ - available_modes = excluded.available_modes, \ - available_models = excluded.available_models, \ - available_commands = excluded.available_commands, \ - sort_order = excluded.sort_order, \ - updated_at = excluded.updated_at", + "UPDATE agent_metadata SET \ + agent_capabilities = ?, \ + auth_methods = ?, \ + config_options = ?, \ + available_modes = ?, \ + available_models = ?, \ + available_commands = ?, \ + updated_at = ? \ + WHERE user_id IS NULL AND id = ?", ) - .bind(params.id) - .bind(params.icon) - .bind(params.name) - .bind(params.name_i18n) - .bind(params.description) - .bind(params.description_i18n) - .bind(params.backend) - .bind(params.agent_type) - .bind(params.agent_source) - .bind(params.agent_source_info) - .bind(params.enabled) - .bind(params.command) - .bind(params.args) - .bind(params.env) - .bind(params.native_skills_dirs) - .bind(params.behavior_policy) - .bind(params.yolo_id) - .bind(params.agent_capabilities) - .bind(params.auth_methods) - .bind(params.config_options) - .bind(params.available_modes) - .bind(params.available_models) - .bind(params.available_commands) - .bind(params.sort_order) - .bind(now) + .bind(&agent_capabilities) + .bind(&auth_methods) + .bind(&config_options) + .bind(&available_modes) + .bind(&available_models) + .bind(&available_commands) .bind(now) + .bind(id) .execute(&self.pool) .await?; - let row = self - .get(params.id) - .await? - .ok_or_else(|| DbError::Init(format!("upsert did not produce row for id '{}'", params.id)))?; - Ok(row) + self.fetch_optional_safe( + &format!("SELECT {AGENT_METADATA_SAFE_COLUMNS} FROM agent_metadata WHERE user_id IS NULL AND id = ?"), + id, + ) + .await } - async fn apply_handshake( + async fn apply_handshake_for_user( &self, + user_id: &str, id: &str, params: &UpdateAgentHandshakeParams<'_>, ) -> Result, DbError> { - let Some(existing) = self.get(id).await? else { + if !self.ensure_user_row_for_write(user_id, id).await? { + return Ok(None); + } + + let Some(existing) = self.get_for_user(user_id, id).await? else { return Ok(None); }; @@ -465,7 +629,7 @@ impl IAgentMetadataRepository for SqliteAgentMetadataRepository { available_models = ?, \ available_commands = ?, \ updated_at = ? \ - WHERE id = ?", + WHERE user_id = ? AND id = ?", ) .bind(&agent_capabilities) .bind(&auth_methods) @@ -474,11 +638,12 @@ impl IAgentMetadataRepository for SqliteAgentMetadataRepository { .bind(&available_models) .bind(&available_commands) .bind(now) + .bind(user_id) .bind(id) .execute(&self.pool) .await?; - self.get(id).await + self.get_for_user(user_id, id).await } async fn update_availability_snapshot( @@ -486,6 +651,20 @@ impl IAgentMetadataRepository for SqliteAgentMetadataRepository { id: &str, params: &UpdateAgentAvailabilitySnapshotParams<'_>, ) -> Result, DbError> { + self.update_availability_snapshot_for_user(DEFAULT_USER_ID, id, params) + .await + } + + async fn update_availability_snapshot_for_user( + &self, + user_id: &str, + id: &str, + params: &UpdateAgentAvailabilitySnapshotParams<'_>, + ) -> Result, DbError> { + if !self.ensure_user_row_for_write(user_id, id).await? { + return Ok(None); + } + let now = now_ms(); let result = sqlx::query( "UPDATE agent_metadata SET \ @@ -499,7 +678,7 @@ impl IAgentMetadataRepository for SqliteAgentMetadataRepository { last_success_at = ?, \ last_failure_at = ?, \ updated_at = ? \ - WHERE id = ?", + WHERE user_id = ? AND id = ?", ) .bind(params.last_check_status) .bind(params.last_check_kind) @@ -511,6 +690,7 @@ impl IAgentMetadataRepository for SqliteAgentMetadataRepository { .bind(params.last_success_at) .bind(params.last_failure_at) .bind(now) + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -519,7 +699,7 @@ impl IAgentMetadataRepository for SqliteAgentMetadataRepository { return Ok(None); } - self.get(id).await + self.get_for_user(user_id, id).await } async fn update_agent_overrides( @@ -528,13 +708,29 @@ impl IAgentMetadataRepository for SqliteAgentMetadataRepository { command_override: Option<&str>, env_override: Option<&str>, ) -> Result<(), DbError> { + self.update_agent_overrides_for_user(DEFAULT_USER_ID, id, command_override, env_override) + .await + } + + async fn update_agent_overrides_for_user( + &self, + user_id: &str, + id: &str, + command_override: Option<&str>, + env_override: Option<&str>, + ) -> Result<(), DbError> { + if !self.ensure_user_row_for_write(user_id, id).await? { + return Ok(()); + } + sqlx::query( "UPDATE agent_metadata SET command_override = ?, env_override = ?, \ - updated_at = ? WHERE id = ?", + updated_at = ? WHERE user_id = ? AND id = ?", ) .bind(command_override) .bind(env_override) .bind(aionui_common::now_ms()) + .bind(user_id) .bind(id) .execute(&self.pool) .await @@ -543,10 +739,19 @@ impl IAgentMetadataRepository for SqliteAgentMetadataRepository { } async fn set_enabled(&self, id: &str, enabled: bool) -> Result { + self.set_enabled_for_user(DEFAULT_USER_ID, id, enabled).await + } + + async fn set_enabled_for_user(&self, user_id: &str, id: &str, enabled: bool) -> Result { + if !self.ensure_user_row_for_write(user_id, id).await? { + return Ok(false); + } + let now = now_ms(); - let result = sqlx::query("UPDATE agent_metadata SET enabled = ?, updated_at = ? WHERE id = ?") + let result = sqlx::query("UPDATE agent_metadata SET enabled = ?, updated_at = ? WHERE user_id = ? AND id = ?") .bind(enabled) .bind(now) + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -554,7 +759,12 @@ impl IAgentMetadataRepository for SqliteAgentMetadataRepository { } async fn delete(&self, id: &str) -> Result { - let result = sqlx::query("DELETE FROM agent_metadata WHERE id = ?") + self.delete_for_user(DEFAULT_USER_ID, id).await + } + + async fn delete_for_user(&self, user_id: &str, id: &str) -> Result { + let result = sqlx::query("DELETE FROM agent_metadata WHERE user_id = ? AND id = ?") + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -562,17 +772,132 @@ impl IAgentMetadataRepository for SqliteAgentMetadataRepository { } } +impl SqliteAgentMetadataRepository { + async fn upsert_with_user_id( + &self, + user_id: Option<&str>, + params: &UpsertAgentMetadataParams<'_>, + ) -> Result { + let now = now_ms(); + + let conflict_target = if user_id.is_some() { + "ON CONFLICT(user_id, id) WHERE user_id IS NOT NULL DO UPDATE SET" + } else { + "ON CONFLICT(id) WHERE user_id IS NULL DO UPDATE SET" + }; + + let sql = format!( + "INSERT INTO agent_metadata \ + (id, user_id, icon, name, name_i18n, description, description_i18n, \ + backend, agent_type, agent_source, agent_source_info, \ + enabled, command, args, env, native_skills_dirs, \ + behavior_policy, yolo_id, \ + agent_capabilities, auth_methods, config_options, \ + available_modes, available_models, available_commands, \ + sort_order, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \ + {conflict_target} \ + icon = excluded.icon, \ + name = excluded.name, \ + name_i18n = excluded.name_i18n, \ + description = excluded.description, \ + description_i18n = excluded.description_i18n, \ + backend = excluded.backend, \ + agent_type = excluded.agent_type, \ + agent_source = excluded.agent_source, \ + agent_source_info = excluded.agent_source_info, \ + enabled = excluded.enabled, \ + command = excluded.command, \ + args = excluded.args, \ + env = excluded.env, \ + native_skills_dirs = excluded.native_skills_dirs, \ + behavior_policy = excluded.behavior_policy, \ + yolo_id = excluded.yolo_id, \ + agent_capabilities = excluded.agent_capabilities, \ + auth_methods = excluded.auth_methods, \ + config_options = excluded.config_options, \ + available_modes = excluded.available_modes, \ + available_models = excluded.available_models, \ + available_commands = excluded.available_commands, \ + sort_order = excluded.sort_order, \ + updated_at = excluded.updated_at" + ); + + sqlx::query(&sql) + .bind(params.id) + .bind(user_id) + .bind(params.icon) + .bind(params.name) + .bind(params.name_i18n) + .bind(params.description) + .bind(params.description_i18n) + .bind(params.backend) + .bind(params.agent_type) + .bind(params.agent_source) + .bind(params.agent_source_info) + .bind(params.enabled) + .bind(params.command) + .bind(params.args) + .bind(params.env) + .bind(params.native_skills_dirs) + .bind(params.behavior_policy) + .bind(params.yolo_id) + .bind(params.agent_capabilities) + .bind(params.auth_methods) + .bind(params.config_options) + .bind(params.available_modes) + .bind(params.available_models) + .bind(params.available_commands) + .bind(params.sort_order) + .bind(now) + .bind(now) + .execute(&self.pool) + .await?; + + let row = match user_id { + Some(user_id) => self.get_for_user(user_id, params.id).await?, + None => { + self.fetch_optional_safe( + &format!( + "SELECT {AGENT_METADATA_SAFE_COLUMNS} FROM agent_metadata WHERE user_id IS NULL AND id = ?" + ), + params.id, + ) + .await? + } + } + .ok_or_else(|| DbError::Init(format!("upsert did not produce row for id '{}'", params.id)))?; + Ok(row) + } +} + #[cfg(test)] mod tests { use super::*; use crate::init_database_memory; + const USER_A: &str = "system_default_user"; + const USER_B: &str = "agent_user_b"; + async fn setup() -> (SqliteAgentMetadataRepository, crate::Database) { let db = init_database_memory().await.unwrap(); let repo = SqliteAgentMetadataRepository::new(db.pool().clone()); + ensure_user(&db, USER_B).await; (repo, db) } + async fn ensure_user(db: &crate::Database, user_id: &str) { + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, 1, 1)", + ) + .bind(user_id) + .bind(user_id) + .execute(db.pool()) + .await + .unwrap(); + } + async fn corrupt_cache_field(db: &crate::Database, id: &str, field: &str, bytes_hex: &str) { assert!( matches!( @@ -1035,4 +1360,89 @@ mod tests { // seed columns untouched assert_eq!(row.name, "agent-x"); } + + #[tokio::test] + async fn global_builtin_rows_are_visible_to_all_users() { + let (repo, _db) = setup().await; + + let user_a = repo.get_for_user(USER_A, "2d23ff1c").await.unwrap().unwrap(); + let user_b = repo.get_for_user(USER_B, "2d23ff1c").await.unwrap().unwrap(); + + assert_eq!(user_a.name, "Claude Code"); + assert_eq!(user_b.name, "Claude Code"); + assert_eq!(user_a.agent_source, "builtin"); + assert_eq!(user_b.agent_source, "builtin"); + } + + #[tokio::test] + async fn custom_rows_with_same_id_are_isolated_by_user() { + let (repo, _db) = setup().await; + + repo.upsert_for_user(USER_A, &custom_params("custom-shared", "User A Custom")) + .await + .unwrap(); + repo.upsert_for_user(USER_B, &custom_params("custom-shared", "User B Custom")) + .await + .unwrap(); + + let user_a = repo.get_for_user(USER_A, "custom-shared").await.unwrap().unwrap(); + let user_b = repo.get_for_user(USER_B, "custom-shared").await.unwrap().unwrap(); + assert_eq!(user_a.name, "User A Custom"); + assert_eq!(user_b.name, "User B Custom"); + + let user_a_names: Vec = repo + .list_all_for_user(USER_A) + .await + .unwrap() + .into_iter() + .filter(|row| row.id == "custom-shared") + .map(|row| row.name) + .collect(); + let user_b_names: Vec = repo + .list_all_for_user(USER_B) + .await + .unwrap() + .into_iter() + .filter(|row| row.id == "custom-shared") + .map(|row| row.name) + .collect(); + assert_eq!(user_a_names, vec!["User A Custom"]); + assert_eq!(user_b_names, vec!["User B Custom"]); + } + + #[tokio::test] + async fn user_enabled_toggle_shadows_global_builtin_without_mutating_global() { + let (repo, db) = setup().await; + + assert!(repo.set_enabled_for_user(USER_B, "2d23ff1c", false).await.unwrap()); + + let user_b = repo.get_for_user(USER_B, "2d23ff1c").await.unwrap().unwrap(); + assert!(!user_b.enabled); + + let global_enabled: bool = + sqlx::query_scalar("SELECT enabled FROM agent_metadata WHERE user_id IS NULL AND id = '2d23ff1c'") + .fetch_one(db.pool()) + .await + .unwrap(); + assert!(global_enabled); + } + + #[tokio::test] + async fn user_overrides_shadow_global_builtin_without_mutating_global() { + let (repo, db) = setup().await; + + repo.update_agent_overrides_for_user(USER_B, "2d23ff1c", Some("/tmp/claude"), Some("[]")) + .await + .unwrap(); + + let user_b = repo.get_for_user(USER_B, "2d23ff1c").await.unwrap().unwrap(); + assert_eq!(user_b.command_override.as_deref(), Some("/tmp/claude")); + + let global_override: Option = + sqlx::query_scalar("SELECT command_override FROM agent_metadata WHERE user_id IS NULL AND id = '2d23ff1c'") + .fetch_one(db.pool()) + .await + .unwrap(); + assert!(global_override.is_none()); + } } diff --git a/crates/aionui-db/src/repository/sqlite_assistant.rs b/crates/aionui-db/src/repository/sqlite_assistant.rs index afd11b38c..fd59bf401 100644 --- a/crates/aionui-db/src/repository/sqlite_assistant.rs +++ b/crates/aionui-db/src/repository/sqlite_assistant.rs @@ -14,6 +14,8 @@ use crate::repository::assistant::{ IAssistantPreferenceRepository, IAssistantRepository, }; +const DEFAULT_USER_ID: &str = "system_default_user"; + /// SQLite-backed implementation of [`IAssistantRepository`]. #[derive(Clone, Debug)] pub struct SqliteAssistantRepository { @@ -33,14 +35,25 @@ fn is_unique_violation(err: &dyn sqlx::error::DatabaseError) -> bool { #[async_trait::async_trait] impl IAssistantRepository for SqliteAssistantRepository { async fn list(&self) -> Result, DbError> { - let rows = sqlx::query_as::<_, AssistantRow>("SELECT * FROM assistants ORDER BY updated_at DESC") - .fetch_all(&self.pool) - .await?; + self.list_for_user(DEFAULT_USER_ID).await + } + + async fn list_for_user(&self, user_id: &str) -> Result, DbError> { + let rows = + sqlx::query_as::<_, AssistantRow>("SELECT * FROM assistants WHERE user_id = ? ORDER BY updated_at DESC") + .bind(user_id) + .fetch_all(&self.pool) + .await?; Ok(rows) } async fn get(&self, id: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, AssistantRow>("SELECT * FROM assistants WHERE id = ?") + self.get_for_user(DEFAULT_USER_ID, id).await + } + + async fn get_for_user(&self, user_id: &str, id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, AssistantRow>("SELECT * FROM assistants WHERE user_id = ? AND id = ?") + .bind(user_id) .bind(id) .fetch_optional(&self.pool) .await?; @@ -48,16 +61,25 @@ impl IAssistantRepository for SqliteAssistantRepository { } async fn create(&self, params: &CreateAssistantParams<'_>) -> Result { + self.create_for_user(DEFAULT_USER_ID, params).await + } + + async fn create_for_user( + &self, + user_id: &str, + params: &CreateAssistantParams<'_>, + ) -> Result { let now = now_ms(); sqlx::query( "INSERT INTO assistants \ - (id, name, description, avatar, enabled_skills, \ + (id, user_id, name, description, avatar, enabled_skills, \ custom_skill_names, disabled_builtin_skills, prompts, models, \ name_i18n, description_i18n, prompts_i18n, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(params.id) + .bind(user_id) .bind(params.name) .bind(params.description) .bind(params.avatar) @@ -99,7 +121,16 @@ impl IAssistantRepository for SqliteAssistantRepository { } async fn update(&self, id: &str, params: &UpdateAssistantParams<'_>) -> Result, DbError> { - let Some(existing) = self.get(id).await? else { + self.update_for_user(DEFAULT_USER_ID, id, params).await + } + + async fn update_for_user( + &self, + user_id: &str, + id: &str, + params: &UpdateAssistantParams<'_>, + ) -> Result, DbError> { + let Some(existing) = self.get_for_user(user_id, id).await? else { return Ok(None); }; @@ -111,7 +142,7 @@ impl IAssistantRepository for SqliteAssistantRepository { enabled_skills = ?, custom_skill_names = ?, disabled_builtin_skills = ?, \ prompts = ?, models = ?, name_i18n = ?, description_i18n = ?, \ prompts_i18n = ?, updated_at = ? \ - WHERE id = ?", + WHERE user_id = ? AND id = ?", ) .bind(&merged.name) .bind(&merged.description) @@ -125,6 +156,7 @@ impl IAssistantRepository for SqliteAssistantRepository { .bind(&merged.description_i18n) .bind(&merged.prompts_i18n) .bind(merged.updated_at) + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -133,7 +165,12 @@ impl IAssistantRepository for SqliteAssistantRepository { } async fn delete(&self, id: &str) -> Result { - let result = sqlx::query("DELETE FROM assistants WHERE id = ?") + self.delete_for_user(DEFAULT_USER_ID, id).await + } + + async fn delete_for_user(&self, user_id: &str, id: &str) -> Result { + let result = sqlx::query("DELETE FROM assistants WHERE user_id = ? AND id = ?") + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -141,14 +178,22 @@ impl IAssistantRepository for SqliteAssistantRepository { } async fn upsert(&self, params: &CreateAssistantParams<'_>) -> Result { + self.upsert_for_user(DEFAULT_USER_ID, params).await + } + + async fn upsert_for_user( + &self, + user_id: &str, + params: &CreateAssistantParams<'_>, + ) -> Result { let now = now_ms(); - sqlx::query( + let result = sqlx::query( "INSERT INTO assistants \ - (id, name, description, avatar, enabled_skills, \ + (id, user_id, name, description, avatar, enabled_skills, \ custom_skill_names, disabled_builtin_skills, prompts, models, \ name_i18n, description_i18n, prompts_i18n, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \ ON CONFLICT(id) DO UPDATE SET \ name = excluded.name, \ description = excluded.description, \ @@ -161,9 +206,11 @@ impl IAssistantRepository for SqliteAssistantRepository { name_i18n = excluded.name_i18n, \ description_i18n = excluded.description_i18n, \ prompts_i18n = excluded.prompts_i18n, \ - updated_at = excluded.updated_at", + updated_at = excluded.updated_at \ + WHERE assistants.user_id = excluded.user_id", ) .bind(params.id) + .bind(user_id) .bind(params.name) .bind(params.description) .bind(params.avatar) @@ -180,8 +227,15 @@ impl IAssistantRepository for SqliteAssistantRepository { .execute(&self.pool) .await?; + if result.rows_affected() == 0 { + return Err(DbError::Conflict(format!( + "Assistant with id '{}' already exists for another user", + params.id + ))); + } + let row = self - .get(params.id) + .get_for_user(user_id, params.id) .await? .ok_or_else(|| DbError::Init(format!("upsert did not produce row for id '{}'", params.id)))?; Ok(row) @@ -234,6 +288,112 @@ impl SqliteAssistantDefinitionRepository { pub fn new(pool: SqlitePool) -> Self { Self { pool } } + + async fn upsert_with_user_id( + &self, + user_id: Option<&str>, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + let now = now_ms(); + + let result = sqlx::query( + "INSERT INTO assistant_definitions ( + id, user_id, assistant_id, source, owner_type, source_ref, + name, name_i18n, description, description_i18n, avatar_type, avatar_value, + agent_id, rule_resource_type, rule_resource_ref, + recommended_prompts, recommended_prompts_i18n, + default_model_mode, default_model_value, + default_permission_mode, default_permission_value, + default_thought_level_mode, default_thought_level_value, + default_skills_mode, default_skill_ids, custom_skill_names, default_disabled_builtin_skill_ids, + default_mcps_mode, default_mcp_ids, + created_at, updated_at, deleted_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL) + ON CONFLICT(id) DO UPDATE SET + assistant_id = excluded.assistant_id, + source = excluded.source, + owner_type = excluded.owner_type, + source_ref = excluded.source_ref, + name = excluded.name, + name_i18n = excluded.name_i18n, + description = excluded.description, + description_i18n = excluded.description_i18n, + avatar_type = excluded.avatar_type, + avatar_value = excluded.avatar_value, + agent_id = excluded.agent_id, + rule_resource_type = excluded.rule_resource_type, + rule_resource_ref = excluded.rule_resource_ref, + recommended_prompts = excluded.recommended_prompts, + recommended_prompts_i18n = excluded.recommended_prompts_i18n, + default_model_mode = excluded.default_model_mode, + default_model_value = excluded.default_model_value, + default_permission_mode = excluded.default_permission_mode, + default_permission_value = excluded.default_permission_value, + default_thought_level_mode = excluded.default_thought_level_mode, + default_thought_level_value = excluded.default_thought_level_value, + default_skills_mode = excluded.default_skills_mode, + default_skill_ids = excluded.default_skill_ids, + custom_skill_names = excluded.custom_skill_names, + default_disabled_builtin_skill_ids = excluded.default_disabled_builtin_skill_ids, + default_mcps_mode = excluded.default_mcps_mode, + default_mcp_ids = excluded.default_mcp_ids, + updated_at = excluded.updated_at, + deleted_at = NULL + WHERE assistant_definitions.user_id IS excluded.user_id", + ) + .bind(params.id) + .bind(user_id) + .bind(params.assistant_id) + .bind(params.source) + .bind(params.owner_type) + .bind(params.source_ref) + .bind(params.name) + .bind(params.name_i18n) + .bind(params.description) + .bind(params.description_i18n) + .bind(params.avatar_type) + .bind(params.avatar_value) + .bind(params.agent_id) + .bind(params.rule_resource_type) + .bind(params.rule_resource_ref) + .bind(params.recommended_prompts) + .bind(params.recommended_prompts_i18n) + .bind(params.default_model_mode) + .bind(params.default_model_value) + .bind(params.default_permission_mode) + .bind(params.default_permission_value) + .bind(params.default_thought_level_mode) + .bind(params.default_thought_level_value) + .bind(params.default_skills_mode) + .bind(params.default_skill_ids) + .bind(params.custom_skill_names) + .bind(params.default_disabled_builtin_skill_ids) + .bind(params.default_mcps_mode) + .bind(params.default_mcp_ids) + .bind(now) + .bind(now) + .execute(&self.pool) + .await?; + + if result.rows_affected() == 0 { + return Err(DbError::Conflict(format!( + "Assistant definition with id '{}' already exists for another scope", + params.id + ))); + } + + let row = match user_id { + Some(user_id) => self.get_by_id_for_user(user_id, params.id).await?, + None => self.get_by_id(params.id).await?, + } + .ok_or_else(|| { + DbError::Init(format!( + "upsert did not produce assistant definition row for id '{}'", + params.id + )) + })?; + Ok(row) + } } /// SQLite-backed implementation of [`IAssistantOverlayRepository`]. @@ -269,34 +429,55 @@ impl SqliteAssistantOverrideRepository { #[async_trait::async_trait] impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { async fn get(&self, assistant_id: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, AssistantOverrideRow>("SELECT * FROM assistant_overrides WHERE assistant_id = ?") - .bind(assistant_id) - .fetch_optional(&self.pool) - .await?; + self.get_for_user(DEFAULT_USER_ID, assistant_id).await + } + + async fn get_for_user(&self, user_id: &str, assistant_id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, AssistantOverrideRow>( + "SELECT * FROM assistant_overrides WHERE user_id = ? AND assistant_id = ?", + ) + .bind(user_id) + .bind(assistant_id) + .fetch_optional(&self.pool) + .await?; Ok(row) } async fn get_all(&self) -> Result, DbError> { - let rows = sqlx::query_as::<_, AssistantOverrideRow>("SELECT * FROM assistant_overrides") + self.get_all_for_user(DEFAULT_USER_ID).await + } + + async fn get_all_for_user(&self, user_id: &str) -> Result, DbError> { + let rows = sqlx::query_as::<_, AssistantOverrideRow>("SELECT * FROM assistant_overrides WHERE user_id = ?") + .bind(user_id) .fetch_all(&self.pool) .await?; Ok(rows) } async fn upsert(&self, params: &UpsertOverrideParams<'_>) -> Result { + self.upsert_for_user(DEFAULT_USER_ID, params).await + } + + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertOverrideParams<'_>, + ) -> Result { let now = now_ms(); let last_used_at: Option = params.last_used_at; sqlx::query( "INSERT INTO assistant_overrides \ - (assistant_id, enabled, sort_order, last_used_at, updated_at) \ - VALUES (?, ?, ?, ?, ?) \ - ON CONFLICT(assistant_id) DO UPDATE SET \ + (user_id, assistant_id, enabled, sort_order, last_used_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?) \ + ON CONFLICT(user_id, assistant_id) DO UPDATE SET \ enabled = excluded.enabled, \ sort_order = excluded.sort_order, \ last_used_at = COALESCE(excluded.last_used_at, assistant_overrides.last_used_at), \ updated_at = excluded.updated_at", ) + .bind(user_id) .bind(params.assistant_id) .bind(params.enabled) .bind(params.sort_order) @@ -305,7 +486,7 @@ impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { .execute(&self.pool) .await?; - let row = self.get(params.assistant_id).await?.ok_or_else(|| { + let row = self.get_for_user(user_id, params.assistant_id).await?.ok_or_else(|| { DbError::Init(format!( "upsert did not produce override row for id '{}'", params.assistant_id @@ -315,7 +496,12 @@ impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { } async fn delete(&self, assistant_id: &str) -> Result { - let result = sqlx::query("DELETE FROM assistant_overrides WHERE assistant_id = ?") + self.delete_for_user(DEFAULT_USER_ID, assistant_id).await + } + + async fn delete_for_user(&self, user_id: &str, assistant_id: &str) -> Result { + let result = sqlx::query("DELETE FROM assistant_overrides WHERE user_id = ? AND assistant_id = ?") + .bind(user_id) .bind(assistant_id) .execute(&self.pool) .await?; @@ -323,16 +509,21 @@ impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { } async fn delete_orphans(&self, valid_ids: &[&str]) -> Result { + self.delete_orphans_for_user(DEFAULT_USER_ID, valid_ids).await + } + + async fn delete_orphans_for_user(&self, user_id: &str, valid_ids: &[&str]) -> Result { if valid_ids.is_empty() { - let result = sqlx::query("DELETE FROM assistant_overrides") + let result = sqlx::query("DELETE FROM assistant_overrides WHERE user_id = ?") + .bind(user_id) .execute(&self.pool) .await?; return Ok(result.rows_affected()); } let placeholders = std::iter::repeat_n("?", valid_ids.len()).collect::>().join(","); - let sql = format!("DELETE FROM assistant_overrides WHERE assistant_id NOT IN ({placeholders})"); - let mut q = sqlx::query(&sql); + let sql = format!("DELETE FROM assistant_overrides WHERE user_id = ? AND assistant_id NOT IN ({placeholders})"); + let mut q = sqlx::query(&sql).bind(user_id); for id in valid_ids { q = q.bind(*id); } @@ -344,26 +535,94 @@ impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { #[async_trait::async_trait] impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { async fn list(&self) -> Result, DbError> { + self.list_for_user(DEFAULT_USER_ID).await + } + + async fn list_for_user(&self, user_id: &str) -> Result, DbError> { let rows = sqlx::query_as::<_, AssistantDefinitionRow>( - "SELECT * FROM assistant_definitions WHERE deleted_at IS NULL ORDER BY updated_at DESC", + "SELECT * + FROM assistant_definitions d + WHERE d.user_id = ? AND d.deleted_at IS NULL + UNION ALL + SELECT * + FROM assistant_definitions d + WHERE d.user_id IS NULL + AND d.deleted_at IS NULL + AND NOT EXISTS ( + SELECT 1 + FROM assistant_definitions u + WHERE u.user_id = ? + AND ( + u.assistant_id = d.assistant_id + OR ( + u.source = d.source + AND u.source_ref IS NOT NULL + AND d.source_ref IS NOT NULL + AND u.source_ref = d.source_ref + ) + ) + ) + ORDER BY updated_at DESC", ) + .bind(user_id) + .bind(user_id) .fetch_all(&self.pool) .await?; Ok(rows) } async fn list_including_deleted(&self) -> Result, DbError> { - let rows = - sqlx::query_as::<_, AssistantDefinitionRow>("SELECT * FROM assistant_definitions ORDER BY updated_at DESC") - .fetch_all(&self.pool) - .await?; + self.list_including_deleted_for_user(DEFAULT_USER_ID).await + } + + async fn list_including_deleted_for_user(&self, user_id: &str) -> Result, DbError> { + let rows = sqlx::query_as::<_, AssistantDefinitionRow>( + "SELECT * + FROM assistant_definitions d + WHERE d.user_id = ? + UNION ALL + SELECT * + FROM assistant_definitions d + WHERE d.user_id IS NULL + AND NOT EXISTS ( + SELECT 1 + FROM assistant_definitions u + WHERE u.user_id = ? + AND ( + u.assistant_id = d.assistant_id + OR ( + u.source = d.source + AND u.source_ref IS NOT NULL + AND d.source_ref IS NOT NULL + AND u.source_ref = d.source_ref + ) + ) + ) + ORDER BY updated_at DESC", + ) + .bind(user_id) + .bind(user_id) + .fetch_all(&self.pool) + .await?; Ok(rows) } async fn get_by_assistant_id(&self, assistant_id: &str) -> Result, DbError> { + self.get_by_assistant_id_for_user(DEFAULT_USER_ID, assistant_id).await + } + + async fn get_by_assistant_id_for_user( + &self, + user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { let row = sqlx::query_as::<_, AssistantDefinitionRow>( - "SELECT * FROM assistant_definitions WHERE assistant_id = ? AND deleted_at IS NULL", + "SELECT * FROM assistant_definitions + WHERE (user_id IS NULL OR user_id = ?) AND assistant_id = ? AND deleted_at IS NULL + ORDER BY user_id IS NULL ASC + LIMIT 1", ) + .bind(user_id) .bind(assistant_id) .fetch_optional(&self.pool) .await?; @@ -374,18 +633,38 @@ impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { &self, assistant_id: &str, ) -> Result, DbError> { - let row = - sqlx::query_as::<_, AssistantDefinitionRow>("SELECT * FROM assistant_definitions WHERE assistant_id = ?") - .bind(assistant_id) - .fetch_optional(&self.pool) - .await?; + self.get_by_assistant_id_including_deleted_for_user(DEFAULT_USER_ID, assistant_id) + .await + } + + async fn get_by_assistant_id_including_deleted_for_user( + &self, + user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + let row = sqlx::query_as::<_, AssistantDefinitionRow>( + "SELECT * FROM assistant_definitions + WHERE (user_id IS NULL OR user_id = ?) AND assistant_id = ? + ORDER BY user_id IS NULL ASC + LIMIT 1", + ) + .bind(user_id) + .bind(assistant_id) + .fetch_optional(&self.pool) + .await?; Ok(row) } async fn get_by_id(&self, id: &str) -> Result, DbError> { + self.get_by_id_for_user(DEFAULT_USER_ID, id).await + } + + async fn get_by_id_for_user(&self, user_id: &str, id: &str) -> Result, DbError> { let row = sqlx::query_as::<_, AssistantDefinitionRow>( - "SELECT * FROM assistant_definitions WHERE id = ? AND deleted_at IS NULL", + "SELECT * FROM assistant_definitions + WHERE (user_id IS NULL OR user_id = ?) AND id = ? AND deleted_at IS NULL", ) + .bind(user_id) .bind(id) .fetch_optional(&self.pool) .await?; @@ -396,10 +675,24 @@ impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { &self, source: &str, source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref_for_user(DEFAULT_USER_ID, source, source_ref) + .await + } + + async fn get_by_source_ref_for_user( + &self, + user_id: &str, + source: &str, + source_ref: &str, ) -> Result, DbError> { let row = sqlx::query_as::<_, AssistantDefinitionRow>( - "SELECT * FROM assistant_definitions WHERE source = ? AND source_ref = ? AND deleted_at IS NULL", + "SELECT * FROM assistant_definitions + WHERE (user_id IS NULL OR user_id = ?) AND source = ? AND source_ref = ? AND deleted_at IS NULL + ORDER BY user_id IS NULL ASC + LIMIT 1", ) + .bind(user_id) .bind(source) .bind(source_ref) .fetch_optional(&self.pool) @@ -411,10 +704,24 @@ impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { &self, source: &str, source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref_including_deleted_for_user(DEFAULT_USER_ID, source, source_ref) + .await + } + + async fn get_by_source_ref_including_deleted_for_user( + &self, + user_id: &str, + source: &str, + source_ref: &str, ) -> Result, DbError> { let row = sqlx::query_as::<_, AssistantDefinitionRow>( - "SELECT * FROM assistant_definitions WHERE source = ? AND source_ref = ?", + "SELECT * FROM assistant_definitions + WHERE (user_id IS NULL OR user_id = ?) AND source = ? AND source_ref = ? + ORDER BY user_id IS NULL ASC + LIMIT 1", ) + .bind(user_id) .bind(source) .bind(source_ref) .fetch_optional(&self.pool) @@ -422,92 +729,59 @@ impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { Ok(row) } - async fn upsert(&self, params: &UpsertAssistantDefinitionParams<'_>) -> Result { - let now = now_ms(); + async fn get_global_by_source_ref_including_deleted( + &self, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + let row = sqlx::query_as::<_, AssistantDefinitionRow>( + "SELECT * FROM assistant_definitions + WHERE user_id IS NULL AND source = ? AND source_ref = ? + LIMIT 1", + ) + .bind(source) + .bind(source_ref) + .fetch_optional(&self.pool) + .await?; + Ok(row) + } - sqlx::query( - "INSERT INTO assistant_definitions ( - id, assistant_id, source, owner_type, source_ref, - name, name_i18n, description, description_i18n, avatar_type, avatar_value, - agent_id, rule_resource_type, rule_resource_ref, - recommended_prompts, recommended_prompts_i18n, - default_model_mode, default_model_value, - default_permission_mode, default_permission_value, - default_thought_level_mode, default_thought_level_value, - default_skills_mode, default_skill_ids, custom_skill_names, default_disabled_builtin_skill_ids, - default_mcps_mode, default_mcp_ids, - created_at, updated_at, deleted_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL) - ON CONFLICT(id) DO UPDATE SET - assistant_id = excluded.assistant_id, - source = excluded.source, - owner_type = excluded.owner_type, - source_ref = excluded.source_ref, - name = excluded.name, - name_i18n = excluded.name_i18n, - description = excluded.description, - description_i18n = excluded.description_i18n, - avatar_type = excluded.avatar_type, - avatar_value = excluded.avatar_value, - agent_id = excluded.agent_id, - rule_resource_type = excluded.rule_resource_type, - rule_resource_ref = excluded.rule_resource_ref, - recommended_prompts = excluded.recommended_prompts, - recommended_prompts_i18n = excluded.recommended_prompts_i18n, - default_model_mode = excluded.default_model_mode, - default_model_value = excluded.default_model_value, - default_permission_mode = excluded.default_permission_mode, - default_permission_value = excluded.default_permission_value, - default_thought_level_mode = excluded.default_thought_level_mode, - default_thought_level_value = excluded.default_thought_level_value, - default_skills_mode = excluded.default_skills_mode, - default_skill_ids = excluded.default_skill_ids, - custom_skill_names = excluded.custom_skill_names, - default_disabled_builtin_skill_ids = excluded.default_disabled_builtin_skill_ids, - default_mcps_mode = excluded.default_mcps_mode, - default_mcp_ids = excluded.default_mcp_ids, - updated_at = excluded.updated_at, - deleted_at = NULL", + async fn get_global_by_assistant_id_including_deleted( + &self, + assistant_id: &str, + ) -> Result, DbError> { + let row = sqlx::query_as::<_, AssistantDefinitionRow>( + "SELECT * FROM assistant_definitions + WHERE user_id IS NULL AND assistant_id = ? + LIMIT 1", ) - .bind(params.id) - .bind(params.assistant_id) - .bind(params.source) - .bind(params.owner_type) - .bind(params.source_ref) - .bind(params.name) - .bind(params.name_i18n) - .bind(params.description) - .bind(params.description_i18n) - .bind(params.avatar_type) - .bind(params.avatar_value) - .bind(params.agent_id) - .bind(params.rule_resource_type) - .bind(params.rule_resource_ref) - .bind(params.recommended_prompts) - .bind(params.recommended_prompts_i18n) - .bind(params.default_model_mode) - .bind(params.default_model_value) - .bind(params.default_permission_mode) - .bind(params.default_permission_value) - .bind(params.default_thought_level_mode) - .bind(params.default_thought_level_value) - .bind(params.default_skills_mode) - .bind(params.default_skill_ids) - .bind(params.custom_skill_names) - .bind(params.default_disabled_builtin_skill_ids) - .bind(params.default_mcps_mode) - .bind(params.default_mcp_ids) - .bind(now) - .bind(now) - .execute(&self.pool) + .bind(assistant_id) + .fetch_optional(&self.pool) .await?; + Ok(row) + } - self.get_by_id(params.id).await?.ok_or_else(|| { - DbError::Init(format!( - "upsert did not produce assistant definition row for id '{}'", - params.id - )) - }) + async fn upsert(&self, params: &UpsertAssistantDefinitionParams<'_>) -> Result { + if params.source == "builtin" && params.owner_type == "system" { + self.upsert_global(params).await + } else { + self.upsert_for_user(DEFAULT_USER_ID, params).await + } + } + + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert_with_user_id(Some(user_id), params).await + } + + async fn upsert_global( + &self, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert_with_user_id(None, params).await } async fn update_avatar_fields_preserving_deleted( @@ -544,14 +818,38 @@ impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { .await?; Ok(result.rows_affected() > 0) } + + async fn soft_delete_for_user(&self, user_id: &str, id: &str, deleted_at: i64) -> Result { + let result = sqlx::query( + "UPDATE assistant_definitions + SET deleted_at = ?, updated_at = ? + WHERE user_id = ? AND id = ? AND deleted_at IS NULL", + ) + .bind(deleted_at) + .bind(now_ms()) + .bind(user_id) + .bind(id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } } #[async_trait::async_trait] impl IAssistantOverlayRepository for SqliteAssistantOverlayRepository { async fn get(&self, assistant_definition_id: &str) -> Result, DbError> { + self.get_for_user(DEFAULT_USER_ID, assistant_definition_id).await + } + + async fn get_for_user( + &self, + user_id: &str, + assistant_definition_id: &str, + ) -> Result, DbError> { let row = sqlx::query_as::<_, AssistantOverlayRow>( - "SELECT * FROM assistant_overlays WHERE assistant_definition_id = ?", + "SELECT * FROM assistant_overlays WHERE user_id = ? AND assistant_definition_id = ?", ) + .bind(user_id) .bind(assistant_definition_id) .fetch_optional(&self.pool) .await?; @@ -559,27 +857,42 @@ impl IAssistantOverlayRepository for SqliteAssistantOverlayRepository { } async fn list(&self) -> Result, DbError> { + self.list_for_user(DEFAULT_USER_ID).await + } + + async fn list_for_user(&self, user_id: &str) -> Result, DbError> { let rows = sqlx::query_as::<_, AssistantOverlayRow>( - "SELECT * FROM assistant_overlays ORDER BY sort_order, updated_at", + "SELECT * FROM assistant_overlays WHERE user_id = ? ORDER BY sort_order, updated_at", ) + .bind(user_id) .fetch_all(&self.pool) .await?; Ok(rows) } async fn upsert(&self, params: &UpsertAssistantOverlayParams<'_>) -> Result { + self.upsert_for_user(DEFAULT_USER_ID, params).await + } + + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAssistantOverlayParams<'_>, + ) -> Result { let now = now_ms(); sqlx::query( "INSERT INTO assistant_overlays ( - assistant_definition_id, enabled, sort_order, agent_id_override, last_used_at, created_at, updated_at - ) VALUES (?, ?, ?, ?, ?, ?, ?) - ON CONFLICT(assistant_definition_id) DO UPDATE SET + user_id, assistant_definition_id, enabled, sort_order, agent_id_override, + last_used_at, created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(user_id, assistant_definition_id) DO UPDATE SET enabled = excluded.enabled, sort_order = excluded.sort_order, agent_id_override = excluded.agent_id_override, last_used_at = excluded.last_used_at, updated_at = excluded.updated_at", ) + .bind(user_id) .bind(params.assistant_definition_id) .bind(params.enabled) .bind(params.sort_order) @@ -590,16 +903,23 @@ impl IAssistantOverlayRepository for SqliteAssistantOverlayRepository { .execute(&self.pool) .await?; - self.get(params.assistant_definition_id).await?.ok_or_else(|| { - DbError::Init(format!( - "upsert did not produce overlay row for assistant_definition_id '{}'", - params.assistant_definition_id - )) - }) + self.get_for_user(user_id, params.assistant_definition_id) + .await? + .ok_or_else(|| { + DbError::Init(format!( + "upsert did not produce overlay row for assistant_definition_id '{}'", + params.assistant_definition_id + )) + }) } async fn delete(&self, assistant_definition_id: &str) -> Result { - let result = sqlx::query("DELETE FROM assistant_overlays WHERE assistant_definition_id = ?") + self.delete_for_user(DEFAULT_USER_ID, assistant_definition_id).await + } + + async fn delete_for_user(&self, user_id: &str, assistant_definition_id: &str) -> Result { + let result = sqlx::query("DELETE FROM assistant_overlays WHERE user_id = ? AND assistant_definition_id = ?") + .bind(user_id) .bind(assistant_definition_id) .execute(&self.pool) .await?; @@ -610,9 +930,18 @@ impl IAssistantOverlayRepository for SqliteAssistantOverlayRepository { #[async_trait::async_trait] impl IAssistantPreferenceRepository for SqliteAssistantPreferenceRepository { async fn get(&self, assistant_definition_id: &str) -> Result, DbError> { + self.get_for_user(DEFAULT_USER_ID, assistant_definition_id).await + } + + async fn get_for_user( + &self, + user_id: &str, + assistant_definition_id: &str, + ) -> Result, DbError> { let row = sqlx::query_as::<_, AssistantPreferenceRow>( - "SELECT * FROM assistant_preferences WHERE assistant_definition_id = ?", + "SELECT * FROM assistant_preferences WHERE user_id = ? AND assistant_definition_id = ?", ) + .bind(user_id) .bind(assistant_definition_id) .fetch_optional(&self.pool) .await?; @@ -620,13 +949,21 @@ impl IAssistantPreferenceRepository for SqliteAssistantPreferenceRepository { } async fn upsert(&self, params: &UpsertAssistantPreferenceParams<'_>) -> Result { + self.upsert_for_user(DEFAULT_USER_ID, params).await + } + + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAssistantPreferenceParams<'_>, + ) -> Result { let now = now_ms(); sqlx::query( "INSERT INTO assistant_preferences ( - assistant_definition_id, last_model_id, last_permission_value, last_thought_level_value, last_skill_ids, + user_id, assistant_definition_id, last_model_id, last_permission_value, last_thought_level_value, last_skill_ids, last_disabled_builtin_skill_ids, last_mcp_ids, created_at, updated_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) - ON CONFLICT(assistant_definition_id) DO UPDATE SET + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(user_id, assistant_definition_id) DO UPDATE SET last_model_id = excluded.last_model_id, last_permission_value = excluded.last_permission_value, last_thought_level_value = excluded.last_thought_level_value, @@ -635,6 +972,7 @@ impl IAssistantPreferenceRepository for SqliteAssistantPreferenceRepository { last_mcp_ids = excluded.last_mcp_ids, updated_at = excluded.updated_at", ) + .bind(user_id) .bind(params.assistant_definition_id) .bind(params.last_model_id) .bind(params.last_permission_value) @@ -647,16 +985,23 @@ impl IAssistantPreferenceRepository for SqliteAssistantPreferenceRepository { .execute(&self.pool) .await?; - self.get(params.assistant_definition_id).await?.ok_or_else(|| { - DbError::Init(format!( - "upsert did not produce preference row for assistant_definition_id '{}'", - params.assistant_definition_id - )) - }) + self.get_for_user(user_id, params.assistant_definition_id) + .await? + .ok_or_else(|| { + DbError::Init(format!( + "upsert did not produce preference row for assistant_definition_id '{}'", + params.assistant_definition_id + )) + }) } async fn delete(&self, assistant_definition_id: &str) -> Result { - let result = sqlx::query("DELETE FROM assistant_preferences WHERE assistant_definition_id = ?") + self.delete_for_user(DEFAULT_USER_ID, assistant_definition_id).await + } + + async fn delete_for_user(&self, user_id: &str, assistant_definition_id: &str) -> Result { + let result = sqlx::query("DELETE FROM assistant_preferences WHERE user_id = ? AND assistant_definition_id = ?") + .bind(user_id) .bind(assistant_definition_id) .execute(&self.pool) .await?; @@ -673,6 +1018,9 @@ mod tests { use super::*; use crate::init_database_memory; + const USER_A: &str = "system_default_user"; + const USER_B: &str = "user_b"; + async fn setup() -> ( SqliteAssistantRepository, SqliteAssistantOverrideRepository, @@ -681,6 +1029,7 @@ mod tests { let db = init_database_memory().await.unwrap(); let a = SqliteAssistantRepository::new(db.pool().clone()); let o = SqliteAssistantOverrideRepository::new(db.pool().clone()); + ensure_user(&db, USER_B).await; (a, o, db) } @@ -694,9 +1043,22 @@ mod tests { let d = SqliteAssistantDefinitionRepository::new(db.pool().clone()); let s = SqliteAssistantOverlayRepository::new(db.pool().clone()); let p = SqliteAssistantPreferenceRepository::new(db.pool().clone()); + ensure_user(&db, USER_B).await; (d, s, p, db) } + async fn ensure_user(db: &crate::Database, user_id: &str) { + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, 1, 1)", + ) + .bind(user_id) + .bind(user_id) + .execute(db.pool()) + .await + .unwrap(); + } + fn params<'a>(id: &'a str, name: &'a str) -> CreateAssistantParams<'a> { CreateAssistantParams { id, @@ -715,12 +1077,21 @@ mod tests { } fn definition_params<'a>(id: &'a str, name: &'a str) -> UpsertAssistantDefinitionParams<'a> { + definition_params_with_id("asstdef_u1", id, Some(id), name) + } + + fn definition_params_with_id<'a>( + definition_id: &'a str, + assistant_id: &'a str, + source_ref: Option<&'a str>, + name: &'a str, + ) -> UpsertAssistantDefinitionParams<'a> { UpsertAssistantDefinitionParams { - id: "asstdef_u1", - assistant_id: id, + id: definition_id, + assistant_id, source: "user", owner_type: "user", - source_ref: Some(id), + source_ref, name, name_i18n: r#"{"zh-CN":"助手"}"#, description: Some("desc"), @@ -747,6 +1118,19 @@ mod tests { } } + fn builtin_definition_params<'a>( + definition_id: &'a str, + assistant_id: &'a str, + source_ref: Option<&'a str>, + name: &'a str, + ) -> UpsertAssistantDefinitionParams<'a> { + UpsertAssistantDefinitionParams { + source: "builtin", + owner_type: "system", + ..definition_params_with_id(definition_id, assistant_id, source_ref, name) + } + } + #[tokio::test] async fn assistant_list_empty() { let (a, _o, _db) = setup().await; @@ -869,6 +1253,57 @@ mod tests { assert_eq!(list.len(), 1); } + #[tokio::test] + async fn assistants_are_scoped_by_user() { + let (a, _o, _db) = setup().await; + a.create_for_user(USER_A, ¶ms("a1", "User A Assistant")) + .await + .unwrap(); + a.create_for_user(USER_B, ¶ms("b1", "User B Assistant")) + .await + .unwrap(); + + assert!(a.get_for_user(USER_A, "a1").await.unwrap().is_some()); + assert!(a.get_for_user(USER_A, "b1").await.unwrap().is_none()); + assert!(a.get_for_user(USER_B, "a1").await.unwrap().is_none()); + assert!(a.get_for_user(USER_B, "b1").await.unwrap().is_some()); + + let user_a_ids: Vec = a + .list_for_user(USER_A) + .await + .unwrap() + .into_iter() + .map(|row| row.id) + .collect(); + let user_b_ids: Vec = a + .list_for_user(USER_B) + .await + .unwrap() + .into_iter() + .map(|row| row.id) + .collect(); + assert_eq!(user_a_ids, vec!["a1"]); + assert_eq!(user_b_ids, vec!["b1"]); + } + + #[tokio::test] + async fn assistant_upsert_rejects_cross_user_id_takeover() { + let (a, _o, _db) = setup().await; + a.upsert_for_user(USER_A, ¶ms("shared", "User A Assistant")) + .await + .unwrap(); + + let err = a + .upsert_for_user(USER_B, ¶ms("shared", "User B Assistant")) + .await + .expect_err("cross-user upsert must not take over an existing assistant id"); + assert!(matches!(err, DbError::Conflict(_))); + + let user_a = a.get_for_user(USER_A, "shared").await.unwrap().unwrap(); + assert_eq!(user_a.name, "User A Assistant"); + assert!(a.get_for_user(USER_B, "shared").await.unwrap().is_none()); + } + #[tokio::test] async fn override_get_missing_returns_none() { let (_a, o, _db) = setup().await; @@ -997,6 +1432,42 @@ mod tests { assert!(o.get_all().await.unwrap().is_empty()); } + #[tokio::test] + async fn overrides_are_scoped_by_user() { + let (_a, o, _db) = setup().await; + o.upsert_for_user( + USER_A, + &UpsertOverrideParams { + assistant_id: "shared", + enabled: true, + sort_order: 1, + last_used_at: Some(100), + }, + ) + .await + .unwrap(); + o.upsert_for_user( + USER_B, + &UpsertOverrideParams { + assistant_id: "shared", + enabled: false, + sort_order: 9, + last_used_at: Some(900), + }, + ) + .await + .unwrap(); + + let user_a = o.get_for_user(USER_A, "shared").await.unwrap().unwrap(); + let user_b = o.get_for_user(USER_B, "shared").await.unwrap().unwrap(); + assert!(user_a.enabled); + assert_eq!(user_a.sort_order, 1); + assert!(!user_b.enabled); + assert_eq!(user_b.sort_order, 9); + assert_eq!(o.get_all_for_user(USER_A).await.unwrap().len(), 1); + assert_eq!(o.get_all_for_user(USER_B).await.unwrap().len(), 1); + } + #[tokio::test] async fn definition_upsert_then_get() { let (d, _s, _p, _db) = setup_v2().await; @@ -1013,6 +1484,159 @@ mod tests { assert_eq!(fetched.avatar_value.as_deref(), Some("🤖")); } + #[tokio::test] + async fn definitions_with_same_refs_are_isolated_by_user() { + let (d, _s, _p, _db) = setup_v2().await; + d.upsert_for_user( + USER_A, + &definition_params_with_id("def_a", "shared-assistant", Some("shared-ref"), "User A Definition"), + ) + .await + .unwrap(); + d.upsert_for_user( + USER_B, + &definition_params_with_id("def_b", "shared-assistant", Some("shared-ref"), "User B Definition"), + ) + .await + .unwrap(); + + let user_a = d + .get_by_source_ref_for_user(USER_A, "user", "shared-ref") + .await + .unwrap() + .unwrap(); + let user_b = d + .get_by_source_ref_for_user(USER_B, "user", "shared-ref") + .await + .unwrap() + .unwrap(); + assert_eq!(user_a.id, "def_a"); + assert_eq!(user_a.name, "User A Definition"); + assert_eq!(user_b.id, "def_b"); + assert_eq!(user_b.name, "User B Definition"); + } + + #[tokio::test] + async fn definition_upsert_rejects_cross_scope_id_takeover() { + let (d, _s, _p, _db) = setup_v2().await; + d.upsert_for_user( + USER_A, + &definition_params_with_id("shared_def", "assistant-a", Some("ref-a"), "User A Definition"), + ) + .await + .unwrap(); + + let err = d + .upsert_for_user( + USER_B, + &definition_params_with_id("shared_def", "assistant-b", Some("ref-b"), "User B Definition"), + ) + .await + .unwrap_err(); + assert!(matches!(err, DbError::Conflict(_))); + + let user_a = d.get_by_id_for_user(USER_A, "shared_def").await.unwrap().unwrap(); + assert_eq!(user_a.assistant_id, "assistant-a"); + assert_eq!(user_a.name, "User A Definition"); + assert!(d.get_by_id_for_user(USER_B, "shared_def").await.unwrap().is_none()); + + d.upsert_global(&builtin_definition_params( + "builtin_shared_def", + "builtin-a", + Some("builtin-a"), + "Builtin Definition", + )) + .await + .unwrap(); + d.upsert_global(&builtin_definition_params( + "builtin_shared_def", + "builtin-a", + Some("builtin-a"), + "Builtin Definition Updated", + )) + .await + .unwrap(); + + let err = d + .upsert_for_user( + USER_A, + &definition_params_with_id("builtin_shared_def", "user-a", Some("user-a"), "User Definition"), + ) + .await + .unwrap_err(); + assert!(matches!(err, DbError::Conflict(_))); + + let global = d.get_by_id("builtin_shared_def").await.unwrap().unwrap(); + assert_eq!(global.assistant_id, "builtin-a"); + assert_eq!(global.name, "Builtin Definition Updated"); + } + + #[tokio::test] + async fn global_builtin_definition_is_visible_to_all_users() { + let (d, _s, _p, _db) = setup_v2().await; + d.upsert_global(&builtin_definition_params( + "builtin_def", + "builtin-assistant", + Some("builtin-ref"), + "Builtin Definition", + )) + .await + .unwrap(); + + let user_a = d + .get_by_assistant_id_for_user(USER_A, "builtin-assistant") + .await + .unwrap(); + let user_b = d + .get_by_assistant_id_for_user(USER_B, "builtin-assistant") + .await + .unwrap(); + assert_eq!(user_a.unwrap().id, "builtin_def"); + assert_eq!(user_b.unwrap().id, "builtin_def"); + } + + #[tokio::test] + async fn user_definition_overrides_global_definition() { + let (d, _s, _p, _db) = setup_v2().await; + d.upsert_global(&builtin_definition_params( + "global_def", + "shared-assistant", + Some("shared-ref"), + "Global Definition", + )) + .await + .unwrap(); + d.upsert_for_user( + USER_B, + &definition_params_with_id("user_def", "shared-assistant", Some("shared-ref"), "User Definition"), + ) + .await + .unwrap(); + + let user_a = d + .get_by_assistant_id_for_user(USER_A, "shared-assistant") + .await + .unwrap() + .unwrap(); + let user_b_by_assistant = d + .get_by_assistant_id_for_user(USER_B, "shared-assistant") + .await + .unwrap() + .unwrap(); + let user_b_by_ref = d + .get_by_source_ref_for_user(USER_B, "user", "shared-ref") + .await + .unwrap() + .unwrap(); + let user_b_list = d.list_for_user(USER_B).await.unwrap(); + + assert_eq!(user_a.id, "global_def"); + assert_eq!(user_b_by_assistant.id, "user_def"); + assert_eq!(user_b_by_ref.id, "user_def"); + assert_eq!(user_b_list.len(), 1); + assert_eq!(user_b_list[0].id, "user_def"); + } + #[tokio::test] async fn state_upsert_then_list() { let (d, s, _p, _db) = setup_v2().await; @@ -1035,6 +1659,45 @@ mod tests { assert_eq!(list[0].agent_id_override.as_deref(), Some("claude")); } + #[tokio::test] + async fn overlays_are_scoped_by_user() { + let (d, s, _p, _db) = setup_v2().await; + let definition = d.upsert(&definition_params("u1", "User One")).await.unwrap(); + s.upsert_for_user( + USER_A, + &UpsertAssistantOverlayParams { + assistant_definition_id: &definition.id, + enabled: true, + sort_order: 1, + agent_id_override: Some("agent-a"), + last_used_at: Some(10), + }, + ) + .await + .unwrap(); + s.upsert_for_user( + USER_B, + &UpsertAssistantOverlayParams { + assistant_definition_id: &definition.id, + enabled: false, + sort_order: 7, + agent_id_override: Some("agent-b"), + last_used_at: Some(70), + }, + ) + .await + .unwrap(); + + let user_a = s.get_for_user(USER_A, &definition.id).await.unwrap().unwrap(); + let user_b = s.get_for_user(USER_B, &definition.id).await.unwrap().unwrap(); + assert!(user_a.enabled); + assert_eq!(user_a.agent_id_override.as_deref(), Some("agent-a")); + assert!(!user_b.enabled); + assert_eq!(user_b.agent_id_override.as_deref(), Some("agent-b")); + assert_eq!(s.list_for_user(USER_A).await.unwrap().len(), 1); + assert_eq!(s.list_for_user(USER_B).await.unwrap().len(), 1); + } + #[tokio::test] async fn preference_upsert_then_get() { let (d, _s, p, _db) = setup_v2().await; @@ -1058,4 +1721,45 @@ mod tests { assert_eq!(fetched.last_skill_ids, r#"["pdf"]"#); assert_eq!(fetched.last_thought_level_value.as_deref(), Some("high")); } + + #[tokio::test] + async fn preferences_are_scoped_by_user() { + let (d, _s, p, _db) = setup_v2().await; + let definition = d.upsert(&definition_params("u1", "User One")).await.unwrap(); + p.upsert_for_user( + USER_A, + &UpsertAssistantPreferenceParams { + assistant_definition_id: &definition.id, + last_model_id: Some("model-a"), + last_permission_value: Some("read-only"), + last_thought_level_value: Some("low"), + last_skill_ids: r#"["a"]"#, + last_disabled_builtin_skill_ids: "[]", + last_mcp_ids: "[]", + }, + ) + .await + .unwrap(); + p.upsert_for_user( + USER_B, + &UpsertAssistantPreferenceParams { + assistant_definition_id: &definition.id, + last_model_id: Some("model-b"), + last_permission_value: Some("workspace-write"), + last_thought_level_value: Some("high"), + last_skill_ids: r#"["b"]"#, + last_disabled_builtin_skill_ids: "[]", + last_mcp_ids: "[]", + }, + ) + .await + .unwrap(); + + let user_a = p.get_for_user(USER_A, &definition.id).await.unwrap().unwrap(); + let user_b = p.get_for_user(USER_B, &definition.id).await.unwrap().unwrap(); + assert_eq!(user_a.last_model_id.as_deref(), Some("model-a")); + assert_eq!(user_a.last_skill_ids, r#"["a"]"#); + assert_eq!(user_b.last_model_id.as_deref(), Some("model-b")); + assert_eq!(user_b.last_skill_ids, r#"["b"]"#); + } } diff --git a/crates/aionui-db/src/repository/sqlite_channel.rs b/crates/aionui-db/src/repository/sqlite_channel.rs index 1052df95e..1e33b4911 100644 --- a/crates/aionui-db/src/repository/sqlite_channel.rs +++ b/crates/aionui-db/src/repository/sqlite_channel.rs @@ -20,27 +20,32 @@ impl SqliteChannelRepository { impl IChannelRepository for SqliteChannelRepository { // ── Plugin CRUD ────────────────────────────────────────────────── - async fn get_all_plugins(&self) -> Result, DbError> { - let rows = sqlx::query_as::<_, ChannelPluginRow>("SELECT * FROM assistant_plugins ORDER BY created_at ASC") - .fetch_all(&self.pool) - .await?; + async fn get_all_plugins(&self, owner_user_id: &str) -> Result, DbError> { + let rows = sqlx::query_as::<_, ChannelPluginRow>( + "SELECT * FROM assistant_plugins WHERE owner_user_id = ? ORDER BY created_at ASC", + ) + .bind(owner_user_id) + .fetch_all(&self.pool) + .await?; Ok(rows) } - async fn get_plugin(&self, id: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, ChannelPluginRow>("SELECT * FROM assistant_plugins WHERE id = ?") - .bind(id) - .fetch_optional(&self.pool) - .await?; + async fn get_plugin(&self, owner_user_id: &str, id: &str) -> Result, DbError> { + let row = + sqlx::query_as::<_, ChannelPluginRow>("SELECT * FROM assistant_plugins WHERE owner_user_id = ? AND id = ?") + .bind(owner_user_id) + .bind(id) + .fetch_optional(&self.pool) + .await?; Ok(row) } - async fn upsert_plugin(&self, row: &ChannelPluginRow) -> Result<(), DbError> { + async fn upsert_plugin(&self, owner_user_id: &str, row: &ChannelPluginRow) -> Result<(), DbError> { sqlx::query( "INSERT INTO assistant_plugins \ - (id, type, name, enabled, config, status, last_connected, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) \ - ON CONFLICT(id) DO UPDATE SET \ + (id, owner_user_id, type, name, enabled, config, status, last_connected, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \ + ON CONFLICT(owner_user_id, id) DO UPDATE SET \ type = excluded.type, \ name = excluded.name, \ enabled = excluded.enabled, \ @@ -50,6 +55,7 @@ impl IChannelRepository for SqliteChannelRepository { updated_at = excluded.updated_at", ) .bind(&row.id) + .bind(owner_user_id) .bind(&row.r#type) .bind(&row.name) .bind(row.enabled) @@ -63,7 +69,12 @@ impl IChannelRepository for SqliteChannelRepository { Ok(()) } - async fn update_plugin_status(&self, id: &str, params: &UpdatePluginStatusParams) -> Result<(), DbError> { + async fn update_plugin_status( + &self, + owner_user_id: &str, + id: &str, + params: &UpdatePluginStatusParams, + ) -> Result<(), DbError> { let mut set_clauses = Vec::new(); if params.status.is_some() { set_clauses.push("status = ?"); @@ -80,7 +91,10 @@ impl IChannelRepository for SqliteChannelRepository { } set_clauses.push("updated_at = ?"); - let sql = format!("UPDATE assistant_plugins SET {} WHERE id = ?", set_clauses.join(", ")); + let sql = format!( + "UPDATE assistant_plugins SET {} WHERE owner_user_id = ? AND id = ?", + set_clauses.join(", ") + ); let now = aionui_common::now_ms(); let mut query = sqlx::query(&sql); @@ -95,6 +109,7 @@ impl IChannelRepository for SqliteChannelRepository { query = query.bind(enabled); } query = query.bind(now); + query = query.bind(owner_user_id); query = query.bind(id); let result = query.execute(&self.pool).await?; @@ -104,8 +119,9 @@ impl IChannelRepository for SqliteChannelRepository { Ok(()) } - async fn delete_plugin(&self, id: &str) -> Result<(), DbError> { - let result = sqlx::query("DELETE FROM assistant_plugins WHERE id = ?") + async fn delete_plugin(&self, owner_user_id: &str, id: &str) -> Result<(), DbError> { + let result = sqlx::query("DELETE FROM assistant_plugins WHERE owner_user_id = ? AND id = ?") + .bind(owner_user_id) .bind(id) .execute(&self.pool) .await?; @@ -117,22 +133,27 @@ impl IChannelRepository for SqliteChannelRepository { // ── User CRUD ──────────────────────────────────────────────────── - async fn get_all_users(&self) -> Result, DbError> { - let rows = sqlx::query_as::<_, AssistantUserRow>("SELECT * FROM assistant_users ORDER BY authorized_at DESC") - .fetch_all(&self.pool) - .await?; + async fn get_all_users(&self, owner_user_id: &str) -> Result, DbError> { + let rows = sqlx::query_as::<_, AssistantUserRow>( + "SELECT * FROM assistant_users WHERE owner_user_id = ? ORDER BY authorized_at DESC", + ) + .bind(owner_user_id) + .fetch_all(&self.pool) + .await?; Ok(rows) } async fn get_user_by_platform( &self, + owner_user_id: &str, platform_user_id: &str, platform_type: &str, ) -> Result, DbError> { let row = sqlx::query_as::<_, AssistantUserRow>( "SELECT * FROM assistant_users \ - WHERE platform_user_id = ? AND platform_type = ?", + WHERE owner_user_id = ? AND platform_user_id = ? AND platform_type = ?", ) + .bind(owner_user_id) .bind(platform_user_id) .bind(platform_type) .fetch_optional(&self.pool) @@ -140,14 +161,15 @@ impl IChannelRepository for SqliteChannelRepository { Ok(row) } - async fn create_user(&self, row: &AssistantUserRow) -> Result<(), DbError> { + async fn create_user(&self, owner_user_id: &str, row: &AssistantUserRow) -> Result<(), DbError> { sqlx::query( "INSERT INTO assistant_users \ - (id, platform_user_id, platform_type, display_name, \ + (id, owner_user_id, platform_user_id, platform_type, display_name, \ authorized_at, last_active, session_id) \ - VALUES (?, ?, ?, ?, ?, ?, ?)", + VALUES (?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(&row.id) + .bind(owner_user_id) .bind(&row.platform_user_id) .bind(&row.platform_type) .bind(&row.display_name) @@ -169,9 +191,15 @@ impl IChannelRepository for SqliteChannelRepository { Ok(()) } - async fn update_user_last_active(&self, id: &str, last_active: aionui_common::TimestampMs) -> Result<(), DbError> { - let result = sqlx::query("UPDATE assistant_users SET last_active = ? WHERE id = ?") + async fn update_user_last_active( + &self, + owner_user_id: &str, + id: &str, + last_active: aionui_common::TimestampMs, + ) -> Result<(), DbError> { + let result = sqlx::query("UPDATE assistant_users SET last_active = ? WHERE owner_user_id = ? AND id = ?") .bind(last_active) + .bind(owner_user_id) .bind(id) .execute(&self.pool) .await?; @@ -181,8 +209,9 @@ impl IChannelRepository for SqliteChannelRepository { Ok(()) } - async fn delete_user(&self, id: &str) -> Result<(), DbError> { - let result = sqlx::query("DELETE FROM assistant_users WHERE id = ?") + async fn delete_user(&self, owner_user_id: &str, id: &str) -> Result<(), DbError> { + let result = sqlx::query("DELETE FROM assistant_users WHERE owner_user_id = ? AND id = ?") + .bind(owner_user_id) .bind(id) .execute(&self.pool) .await?; @@ -194,34 +223,47 @@ impl IChannelRepository for SqliteChannelRepository { // ── Session CRUD ───────────────────────────────────────────────── - async fn get_all_sessions(&self) -> Result, DbError> { - let rows = - sqlx::query_as::<_, AssistantSessionRow>("SELECT * FROM assistant_sessions ORDER BY last_activity DESC") - .fetch_all(&self.pool) - .await?; + async fn get_all_sessions(&self, owner_user_id: &str) -> Result, DbError> { + let rows = sqlx::query_as::<_, AssistantSessionRow>( + "SELECT s.* FROM assistant_sessions s \ + JOIN assistant_users u ON u.id = s.user_id \ + WHERE u.owner_user_id = ? \ + ORDER BY s.last_activity DESC", + ) + .bind(owner_user_id) + .fetch_all(&self.pool) + .await?; Ok(rows) } - async fn get_session(&self, id: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, AssistantSessionRow>("SELECT * FROM assistant_sessions WHERE id = ?") - .bind(id) - .fetch_optional(&self.pool) - .await?; + async fn get_session(&self, owner_user_id: &str, id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, AssistantSessionRow>( + "SELECT s.* FROM assistant_sessions s \ + JOIN assistant_users u ON u.id = s.user_id \ + WHERE u.owner_user_id = ? AND s.id = ?", + ) + .bind(owner_user_id) + .bind(id) + .fetch_optional(&self.pool) + .await?; Ok(row) } async fn get_or_create_session( &self, - user_id: &str, + owner_user_id: &str, + channel_user_id: &str, chat_id: &str, new_row: &AssistantSessionRow, ) -> Result { // Try to find an existing session first. let existing = sqlx::query_as::<_, AssistantSessionRow>( - "SELECT * FROM assistant_sessions \ - WHERE user_id = ? AND chat_id = ?", + "SELECT s.* FROM assistant_sessions s \ + JOIN assistant_users u ON u.id = s.user_id \ + WHERE u.owner_user_id = ? AND s.user_id = ? AND s.chat_id = ?", ) - .bind(user_id) + .bind(owner_user_id) + .bind(channel_user_id) .bind(chat_id) .fetch_optional(&self.pool) .await?; @@ -246,66 +288,128 @@ impl IChannelRepository for SqliteChannelRepository { "INSERT INTO assistant_sessions \ (id, user_id, agent_type, conversation_id, workspace, \ chat_id, created_at, last_activity) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + SELECT ?, ?, ?, ?, ?, ?, ?, ? \ + WHERE EXISTS ( + SELECT 1 FROM assistant_users + WHERE owner_user_id = ? AND id = ? + )", ) .bind(&new_row.id) - .bind(&new_row.user_id) + .bind(channel_user_id) .bind(&new_row.agent_type) .bind(&new_row.conversation_id) .bind(&new_row.workspace) .bind(&new_row.chat_id) .bind(new_row.created_at) .bind(new_row.last_activity) + .bind(owner_user_id) + .bind(channel_user_id) .execute(&self.pool) .await?; - Ok(new_row.clone()) + self.get_session(owner_user_id, &new_row.id) + .await? + .ok_or_else(|| DbError::NotFound(format!("Channel user '{channel_user_id}' not found"))) } async fn update_session_activity( &self, + owner_user_id: &str, id: &str, last_activity: aionui_common::TimestampMs, ) -> Result<(), DbError> { - let result = sqlx::query("UPDATE assistant_sessions SET last_activity = ? WHERE id = ?") - .bind(last_activity) - .bind(id) - .execute(&self.pool) - .await?; + let result = sqlx::query( + "UPDATE assistant_sessions SET last_activity = ? \ + WHERE id = ? AND EXISTS ( + SELECT 1 FROM assistant_users u + WHERE u.id = assistant_sessions.user_id AND u.owner_user_id = ? + )", + ) + .bind(last_activity) + .bind(id) + .bind(owner_user_id) + .execute(&self.pool) + .await?; if result.rows_affected() == 0 { return Err(DbError::NotFound(format!("Session '{id}' not found"))); } Ok(()) } - async fn update_session_conversation(&self, id: &str, conversation_id: &str) -> Result<(), DbError> { + async fn update_session_conversation( + &self, + owner_user_id: &str, + id: &str, + conversation_id: &str, + ) -> Result<(), DbError> { let now = aionui_common::now_ms(); let result = sqlx::query( "UPDATE assistant_sessions \ SET conversation_id = ?, last_activity = ? \ - WHERE id = ?", + WHERE id = ? \ + AND EXISTS ( + SELECT 1 FROM assistant_users u + WHERE u.id = assistant_sessions.user_id AND u.owner_user_id = ? + ) + AND EXISTS ( + SELECT 1 FROM conversations c + WHERE c.id = ? AND c.user_id = ? + )", ) .bind(conversation_id) .bind(now) .bind(id) + .bind(owner_user_id) + .bind(conversation_id) + .bind(owner_user_id) .execute(&self.pool) .await?; if result.rows_affected() == 0 { + let session_exists = sqlx::query_scalar::<_, i64>( + "SELECT COUNT(*) FROM assistant_sessions s \ + JOIN assistant_users u ON u.id = s.user_id \ + WHERE s.id = ? AND u.owner_user_id = ?", + ) + .bind(id) + .bind(owner_user_id) + .fetch_one(&self.pool) + .await?; + + if session_exists > 0 { + let conversation_owner = + sqlx::query_scalar::<_, String>("SELECT user_id FROM conversations WHERE id = ?") + .bind(conversation_id) + .fetch_optional(&self.pool) + .await?; + + if conversation_owner + .as_deref() + .is_some_and(|user_id| user_id != owner_user_id) + { + return Err(DbError::Conflict( + "CROSS_ACCOUNT_REFERENCE: channel session conversation belongs to another user".into(), + )); + } + } return Err(DbError::NotFound(format!("Session '{id}' not found"))); } Ok(()) } - async fn update_session_agent_type(&self, id: &str, agent_type: &str) -> Result<(), DbError> { + async fn update_session_agent_type(&self, owner_user_id: &str, id: &str, agent_type: &str) -> Result<(), DbError> { let now = aionui_common::now_ms(); let result = sqlx::query( "UPDATE assistant_sessions \ SET agent_type = ?, last_activity = ? \ - WHERE id = ?", + WHERE id = ? AND EXISTS ( + SELECT 1 FROM assistant_users u + WHERE u.id = assistant_sessions.user_id AND u.owner_user_id = ? + )", ) .bind(agent_type) .bind(now) .bind(id) + .bind(owner_user_id) .execute(&self.pool) .await?; if result.rows_affected() == 0 { @@ -314,21 +418,37 @@ impl IChannelRepository for SqliteChannelRepository { Ok(()) } - async fn delete_sessions_by_user(&self, user_id: &str) -> Result<(), DbError> { - sqlx::query("DELETE FROM assistant_sessions WHERE user_id = ?") - .bind(user_id) - .execute(&self.pool) - .await?; + async fn delete_sessions_by_user(&self, owner_user_id: &str, channel_user_id: &str) -> Result<(), DbError> { + sqlx::query( + "DELETE FROM assistant_sessions \ + WHERE user_id = ? AND EXISTS ( + SELECT 1 FROM assistant_users u + WHERE u.id = assistant_sessions.user_id AND u.owner_user_id = ? + )", + ) + .bind(channel_user_id) + .bind(owner_user_id) + .execute(&self.pool) + .await?; Ok(()) } - async fn delete_session_by_user_chat(&self, user_id: &str, chat_id: &str) -> Result<(), DbError> { + async fn delete_session_by_user_chat( + &self, + owner_user_id: &str, + channel_user_id: &str, + chat_id: &str, + ) -> Result<(), DbError> { sqlx::query( "DELETE FROM assistant_sessions \ - WHERE user_id = ? AND chat_id = ?", + WHERE user_id = ? AND chat_id = ? AND EXISTS ( + SELECT 1 FROM assistant_users u + WHERE u.id = assistant_sessions.user_id AND u.owner_user_id = ? + )", ) - .bind(user_id) + .bind(channel_user_id) .bind(chat_id) + .bind(owner_user_id) .execute(&self.pool) .await?; Ok(()) @@ -336,14 +456,15 @@ impl IChannelRepository for SqliteChannelRepository { // ── Pairing Codes ──────────────────────────────────────────────── - async fn create_pairing(&self, row: &PairingCodeRow) -> Result<(), DbError> { + async fn create_pairing(&self, owner_user_id: &str, row: &PairingCodeRow) -> Result<(), DbError> { sqlx::query( "INSERT INTO assistant_pairing_codes \ - (code, platform_user_id, platform_type, display_name, \ + (code, owner_user_id, platform_user_id, platform_type, display_name, \ requested_at, expires_at, status) \ - VALUES (?, ?, ?, ?, ?, ?, ?)", + VALUES (?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(&row.code) + .bind(owner_user_id) .bind(&row.platform_user_id) .bind(&row.platform_type) .bind(&row.display_name) @@ -362,28 +483,33 @@ impl IChannelRepository for SqliteChannelRepository { Ok(()) } - async fn get_pending_pairings(&self) -> Result, DbError> { + async fn get_pending_pairings(&self, owner_user_id: &str) -> Result, DbError> { let rows = sqlx::query_as::<_, PairingCodeRow>( "SELECT * FROM assistant_pairing_codes \ - WHERE status = 'pending' \ + WHERE owner_user_id = ? AND status = 'pending' \ ORDER BY requested_at DESC", ) + .bind(owner_user_id) .fetch_all(&self.pool) .await?; Ok(rows) } - async fn get_pairing_by_code(&self, code: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, PairingCodeRow>("SELECT * FROM assistant_pairing_codes WHERE code = ?") - .bind(code) - .fetch_optional(&self.pool) - .await?; + async fn get_pairing_by_code(&self, owner_user_id: &str, code: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, PairingCodeRow>( + "SELECT * FROM assistant_pairing_codes WHERE owner_user_id = ? AND code = ?", + ) + .bind(owner_user_id) + .bind(code) + .fetch_optional(&self.pool) + .await?; Ok(row) } - async fn update_pairing_status(&self, code: &str, status: &str) -> Result<(), DbError> { - let result = sqlx::query("UPDATE assistant_pairing_codes SET status = ? WHERE code = ?") + async fn update_pairing_status(&self, owner_user_id: &str, code: &str, status: &str) -> Result<(), DbError> { + let result = sqlx::query("UPDATE assistant_pairing_codes SET status = ? WHERE owner_user_id = ? AND code = ?") .bind(status) + .bind(owner_user_id) .bind(code) .execute(&self.pool) .await?; @@ -393,12 +519,17 @@ impl IChannelRepository for SqliteChannelRepository { Ok(()) } - async fn cleanup_expired_pairings(&self, now: aionui_common::TimestampMs) -> Result { + async fn cleanup_expired_pairings( + &self, + owner_user_id: &str, + now: aionui_common::TimestampMs, + ) -> Result { let result = sqlx::query( "UPDATE assistant_pairing_codes \ SET status = 'expired' \ - WHERE status = 'pending' AND expires_at <= ?", + WHERE owner_user_id = ? AND status = 'pending' AND expires_at <= ?", ) + .bind(owner_user_id) .bind(now) .execute(&self.pool) .await?; @@ -419,16 +550,36 @@ mod tests { use super::*; use crate::init_database_memory; + const OWNER_A: &str = "system_default_user"; + const OWNER_B: &str = "other_user"; + async fn setup() -> (SqliteChannelRepository, crate::Database) { let db = init_database_memory().await.unwrap(); let repo = SqliteChannelRepository::new(db.pool().clone()); (repo, db) } + async fn create_owner(pool: &SqlitePool, user_id: &str) { + let now = aionui_common::now_ms(); + sqlx::query( + "INSERT OR IGNORE INTO users \ + (id, username, password_hash, created_at, updated_at) \ + VALUES (?, ?, 'hash', ?, ?)", + ) + .bind(user_id) + .bind(user_id) + .bind(now) + .bind(now) + .execute(pool) + .await + .unwrap(); + } + fn sample_plugin() -> ChannelPluginRow { let now = aionui_common::now_ms(); ChannelPluginRow { id: "tg-1".into(), + owner_user_id: OWNER_A.into(), r#type: "telegram".into(), name: "My Telegram Bot".into(), enabled: false, @@ -444,6 +595,7 @@ mod tests { let now = aionui_common::now_ms(); AssistantUserRow { id: "usr-1".into(), + owner_user_id: OWNER_A.into(), platform_user_id: "tg_12345".into(), platform_type: "telegram".into(), display_name: Some("Alice".into()), @@ -471,6 +623,7 @@ mod tests { let now = aionui_common::now_ms(); PairingCodeRow { code: "123456".into(), + owner_user_id: OWNER_A.into(), platform_user_id: "tg_99".into(), platform_type: "telegram".into(), display_name: Some("Bob".into()), @@ -485,7 +638,7 @@ mod tests { #[tokio::test] async fn get_all_plugins_empty() { let (repo, _db) = setup().await; - let plugins = repo.get_all_plugins().await.unwrap(); + let plugins = repo.get_all_plugins(OWNER_A).await.unwrap(); assert!(plugins.is_empty()); } @@ -493,9 +646,9 @@ mod tests { async fn upsert_and_get_plugin() { let (repo, _db) = setup().await; let plugin = sample_plugin(); - repo.upsert_plugin(&plugin).await.unwrap(); + repo.upsert_plugin(OWNER_A, &plugin).await.unwrap(); - let found = repo.get_plugin("tg-1").await.unwrap().unwrap(); + let found = repo.get_plugin(OWNER_A, "tg-1").await.unwrap().unwrap(); assert_eq!(found.id, "tg-1"); assert_eq!(found.r#type, "telegram"); assert_eq!(found.name, "My Telegram Bot"); @@ -506,7 +659,7 @@ mod tests { async fn upsert_plugin_updates_existing() { let (repo, _db) = setup().await; let plugin = sample_plugin(); - repo.upsert_plugin(&plugin).await.unwrap(); + repo.upsert_plugin(OWNER_A, &plugin).await.unwrap(); let updated = ChannelPluginRow { name: "Updated Bot".into(), @@ -514,9 +667,9 @@ mod tests { updated_at: aionui_common::now_ms(), ..plugin }; - repo.upsert_plugin(&updated).await.unwrap(); + repo.upsert_plugin(OWNER_A, &updated).await.unwrap(); - let found = repo.get_plugin("tg-1").await.unwrap().unwrap(); + let found = repo.get_plugin(OWNER_A, "tg-1").await.unwrap().unwrap(); assert_eq!(found.name, "Updated Bot"); assert!(found.enabled); } @@ -524,11 +677,12 @@ mod tests { #[tokio::test] async fn get_all_plugins_returns_multiple() { let (repo, _db) = setup().await; - repo.upsert_plugin(&sample_plugin()).await.unwrap(); + repo.upsert_plugin(OWNER_A, &sample_plugin()).await.unwrap(); let now = aionui_common::now_ms(); let lark = ChannelPluginRow { id: "lark-1".into(), + owner_user_id: OWNER_A.into(), r#type: "lark".into(), name: "Lark Bot".into(), enabled: true, @@ -538,19 +692,72 @@ mod tests { created_at: now, updated_at: now, }; - repo.upsert_plugin(&lark).await.unwrap(); + repo.upsert_plugin(OWNER_A, &lark).await.unwrap(); - let all = repo.get_all_plugins().await.unwrap(); + let all = repo.get_all_plugins(OWNER_A).await.unwrap(); assert_eq!(all.len(), 2); } + #[tokio::test] + async fn plugins_are_filtered_by_owner() { + let (repo, db) = setup().await; + create_owner(db.pool(), OWNER_B).await; + + repo.upsert_plugin(OWNER_A, &sample_plugin()).await.unwrap(); + + assert!(repo.get_plugin(OWNER_B, "tg-1").await.unwrap().is_none()); + assert!(repo.get_all_plugins(OWNER_B).await.unwrap().is_empty()); + } + + #[tokio::test] + async fn same_plugin_id_can_exist_for_different_owners() { + let (repo, db) = setup().await; + create_owner(db.pool(), OWNER_B).await; + + let owner_a_plugin = sample_plugin(); + let owner_b_plugin = ChannelPluginRow { + owner_user_id: OWNER_B.into(), + name: "Owner B Telegram Bot".into(), + enabled: true, + ..sample_plugin() + }; + + repo.upsert_plugin(OWNER_A, &owner_a_plugin).await.unwrap(); + repo.upsert_plugin(OWNER_B, &owner_b_plugin).await.unwrap(); + + let owner_a_found = repo.get_plugin(OWNER_A, "tg-1").await.unwrap().unwrap(); + let owner_b_found = repo.get_plugin(OWNER_B, "tg-1").await.unwrap().unwrap(); + assert_eq!(owner_a_found.id, "tg-1"); + assert_eq!(owner_b_found.id, "tg-1"); + assert_eq!(owner_a_found.name, "My Telegram Bot"); + assert_eq!(owner_b_found.name, "Owner B Telegram Bot"); + + repo.update_plugin_status( + OWNER_B, + "tg-1", + &UpdatePluginStatusParams { + status: Some("running".into()), + last_connected: None, + enabled: None, + }, + ) + .await + .unwrap(); + + let owner_a_after = repo.get_plugin(OWNER_A, "tg-1").await.unwrap().unwrap(); + let owner_b_after = repo.get_plugin(OWNER_B, "tg-1").await.unwrap().unwrap(); + assert_eq!(owner_a_after.status, None); + assert_eq!(owner_b_after.status, Some("running".into())); + } + #[tokio::test] async fn update_plugin_status_sets_fields() { let (repo, _db) = setup().await; - repo.upsert_plugin(&sample_plugin()).await.unwrap(); + repo.upsert_plugin(OWNER_A, &sample_plugin()).await.unwrap(); let now = aionui_common::now_ms(); repo.update_plugin_status( + OWNER_A, "tg-1", &UpdatePluginStatusParams { status: Some("running".into()), @@ -561,7 +768,7 @@ mod tests { .await .unwrap(); - let found = repo.get_plugin("tg-1").await.unwrap().unwrap(); + let found = repo.get_plugin(OWNER_A, "tg-1").await.unwrap().unwrap(); assert_eq!(found.status.as_deref(), Some("running")); assert_eq!(found.last_connected, Some(now)); assert!(found.enabled); @@ -572,6 +779,7 @@ mod tests { let (repo, _db) = setup().await; let err = repo .update_plugin_status( + OWNER_A, "nope", &UpdatePluginStatusParams { status: Some("error".into()), @@ -586,9 +794,9 @@ mod tests { #[tokio::test] async fn update_plugin_status_empty_params_is_noop() { let (repo, _db) = setup().await; - repo.upsert_plugin(&sample_plugin()).await.unwrap(); + repo.upsert_plugin(OWNER_A, &sample_plugin()).await.unwrap(); // No fields to update → no-op, no error. - repo.update_plugin_status("tg-1", &UpdatePluginStatusParams::default()) + repo.update_plugin_status(OWNER_A, "tg-1", &UpdatePluginStatusParams::default()) .await .unwrap(); } @@ -596,15 +804,15 @@ mod tests { #[tokio::test] async fn delete_plugin_removes_row() { let (repo, _db) = setup().await; - repo.upsert_plugin(&sample_plugin()).await.unwrap(); - repo.delete_plugin("tg-1").await.unwrap(); - assert!(repo.get_plugin("tg-1").await.unwrap().is_none()); + repo.upsert_plugin(OWNER_A, &sample_plugin()).await.unwrap(); + repo.delete_plugin(OWNER_A, "tg-1").await.unwrap(); + assert!(repo.get_plugin(OWNER_A, "tg-1").await.unwrap().is_none()); } #[tokio::test] async fn delete_plugin_not_found() { let (repo, _db) = setup().await; - let err = repo.delete_plugin("nope").await.unwrap_err(); + let err = repo.delete_plugin(OWNER_A, "nope").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -613,7 +821,7 @@ mod tests { #[tokio::test] async fn get_all_users_empty() { let (repo, _db) = setup().await; - let users = repo.get_all_users().await.unwrap(); + let users = repo.get_all_users(OWNER_A).await.unwrap(); assert!(users.is_empty()); } @@ -621,10 +829,10 @@ mod tests { async fn create_and_get_user_by_platform() { let (repo, _db) = setup().await; let user = sample_user(); - repo.create_user(&user).await.unwrap(); + repo.create_user(OWNER_A, &user).await.unwrap(); let found = repo - .get_user_by_platform("tg_12345", "telegram") + .get_user_by_platform(OWNER_A, "tg_12345", "telegram") .await .unwrap() .unwrap(); @@ -635,32 +843,99 @@ mod tests { #[tokio::test] async fn create_duplicate_user_returns_conflict() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let dup = AssistantUserRow { id: "usr-2".into(), ..sample_user() }; - let err = repo.create_user(&dup).await.unwrap_err(); + let err = repo.create_user(OWNER_A, &dup).await.unwrap_err(); assert!(matches!(err, DbError::Conflict(_))); } + #[tokio::test] + async fn platform_users_are_filtered_by_owner() { + let (repo, db) = setup().await; + create_owner(db.pool(), OWNER_B).await; + + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); + let other = AssistantUserRow { + id: "usr-2".into(), + owner_user_id: OWNER_B.into(), + platform_user_id: "tg_other".into(), + ..sample_user() + }; + repo.create_user(OWNER_B, &other).await.unwrap(); + + let owner_a = repo + .get_user_by_platform(OWNER_A, "tg_12345", "telegram") + .await + .unwrap() + .unwrap(); + let owner_b = repo + .get_user_by_platform(OWNER_B, "tg_other", "telegram") + .await + .unwrap() + .unwrap(); + + assert_eq!(owner_a.id, "usr-1"); + assert_eq!(owner_b.id, "usr-2"); + assert!( + repo.get_user_by_platform(OWNER_A, "tg_other", "telegram") + .await + .unwrap() + .is_none() + ); + } + + #[tokio::test] + async fn same_platform_user_can_exist_for_different_owners() { + let (repo, db) = setup().await; + create_owner(db.pool(), OWNER_B).await; + + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); + let other_owner_user = AssistantUserRow { + id: "usr-2".into(), + owner_user_id: OWNER_B.into(), + ..sample_user() + }; + repo.create_user(OWNER_B, &other_owner_user).await.unwrap(); + + let owner_a = repo + .get_user_by_platform(OWNER_A, "tg_12345", "telegram") + .await + .unwrap() + .unwrap(); + let owner_b = repo + .get_user_by_platform(OWNER_B, "tg_12345", "telegram") + .await + .unwrap() + .unwrap(); + assert_eq!(owner_a.id, "usr-1"); + assert_eq!(owner_b.id, "usr-2"); + } + #[tokio::test] async fn get_user_by_platform_not_found() { let (repo, _db) = setup().await; - assert!(repo.get_user_by_platform("nope", "telegram").await.unwrap().is_none()); + assert!( + repo.get_user_by_platform(OWNER_A, "nope", "telegram") + .await + .unwrap() + .is_none() + ); } #[tokio::test] async fn update_user_last_active_updates_timestamp() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let new_ts = aionui_common::now_ms() + 5000; - repo.update_user_last_active("usr-1", new_ts).await.unwrap(); + repo.update_user_last_active(OWNER_A, "usr-1", new_ts).await.unwrap(); let found = repo - .get_user_by_platform("tg_12345", "telegram") + .get_user_by_platform(OWNER_A, "tg_12345", "telegram") .await .unwrap() .unwrap(); @@ -670,17 +945,17 @@ mod tests { #[tokio::test] async fn update_user_last_active_not_found() { let (repo, _db) = setup().await; - let err = repo.update_user_last_active("nope", 123).await.unwrap_err(); + let err = repo.update_user_last_active(OWNER_A, "nope", 123).await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } #[tokio::test] async fn delete_user_removes_row() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); - repo.delete_user("usr-1").await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); + repo.delete_user(OWNER_A, "usr-1").await.unwrap(); assert!( - repo.get_user_by_platform("tg_12345", "telegram") + repo.get_user_by_platform(OWNER_A, "tg_12345", "telegram") .await .unwrap() .is_none() @@ -690,25 +965,27 @@ mod tests { #[tokio::test] async fn delete_user_not_found() { let (repo, _db) = setup().await; - let err = repo.delete_user("nope").await.unwrap_err(); + let err = repo.delete_user(OWNER_A, "nope").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } #[tokio::test] async fn delete_user_cascades_sessions() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let session = sample_session("usr-1"); - repo.get_or_create_session("usr-1", "chat-abc", &session).await.unwrap(); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-abc", &session) + .await + .unwrap(); // Sessions exist before delete. - assert_eq!(repo.get_all_sessions().await.unwrap().len(), 1); + assert_eq!(repo.get_all_sessions(OWNER_A).await.unwrap().len(), 1); - repo.delete_user("usr-1").await.unwrap(); + repo.delete_user(OWNER_A, "usr-1").await.unwrap(); // Sessions cascade-deleted. - assert!(repo.get_all_sessions().await.unwrap().is_empty()); + assert!(repo.get_all_sessions(OWNER_A).await.unwrap().is_empty()); } // ── Session tests ──────────────────────────────────────────────── @@ -716,16 +993,19 @@ mod tests { #[tokio::test] async fn get_all_sessions_empty() { let (repo, _db) = setup().await; - assert!(repo.get_all_sessions().await.unwrap().is_empty()); + assert!(repo.get_all_sessions(OWNER_A).await.unwrap().is_empty()); } #[tokio::test] async fn get_or_create_session_creates_new() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let new = sample_session("usr-1"); - let result = repo.get_or_create_session("usr-1", "chat-abc", &new).await.unwrap(); + let result = repo + .get_or_create_session(OWNER_A, "usr-1", "chat-abc", &new) + .await + .unwrap(); assert_eq!(result.id, "sess-1"); assert_eq!(result.user_id, "usr-1"); assert_eq!(result.chat_id.as_deref(), Some("chat-abc")); @@ -734,17 +1014,23 @@ mod tests { #[tokio::test] async fn get_or_create_session_reuses_existing() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let new = sample_session("usr-1"); - let first = repo.get_or_create_session("usr-1", "chat-abc", &new).await.unwrap(); + let first = repo + .get_or_create_session(OWNER_A, "usr-1", "chat-abc", &new) + .await + .unwrap(); // Second call with different new_row id should still return the first. let another = AssistantSessionRow { id: "sess-2".into(), ..new }; - let second = repo.get_or_create_session("usr-1", "chat-abc", &another).await.unwrap(); + let second = repo + .get_or_create_session(OWNER_A, "usr-1", "chat-abc", &another) + .await + .unwrap(); assert_eq!(second.id, first.id); // last_activity should be updated. assert!(second.last_activity >= first.last_activity); @@ -753,106 +1039,109 @@ mod tests { #[tokio::test] async fn per_chat_isolation_different_chats() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let s1 = sample_session("usr-1"); - repo.get_or_create_session("usr-1", "chat-abc", &s1).await.unwrap(); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-abc", &s1) + .await + .unwrap(); let s2 = AssistantSessionRow { id: "sess-2".into(), chat_id: Some("chat-xyz".into()), ..sample_session("usr-1") }; - repo.get_or_create_session("usr-1", "chat-xyz", &s2).await.unwrap(); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-xyz", &s2) + .await + .unwrap(); - assert_eq!(repo.get_all_sessions().await.unwrap().len(), 2); + assert_eq!(repo.get_all_sessions(OWNER_A).await.unwrap().len(), 2); } #[tokio::test] async fn get_session_by_id() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let new = sample_session("usr-1"); - repo.get_or_create_session("usr-1", "chat-abc", &new).await.unwrap(); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-abc", &new) + .await + .unwrap(); - let found = repo.get_session("sess-1").await.unwrap().unwrap(); + let found = repo.get_session(OWNER_A, "sess-1").await.unwrap().unwrap(); assert_eq!(found.agent_type, "gemini"); } #[tokio::test] async fn get_session_not_found() { let (repo, _db) = setup().await; - assert!(repo.get_session("nope").await.unwrap().is_none()); + assert!(repo.get_session(OWNER_A, "nope").await.unwrap().is_none()); } #[tokio::test] async fn update_session_activity_updates_timestamp() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let new = sample_session("usr-1"); - repo.get_or_create_session("usr-1", "chat-abc", &new).await.unwrap(); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-abc", &new) + .await + .unwrap(); let new_ts = aionui_common::now_ms() + 5000; - repo.update_session_activity("sess-1", new_ts).await.unwrap(); + repo.update_session_activity(OWNER_A, "sess-1", new_ts).await.unwrap(); - let found = repo.get_session("sess-1").await.unwrap().unwrap(); + let found = repo.get_session(OWNER_A, "sess-1").await.unwrap().unwrap(); assert_eq!(found.last_activity, new_ts); } #[tokio::test] async fn update_session_activity_not_found() { let (repo, _db) = setup().await; - let err = repo.update_session_activity("nope", 123).await.unwrap_err(); + let err = repo.update_session_activity(OWNER_A, "nope", 123).await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } #[tokio::test] async fn delete_sessions_by_user_removes_all() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let s1 = sample_session("usr-1"); - repo.get_or_create_session("usr-1", "chat-abc", &s1).await.unwrap(); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-abc", &s1) + .await + .unwrap(); let s2 = AssistantSessionRow { id: "sess-2".into(), chat_id: Some("chat-xyz".into()), ..sample_session("usr-1") }; - repo.get_or_create_session("usr-1", "chat-xyz", &s2).await.unwrap(); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-xyz", &s2) + .await + .unwrap(); - repo.delete_sessions_by_user("usr-1").await.unwrap(); - assert!(repo.get_all_sessions().await.unwrap().is_empty()); + repo.delete_sessions_by_user(OWNER_A, "usr-1").await.unwrap(); + assert!(repo.get_all_sessions(OWNER_A).await.unwrap().is_empty()); } #[tokio::test] async fn delete_sessions_by_user_no_sessions_is_ok() { let (repo, _db) = setup().await; // No sessions exist for this user — should not error. - repo.delete_sessions_by_user("usr-1").await.unwrap(); + repo.delete_sessions_by_user(OWNER_A, "usr-1").await.unwrap(); } /// Helper to create a stub conversation for FK-constrained tests. - async fn create_stub_conversation(pool: &SqlitePool, conv_id: &str) { + async fn create_stub_conversation(pool: &SqlitePool, user_id: &str, conv_id: &str) { let now = aionui_common::now_ms(); - // Create prerequisite user first (matches `users` table schema) - sqlx::query( - "INSERT OR IGNORE INTO users \ - (id, username, password_hash, created_at, updated_at) \ - VALUES ('sys_test', 'test_user', 'hash', ?1, ?1)", - ) - .bind(now) - .execute(pool) - .await - .unwrap(); sqlx::query( "INSERT INTO conversations (id, user_id, name, type, created_at, updated_at) \ - VALUES (?1, 'sys_test', 'Test Conv', 'chat', ?2, ?2)", + VALUES (?1, ?2, 'Test Conv', 'chat', ?3, ?3)", ) .bind(conv_id) + .bind(user_id) .bind(now) .execute(pool) .await @@ -862,67 +1151,119 @@ mod tests { #[tokio::test] async fn update_session_conversation_persists() { let (repo, db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let new = sample_session("usr-1"); - repo.get_or_create_session("usr-1", "chat-abc", &new).await.unwrap(); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-abc", &new) + .await + .unwrap(); - create_stub_conversation(db.pool(), "conv-42").await; + create_stub_conversation(db.pool(), OWNER_A, "conv-42").await; - repo.update_session_conversation("sess-1", "conv-42").await.unwrap(); + repo.update_session_conversation(OWNER_A, "sess-1", "conv-42") + .await + .unwrap(); - let found = repo.get_session("sess-1").await.unwrap().unwrap(); + let found = repo.get_session(OWNER_A, "sess-1").await.unwrap().unwrap(); assert_eq!(found.conversation_id.as_deref(), Some("conv-42")); } + #[tokio::test] + async fn update_session_conversation_rejects_cross_owner_conversation() { + let (repo, db) = setup().await; + create_owner(db.pool(), OWNER_B).await; + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); + + let new = sample_session("usr-1"); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-abc", &new) + .await + .unwrap(); + + create_stub_conversation(db.pool(), OWNER_B, "conv-other").await; + + let err = repo + .update_session_conversation(OWNER_A, "sess-1", "conv-other") + .await + .unwrap_err(); + assert!(matches!( + err, + DbError::Conflict(msg) if msg.starts_with("CROSS_ACCOUNT_REFERENCE:") + )); + assert!( + repo.get_session(OWNER_A, "sess-1") + .await + .unwrap() + .unwrap() + .conversation_id + .is_none() + ); + } + #[tokio::test] async fn update_session_conversation_not_found() { let (repo, _db) = setup().await; - let err = repo.update_session_conversation("nope", "conv-1").await.unwrap_err(); + let err = repo + .update_session_conversation(OWNER_A, "nope", "conv-1") + .await + .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } #[tokio::test] async fn update_session_agent_type_persists() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let new = sample_session("usr-1"); - repo.get_or_create_session("usr-1", "chat-abc", &new).await.unwrap(); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-abc", &new) + .await + .unwrap(); - assert_eq!(repo.get_session("sess-1").await.unwrap().unwrap().agent_type, "gemini"); + assert_eq!( + repo.get_session(OWNER_A, "sess-1").await.unwrap().unwrap().agent_type, + "gemini" + ); - repo.update_session_agent_type("sess-1", "acp").await.unwrap(); + repo.update_session_agent_type(OWNER_A, "sess-1", "acp").await.unwrap(); - let found = repo.get_session("sess-1").await.unwrap().unwrap(); + let found = repo.get_session(OWNER_A, "sess-1").await.unwrap().unwrap(); assert_eq!(found.agent_type, "acp"); } #[tokio::test] async fn update_session_agent_type_not_found() { let (repo, _db) = setup().await; - let err = repo.update_session_agent_type("nope", "acp").await.unwrap_err(); + let err = repo + .update_session_agent_type(OWNER_A, "nope", "acp") + .await + .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } #[tokio::test] async fn delete_session_by_user_chat_removes_only_target() { let (repo, _db) = setup().await; - repo.create_user(&sample_user()).await.unwrap(); + repo.create_user(OWNER_A, &sample_user()).await.unwrap(); let s1 = sample_session("usr-1"); - repo.get_or_create_session("usr-1", "chat-abc", &s1).await.unwrap(); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-abc", &s1) + .await + .unwrap(); let s2 = AssistantSessionRow { id: "sess-2".into(), chat_id: Some("chat-xyz".into()), ..sample_session("usr-1") }; - repo.get_or_create_session("usr-1", "chat-xyz", &s2).await.unwrap(); + repo.get_or_create_session(OWNER_A, "usr-1", "chat-xyz", &s2) + .await + .unwrap(); - repo.delete_session_by_user_chat("usr-1", "chat-abc").await.unwrap(); + repo.delete_session_by_user_chat(OWNER_A, "usr-1", "chat-abc") + .await + .unwrap(); - let remaining = repo.get_all_sessions().await.unwrap(); + let remaining = repo.get_all_sessions(OWNER_A).await.unwrap(); assert_eq!(remaining.len(), 1); assert_eq!(remaining[0].chat_id.as_deref(), Some("chat-xyz")); } @@ -931,7 +1272,9 @@ mod tests { async fn delete_session_by_user_chat_no_match_is_ok() { let (repo, _db) = setup().await; // No sessions exist — should not error. - repo.delete_session_by_user_chat("usr-1", "chat-abc").await.unwrap(); + repo.delete_session_by_user_chat(OWNER_A, "usr-1", "chat-abc") + .await + .unwrap(); } // ── Pairing tests ──────────────────────────────────────────────── @@ -940,9 +1283,9 @@ mod tests { async fn create_and_get_pairing() { let (repo, _db) = setup().await; let pairing = sample_pairing(); - repo.create_pairing(&pairing).await.unwrap(); + repo.create_pairing(OWNER_A, &pairing).await.unwrap(); - let found = repo.get_pairing_by_code("123456").await.unwrap().unwrap(); + let found = repo.get_pairing_by_code(OWNER_A, "123456").await.unwrap().unwrap(); assert_eq!(found.platform_user_id, "tg_99"); assert_eq!(found.status, "pending"); } @@ -950,25 +1293,55 @@ mod tests { #[tokio::test] async fn create_duplicate_pairing_returns_conflict() { let (repo, _db) = setup().await; - repo.create_pairing(&sample_pairing()).await.unwrap(); - let err = repo.create_pairing(&sample_pairing()).await.unwrap_err(); + repo.create_pairing(OWNER_A, &sample_pairing()).await.unwrap(); + let err = repo.create_pairing(OWNER_A, &sample_pairing()).await.unwrap_err(); assert!(matches!(err, DbError::Conflict(_))); } + #[tokio::test] + async fn pairing_lookup_is_filtered_by_owner() { + let (repo, db) = setup().await; + create_owner(db.pool(), OWNER_B).await; + + repo.create_pairing(OWNER_A, &sample_pairing()).await.unwrap(); + + assert!(repo.get_pairing_by_code(OWNER_B, "123456").await.unwrap().is_none()); + assert!(repo.get_pending_pairings(OWNER_B).await.unwrap().is_empty()); + } + + #[tokio::test] + async fn same_pairing_code_can_exist_for_different_owners() { + let (repo, db) = setup().await; + create_owner(db.pool(), OWNER_B).await; + + repo.create_pairing(OWNER_A, &sample_pairing()).await.unwrap(); + let owner_b_pairing = PairingCodeRow { + owner_user_id: OWNER_B.into(), + platform_user_id: "tg_owner_b".into(), + ..sample_pairing() + }; + repo.create_pairing(OWNER_B, &owner_b_pairing).await.unwrap(); + + let owner_a = repo.get_pairing_by_code(OWNER_A, "123456").await.unwrap().unwrap(); + let owner_b = repo.get_pairing_by_code(OWNER_B, "123456").await.unwrap().unwrap(); + assert_eq!(owner_a.platform_user_id, "tg_99"); + assert_eq!(owner_b.platform_user_id, "tg_owner_b"); + } + #[tokio::test] async fn get_pending_pairings_filters_by_status() { let (repo, _db) = setup().await; let p1 = sample_pairing(); - repo.create_pairing(&p1).await.unwrap(); + repo.create_pairing(OWNER_A, &p1).await.unwrap(); let p2 = PairingCodeRow { code: "654321".into(), status: "approved".into(), ..sample_pairing() }; - repo.create_pairing(&p2).await.unwrap(); + repo.create_pairing(OWNER_A, &p2).await.unwrap(); - let pending = repo.get_pending_pairings().await.unwrap(); + let pending = repo.get_pending_pairings(OWNER_A).await.unwrap(); assert_eq!(pending.len(), 1); assert_eq!(pending[0].code, "123456"); } @@ -976,24 +1349,27 @@ mod tests { #[tokio::test] async fn get_pairing_by_code_not_found() { let (repo, _db) = setup().await; - assert!(repo.get_pairing_by_code("000000").await.unwrap().is_none()); + assert!(repo.get_pairing_by_code(OWNER_A, "000000").await.unwrap().is_none()); } #[tokio::test] async fn update_pairing_status_changes_status() { let (repo, _db) = setup().await; - repo.create_pairing(&sample_pairing()).await.unwrap(); + repo.create_pairing(OWNER_A, &sample_pairing()).await.unwrap(); - repo.update_pairing_status("123456", "approved").await.unwrap(); + repo.update_pairing_status(OWNER_A, "123456", "approved").await.unwrap(); - let found = repo.get_pairing_by_code("123456").await.unwrap().unwrap(); + let found = repo.get_pairing_by_code(OWNER_A, "123456").await.unwrap().unwrap(); assert_eq!(found.status, "approved"); } #[tokio::test] async fn update_pairing_status_not_found() { let (repo, _db) = setup().await; - let err = repo.update_pairing_status("000000", "approved").await.unwrap_err(); + let err = repo + .update_pairing_status(OWNER_A, "000000", "approved") + .await + .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -1008,7 +1384,7 @@ mod tests { expires_at: now - 1000, ..sample_pairing() }; - repo.create_pairing(&expired).await.unwrap(); + repo.create_pairing(OWNER_A, &expired).await.unwrap(); // Create a still-valid pairing. let valid = PairingCodeRow { @@ -1016,15 +1392,15 @@ mod tests { expires_at: now + 600_000, ..sample_pairing() }; - repo.create_pairing(&valid).await.unwrap(); + repo.create_pairing(OWNER_A, &valid).await.unwrap(); - let cleaned = repo.cleanup_expired_pairings(now).await.unwrap(); + let cleaned = repo.cleanup_expired_pairings(OWNER_A, now).await.unwrap(); assert_eq!(cleaned, 1); - let found_expired = repo.get_pairing_by_code("111111").await.unwrap().unwrap(); + let found_expired = repo.get_pairing_by_code(OWNER_A, "111111").await.unwrap().unwrap(); assert_eq!(found_expired.status, "expired"); - let found_valid = repo.get_pairing_by_code("222222").await.unwrap().unwrap(); + let found_valid = repo.get_pairing_by_code(OWNER_A, "222222").await.unwrap().unwrap(); assert_eq!(found_valid.status, "pending"); } @@ -1040,12 +1416,12 @@ mod tests { status: "approved".into(), ..sample_pairing() }; - repo.create_pairing(&approved).await.unwrap(); + repo.create_pairing(OWNER_A, &approved).await.unwrap(); - let cleaned = repo.cleanup_expired_pairings(now).await.unwrap(); + let cleaned = repo.cleanup_expired_pairings(OWNER_A, now).await.unwrap(); assert_eq!(cleaned, 0); - let found = repo.get_pairing_by_code("333333").await.unwrap().unwrap(); + let found = repo.get_pairing_by_code(OWNER_A, "333333").await.unwrap().unwrap(); assert_eq!(found.status, "approved"); } } diff --git a/crates/aionui-db/src/repository/sqlite_client_preference.rs b/crates/aionui-db/src/repository/sqlite_client_preference.rs index f5996fa7b..cfeb8cfaa 100644 --- a/crates/aionui-db/src/repository/sqlite_client_preference.rs +++ b/crates/aionui-db/src/repository/sqlite_client_preference.rs @@ -18,15 +18,17 @@ impl SqliteClientPreferenceRepository { #[async_trait::async_trait] impl IClientPreferenceRepository for SqliteClientPreferenceRepository { - async fn get_all(&self) -> Result, DbError> { - let rows = sqlx::query_as::<_, ClientPreference>("SELECT * FROM client_preferences ORDER BY key") - .fetch_all(&self.pool) - .await?; + async fn get_all(&self, user_id: &str) -> Result, DbError> { + let rows = + sqlx::query_as::<_, ClientPreference>("SELECT * FROM client_preferences WHERE user_id = ? ORDER BY key") + .bind(user_id) + .fetch_all(&self.pool) + .await?; Ok(rows) } - async fn get_by_keys(&self, keys: &[&str]) -> Result, DbError> { + async fn get_by_keys(&self, user_id: &str, keys: &[&str]) -> Result, DbError> { if keys.is_empty() { return Ok(vec![]); } @@ -34,11 +36,11 @@ impl IClientPreferenceRepository for SqliteClientPreferenceRepository { // Build dynamic IN clause with positional placeholders let placeholders: Vec<&str> = keys.iter().map(|_| "?").collect(); let sql = format!( - "SELECT * FROM client_preferences WHERE key IN ({}) ORDER BY key", + "SELECT * FROM client_preferences WHERE user_id = ? AND key IN ({}) ORDER BY key", placeholders.join(", ") ); - let mut query = sqlx::query_as::<_, ClientPreference>(&sql); + let mut query = sqlx::query_as::<_, ClientPreference>(&sql).bind(user_id); for key in keys { query = query.bind(*key); } @@ -47,7 +49,7 @@ impl IClientPreferenceRepository for SqliteClientPreferenceRepository { Ok(rows) } - async fn upsert_batch(&self, entries: &[(&str, &str)]) -> Result<(), DbError> { + async fn upsert_batch(&self, user_id: &str, entries: &[(&str, &str)]) -> Result<(), DbError> { if entries.is_empty() { return Ok(()); } @@ -59,12 +61,13 @@ impl IClientPreferenceRepository for SqliteClientPreferenceRepository { for (key, value) in entries { sqlx::query( - "INSERT INTO client_preferences (key, value, updated_at) \ - VALUES (?, ?, ?) \ - ON CONFLICT(key) DO UPDATE SET \ + "INSERT INTO client_preferences (user_id, key, value, updated_at) \ + VALUES (?, ?, ?, ?) \ + ON CONFLICT(user_id, key) DO UPDATE SET \ value = excluded.value, \ updated_at = excluded.updated_at", ) + .bind(user_id) .bind(*key) .bind(*value) .bind(now) @@ -76,18 +79,18 @@ impl IClientPreferenceRepository for SqliteClientPreferenceRepository { Ok(()) } - async fn delete_keys(&self, keys: &[&str]) -> Result<(), DbError> { + async fn delete_keys(&self, user_id: &str, keys: &[&str]) -> Result<(), DbError> { if keys.is_empty() { return Ok(()); } let placeholders: Vec<&str> = keys.iter().map(|_| "?").collect(); let sql = format!( - "DELETE FROM client_preferences WHERE key IN ({})", + "DELETE FROM client_preferences WHERE user_id = ? AND key IN ({})", placeholders.join(", ") ); - let mut query = sqlx::query(&sql); + let mut query = sqlx::query(&sql).bind(user_id); for key in keys { query = query.bind(*key); } @@ -102,8 +105,20 @@ mod tests { use super::*; use crate::init_database_memory; + const USER_A: &str = "system_default_user"; + const USER_B: &str = "user_b"; + async fn setup() -> (SqliteClientPreferenceRepository, crate::Database) { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, 1, 1)", + ) + .bind(USER_B) + .bind(USER_B) + .execute(db.pool()) + .await + .unwrap(); let repo = SqliteClientPreferenceRepository::new(db.pool().clone()); (repo, db) } @@ -111,18 +126,18 @@ mod tests { #[tokio::test] async fn get_all_empty() { let (repo, _db) = setup().await; - let prefs = repo.get_all().await.unwrap(); + let prefs = repo.get_all(USER_A).await.unwrap(); assert!(prefs.is_empty()); } #[tokio::test] async fn upsert_and_get_all() { let (repo, _db) = setup().await; - repo.upsert_batch(&[("theme", "\"dark\""), ("pet.size", "360")]) + repo.upsert_batch(USER_A, &[("theme", "\"dark\""), ("pet.size", "360")]) .await .unwrap(); - let prefs = repo.get_all().await.unwrap(); + let prefs = repo.get_all(USER_A).await.unwrap(); assert_eq!(prefs.len(), 2); assert_eq!(prefs[0].key, "pet.size"); assert_eq!(prefs[0].value, "360"); @@ -133,9 +148,11 @@ mod tests { #[tokio::test] async fn get_by_keys_filters_correctly() { let (repo, _db) = setup().await; - repo.upsert_batch(&[("a", "1"), ("b", "2"), ("c", "3")]).await.unwrap(); + repo.upsert_batch(USER_A, &[("a", "1"), ("b", "2"), ("c", "3")]) + .await + .unwrap(); - let prefs = repo.get_by_keys(&["a", "c", "nonexistent"]).await.unwrap(); + let prefs = repo.get_by_keys(USER_A, &["a", "c", "nonexistent"]).await.unwrap(); assert_eq!(prefs.len(), 2); let keys: Vec<&str> = prefs.iter().map(|p| p.key.as_str()).collect(); @@ -146,17 +163,17 @@ mod tests { #[tokio::test] async fn get_by_keys_empty_input() { let (repo, _db) = setup().await; - let prefs = repo.get_by_keys(&[]).await.unwrap(); + let prefs = repo.get_by_keys(USER_A, &[]).await.unwrap(); assert!(prefs.is_empty()); } #[tokio::test] async fn upsert_overwrites_existing_key() { let (repo, _db) = setup().await; - repo.upsert_batch(&[("k", "v1")]).await.unwrap(); - repo.upsert_batch(&[("k", "v2")]).await.unwrap(); + repo.upsert_batch(USER_A, &[("k", "v1")]).await.unwrap(); + repo.upsert_batch(USER_A, &[("k", "v2")]).await.unwrap(); - let prefs = repo.get_all().await.unwrap(); + let prefs = repo.get_all(USER_A).await.unwrap(); assert_eq!(prefs.len(), 1); assert_eq!(prefs[0].value, "v2"); } @@ -164,11 +181,13 @@ mod tests { #[tokio::test] async fn delete_keys_removes_entries() { let (repo, _db) = setup().await; - repo.upsert_batch(&[("a", "1"), ("b", "2"), ("c", "3")]).await.unwrap(); + repo.upsert_batch(USER_A, &[("a", "1"), ("b", "2"), ("c", "3")]) + .await + .unwrap(); - repo.delete_keys(&["a", "c"]).await.unwrap(); + repo.delete_keys(USER_A, &["a", "c"]).await.unwrap(); - let prefs = repo.get_all().await.unwrap(); + let prefs = repo.get_all(USER_A).await.unwrap(); assert_eq!(prefs.len(), 1); assert_eq!(prefs[0].key, "b"); } @@ -176,22 +195,39 @@ mod tests { #[tokio::test] async fn delete_keys_nonexistent_is_noop() { let (repo, _db) = setup().await; - repo.delete_keys(&["ghost"]).await.unwrap(); - assert!(repo.get_all().await.unwrap().is_empty()); + repo.delete_keys(USER_A, &["ghost"]).await.unwrap(); + assert!(repo.get_all(USER_A).await.unwrap().is_empty()); } #[tokio::test] async fn delete_keys_empty_input() { let (repo, _db) = setup().await; - repo.upsert_batch(&[("x", "1")]).await.unwrap(); - repo.delete_keys(&[]).await.unwrap(); - assert_eq!(repo.get_all().await.unwrap().len(), 1); + repo.upsert_batch(USER_A, &[("x", "1")]).await.unwrap(); + repo.delete_keys(USER_A, &[]).await.unwrap(); + assert_eq!(repo.get_all(USER_A).await.unwrap().len(), 1); } #[tokio::test] async fn upsert_empty_batch_is_noop() { let (repo, _db) = setup().await; - repo.upsert_batch(&[]).await.unwrap(); - assert!(repo.get_all().await.unwrap().is_empty()); + repo.upsert_batch(USER_A, &[]).await.unwrap(); + assert!(repo.get_all(USER_A).await.unwrap().is_empty()); + } + + #[tokio::test] + async fn preferences_are_scoped_by_user() { + let (repo, _db) = setup().await; + repo.upsert_batch(USER_A, &[("theme", "\"dark\"")]).await.unwrap(); + repo.upsert_batch(USER_B, &[("theme", "\"light\"")]).await.unwrap(); + + assert_eq!(repo.get_by_keys(USER_A, &["theme"]).await.unwrap()[0].value, "\"dark\""); + assert_eq!( + repo.get_by_keys(USER_B, &["theme"]).await.unwrap()[0].value, + "\"light\"" + ); + + repo.delete_keys(USER_B, &["theme"]).await.unwrap(); + assert!(repo.get_by_keys(USER_B, &["theme"]).await.unwrap().is_empty()); + assert_eq!(repo.get_by_keys(USER_A, &["theme"]).await.unwrap().len(), 1); } } diff --git a/crates/aionui-db/src/repository/sqlite_conversation.rs b/crates/aionui-db/src/repository/sqlite_conversation.rs index fdc9acfda..d341c48ca 100644 --- a/crates/aionui-db/src/repository/sqlite_conversation.rs +++ b/crates/aionui-db/src/repository/sqlite_conversation.rs @@ -1,4 +1,4 @@ -use sqlx::SqlitePool; +use sqlx::{Row, SqlitePool}; use aionui_common::PaginatedResult; @@ -9,7 +9,7 @@ use crate::models::{ }; use crate::repository::conversation::{ ConversationFilters, ConversationRowUpdate, IConversationRepository, MessagePageCursor, MessagePageDirection, - MessagePageParams, MessagePageResult, MessageRowUpdate, MessageSearchRow, + MessagePageParams, MessagePageResult, MessageRowUpdate, MessageSearchRow, StaleRuntimeMessageRow, }; /// Bump `conversations.updated_at` so the conversation-list sort @@ -20,19 +20,6 @@ use crate::repository::conversation::{ /// out-of-order streaming upsert (older event time) can never move a /// conversation backward in the list. Runs inside the caller's transaction so /// the message write and the bump commit atomically. -async fn bump_conversation_updated_at( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - conversation_id: &str, - at: i64, -) -> Result<(), sqlx::Error> { - sqlx::query("UPDATE conversations SET updated_at = MAX(updated_at, ?) WHERE id = ?") - .bind(at) - .bind(conversation_id) - .execute(&mut **tx) - .await?; - Ok(()) -} - /// SQLite-backed implementation of [`IConversationRepository`]. #[derive(Clone, Debug)] pub struct SqliteConversationRepository { @@ -44,36 +31,97 @@ impl SqliteConversationRepository { Self { pool } } - async fn insert_message_once(&self, message: &MessageRow) -> Result<(), sqlx::Error> { - let mut tx = self.pool.begin().await?; - sqlx::query( - "INSERT INTO messages \ - (id, conversation_id, msg_id, type, content, position, \ - status, hidden, created_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind(&message.id) - .bind(&message.conversation_id) - .bind(&message.msg_id) - .bind(&message.r#type) - .bind(&message.content) - .bind(&message.position) - .bind(&message.status) - .bind(message.hidden) - .bind(message.created_at) - .execute(&mut *tx) - .await?; + async fn conversation_exists_for_user(&self, user_id: &str, conversation_id: &str) -> Result { + let exists: i64 = sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM conversations WHERE user_id = ? AND id = ?)") + .bind(user_id) + .bind(conversation_id) + .fetch_one(&self.pool) + .await?; - bump_conversation_updated_at(&mut tx, &message.conversation_id, message.created_at).await?; - tx.commit().await?; + Ok(exists != 0) + } - Ok(()) + async fn ensure_conversation_for_user(&self, user_id: &str, conversation_id: &str) -> Result<(), DbError> { + if self.conversation_exists_for_user(user_id, conversation_id).await? { + Ok(()) + } else { + Err(DbError::NotFound(format!("Conversation '{conversation_id}' not found"))) + } } - async fn upsert_message_once(&self, message: &MessageRow) -> Result<(), sqlx::Error> { - let mut tx = self.pool.begin().await?; - sqlx::query( - "INSERT INTO messages \ + async fn insert_message_once(&self, user_id: &str, message: &MessageRow) -> Result<(), DbError> { + // BEGIN IMMEDIATE claims the writer lock up front (same pattern as + // `claim_run`) so concurrent inserters queue on SQLite's busy handler + // instead of a read-then-write transaction failing with "database is + // locked" when it tries to upgrade to the write lock. + let mut connection = self.pool.acquire().await?; + sqlx::query("BEGIN IMMEDIATE").execute(&mut *connection).await?; + + let result: Result<(), DbError> = async { + // Ownership check inside the same transaction as the insert + + // bump, so parent-chain authorization and the write are atomic. + let exists: i64 = + sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM conversations WHERE user_id = ? AND id = ?)") + .bind(user_id) + .bind(&message.conversation_id) + .fetch_one(&mut *connection) + .await?; + if exists == 0 { + return Err(DbError::NotFound(format!( + "Conversation '{}' not found", + message.conversation_id + ))); + } + sqlx::query( + "INSERT INTO messages \ + (id, conversation_id, msg_id, type, content, position, \ + status, hidden, created_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(&message.id) + .bind(&message.conversation_id) + .bind(&message.msg_id) + .bind(&message.r#type) + .bind(&message.content) + .bind(&message.position) + .bind(&message.status) + .bind(message.hidden) + .bind(message.created_at) + .execute(&mut *connection) + .await?; + + // Persisting a message bumps the parent conversation's recency so + // the conversation-list sort floats fresh activity to the top; + // MAX() keeps recency monotonic under out-of-order upserts. + sqlx::query("UPDATE conversations SET updated_at = MAX(updated_at, ?) WHERE id = ?") + .bind(message.created_at) + .bind(&message.conversation_id) + .execute(&mut *connection) + .await?; + Ok(()) + } + .await; + + match result { + Ok(()) => { + sqlx::query("COMMIT").execute(&mut *connection).await?; + Ok(()) + } + Err(error) => { + let _ = sqlx::query("ROLLBACK").execute(&mut *connection).await; + Err(error) + } + } + } + + async fn upsert_message_once(&self, user_id: &str, message: &MessageRow) -> Result<(), DbError> { + // Writer-lock-first transaction; see `insert_message_once` for why. + let mut connection = self.pool.acquire().await?; + sqlx::query("BEGIN IMMEDIATE").execute(&mut *connection).await?; + + let result: Result<(), DbError> = async { + let result = sqlx::query( + "INSERT INTO messages \ (id, conversation_id, msg_id, type, content, position, \ status, hidden, created_at) \ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) \ @@ -100,35 +148,74 @@ impl SqliteConversationRepository { END, \ position = COALESCE(messages.position, excluded.position), \ hidden = excluded.hidden, \ - created_at = MIN(messages.created_at, excluded.created_at)", - ) - .bind(&message.id) - .bind(&message.conversation_id) - .bind(&message.msg_id) - .bind(&message.r#type) - .bind(&message.content) - .bind(&message.position) - .bind(&message.status) - .bind(message.hidden) - .bind(message.created_at) - .execute(&mut *tx) - .await?; + created_at = MIN(messages.created_at, excluded.created_at) \ + WHERE messages.conversation_id = excluded.conversation_id \ + AND EXISTS ( \ + SELECT 1 FROM conversations c \ + WHERE c.id = messages.conversation_id AND c.user_id = ? \ + )", + ) + .bind(&message.id) + .bind(&message.conversation_id) + .bind(&message.msg_id) + .bind(&message.r#type) + .bind(&message.content) + .bind(&message.position) + .bind(&message.status) + .bind(message.hidden) + .bind(message.created_at) + .bind(user_id) + .execute(&mut *connection) + .await?; - bump_conversation_updated_at(&mut tx, &message.conversation_id, message.created_at).await?; - tx.commit().await?; + // 0 rows: either the id exists under another conversation, or the + // user-scope EXISTS guard rejected a foreign owner. The rollback + // below means the timestamp bump never applies either way. + if result.rows_affected() == 0 { + return Err(DbError::Conflict(format!( + "Message with id '{}' already exists outside the requested conversation", + message.id + ))); + } - Ok(()) + sqlx::query("UPDATE conversations SET updated_at = MAX(updated_at, ?) WHERE id = ?") + .bind(message.created_at) + .bind(&message.conversation_id) + .execute(&mut *connection) + .await?; + Ok(()) + } + .await; + + match result { + Ok(()) => { + sqlx::query("COMMIT").execute(&mut *connection).await?; + Ok(()) + } + Err(error) => { + let _ = sqlx::query("ROLLBACK").execute(&mut *connection).await; + Err(error) + } + } } - async fn visible_message_exists_before(&self, conv_id: &str, cursor: &MessagePageCursor) -> Result { + async fn visible_message_exists_before( + &self, + user_id: &str, + conv_id: &str, + cursor: &MessagePageCursor, + ) -> Result { let exists: i64 = sqlx::query_scalar( "SELECT EXISTS( \ - SELECT 1 FROM messages \ - WHERE conversation_id = ? \ - AND (created_at < ? OR (created_at = ? AND id < ?)) \ - AND type NOT IN ('cron_trigger', 'skill_suggest') \ + SELECT 1 FROM messages m \ + INNER JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? \ + AND m.conversation_id = ? \ + AND (m.created_at < ? OR (m.created_at = ? AND m.id < ?)) \ + AND m.type NOT IN ('cron_trigger', 'skill_suggest') \ )", ) + .bind(user_id) .bind(conv_id) .bind(cursor.created_at) .bind(cursor.created_at) @@ -139,15 +226,23 @@ impl SqliteConversationRepository { Ok(exists != 0) } - async fn visible_message_exists_after(&self, conv_id: &str, cursor: &MessagePageCursor) -> Result { + async fn visible_message_exists_after( + &self, + user_id: &str, + conv_id: &str, + cursor: &MessagePageCursor, + ) -> Result { let exists: i64 = sqlx::query_scalar( "SELECT EXISTS( \ - SELECT 1 FROM messages \ - WHERE conversation_id = ? \ - AND (created_at > ? OR (created_at = ? AND id > ?)) \ - AND type NOT IN ('cron_trigger', 'skill_suggest') \ + SELECT 1 FROM messages m \ + INNER JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? \ + AND m.conversation_id = ? \ + AND (m.created_at > ? OR (m.created_at = ? AND m.id > ?)) \ + AND m.type NOT IN ('cron_trigger', 'skill_suggest') \ )", ) + .bind(user_id) .bind(conv_id) .bind(cursor.created_at) .bind(cursor.created_at) @@ -158,7 +253,12 @@ impl SqliteConversationRepository { Ok(exists != 0) } - async fn page_with_flags(&self, conv_id: &str, items: Vec) -> Result { + async fn page_with_flags( + &self, + user_id: &str, + conv_id: &str, + items: Vec, + ) -> Result { let Some(first) = items.first() else { return Ok(MessagePageResult { items, @@ -171,8 +271,12 @@ impl SqliteConversationRepository { .last() .map(MessagePageCursor::from) .unwrap_or(first_cursor.clone()); - let has_more_before = self.visible_message_exists_before(conv_id, &first_cursor).await?; - let has_more_after = self.visible_message_exists_after(conv_id, &last_cursor).await?; + let has_more_before = self + .visible_message_exists_before(user_id, conv_id, &first_cursor) + .await?; + let has_more_after = self + .visible_message_exists_after(user_id, conv_id, &last_cursor) + .await?; Ok(MessagePageResult { items, @@ -186,8 +290,9 @@ impl SqliteConversationRepository { impl IConversationRepository for SqliteConversationRepository { // ── Conversation CRUD ─────────────────────────────────────────── - async fn get(&self, id: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, ConversationRow>("SELECT * FROM conversations WHERE id = ?") + async fn get(&self, user_id: &str, id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, ConversationRow>("SELECT * FROM conversations WHERE user_id = ? AND id = ?") + .bind(user_id) .bind(id) .fetch_optional(&self.pool) .await?; @@ -195,6 +300,15 @@ impl IConversationRepository for SqliteConversationRepository { Ok(row) } + async fn owner_user_id(&self, id: &str) -> Result, DbError> { + let user_id = sqlx::query_scalar::<_, String>("SELECT user_id FROM conversations WHERE id = ?") + .bind(id) + .fetch_optional(&self.pool) + .await?; + + Ok(user_id) + } + async fn create(&self, row: &ConversationRow) -> Result<(), DbError> { sqlx::query( "INSERT INTO conversations \ @@ -221,7 +335,7 @@ impl IConversationRepository for SqliteConversationRepository { Ok(()) } - async fn update(&self, id: &str, updates: &ConversationRowUpdate) -> Result<(), DbError> { + async fn update(&self, user_id: &str, id: &str, updates: &ConversationRowUpdate) -> Result<(), DbError> { // Build dynamic SET clause let mut set_parts: Vec = Vec::new(); let mut binds: Vec = Vec::new(); @@ -259,12 +373,16 @@ impl IConversationRepository for SqliteConversationRepository { return Ok(()); } - let sql = format!("UPDATE conversations SET {} WHERE id = ?", set_parts.join(", ")); + let sql = format!( + "UPDATE conversations SET {} WHERE user_id = ? AND id = ?", + set_parts.join(", ") + ); let mut query = sqlx::query(&sql); for bind in &binds { query = bind_value(query, bind); } + query = query.bind(user_id); query = query.bind(id); let result = query.execute(&self.pool).await?; @@ -276,8 +394,9 @@ impl IConversationRepository for SqliteConversationRepository { Ok(()) } - async fn delete(&self, id: &str) -> Result<(), DbError> { - let result = sqlx::query("DELETE FROM conversations WHERE id = ?") + async fn delete(&self, user_id: &str, id: &str) -> Result<(), DbError> { + let result = sqlx::query("DELETE FROM conversations WHERE user_id = ? AND id = ?") + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -432,11 +551,15 @@ impl IConversationRepository for SqliteConversationRepository { async fn get_assistant_snapshot( &self, + user_id: &str, conversation_id: &str, ) -> Result, DbError> { let row = sqlx::query_as::<_, ConversationAssistantSnapshotRow>( - "SELECT * FROM conversation_assistant_snapshots WHERE conversation_id = ?", + "SELECT s.* FROM conversation_assistant_snapshots s \ + INNER JOIN conversations c ON c.id = s.conversation_id \ + WHERE c.user_id = ? AND s.conversation_id = ?", ) + .bind(user_id) .bind(conversation_id) .fetch_optional(&self.pool) .await?; @@ -446,8 +569,11 @@ impl IConversationRepository for SqliteConversationRepository { async fn upsert_assistant_snapshot( &self, + user_id: &str, params: &UpsertConversationAssistantSnapshotParams<'_>, ) -> Result, DbError> { + self.ensure_conversation_for_user(user_id, params.conversation_id) + .await?; let now = aionui_common::now_ms(); sqlx::query( "INSERT INTO conversation_assistant_snapshots ( @@ -512,14 +638,23 @@ impl IConversationRepository for SqliteConversationRepository { .execute(&self.pool) .await?; - self.get_assistant_snapshot(params.conversation_id).await + self.get_assistant_snapshot(user_id, params.conversation_id).await } - async fn delete_assistant_snapshot(&self, conversation_id: &str) -> Result { - let result = sqlx::query("DELETE FROM conversation_assistant_snapshots WHERE conversation_id = ?") - .bind(conversation_id) - .execute(&self.pool) - .await?; + async fn delete_assistant_snapshot(&self, user_id: &str, conversation_id: &str) -> Result { + let result = sqlx::query( + "DELETE FROM conversation_assistant_snapshots \ + WHERE conversation_id = ? \ + AND EXISTS ( \ + SELECT 1 FROM conversations c \ + WHERE c.id = conversation_assistant_snapshots.conversation_id \ + AND c.user_id = ? \ + )", + ) + .bind(conversation_id) + .bind(user_id) + .execute(&self.pool) + .await?; Ok(result.rows_affected() > 0) } @@ -528,21 +663,26 @@ impl IConversationRepository for SqliteConversationRepository { async fn list_messages_page( &self, + user_id: &str, conv_id: &str, params: &MessagePageParams, ) -> Result { + self.ensure_conversation_for_user(user_id, conv_id).await?; let limit = params.limit.max(1) as i64; let fetch_limit = limit + 1; let mut rows = match ¶ms.direction { MessagePageDirection::InitialLatest => { let mut rows = sqlx::query_as::<_, MessageRow>( - "SELECT * FROM messages \ - WHERE conversation_id = ? \ - AND type NOT IN ('cron_trigger', 'skill_suggest') \ - ORDER BY created_at DESC, id DESC \ + "SELECT m.* FROM messages m \ + INNER JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? \ + AND m.conversation_id = ? \ + AND m.type NOT IN ('cron_trigger', 'skill_suggest') \ + ORDER BY m.created_at DESC, m.id DESC \ LIMIT ?", ) + .bind(user_id) .bind(conv_id) .bind(fetch_limit) .fetch_all(&self.pool) @@ -553,13 +693,16 @@ impl IConversationRepository for SqliteConversationRepository { } MessagePageDirection::Before { cursor } => { let mut rows = sqlx::query_as::<_, MessageRow>( - "SELECT * FROM messages \ - WHERE conversation_id = ? \ - AND (created_at < ? OR (created_at = ? AND id < ?)) \ - AND type NOT IN ('cron_trigger', 'skill_suggest') \ - ORDER BY created_at DESC, id DESC \ + "SELECT m.* FROM messages m \ + INNER JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? \ + AND m.conversation_id = ? \ + AND (m.created_at < ? OR (m.created_at = ? AND m.id < ?)) \ + AND m.type NOT IN ('cron_trigger', 'skill_suggest') \ + ORDER BY m.created_at DESC, m.id DESC \ LIMIT ?", ) + .bind(user_id) .bind(conv_id) .bind(cursor.created_at) .bind(cursor.created_at) @@ -573,13 +716,16 @@ impl IConversationRepository for SqliteConversationRepository { } MessagePageDirection::After { cursor } => { let mut rows = sqlx::query_as::<_, MessageRow>( - "SELECT * FROM messages \ - WHERE conversation_id = ? \ - AND (created_at > ? OR (created_at = ? AND id > ?)) \ - AND type NOT IN ('cron_trigger', 'skill_suggest') \ - ORDER BY created_at ASC, id ASC \ + "SELECT m.* FROM messages m \ + INNER JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? \ + AND m.conversation_id = ? \ + AND (m.created_at > ? OR (m.created_at = ? AND m.id > ?)) \ + AND m.type NOT IN ('cron_trigger', 'skill_suggest') \ + ORDER BY m.created_at ASC, m.id ASC \ LIMIT ?", ) + .bind(user_id) .bind(conv_id) .bind(cursor.created_at) .bind(cursor.created_at) @@ -592,11 +738,14 @@ impl IConversationRepository for SqliteConversationRepository { } MessagePageDirection::Anchor { message_id } => { let anchor = sqlx::query_as::<_, MessageRow>( - "SELECT * FROM messages \ - WHERE conversation_id = ? \ - AND id = ? \ - AND type NOT IN ('cron_trigger', 'skill_suggest')", + "SELECT m.* FROM messages m \ + INNER JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? \ + AND m.conversation_id = ? \ + AND m.id = ? \ + AND m.type NOT IN ('cron_trigger', 'skill_suggest')", ) + .bind(user_id) .bind(conv_id) .bind(message_id) .fetch_optional(&self.pool) @@ -605,13 +754,16 @@ impl IConversationRepository for SqliteConversationRepository { let side_limit = limit; let mut before = sqlx::query_as::<_, MessageRow>( - "SELECT * FROM messages \ - WHERE conversation_id = ? \ - AND (created_at < ? OR (created_at = ? AND id < ?)) \ - AND type NOT IN ('cron_trigger', 'skill_suggest') \ - ORDER BY created_at DESC, id DESC \ + "SELECT m.* FROM messages m \ + INNER JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? \ + AND m.conversation_id = ? \ + AND (m.created_at < ? OR (m.created_at = ? AND m.id < ?)) \ + AND m.type NOT IN ('cron_trigger', 'skill_suggest') \ + ORDER BY m.created_at DESC, m.id DESC \ LIMIT ?", ) + .bind(user_id) .bind(conv_id) .bind(anchor.created_at) .bind(anchor.created_at) @@ -622,13 +774,16 @@ impl IConversationRepository for SqliteConversationRepository { before.reverse(); let after = sqlx::query_as::<_, MessageRow>( - "SELECT * FROM messages \ - WHERE conversation_id = ? \ - AND (created_at > ? OR (created_at = ? AND id > ?)) \ - AND type NOT IN ('cron_trigger', 'skill_suggest') \ - ORDER BY created_at ASC, id ASC \ + "SELECT m.* FROM messages m \ + INNER JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? \ + AND m.conversation_id = ? \ + AND (m.created_at > ? OR (m.created_at = ? AND m.id > ?)) \ + AND m.type NOT IN ('cron_trigger', 'skill_suggest') \ + ORDER BY m.created_at ASC, m.id ASC \ LIMIT ?", ) + .bind(user_id) .bind(conv_id) .bind(anchor.created_at) .bind(anchor.created_at) @@ -659,16 +814,19 @@ impl IConversationRepository for SqliteConversationRepository { } }; rows.sort_by(|a, b| (a.created_at, &a.id).cmp(&(b.created_at, &b.id))); - self.page_with_flags(conv_id, rows).await + self.page_with_flags(user_id, conv_id, rows).await } - async fn get_message(&self, conv_id: &str, message_id: &str) -> Result, DbError> { + async fn get_message(&self, user_id: &str, conv_id: &str, message_id: &str) -> Result, DbError> { let row = sqlx::query_as::<_, MessageRow>( - "SELECT * FROM messages \ - WHERE conversation_id = ? \ - AND id = ? \ - AND type NOT IN ('cron_trigger', 'skill_suggest')", + "SELECT m.* FROM messages m \ + INNER JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? \ + AND m.conversation_id = ? \ + AND m.id = ? \ + AND m.type NOT IN ('cron_trigger', 'skill_suggest')", ) + .bind(user_id) .bind(conv_id) .bind(message_id) .fetch_optional(&self.pool) @@ -677,15 +835,21 @@ impl IConversationRepository for SqliteConversationRepository { Ok(row) } - async fn insert_message(&self, message: &MessageRow) -> Result<(), DbError> { - self.insert_message_once(message).await.map_err(DbError::from) + async fn insert_message(&self, user_id: &str, message: &MessageRow) -> Result<(), DbError> { + self.insert_message_once(user_id, message).await } - async fn upsert_message(&self, message: &MessageRow) -> Result<(), DbError> { - self.upsert_message_once(message).await.map_err(DbError::from) + async fn upsert_message(&self, user_id: &str, message: &MessageRow) -> Result<(), DbError> { + self.upsert_message_once(user_id, message).await } - async fn update_message(&self, id: &str, updates: &MessageRowUpdate) -> Result<(), DbError> { + async fn update_message( + &self, + user_id: &str, + conversation_id: &str, + id: &str, + updates: &MessageRowUpdate, + ) -> Result<(), DbError> { let mut set_parts: Vec = Vec::new(); let mut binds: Vec = Vec::new(); @@ -706,13 +870,23 @@ impl IConversationRepository for SqliteConversationRepository { return Ok(()); } - let sql = format!("UPDATE messages SET {} WHERE id = ?", set_parts.join(", ")); + let sql = format!( + "UPDATE messages SET {} \ + WHERE conversation_id = ? AND id = ? \ + AND EXISTS ( \ + SELECT 1 FROM conversations c \ + WHERE c.id = messages.conversation_id AND c.user_id = ? \ + )", + set_parts.join(", ") + ); let mut query = sqlx::query(&sql); for bind in &binds { query = bind_value(query, bind); } + query = query.bind(conversation_id); query = query.bind(id); + query = query.bind(user_id); let result = query.execute(&self.pool).await?; @@ -723,25 +897,36 @@ impl IConversationRepository for SqliteConversationRepository { Ok(()) } - async fn delete_messages_by_conversation(&self, conv_id: &str) -> Result<(), DbError> { - sqlx::query("DELETE FROM messages WHERE conversation_id = ?") - .bind(conv_id) - .execute(&self.pool) - .await?; + async fn delete_messages_by_conversation(&self, user_id: &str, conv_id: &str) -> Result<(), DbError> { + sqlx::query( + "DELETE FROM messages \ + WHERE conversation_id = ? \ + AND EXISTS ( \ + SELECT 1 FROM conversations c \ + WHERE c.id = messages.conversation_id AND c.user_id = ? \ + )", + ) + .bind(conv_id) + .bind(user_id) + .execute(&self.pool) + .await?; Ok(()) } async fn get_message_by_msg_id( &self, + user_id: &str, conv_id: &str, msg_id: &str, msg_type: &str, ) -> Result, DbError> { let row = sqlx::query_as::<_, MessageRow>( - "SELECT * FROM messages \ - WHERE conversation_id = ? AND msg_id = ? AND type = ?", + "SELECT m.* FROM messages m \ + INNER JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? AND m.conversation_id = ? AND m.msg_id = ? AND m.type = ?", ) + .bind(user_id) .bind(conv_id) .bind(msg_id) .bind(msg_type) @@ -751,9 +936,9 @@ impl IConversationRepository for SqliteConversationRepository { Ok(row) } - async fn list_stale_runtime_messages(&self) -> Result, DbError> { - let rows = sqlx::query_as::<_, MessageRow>( - "SELECT m.* FROM messages m \ + async fn list_stale_runtime_messages(&self) -> Result, DbError> { + let rows = sqlx::query( + "SELECT c.user_id, m.* FROM messages m \ INNER JOIN conversations c ON c.id = m.conversation_id \ WHERE m.position = 'left' \ AND m.status IN ('work', 'pending') \ @@ -763,7 +948,25 @@ impl IConversationRepository for SqliteConversationRepository { .fetch_all(&self.pool) .await?; - Ok(rows) + Ok(rows + .into_iter() + .map(|row| { + Ok(StaleRuntimeMessageRow { + user_id: row.try_get("user_id")?, + message: MessageRow { + id: row.try_get("id")?, + conversation_id: row.try_get("conversation_id")?, + msg_id: row.try_get("msg_id")?, + r#type: row.try_get("type")?, + content: row.try_get("content")?, + position: row.try_get("position")?, + status: row.try_get("status")?, + hidden: row.try_get("hidden")?, + created_at: row.try_get("created_at")?, + }, + }) + }) + .collect::, sqlx::Error>>()?) } async fn search_messages( @@ -832,12 +1035,18 @@ impl IConversationRepository for SqliteConversationRepository { Ok(PaginatedResult { items, total, has_more }) } - async fn list_artifacts(&self, conversation_id: &str) -> Result, DbError> { + async fn list_artifacts( + &self, + user_id: &str, + conversation_id: &str, + ) -> Result, DbError> { let rows = sqlx::query_as::<_, ConversationArtifactRow>( - "SELECT * FROM conversation_artifacts \ - WHERE conversation_id = ? \ - ORDER BY created_at ASC, id ASC", + "SELECT a.* FROM conversation_artifacts a \ + INNER JOIN conversations c ON c.id = a.conversation_id \ + WHERE c.user_id = ? AND a.conversation_id = ? \ + ORDER BY a.created_at ASC, a.id ASC", ) + .bind(user_id) .bind(conversation_id) .fetch_all(&self.pool) .await?; @@ -847,12 +1056,16 @@ impl IConversationRepository for SqliteConversationRepository { async fn get_artifact( &self, + user_id: &str, conversation_id: &str, artifact_id: &str, ) -> Result, DbError> { let row = sqlx::query_as::<_, ConversationArtifactRow>( - "SELECT * FROM conversation_artifacts WHERE conversation_id = ? AND id = ?", + "SELECT a.* FROM conversation_artifacts a \ + INNER JOIN conversations c ON c.id = a.conversation_id \ + WHERE c.user_id = ? AND a.conversation_id = ? AND a.id = ?", ) + .bind(user_id) .bind(conversation_id) .bind(artifact_id) .fetch_optional(&self.pool) @@ -861,8 +1074,14 @@ impl IConversationRepository for SqliteConversationRepository { Ok(row) } - async fn upsert_artifact(&self, artifact: &ConversationArtifactRow) -> Result { - sqlx::query( + async fn upsert_artifact( + &self, + user_id: &str, + artifact: &ConversationArtifactRow, + ) -> Result { + self.ensure_conversation_for_user(user_id, &artifact.conversation_id) + .await?; + let result = sqlx::query( "INSERT INTO conversation_artifacts \ (id, conversation_id, cron_job_id, kind, status, payload, created_at, updated_at) \ VALUES (?, ?, ?, ?, ?, ?, ?, ?) \ @@ -872,7 +1091,12 @@ impl IConversationRepository for SqliteConversationRepository { kind = excluded.kind, \ status = excluded.status, \ payload = excluded.payload, \ - updated_at = excluded.updated_at", + updated_at = excluded.updated_at \ + WHERE conversation_artifacts.conversation_id = excluded.conversation_id \ + AND EXISTS ( \ + SELECT 1 FROM conversations c \ + WHERE c.id = conversation_artifacts.conversation_id AND c.user_id = ? \ + )", ) .bind(&artifact.id) .bind(&artifact.conversation_id) @@ -882,16 +1106,25 @@ impl IConversationRepository for SqliteConversationRepository { .bind(&artifact.payload) .bind(artifact.created_at) .bind(artifact.updated_at) + .bind(user_id) .execute(&self.pool) .await?; - self.get_artifact(&artifact.conversation_id, &artifact.id) + if result.rows_affected() == 0 { + return Err(DbError::Conflict(format!( + "Conversation artifact with id '{}' already exists outside the requested conversation", + artifact.id + ))); + } + + self.get_artifact(user_id, &artifact.conversation_id, &artifact.id) .await? .ok_or_else(|| DbError::Init(format!("upsert artifact did not produce row for id '{}'", artifact.id))) } async fn update_artifact_status( &self, + user_id: &str, conversation_id: &str, artifact_id: &str, status: &str, @@ -900,12 +1133,17 @@ impl IConversationRepository for SqliteConversationRepository { let result = sqlx::query( "UPDATE conversation_artifacts \ SET status = ?, updated_at = ? \ - WHERE conversation_id = ? AND id = ?", + WHERE conversation_id = ? AND id = ? \ + AND EXISTS ( \ + SELECT 1 FROM conversations c \ + WHERE c.id = conversation_artifacts.conversation_id AND c.user_id = ? \ + )", ) .bind(status) .bind(updated_at) .bind(conversation_id) .bind(artifact_id) + .bind(user_id) .execute(&self.pool) .await?; @@ -913,29 +1151,37 @@ impl IConversationRepository for SqliteConversationRepository { return Ok(None); } - self.get_artifact(conversation_id, artifact_id).await + self.get_artifact(user_id, conversation_id, artifact_id).await } async fn mark_skill_suggest_artifacts_saved( &self, + user_id: &str, cron_job_id: &str, updated_at: i64, ) -> Result, DbError> { sqlx::query( "UPDATE conversation_artifacts \ SET status = 'saved', updated_at = ? \ - WHERE kind = 'skill_suggest' AND cron_job_id = ? AND status != 'saved'", + WHERE kind = 'skill_suggest' AND cron_job_id = ? AND status != 'saved' \ + AND EXISTS ( \ + SELECT 1 FROM conversations c \ + WHERE c.id = conversation_artifacts.conversation_id AND c.user_id = ? \ + )", ) .bind(updated_at) .bind(cron_job_id) + .bind(user_id) .execute(&self.pool) .await?; let rows = sqlx::query_as::<_, ConversationArtifactRow>( - "SELECT * FROM conversation_artifacts \ - WHERE kind = 'skill_suggest' AND cron_job_id = ? \ - ORDER BY created_at ASC, id ASC", + "SELECT a.* FROM conversation_artifacts a \ + INNER JOIN conversations c ON c.id = a.conversation_id \ + WHERE c.user_id = ? AND a.kind = 'skill_suggest' AND a.cron_job_id = ? \ + ORDER BY a.created_at ASC, a.id ASC", ) + .bind(user_id) .bind(cron_job_id) .fetch_all(&self.pool) .await?; @@ -943,21 +1189,35 @@ impl IConversationRepository for SqliteConversationRepository { Ok(rows) } - async fn delete_artifacts_by_conversation(&self, conversation_id: &str) -> Result<(), DbError> { - sqlx::query("DELETE FROM conversation_artifacts WHERE conversation_id = ?") - .bind(conversation_id) - .execute(&self.pool) - .await?; + async fn delete_artifacts_by_conversation(&self, user_id: &str, conversation_id: &str) -> Result<(), DbError> { + sqlx::query( + "DELETE FROM conversation_artifacts \ + WHERE conversation_id = ? \ + AND EXISTS ( \ + SELECT 1 FROM conversations c \ + WHERE c.id = conversation_artifacts.conversation_id AND c.user_id = ? \ + )", + ) + .bind(conversation_id) + .bind(user_id) + .execute(&self.pool) + .await?; Ok(()) } - async fn list_legacy_cron_trigger_messages(&self, conversation_id: &str) -> Result, DbError> { + async fn list_legacy_cron_trigger_messages( + &self, + user_id: &str, + conversation_id: &str, + ) -> Result, DbError> { let rows = sqlx::query_as::<_, MessageRow>( - "SELECT * FROM messages \ - WHERE conversation_id = ? AND type = 'cron_trigger' \ - ORDER BY created_at ASC, id ASC", + "SELECT m.* FROM messages m \ + INNER JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? AND m.conversation_id = ? AND m.type = 'cron_trigger' \ + ORDER BY m.created_at ASC, m.id ASC", ) + .bind(user_id) .bind(conversation_id) .fetch_all(&self.pool) .await?; @@ -1107,7 +1367,7 @@ mod tests { let conv = sample_conversation(SYSTEM_USER_ID); repo.create(&conv).await.unwrap(); - let found = repo.get(&conv.id).await.unwrap().unwrap(); + let found = repo.get(&conv.user_id, &conv.id).await.unwrap().unwrap(); assert_eq!(found.id, conv.id); assert_eq!(found.name, "Test Conversation"); @@ -1119,7 +1379,7 @@ mod tests { #[tokio::test] async fn get_nonexistent_returns_none() { let (repo, _db) = setup().await; - assert!(repo.get("no_such_id").await.unwrap().is_none()); + assert!(repo.get("user_1", "no_such_id").await.unwrap().is_none()); } /// Persisting a message must bump the parent conversation's updated_at so the @@ -1136,9 +1396,9 @@ mod tests { // insert_message with a newer event time bumps updated_at forward. let mut msg = sample_message(&conv.id); msg.created_at = 5_000; - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(SYSTEM_USER_ID, &msg).await.unwrap(); assert_eq!( - repo.get(&conv.id).await.unwrap().unwrap().updated_at, + repo.get(SYSTEM_USER_ID, &conv.id).await.unwrap().unwrap().updated_at, 5_000, "insert must bump updated_at to the message time" ); @@ -1147,9 +1407,9 @@ mod tests { let mut newer = sample_message(&conv.id); newer.id = "tool-1".to_string(); newer.created_at = 9_000; - repo.upsert_message(&newer).await.unwrap(); + repo.upsert_message(SYSTEM_USER_ID, &newer).await.unwrap(); assert_eq!( - repo.get(&conv.id).await.unwrap().unwrap().updated_at, + repo.get(SYSTEM_USER_ID, &conv.id).await.unwrap().unwrap().updated_at, 9_000, "newer upsert must advance updated_at" ); @@ -1158,9 +1418,9 @@ mod tests { let mut older = sample_message(&conv.id); older.id = "tool-2".to_string(); older.created_at = 3_000; - repo.upsert_message(&older).await.unwrap(); + repo.upsert_message(SYSTEM_USER_ID, &older).await.unwrap(); assert_eq!( - repo.get(&conv.id).await.unwrap().unwrap().updated_at, + repo.get(SYSTEM_USER_ID, &conv.id).await.unwrap().unwrap().updated_at, 9_000, "older upsert must not move updated_at backward" ); @@ -1174,6 +1434,7 @@ mod tests { let now = aionui_common::now_ms(); repo.update( + &conv.user_id, &conv.id, &ConversationRowUpdate { name: Some("Updated Name".to_string()), @@ -1184,7 +1445,7 @@ mod tests { .await .unwrap(); - let found = repo.get(&conv.id).await.unwrap().unwrap(); + let found = repo.get(&conv.user_id, &conv.id).await.unwrap().unwrap(); assert_eq!(found.name, "Updated Name"); assert!(found.updated_at >= conv.updated_at); } @@ -1197,6 +1458,7 @@ mod tests { let pin_time = aionui_common::now_ms(); repo.update( + &conv.user_id, &conv.id, &ConversationRowUpdate { pinned: Some(true), @@ -1208,7 +1470,7 @@ mod tests { .await .unwrap(); - let found = repo.get(&conv.id).await.unwrap().unwrap(); + let found = repo.get(&conv.user_id, &conv.id).await.unwrap().unwrap(); assert!(found.pinned); assert_eq!(found.pinned_at, Some(pin_time)); } @@ -1218,6 +1480,7 @@ mod tests { let (repo, _db) = setup().await; let err = repo .update( + "user_1", "no_id", &ConversationRowUpdate { name: Some("x".to_string()), @@ -1236,7 +1499,9 @@ mod tests { repo.create(&conv).await.unwrap(); // Empty update should succeed without error - repo.update(&conv.id, &ConversationRowUpdate::default()).await.unwrap(); + repo.update(&conv.user_id, &conv.id, &ConversationRowUpdate::default()) + .await + .unwrap(); } #[tokio::test] @@ -1245,39 +1510,33 @@ mod tests { let conv = sample_conversation(SYSTEM_USER_ID); repo.create(&conv).await.unwrap(); - repo.delete(&conv.id).await.unwrap(); - assert!(repo.get(&conv.id).await.unwrap().is_none()); + repo.delete(&conv.user_id, &conv.id).await.unwrap(); + assert!(repo.get(&conv.user_id, &conv.id).await.unwrap().is_none()); } #[tokio::test] async fn delete_cascades_messages() { - let (repo, _db) = setup().await; + let (repo, db) = setup().await; let conv = sample_conversation(SYSTEM_USER_ID); repo.create(&conv).await.unwrap(); let msg = sample_message(&conv.id); - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); - repo.delete(&conv.id).await.unwrap(); + repo.delete(&conv.user_id, &conv.id).await.unwrap(); - // Messages should be gone due to CASCADE - let result = repo - .list_messages_page( - &conv.id, - &MessagePageParams { - limit: 50, - direction: MessagePageDirection::InitialLatest, - }, - ) + let remaining: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM messages WHERE conversation_id = ?") + .bind(&conv.id) + .fetch_one(db.pool()) .await .unwrap(); - assert!(result.items.is_empty()); + assert_eq!(remaining, 0); } #[tokio::test] async fn delete_nonexistent_returns_not_found() { let (repo, _db) = setup().await; - let err = repo.delete("no_id").await.unwrap_err(); + let err = repo.delete("user_1", "no_id").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -1591,10 +1850,11 @@ mod tests { repo.create(&conv).await.unwrap(); let msg = sample_message(&conv.id); - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); let result = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 50, @@ -1618,11 +1878,12 @@ mod tests { let mut msg = sample_message(&conv.id); msg.id = aionui_common::generate_prefixed_id("msg"); msg.created_at = (i + 1) * 1000; - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); } let page1 = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 3, @@ -1650,11 +1911,12 @@ mod tests { let mut msg = sample_message(&conv.id); msg.id = aionui_common::generate_prefixed_id("msg"); msg.created_at = (i + 1) * 1000; - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); } let latest = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 3, @@ -1665,6 +1927,7 @@ mod tests { .unwrap(); let older = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 3, @@ -1691,9 +1954,11 @@ mod tests { repo.create(&conv).await.unwrap(); let msg = sample_message(&conv.id); - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); repo.update_message( + &conv.user_id, + &conv.id, &msg.id, &MessageRowUpdate { content: Some(r#"{"content":"Updated"}"#.to_string()), @@ -1705,6 +1970,7 @@ mod tests { let result = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 50, @@ -1721,6 +1987,8 @@ mod tests { let (repo, _db) = setup().await; let err = repo .update_message( + "user_1", + "conv_1", "no_id", &MessageRowUpdate { hidden: Some(true), @@ -1741,13 +2009,16 @@ mod tests { for _ in 0..3 { let mut msg = sample_message(&conv.id); msg.id = aionui_common::generate_prefixed_id("msg"); - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); } - repo.delete_messages_by_conversation(&conv.id).await.unwrap(); + repo.delete_messages_by_conversation(&conv.user_id, &conv.id) + .await + .unwrap(); let result = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 50, @@ -1766,10 +2037,10 @@ mod tests { repo.create(&conv).await.unwrap(); let msg = sample_message(&conv.id); - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); let found = repo - .get_message_by_msg_id(&conv.id, "client_msg_1", "text") + .get_message_by_msg_id(&conv.user_id, &conv.id, "client_msg_1", "text") .await .unwrap(); assert!(found.is_some()); @@ -1777,7 +2048,7 @@ mod tests { // Wrong type → not found let not_found = repo - .get_message_by_msg_id(&conv.id, "client_msg_1", "tips") + .get_message_by_msg_id(&conv.user_id, &conv.id, "client_msg_1", "tips") .await .unwrap(); assert!(not_found.is_none()); @@ -1791,12 +2062,12 @@ mod tests { let mut msg1 = sample_message(&conv.id); msg1.content = r#"{"content":"Rust 审查报告"}"#.to_string(); - repo.insert_message(&msg1).await.unwrap(); + repo.insert_message(&conv.user_id, &msg1).await.unwrap(); let mut msg2 = sample_message(&conv.id); msg2.id = aionui_common::generate_prefixed_id("msg"); msg2.content = r#"{"content":"Python 测试"}"#.to_string(); - repo.insert_message(&msg2).await.unwrap(); + repo.insert_message(&conv.user_id, &msg2).await.unwrap(); let result = repo.search_messages(SYSTEM_USER_ID, "审查", 1, 20).await.unwrap(); assert_eq!(result.items.len(), 1); @@ -1811,7 +2082,7 @@ mod tests { repo.create(&conv).await.unwrap(); let msg = sample_message(&conv.id); - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); let result = repo .search_messages(SYSTEM_USER_ID, "xxxxnotexist", 1, 20) @@ -1832,7 +2103,7 @@ mod tests { msg.id = aionui_common::generate_prefixed_id("msg"); msg.content = format!(r#"{{"content":"match keyword item {i}"}}"#); msg.created_at = (i + 1) * 1000; - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); } let result = repo.search_messages(SYSTEM_USER_ID, "keyword", 1, 2).await.unwrap(); diff --git a/crates/aionui-db/src/repository/sqlite_cron.rs b/crates/aionui-db/src/repository/sqlite_cron.rs index 916e133da..5c6b87d16 100644 --- a/crates/aionui-db/src/repository/sqlite_cron.rs +++ b/crates/aionui-db/src/repository/sqlite_cron.rs @@ -22,18 +22,29 @@ impl SqliteCronRepository { #[async_trait::async_trait] impl ICronRepository for SqliteCronRepository { async fn insert(&self, row: &CronJobRow) -> Result<(), DbError> { + // Existing-mode jobs bind to a live conversation, so the referenced + // conversation must belong to the job owner. New-conversation jobs + // keep a stale anchor id that is re-resolved at run time and stays + // unanchored until then; the executor re-validates ownership before + // dispatch. + if row.execution_mode == "existing" { + self.validate_conversation_owner(&row.user_id, &row.conversation_id) + .await?; + } + sqlx::query( "INSERT INTO cron_jobs (\ - id, name, enabled, schedule_kind, schedule_value, schedule_tz, \ + id, user_id, name, enabled, schedule_kind, schedule_value, schedule_tz, \ schedule_description, payload_message, execution_mode, agent_config, \ conversation_id, conversation_title, created_by, \ skill_content, description, created_at, updated_at, next_run_at, last_run_at, \ last_status, last_error, run_count, retry_count, max_retries, queue_enabled\ ) VALUES (\ - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?\ + ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?\ )", ) .bind(&row.id) + .bind(&row.user_id) .bind(&row.name) .bind(row.enabled) .bind(&row.schedule_kind) @@ -63,108 +74,42 @@ impl ICronRepository for SqliteCronRepository { Ok(()) } - async fn update(&self, id: &str, params: &UpdateCronJobParams) -> Result<(), DbError> { - let mut set_parts: Vec = Vec::new(); - let mut binds: Vec = Vec::new(); - - macro_rules! push_str { - ($field:ident) => { - if let Some(ref v) = params.$field { - set_parts.push(concat!(stringify!($field), " = ?").to_string()); - binds.push(BindValue::Str(v.clone())); - } - }; - } - - macro_rules! push_opt_str { - ($field:ident) => { - if let Some(ref v) = params.$field { - set_parts.push(concat!(stringify!($field), " = ?").to_string()); - binds.push(BindValue::OptStr(v.clone())); - } - }; - } - - macro_rules! push_opt_i64 { - ($field:ident) => { - if let Some(ref v) = params.$field { - set_parts.push(concat!(stringify!($field), " = ?").to_string()); - binds.push(BindValue::OptI64(*v)); - } - }; - } - - macro_rules! push_i64 { - ($field:ident) => { - if let Some(v) = params.$field { - set_parts.push(concat!(stringify!($field), " = ?").to_string()); - binds.push(BindValue::I64(v)); - } - }; - } - - if let Some(v) = params.enabled { - set_parts.push("enabled = ?".to_string()); - binds.push(BindValue::Bool(v)); - } - if let Some(v) = params.queue_enabled { - set_parts.push("queue_enabled = ?".to_string()); - binds.push(BindValue::Bool(v)); - } - - push_str!(name); - push_str!(schedule_kind); - push_str!(schedule_value); - push_opt_str!(schedule_tz); - push_opt_str!(schedule_description); - push_str!(payload_message); - push_str!(execution_mode); - push_opt_str!(agent_config); - push_str!(conversation_id); - push_opt_str!(conversation_title); - push_opt_str!(skill_content); - push_opt_str!(description); - push_opt_i64!(next_run_at); - push_opt_i64!(last_run_at); - push_opt_str!(last_status); - push_opt_str!(last_error); - push_i64!(run_count); - push_i64!(retry_count); - - if set_parts.is_empty() { - return Ok(()); - } - - set_parts.push("updated_at = ?".to_string()); - binds.push(BindValue::I64(now_ms())); - - let sql = format!("UPDATE cron_jobs SET {} WHERE id = ?", set_parts.join(", ")); + async fn update_system(&self, id: &str, params: &UpdateCronJobParams) -> Result<(), DbError> { + self.update_inner(None, id, params).await + } - let mut query = sqlx::query(&sql); - for bind in &binds { - query = bind_value(query, bind); - } - query = query.bind(id); + async fn update_for_user(&self, user_id: &str, id: &str, params: &UpdateCronJobParams) -> Result<(), DbError> { + self.update_inner(Some(user_id), id, params).await + } - let result = query.execute(&self.pool).await?; + async fn delete_system(&self, id: &str) -> Result<(), DbError> { + let result = sqlx::query("DELETE FROM cron_jobs WHERE id = ?") + .bind(id) + .execute(&self.pool) + .await?; if result.rows_affected() == 0 { return Err(DbError::NotFound(format!("cron job '{id}'"))); } Ok(()) } - async fn delete(&self, id: &str) -> Result<(), DbError> { - let result = sqlx::query("DELETE FROM cron_jobs WHERE id = ?") - .bind(id) - .execute(&self.pool) - .await?; + async fn delete_for_user(&self, user_id: &str, id: &str) -> Result<(), DbError> { + let result = sqlx::query( + "DELETE FROM cron_jobs \ + WHERE id = ? \ + AND user_id = ?", + ) + .bind(id) + .bind(user_id) + .execute(&self.pool) + .await?; if result.rows_affected() == 0 { return Err(DbError::NotFound(format!("cron job '{id}'"))); } Ok(()) } - async fn get_by_id(&self, id: &str) -> Result, DbError> { + async fn get_by_id_system(&self, id: &str) -> Result, DbError> { let row = sqlx::query_as::<_, CronJobRow>("SELECT * FROM cron_jobs WHERE id = ?") .bind(id) .fetch_optional(&self.pool) @@ -172,21 +117,38 @@ impl ICronRepository for SqliteCronRepository { Ok(row) } - async fn list_all(&self) -> Result, DbError> { + async fn get_by_id_for_user(&self, user_id: &str, id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, CronJobRow>("SELECT * FROM cron_jobs WHERE user_id = ? AND id = ?") + .bind(user_id) + .bind(id) + .fetch_optional(&self.pool) + .await?; + Ok(row) + } + + async fn list_all_system(&self) -> Result, DbError> { let rows = sqlx::query_as::<_, CronJobRow>("SELECT * FROM cron_jobs ORDER BY created_at ASC") .fetch_all(&self.pool) .await?; Ok(rows) } - async fn list_enabled(&self) -> Result, DbError> { + async fn list_all_for_user(&self, user_id: &str) -> Result, DbError> { + let rows = sqlx::query_as::<_, CronJobRow>("SELECT * FROM cron_jobs WHERE user_id = ? ORDER BY created_at ASC") + .bind(user_id) + .fetch_all(&self.pool) + .await?; + Ok(rows) + } + + async fn list_enabled_system(&self) -> Result, DbError> { let rows = sqlx::query_as::<_, CronJobRow>("SELECT * FROM cron_jobs WHERE enabled = 1 ORDER BY created_at ASC") .fetch_all(&self.pool) .await?; Ok(rows) } - async fn list_by_conversation(&self, conversation_id: &str) -> Result, DbError> { + async fn list_by_conversation_system(&self, conversation_id: &str) -> Result, DbError> { let rows = sqlx::query_as::<_, CronJobRow>( "SELECT * FROM cron_jobs WHERE conversation_id = ? ORDER BY created_at ASC", ) @@ -196,7 +158,24 @@ impl ICronRepository for SqliteCronRepository { Ok(rows) } - async fn delete_by_conversation(&self, conversation_id: &str) -> Result { + async fn list_by_conversation_for_user( + &self, + user_id: &str, + conversation_id: &str, + ) -> Result, DbError> { + let rows = sqlx::query_as::<_, CronJobRow>( + "SELECT * FROM cron_jobs \ + WHERE user_id = ? AND conversation_id = ? \ + ORDER BY created_at ASC", + ) + .bind(user_id) + .bind(conversation_id) + .fetch_all(&self.pool) + .await?; + Ok(rows) + } + + async fn delete_by_conversation_system(&self, conversation_id: &str) -> Result { let result = sqlx::query("DELETE FROM cron_jobs WHERE conversation_id = ?") .bind(conversation_id) .execute(&self.pool) @@ -209,6 +188,19 @@ impl ICronRepository for SqliteCronRepository { sqlx::query("BEGIN IMMEDIATE").execute(&mut *connection).await?; let result = async { + let job_has_owner: bool = sqlx::query_scalar( + "SELECT EXISTS(\ + SELECT 1 FROM cron_jobs \ + WHERE id = ? AND user_id IS NOT NULL AND user_id != ''\ + )", + ) + .bind(params.job_id) + .fetch_one(&mut *connection) + .await?; + if !job_has_owner { + return Ok(CronRunClaimResult::Duplicate); + } + let existing = sqlx::query_as::<_, (String, Option, Option)>( "SELECT status, owner_id, lease_until FROM cron_job_runs WHERE job_id = ? AND scheduled_at = ?", ) @@ -312,7 +304,11 @@ impl ICronRepository for SqliteCronRepository { ) -> Result { let result = sqlx::query( "UPDATE cron_job_runs SET lease_until = ?, updated_at = ? \ - WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ?", + WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ? \ + AND EXISTS (\ + SELECT 1 FROM cron_jobs j \ + WHERE j.id = cron_job_runs.job_id AND j.user_id IS NOT NULL AND j.user_id != ''\ + )", ) .bind(lease_until) .bind(updated_at) @@ -334,7 +330,11 @@ impl ICronRepository for SqliteCronRepository { ) -> Result { let result = sqlx::query( "UPDATE cron_job_runs SET status = 'retrying', owner_id = NULL, lease_until = ?, updated_at = ? \ - WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ?", + WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ? \ + AND EXISTS (\ + SELECT 1 FROM cron_jobs j \ + WHERE j.id = cron_job_runs.job_id AND j.user_id IS NOT NULL AND j.user_id != ''\ + )", ) .bind(retry_at) .bind(updated_at) @@ -350,7 +350,11 @@ impl ICronRepository for SqliteCronRepository { let result = sqlx::query( "UPDATE cron_job_runs SET status = ?, conversation_id = ?, error = ?, lease_until = NULL, \ finished_at = ?, updated_at = ? \ - WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ?", + WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ? \ + AND EXISTS (\ + SELECT 1 FROM cron_jobs j \ + WHERE j.id = cron_job_runs.job_id AND j.user_id IS NOT NULL AND j.user_id != ''\ + )", ) .bind(params.status) .bind(params.conversation_id) @@ -376,9 +380,11 @@ impl ICronRepository for SqliteCronRepository { async fn get_recoverable_run(&self, job_id: &str, now: TimestampMs) -> Result, DbError> { let row = sqlx::query_as::<_, (TimestampMs, TimestampMs)>( - "SELECT scheduled_at, MAX(lease_until, ?) AS wake_at FROM cron_job_runs \ - WHERE job_id = ? AND status IN ('running', 'retrying') AND lease_until IS NOT NULL \ - ORDER BY lease_until ASC LIMIT 1", + "SELECT r.scheduled_at, MAX(r.lease_until, ?) AS wake_at FROM cron_job_runs r \ + JOIN cron_jobs j ON j.id = r.job_id \ + WHERE r.job_id = ? AND j.user_id IS NOT NULL AND j.user_id != '' \ + AND r.status IN ('running', 'retrying') AND r.lease_until IS NOT NULL \ + ORDER BY r.lease_until ASC LIMIT 1", ) .bind(now) .bind(job_id) @@ -389,6 +395,131 @@ impl ICronRepository for SqliteCronRepository { } } +impl SqliteCronRepository { + async fn validate_conversation_owner(&self, user_id: &str, conversation_id: &str) -> Result<(), DbError> { + if conversation_id.trim().is_empty() { + return Ok(()); + } + + let owns_conversation: bool = + sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM conversations WHERE user_id = ? AND id = ?)") + .bind(user_id) + .bind(conversation_id) + .fetch_one(&self.pool) + .await?; + if !owns_conversation { + return Err(DbError::NotFound(format!("conversation '{conversation_id}'"))); + } + Ok(()) + } + + async fn update_inner(&self, user_id: Option<&str>, id: &str, params: &UpdateCronJobParams) -> Result<(), DbError> { + let mut set_parts: Vec = Vec::new(); + let mut binds: Vec = Vec::new(); + + macro_rules! push_str { + ($field:ident) => { + if let Some(ref v) = params.$field { + set_parts.push(concat!(stringify!($field), " = ?").to_string()); + binds.push(BindValue::Str(v.clone())); + } + }; + } + + macro_rules! push_opt_str { + ($field:ident) => { + if let Some(ref v) = params.$field { + set_parts.push(concat!(stringify!($field), " = ?").to_string()); + binds.push(BindValue::OptStr(v.clone())); + } + }; + } + + macro_rules! push_opt_i64 { + ($field:ident) => { + if let Some(ref v) = params.$field { + set_parts.push(concat!(stringify!($field), " = ?").to_string()); + binds.push(BindValue::OptI64(*v)); + } + }; + } + + macro_rules! push_i64 { + ($field:ident) => { + if let Some(v) = params.$field { + set_parts.push(concat!(stringify!($field), " = ?").to_string()); + binds.push(BindValue::I64(v)); + } + }; + } + + if let Some(v) = params.enabled { + set_parts.push("enabled = ?".to_string()); + binds.push(BindValue::Bool(v)); + } + if let Some(v) = params.queue_enabled { + set_parts.push("queue_enabled = ?".to_string()); + binds.push(BindValue::Bool(v)); + } + + push_str!(name); + push_str!(schedule_kind); + push_str!(schedule_value); + push_opt_str!(schedule_tz); + push_opt_str!(schedule_description); + push_str!(payload_message); + push_str!(execution_mode); + push_opt_str!(agent_config); + push_str!(conversation_id); + push_opt_str!(conversation_title); + push_opt_str!(skill_content); + push_opt_str!(description); + push_opt_i64!(next_run_at); + push_opt_i64!(last_run_at); + push_opt_str!(last_status); + push_opt_str!(last_error); + push_i64!(run_count); + push_i64!(retry_count); + + if set_parts.is_empty() { + return Ok(()); + } + + set_parts.push("updated_at = ?".to_string()); + binds.push(BindValue::I64(now_ms())); + + if let (Some(user_id), Some(conversation_id)) = (user_id, params.conversation_id.as_deref()) + && !conversation_id.trim().is_empty() + { + self.validate_conversation_owner(user_id, conversation_id).await?; + } + + let sql = if user_id.is_some() { + format!( + "UPDATE cron_jobs SET {} WHERE id = ? AND user_id = ?", + set_parts.join(", ") + ) + } else { + format!("UPDATE cron_jobs SET {} WHERE id = ?", set_parts.join(", ")) + }; + + let mut query = sqlx::query(&sql); + for bind in &binds { + query = bind_value(query, bind); + } + query = query.bind(id); + if let Some(user_id) = user_id { + query = query.bind(user_id); + } + + let result = query.execute(&self.pool).await?; + if result.rows_affected() == 0 { + return Err(DbError::NotFound(format!("cron job '{id}'"))); + } + Ok(()) + } +} + // ── Dynamic bind helpers ──────────────────────────────────────────── #[derive(Debug, Clone)] @@ -445,6 +576,7 @@ mod tests { let now = now_ms(); CronJobRow { id: id.into(), + user_id: "user_1".into(), name: "Test Job".into(), enabled: true, schedule_kind: "every".into(), @@ -478,7 +610,7 @@ mod tests { let row = make_row("cron_1"); repo.insert(&row).await.unwrap(); - let found = repo.get_by_id("cron_1").await.unwrap().expect("found"); + let found = repo.get_by_id_system("cron_1").await.unwrap().expect("found"); assert_eq!(found.id, "cron_1"); assert_eq!(found.name, "Test Job"); assert!(found.enabled); @@ -489,7 +621,7 @@ mod tests { #[tokio::test] async fn get_by_id_returns_none_for_missing() { let (repo, _db) = setup().await; - let result = repo.get_by_id("cron_missing").await.unwrap(); + let result = repo.get_by_id_system("cron_missing").await.unwrap(); assert!(result.is_none()); } @@ -499,7 +631,7 @@ mod tests { repo.insert(&make_row("cron_a")).await.unwrap(); repo.insert(&make_row("cron_b")).await.unwrap(); - let all = repo.list_all().await.unwrap(); + let all = repo.list_all_system().await.unwrap(); assert_eq!(all.len(), 2); } @@ -512,7 +644,7 @@ mod tests { disabled.enabled = false; repo.insert(&disabled).await.unwrap(); - let enabled = repo.list_enabled().await.unwrap(); + let enabled = repo.list_enabled_system().await.unwrap(); assert_eq!(enabled.len(), 1); assert_eq!(enabled[0].id, "cron_e1"); } @@ -533,11 +665,11 @@ mod tests { other.conversation_id = "conv_2".into(); repo.insert(&other).await.unwrap(); - let conv1_jobs = repo.list_by_conversation("conv_1").await.unwrap(); + let conv1_jobs = repo.list_by_conversation_system("conv_1").await.unwrap(); assert_eq!(conv1_jobs.len(), 1); assert_eq!(conv1_jobs[0].id, "cron_c1"); - let conv2_jobs = repo.list_by_conversation("conv_2").await.unwrap(); + let conv2_jobs = repo.list_by_conversation_system("conv_2").await.unwrap(); assert_eq!(conv2_jobs.len(), 1); assert_eq!(conv2_jobs[0].id, "cron_c2"); } @@ -553,9 +685,9 @@ mod tests { run_count: Some(42), ..Default::default() }; - repo.update("cron_u1", ¶ms).await.unwrap(); + repo.update_system("cron_u1", ¶ms).await.unwrap(); - let updated = repo.get_by_id("cron_u1").await.unwrap().unwrap(); + let updated = repo.get_by_id_system("cron_u1").await.unwrap().unwrap(); assert_eq!(updated.name, "Renamed"); assert!(!updated.enabled); assert_eq!(updated.run_count, 42); @@ -567,7 +699,7 @@ mod tests { let (repo, _db) = setup().await; repo.insert(&make_row("cron_queue")).await.unwrap(); - repo.update( + repo.update_system( "cron_queue", &UpdateCronJobParams { queue_enabled: Some(true), @@ -577,7 +709,13 @@ mod tests { .await .unwrap(); - assert!(repo.get_by_id("cron_queue").await.unwrap().unwrap().queue_enabled); + assert!( + repo.get_by_id_system("cron_queue") + .await + .unwrap() + .unwrap() + .queue_enabled + ); } #[tokio::test] @@ -842,9 +980,9 @@ mod tests { skill_content: Some(Some("---\nname: skill\n---\nDo it".into())), ..Default::default() }; - repo.update("cron_u2", ¶ms).await.unwrap(); + repo.update_system("cron_u2", ¶ms).await.unwrap(); - let updated = repo.get_by_id("cron_u2").await.unwrap().unwrap(); + let updated = repo.get_by_id_system("cron_u2").await.unwrap().unwrap(); assert_eq!(updated.last_status.as_deref(), Some("ok")); assert_eq!(updated.last_error.as_deref(), Some("timeout")); assert!(updated.skill_content.is_some()); @@ -855,9 +993,9 @@ mod tests { skill_content: Some(None), ..Default::default() }; - repo.update("cron_u2", &clear_params).await.unwrap(); + repo.update_system("cron_u2", &clear_params).await.unwrap(); - let cleared = repo.get_by_id("cron_u2").await.unwrap().unwrap(); + let cleared = repo.get_by_id_system("cron_u2").await.unwrap().unwrap(); assert!(cleared.last_status.is_none()); assert!(cleared.last_error.is_none()); assert!(cleared.skill_content.is_none()); @@ -870,7 +1008,7 @@ mod tests { name: Some("x".into()), ..Default::default() }; - let err = repo.update("cron_nope", ¶ms).await.unwrap_err(); + let err = repo.update_system("cron_nope", ¶ms).await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -879,9 +1017,11 @@ mod tests { let (repo, _db) = setup().await; repo.insert(&make_row("cron_noop")).await.unwrap(); - let before = repo.get_by_id("cron_noop").await.unwrap().unwrap(); - repo.update("cron_noop", &UpdateCronJobParams::default()).await.unwrap(); - let after = repo.get_by_id("cron_noop").await.unwrap().unwrap(); + let before = repo.get_by_id_system("cron_noop").await.unwrap().unwrap(); + repo.update_system("cron_noop", &UpdateCronJobParams::default()) + .await + .unwrap(); + let after = repo.get_by_id_system("cron_noop").await.unwrap().unwrap(); assert_eq!(before.updated_at, after.updated_at); } @@ -891,15 +1031,15 @@ mod tests { let (repo, _db) = setup().await; repo.insert(&make_row("cron_d1")).await.unwrap(); - repo.delete("cron_d1").await.unwrap(); - let result = repo.get_by_id("cron_d1").await.unwrap(); + repo.delete_system("cron_d1").await.unwrap(); + let result = repo.get_by_id_system("cron_d1").await.unwrap(); assert!(result.is_none()); } #[tokio::test] async fn delete_nonexistent_returns_not_found() { let (repo, _db) = setup().await; - let err = repo.delete("cron_nope").await.unwrap_err(); + let err = repo.delete_system("cron_nope").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -909,17 +1049,17 @@ mod tests { repo.insert(&make_row("cron_dc1")).await.unwrap(); repo.insert(&make_row("cron_dc2")).await.unwrap(); - let deleted = repo.delete_by_conversation("conv_1").await.unwrap(); + let deleted = repo.delete_by_conversation_system("conv_1").await.unwrap(); assert_eq!(deleted, 2); - let remaining = repo.list_all().await.unwrap(); + let remaining = repo.list_all_system().await.unwrap(); assert!(remaining.is_empty()); } #[tokio::test] async fn delete_by_conversation_returns_zero_for_no_match() { let (repo, _db) = setup().await; - let deleted = repo.delete_by_conversation("conv_none").await.unwrap(); + let deleted = repo.delete_by_conversation_system("conv_none").await.unwrap(); assert_eq!(deleted, 0); } @@ -936,9 +1076,9 @@ mod tests { next_run_at: Some(Some(9999999)), ..Default::default() }; - repo.update("cron_s1", ¶ms).await.unwrap(); + repo.update_system("cron_s1", ¶ms).await.unwrap(); - let updated = repo.get_by_id("cron_s1").await.unwrap().unwrap(); + let updated = repo.get_by_id_system("cron_s1").await.unwrap().unwrap(); assert_eq!(updated.schedule_kind, "cron"); assert_eq!(updated.schedule_value, "0 0 9 * * *"); assert_eq!(updated.schedule_tz.as_deref(), Some("Asia/Shanghai")); @@ -960,7 +1100,7 @@ mod tests { cron_job.schedule_tz = Some("UTC".into()); repo.insert(&cron_job).await.unwrap(); - let all = repo.list_all().await.unwrap(); + let all = repo.list_all_system().await.unwrap(); assert_eq!(all.len(), 2); } @@ -971,7 +1111,7 @@ mod tests { row.skill_content = Some("---\nname: My Skill\ndescription: A test\n---\nDo X".into()); repo.insert(&row).await.unwrap(); - let found = repo.get_by_id("cron_sk").await.unwrap().unwrap(); + let found = repo.get_by_id_system("cron_sk").await.unwrap().unwrap(); assert!(found.skill_content.unwrap().contains("My Skill")); } @@ -982,7 +1122,7 @@ mod tests { row.agent_config = Some(r#"{"name":"GPT","model_id":"gpt-4"}"#.into()); repo.insert(&row).await.unwrap(); - let found = repo.get_by_id("cron_ac").await.unwrap().unwrap(); + let found = repo.get_by_id_system("cron_ac").await.unwrap().unwrap(); let config = found.agent_config.unwrap(); assert!(config.contains("gpt-4")); } diff --git a/crates/aionui-db/src/repository/sqlite_diagnostics.rs b/crates/aionui-db/src/repository/sqlite_diagnostics.rs index 8793717cf..0a3753749 100644 --- a/crates/aionui-db/src/repository/sqlite_diagnostics.rs +++ b/crates/aionui-db/src/repository/sqlite_diagnostics.rs @@ -85,14 +85,14 @@ impl SqliteFeedbackDiagnosticsRepository { "created_at": conversation.try_get::("created_at")?, "updated_at": conversation.try_get::("updated_at")?, }, - "messages": self.collect_message_diagnostics(conversation_id).await?, + "messages": self.collect_message_diagnostics(&request.user_id, conversation_id).await?, "recent_conversations": self.collect_recent_conversations( &request.user_id, conversation_id, conversation.try_get::("updated_at")?, ).await?, - "acp_session": self.collect_acp_session(conversation_id).await?, - "agent_metadata": self.collect_agent_metadata(agent_id.as_deref()).await?, + "acp_session": self.collect_acp_session(&request.user_id, conversation_id).await?, + "agent_metadata": self.collect_agent_metadata(&request.user_id, agent_id.as_deref()).await?, "assistant_snapshot": self.collect_assistant_snapshot(&request.user_id, conversation_id).await?, }); @@ -218,13 +218,15 @@ impl SqliteFeedbackDiagnosticsRepository { })) } - async fn collect_message_diagnostics(&self, conversation_id: &str) -> Result { + async fn collect_message_diagnostics(&self, user_id: &str, conversation_id: &str) -> Result { let aggregate_rows = sqlx::query( - "SELECT type, status, hidden, COUNT(*) AS count, SUM(length(content)) AS content_bytes \ - FROM messages \ - WHERE conversation_id = ? \ - GROUP BY type, status, hidden", + "SELECT m.type, m.status, m.hidden, COUNT(*) AS count, SUM(length(m.content)) AS content_bytes \ + FROM messages m \ + JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? AND m.conversation_id = ? \ + GROUP BY m.type, m.status, m.hidden", ) + .bind(user_id) .bind(conversation_id) .fetch_all(&self.pool) .await?; @@ -253,16 +255,18 @@ impl SqliteFeedbackDiagnosticsRepository { let recent_messages = sqlx::query( "SELECT \ - id, msg_id, type, position, status, hidden, created_at, length(content) AS content_bytes, \ - CASE WHEN json_valid(content) THEN length(json_extract(content, '$.text')) END AS text_length, \ - CASE WHEN json_valid(content) THEN json_array_length(content, '$.attachments') END AS attachment_count, \ - CASE WHEN json_valid(content) THEN json_array_length(content, '$.images') END AS image_count, \ - CASE WHEN json_valid(content) THEN json_array_length(content, '$.toolCalls') END AS tool_call_count \ - FROM messages \ - WHERE conversation_id = ? \ - ORDER BY created_at DESC, id DESC \ + m.id, m.msg_id, m.type, m.position, m.status, m.hidden, m.created_at, length(m.content) AS content_bytes, \ + CASE WHEN json_valid(m.content) THEN length(json_extract(m.content, '$.text')) END AS text_length, \ + CASE WHEN json_valid(m.content) THEN json_array_length(m.content, '$.attachments') END AS attachment_count, \ + CASE WHEN json_valid(m.content) THEN json_array_length(m.content, '$.images') END AS image_count, \ + CASE WHEN json_valid(m.content) THEN json_array_length(m.content, '$.toolCalls') END AS tool_call_count \ + FROM messages m \ + JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? AND m.conversation_id = ? \ + ORDER BY m.created_at DESC, m.id DESC \ LIMIT 20", ) + .bind(user_id) .bind(conversation_id) .fetch_all(&self.pool) .await?; @@ -289,19 +293,21 @@ impl SqliteFeedbackDiagnosticsRepository { let recent_errors = sqlx::query( "SELECT \ - id, type, status, position, created_at, length(content) AS content_bytes, \ - CASE WHEN json_valid(content) THEN COALESCE(json_extract(content, '$.error.code'), json_extract(content, '$.code')) END AS code, \ - CASE WHEN json_valid(content) THEN COALESCE(json_extract(content, '$.error.ownership'), json_extract(content, '$.ownership')) END AS ownership, \ - CASE WHEN json_valid(content) THEN COALESCE(json_extract(content, '$.error.retryable'), json_extract(content, '$.retryable')) END AS retryable, \ - CASE WHEN json_valid(content) THEN json_extract(content, '$.resolution.kind') END AS resolution_kind, \ - CASE WHEN json_valid(content) THEN json_extract(content, '$.resolution.targetId') END AS resolution_target_id, \ - CASE WHEN json_valid(content) THEN json_extract(content, '$.feedbackRecommended') END AS feedback_recommended \ - FROM messages \ - WHERE conversation_id = ? \ - AND (status = 'error' OR type = 'tips' OR (json_valid(content) AND json_extract(content, '$.error.code') IS NOT NULL)) \ - ORDER BY created_at DESC, id DESC \ + m.id, m.type, m.status, m.position, m.created_at, length(m.content) AS content_bytes, \ + CASE WHEN json_valid(m.content) THEN COALESCE(json_extract(m.content, '$.error.code'), json_extract(m.content, '$.code')) END AS code, \ + CASE WHEN json_valid(m.content) THEN COALESCE(json_extract(m.content, '$.error.ownership'), json_extract(m.content, '$.ownership')) END AS ownership, \ + CASE WHEN json_valid(m.content) THEN COALESCE(json_extract(m.content, '$.error.retryable'), json_extract(m.content, '$.retryable')) END AS retryable, \ + CASE WHEN json_valid(m.content) THEN json_extract(m.content, '$.resolution.kind') END AS resolution_kind, \ + CASE WHEN json_valid(m.content) THEN json_extract(m.content, '$.resolution.targetId') END AS resolution_target_id, \ + CASE WHEN json_valid(m.content) THEN json_extract(m.content, '$.feedbackRecommended') END AS feedback_recommended \ + FROM messages m \ + JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? AND m.conversation_id = ? \ + AND (m.status = 'error' OR m.type = 'tips' OR (json_valid(m.content) AND json_extract(m.content, '$.error.code') IS NOT NULL)) \ + ORDER BY m.created_at DESC, m.id DESC \ LIMIT 10", ) + .bind(user_id) .bind(conversation_id) .fetch_all(&self.pool) .await?; @@ -328,16 +334,18 @@ impl SqliteFeedbackDiagnosticsRepository { let recent_error_detail_rows = sqlx::query( "SELECT \ - id, msg_id, type, status, position, hidden, created_at, length(content) AS content_bytes, content, \ - CASE WHEN json_valid(content) THEN COALESCE(json_extract(content, '$.error.code'), json_extract(content, '$.code')) END AS code, \ - CASE WHEN json_valid(content) THEN COALESCE(json_extract(content, '$.error.ownership'), json_extract(content, '$.ownership')) END AS ownership, \ - CASE WHEN json_valid(content) THEN COALESCE(json_extract(content, '$.error.retryable'), json_extract(content, '$.retryable')) END AS retryable \ - FROM messages \ - WHERE conversation_id = ? \ - AND (status = 'error' OR type = 'tips' OR (json_valid(content) AND json_extract(content, '$.error.code') IS NOT NULL)) \ - ORDER BY created_at DESC, id DESC \ + m.id, m.msg_id, m.type, m.status, m.position, m.hidden, m.created_at, length(m.content) AS content_bytes, m.content, \ + CASE WHEN json_valid(m.content) THEN COALESCE(json_extract(m.content, '$.error.code'), json_extract(m.content, '$.code')) END AS code, \ + CASE WHEN json_valid(m.content) THEN COALESCE(json_extract(m.content, '$.error.ownership'), json_extract(m.content, '$.ownership')) END AS ownership, \ + CASE WHEN json_valid(m.content) THEN COALESCE(json_extract(m.content, '$.error.retryable'), json_extract(m.content, '$.retryable')) END AS retryable \ + FROM messages m \ + JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? AND m.conversation_id = ? \ + AND (m.status = 'error' OR m.type = 'tips' OR (json_valid(m.content) AND json_extract(m.content, '$.error.code') IS NOT NULL)) \ + ORDER BY m.created_at DESC, m.id DESC \ LIMIT 10", ) + .bind(user_id) .bind(conversation_id) .fetch_all(&self.pool) .await?; @@ -388,17 +396,19 @@ impl SqliteFeedbackDiagnosticsRepository { })) } - async fn collect_acp_session(&self, conversation_id: &str) -> Result { + async fn collect_acp_session(&self, user_id: &str, conversation_id: &str) -> Result { let row = sqlx::query( "SELECT \ - conversation_id, agent_source, agent_id, session_id, session_status, \ - session_config, length(session_config) AS session_config_bytes, \ - last_active_at, suspended_at, \ - CASE WHEN json_valid(session_config) THEN json_extract(session_config, '$.runtime.current_mode_id') END AS current_mode_id, \ - CASE WHEN json_valid(session_config) THEN json_extract(session_config, '$.runtime.current_model_id') END AS current_model_id \ - FROM acp_session \ - WHERE conversation_id = ?", + s.conversation_id, s.agent_source, s.agent_id, s.session_id, s.session_status, \ + s.session_config, length(s.session_config) AS session_config_bytes, \ + s.last_active_at, s.suspended_at, \ + CASE WHEN json_valid(s.session_config) THEN json_extract(s.session_config, '$.runtime.current_mode_id') END AS current_mode_id, \ + CASE WHEN json_valid(s.session_config) THEN json_extract(s.session_config, '$.runtime.current_model_id') END AS current_model_id \ + FROM acp_session s \ + JOIN conversations c ON c.id = s.conversation_id \ + WHERE c.user_id = ? AND s.conversation_id = ?", ) + .bind(user_id) .bind(conversation_id) .fetch_optional(&self.pool) .await?; @@ -428,7 +438,7 @@ impl SqliteFeedbackDiagnosticsRepository { })) } - async fn collect_agent_metadata(&self, agent_id: Option<&str>) -> Result { + async fn collect_agent_metadata(&self, user_id: &str, agent_id: Option<&str>) -> Result { let Some(agent_id) = agent_id else { return Ok(Value::Null); }; @@ -441,9 +451,10 @@ impl SqliteFeedbackDiagnosticsRepository { last_check_status, last_check_kind, last_check_error_code, last_check_latency_ms, \ last_check_at, last_success_at, last_failure_at \ FROM agent_metadata \ - WHERE id = ?", + WHERE id = ? AND (user_id IS NULL OR user_id = ?)", ) .bind(agent_id) + .bind(user_id) .fetch_optional(&self.pool) .await?; @@ -528,17 +539,22 @@ impl SqliteFeedbackDiagnosticsRepository { let provider_id = self.resolve_provider_id(request).await?; let mut query = "SELECT id, platform, name, base_url, api_key_encrypted, models, enabled, capabilities, \ context_limit, model_enabled, model_health, is_full_url, created_at, updated_at \ - FROM providers" + FROM providers \ + WHERE user_id = ?" .to_owned(); if provider_id.is_some() { - query.push_str(" WHERE id = ?"); + query.push_str(" AND id = ?"); } query.push_str(" ORDER BY updated_at DESC LIMIT 20"); let rows = if let Some(provider_id) = provider_id.as_deref() { - sqlx::query(&query).bind(provider_id).fetch_all(&self.pool).await? + sqlx::query(&query) + .bind(&request.user_id) + .bind(provider_id) + .fetch_all(&self.pool) + .await? } else { - sqlx::query(&query).fetch_all(&self.pool).await? + sqlx::query(&query).bind(&request.user_id).fetch_all(&self.pool).await? }; let providers = rows @@ -726,7 +742,7 @@ impl SqliteFeedbackDiagnosticsRepository { original_json, builtin, deleted_at, created_at, updated_at, length(transport_config) AS transport_config_bytes, \ length(original_json) AS original_json_bytes \ FROM mcp_servers \ - WHERE deleted_at IS NULL" + WHERE user_id = ? AND deleted_at IS NULL" .to_owned(); if request.context.mcp_server_id.is_some() { query.push_str(" AND id = ?"); @@ -734,9 +750,13 @@ impl SqliteFeedbackDiagnosticsRepository { query.push_str(" ORDER BY updated_at DESC LIMIT 30"); let rows = if let Some(mcp_server_id) = request.context.mcp_server_id.as_deref() { - sqlx::query(&query).bind(mcp_server_id).fetch_all(&self.pool).await? + sqlx::query(&query) + .bind(&request.user_id) + .bind(mcp_server_id) + .fetch_all(&self.pool) + .await? } else { - sqlx::query(&query).fetch_all(&self.pool).await? + sqlx::query(&query).bind(&request.user_id).fetch_all(&self.pool).await? }; let servers = rows @@ -774,18 +794,20 @@ impl SqliteFeedbackDiagnosticsRepository { )) } - async fn collect_client_ui_settings(&self) -> Result { + async fn collect_client_ui_settings(&self, user_id: &str) -> Result { let preference_rows = sqlx::query( "SELECT key, value, updated_at \ FROM client_preferences \ - WHERE key LIKE 'appearance.%' \ + WHERE user_id = ? \ + AND (key LIKE 'appearance.%' \ OR key LIKE 'window.%' \ OR key LIKE 'display.%' \ OR key LIKE 'workspace.%' \ - OR key LIKE 'settings.%' \ + OR key LIKE 'settings.%') \ ORDER BY updated_at DESC, key ASC \ LIMIT 50", ) + .bind(user_id) .fetch_all(&self.pool) .await?; @@ -805,8 +827,9 @@ impl SqliteFeedbackDiagnosticsRepository { let settings = sqlx::query( "SELECT language, notification_enabled, cron_notification_enabled, command_queue_enabled, save_upload_to_workspace, updated_at \ FROM system_settings \ - WHERE id = 1", + WHERE user_id = ?", ) + .bind(user_id) .fetch_optional(&self.pool) .await?; @@ -962,15 +985,20 @@ impl SqliteFeedbackDiagnosticsRepository { .bind(&request.user_id) .fetch_one(&self.pool) .await?; - let provider_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM providers") - .fetch_one(&self.pool) - .await?; - let agent_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM agent_metadata") - .fetch_one(&self.pool) - .await?; - let active_mcp_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM mcp_servers WHERE deleted_at IS NULL") + let provider_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM providers WHERE user_id = ?") + .bind(&request.user_id) .fetch_one(&self.pool) .await?; + let agent_count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM agent_metadata WHERE user_id IS NULL OR user_id = ?") + .bind(&request.user_id) + .fetch_one(&self.pool) + .await?; + let active_mcp_count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM mcp_servers WHERE user_id = ? AND deleted_at IS NULL") + .bind(&request.user_id) + .fetch_one(&self.pool) + .await?; Ok(profile_result( FeedbackDiagnosticsProfile::GlobalSummary, @@ -984,8 +1012,8 @@ impl SqliteFeedbackDiagnosticsRepository { "conversation_status_counts": self.collect_global_conversation_status_counts(&request.user_id).await?, "recent_conversations": self.collect_global_recent_conversations(&request.user_id).await?, "recent_errors": self.collect_global_recent_errors(&request.user_id).await?, - "agent_health": self.collect_global_agent_health().await?, - "provider_health": self.collect_global_provider_health().await?, + "agent_health": self.collect_global_agent_health(&request.user_id).await?, + "provider_health": self.collect_global_provider_health(&request.user_id).await?, }), )) } @@ -1116,7 +1144,7 @@ impl SqliteFeedbackDiagnosticsRepository { let conversation_message_count = row.try_get::("message_count")?; let conversation_error_message_count = row.try_get::("error_message_count")?; let conversation_updated_at = row.try_get::("updated_at")?; - let messages = self.collect_global_conversation_message_samples(&id).await?; + let messages = self.collect_global_conversation_message_samples(user_id, &id).await?; let item = json!({ "id": id.clone(), "conversation_id": id, @@ -1243,22 +1271,28 @@ impl SqliteFeedbackDiagnosticsRepository { })) } - async fn collect_global_conversation_message_samples(&self, conversation_id: &str) -> Result { + async fn collect_global_conversation_message_samples( + &self, + user_id: &str, + conversation_id: &str, + ) -> Result { let error_rows = sqlx::query( "SELECT \ - id, msg_id, type, status, position, hidden, created_at, length(content) AS content_bytes, content, \ - CASE WHEN json_valid(content) THEN COALESCE(json_extract(content, '$.error.code'), json_extract(content, '$.code')) END AS code, \ - CASE WHEN json_valid(content) THEN COALESCE(json_extract(content, '$.error.ownership'), json_extract(content, '$.ownership')) END AS ownership, \ - CASE WHEN json_valid(content) THEN COALESCE(json_extract(content, '$.error.retryable'), json_extract(content, '$.retryable')) END AS retryable, \ - CASE WHEN json_valid(content) THEN json_extract(content, '$.resolution.kind') END AS resolution_kind, \ - CASE WHEN json_valid(content) THEN json_extract(content, '$.resolution.targetId') END AS resolution_target_id, \ - CASE WHEN json_valid(content) THEN json_extract(content, '$.feedbackRecommended') END AS feedback_recommended \ - FROM messages \ - WHERE conversation_id = ? \ - AND (status = 'error' OR type = 'tips' OR (json_valid(content) AND json_extract(content, '$.error.code') IS NOT NULL)) \ - ORDER BY created_at DESC, id DESC \ + m.id, m.msg_id, m.type, m.status, m.position, m.hidden, m.created_at, length(m.content) AS content_bytes, m.content, \ + CASE WHEN json_valid(m.content) THEN COALESCE(json_extract(m.content, '$.error.code'), json_extract(m.content, '$.code')) END AS code, \ + CASE WHEN json_valid(m.content) THEN COALESCE(json_extract(m.content, '$.error.ownership'), json_extract(m.content, '$.ownership')) END AS ownership, \ + CASE WHEN json_valid(m.content) THEN COALESCE(json_extract(m.content, '$.error.retryable'), json_extract(m.content, '$.retryable')) END AS retryable, \ + CASE WHEN json_valid(m.content) THEN json_extract(m.content, '$.resolution.kind') END AS resolution_kind, \ + CASE WHEN json_valid(m.content) THEN json_extract(m.content, '$.resolution.targetId') END AS resolution_target_id, \ + CASE WHEN json_valid(m.content) THEN json_extract(m.content, '$.feedbackRecommended') END AS feedback_recommended \ + FROM messages m \ + JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? AND m.conversation_id = ? \ + AND (m.status = 'error' OR m.type = 'tips' OR (json_valid(m.content) AND json_extract(m.content, '$.error.code') IS NOT NULL)) \ + ORDER BY m.created_at DESC, m.id DESC \ LIMIT ?", ) + .bind(user_id) .bind(conversation_id) .bind(GLOBAL_CONVERSATION_ERROR_LIMIT) .fetch_all(&self.pool) @@ -1293,17 +1327,19 @@ impl SqliteFeedbackDiagnosticsRepository { let message_rows = sqlx::query( "SELECT \ - id, msg_id, type, status, position, hidden, created_at, length(content) AS content_bytes, content, \ - CASE WHEN json_valid(content) THEN length(json_extract(content, '$.text')) END AS text_length, \ - CASE WHEN json_valid(content) THEN json_array_length(content, '$.attachments') END AS attachment_count, \ - CASE WHEN json_valid(content) THEN json_array_length(content, '$.images') END AS image_count, \ - CASE WHEN json_valid(content) THEN json_array_length(content, '$.toolCalls') END AS tool_call_count \ - FROM messages \ - WHERE conversation_id = ? \ - AND NOT (status = 'error' OR type = 'tips' OR (json_valid(content) AND json_extract(content, '$.error.code') IS NOT NULL)) \ - ORDER BY created_at DESC, id DESC \ + m.id, m.msg_id, m.type, m.status, m.position, m.hidden, m.created_at, length(m.content) AS content_bytes, m.content, \ + CASE WHEN json_valid(m.content) THEN length(json_extract(m.content, '$.text')) END AS text_length, \ + CASE WHEN json_valid(m.content) THEN json_array_length(m.content, '$.attachments') END AS attachment_count, \ + CASE WHEN json_valid(m.content) THEN json_array_length(m.content, '$.images') END AS image_count, \ + CASE WHEN json_valid(m.content) THEN json_array_length(m.content, '$.toolCalls') END AS tool_call_count \ + FROM messages m \ + JOIN conversations c ON c.id = m.conversation_id \ + WHERE c.user_id = ? AND m.conversation_id = ? \ + AND NOT (m.status = 'error' OR m.type = 'tips' OR (json_valid(m.content) AND json_extract(m.content, '$.error.code') IS NOT NULL)) \ + ORDER BY m.created_at DESC, m.id DESC \ LIMIT ?", ) + .bind(user_id) .bind(conversation_id) .bind(GLOBAL_CONVERSATION_MESSAGE_LIMIT) .fetch_all(&self.pool) @@ -1404,7 +1440,7 @@ impl SqliteFeedbackDiagnosticsRepository { })) } - async fn collect_global_agent_health(&self) -> Result { + async fn collect_global_agent_health(&self, user_id: &str) -> Result { let rows = sqlx::query( "SELECT \ id, name, backend, agent_type, agent_source, enabled, sort_order, \ @@ -1412,9 +1448,11 @@ impl SqliteFeedbackDiagnosticsRepository { last_check_status, last_check_kind, last_check_error_code, last_check_latency_ms, \ last_check_at, last_success_at, last_failure_at, updated_at \ FROM agent_metadata \ + WHERE user_id IS NULL OR user_id = ? \ ORDER BY updated_at DESC, id DESC \ LIMIT ?", ) + .bind(user_id) .bind(GLOBAL_HEALTH_LIMIT) .fetch_all(&self.pool) .await?; @@ -1451,14 +1489,16 @@ impl SqliteFeedbackDiagnosticsRepository { })) } - async fn collect_global_provider_health(&self) -> Result { + async fn collect_global_provider_health(&self, user_id: &str) -> Result { let rows = sqlx::query( "SELECT id, platform, name, base_url, api_key_encrypted, models, enabled, capabilities, \ context_limit, model_enabled, model_health, is_full_url, created_at, updated_at \ FROM providers \ + WHERE user_id = ? \ ORDER BY updated_at DESC, id DESC \ LIMIT ?", ) + .bind(user_id) .bind(GLOBAL_HEALTH_LIMIT) .fetch_all(&self.pool) .await?; @@ -1514,7 +1554,9 @@ impl IFeedbackDiagnosticsRepository for SqliteFeedbackDiagnosticsRepository { FeedbackDiagnosticsProfile::ModelAuth => self.collect_model_auth(request).await?, FeedbackDiagnosticsProfile::AgentTeam => self.collect_agent_team(request).await?, FeedbackDiagnosticsProfile::McpTools => self.collect_mcp_tools(request).await?, - FeedbackDiagnosticsProfile::ClientUiSettings => self.collect_client_ui_settings().await?, + FeedbackDiagnosticsProfile::ClientUiSettings => { + self.collect_client_ui_settings(&request.user_id).await? + } FeedbackDiagnosticsProfile::WorkspaceSummary => self.collect_workspace_summary(request).await?, FeedbackDiagnosticsProfile::GlobalSummary => self.collect_global_summary(request).await?, }; diff --git a/crates/aionui-db/src/repository/sqlite_mcp_server.rs b/crates/aionui-db/src/repository/sqlite_mcp_server.rs index 3214e7415..6dcbd529a 100644 --- a/crates/aionui-db/src/repository/sqlite_mcp_server.rs +++ b/crates/aionui-db/src/repository/sqlite_mcp_server.rs @@ -22,36 +22,44 @@ impl SqliteMcpServerRepository { #[async_trait::async_trait] impl IMcpServerRepository for SqliteMcpServerRepository { - async fn list(&self) -> Result, DbError> { + async fn list(&self, user_id: &str) -> Result, DbError> { let rows = sqlx::query_as::<_, McpServerRow>( - "SELECT * FROM mcp_servers WHERE deleted_at IS NULL ORDER BY created_at ASC", + "SELECT * FROM mcp_servers WHERE user_id = ? AND deleted_at IS NULL ORDER BY created_at ASC, rowid ASC", ) + .bind(user_id) .fetch_all(&self.pool) .await?; Ok(rows) } - async fn find_by_id(&self, id: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, McpServerRow>("SELECT * FROM mcp_servers WHERE id = ? AND deleted_at IS NULL") - .bind(id) - .fetch_optional(&self.pool) - .await?; + async fn find_by_id(&self, user_id: &str, id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, McpServerRow>( + "SELECT * FROM mcp_servers WHERE user_id = ? AND id = ? AND deleted_at IS NULL", + ) + .bind(user_id) + .bind(id) + .fetch_optional(&self.pool) + .await?; Ok(row) } - async fn find_by_name(&self, name: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, McpServerRow>("SELECT * FROM mcp_servers WHERE name = ? AND deleted_at IS NULL") - .bind(name) - .fetch_optional(&self.pool) - .await?; + async fn find_by_name(&self, user_id: &str, name: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, McpServerRow>( + "SELECT * FROM mcp_servers WHERE user_id = ? AND name = ? AND deleted_at IS NULL", + ) + .bind(user_id) + .bind(name) + .fetch_optional(&self.pool) + .await?; Ok(row) } - async fn find_by_id_any(&self, id: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, McpServerRow>("SELECT * FROM mcp_servers WHERE id = ?") + async fn find_by_id_any(&self, user_id: &str, id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, McpServerRow>("SELECT * FROM mcp_servers WHERE user_id = ? AND id = ?") + .bind(user_id) .bind(id) .fetch_optional(&self.pool) .await?; @@ -59,8 +67,9 @@ impl IMcpServerRepository for SqliteMcpServerRepository { Ok(row) } - async fn find_by_name_any(&self, name: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, McpServerRow>("SELECT * FROM mcp_servers WHERE name = ?") + async fn find_by_name_any(&self, user_id: &str, name: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, McpServerRow>("SELECT * FROM mcp_servers WHERE user_id = ? AND name = ?") + .bind(user_id) .bind(name) .fetch_optional(&self.pool) .await?; @@ -68,17 +77,19 @@ impl IMcpServerRepository for SqliteMcpServerRepository { Ok(row) } - async fn list_by_ids_any(&self, ids: &[String]) -> Result, DbError> { + async fn list_by_ids_any(&self, user_id: &str, ids: &[String]) -> Result, DbError> { if ids.is_empty() { return Ok(Vec::new()); } - let mut query = QueryBuilder::new("SELECT * FROM mcp_servers WHERE id IN ("); + let mut query = QueryBuilder::new("SELECT * FROM mcp_servers WHERE user_id = "); + query.push_bind(user_id); + query.push(" AND id IN ("); let mut separated = query.separated(", "); for id in ids { separated.push_bind(id); } - separated.push_unseparated(") ORDER BY created_at ASC"); + separated.push_unseparated(") ORDER BY created_at ASC, rowid ASC"); let rows = query.build_query_as::().fetch_all(&self.pool).await?; let rows_by_id: HashMap<_, _> = rows.into_iter().map(|row| (row.id.clone(), row)).collect(); @@ -93,12 +104,13 @@ impl IMcpServerRepository for SqliteMcpServerRepository { sqlx::query( "INSERT INTO mcp_servers \ - (id, name, description, enabled, transport_type, transport_config, \ + (id, user_id, name, description, enabled, transport_type, transport_config, \ tools, last_test_status, last_connected, original_json, builtin, \ deleted_at, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(&id) + .bind(params.user_id) .bind(params.name) .bind(params.description) .bind(params.enabled) @@ -123,6 +135,7 @@ impl IMcpServerRepository for SqliteMcpServerRepository { Ok(McpServerRow { id, + user_id: params.user_id.to_string(), name: params.name.to_string(), description: params.description.map(String::from), enabled: params.enabled, @@ -139,9 +152,14 @@ impl IMcpServerRepository for SqliteMcpServerRepository { }) } - async fn update(&self, id: &str, params: UpdateMcpServerParams<'_>) -> Result { + async fn update( + &self, + user_id: &str, + id: &str, + params: UpdateMcpServerParams<'_>, + ) -> Result { let existing = self - .find_by_id_any(id) + .find_by_id_any(user_id, id) .await? .ok_or_else(|| DbError::NotFound(format!("MCP server '{id}' not found")))?; @@ -152,7 +170,7 @@ impl IMcpServerRepository for SqliteMcpServerRepository { name = ?, description = ?, enabled = ?, transport_type = ?, \ transport_config = ?, tools = ?, original_json = ?, \ builtin = ?, deleted_at = ?, updated_at = ? \ - WHERE id = ?", + WHERE user_id = ? AND id = ?", ) .bind(&merged.name) .bind(&merged.description) @@ -164,6 +182,7 @@ impl IMcpServerRepository for SqliteMcpServerRepository { .bind(merged.builtin) .bind(merged.deleted_at) .bind(merged.updated_at) + .bind(user_id) .bind(id) .execute(&self.pool) .await @@ -177,13 +196,15 @@ impl IMcpServerRepository for SqliteMcpServerRepository { Ok(merged) } - async fn delete(&self, id: &str) -> Result<(), DbError> { + async fn delete(&self, user_id: &str, id: &str) -> Result<(), DbError> { let now = aionui_common::now_ms(); let result = sqlx::query( - "UPDATE mcp_servers SET enabled = 0, deleted_at = ?, updated_at = ? WHERE id = ? AND deleted_at IS NULL", + "UPDATE mcp_servers SET enabled = 0, deleted_at = ?, updated_at = ? \ + WHERE user_id = ? AND id = ? AND deleted_at IS NULL", ) .bind(now) .bind(now) + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -195,11 +216,15 @@ impl IMcpServerRepository for SqliteMcpServerRepository { Ok(()) } - async fn batch_upsert(&self, servers: &[CreateMcpServerParams<'_>]) -> Result, DbError> { + async fn batch_upsert( + &self, + user_id: &str, + servers: &[CreateMcpServerParams<'_>], + ) -> Result, DbError> { let mut results = Vec::with_capacity(servers.len()); for params in servers { - let row = match self.find_by_name(params.name).await? { + let row = match self.find_by_name(user_id, params.name).await? { Some(existing) => { let update_params = UpdateMcpServerParams { description: Some(params.description), @@ -211,7 +236,7 @@ impl IMcpServerRepository for SqliteMcpServerRepository { builtin: Some(params.builtin), ..Default::default() }; - self.update(&existing.id, update_params).await? + self.update(user_id, &existing.id, update_params).await? } None => self.create(params.clone()).await?, }; @@ -221,17 +246,24 @@ impl IMcpServerRepository for SqliteMcpServerRepository { Ok(results) } - async fn update_status(&self, id: &str, status: &str, last_connected: Option) -> Result<(), DbError> { + async fn update_status( + &self, + user_id: &str, + id: &str, + status: &str, + last_connected: Option, + ) -> Result<(), DbError> { let now = aionui_common::now_ms(); let result = sqlx::query( "UPDATE mcp_servers SET last_test_status = ?, \ last_connected = COALESCE(?, last_connected), \ - updated_at = ? WHERE id = ? AND deleted_at IS NULL", + updated_at = ? WHERE user_id = ? AND id = ? AND deleted_at IS NULL", ) .bind(status) .bind(last_connected) .bind(now) + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -243,16 +275,19 @@ impl IMcpServerRepository for SqliteMcpServerRepository { Ok(()) } - async fn update_tools(&self, id: &str, tools: Option<&str>) -> Result<(), DbError> { + async fn update_tools(&self, user_id: &str, id: &str, tools: Option<&str>) -> Result<(), DbError> { let now = aionui_common::now_ms(); - let result = - sqlx::query("UPDATE mcp_servers SET tools = ?, updated_at = ? WHERE id = ? AND deleted_at IS NULL") - .bind(tools) - .bind(now) - .bind(id) - .execute(&self.pool) - .await?; + let result = sqlx::query( + "UPDATE mcp_servers SET tools = ?, updated_at = ? \ + WHERE user_id = ? AND id = ? AND deleted_at IS NULL", + ) + .bind(tools) + .bind(now) + .bind(user_id) + .bind(id) + .execute(&self.pool) + .await?; if result.rows_affected() == 0 { return Err(DbError::NotFound(format!("MCP server '{id}' not found"))); @@ -267,6 +302,7 @@ fn merge_update(existing: McpServerRow, params: UpdateMcpServerParams<'_>) -> Mc let now = aionui_common::now_ms(); McpServerRow { id: existing.id, + user_id: existing.user_id, name: params.name.unwrap_or(&existing.name).to_string(), description: params.description.map_or(existing.description, |v| v.map(String::from)), enabled: params.enabled.unwrap_or(existing.enabled), @@ -297,14 +333,27 @@ mod tests { use super::*; use crate::init_database_memory; + const USER_A: &str = "system_default_user"; + const USER_B: &str = "user_b"; + async fn setup() -> (SqliteMcpServerRepository, crate::Database) { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, 1, 1)", + ) + .bind(USER_B) + .bind(USER_B) + .execute(db.pool()) + .await + .unwrap(); let repo = SqliteMcpServerRepository::new(db.pool().clone()); (repo, db) } fn stdio_params() -> CreateMcpServerParams<'static> { CreateMcpServerParams { + user_id: USER_A, name: "test-mcp", description: Some("A test MCP server"), enabled: false, @@ -318,6 +367,7 @@ mod tests { fn http_params() -> CreateMcpServerParams<'static> { CreateMcpServerParams { + user_id: USER_A, name: "http-mcp", description: None, enabled: true, @@ -332,7 +382,7 @@ mod tests { #[tokio::test] async fn list_empty() { let (repo, _db) = setup().await; - let servers = repo.list().await.unwrap(); + let servers = repo.list(USER_A).await.unwrap(); assert!(servers.is_empty()); } @@ -370,7 +420,7 @@ mod tests { let (repo, _db) = setup().await; let created = repo.create(stdio_params()).await.unwrap(); - let found = repo.find_by_id(&created.id).await.unwrap().unwrap(); + let found = repo.find_by_id(USER_A, &created.id).await.unwrap().unwrap(); assert_eq!(found.id, created.id); assert_eq!(found.name, "test-mcp"); } @@ -378,7 +428,7 @@ mod tests { #[tokio::test] async fn find_by_id_nonexistent() { let (repo, _db) = setup().await; - assert!(repo.find_by_id("no_such_id").await.unwrap().is_none()); + assert!(repo.find_by_id(USER_A, "no_such_id").await.unwrap().is_none()); } #[tokio::test] @@ -386,14 +436,14 @@ mod tests { let (repo, _db) = setup().await; let created = repo.create(stdio_params()).await.unwrap(); - let found = repo.find_by_name("test-mcp").await.unwrap().unwrap(); + let found = repo.find_by_name(USER_A, "test-mcp").await.unwrap().unwrap(); assert_eq!(found.id, created.id); } #[tokio::test] async fn find_by_name_nonexistent() { let (repo, _db) = setup().await; - assert!(repo.find_by_name("nope").await.unwrap().is_none()); + assert!(repo.find_by_name(USER_A, "nope").await.unwrap().is_none()); } #[tokio::test] @@ -402,7 +452,7 @@ mod tests { let s1 = repo.create(stdio_params()).await.unwrap(); let s2 = repo.create(http_params()).await.unwrap(); - let all = repo.list().await.unwrap(); + let all = repo.list(USER_A).await.unwrap(); assert_eq!(all.len(), 2); assert_eq!(all[0].id, s1.id); assert_eq!(all[1].id, s2.id); @@ -415,6 +465,7 @@ mod tests { let updated = repo .update( + USER_A, &created.id, UpdateMcpServerParams { enabled: Some(true), @@ -438,6 +489,7 @@ mod tests { let err = repo .update( + USER_A, &s2.id, UpdateMcpServerParams { name: Some("test-mcp"), @@ -457,6 +509,7 @@ mod tests { let updated = repo .update( + USER_A, &created.id, UpdateMcpServerParams { description: Some(None), @@ -475,7 +528,7 @@ mod tests { async fn update_nonexistent_returns_not_found() { let (repo, _db) = setup().await; let err = repo - .update("no_id", UpdateMcpServerParams::default()) + .update(USER_A, "no_id", UpdateMcpServerParams::default()) .await .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); @@ -486,14 +539,14 @@ mod tests { let (repo, _db) = setup().await; let created = repo.create(stdio_params()).await.unwrap(); - repo.delete(&created.id).await.unwrap(); - assert!(repo.find_by_id(&created.id).await.unwrap().is_none()); + repo.delete(USER_A, &created.id).await.unwrap(); + assert!(repo.find_by_id(USER_A, &created.id).await.unwrap().is_none()); } #[tokio::test] async fn delete_nonexistent_returns_not_found() { let (repo, _db) = setup().await; - let err = repo.delete("no_id").await.unwrap_err(); + let err = repo.delete(USER_A, "no_id").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -504,13 +557,16 @@ mod tests { assert!(!existing.enabled); let results = repo - .batch_upsert(&[ - CreateMcpServerParams { - enabled: true, - ..stdio_params() - }, - http_params(), - ]) + .batch_upsert( + USER_A, + &[ + CreateMcpServerParams { + enabled: true, + ..stdio_params() + }, + http_params(), + ], + ) .await .unwrap(); @@ -529,9 +585,11 @@ mod tests { let created = repo.create(stdio_params()).await.unwrap(); let ts = aionui_common::now_ms(); - repo.update_status(&created.id, "connected", Some(ts)).await.unwrap(); + repo.update_status(USER_A, &created.id, "connected", Some(ts)) + .await + .unwrap(); - let found = repo.find_by_id(&created.id).await.unwrap().unwrap(); + let found = repo.find_by_id(USER_A, &created.id).await.unwrap().unwrap(); assert_eq!(found.last_test_status, "connected"); assert_eq!(found.last_connected, Some(ts)); } @@ -542,11 +600,13 @@ mod tests { let created = repo.create(stdio_params()).await.unwrap(); let ts = aionui_common::now_ms(); - repo.update_status(&created.id, "connected", Some(ts)).await.unwrap(); + repo.update_status(USER_A, &created.id, "connected", Some(ts)) + .await + .unwrap(); - repo.update_status(&created.id, "error", None).await.unwrap(); + repo.update_status(USER_A, &created.id, "error", None).await.unwrap(); - let found = repo.find_by_id(&created.id).await.unwrap().unwrap(); + let found = repo.find_by_id(USER_A, &created.id).await.unwrap().unwrap(); assert_eq!(found.last_test_status, "error"); assert_eq!(found.last_connected, Some(ts)); } @@ -554,7 +614,10 @@ mod tests { #[tokio::test] async fn update_status_nonexistent_returns_not_found() { let (repo, _db) = setup().await; - let err = repo.update_status("no_id", "connected", None).await.unwrap_err(); + let err = repo + .update_status(USER_A, "no_id", "connected", None) + .await + .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -565,9 +628,9 @@ mod tests { assert!(created.tools.is_none()); let tools_json = r#"[{"name":"read_file","description":"Read a file"}]"#; - repo.update_tools(&created.id, Some(tools_json)).await.unwrap(); + repo.update_tools(USER_A, &created.id, Some(tools_json)).await.unwrap(); - let found = repo.find_by_id(&created.id).await.unwrap().unwrap(); + let found = repo.find_by_id(USER_A, &created.id).await.unwrap().unwrap(); assert_eq!(found.tools.as_deref(), Some(tools_json)); } @@ -583,16 +646,55 @@ mod tests { .unwrap(); assert!(created.tools.is_some()); - repo.update_tools(&created.id, None).await.unwrap(); + repo.update_tools(USER_A, &created.id, None).await.unwrap(); - let found = repo.find_by_id(&created.id).await.unwrap().unwrap(); + let found = repo.find_by_id(USER_A, &created.id).await.unwrap().unwrap(); assert!(found.tools.is_none()); } #[tokio::test] async fn update_tools_nonexistent_returns_not_found() { let (repo, _db) = setup().await; - let err = repo.update_tools("no_id", Some("[]")).await.unwrap_err(); + let err = repo.update_tools(USER_A, "no_id", Some("[]")).await.unwrap_err(); + assert!(matches!(err, DbError::NotFound(_))); + } + + #[tokio::test] + async fn mcp_server_operations_are_scoped_by_user() { + let (repo, _db) = setup().await; + let server_a = repo.create(stdio_params()).await.unwrap(); + let server_b = repo + .create(CreateMcpServerParams { + user_id: USER_B, + ..stdio_params() + }) + .await + .unwrap(); + + assert_ne!(server_a.id, server_b.id); + assert_eq!(repo.list(USER_A).await.unwrap().len(), 1); + assert_eq!(repo.list(USER_B).await.unwrap().len(), 1); + assert!(repo.find_by_id(USER_B, &server_a.id).await.unwrap().is_none()); + assert_eq!( + repo.find_by_name(USER_B, "test-mcp").await.unwrap().unwrap().id, + server_b.id + ); + + let err = repo + .update( + USER_B, + &server_a.id, + UpdateMcpServerParams { + enabled: Some(true), + ..Default::default() + }, + ) + .await + .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); + + repo.delete(USER_B, &server_b.id).await.unwrap(); + assert_eq!(repo.list(USER_A).await.unwrap().len(), 1); + assert!(repo.list(USER_B).await.unwrap().is_empty()); } } diff --git a/crates/aionui-db/src/repository/sqlite_oauth_token.rs b/crates/aionui-db/src/repository/sqlite_oauth_token.rs index 9a596975d..76926dbcb 100644 --- a/crates/aionui-db/src/repository/sqlite_oauth_token.rs +++ b/crates/aionui-db/src/repository/sqlite_oauth_token.rs @@ -18,8 +18,9 @@ impl SqliteOAuthTokenRepository { #[async_trait::async_trait] impl IOAuthTokenRepository for SqliteOAuthTokenRepository { - async fn get_by_url(&self, server_url: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, OAuthTokenRow>("SELECT * FROM oauth_tokens WHERE server_url = ?") + async fn get_by_url(&self, user_id: &str, server_url: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, OAuthTokenRow>("SELECT * FROM oauth_tokens WHERE user_id = ? AND server_url = ?") + .bind(user_id) .bind(server_url) .fetch_optional(&self.pool) .await?; @@ -32,16 +33,17 @@ impl IOAuthTokenRepository for SqliteOAuthTokenRepository { sqlx::query( "INSERT INTO oauth_tokens \ - (server_url, access_token, refresh_token, token_type, \ + (user_id, server_url, access_token, refresh_token, token_type, \ expires_at, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?) \ - ON CONFLICT(server_url) DO UPDATE SET \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?) \ + ON CONFLICT(user_id, server_url) DO UPDATE SET \ access_token = excluded.access_token, \ refresh_token = excluded.refresh_token, \ token_type = excluded.token_type, \ expires_at = excluded.expires_at, \ updated_at = excluded.updated_at", ) + .bind(params.user_id) .bind(params.server_url) .bind(params.access_token) .bind(params.refresh_token) @@ -54,15 +56,16 @@ impl IOAuthTokenRepository for SqliteOAuthTokenRepository { // Fetch the row to get the correct created_at (preserved on conflict). let row = self - .get_by_url(params.server_url) + .get_by_url(params.user_id, params.server_url) .await? .ok_or_else(|| DbError::Init("Upsert succeeded but row not found".to_string()))?; Ok(row) } - async fn delete(&self, server_url: &str) -> Result<(), DbError> { - let result = sqlx::query("DELETE FROM oauth_tokens WHERE server_url = ?") + async fn delete(&self, user_id: &str, server_url: &str) -> Result<(), DbError> { + let result = sqlx::query("DELETE FROM oauth_tokens WHERE user_id = ? AND server_url = ?") + .bind(user_id) .bind(server_url) .execute(&self.pool) .await?; @@ -74,10 +77,12 @@ impl IOAuthTokenRepository for SqliteOAuthTokenRepository { Ok(()) } - async fn list_authenticated_urls(&self) -> Result, DbError> { - let rows: Vec<(String,)> = sqlx::query_as("SELECT server_url FROM oauth_tokens ORDER BY created_at ASC") - .fetch_all(&self.pool) - .await?; + async fn list_authenticated_urls(&self, user_id: &str) -> Result, DbError> { + let rows: Vec<(String,)> = + sqlx::query_as("SELECT server_url FROM oauth_tokens WHERE user_id = ? ORDER BY created_at ASC") + .bind(user_id) + .fetch_all(&self.pool) + .await?; Ok(rows.into_iter().map(|(url,)| url).collect()) } @@ -88,14 +93,27 @@ mod tests { use super::*; use crate::init_database_memory; + const USER_A: &str = "system_default_user"; + const USER_B: &str = "user_b"; + async fn setup() -> (SqliteOAuthTokenRepository, crate::Database) { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, 1, 1)", + ) + .bind(USER_B) + .bind(USER_B) + .execute(db.pool()) + .await + .unwrap(); let repo = SqliteOAuthTokenRepository::new(db.pool().clone()); (repo, db) } fn sample_params() -> UpsertOAuthTokenParams<'static> { UpsertOAuthTokenParams { + user_id: USER_A, server_url: "https://mcp.example.com", access_token: "enc_access_token_123", refresh_token: Some("enc_refresh_token_456"), @@ -107,7 +125,7 @@ mod tests { #[tokio::test] async fn get_by_url_nonexistent() { let (repo, _db) = setup().await; - assert!(repo.get_by_url("https://nope.com").await.unwrap().is_none()); + assert!(repo.get_by_url(USER_A, "https://nope.com").await.unwrap().is_none()); } #[tokio::test] @@ -131,6 +149,7 @@ mod tests { let updated = repo .upsert(UpsertOAuthTokenParams { + user_id: USER_A, server_url: "https://mcp.example.com", access_token: "new_access_token", refresh_token: None, @@ -153,7 +172,11 @@ mod tests { let (repo, _db) = setup().await; repo.upsert(sample_params()).await.unwrap(); - let found = repo.get_by_url("https://mcp.example.com").await.unwrap().unwrap(); + let found = repo + .get_by_url(USER_A, "https://mcp.example.com") + .await + .unwrap() + .unwrap(); assert_eq!(found.access_token, "enc_access_token_123"); } @@ -162,21 +185,26 @@ mod tests { let (repo, _db) = setup().await; repo.upsert(sample_params()).await.unwrap(); - repo.delete("https://mcp.example.com").await.unwrap(); - assert!(repo.get_by_url("https://mcp.example.com").await.unwrap().is_none()); + repo.delete(USER_A, "https://mcp.example.com").await.unwrap(); + assert!( + repo.get_by_url(USER_A, "https://mcp.example.com") + .await + .unwrap() + .is_none() + ); } #[tokio::test] async fn delete_nonexistent_returns_not_found() { let (repo, _db) = setup().await; - let err = repo.delete("https://nope.com").await.unwrap_err(); + let err = repo.delete(USER_A, "https://nope.com").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } #[tokio::test] async fn list_authenticated_urls_empty() { let (repo, _db) = setup().await; - let urls = repo.list_authenticated_urls().await.unwrap(); + let urls = repo.list_authenticated_urls(USER_A).await.unwrap(); assert!(urls.is_empty()); } @@ -185,6 +213,7 @@ mod tests { let (repo, _db) = setup().await; repo.upsert(sample_params()).await.unwrap(); repo.upsert(UpsertOAuthTokenParams { + user_id: USER_A, server_url: "https://other.example.com", access_token: "token2", refresh_token: None, @@ -194,9 +223,54 @@ mod tests { .await .unwrap(); - let urls = repo.list_authenticated_urls().await.unwrap(); + let urls = repo.list_authenticated_urls(USER_A).await.unwrap(); assert_eq!(urls.len(), 2); assert!(urls.contains(&"https://mcp.example.com".to_string())); assert!(urls.contains(&"https://other.example.com".to_string())); } + + #[tokio::test] + async fn oauth_tokens_are_scoped_by_user() { + let (repo, _db) = setup().await; + repo.upsert(sample_params()).await.unwrap(); + repo.upsert(UpsertOAuthTokenParams { + user_id: USER_B, + access_token: "user_b_token", + refresh_token: None, + ..sample_params() + }) + .await + .unwrap(); + + assert_eq!( + repo.get_by_url(USER_A, "https://mcp.example.com") + .await + .unwrap() + .unwrap() + .access_token, + "enc_access_token_123" + ); + assert_eq!( + repo.get_by_url(USER_B, "https://mcp.example.com") + .await + .unwrap() + .unwrap() + .access_token, + "user_b_token" + ); + + repo.delete(USER_B, "https://mcp.example.com").await.unwrap(); + assert!( + repo.get_by_url(USER_B, "https://mcp.example.com") + .await + .unwrap() + .is_none() + ); + assert!( + repo.get_by_url(USER_A, "https://mcp.example.com") + .await + .unwrap() + .is_some() + ); + } } diff --git a/crates/aionui-db/src/repository/sqlite_provider.rs b/crates/aionui-db/src/repository/sqlite_provider.rs index e943c88f6..d18045fe0 100644 --- a/crates/aionui-db/src/repository/sqlite_provider.rs +++ b/crates/aionui-db/src/repository/sqlite_provider.rs @@ -19,16 +19,18 @@ impl SqliteProviderRepository { #[async_trait::async_trait] impl IProviderRepository for SqliteProviderRepository { - async fn list(&self) -> Result, DbError> { - let rows = sqlx::query_as::<_, Provider>("SELECT * FROM providers ORDER BY created_at ASC") + async fn list(&self, user_id: &str) -> Result, DbError> { + let rows = sqlx::query_as::<_, Provider>("SELECT * FROM providers WHERE user_id = ? ORDER BY created_at ASC") + .bind(user_id) .fetch_all(&self.pool) .await?; Ok(rows) } - async fn find_by_id(&self, id: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, Provider>("SELECT * FROM providers WHERE id = ?") + async fn find_by_id(&self, user_id: &str, id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, Provider>("SELECT * FROM providers WHERE user_id = ? AND id = ?") + .bind(user_id) .bind(id) .fetch_optional(&self.pool) .await?; @@ -45,12 +47,13 @@ impl IProviderRepository for SqliteProviderRepository { sqlx::query( "INSERT INTO providers \ - (id, platform, name, base_url, api_key_encrypted, models, enabled, \ + (id, user_id, platform, name, base_url, api_key_encrypted, models, enabled, \ capabilities, context_limit, model_protocols, model_enabled, \ model_health, model_settings, bedrock_config, is_full_url, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(&id) + .bind(params.user_id) .bind(params.platform) .bind(params.name) .bind(params.base_url) @@ -78,6 +81,7 @@ impl IProviderRepository for SqliteProviderRepository { Ok(Provider { id, + user_id: params.user_id.to_string(), platform: params.platform.to_string(), name: params.name.to_string(), base_url: params.base_url.to_string(), @@ -97,9 +101,9 @@ impl IProviderRepository for SqliteProviderRepository { }) } - async fn update(&self, id: &str, params: UpdateProviderParams<'_>) -> Result { + async fn update(&self, user_id: &str, id: &str, params: UpdateProviderParams<'_>) -> Result { let existing = self - .find_by_id(id) + .find_by_id(user_id, id) .await? .ok_or_else(|| DbError::NotFound(format!("Provider '{id}' not found")))?; @@ -111,7 +115,7 @@ impl IProviderRepository for SqliteProviderRepository { models = ?, enabled = ?, capabilities = ?, context_limit = ?, \ model_protocols = ?, model_enabled = ?, model_health = ?, \ model_settings = ?, bedrock_config = ?, is_full_url = ?, updated_at = ? \ - WHERE id = ?", + WHERE user_id = ? AND id = ?", ) .bind(&merged.platform) .bind(&merged.name) @@ -128,6 +132,7 @@ impl IProviderRepository for SqliteProviderRepository { .bind(&merged.bedrock_config) .bind(merged.is_full_url) .bind(merged.updated_at) + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -135,8 +140,9 @@ impl IProviderRepository for SqliteProviderRepository { Ok(merged) } - async fn delete(&self, id: &str) -> Result<(), DbError> { - let result = sqlx::query("DELETE FROM providers WHERE id = ?") + async fn delete(&self, user_id: &str, id: &str) -> Result<(), DbError> { + let result = sqlx::query("DELETE FROM providers WHERE user_id = ? AND id = ?") + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -159,6 +165,7 @@ fn merge_update(existing: Provider, params: UpdateProviderParams<'_>) -> Provide let now = aionui_common::now_ms(); Provider { id: existing.id, + user_id: existing.user_id, platform: params.platform.unwrap_or(&existing.platform).to_string(), name: params.name.unwrap_or(&existing.name).to_string(), base_url: params.base_url.unwrap_or(&existing.base_url).to_string(), @@ -194,8 +201,20 @@ mod tests { use super::*; use crate::init_database_memory; + const USER_A: &str = "system_default_user"; + const USER_B: &str = "user_b"; + async fn setup() -> (SqliteProviderRepository, crate::Database) { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, 1, 1)", + ) + .bind(USER_B) + .bind(USER_B) + .execute(db.pool()) + .await + .unwrap(); let repo = SqliteProviderRepository::new(db.pool().clone()); (repo, db) } @@ -203,6 +222,7 @@ mod tests { fn sample_params() -> CreateProviderParams<'static> { CreateProviderParams { id: None, + user_id: USER_A, platform: "anthropic", name: "Anthropic", base_url: "https://api.anthropic.com", @@ -223,7 +243,7 @@ mod tests { #[tokio::test] async fn list_empty() { let (repo, _db) = setup().await; - let providers = repo.list().await.unwrap(); + let providers = repo.list(USER_A).await.unwrap(); assert!(providers.is_empty()); } @@ -259,7 +279,7 @@ mod tests { assert_eq!(p.id, "my-custom-id-1"); assert_eq!(p.platform, "anthropic"); - let found = repo.find_by_id("my-custom-id-1").await.unwrap().unwrap(); + let found = repo.find_by_id(USER_A, "my-custom-id-1").await.unwrap().unwrap(); assert_eq!(found.id, "my-custom-id-1"); } @@ -288,7 +308,7 @@ mod tests { let (repo, _db) = setup().await; let created = repo.create(sample_params()).await.unwrap(); - let found = repo.find_by_id(&created.id).await.unwrap().unwrap(); + let found = repo.find_by_id(USER_A, &created.id).await.unwrap().unwrap(); assert_eq!(found.id, created.id); assert_eq!(found.platform, "anthropic"); assert_eq!(found.models, r#"["claude-sonnet-4-20250514"]"#); @@ -297,7 +317,7 @@ mod tests { #[tokio::test] async fn find_by_id_nonexistent() { let (repo, _db) = setup().await; - assert!(repo.find_by_id("no_such_id").await.unwrap().is_none()); + assert!(repo.find_by_id(USER_A, "no_such_id").await.unwrap().is_none()); } #[tokio::test] @@ -314,7 +334,7 @@ mod tests { .await .unwrap(); - let all = repo.list().await.unwrap(); + let all = repo.list(USER_A).await.unwrap(); assert_eq!(all.len(), 2); assert_eq!(all[0].id, p1.id); assert_eq!(all[1].id, p2.id); @@ -327,6 +347,7 @@ mod tests { let updated = repo .update( + USER_A, &created.id, UpdateProviderParams { name: Some("Anthropic Updated"), @@ -352,6 +373,7 @@ mod tests { let updated = repo .update( + USER_A, &created.id, UpdateProviderParams { api_key_encrypted: Some("new_encrypted_key"), @@ -367,7 +389,10 @@ mod tests { #[tokio::test] async fn update_nonexistent_returns_not_found() { let (repo, _db) = setup().await; - let err = repo.update("no_id", UpdateProviderParams::default()).await.unwrap_err(); + let err = repo + .update(USER_A, "no_id", UpdateProviderParams::default()) + .await + .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -380,6 +405,7 @@ mod tests { // Set optional field let updated = repo .update( + USER_A, &created.id, UpdateProviderParams { model_protocols: Some(Some(r#"{"model1":"openai"}"#)), @@ -396,6 +422,7 @@ mod tests { // Clear optional field let cleared = repo .update( + USER_A, &created.id, UpdateProviderParams { model_protocols: Some(None), @@ -415,14 +442,14 @@ mod tests { let (repo, _db) = setup().await; let created = repo.create(sample_params()).await.unwrap(); - repo.delete(&created.id).await.unwrap(); - assert!(repo.find_by_id(&created.id).await.unwrap().is_none()); + repo.delete(USER_A, &created.id).await.unwrap(); + assert!(repo.find_by_id(USER_A, &created.id).await.unwrap().is_none()); } #[tokio::test] async fn delete_nonexistent_returns_not_found() { let (repo, _db) = setup().await; - let err = repo.delete("no_id").await.unwrap_err(); + let err = repo.delete(USER_A, "no_id").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -438,10 +465,46 @@ mod tests { .await .unwrap(); - repo.delete(&p1.id).await.unwrap(); + repo.delete(USER_A, &p1.id).await.unwrap(); - let all = repo.list().await.unwrap(); + let all = repo.list(USER_A).await.unwrap(); assert_eq!(all.len(), 1); assert_eq!(all[0].id, p2.id); } + + #[tokio::test] + async fn provider_operations_are_scoped_by_user() { + let (repo, _db) = setup().await; + let provider_a = repo.create(sample_params()).await.unwrap(); + let provider_b = repo + .create(CreateProviderParams { + id: Some("same-visible-provider-id"), + user_id: USER_B, + name: "Other User Provider", + ..sample_params() + }) + .await + .unwrap(); + + assert_eq!(repo.list(USER_A).await.unwrap().len(), 1); + assert_eq!(repo.list(USER_B).await.unwrap().len(), 1); + assert!(repo.find_by_id(USER_B, &provider_a.id).await.unwrap().is_none()); + + let err = repo + .update( + USER_B, + &provider_a.id, + UpdateProviderParams { + name: Some("cross-user update"), + ..Default::default() + }, + ) + .await + .unwrap_err(); + assert!(matches!(err, DbError::NotFound(_))); + + repo.delete(USER_B, &provider_b.id).await.unwrap(); + assert_eq!(repo.list(USER_A).await.unwrap().len(), 1); + assert!(repo.list(USER_B).await.unwrap().is_empty()); + } } diff --git a/crates/aionui-db/src/repository/sqlite_remote_agent.rs b/crates/aionui-db/src/repository/sqlite_remote_agent.rs index c98897c5b..e7fbc4867 100644 --- a/crates/aionui-db/src/repository/sqlite_remote_agent.rs +++ b/crates/aionui-db/src/repository/sqlite_remote_agent.rs @@ -19,16 +19,20 @@ impl SqliteRemoteAgentRepository { #[async_trait::async_trait] impl IRemoteAgentRepository for SqliteRemoteAgentRepository { - async fn list(&self) -> Result, DbError> { - let rows = sqlx::query_as::<_, RemoteAgentRow>("SELECT * FROM remote_agents ORDER BY created_at ASC") - .fetch_all(&self.pool) - .await?; + async fn list(&self, user_id: &str) -> Result, DbError> { + let rows = sqlx::query_as::<_, RemoteAgentRow>( + "SELECT * FROM remote_agents WHERE user_id = ? ORDER BY created_at ASC", + ) + .bind(user_id) + .fetch_all(&self.pool) + .await?; Ok(rows) } - async fn find_by_id(&self, id: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, RemoteAgentRow>("SELECT * FROM remote_agents WHERE id = ?") + async fn find_by_id(&self, user_id: &str, id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, RemoteAgentRow>("SELECT * FROM remote_agents WHERE user_id = ? AND id = ?") + .bind(user_id) .bind(id) .fetch_optional(&self.pool) .await?; @@ -43,12 +47,13 @@ impl IRemoteAgentRepository for SqliteRemoteAgentRepository { sqlx::query( "INSERT INTO remote_agents \ - (id, name, protocol, url, auth_type, auth_token, allow_insecure, \ + (id, user_id, name, protocol, url, auth_type, auth_token, allow_insecure, \ avatar, description, device_id, device_public_key, device_private_key, \ device_token, status, last_connected_at, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(&id) + .bind(params.user_id) .bind(params.name) .bind(params.protocol) .bind(params.url) @@ -70,6 +75,7 @@ impl IRemoteAgentRepository for SqliteRemoteAgentRepository { Ok(RemoteAgentRow { id, + user_id: params.user_id.to_string(), name: params.name.to_string(), protocol: params.protocol.to_string(), url: params.url.to_string(), @@ -89,9 +95,14 @@ impl IRemoteAgentRepository for SqliteRemoteAgentRepository { }) } - async fn update(&self, id: &str, params: UpdateRemoteAgentParams<'_>) -> Result { + async fn update( + &self, + user_id: &str, + id: &str, + params: UpdateRemoteAgentParams<'_>, + ) -> Result { let existing = self - .find_by_id(id) + .find_by_id(user_id, id) .await? .ok_or_else(|| DbError::NotFound(format!("Remote agent '{id}' not found")))?; @@ -101,7 +112,7 @@ impl IRemoteAgentRepository for SqliteRemoteAgentRepository { "UPDATE remote_agents SET \ name = ?, protocol = ?, url = ?, auth_type = ?, auth_token = ?, \ allow_insecure = ?, avatar = ?, description = ?, updated_at = ? \ - WHERE id = ?", + WHERE user_id = ? AND id = ?", ) .bind(&merged.name) .bind(&merged.protocol) @@ -112,6 +123,7 @@ impl IRemoteAgentRepository for SqliteRemoteAgentRepository { .bind(&merged.avatar) .bind(&merged.description) .bind(merged.updated_at) + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -119,8 +131,9 @@ impl IRemoteAgentRepository for SqliteRemoteAgentRepository { Ok(merged) } - async fn delete(&self, id: &str) -> Result<(), DbError> { - let result = sqlx::query("DELETE FROM remote_agents WHERE id = ?") + async fn delete(&self, user_id: &str, id: &str) -> Result<(), DbError> { + let result = sqlx::query("DELETE FROM remote_agents WHERE user_id = ? AND id = ?") + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -134,6 +147,7 @@ impl IRemoteAgentRepository for SqliteRemoteAgentRepository { async fn update_status( &self, + user_id: &str, id: &str, status: &str, last_connected_at: Option, @@ -143,11 +157,12 @@ impl IRemoteAgentRepository for SqliteRemoteAgentRepository { let result = sqlx::query( "UPDATE remote_agents SET status = ?, \ last_connected_at = COALESCE(?, last_connected_at), \ - updated_at = ? WHERE id = ?", + updated_at = ? WHERE user_id = ? AND id = ?", ) .bind(status) .bind(last_connected_at) .bind(now) + .bind(user_id) .bind(id) .execute(&self.pool) .await?; @@ -165,6 +180,7 @@ fn merge_update(existing: RemoteAgentRow, params: UpdateRemoteAgentParams<'_>) - let now = aionui_common::now_ms(); RemoteAgentRow { id: existing.id, + user_id: existing.user_id, name: params.name.unwrap_or(&existing.name).to_string(), protocol: params.protocol.unwrap_or(&existing.protocol).to_string(), url: params.url.unwrap_or(&existing.url).to_string(), @@ -190,14 +206,27 @@ mod tests { use super::*; use crate::init_database_memory; + const USER_A: &str = "system_default_user"; + const USER_B: &str = "user_b"; + async fn setup() -> (SqliteRemoteAgentRepository, crate::Database) { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, 1, 1)", + ) + .bind(USER_B) + .bind(USER_B) + .execute(db.pool()) + .await + .unwrap(); let repo = SqliteRemoteAgentRepository::new(db.pool().clone()); (repo, db) } fn sample_params() -> CreateRemoteAgentParams<'static> { CreateRemoteAgentParams { + user_id: USER_A, name: "Test Agent", protocol: "acp", url: "wss://remote.example.com", @@ -216,7 +245,7 @@ mod tests { #[tokio::test] async fn list_empty() { let (repo, _db) = setup().await; - let agents = repo.list().await.unwrap(); + let agents = repo.list(USER_A).await.unwrap(); assert!(agents.is_empty()); } @@ -267,7 +296,7 @@ mod tests { let (repo, _db) = setup().await; let created = repo.create(sample_params()).await.unwrap(); - let found = repo.find_by_id(&created.id).await.unwrap().unwrap(); + let found = repo.find_by_id(USER_A, &created.id).await.unwrap().unwrap(); assert_eq!(found.id, created.id); assert_eq!(found.name, "Test Agent"); assert_eq!(found.protocol, "acp"); @@ -277,7 +306,7 @@ mod tests { #[tokio::test] async fn find_by_id_nonexistent() { let (repo, _db) = setup().await; - assert!(repo.find_by_id("no_such_id").await.unwrap().is_none()); + assert!(repo.find_by_id(USER_A, "no_such_id").await.unwrap().is_none()); } #[tokio::test] @@ -292,7 +321,7 @@ mod tests { .await .unwrap(); - let all = repo.list().await.unwrap(); + let all = repo.list(USER_A).await.unwrap(); assert_eq!(all.len(), 2); assert_eq!(all[0].id, a1.id); assert_eq!(all[1].id, a2.id); @@ -305,6 +334,7 @@ mod tests { let updated = repo .update( + USER_A, &created.id, UpdateRemoteAgentParams { name: Some("Updated Agent"), @@ -332,6 +362,7 @@ mod tests { let updated = repo .update( + USER_A, &created.id, UpdateRemoteAgentParams { description: Some(None), @@ -350,7 +381,7 @@ mod tests { async fn update_nonexistent_returns_not_found() { let (repo, _db) = setup().await; let err = repo - .update("no_id", UpdateRemoteAgentParams::default()) + .update(USER_A, "no_id", UpdateRemoteAgentParams::default()) .await .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); @@ -361,14 +392,14 @@ mod tests { let (repo, _db) = setup().await; let created = repo.create(sample_params()).await.unwrap(); - repo.delete(&created.id).await.unwrap(); - assert!(repo.find_by_id(&created.id).await.unwrap().is_none()); + repo.delete(USER_A, &created.id).await.unwrap(); + assert!(repo.find_by_id(USER_A, &created.id).await.unwrap().is_none()); } #[tokio::test] async fn delete_nonexistent_returns_not_found() { let (repo, _db) = setup().await; - let err = repo.delete("no_id").await.unwrap_err(); + let err = repo.delete(USER_A, "no_id").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -380,11 +411,11 @@ mod tests { assert!(created.last_connected_at.is_none()); let connect_time = aionui_common::now_ms(); - repo.update_status(&created.id, "connected", Some(connect_time)) + repo.update_status(USER_A, &created.id, "connected", Some(connect_time)) .await .unwrap(); - let found = repo.find_by_id(&created.id).await.unwrap().unwrap(); + let found = repo.find_by_id(USER_A, &created.id).await.unwrap().unwrap(); assert_eq!(found.status, "connected"); assert_eq!(found.last_connected_at, Some(connect_time)); assert!(found.updated_at >= created.updated_at); @@ -397,14 +428,14 @@ mod tests { // First set a connected timestamp let connect_time = aionui_common::now_ms(); - repo.update_status(&created.id, "connected", Some(connect_time)) + repo.update_status(USER_A, &created.id, "connected", Some(connect_time)) .await .unwrap(); // Now update status to error without providing last_connected_at - repo.update_status(&created.id, "error", None).await.unwrap(); + repo.update_status(USER_A, &created.id, "error", None).await.unwrap(); - let found = repo.find_by_id(&created.id).await.unwrap().unwrap(); + let found = repo.find_by_id(USER_A, &created.id).await.unwrap().unwrap(); assert_eq!(found.status, "error"); // COALESCE preserves the existing last_connected_at assert_eq!(found.last_connected_at, Some(connect_time)); @@ -416,9 +447,9 @@ mod tests { let created = repo.create(sample_params()).await.unwrap(); assert!(created.last_connected_at.is_none()); - repo.update_status(&created.id, "error", None).await.unwrap(); + repo.update_status(USER_A, &created.id, "error", None).await.unwrap(); - let found = repo.find_by_id(&created.id).await.unwrap().unwrap(); + let found = repo.find_by_id(USER_A, &created.id).await.unwrap().unwrap(); assert_eq!(found.status, "error"); assert!(found.last_connected_at.is_none()); } @@ -426,7 +457,10 @@ mod tests { #[tokio::test] async fn update_status_nonexistent_returns_not_found() { let (repo, _db) = setup().await; - let err = repo.update_status("no_id", "connected", None).await.unwrap_err(); + let err = repo + .update_status(USER_A, "no_id", "connected", None) + .await + .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -442,10 +476,38 @@ mod tests { .await .unwrap(); - repo.delete(&a1.id).await.unwrap(); + repo.delete(USER_A, &a1.id).await.unwrap(); - let all = repo.list().await.unwrap(); + let all = repo.list(USER_A).await.unwrap(); assert_eq!(all.len(), 1); assert_eq!(all[0].id, a2.id); } + + #[tokio::test] + async fn remote_agent_operations_are_scoped_by_user() { + let (repo, _db) = setup().await; + let agent_a = repo.create(sample_params()).await.unwrap(); + let agent_b = repo + .create(CreateRemoteAgentParams { + user_id: USER_B, + name: "User B Agent", + ..sample_params() + }) + .await + .unwrap(); + + assert_eq!(repo.list(USER_A).await.unwrap().len(), 1); + assert_eq!(repo.list(USER_B).await.unwrap().len(), 1); + assert!(repo.find_by_id(USER_B, &agent_a.id).await.unwrap().is_none()); + + let err = repo + .update_status(USER_B, &agent_a.id, "connected", Some(aionui_common::now_ms())) + .await + .unwrap_err(); + assert!(matches!(err, DbError::NotFound(_))); + + repo.delete(USER_B, &agent_b.id).await.unwrap(); + assert_eq!(repo.list(USER_A).await.unwrap().len(), 1); + assert!(repo.list(USER_B).await.unwrap().is_empty()); + } } diff --git a/crates/aionui-db/src/repository/sqlite_settings.rs b/crates/aionui-db/src/repository/sqlite_settings.rs index ffcd15c85..519ad0022 100644 --- a/crates/aionui-db/src/repository/sqlite_settings.rs +++ b/crates/aionui-db/src/repository/sqlite_settings.rs @@ -18,8 +18,9 @@ impl SqliteSettingsRepository { #[async_trait::async_trait] impl ISettingsRepository for SqliteSettingsRepository { - async fn get_settings(&self) -> Result, DbError> { - let row = sqlx::query_as::<_, SystemSettings>("SELECT * FROM system_settings WHERE id = 1") + async fn get_settings(&self, user_id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, SystemSettings>("SELECT * FROM system_settings WHERE user_id = ?") + .bind(user_id) .fetch_optional(&self.pool) .await?; @@ -28,6 +29,7 @@ impl ISettingsRepository for SqliteSettingsRepository { async fn upsert_settings( &self, + user_id: &str, language: &str, notification_enabled: bool, cron_notification_enabled: bool, @@ -38,10 +40,10 @@ impl ISettingsRepository for SqliteSettingsRepository { sqlx::query( "INSERT INTO system_settings \ - (id, language, notification_enabled, cron_notification_enabled, \ + (user_id, language, notification_enabled, cron_notification_enabled, \ command_queue_enabled, save_upload_to_workspace, updated_at) \ - VALUES (1, ?, ?, ?, ?, ?, ?) \ - ON CONFLICT(id) DO UPDATE SET \ + VALUES (?, ?, ?, ?, ?, ?, ?) \ + ON CONFLICT(user_id) DO UPDATE SET \ language = excluded.language, \ notification_enabled = excluded.notification_enabled, \ cron_notification_enabled = excluded.cron_notification_enabled, \ @@ -49,6 +51,7 @@ impl ISettingsRepository for SqliteSettingsRepository { save_upload_to_workspace = excluded.save_upload_to_workspace, \ updated_at = excluded.updated_at", ) + .bind(user_id) .bind(language) .bind(notification_enabled) .bind(cron_notification_enabled) @@ -59,7 +62,7 @@ impl ISettingsRepository for SqliteSettingsRepository { .await?; Ok(SystemSettings { - id: 1, + user_id: user_id.to_string(), language: language.to_string(), notification_enabled, cron_notification_enabled, @@ -75,8 +78,20 @@ mod tests { use super::*; use crate::init_database_memory; + const USER_A: &str = "system_default_user"; + const USER_B: &str = "user_b"; + async fn setup() -> (SqliteSettingsRepository, crate::Database) { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, 1, 1)", + ) + .bind(USER_B) + .bind(USER_B) + .execute(db.pool()) + .await + .unwrap(); let repo = SqliteSettingsRepository::new(db.pool().clone()); (repo, db) } @@ -84,15 +99,18 @@ mod tests { #[tokio::test] async fn get_settings_returns_none_when_empty() { let (repo, _db) = setup().await; - assert!(repo.get_settings().await.unwrap().is_none()); + assert!(repo.get_settings(USER_A).await.unwrap().is_none()); } #[tokio::test] async fn upsert_creates_settings() { let (repo, _db) = setup().await; - let s = repo.upsert_settings("zh-CN", false, true, true, false).await.unwrap(); + let s = repo + .upsert_settings(USER_A, "zh-CN", false, true, true, false) + .await + .unwrap(); - assert_eq!(s.id, 1); + assert_eq!(s.user_id, USER_A); assert_eq!(s.language, "zh-CN"); assert!(!s.notification_enabled); assert!(s.cron_notification_enabled); @@ -104,9 +122,11 @@ mod tests { #[tokio::test] async fn upsert_then_get_returns_same() { let (repo, _db) = setup().await; - repo.upsert_settings("en-US", true, false, false, true).await.unwrap(); + repo.upsert_settings(USER_A, "en-US", true, false, false, true) + .await + .unwrap(); - let s = repo.get_settings().await.unwrap().unwrap(); + let s = repo.get_settings(USER_A).await.unwrap().unwrap(); assert_eq!(s.language, "en-US"); assert!(s.notification_enabled); assert!(!s.cron_notification_enabled); @@ -117,8 +137,13 @@ mod tests { #[tokio::test] async fn upsert_overwrites_existing() { let (repo, _db) = setup().await; - repo.upsert_settings("en-US", true, false, false, false).await.unwrap(); - let s = repo.upsert_settings("ja-JP", false, true, true, true).await.unwrap(); + repo.upsert_settings(USER_A, "en-US", true, false, false, false) + .await + .unwrap(); + let s = repo + .upsert_settings(USER_A, "ja-JP", false, true, true, true) + .await + .unwrap(); assert_eq!(s.language, "ja-JP"); assert!(!s.notification_enabled); @@ -127,7 +152,25 @@ mod tests { assert!(s.save_upload_to_workspace); // Verify persisted via get - let fetched = repo.get_settings().await.unwrap().unwrap(); + let fetched = repo.get_settings(USER_A).await.unwrap().unwrap(); assert_eq!(fetched.language, "ja-JP"); } + + #[tokio::test] + async fn settings_are_scoped_by_user() { + let (repo, _db) = setup().await; + repo.upsert_settings(USER_A, "en-US", true, false, false, false) + .await + .unwrap(); + repo.upsert_settings(USER_B, "zh-CN", false, true, true, true) + .await + .unwrap(); + + let a = repo.get_settings(USER_A).await.unwrap().unwrap(); + let b = repo.get_settings(USER_B).await.unwrap().unwrap(); + assert_eq!(a.language, "en-US"); + assert_eq!(b.language, "zh-CN"); + assert!(a.notification_enabled); + assert!(!b.notification_enabled); + } } diff --git a/crates/aionui-db/src/repository/sqlite_skill.rs b/crates/aionui-db/src/repository/sqlite_skill.rs index 72281d8b3..a40452e68 100644 --- a/crates/aionui-db/src/repository/sqlite_skill.rs +++ b/crates/aionui-db/src/repository/sqlite_skill.rs @@ -4,6 +4,8 @@ use crate::error::DbError; use crate::models::{SkillImportRecordRow, SkillRow}; use crate::repository::skill::{CreateSkillImportRecordParams, ISkillRepository, UpsertSkillParams}; +const DEFAULT_USER_ID: &str = "system_default_user"; + /// SQLite-backed implementation of [`ISkillRepository`]. #[derive(Clone, Debug)] pub struct SqliteSkillRepository { @@ -19,34 +21,68 @@ impl SqliteSkillRepository { #[async_trait::async_trait] impl ISkillRepository for SqliteSkillRepository { async fn list(&self) -> Result, DbError> { + self.list_for_user(DEFAULT_USER_ID).await + } + + async fn list_for_user(&self, user_id: &str) -> Result, DbError> { let rows = sqlx::query_as::<_, SkillRow>( - "SELECT * FROM skills WHERE deleted_at IS NULL AND enabled = 1 ORDER BY updated_at DESC, name ASC", + "SELECT * FROM skills \ + WHERE (user_id IS NULL OR user_id = ?) AND deleted_at IS NULL AND enabled = 1 \ + ORDER BY updated_at DESC, name ASC", ) + .bind(user_id) .fetch_all(&self.pool) .await?; Ok(rows) } async fn find_by_name(&self, name: &str) -> Result, DbError> { - let row = - sqlx::query_as::<_, SkillRow>("SELECT * FROM skills WHERE name = ? AND deleted_at IS NULL AND enabled = 1") - .bind(name) - .fetch_optional(&self.pool) - .await?; + self.find_by_name_for_user(DEFAULT_USER_ID, name).await + } + + async fn find_by_name_for_user(&self, user_id: &str, name: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, SkillRow>( + "SELECT * FROM skills \ + WHERE (user_id IS NULL OR user_id = ?) AND name = ? AND deleted_at IS NULL AND enabled = 1 \ + ORDER BY user_id IS NULL ASC \ + LIMIT 1", + ) + .bind(user_id) + .bind(name) + .fetch_optional(&self.pool) + .await?; Ok(row) } async fn find_by_name_any(&self, name: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, SkillRow>("SELECT * FROM skills WHERE name = ?") - .bind(name) - .fetch_optional(&self.pool) - .await?; + self.find_by_name_any_for_user(DEFAULT_USER_ID, name).await + } + + async fn find_by_name_any_for_user(&self, user_id: &str, name: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, SkillRow>( + "SELECT * FROM skills \ + WHERE (user_id IS NULL OR user_id = ?) AND name = ? \ + ORDER BY user_id IS NULL ASC \ + LIMIT 1", + ) + .bind(user_id) + .bind(name) + .fetch_optional(&self.pool) + .await?; Ok(row) } async fn upsert(&self, params: UpsertSkillParams<'_>) -> Result { + self.upsert_for_user(DEFAULT_USER_ID, params).await + } + + async fn upsert_for_user(&self, user_id: &str, params: UpsertSkillParams<'_>) -> Result { let now = aionui_common::now_ms(); - let existing = self.find_by_name_any(params.name).await?; + let existing = sqlx::query_as::<_, SkillRow>("SELECT * FROM skills WHERE user_id = ? AND name = ?") + .bind(user_id) + .bind(params.name) + .fetch_optional(&self.pool) + .await?; let id = existing .as_ref() .map(|row| row.id.clone()) @@ -55,9 +91,9 @@ impl ISkillRepository for SqliteSkillRepository { sqlx::query( "INSERT INTO skills \ - (id, name, description, path, source, enabled, deleted_at, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, NULL, ?, ?) \ - ON CONFLICT(name) DO UPDATE SET \ + (id, user_id, name, description, path, source, enabled, deleted_at, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?, NULL, ?, ?) \ + ON CONFLICT(user_id, name) WHERE user_id IS NOT NULL DO UPDATE SET \ description = excluded.description, \ path = excluded.path, \ source = excluded.source, \ @@ -66,6 +102,7 @@ impl ISkillRepository for SqliteSkillRepository { updated_at = excluded.updated_at", ) .bind(&id) + .bind(user_id) .bind(params.name) .bind(params.description) .bind(params.path) @@ -76,18 +113,66 @@ impl ISkillRepository for SqliteSkillRepository { .execute(&self.pool) .await?; - self.find_by_name_any(params.name) + self.find_by_name_any_for_user(user_id, params.name) .await? .ok_or_else(|| DbError::NotFound(format!("skill '{}' was not found after upsert", params.name))) } + async fn upsert_global(&self, params: UpsertSkillParams<'_>) -> Result { + let now = aionui_common::now_ms(); + let existing = sqlx::query_as::<_, SkillRow>("SELECT * FROM skills WHERE user_id IS NULL AND name = ?") + .bind(params.name) + .fetch_optional(&self.pool) + .await?; + let id = existing + .as_ref() + .map(|row| row.id.clone()) + .unwrap_or_else(|| aionui_common::generate_prefixed_id("skill")); + let created_at = existing.as_ref().map(|row| row.created_at).unwrap_or(now); + + sqlx::query( + "INSERT INTO skills \ + (id, user_id, name, description, path, source, enabled, deleted_at, created_at, updated_at) \ + VALUES (?, NULL, ?, ?, ?, ?, ?, NULL, ?, ?) \ + ON CONFLICT(name) WHERE user_id IS NULL DO UPDATE SET \ + description = excluded.description, \ + path = excluded.path, \ + source = excluded.source, \ + enabled = excluded.enabled, \ + deleted_at = NULL, \ + updated_at = excluded.updated_at", + ) + .bind(&id) + .bind(params.name) + .bind(params.description) + .bind(params.path) + .bind(params.source) + .bind(params.enabled) + .bind(created_at) + .bind(now) + .execute(&self.pool) + .await?; + + sqlx::query_as::<_, SkillRow>("SELECT * FROM skills WHERE user_id IS NULL AND name = ?") + .bind(params.name) + .fetch_optional(&self.pool) + .await? + .ok_or_else(|| DbError::NotFound(format!("global skill '{}' was not found after upsert", params.name))) + } + async fn delete_by_name(&self, name: &str) -> Result { + self.delete_by_name_for_user(DEFAULT_USER_ID, name).await + } + + async fn delete_by_name_for_user(&self, user_id: &str, name: &str) -> Result { let now = aionui_common::now_ms(); let result = sqlx::query( - "UPDATE skills SET enabled = 0, deleted_at = ?, updated_at = ? WHERE name = ? AND deleted_at IS NULL", + "UPDATE skills SET enabled = 0, deleted_at = ?, updated_at = ? \ + WHERE user_id = ? AND name = ? AND deleted_at IS NULL", ) .bind(now) .bind(now) + .bind(user_id) .bind(name) .execute(&self.pool) .await?; @@ -96,7 +181,7 @@ impl ISkillRepository for SqliteSkillRepository { return Err(DbError::NotFound(format!("skill '{name}'"))); } - self.find_by_name_any(name) + self.find_by_name_any_for_user(user_id, name) .await? .ok_or_else(|| DbError::NotFound(format!("skill '{name}'"))) } @@ -104,6 +189,14 @@ impl ISkillRepository for SqliteSkillRepository { async fn create_import_record( &self, params: CreateSkillImportRecordParams<'_>, + ) -> Result { + self.create_import_record_for_user(DEFAULT_USER_ID, params).await + } + + async fn create_import_record_for_user( + &self, + user_id: &str, + params: CreateSkillImportRecordParams<'_>, ) -> Result { let id = aionui_common::generate_prefixed_id("skill_import"); let now = aionui_common::now_ms(); @@ -111,8 +204,8 @@ impl ISkillRepository for SqliteSkillRepository { sqlx::query( "INSERT INTO skill_import_records \ (id, operation_id, source_label, source_path, source_name, skill_id, skill_name, \ - status, error_code, error_path, actual_bytes, limit_bytes, line, column, created_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + status, error_code, error_path, actual_bytes, limit_bytes, line, column, created_at, user_id) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(&id) .bind(params.operation_id) @@ -129,6 +222,7 @@ impl ISkillRepository for SqliteSkillRepository { .bind(params.line) .bind(params.column) .bind(now) + .bind(user_id) .execute(&self.pool) .await?; @@ -140,9 +234,18 @@ impl ISkillRepository for SqliteSkillRepository { } async fn list_import_records(&self, limit: i64) -> Result, DbError> { + self.list_import_records_for_user(DEFAULT_USER_ID, limit).await + } + + async fn list_import_records_for_user( + &self, + user_id: &str, + limit: i64, + ) -> Result, DbError> { let rows = sqlx::query_as::<_, SkillImportRecordRow>( - "SELECT * FROM skill_import_records ORDER BY created_at DESC, id DESC LIMIT ?", + "SELECT * FROM skill_import_records WHERE user_id = ? ORDER BY created_at DESC, id DESC LIMIT ?", ) + .bind(user_id) .bind(limit.max(0)) .fetch_all(&self.pool) .await?; @@ -155,12 +258,28 @@ mod tests { use super::*; use crate::init_database_memory; + const USER_A: &str = "system_default_user"; + const USER_B: &str = "user_b"; + async fn setup() -> (SqliteSkillRepository, crate::Database) { let db = init_database_memory().await.unwrap(); let repo = SqliteSkillRepository::new(db.pool().clone()); + ensure_user(&db, USER_B).await; (repo, db) } + async fn ensure_user(db: &crate::Database, user_id: &str) { + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, 'hash', 'active', 0, 1, 1)", + ) + .bind(user_id) + .bind(user_id) + .execute(db.pool()) + .await + .unwrap(); + } + #[tokio::test] async fn upsert_restores_soft_deleted_skill() { let (repo, _db) = setup().await; @@ -254,4 +373,148 @@ mod tests { let records = repo.list_import_records(10).await.unwrap(); assert_eq!(records.len(), 1); } + + #[tokio::test] + async fn skill_operations_are_scoped_by_user() { + let (repo, _db) = setup().await; + + let first = repo + .upsert_for_user( + USER_A, + UpsertSkillParams { + name: "shared", + description: Some("A"), + path: "/tmp/a", + source: "user", + enabled: true, + }, + ) + .await + .unwrap(); + let second = repo + .upsert_for_user( + USER_B, + UpsertSkillParams { + name: "shared", + description: Some("B"), + path: "/tmp/b", + source: "user", + enabled: true, + }, + ) + .await + .unwrap(); + + assert_ne!(first.id, second.id); + assert_eq!( + repo.find_by_name_for_user(USER_A, "shared") + .await + .unwrap() + .unwrap() + .path, + "/tmp/a" + ); + assert_eq!( + repo.find_by_name_for_user(USER_B, "shared") + .await + .unwrap() + .unwrap() + .path, + "/tmp/b" + ); + + repo.delete_by_name_for_user(USER_A, "shared").await.unwrap(); + + assert!(repo.find_by_name_for_user(USER_A, "shared").await.unwrap().is_none()); + assert!(repo.find_by_name_for_user(USER_B, "shared").await.unwrap().is_some()); + } + + #[tokio::test] + async fn global_builtin_skill_is_visible_to_all_users_and_user_skill_overrides_it() { + let (repo, _db) = setup().await; + + repo.upsert_global(UpsertSkillParams { + name: "shared", + description: Some("Global"), + path: "/tmp/global", + source: "builtin", + enabled: true, + }) + .await + .unwrap(); + + assert_eq!( + repo.find_by_name_for_user(USER_B, "shared") + .await + .unwrap() + .unwrap() + .path, + "/tmp/global" + ); + + repo.upsert_for_user( + USER_B, + UpsertSkillParams { + name: "shared", + description: Some("B"), + path: "/tmp/b", + source: "user", + enabled: true, + }, + ) + .await + .unwrap(); + + assert_eq!( + repo.find_by_name_for_user(USER_B, "shared") + .await + .unwrap() + .unwrap() + .path, + "/tmp/b" + ); + assert_eq!( + repo.find_by_name_for_user(USER_A, "shared") + .await + .unwrap() + .unwrap() + .path, + "/tmp/global" + ); + } + + #[tokio::test] + async fn import_records_are_scoped_by_user() { + let (repo, _db) = setup().await; + + for user_id in [USER_A, USER_B] { + repo.create_import_record_for_user( + user_id, + CreateSkillImportRecordParams { + operation_id: user_id, + source_label: "pack", + source_path: None, + source_name: "skill", + skill_id: None, + skill_name: None, + status: "imported", + error_code: None, + error_path: None, + actual_bytes: None, + limit_bytes: None, + line: None, + column: None, + }, + ) + .await + .unwrap(); + } + + let a_records = repo.list_import_records_for_user(USER_A, 10).await.unwrap(); + let b_records = repo.list_import_records_for_user(USER_B, 10).await.unwrap(); + assert_eq!(a_records.len(), 1); + assert_eq!(a_records[0].operation_id, USER_A); + assert_eq!(b_records.len(), 1); + assert_eq!(b_records[0].operation_id, USER_B); + } } diff --git a/crates/aionui-db/src/repository/sqlite_team.rs b/crates/aionui-db/src/repository/sqlite_team.rs index 63990280c..1ab168274 100644 --- a/crates/aionui-db/src/repository/sqlite_team.rs +++ b/crates/aionui-db/src/repository/sqlite_team.rs @@ -42,7 +42,7 @@ impl ITeamRepository for SqliteTeamRepository { Ok(()) } - async fn list_teams(&self) -> Result, DbError> { + async fn list_teams_for_restore(&self) -> Result, DbError> { let rows = sqlx::query_as::<_, TeamRow>("SELECT * FROM teams ORDER BY created_at ASC") .fetch_all(&self.pool) .await?; @@ -57,7 +57,16 @@ impl ITeamRepository for SqliteTeamRepository { Ok(rows) } - async fn get_team(&self, team_id: &str) -> Result, DbError> { + async fn get_team(&self, user_id: &str, team_id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, TeamRow>("SELECT * FROM teams WHERE user_id = ? AND id = ?") + .bind(user_id) + .bind(team_id) + .fetch_optional(&self.pool) + .await?; + Ok(row) + } + + async fn get_team_for_restore(&self, team_id: &str) -> Result, DbError> { let row = sqlx::query_as::<_, TeamRow>("SELECT * FROM teams WHERE id = ?") .bind(team_id) .fetch_optional(&self.pool) @@ -65,7 +74,7 @@ impl ITeamRepository for SqliteTeamRepository { Ok(row) } - async fn update_team(&self, team_id: &str, params: &UpdateTeamParams) -> Result<(), DbError> { + async fn update_team(&self, user_id: &str, team_id: &str, params: &UpdateTeamParams) -> Result<(), DbError> { let mut set_clauses = Vec::new(); if params.name.is_some() { set_clauses.push("name = ?"); @@ -88,7 +97,10 @@ impl ITeamRepository for SqliteTeamRepository { } set_clauses.push("updated_at = ?"); - let sql = format!("UPDATE teams SET {} WHERE id = ?", set_clauses.join(", ")); + let sql = format!( + "UPDATE teams SET {} WHERE user_id = ? AND id = ?", + set_clauses.join(", ") + ); let mut query = sqlx::query(&sql); if let Some(ref name) = params.name { @@ -107,6 +119,7 @@ impl ITeamRepository for SqliteTeamRepository { query = query.bind(session_mode); } query = query.bind(now_ms()); + query = query.bind(user_id); query = query.bind(team_id); let result = query.execute(&self.pool).await?; @@ -116,8 +129,9 @@ impl ITeamRepository for SqliteTeamRepository { Ok(()) } - async fn delete_team(&self, team_id: &str) -> Result<(), DbError> { - let result = sqlx::query("DELETE FROM teams WHERE id = ?") + async fn delete_team(&self, user_id: &str, team_id: &str) -> Result<(), DbError> { + let result = sqlx::query("DELETE FROM teams WHERE user_id = ? AND id = ?") + .bind(user_id) .bind(team_id) .execute(&self.pool) .await?; @@ -129,11 +143,12 @@ impl ITeamRepository for SqliteTeamRepository { // ── Mailbox ────────────────────────────────────────────────────── - async fn write_message(&self, row: &MailboxMessageRow) -> Result<(), DbError> { - sqlx::query( + async fn write_message(&self, user_id: &str, row: &MailboxMessageRow) -> Result<(), DbError> { + let result = sqlx::query( "INSERT INTO mailbox \ (id, team_id, to_agent_id, from_agent_id, type, content, summary, files, read, created_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + SELECT ?, ?, ?, ?, ?, ?, ?, ?, ?, ? \ + WHERE EXISTS (SELECT 1 FROM teams t WHERE t.id = ? AND t.user_id = ?)", ) .bind(&row.id) .bind(&row.team_id) @@ -145,12 +160,22 @@ impl ITeamRepository for SqliteTeamRepository { .bind(&row.files) .bind(row.read) .bind(row.created_at) + .bind(&row.team_id) + .bind(user_id) .execute(&self.pool) .await?; + if result.rows_affected() == 0 { + return Err(DbError::NotFound(format!("team {}", row.team_id))); + } Ok(()) } - async fn read_unread_and_mark(&self, team_id: &str, to_agent_id: &str) -> Result, DbError> { + async fn read_unread_and_mark( + &self, + user_id: &str, + team_id: &str, + to_agent_id: &str, + ) -> Result, DbError> { // Use BEGIN IMMEDIATE for atomicity: prevents concurrent readers // from seeing the same unread messages. let mut tx = self.pool.begin().await?; @@ -164,20 +189,24 @@ impl ITeamRepository for SqliteTeamRepository { type, content, summary, files, read, created_at \ FROM mailbox \ WHERE team_id = ? AND to_agent_id = ? AND read = 0 \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = mailbox.team_id AND t.user_id = ?) \ ORDER BY created_at ASC", ) .bind(team_id) .bind(to_agent_id) + .bind(user_id) .fetch_all(&mut *tx) .await?; if !rows.is_empty() { sqlx::query( "UPDATE mailbox SET read = 1 \ - WHERE team_id = ? AND to_agent_id = ? AND read = 0", + WHERE team_id = ? AND to_agent_id = ? AND read = 0 \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = mailbox.team_id AND t.user_id = ?)", ) .bind(team_id) .bind(to_agent_id) + .bind(user_id) .execute(&mut *tx) .await?; } @@ -186,33 +215,46 @@ impl ITeamRepository for SqliteTeamRepository { Ok(rows) } - async fn peek_unread(&self, team_id: &str, to_agent_id: &str) -> Result, DbError> { + async fn peek_unread( + &self, + user_id: &str, + team_id: &str, + to_agent_id: &str, + ) -> Result, DbError> { let rows = sqlx::query_as::<_, MailboxMessageRow>( "SELECT id, team_id, to_agent_id, from_agent_id, \ type, content, summary, files, read, created_at \ FROM mailbox \ WHERE team_id = ? AND to_agent_id = ? AND read = 0 \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = mailbox.team_id AND t.user_id = ?) \ ORDER BY created_at ASC", ) .bind(team_id) .bind(to_agent_id) + .bind(user_id) .fetch_all(&self.pool) .await?; Ok(rows) } - async fn mark_read_batch(&self, ids: &[String]) -> Result<(), DbError> { + async fn mark_read_batch(&self, user_id: &str, team_id: &str, ids: &[String]) -> Result<(), DbError> { if ids.is_empty() { return Ok(()); } // SQLite placeholder limit is 999; batch if needed. for chunk in ids.chunks(500) { let placeholders: String = chunk.iter().map(|_| "?").collect::>().join(","); - let sql = format!("UPDATE mailbox SET read = 1 WHERE id IN ({placeholders})"); + let sql = format!( + "UPDATE mailbox SET read = 1 \ + WHERE team_id = ? AND id IN ({placeholders}) \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = mailbox.team_id AND t.user_id = ?)" + ); let mut query = sqlx::query(&sql); + query = query.bind(team_id); for id in chunk { query = query.bind(id); } + query = query.bind(user_id); query.execute(&self.pool).await?; } Ok(()) @@ -220,6 +262,7 @@ impl ITeamRepository for SqliteTeamRepository { async fn get_history( &self, + user_id: &str, team_id: &str, to_agent_id: &str, limit: Option, @@ -230,11 +273,13 @@ impl ITeamRepository for SqliteTeamRepository { type, content, summary, files, read, created_at \ FROM mailbox \ WHERE team_id = ? AND to_agent_id = ? \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = mailbox.team_id AND t.user_id = ?) \ ORDER BY created_at ASC \ LIMIT ?", ) .bind(team_id) .bind(to_agent_id) + .bind(user_id) .bind(limit) .fetch_all(&self.pool) .await? @@ -244,32 +289,40 @@ impl ITeamRepository for SqliteTeamRepository { type, content, summary, files, read, created_at \ FROM mailbox \ WHERE team_id = ? AND to_agent_id = ? \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = mailbox.team_id AND t.user_id = ?) \ ORDER BY created_at ASC", ) .bind(team_id) .bind(to_agent_id) + .bind(user_id) .fetch_all(&self.pool) .await? }; Ok(rows) } - async fn delete_mailbox_by_team(&self, team_id: &str) -> Result<(), DbError> { - sqlx::query("DELETE FROM mailbox WHERE team_id = ?") - .bind(team_id) - .execute(&self.pool) - .await?; + async fn delete_mailbox_by_team(&self, user_id: &str, team_id: &str) -> Result<(), DbError> { + sqlx::query( + "DELETE FROM mailbox \ + WHERE team_id = ? \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = mailbox.team_id AND t.user_id = ?)", + ) + .bind(team_id) + .bind(user_id) + .execute(&self.pool) + .await?; Ok(()) } // ── Tasks ──────────────────────────────────────────────────────── - async fn create_task(&self, row: &TeamTaskRow) -> Result<(), DbError> { - sqlx::query( + async fn create_task(&self, user_id: &str, row: &TeamTaskRow) -> Result<(), DbError> { + let result = sqlx::query( "INSERT INTO team_tasks \ (id, team_id, subject, description, status, owner, \ blocked_by, blocks, metadata, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + SELECT ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? \ + WHERE EXISTS (SELECT 1 FROM teams t WHERE t.id = ? AND t.user_id = ?)", ) .bind(&row.id) .bind(&row.team_id) @@ -282,21 +335,42 @@ impl ITeamRepository for SqliteTeamRepository { .bind(&row.metadata) .bind(row.created_at) .bind(row.updated_at) + .bind(&row.team_id) + .bind(user_id) .execute(&self.pool) .await?; + if result.rows_affected() == 0 { + return Err(DbError::NotFound(format!("team {}", row.team_id))); + } Ok(()) } - async fn find_task_by_id(&self, team_id: &str, task_id: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, TeamTaskRow>("SELECT * FROM team_tasks WHERE team_id = ? AND id = ?") - .bind(team_id) - .bind(task_id) - .fetch_optional(&self.pool) - .await?; + async fn find_task_by_id( + &self, + user_id: &str, + team_id: &str, + task_id: &str, + ) -> Result, DbError> { + let row = sqlx::query_as::<_, TeamTaskRow>( + "SELECT * FROM team_tasks \ + WHERE team_id = ? AND id = ? \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = team_tasks.team_id AND t.user_id = ?)", + ) + .bind(team_id) + .bind(task_id) + .bind(user_id) + .fetch_optional(&self.pool) + .await?; Ok(row) } - async fn update_task(&self, task_id: &str, params: &UpdateTaskParams) -> Result<(), DbError> { + async fn update_task( + &self, + user_id: &str, + team_id: &str, + task_id: &str, + params: &UpdateTaskParams, + ) -> Result<(), DbError> { let mut set_clauses = Vec::new(); if params.status.is_some() { set_clauses.push("status = ?"); @@ -319,7 +393,12 @@ impl ITeamRepository for SqliteTeamRepository { } set_clauses.push("updated_at = ?"); - let sql = format!("UPDATE team_tasks SET {} WHERE id = ?", set_clauses.join(", ")); + let sql = format!( + "UPDATE team_tasks SET {} \ + WHERE team_id = ? AND id = ? \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = team_tasks.team_id AND t.user_id = ?)", + set_clauses.join(", ") + ); let mut query = sqlx::query(&sql); if let Some(ref status) = params.status { @@ -338,7 +417,9 @@ impl ITeamRepository for SqliteTeamRepository { query = query.bind(metadata); } query = query.bind(now_ms()); + query = query.bind(team_id); query = query.bind(task_id); + query = query.bind(user_id); let result = query.execute(&self.pool).await?; if result.rows_affected() == 0 { @@ -347,24 +428,41 @@ impl ITeamRepository for SqliteTeamRepository { Ok(()) } - async fn list_tasks(&self, team_id: &str) -> Result, DbError> { - let rows = - sqlx::query_as::<_, TeamTaskRow>("SELECT * FROM team_tasks WHERE team_id = ? ORDER BY created_at ASC") - .bind(team_id) - .fetch_all(&self.pool) - .await?; + async fn list_tasks(&self, user_id: &str, team_id: &str) -> Result, DbError> { + let rows = sqlx::query_as::<_, TeamTaskRow>( + "SELECT * FROM team_tasks \ + WHERE team_id = ? \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = team_tasks.team_id AND t.user_id = ?) \ + ORDER BY created_at ASC", + ) + .bind(team_id) + .bind(user_id) + .fetch_all(&self.pool) + .await?; Ok(rows) } - async fn append_to_blocks(&self, task_id: &str, blocked_task_id: &str) -> Result<(), DbError> { + async fn append_to_blocks( + &self, + user_id: &str, + team_id: &str, + task_id: &str, + blocked_task_id: &str, + ) -> Result<(), DbError> { // Read current blocks, append, and write back within a transaction. let mut tx = self.pool.begin().await?; - let row = sqlx::query_as::<_, TeamTaskRow>("SELECT * FROM team_tasks WHERE id = ?") - .bind(task_id) - .fetch_optional(&mut *tx) - .await? - .ok_or_else(|| DbError::NotFound(format!("task {task_id}")))?; + let row = sqlx::query_as::<_, TeamTaskRow>( + "SELECT * FROM team_tasks \ + WHERE team_id = ? AND id = ? \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = team_tasks.team_id AND t.user_id = ?)", + ) + .bind(team_id) + .bind(task_id) + .bind(user_id) + .fetch_optional(&mut *tx) + .await? + .ok_or_else(|| DbError::NotFound(format!("task {task_id}")))?; let mut blocks: Vec = serde_json::from_str(&row.blocks).unwrap_or_default(); if !blocks.contains(&blocked_task_id.to_string()) { @@ -372,47 +470,76 @@ impl ITeamRepository for SqliteTeamRepository { } let new_blocks = serde_json::to_string(&blocks).unwrap_or_else(|_| "[]".to_string()); - sqlx::query("UPDATE team_tasks SET blocks = ?, updated_at = ? WHERE id = ?") - .bind(&new_blocks) - .bind(now_ms()) - .bind(task_id) - .execute(&mut *tx) - .await?; + sqlx::query( + "UPDATE team_tasks SET blocks = ?, updated_at = ? \ + WHERE team_id = ? AND id = ? \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = team_tasks.team_id AND t.user_id = ?)", + ) + .bind(&new_blocks) + .bind(now_ms()) + .bind(team_id) + .bind(task_id) + .bind(user_id) + .execute(&mut *tx) + .await?; tx.commit().await?; Ok(()) } - async fn remove_from_blocked_by(&self, task_id: &str, unblocked_task_id: &str) -> Result<(), DbError> { + async fn remove_from_blocked_by( + &self, + user_id: &str, + team_id: &str, + task_id: &str, + unblocked_task_id: &str, + ) -> Result<(), DbError> { // Read current blocked_by, remove, and write back within a transaction. let mut tx = self.pool.begin().await?; - let row = sqlx::query_as::<_, TeamTaskRow>("SELECT * FROM team_tasks WHERE id = ?") - .bind(task_id) - .fetch_optional(&mut *tx) - .await? - .ok_or_else(|| DbError::NotFound(format!("task {task_id}")))?; + let row = sqlx::query_as::<_, TeamTaskRow>( + "SELECT * FROM team_tasks \ + WHERE team_id = ? AND id = ? \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = team_tasks.team_id AND t.user_id = ?)", + ) + .bind(team_id) + .bind(task_id) + .bind(user_id) + .fetch_optional(&mut *tx) + .await? + .ok_or_else(|| DbError::NotFound(format!("task {task_id}")))?; let mut blocked_by: Vec = serde_json::from_str(&row.blocked_by).unwrap_or_default(); blocked_by.retain(|id| id != unblocked_task_id); let new_blocked_by = serde_json::to_string(&blocked_by).unwrap_or_else(|_| "[]".to_string()); - sqlx::query("UPDATE team_tasks SET blocked_by = ?, updated_at = ? WHERE id = ?") - .bind(&new_blocked_by) - .bind(now_ms()) - .bind(task_id) - .execute(&mut *tx) - .await?; + sqlx::query( + "UPDATE team_tasks SET blocked_by = ?, updated_at = ? \ + WHERE team_id = ? AND id = ? \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = team_tasks.team_id AND t.user_id = ?)", + ) + .bind(&new_blocked_by) + .bind(now_ms()) + .bind(team_id) + .bind(task_id) + .bind(user_id) + .execute(&mut *tx) + .await?; tx.commit().await?; Ok(()) } - async fn delete_tasks_by_team(&self, team_id: &str) -> Result<(), DbError> { - sqlx::query("DELETE FROM team_tasks WHERE team_id = ?") - .bind(team_id) - .execute(&self.pool) - .await?; + async fn delete_tasks_by_team(&self, user_id: &str, team_id: &str) -> Result<(), DbError> { + sqlx::query( + "DELETE FROM team_tasks \ + WHERE team_id = ? \ + AND EXISTS (SELECT 1 FROM teams t WHERE t.id = team_tasks.team_id AND t.user_id = ?)", + ) + .bind(team_id) + .bind(user_id) + .execute(&self.pool) + .await?; Ok(()) } } diff --git a/crates/aionui-db/src/repository/sqlite_user.rs b/crates/aionui-db/src/repository/sqlite_user.rs index cf55bcf0f..ee4e46ade 100644 --- a/crates/aionui-db/src/repository/sqlite_user.rs +++ b/crates/aionui-db/src/repository/sqlite_user.rs @@ -1,7 +1,7 @@ use sqlx::SqlitePool; use crate::error::DbError; -use crate::models::User; +use crate::models::{ExternalUserProjection, User, UserStatus, UserType}; use crate::repository::IUserRepository; /// SQLite-backed implementation of [`IUserRepository`]. @@ -19,9 +19,12 @@ impl SqliteUserRepository { #[async_trait::async_trait] impl IUserRepository for SqliteUserRepository { async fn has_users(&self) -> Result { - let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM users WHERE password_hash != ''") - .fetch_one(&self.pool) - .await?; + let row: (i64,) = sqlx::query_as( + "SELECT COUNT(*) FROM users \ + WHERE user_type = 'local' AND password_hash IS NOT NULL AND password_hash != ''", + ) + .fetch_one(&self.pool) + .await?; Ok(row.0 > 0) } @@ -52,7 +55,7 @@ impl IUserRepository for SqliteUserRepository { let now = aionui_common::now_ms(); let result = sqlx::query( "UPDATE users SET username = ?, password_hash = ?, updated_at = ? \ - WHERE id = 'system_default_user'", + WHERE id = 'system_default_user' AND user_type = 'local'", ) .bind(username) .bind(password_hash) @@ -78,8 +81,8 @@ impl IUserRepository for SqliteUserRepository { let now = aionui_common::now_ms(); sqlx::query( - "INSERT INTO users (id, username, password_hash, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?)", + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, ?, 'active', 0, ?, ?)", ) .bind(&id) .bind(username) @@ -97,11 +100,15 @@ impl IUserRepository for SqliteUserRepository { Ok(User { id, - username: username.to_string(), + user_type: UserType::Local, + external_user_id: None, + username: Some(username.to_string()), email: None, - password_hash: password_hash.to_string(), + password_hash: Some(password_hash.to_string()), avatar_path: None, jwt_secret: None, + status: UserStatus::Active, + session_generation: 0, created_at: now, updated_at: now, last_login: None, @@ -109,14 +116,133 @@ impl IUserRepository for SqliteUserRepository { } async fn find_by_username(&self, username: &str) -> Result, DbError> { - let user = sqlx::query_as::<_, User>("SELECT * FROM users WHERE username = ?") - .bind(username) + let user = sqlx::query_as::<_, User>( + "SELECT * FROM users \ + WHERE user_type = 'local' AND password_hash IS NOT NULL AND username = ?", + ) + .bind(username) + .fetch_optional(&self.pool) + .await?; + + Ok(user) + } + + async fn ensure_external_user( + &self, + user_type: UserType, + external_user_id: &str, + projection: ExternalUserProjection, + ) -> Result { + if user_type == UserType::Local { + return Err(DbError::Conflict( + "External identity projection requires a non-local user type".to_string(), + )); + } + if external_user_id.trim().is_empty() { + return Err(DbError::Conflict("external_user_id must not be empty".to_string())); + } + + if let Some(existing) = self.find_by_external_user_id(user_type, external_user_id).await? { + return Ok(existing); + } + + let id = aionui_common::generate_prefixed_id("user"); + let now = aionui_common::now_ms(); + + sqlx::query( + "INSERT INTO users \ + (id, user_type, external_user_id, username, email, password_hash, avatar_path, status, session_generation, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, NULL, ?, 'active', 0, ?, ?)", + ) + .bind(&id) + .bind(user_type.as_str()) + .bind(external_user_id) + .bind(projection.username.as_deref()) + .bind(projection.email.as_deref()) + .bind(projection.avatar_path.as_deref()) + .bind(now) + .bind(now) + .execute(&self.pool) + .await + .map_err(|e| match &e { + sqlx::Error::Database(db_err) if is_unique_violation(db_err.as_ref()) => { + DbError::Conflict("External user mapping already exists".to_string()) + } + _ => DbError::Query(e), + })?; + + self.find_by_id(&id) + .await? + .ok_or_else(|| DbError::NotFound(format!("User '{id}' not found after insert"))) + } + + async fn find_by_external_user_id( + &self, + user_type: UserType, + external_user_id: &str, + ) -> Result, DbError> { + let user = sqlx::query_as::<_, User>("SELECT * FROM users WHERE user_type = ? AND external_user_id = ?") + .bind(user_type.as_str()) + .bind(external_user_id) .fetch_optional(&self.pool) .await?; Ok(user) } + async fn adopt_system_default_data(&self, owner_id: &str) -> Result { + let mut tx = self.pool.begin().await?; + + // Adoption window: exactly one external user, and it is the caller. + // A second provisioned account closes the window forever — later + // accounts must never inherit another machine user's data. + let (external_count,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM users WHERE user_type != 'local'") + .fetch_one(&mut *tx) + .await?; + if external_count != 1 { + return Ok(0); + } + let (is_owner,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM users WHERE id = ? AND user_type != 'local'") + .bind(owner_id) + .fetch_one(&mut *tx) + .await?; + if is_owner != 1 { + return Ok(0); + } + + // Discover ownership tables from the live schema rather than a + // hand-maintained list, so user-scoped tables added by future + // migrations are adopted automatically. Convention (root-scope + // design): a `user_id` column IS the ownership column. + let tables: Vec<(String,)> = sqlx::query_as( + "SELECT m.name FROM sqlite_master m \ + WHERE m.type = 'table' \ + AND m.name NOT LIKE 'sqlite_%' \ + AND m.name != 'users' \ + AND EXISTS (SELECT 1 FROM pragma_table_info(m.name) p WHERE p.name = 'user_id')", + ) + .fetch_all(&mut *tx) + .await?; + + // Global template rows (`user_id IS NULL`) are shared and stay put; + // only rows owned by the local default user move. `UPDATE OR IGNORE` + // skips rows that would collide with the new owner's existing rows on + // per-user PK/UNIQUE tables (e.g. `system_settings.user_id`). + let mut moved: u64 = 0; + for (table,) in &tables { + let result = sqlx::query(&format!( + "UPDATE OR IGNORE \"{table}\" SET user_id = ? WHERE user_id = 'system_default_user'" + )) + .bind(owner_id) + .execute(&mut *tx) + .await?; + moved += result.rows_affected(); + } + + tx.commit().await?; + Ok(moved) + } + async fn find_by_id(&self, id: &str) -> Result, DbError> { let user = sqlx::query_as::<_, User>("SELECT * FROM users WHERE id = ?") .bind(id) @@ -126,6 +252,15 @@ impl IUserRepository for SqliteUserRepository { Ok(user) } + async fn find_active_by_id(&self, id: &str) -> Result, DbError> { + let user = sqlx::query_as::<_, User>("SELECT * FROM users WHERE id = ? AND status = 'active'") + .bind(id) + .fetch_optional(&self.pool) + .await?; + + Ok(user) + } + async fn list_users(&self) -> Result, DbError> { let users = sqlx::query_as::<_, User>("SELECT * FROM users") .fetch_all(&self.pool) @@ -144,12 +279,15 @@ impl IUserRepository for SqliteUserRepository { async fn update_password(&self, user_id: &str, password_hash: &str) -> Result<(), DbError> { let now = aionui_common::now_ms(); - let result = sqlx::query("UPDATE users SET password_hash = ?, updated_at = ? WHERE id = ?") - .bind(password_hash) - .bind(now) - .bind(user_id) - .execute(&self.pool) - .await?; + let result = sqlx::query( + "UPDATE users SET password_hash = ?, updated_at = ? \ + WHERE id = ? AND user_type = 'local'", + ) + .bind(password_hash) + .bind(now) + .bind(user_id) + .execute(&self.pool) + .await?; if result.rows_affected() == 0 { return Err(DbError::NotFound(format!("User '{user_id}' not found"))); @@ -160,18 +298,21 @@ impl IUserRepository for SqliteUserRepository { async fn update_username(&self, user_id: &str, username: &str) -> Result<(), DbError> { let now = aionui_common::now_ms(); - let result = sqlx::query("UPDATE users SET username = ?, updated_at = ? WHERE id = ?") - .bind(username) - .bind(now) - .bind(user_id) - .execute(&self.pool) - .await - .map_err(|e| match &e { - sqlx::Error::Database(db_err) if is_unique_violation(db_err.as_ref()) => { - DbError::Conflict(format!("Username '{username}' already exists")) - } - _ => DbError::Query(e), - })?; + let result = sqlx::query( + "UPDATE users SET username = ?, updated_at = ? \ + WHERE id = ? AND user_type = 'local'", + ) + .bind(username) + .bind(now) + .bind(user_id) + .execute(&self.pool) + .await + .map_err(|e| match &e { + sqlx::Error::Database(db_err) if is_unique_violation(db_err.as_ref()) => { + DbError::Conflict(format!("Username '{username}' already exists")) + } + _ => DbError::Query(e), + })?; if result.rows_affected() == 0 { return Err(DbError::NotFound(format!("User '{user_id}' not found"))); @@ -211,12 +352,62 @@ impl IUserRepository for SqliteUserRepository { Ok(()) } + + async fn set_status(&self, user_id: &str, status: UserStatus) -> Result<(), DbError> { + let now = aionui_common::now_ms(); + let result = sqlx::query( + "UPDATE users \ + SET status = ?, \ + session_generation = CASE \ + WHEN ? = 'disabled' AND status != 'disabled' THEN session_generation + 1 \ + ELSE session_generation \ + END, \ + updated_at = ? \ + WHERE id = ?", + ) + .bind(status.as_str()) + .bind(status.as_str()) + .bind(now) + .bind(user_id) + .execute(&self.pool) + .await?; + + if result.rows_affected() == 0 { + return Err(DbError::NotFound(format!("User '{user_id}' not found"))); + } + + Ok(()) + } + + async fn increment_session_generation(&self, user_id: &str) -> Result { + let now = aionui_common::now_ms(); + let result = sqlx::query( + "UPDATE users \ + SET session_generation = session_generation + 1, updated_at = ? \ + WHERE id = ?", + ) + .bind(now) + .bind(user_id) + .execute(&self.pool) + .await?; + + if result.rows_affected() == 0 { + return Err(DbError::NotFound(format!("User '{user_id}' not found"))); + } + + let generation: i64 = sqlx::query_scalar("SELECT session_generation FROM users WHERE id = ?") + .bind(user_id) + .fetch_one(&self.pool) + .await?; + + Ok(generation) + } } /// Checks if a SQLite database error is a UNIQUE constraint violation. fn is_unique_violation(err: &dyn sqlx::error::DatabaseError) -> bool { - // SQLite error code 2067 = SQLITE_CONSTRAINT_UNIQUE - err.code().is_some_and(|c| c == "2067") + // SQLite error code 2067 = SQLITE_CONSTRAINT_UNIQUE, 1555 = SQLITE_CONSTRAINT_PRIMARYKEY. + err.code().is_some_and(|c| c == "2067" || c == "1555") } #[cfg(test)] @@ -240,7 +431,7 @@ mod tests { #[test] fn non_unique_violation_code_rejected() { - assert!(!is_unique_violation(&FakeDbError("1555"))); + assert!(!is_unique_violation(&FakeDbError("1299"))); } /// Minimal fake for testing is_unique_violation. @@ -289,8 +480,11 @@ mod tests { let user = repo.create_user("alice", "hash123").await.unwrap(); assert!(user.id.starts_with("user_")); - assert_eq!(user.username, "alice"); - assert_eq!(user.password_hash, "hash123"); + assert_eq!(user.user_type, UserType::Local); + assert_eq!(user.status, UserStatus::Active); + assert_eq!(user.session_generation, 0); + assert_eq!(user.username.as_deref(), Some("alice")); + assert_eq!(user.password_hash.as_deref(), Some("hash123")); assert!(user.email.is_none()); assert!(user.avatar_path.is_none()); assert!(user.jwt_secret.is_none()); @@ -326,7 +520,9 @@ mod tests { let (repo, _db) = setup().await; let user = repo.get_system_user().await.unwrap().unwrap(); assert_eq!(user.id, "system_default_user"); - assert_eq!(user.username, "admin"); + assert_eq!(user.user_type, UserType::Local); + assert_eq!(user.status, UserStatus::Active); + assert_eq!(user.username.as_deref(), Some("admin")); } #[tokio::test] @@ -347,7 +543,7 @@ mod tests { let found = repo.find_by_username("charlie").await.unwrap(); assert!(found.is_some()); - assert_eq!(found.unwrap().username, "charlie"); + assert_eq!(found.unwrap().username.as_deref(), Some("charlie")); } #[tokio::test] @@ -400,7 +596,7 @@ mod tests { repo.update_password(&user.id, "new_hash").await.unwrap(); let updated = repo.find_by_id(&user.id).await.unwrap().unwrap(); - assert_eq!(updated.password_hash, "new_hash"); + assert_eq!(updated.password_hash.as_deref(), Some("new_hash")); assert!(updated.updated_at >= user.updated_at); } @@ -419,7 +615,7 @@ mod tests { repo.update_username(&user.id, "ivan_new").await.unwrap(); let updated = repo.find_by_id(&user.id).await.unwrap().unwrap(); - assert_eq!(updated.username, "ivan_new"); + assert_eq!(updated.username.as_deref(), Some("ivan_new")); } #[tokio::test] @@ -473,7 +669,229 @@ mod tests { repo.set_system_user_credentials("admin", "secure_hash").await.unwrap(); let user = repo.get_system_user().await.unwrap().unwrap(); - assert_eq!(user.username, "admin"); - assert_eq!(user.password_hash, "secure_hash"); + assert_eq!(user.username.as_deref(), Some("admin")); + assert_eq!(user.password_hash.as_deref(), Some("secure_hash")); + } + + #[tokio::test] + async fn ensure_external_user_is_idempotent_and_has_no_password() { + let (repo, _db) = setup().await; + let projection = ExternalUserProjection { + username: Some("AionPro User".to_string()), + email: Some("user@example.com".to_string()), + avatar_path: Some("/avatar.png".to_string()), + }; + + let first = repo + .ensure_external_user(UserType::Aionpro, "external-123", projection) + .await + .unwrap(); + let second = repo + .ensure_external_user( + UserType::Aionpro, + "external-123", + ExternalUserProjection { + username: Some("Different".to_string()), + email: None, + avatar_path: None, + }, + ) + .await + .unwrap(); + + assert_eq!(first.id, second.id); + assert_eq!(first.user_type, UserType::Aionpro); + assert_eq!(first.external_user_id.as_deref(), Some("external-123")); + assert_eq!(first.status, UserStatus::Active); + assert_eq!(first.session_generation, 0); + assert!(first.password_hash.is_none()); + assert!(repo.find_by_username("AionPro User").await.unwrap().is_none()); + } + + #[tokio::test] + async fn disabled_user_is_excluded_from_active_lookup() { + let (repo, _db) = setup().await; + let user = repo + .ensure_external_user( + UserType::Aionpro, + "external-disabled", + ExternalUserProjection::default(), + ) + .await + .unwrap(); + + assert!(repo.find_active_by_id(&user.id).await.unwrap().is_some()); + repo.set_status(&user.id, UserStatus::Disabled).await.unwrap(); + + let disabled = repo.find_by_id(&user.id).await.unwrap().unwrap(); + assert_eq!(disabled.status, UserStatus::Disabled); + assert_eq!(disabled.session_generation, 1); + assert!(repo.find_active_by_id(&user.id).await.unwrap().is_none()); + + repo.set_status(&user.id, UserStatus::Disabled).await.unwrap(); + let disabled_again = repo.find_by_id(&user.id).await.unwrap().unwrap(); + assert_eq!(disabled_again.session_generation, 1); + } + + #[tokio::test] + async fn increment_session_generation_returns_new_value() { + let (repo, _db) = setup().await; + let user = repo.create_user("session-user", "h").await.unwrap(); + + let generation = repo.increment_session_generation(&user.id).await.unwrap(); + + assert_eq!(generation, 1); + let updated = repo.find_by_id(&user.id).await.unwrap().unwrap(); + assert_eq!(updated.session_generation, 1); + } + + async fn seed_legacy_conversation(db: &crate::Database, id: &str) { + let now = aionui_common::now_ms(); + sqlx::query( + "INSERT INTO conversations (id, user_id, name, type, created_at, updated_at) \ + VALUES (?, 'system_default_user', 'legacy', 'aionrs', ?, ?)", + ) + .bind(id) + .bind(now) + .bind(now) + .execute(db.pool()) + .await + .unwrap(); + } + + async fn conversation_owner(db: &crate::Database, id: &str) -> String { + let (owner,): (String,) = sqlx::query_as("SELECT user_id FROM conversations WHERE id = ?") + .bind(id) + .fetch_one(db.pool()) + .await + .unwrap(); + owner + } + + #[tokio::test] + async fn adopt_moves_legacy_data_to_sole_external_user_once() { + let (repo, db) = setup().await; + seed_legacy_conversation(&db, "conv-legacy").await; + let user = repo + .ensure_external_user(UserType::Aionpro, "ext-1", ExternalUserProjection::default()) + .await + .unwrap(); + + let moved = repo.adopt_system_default_data(&user.id).await.unwrap(); + assert!(moved >= 1); + assert_eq!(conversation_owner(&db, "conv-legacy").await, user.id); + + // Self-idempotent: the source set is now empty. + let moved_again = repo.adopt_system_default_data(&user.id).await.unwrap(); + assert_eq!(moved_again, 0); + } + + #[tokio::test] + async fn adopt_is_refused_once_a_second_external_user_exists() { + let (repo, db) = setup().await; + seed_legacy_conversation(&db, "conv-legacy-2").await; + let first = repo + .ensure_external_user(UserType::Aionpro, "ext-a", ExternalUserProjection::default()) + .await + .unwrap(); + let second = repo + .ensure_external_user(UserType::Aionpro, "ext-b", ExternalUserProjection::default()) + .await + .unwrap(); + + assert_eq!(repo.adopt_system_default_data(&first.id).await.unwrap(), 0); + assert_eq!(repo.adopt_system_default_data(&second.id).await.unwrap(), 0); + assert_eq!(conversation_owner(&db, "conv-legacy-2").await, "system_default_user"); + } + + #[tokio::test] + async fn adopt_is_refused_for_local_or_unknown_owner() { + let (repo, db) = setup().await; + seed_legacy_conversation(&db, "conv-legacy-3").await; + + // No external user at all. + assert_eq!(repo.adopt_system_default_data("system_default_user").await.unwrap(), 0); + + // Owner id that is not the external user. + repo.ensure_external_user(UserType::Aionpro, "ext-c", ExternalUserProjection::default()) + .await + .unwrap(); + assert_eq!(repo.adopt_system_default_data("someone-else").await.unwrap(), 0); + assert_eq!(conversation_owner(&db, "conv-legacy-3").await, "system_default_user"); + } + + #[tokio::test] + async fn adopt_discovers_all_known_user_scoped_tables() { + let (_repo, db) = setup().await; + let tables: Vec<(String,)> = sqlx::query_as( + "SELECT m.name FROM sqlite_master m \ + WHERE m.type = 'table' \ + AND m.name NOT LIKE 'sqlite_%' \ + AND m.name != 'users' \ + AND EXISTS (SELECT 1 FROM pragma_table_info(m.name) p WHERE p.name = 'user_id')", + ) + .fetch_all(db.pool()) + .await + .unwrap(); + let names: Vec<&str> = tables.iter().map(|(n,)| n.as_str()).collect(); + + // Sentinel: the discovery convention (`user_id` column == ownership) + // must keep matching the core scope tables. If this fails, either a + // migration renamed an ownership column or the convention broke — + // both must be looked at before shipping. + for expected in [ + "conversations", + "teams", + "cron_jobs", + "skills", + "mcp_servers", + "providers", + "assistants", + "system_settings", + "client_preferences", + ] { + assert!( + names.contains(&expected), + "expected user-scoped table '{expected}' to be discovered" + ); + } + } + + #[tokio::test] + async fn adopt_skips_rows_colliding_with_the_new_owner() { + let (repo, db) = setup().await; + let now = aionui_common::now_ms(); + let user = repo + .ensure_external_user(UserType::Aionpro, "ext-d", ExternalUserProjection::default()) + .await + .unwrap(); + + // Both the legacy user and the new owner have a system_settings row + // (PRIMARY KEY user_id) — the owner's row must win, the legacy row + // must be left behind rather than erroring the whole adoption. + sqlx::query( + "INSERT INTO system_settings (user_id, language, updated_at) VALUES ('system_default_user', 'zh-CN', ?)", + ) + .bind(now) + .execute(db.pool()) + .await + .unwrap(); + sqlx::query("INSERT INTO system_settings (user_id, language, updated_at) VALUES (?, 'en-US', ?)") + .bind(&user.id) + .bind(now) + .execute(db.pool()) + .await + .unwrap(); + seed_legacy_conversation(&db, "conv-legacy-4").await; + + repo.adopt_system_default_data(&user.id).await.unwrap(); + + let (language,): (String,) = sqlx::query_as("SELECT language FROM system_settings WHERE user_id = ?") + .bind(&user.id) + .fetch_one(db.pool()) + .await + .unwrap(); + assert_eq!(language, "en-US"); + assert_eq!(conversation_owner(&db, "conv-legacy-4").await, user.id); } } diff --git a/crates/aionui-db/src/repository/team.rs b/crates/aionui-db/src/repository/team.rs index e7b1c2095..f407d025a 100644 --- a/crates/aionui-db/src/repository/team.rs +++ b/crates/aionui-db/src/repository/team.rs @@ -33,72 +33,109 @@ pub trait ITeamRepository: Send + Sync { /// Inserts a new team record. async fn create_team(&self, row: &TeamRow) -> Result<(), DbError>; - /// Returns all teams ordered by creation time ascending. - async fn list_teams(&self) -> Result, DbError>; + /// Returns all teams for startup/session restore. + async fn list_teams_for_restore(&self) -> Result, DbError>; /// Returns teams owned by `user_id`, ordered by creation time ascending. async fn list_teams_by_user(&self, user_id: &str) -> Result, DbError>; - /// Returns a single team by id, or `None` if not found. - async fn get_team(&self, team_id: &str) -> Result, DbError>; + /// Returns a single team owned by `user_id`, or `None` if not found. + async fn get_team(&self, user_id: &str, team_id: &str) -> Result, DbError>; + + /// Returns a single team for startup/session restore. + async fn get_team_for_restore(&self, team_id: &str) -> Result, DbError>; /// Updates a team by id with the provided fields. /// Returns `DbError::NotFound` if absent. - async fn update_team(&self, team_id: &str, params: &UpdateTeamParams) -> Result<(), DbError>; + async fn update_team(&self, user_id: &str, team_id: &str, params: &UpdateTeamParams) -> Result<(), DbError>; /// Deletes a team by id. Returns `DbError::NotFound` if absent. - async fn delete_team(&self, team_id: &str) -> Result<(), DbError>; + async fn delete_team(&self, user_id: &str, team_id: &str) -> Result<(), DbError>; // ── Mailbox ────────────────────────────────────────────────────── /// Writes a message to the mailbox. - async fn write_message(&self, row: &MailboxMessageRow) -> Result<(), DbError>; + async fn write_message(&self, user_id: &str, row: &MailboxMessageRow) -> Result<(), DbError>; /// Atomically reads all unread messages for `to_agent_id` in a team /// and marks them as read. Uses `BEGIN IMMEDIATE` for atomicity. - async fn read_unread_and_mark(&self, team_id: &str, to_agent_id: &str) -> Result, DbError>; + async fn read_unread_and_mark( + &self, + user_id: &str, + team_id: &str, + to_agent_id: &str, + ) -> Result, DbError>; /// Reads all unread messages for `to_agent_id` without marking them as read. - async fn peek_unread(&self, team_id: &str, to_agent_id: &str) -> Result, DbError>; + async fn peek_unread( + &self, + user_id: &str, + team_id: &str, + to_agent_id: &str, + ) -> Result, DbError>; /// Marks the given message IDs as read. IDs that don't exist are silently ignored. - async fn mark_read_batch(&self, ids: &[String]) -> Result<(), DbError>; + async fn mark_read_batch(&self, user_id: &str, team_id: &str, ids: &[String]) -> Result<(), DbError>; /// Returns message history for an agent, optionally limited. /// Messages are ordered by `created_at` ascending. async fn get_history( &self, + user_id: &str, team_id: &str, to_agent_id: &str, limit: Option, ) -> Result, DbError>; /// Deletes all mailbox messages belonging to a team. - async fn delete_mailbox_by_team(&self, team_id: &str) -> Result<(), DbError>; + async fn delete_mailbox_by_team(&self, user_id: &str, team_id: &str) -> Result<(), DbError>; // ── Tasks ──────────────────────────────────────────────────────── /// Creates a new task. - async fn create_task(&self, row: &TeamTaskRow) -> Result<(), DbError>; + async fn create_task(&self, user_id: &str, row: &TeamTaskRow) -> Result<(), DbError>; /// Finds a task by exact id within a team. - async fn find_task_by_id(&self, team_id: &str, task_id: &str) -> Result, DbError>; + async fn find_task_by_id( + &self, + user_id: &str, + team_id: &str, + task_id: &str, + ) -> Result, DbError>; /// Updates a task by id with the provided fields. /// Returns `DbError::NotFound` if absent. - async fn update_task(&self, task_id: &str, params: &UpdateTaskParams) -> Result<(), DbError>; + async fn update_task( + &self, + user_id: &str, + team_id: &str, + task_id: &str, + params: &UpdateTaskParams, + ) -> Result<(), DbError>; /// Returns all tasks for a team, ordered by `created_at` ascending. - async fn list_tasks(&self, team_id: &str) -> Result, DbError>; + async fn list_tasks(&self, user_id: &str, team_id: &str) -> Result, DbError>; /// Appends `blocked_task_id` to the `blocks` JSON array of `task_id`. /// This is a transactional JSON array append operation. - async fn append_to_blocks(&self, task_id: &str, blocked_task_id: &str) -> Result<(), DbError>; + async fn append_to_blocks( + &self, + user_id: &str, + team_id: &str, + task_id: &str, + blocked_task_id: &str, + ) -> Result<(), DbError>; /// Removes `unblocked_task_id` from the `blocked_by` JSON array of `task_id`. /// This is a transactional JSON array removal operation. - async fn remove_from_blocked_by(&self, task_id: &str, unblocked_task_id: &str) -> Result<(), DbError>; + async fn remove_from_blocked_by( + &self, + user_id: &str, + team_id: &str, + task_id: &str, + unblocked_task_id: &str, + ) -> Result<(), DbError>; /// Deletes all tasks belonging to a team. - async fn delete_tasks_by_team(&self, team_id: &str) -> Result<(), DbError>; + async fn delete_tasks_by_team(&self, user_id: &str, team_id: &str) -> Result<(), DbError>; } diff --git a/crates/aionui-db/src/repository/user.rs b/crates/aionui-db/src/repository/user.rs index 690aa5a85..c547b8045 100644 --- a/crates/aionui-db/src/repository/user.rs +++ b/crates/aionui-db/src/repository/user.rs @@ -1,5 +1,5 @@ use crate::error::DbError; -use crate::models::User; +use crate::models::{ExternalUserProjection, User, UserStatus, UserType}; /// User data access abstraction. /// @@ -32,12 +32,43 @@ pub trait IUserRepository: Send + Sync { /// Returns `DbError::Conflict` if the username already exists. async fn create_user(&self, username: &str, password_hash: &str) -> Result; - /// Finds a user by username. + /// Finds an active local password user by username. async fn find_by_username(&self, username: &str) -> Result, DbError>; + /// Idempotently creates or returns an external identity projection. + async fn ensure_external_user( + &self, + user_type: UserType, + external_user_id: &str, + projection: ExternalUserProjection, + ) -> Result; + + /// Finds a user by external identity mapping. + async fn find_by_external_user_id( + &self, + user_type: UserType, + external_user_id: &str, + ) -> Result, DbError>; + + /// One-time adoption of the machine's pre-multi-account data: while + /// `owner_id` is the ONLY external (aionpro) user in this database, + /// re-own every user-scoped row currently held by `system_default_user` + /// to `owner_id`. Returns the number of rows moved. + /// + /// Self-idempotent: after the first successful adoption the source set is + /// empty, so repeated calls move nothing. Once a second external user is + /// provisioned the adoption window closes permanently. Rows that would + /// collide with an existing row of the new owner (per-user PK/UNIQUE + /// tables such as `system_settings`) are skipped, keeping the owner's own + /// data authoritative. + async fn adopt_system_default_data(&self, owner_id: &str) -> Result; + /// Finds a user by ID. async fn find_by_id(&self, id: &str) -> Result, DbError>; + /// Finds an active user by ID. + async fn find_active_by_id(&self, id: &str) -> Result, DbError>; + /// Lists all users. async fn list_users(&self) -> Result, DbError>; @@ -57,4 +88,11 @@ pub trait IUserRepository: Send + Sync { /// Updates a user's JWT secret. async fn update_jwt_secret(&self, user_id: &str, jwt_secret: &str) -> Result<(), DbError>; + + /// Updates a user's status. Transitioning into `disabled` also revokes + /// existing sessions by incrementing `session_generation`. + async fn set_status(&self, user_id: &str, status: UserStatus) -> Result<(), DbError>; + + /// Increments a user's session generation and returns the new value. + async fn increment_session_generation(&self, user_id: &str) -> Result; } diff --git a/crates/aionui-db/tests/agent_binding_resolver.rs b/crates/aionui-db/tests/agent_binding_resolver.rs index 81511343d..b2214f683 100644 --- a/crates/aionui-db/tests/agent_binding_resolver.rs +++ b/crates/aionui-db/tests/agent_binding_resolver.rs @@ -1,4 +1,7 @@ -use aionui_db::{AgentBindingResolution, init_database_memory, resolve_agent_binding}; +use aionui_db::{ + AgentBindingResolution, IAgentMetadataRepository, SqliteAgentMetadataRepository, UpsertAgentMetadataParams, + init_database_memory, resolve_agent_binding, resolve_agent_binding_for_user, +}; #[tokio::test] async fn resolves_legacy_backend_to_agent_metadata_id() { @@ -34,3 +37,73 @@ async fn resolves_internal_agent_type_when_backend_is_null() { assert_eq!(resolved.agent_type, "aionrs"); assert_eq!(resolved.runtime_backend, "aionrs"); } + +#[tokio::test] +async fn resolves_agent_binding_for_current_user_scope() { + let db = init_database_memory().await.unwrap(); + let repo = SqliteAgentMetadataRepository::new(db.pool().clone()); + for user_id in ["user-a", "user-b"] { + sqlx::query( + "INSERT INTO users (id, username, password_hash, created_at, updated_at) VALUES (?, ?, 'hash', 0, 0)", + ) + .bind(user_id) + .bind(user_id) + .execute(db.pool()) + .await + .unwrap(); + } + + repo.upsert_for_user( + "user-a", + &custom_agent_params("shared-agent", "User A Agent", "a-backend"), + ) + .await + .unwrap(); + repo.upsert_for_user( + "user-b", + &custom_agent_params("shared-agent", "User B Agent", "b-backend"), + ) + .await + .unwrap(); + + let user_a = resolve_agent_binding_for_user(db.pool(), "user-a", "shared-agent") + .await + .unwrap() + .expect("user-a binding"); + let user_b = resolve_agent_binding_for_user(db.pool(), "user-b", "shared-agent") + .await + .unwrap() + .expect("user-b binding"); + + assert_eq!(user_a.runtime_backend, "a-backend"); + assert_eq!(user_b.runtime_backend, "b-backend"); +} + +fn custom_agent_params<'a>(id: &'a str, name: &'a str, backend: &'a str) -> UpsertAgentMetadataParams<'a> { + UpsertAgentMetadataParams { + id, + icon: None, + name, + name_i18n: None, + description: Some("custom"), + description_i18n: None, + backend: Some(backend), + agent_type: "acp", + agent_source: "custom", + agent_source_info: Some("{}"), + enabled: true, + command: Some("custom"), + args: Some("[]"), + env: Some("[]"), + native_skills_dirs: Some("[]"), + behavior_policy: Some("{}"), + yolo_id: None, + agent_capabilities: None, + auth_methods: None, + config_options: None, + available_modes: None, + available_models: None, + available_commands: None, + sort_order: 1000, + } +} diff --git a/crates/aionui-db/tests/channel_repository.rs b/crates/aionui-db/tests/channel_repository.rs index 9febfbbb8..19cc445da 100644 --- a/crates/aionui-db/tests/channel_repository.rs +++ b/crates/aionui-db/tests/channel_repository.rs @@ -9,6 +9,8 @@ use std::sync::Arc; use aionui_db::models::{AssistantSessionRow, AssistantUserRow, ChannelPluginRow, PairingCodeRow}; use aionui_db::{DbError, IChannelRepository, SqliteChannelRepository, UpdatePluginStatusParams, init_database_memory}; +const OWNER_ID: &str = "system_default_user"; + async fn repo() -> (Arc, aionui_db::Database) { let db = init_database_memory().await.unwrap(); let r = Arc::new(SqliteChannelRepository::new(db.pool().clone())); @@ -19,6 +21,7 @@ fn make_plugin(id: &str, plugin_type: &str) -> ChannelPluginRow { let now = aionui_common::now_ms(); ChannelPluginRow { id: id.into(), + owner_user_id: OWNER_ID.into(), r#type: plugin_type.into(), name: format!("{plugin_type} bot"), enabled: false, @@ -34,6 +37,7 @@ fn make_user(id: &str, platform_uid: &str, platform: &str) -> AssistantUserRow { let now = aionui_common::now_ms(); AssistantUserRow { id: id.into(), + owner_user_id: OWNER_ID.into(), platform_user_id: platform_uid.into(), platform_type: platform.into(), display_name: Some(format!("User {id}")), @@ -61,6 +65,7 @@ fn make_pairing(code: &str, platform_uid: &str, expires_offset_ms: i64) -> Pairi let now = aionui_common::now_ms(); PairingCodeRow { code: code.into(), + owner_user_id: OWNER_ID.into(), platform_user_id: platform_uid.into(), platform_type: "telegram".into(), display_name: Some("Tester".into()), @@ -77,15 +82,20 @@ async fn plugin_full_lifecycle() { let (repo, _db) = repo().await; // Empty initially. - assert!(repo.get_all_plugins().await.unwrap().is_empty()); + assert!(repo.get_all_plugins(OWNER_ID).await.unwrap().is_empty()); // Create two plugins. - repo.upsert_plugin(&make_plugin("tg-1", "telegram")).await.unwrap(); - repo.upsert_plugin(&make_plugin("lark-1", "lark")).await.unwrap(); - assert_eq!(repo.get_all_plugins().await.unwrap().len(), 2); + repo.upsert_plugin(OWNER_ID, &make_plugin("tg-1", "telegram")) + .await + .unwrap(); + repo.upsert_plugin(OWNER_ID, &make_plugin("lark-1", "lark")) + .await + .unwrap(); + assert_eq!(repo.get_all_plugins(OWNER_ID).await.unwrap().len(), 2); // Update status. repo.update_plugin_status( + OWNER_ID, "tg-1", &UpdatePluginStatusParams { status: Some("running".into()), @@ -96,13 +106,13 @@ async fn plugin_full_lifecycle() { .await .unwrap(); - let tg = repo.get_plugin("tg-1").await.unwrap().unwrap(); + let tg = repo.get_plugin(OWNER_ID, "tg-1").await.unwrap().unwrap(); assert!(tg.enabled); assert_eq!(tg.status.as_deref(), Some("running")); // Delete one. - repo.delete_plugin("lark-1").await.unwrap(); - assert_eq!(repo.get_all_plugins().await.unwrap().len(), 1); + repo.delete_plugin(OWNER_ID, "lark-1").await.unwrap(); + assert_eq!(repo.get_all_plugins(OWNER_ID).await.unwrap().len(), 1); } // ── DC-3: Same platform user uniqueness constraint ─────────────────── @@ -110,11 +120,13 @@ async fn plugin_full_lifecycle() { #[tokio::test] async fn dc3_duplicate_platform_user_rejected() { let (repo, _db) = repo().await; - repo.create_user(&make_user("u1", "tg_100", "telegram")).await.unwrap(); + repo.create_user(OWNER_ID, &make_user("u1", "tg_100", "telegram")) + .await + .unwrap(); // Same platform_user_id + platform_type with different id. let dup = make_user("u2", "tg_100", "telegram"); - let err = repo.create_user(&dup).await.unwrap_err(); + let err = repo.create_user(OWNER_ID, &dup).await.unwrap_err(); assert!(matches!(err, DbError::Conflict(_))); } @@ -123,20 +135,22 @@ async fn dc3_duplicate_platform_user_rejected() { #[tokio::test] async fn dc1_delete_user_cascades_sessions() { let (repo, _db) = repo().await; - repo.create_user(&make_user("u1", "tg_1", "telegram")).await.unwrap(); + repo.create_user(OWNER_ID, &make_user("u1", "tg_1", "telegram")) + .await + .unwrap(); // Create two sessions for the user. - repo.get_or_create_session("u1", "chat-a", &make_session("s1", "u1", "chat-a")) + repo.get_or_create_session(OWNER_ID, "u1", "chat-a", &make_session("s1", "u1", "chat-a")) .await .unwrap(); - repo.get_or_create_session("u1", "chat-b", &make_session("s2", "u1", "chat-b")) + repo.get_or_create_session(OWNER_ID, "u1", "chat-b", &make_session("s2", "u1", "chat-b")) .await .unwrap(); - assert_eq!(repo.get_all_sessions().await.unwrap().len(), 2); + assert_eq!(repo.get_all_sessions(OWNER_ID).await.unwrap().len(), 2); // Delete user → sessions cascade. - repo.delete_user("u1").await.unwrap(); - assert!(repo.get_all_sessions().await.unwrap().is_empty()); + repo.delete_user(OWNER_ID, "u1").await.unwrap(); + assert!(repo.get_all_sessions(OWNER_ID).await.unwrap().is_empty()); } // ── PC-1: Same user, different chatId → different sessions ─────────── @@ -144,19 +158,21 @@ async fn dc1_delete_user_cascades_sessions() { #[tokio::test] async fn pc1_same_user_different_chat_ids() { let (repo, _db) = repo().await; - repo.create_user(&make_user("u1", "tg_1", "telegram")).await.unwrap(); + repo.create_user(OWNER_ID, &make_user("u1", "tg_1", "telegram")) + .await + .unwrap(); let s1 = repo - .get_or_create_session("u1", "chat-a", &make_session("s1", "u1", "chat-a")) + .get_or_create_session(OWNER_ID, "u1", "chat-a", &make_session("s1", "u1", "chat-a")) .await .unwrap(); let s2 = repo - .get_or_create_session("u1", "chat-b", &make_session("s2", "u1", "chat-b")) + .get_or_create_session(OWNER_ID, "u1", "chat-b", &make_session("s2", "u1", "chat-b")) .await .unwrap(); assert_ne!(s1.id, s2.id); - assert_eq!(repo.get_all_sessions().await.unwrap().len(), 2); + assert_eq!(repo.get_all_sessions(OWNER_ID).await.unwrap().len(), 2); } // ── PC-2: Different users, same chatId → different sessions ────────── @@ -164,15 +180,19 @@ async fn pc1_same_user_different_chat_ids() { #[tokio::test] async fn pc2_different_users_same_chat_id() { let (repo, _db) = repo().await; - repo.create_user(&make_user("u1", "tg_1", "telegram")).await.unwrap(); - repo.create_user(&make_user("u2", "tg_2", "telegram")).await.unwrap(); + repo.create_user(OWNER_ID, &make_user("u1", "tg_1", "telegram")) + .await + .unwrap(); + repo.create_user(OWNER_ID, &make_user("u2", "tg_2", "telegram")) + .await + .unwrap(); let s1 = repo - .get_or_create_session("u1", "chat-x", &make_session("s1", "u1", "chat-x")) + .get_or_create_session(OWNER_ID, "u1", "chat-x", &make_session("s1", "u1", "chat-x")) .await .unwrap(); let s2 = repo - .get_or_create_session("u2", "chat-x", &make_session("s2", "u2", "chat-x")) + .get_or_create_session(OWNER_ID, "u2", "chat-x", &make_session("s2", "u2", "chat-x")) .await .unwrap(); @@ -184,16 +204,18 @@ async fn pc2_different_users_same_chat_id() { #[tokio::test] async fn pc3_same_user_same_chat_reuses_session() { let (repo, _db) = repo().await; - repo.create_user(&make_user("u1", "tg_1", "telegram")).await.unwrap(); + repo.create_user(OWNER_ID, &make_user("u1", "tg_1", "telegram")) + .await + .unwrap(); let s1 = repo - .get_or_create_session("u1", "chat-a", &make_session("s1", "u1", "chat-a")) + .get_or_create_session(OWNER_ID, "u1", "chat-a", &make_session("s1", "u1", "chat-a")) .await .unwrap(); // Second call with a different new_row id but same user+chat. let s2 = repo - .get_or_create_session("u1", "chat-a", &make_session("s999", "u1", "chat-a")) + .get_or_create_session(OWNER_ID, "u1", "chat-a", &make_session("s999", "u1", "chat-a")) .await .unwrap(); @@ -208,9 +230,9 @@ async fn pc3_same_user_same_chat_reuses_session() { async fn pg2_pairing_code_expiry_is_10_minutes() { let (repo, _db) = repo().await; let pairing = make_pairing("123456", "tg_99", 600_000); - repo.create_pairing(&pairing).await.unwrap(); + repo.create_pairing(OWNER_ID, &pairing).await.unwrap(); - let found = repo.get_pairing_by_code("123456").await.unwrap().unwrap(); + let found = repo.get_pairing_by_code(OWNER_ID, "123456").await.unwrap().unwrap(); assert_eq!(found.expires_at - found.requested_at, 600_000); } @@ -222,21 +244,21 @@ async fn expired_pairings_cleaned_up() { let now = aionui_common::now_ms(); // Already expired. - repo.create_pairing(&make_pairing("111111", "tg_1", -1000)) + repo.create_pairing(OWNER_ID, &make_pairing("111111", "tg_1", -1000)) .await .unwrap(); // Still valid. - repo.create_pairing(&make_pairing("222222", "tg_2", 600_000)) + repo.create_pairing(OWNER_ID, &make_pairing("222222", "tg_2", 600_000)) .await .unwrap(); - let cleaned = repo.cleanup_expired_pairings(now).await.unwrap(); + let cleaned = repo.cleanup_expired_pairings(OWNER_ID, now).await.unwrap(); assert_eq!(cleaned, 1); - let expired = repo.get_pairing_by_code("111111").await.unwrap().unwrap(); + let expired = repo.get_pairing_by_code(OWNER_ID, "111111").await.unwrap().unwrap(); assert_eq!(expired.status, "expired"); - let valid = repo.get_pairing_by_code("222222").await.unwrap().unwrap(); + let valid = repo.get_pairing_by_code(OWNER_ID, "222222").await.unwrap().unwrap(); assert_eq!(valid.status, "pending"); } @@ -245,26 +267,38 @@ async fn expired_pairings_cleaned_up() { #[tokio::test] async fn pairing_approve_and_reject() { let (repo, _db) = repo().await; - repo.create_pairing(&make_pairing("100001", "tg_a", 600_000)) + repo.create_pairing(OWNER_ID, &make_pairing("100001", "tg_a", 600_000)) .await .unwrap(); - repo.create_pairing(&make_pairing("100002", "tg_b", 600_000)) + repo.create_pairing(OWNER_ID, &make_pairing("100002", "tg_b", 600_000)) .await .unwrap(); - repo.update_pairing_status("100001", "approved").await.unwrap(); - repo.update_pairing_status("100002", "rejected").await.unwrap(); + repo.update_pairing_status(OWNER_ID, "100001", "approved") + .await + .unwrap(); + repo.update_pairing_status(OWNER_ID, "100002", "rejected") + .await + .unwrap(); // Neither should appear in pending list. - let pending = repo.get_pending_pairings().await.unwrap(); + let pending = repo.get_pending_pairings(OWNER_ID).await.unwrap(); assert!(pending.is_empty()); assert_eq!( - repo.get_pairing_by_code("100001").await.unwrap().unwrap().status, + repo.get_pairing_by_code(OWNER_ID, "100001") + .await + .unwrap() + .unwrap() + .status, "approved" ); assert_eq!( - repo.get_pairing_by_code("100002").await.unwrap().unwrap().status, + repo.get_pairing_by_code(OWNER_ID, "100002") + .await + .unwrap() + .unwrap() + .status, "rejected" ); } @@ -277,13 +311,13 @@ async fn users_ordered_by_authorized_at_desc() { let mut u1 = make_user("u1", "tg_1", "telegram"); u1.authorized_at = 1000; - repo.create_user(&u1).await.unwrap(); + repo.create_user(OWNER_ID, &u1).await.unwrap(); let mut u2 = make_user("u2", "tg_2", "telegram"); u2.authorized_at = 2000; - repo.create_user(&u2).await.unwrap(); + repo.create_user(OWNER_ID, &u2).await.unwrap(); - let users = repo.get_all_users().await.unwrap(); + let users = repo.get_all_users(OWNER_ID).await.unwrap(); assert_eq!(users.len(), 2); assert_eq!(users[0].id, "u2"); // more recent first assert_eq!(users[1].id, "u1"); diff --git a/crates/aionui-db/tests/client_preference_repository.rs b/crates/aionui-db/tests/client_preference_repository.rs index ca59e02c3..4c1d7894b 100644 --- a/crates/aionui-db/tests/client_preference_repository.rs +++ b/crates/aionui-db/tests/client_preference_repository.rs @@ -6,6 +6,8 @@ use std::sync::Arc; use aionui_db::{IClientPreferenceRepository, SqliteClientPreferenceRepository, init_database_memory}; +const USER_ID: &str = "system_default_user"; + async fn repo() -> Arc { let db = init_database_memory().await.unwrap(); Arc::new(SqliteClientPreferenceRepository::new(db.pool().clone())) @@ -16,7 +18,7 @@ async fn repo() -> Arc { #[tokio::test] async fn get_all_returns_empty_when_no_preferences() { let r = repo().await; - assert!(r.get_all().await.unwrap().is_empty()); + assert!(r.get_all(USER_ID).await.unwrap().is_empty()); } // -- Upsert and retrieval -- @@ -24,11 +26,11 @@ async fn get_all_returns_empty_when_no_preferences() { #[tokio::test] async fn upsert_then_get_all_returns_inserted_entries() { let r = repo().await; - r.upsert_batch(&[("theme", "\"dark\""), ("pet.size", "360")]) + r.upsert_batch(USER_ID, &[("theme", "\"dark\""), ("pet.size", "360")]) .await .unwrap(); - let prefs = r.get_all().await.unwrap(); + let prefs = r.get_all(USER_ID).await.unwrap(); assert_eq!(prefs.len(), 2); let keys: Vec<&str> = prefs.iter().map(|p| p.key.as_str()).collect(); @@ -39,10 +41,10 @@ async fn upsert_then_get_all_returns_inserted_entries() { #[tokio::test] async fn upsert_overwrites_existing_key() { let r = repo().await; - r.upsert_batch(&[("k", "v1")]).await.unwrap(); - r.upsert_batch(&[("k", "v2")]).await.unwrap(); + r.upsert_batch(USER_ID, &[("k", "v1")]).await.unwrap(); + r.upsert_batch(USER_ID, &[("k", "v2")]).await.unwrap(); - let prefs = r.get_all().await.unwrap(); + let prefs = r.get_all(USER_ID).await.unwrap(); assert_eq!(prefs.len(), 1); assert_eq!(prefs[0].value, "v2"); } @@ -50,8 +52,8 @@ async fn upsert_overwrites_existing_key() { #[tokio::test] async fn upsert_empty_batch_is_noop() { let r = repo().await; - r.upsert_batch(&[]).await.unwrap(); - assert!(r.get_all().await.unwrap().is_empty()); + r.upsert_batch(USER_ID, &[]).await.unwrap(); + assert!(r.get_all(USER_ID).await.unwrap().is_empty()); } // -- Filtered retrieval -- @@ -59,9 +61,11 @@ async fn upsert_empty_batch_is_noop() { #[tokio::test] async fn get_by_keys_returns_only_matching() { let r = repo().await; - r.upsert_batch(&[("a", "1"), ("b", "2"), ("c", "3")]).await.unwrap(); + r.upsert_batch(USER_ID, &[("a", "1"), ("b", "2"), ("c", "3")]) + .await + .unwrap(); - let prefs = r.get_by_keys(&["a", "c"]).await.unwrap(); + let prefs = r.get_by_keys(USER_ID, &["a", "c"]).await.unwrap(); assert_eq!(prefs.len(), 2); let keys: Vec<&str> = prefs.iter().map(|p| p.key.as_str()).collect(); @@ -72,9 +76,9 @@ async fn get_by_keys_returns_only_matching() { #[tokio::test] async fn get_by_keys_omits_nonexistent() { let r = repo().await; - r.upsert_batch(&[("x", "1")]).await.unwrap(); + r.upsert_batch(USER_ID, &[("x", "1")]).await.unwrap(); - let prefs = r.get_by_keys(&["x", "ghost"]).await.unwrap(); + let prefs = r.get_by_keys(USER_ID, &["x", "ghost"]).await.unwrap(); assert_eq!(prefs.len(), 1); assert_eq!(prefs[0].key, "x"); } @@ -82,9 +86,9 @@ async fn get_by_keys_omits_nonexistent() { #[tokio::test] async fn get_by_keys_empty_input_returns_empty() { let r = repo().await; - r.upsert_batch(&[("x", "1")]).await.unwrap(); + r.upsert_batch(USER_ID, &[("x", "1")]).await.unwrap(); - let prefs = r.get_by_keys(&[]).await.unwrap(); + let prefs = r.get_by_keys(USER_ID, &[]).await.unwrap(); assert!(prefs.is_empty()); } @@ -93,11 +97,13 @@ async fn get_by_keys_empty_input_returns_empty() { #[tokio::test] async fn delete_keys_removes_specified_entries() { let r = repo().await; - r.upsert_batch(&[("a", "1"), ("b", "2"), ("c", "3")]).await.unwrap(); + r.upsert_batch(USER_ID, &[("a", "1"), ("b", "2"), ("c", "3")]) + .await + .unwrap(); - r.delete_keys(&["a", "c"]).await.unwrap(); + r.delete_keys(USER_ID, &["a", "c"]).await.unwrap(); - let prefs = r.get_all().await.unwrap(); + let prefs = r.get_all(USER_ID).await.unwrap(); assert_eq!(prefs.len(), 1); assert_eq!(prefs[0].key, "b"); } @@ -105,19 +111,19 @@ async fn delete_keys_removes_specified_entries() { #[tokio::test] async fn delete_nonexistent_keys_is_noop() { let r = repo().await; - r.upsert_batch(&[("x", "1")]).await.unwrap(); - r.delete_keys(&["ghost"]).await.unwrap(); + r.upsert_batch(USER_ID, &[("x", "1")]).await.unwrap(); + r.delete_keys(USER_ID, &["ghost"]).await.unwrap(); - assert_eq!(r.get_all().await.unwrap().len(), 1); + assert_eq!(r.get_all(USER_ID).await.unwrap().len(), 1); } #[tokio::test] async fn delete_empty_keys_is_noop() { let r = repo().await; - r.upsert_batch(&[("x", "1")]).await.unwrap(); - r.delete_keys(&[]).await.unwrap(); + r.upsert_batch(USER_ID, &[("x", "1")]).await.unwrap(); + r.delete_keys(USER_ID, &[]).await.unwrap(); - assert_eq!(r.get_all().await.unwrap().len(), 1); + assert_eq!(r.get_all(USER_ID).await.unwrap().len(), 1); } // -- Value types -- @@ -125,16 +131,19 @@ async fn delete_empty_keys_is_noop() { #[tokio::test] async fn stores_boolean_number_string_json_values() { let r = repo().await; - r.upsert_batch(&[ - ("bool_key", "true"), - ("num_key", "42"), - ("str_key", "\"hello\""), - ("null_key", "null"), - ]) + r.upsert_batch( + USER_ID, + &[ + ("bool_key", "true"), + ("num_key", "42"), + ("str_key", "\"hello\""), + ("null_key", "null"), + ], + ) .await .unwrap(); - let prefs = r.get_all().await.unwrap(); + let prefs = r.get_all(USER_ID).await.unwrap(); assert_eq!(prefs.len(), 4); let find = |k: &str| prefs.iter().find(|p| p.key == k).unwrap().value.as_str(); diff --git a/crates/aionui-db/tests/conversation_repository.rs b/crates/aionui-db/tests/conversation_repository.rs index 47c2eaab7..f846c587c 100644 --- a/crates/aionui-db/tests/conversation_repository.rs +++ b/crates/aionui-db/tests/conversation_repository.rs @@ -1,7 +1,7 @@ use aionui_db::{ - ConversationFilters, ConversationRowUpdate, IConversationRepository, MessagePageCursor, MessagePageDirection, - MessagePageParams, MessageRowUpdate, SqliteConversationRepository, init_database_memory, models::ConversationRow, - models::MessageRow, + ConversationFilters, ConversationRowUpdate, DbError, IConversationRepository, MessagePageCursor, + MessagePageDirection, MessagePageParams, MessageRowUpdate, SqliteConversationRepository, init_database_memory, + models::ConversationRow, models::MessageRow, }; const USER_ID: &str = "system_default_user"; @@ -12,6 +12,16 @@ async fn setup() -> (SqliteConversationRepository, aionui_db::Database) { (repo, db) } +async fn create_user_2(db: &aionui_db::Database) { + sqlx::query( + "INSERT INTO users (id, username, password_hash, created_at, updated_at) \ + VALUES ('user_2', 'other', 'hash', 1000, 1000)", + ) + .execute(db.pool()) + .await + .unwrap(); +} + fn make_conversation(suffix: &str) -> ConversationRow { let now = aionui_common::now_ms(); ConversationRow { @@ -76,13 +86,14 @@ async fn create_get_update_delete_lifecycle() { repo.create(&conv).await.unwrap(); // Get - let found = repo.get(&conv.id).await.unwrap().unwrap(); + let found = repo.get(&conv.user_id, &conv.id).await.unwrap().unwrap(); assert_eq!(found.name, "Conversation lifecycle"); assert_eq!(found.status.as_deref(), Some("pending")); // Update let now = aionui_common::now_ms(); repo.update( + &conv.user_id, &conv.id, &ConversationRowUpdate { name: Some("Updated Name".to_string()), @@ -94,30 +105,31 @@ async fn create_get_update_delete_lifecycle() { .await .unwrap(); - let updated = repo.get(&conv.id).await.unwrap().unwrap(); + let updated = repo.get(&conv.user_id, &conv.id).await.unwrap().unwrap(); assert_eq!(updated.name, "Updated Name"); assert_eq!(updated.status.as_deref(), Some("running")); // Delete - repo.delete(&conv.id).await.unwrap(); - assert!(repo.get(&conv.id).await.unwrap().is_none()); + repo.delete(&conv.user_id, &conv.id).await.unwrap(); + assert!(repo.get(&conv.user_id, &conv.id).await.unwrap().is_none()); } #[tokio::test] async fn delete_conversation_cascades_messages() { - let (repo, _db) = setup().await; + let (repo, db) = setup().await; let conv = make_conversation("cascade"); repo.create(&conv).await.unwrap(); // Insert messages for i in 0..3 { let msg = make_message(&conv.id, &format!("msg {i}")); - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); } // Verify messages exist let msgs = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 50, @@ -129,19 +141,14 @@ async fn delete_conversation_cascades_messages() { assert_eq!(msgs.items.len(), 3); // Delete conversation → messages cascade - repo.delete(&conv.id).await.unwrap(); + repo.delete(&conv.user_id, &conv.id).await.unwrap(); - let msgs = repo - .list_messages_page( - &conv.id, - &MessagePageParams { - limit: 50, - direction: MessagePageDirection::InitialLatest, - }, - ) + let remaining: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM messages WHERE conversation_id = ?") + .bind(&conv.id) + .fetch_one(db.pool()) .await .unwrap(); - assert!(msgs.items.is_empty()); + assert_eq!(remaining, 0); } // ── Cursor pagination ─────────────────────────────────────────────── @@ -414,11 +421,12 @@ async fn initial_latest_returns_latest_limit_in_ascending_order() { for i in 0..10 { let mut msg = make_message(&conv.id, &format!("item {i}")); msg.created_at = (i + 1) as i64 * 1000; - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); } let p1 = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 3, @@ -446,11 +454,12 @@ async fn before_pages_walk_history_without_duplicates() { let mut msg = make_message(&conv.id, &format!("item {i}")); msg.id = format!("msg-{i}"); msg.created_at = (i + 1) as i64 * 1000; - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); } let latest = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 3, @@ -461,6 +470,7 @@ async fn before_pages_walk_history_without_duplicates() { .unwrap(); let older = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 3, @@ -499,11 +509,12 @@ async fn after_returns_newer_rows_with_same_timestamp_tie_breaker() { let mut msg = make_message(&conv.id, id); msg.id = id.to_string(); msg.created_at = 1000; - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); } let page = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 2, @@ -537,20 +548,21 @@ async fn created_at_id_ordering_is_scoped_per_conversation() { let mut a1 = make_message(&conv_a.id, "a1"); a1.id = "a1".to_string(); a1.created_at = 1000; - repo.insert_message(&a1).await.unwrap(); + repo.insert_message(&conv_a.user_id, &a1).await.unwrap(); let mut a2 = make_message(&conv_a.id, "a2"); a2.id = "a2".to_string(); a2.created_at = 2000; - repo.insert_message(&a2).await.unwrap(); + repo.insert_message(&conv_a.user_id, &a2).await.unwrap(); let mut b1 = make_message(&conv_b.id, "b1"); b1.id = "b1".to_string(); b1.created_at = 1000; - repo.insert_message(&b1).await.unwrap(); + repo.insert_message(&conv_b.user_id, &b1).await.unwrap(); let page_a = repo .list_messages_page( + &conv_a.user_id, &conv_a.id, &MessagePageParams { limit: 10, @@ -561,6 +573,7 @@ async fn created_at_id_ordering_is_scoped_per_conversation() { .unwrap(); let page_b = repo .list_messages_page( + &conv_b.user_id, &conv_b.id, &MessagePageParams { limit: 10, @@ -593,13 +606,14 @@ async fn concurrent_insert_message_does_not_require_sequence_allocation() { for i in 0..48 { let repo = repo.clone(); let conv_id = conv.id.clone(); + let user_id = conv.user_id.clone(); let barrier = barrier.clone(); handles.push(tokio::spawn(async move { barrier.wait().await; let mut msg = make_message(&conv_id, &format!("concurrent {i}")); msg.id = format!("msg-concurrent-{i}"); msg.msg_id = Some(format!("cmsg-concurrent-{i}")); - repo.insert_message(&msg).await + repo.insert_message(&user_id, &msg).await })); } @@ -609,6 +623,7 @@ async fn concurrent_insert_message_does_not_require_sequence_allocation() { let page = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 100, @@ -635,11 +650,12 @@ async fn anchor_window_contains_anchor_and_sets_flags() { let mut msg = make_message(&conv.id, &format!("item {i}")); msg.id = format!("msg-{i}"); msg.created_at = i as i64 * 1000; - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); } let page = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 5, @@ -665,22 +681,26 @@ async fn anchor_rejects_legacy_artifact_rows() { let conv = make_conversation("anchor-artifact"); repo.create(&conv).await.unwrap(); - repo.insert_message(&MessageRow { - id: "legacy-cron".into(), - conversation_id: conv.id.clone(), - msg_id: None, - r#type: "cron_trigger".into(), - content: "{}".into(), - position: Some("center".into()), - status: Some("finish".into()), - hidden: false, - created_at: 1000, - }) + repo.insert_message( + &conv.user_id, + &MessageRow { + id: "legacy-cron".into(), + conversation_id: conv.id.clone(), + msg_id: None, + r#type: "cron_trigger".into(), + content: "{}".into(), + position: Some("center".into()), + status: Some("finish".into()), + hidden: false, + created_at: 1000, + }, + ) .await .unwrap(); let err = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 5, @@ -703,15 +723,16 @@ async fn upsert_preserves_existing_created_at() { let mut msg = make_message(&conv.id, "first"); msg.id = "msg-stable".to_string(); - repo.upsert_message(&msg).await.unwrap(); + repo.upsert_message(&conv.user_id, &msg).await.unwrap(); let mut updated = msg.clone(); updated.content = r#"{"content":"updated"}"#.to_string(); updated.created_at = msg.created_at + 5000; - repo.upsert_message(&updated).await.unwrap(); + repo.upsert_message(&conv.user_id, &updated).await.unwrap(); let page = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 10, @@ -733,9 +754,11 @@ async fn update_message_fields() { repo.create(&conv).await.unwrap(); let msg = make_message(&conv.id, "original"); - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); repo.update_message( + &conv.user_id, + &conv.id, &msg.id, &MessageRowUpdate { content: Some(r#"{"content":"modified"}"#.to_string()), @@ -748,6 +771,7 @@ async fn update_message_fields() { let msgs = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 50, @@ -770,13 +794,16 @@ async fn delete_messages_by_conversation_clears_all() { for i in 0..5 { let msg = make_message(&conv.id, &format!("msg {i}")); - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); } - repo.delete_messages_by_conversation(&conv.id).await.unwrap(); + repo.delete_messages_by_conversation(&conv.user_id, &conv.id) + .await + .unwrap(); let result = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 50, @@ -797,25 +824,25 @@ async fn get_message_by_msg_id_triple() { let mut msg = make_message(&conv.id, "findable"); msg.msg_id = Some("unique_msg_123".to_string()); msg.r#type = "tool_call".to_string(); - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); // Match let found = repo - .get_message_by_msg_id(&conv.id, "unique_msg_123", "tool_call") + .get_message_by_msg_id(&conv.user_id, &conv.id, "unique_msg_123", "tool_call") .await .unwrap(); assert!(found.is_some()); // Wrong type → None let not_found = repo - .get_message_by_msg_id(&conv.id, "unique_msg_123", "text") + .get_message_by_msg_id(&conv.user_id, &conv.id, "unique_msg_123", "text") .await .unwrap(); assert!(not_found.is_none()); // Wrong conv → None let not_found = repo - .get_message_by_msg_id("other_conv", "unique_msg_123", "tool_call") + .get_message_by_msg_id(&conv.user_id, "other_conv", "unique_msg_123", "tool_call") .await .unwrap(); assert!(not_found.is_none()); @@ -833,13 +860,13 @@ async fn search_messages_across_conversations() { repo.create(&c2).await.unwrap(); let msg1 = make_message(&c1.id, "Rust 代码审查报告"); - repo.insert_message(&msg1).await.unwrap(); + repo.insert_message(&c1.user_id, &msg1).await.unwrap(); let msg2 = make_message(&c2.id, "Python 代码审查总结"); - repo.insert_message(&msg2).await.unwrap(); + repo.insert_message(&c2.user_id, &msg2).await.unwrap(); let msg3 = make_message(&c1.id, "unrelated content"); - repo.insert_message(&msg3).await.unwrap(); + repo.insert_message(&c1.user_id, &msg3).await.unwrap(); let result = repo.search_messages(USER_ID, "审查", 1, 20).await.unwrap(); assert_eq!(result.total, 2); @@ -858,7 +885,7 @@ async fn search_messages_empty_result() { repo.create(&conv).await.unwrap(); let msg = make_message(&conv.id, "hello world"); - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); let result = repo .search_messages(USER_ID, "nonexistent_keyword", 1, 20) @@ -878,7 +905,7 @@ async fn search_messages_pagination() { for i in 0..5 { let mut msg = make_message(&conv.id, &format!("searchable item {i}")); msg.created_at = (i + 1) as i64 * 1000; - repo.insert_message(&msg).await.unwrap(); + repo.insert_message(&conv.user_id, &msg).await.unwrap(); } let p1 = repo.search_messages(USER_ID, "searchable", 1, 2).await.unwrap(); @@ -906,6 +933,7 @@ async fn pin_and_unpin_conversation() { // Pin let pin_time = aionui_common::now_ms(); repo.update( + &conv.user_id, &conv.id, &ConversationRowUpdate { pinned: Some(true), @@ -917,13 +945,14 @@ async fn pin_and_unpin_conversation() { .await .unwrap(); - let pinned = repo.get(&conv.id).await.unwrap().unwrap(); + let pinned = repo.get(&conv.user_id, &conv.id).await.unwrap().unwrap(); assert!(pinned.pinned); assert_eq!(pinned.pinned_at, Some(pin_time)); // Unpin let now = aionui_common::now_ms(); repo.update( + &conv.user_id, &conv.id, &ConversationRowUpdate { pinned: Some(false), @@ -935,7 +964,7 @@ async fn pin_and_unpin_conversation() { .await .unwrap(); - let unpinned = repo.get(&conv.id).await.unwrap().unwrap(); + let unpinned = repo.get(&conv.user_id, &conv.id).await.unwrap().unwrap(); assert!(!unpinned.pinned); assert!(unpinned.pinned_at.is_none()); } @@ -947,6 +976,7 @@ async fn update_nonexistent_conversation_returns_not_found() { let (repo, _db) = setup().await; let err = repo .update( + "user_1", "nonexistent_id", &ConversationRowUpdate { name: Some("x".to_string()), @@ -961,7 +991,7 @@ async fn update_nonexistent_conversation_returns_not_found() { #[tokio::test] async fn delete_nonexistent_conversation_returns_not_found() { let (repo, _db) = setup().await; - let err = repo.delete("nonexistent_id").await.unwrap_err(); + let err = repo.delete("user_1", "nonexistent_id").await.unwrap_err(); assert!(matches!(err, aionui_db::DbError::NotFound(_))); } @@ -977,6 +1007,8 @@ async fn update_message_nonexistent_returns_not_found() { let (repo, _db) = setup().await; let err = repo .update_message( + "user_1", + "conv_1", "nonexistent_id", &MessageRowUpdate { hidden: Some(true), @@ -998,6 +1030,7 @@ async fn update_extra_replaces_json() { let now = aionui_common::now_ms(); repo.update( + &conv.user_id, &conv.id, &ConversationRowUpdate { extra: Some(r#"{"workspace":"/new","flag":true}"#.to_string()), @@ -1008,7 +1041,7 @@ async fn update_extra_replaces_json() { .await .unwrap(); - let found = repo.get(&conv.id).await.unwrap().unwrap(); + let found = repo.get(&conv.user_id, &conv.id).await.unwrap().unwrap(); assert_eq!(found.extra, r#"{"workspace":"/new","flag":true}"#); } @@ -1018,26 +1051,32 @@ async fn get_messages_excludes_legacy_cron_and_skill_suggest_rows() { let conv = make_conversation("message-filter"); repo.create(&conv).await.unwrap(); - repo.insert_message(&make_message(&conv.id, "visible")).await.unwrap(); + repo.insert_message(&conv.user_id, &make_message(&conv.id, "visible")) + .await + .unwrap(); for (id, ty) in [("legacy-cron", "cron_trigger"), ("legacy-skill", "skill_suggest")] { - repo.insert_message(&MessageRow { - id: id.into(), - conversation_id: conv.id.clone(), - msg_id: None, - r#type: ty.into(), - content: "{}".into(), - position: Some("center".into()), - status: Some("finish".into()), - hidden: false, - created_at: 2000, - }) + repo.insert_message( + &conv.user_id, + &MessageRow { + id: id.into(), + conversation_id: conv.id.clone(), + msg_id: None, + r#type: ty.into(), + content: "{}".into(), + position: Some("center".into()), + status: Some("finish".into()), + hidden: false, + created_at: 2000, + }, + ) .await .unwrap(); } let rows = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 50, @@ -1056,24 +1095,30 @@ async fn list_legacy_cron_trigger_messages_returns_only_trigger_rows() { let conv = make_conversation("legacy-cron-trigger"); repo.create(&conv).await.unwrap(); - repo.insert_message(&MessageRow { - id: aionui_common::generate_prefixed_id("msg"), - conversation_id: conv.id.clone(), - msg_id: Some("legacy-trigger".into()), - r#type: "cron_trigger".into(), - content: r#"{"cron_job_id":"cron_1","cron_job_name":"Daily Report"}"#.into(), - position: Some("center".into()), - status: Some("finish".into()), - hidden: false, - created_at: 1000, - }) + repo.insert_message( + &conv.user_id, + &MessageRow { + id: aionui_common::generate_prefixed_id("msg"), + conversation_id: conv.id.clone(), + msg_id: Some("legacy-trigger".into()), + r#type: "cron_trigger".into(), + content: r#"{"cron_job_id":"cron_1","cron_job_name":"Daily Report"}"#.into(), + position: Some("center".into()), + status: Some("finish".into()), + hidden: false, + created_at: 1000, + }, + ) .await .unwrap(); - repo.insert_message(&make_message(&conv.id, "plain text")) + repo.insert_message(&conv.user_id, &make_message(&conv.id, "plain text")) .await .unwrap(); - let rows = repo.list_legacy_cron_trigger_messages(&conv.id).await.unwrap(); + let rows = repo + .list_legacy_cron_trigger_messages(&conv.user_id, &conv.id) + .await + .unwrap(); assert_eq!(rows.len(), 1); assert_eq!(rows[0].r#type, "cron_trigger"); } @@ -1086,24 +1131,27 @@ async fn artifact_upsert_list_and_mark_saved() { let artifact_id = format!("{}:skill_suggest:cron_1", conv.id); let inserted = repo - .upsert_artifact(&make_artifact(&conv.id, &artifact_id)) + .upsert_artifact(&conv.user_id, &make_artifact(&conv.id, &artifact_id)) .await .unwrap(); assert_eq!(inserted.status, "pending"); - let listed = repo.list_artifacts(&conv.id).await.unwrap(); + let listed = repo.list_artifacts(&conv.user_id, &conv.id).await.unwrap(); assert_eq!(listed.len(), 1); assert_eq!(listed[0].id, artifact_id); let dismissed = repo - .update_artifact_status(&conv.id, &artifact_id, "dismissed", 2000) + .update_artifact_status(&conv.user_id, &conv.id, &artifact_id, "dismissed", 2000) .await .unwrap() .unwrap(); assert_eq!(dismissed.status, "dismissed"); assert_eq!(dismissed.updated_at, 2000); - let saved = repo.mark_skill_suggest_artifacts_saved("cron_1", 3000).await.unwrap(); + let saved = repo + .mark_skill_suggest_artifacts_saved(&conv.user_id, "cron_1", 3000) + .await + .unwrap(); assert_eq!(saved.len(), 1); assert_eq!(saved[0].status, "saved"); assert_eq!(saved[0].updated_at, 3000); @@ -1116,13 +1164,15 @@ async fn delete_artifacts_by_conversation_removes_rows() { repo.create(&conv).await.unwrap(); let artifact_id = format!("{}:skill_suggest:cron_1", conv.id); - repo.upsert_artifact(&make_artifact(&conv.id, &artifact_id)) + repo.upsert_artifact(&conv.user_id, &make_artifact(&conv.id, &artifact_id)) .await .unwrap(); - repo.delete_artifacts_by_conversation(&conv.id).await.unwrap(); + repo.delete_artifacts_by_conversation(&conv.user_id, &conv.id) + .await + .unwrap(); - let listed = repo.list_artifacts(&conv.id).await.unwrap(); + let listed = repo.list_artifacts(&conv.user_id, &conv.id).await.unwrap(); assert!(listed.is_empty()); } @@ -1132,14 +1182,7 @@ async fn delete_artifacts_by_conversation_removes_rows() { async fn list_paginated_scoped_to_user() { let (repo, db) = setup().await; - // Create a second user - sqlx::query( - "INSERT INTO users (id, username, password_hash, created_at, updated_at) \ - VALUES ('user_2', 'other', 'hash', 1000, 1000)", - ) - .execute(db.pool()) - .await - .unwrap(); + create_user_2(&db).await; let c1 = make_conversation("user1-conv"); repo.create(&c1).await.unwrap(); @@ -1162,3 +1205,80 @@ async fn list_paginated_scoped_to_user() { assert_eq!(result.items.len(), 1); assert_eq!(result.items[0].user_id, USER_ID); } + +#[tokio::test] +async fn upsert_message_rejects_cross_user_id_takeover() { + let (repo, db) = setup().await; + create_user_2(&db).await; + + let c1 = make_conversation("user1-message"); + repo.create(&c1).await.unwrap(); + + let mut c2 = make_conversation("user2-message"); + c2.user_id = "user_2".to_string(); + repo.create(&c2).await.unwrap(); + + let mut user_2_msg = make_message(&c2.id, "user 2 original"); + user_2_msg.id = "shared_msg".to_string(); + repo.upsert_message("user_2", &user_2_msg).await.unwrap(); + + let mut takeover = make_message(&c1.id, "user 1 takeover"); + takeover.id = "shared_msg".to_string(); + let err = repo.upsert_message(USER_ID, &takeover).await.unwrap_err(); + assert!(matches!(err, DbError::Conflict(_))); + + let user_2_messages = repo + .list_messages_page( + "user_2", + &c2.id, + &MessagePageParams { + limit: 20, + direction: MessagePageDirection::InitialLatest, + }, + ) + .await + .unwrap(); + assert_eq!(user_2_messages.items.len(), 1); + assert_eq!(user_2_messages.items[0].content, r#"{"content":"user 2 original"}"#); + + let user_1_messages = repo + .list_messages_page( + USER_ID, + &c1.id, + &MessagePageParams { + limit: 20, + direction: MessagePageDirection::InitialLatest, + }, + ) + .await + .unwrap(); + assert!(user_1_messages.items.is_empty()); +} + +#[tokio::test] +async fn upsert_artifact_rejects_cross_user_id_takeover() { + let (repo, db) = setup().await; + create_user_2(&db).await; + + let c1 = make_conversation("user1-artifact"); + repo.create(&c1).await.unwrap(); + + let mut c2 = make_conversation("user2-artifact"); + c2.user_id = "user_2".to_string(); + repo.create(&c2).await.unwrap(); + + let user_2_artifact = make_artifact(&c2.id, "shared_artifact"); + repo.upsert_artifact("user_2", &user_2_artifact).await.unwrap(); + + let takeover = make_artifact(&c1.id, "shared_artifact"); + let err = repo.upsert_artifact(USER_ID, &takeover).await.unwrap_err(); + assert!(matches!(err, DbError::Conflict(_))); + + let user_2_artifacts = repo.list_artifacts("user_2", &c2.id).await.unwrap(); + assert_eq!(user_2_artifacts.len(), 1); + assert_eq!(user_2_artifacts[0].conversation_id, c2.id); + assert_eq!(user_2_artifacts[0].payload, user_2_artifact.payload); + + let user_1_artifacts = repo.list_artifacts(USER_ID, &c1.id).await.unwrap(); + assert!(user_1_artifacts.is_empty()); +} diff --git a/crates/aionui-db/tests/cron_repository.rs b/crates/aionui-db/tests/cron_repository.rs index 0c90c1e2b..98213b56e 100644 --- a/crates/aionui-db/tests/cron_repository.rs +++ b/crates/aionui-db/tests/cron_repository.rs @@ -44,6 +44,7 @@ fn make_job(id: &str) -> CronJobRow { let now = now_ms(); CronJobRow { id: id.into(), + user_id: "user_1".into(), name: "Test Job".into(), enabled: true, schedule_kind: "every".into(), @@ -71,6 +72,28 @@ fn make_job(id: &str) -> CronJobRow { } } +async fn insert_user_and_conversation(db: &aionui_db::Database, user_id: &str, conversation_id: &str) { + sqlx::query( + "INSERT INTO users (id, username, password_hash, created_at, updated_at) \ + VALUES (?, ?, 'hash', 0, 0)", + ) + .bind(user_id) + .bind(user_id) + .execute(db.pool()) + .await + .unwrap(); + sqlx::query( + "INSERT INTO conversations (id, user_id, name, type, created_at, updated_at) \ + VALUES (?, ?, ?, 'normal', 0, 0)", + ) + .bind(conversation_id) + .bind(user_id) + .bind(conversation_id) + .execute(db.pool()) + .await + .unwrap(); +} + // ── A. CRUD ────────────────────────────────────────────────────────── #[tokio::test] @@ -79,7 +102,7 @@ async fn cj1_insert_returns_all_fields() { let job = make_job("cron_cj1"); r.insert(&job).await.unwrap(); - let found = r.get_by_id("cron_cj1").await.unwrap().expect("found"); + let found = r.get_by_id_system("cron_cj1").await.unwrap().expect("found"); assert_eq!(found.id, "cron_cj1"); assert_eq!(found.name, "Test Job"); assert!(found.enabled); @@ -114,13 +137,13 @@ async fn cj2_three_schedule_kinds() { cron_job.schedule_tz = Some("Asia/Shanghai".into()); r.insert(&cron_job).await.unwrap(); - let at = r.get_by_id("cron_at").await.unwrap().unwrap(); + let at = r.get_by_id_system("cron_at").await.unwrap().unwrap(); assert_eq!(at.schedule_kind, "at"); - let every = r.get_by_id("cron_every").await.unwrap().unwrap(); + let every = r.get_by_id_system("cron_every").await.unwrap().unwrap(); assert_eq!(every.schedule_kind, "every"); - let cron = r.get_by_id("cron_cron").await.unwrap().unwrap(); + let cron = r.get_by_id_system("cron_cron").await.unwrap().unwrap(); assert_eq!(cron.schedule_kind, "cron"); assert_eq!(cron.schedule_tz.as_deref(), Some("Asia/Shanghai")); } @@ -129,14 +152,14 @@ async fn cj2_three_schedule_kinds() { async fn cj4_get_by_id_existing() { let (r, _db) = repo().await; r.insert(&make_job("cron_g1")).await.unwrap(); - let found = r.get_by_id("cron_g1").await.unwrap(); + let found = r.get_by_id_system("cron_g1").await.unwrap(); assert!(found.is_some()); } #[tokio::test] async fn cj5_get_by_id_nonexistent() { let (r, _db) = repo().await; - let found = r.get_by_id("cron_nonexistent").await.unwrap(); + let found = r.get_by_id_system("cron_nonexistent").await.unwrap(); assert!(found.is_none()); } @@ -147,7 +170,7 @@ async fn cj6_list_all() { r.insert(&make_job("cron_l2")).await.unwrap(); r.insert(&make_job("cron_l3")).await.unwrap(); - let all = r.list_all().await.unwrap(); + let all = r.list_all_system().await.unwrap(); assert!(all.len() >= 3); } @@ -170,13 +193,90 @@ async fn cj7_list_by_conversation() { other.conversation_id = "conv_2".into(); r.insert(&other).await.unwrap(); - let conv1 = r.list_by_conversation("conv_1").await.unwrap(); + let conv1 = r.list_by_conversation_system("conv_1").await.unwrap(); assert_eq!(conv1.len(), 2); - let conv2 = r.list_by_conversation("conv_2").await.unwrap(); + let conv2 = r.list_by_conversation_system("conv_2").await.unwrap(); assert_eq!(conv2.len(), 1); } +#[tokio::test] +async fn scoped_crud_filters_by_job_user_id() { + let (r, db) = repo().await; + insert_user_and_conversation(&db, "user_2", "conv_2").await; + + r.insert(&make_job("cron_user_1")).await.unwrap(); + let mut other = make_job("cron_user_2"); + other.user_id = "user_2".into(); + other.conversation_id = "conv_2".into(); + r.insert(&other).await.unwrap(); + + assert!(r.get_by_id_for_user("user_2", "cron_user_1").await.unwrap().is_none()); + assert!(r.get_by_id_for_user("user_1", "cron_user_1").await.unwrap().is_some()); + + let user_1_jobs = r.list_all_for_user("user_1").await.unwrap(); + assert_eq!(user_1_jobs.iter().filter(|job| job.id == "cron_user_1").count(), 1); + assert!(!user_1_jobs.iter().any(|job| job.id == "cron_user_2")); + + let wrong_update = r + .update_for_user( + "user_2", + "cron_user_1", + &UpdateCronJobParams { + name: Some("leak".into()), + ..Default::default() + }, + ) + .await; + assert!(matches!(wrong_update, Err(DbError::NotFound(_)))); + + let wrong_delete = r.delete_for_user("user_2", "cron_user_1").await; + assert!(matches!(wrong_delete, Err(DbError::NotFound(_)))); + + assert_eq!( + r.get_by_id_for_user("user_1", "cron_user_1") + .await + .unwrap() + .unwrap() + .name, + "Test Job" + ); +} + +#[tokio::test] +async fn insert_rejects_conversation_owned_by_another_user() { + let (r, db) = repo().await; + insert_user_and_conversation(&db, "user_2", "conv_2").await; + + let mut cross_user = make_job("cron_cross_user_conversation"); + cross_user.user_id = "user_2".into(); + cross_user.conversation_id = "conv_1".into(); + + let result = r.insert(&cross_user).await; + assert!(matches!(result, Err(DbError::NotFound(_)))); + assert!( + r.get_by_id_for_user("user_2", "cron_cross_user_conversation") + .await + .unwrap() + .is_none() + ); +} + +#[tokio::test] +async fn scheduler_enabled_scan_uses_job_user_id() { + let (r, _db) = repo().await; + r.insert(&make_job("cron_owned")).await.unwrap(); + let mut new_conversation = make_job("cron_new_conversation"); + new_conversation.conversation_id = String::new(); + new_conversation.execution_mode = "new_conversation".into(); + r.insert(&new_conversation).await.unwrap(); + + let enabled = r.list_enabled_system().await.unwrap(); + + assert!(enabled.iter().any(|job| job.id == "cron_owned")); + assert!(enabled.iter().any(|job| job.id == "cron_new_conversation")); +} + #[tokio::test] async fn cj8_update_name_and_enabled() { let (r, _db) = repo().await; @@ -187,9 +287,9 @@ async fn cj8_update_name_and_enabled() { enabled: Some(false), ..Default::default() }; - r.update("cron_u1", ¶ms).await.unwrap(); + r.update_system("cron_u1", ¶ms).await.unwrap(); - let updated = r.get_by_id("cron_u1").await.unwrap().unwrap(); + let updated = r.get_by_id_system("cron_u1").await.unwrap().unwrap(); assert_eq!(updated.name, "Renamed"); assert!(!updated.enabled); assert!(updated.updated_at >= updated.created_at); @@ -207,9 +307,9 @@ async fn cj9_update_schedule_type() { next_run_at: Some(Some(9999999)), ..Default::default() }; - r.update("cron_s1", ¶ms).await.unwrap(); + r.update_system("cron_s1", ¶ms).await.unwrap(); - let updated = r.get_by_id("cron_s1").await.unwrap().unwrap(); + let updated = r.get_by_id_system("cron_s1").await.unwrap().unwrap(); assert_eq!(updated.schedule_kind, "cron"); assert_eq!(updated.schedule_value, "0 0 9 * * *"); assert_eq!(updated.schedule_tz.as_deref(), Some("UTC")); @@ -223,7 +323,7 @@ async fn cj10_update_nonexistent() { name: Some("x".into()), ..Default::default() }; - let err = r.update("cron_nope", ¶ms).await.unwrap_err(); + let err = r.update_system("cron_nope", ¶ms).await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -231,16 +331,16 @@ async fn cj10_update_nonexistent() { async fn cj11_delete() { let (r, _db) = repo().await; r.insert(&make_job("cron_d1")).await.unwrap(); - r.delete("cron_d1").await.unwrap(); + r.delete_system("cron_d1").await.unwrap(); - let found = r.get_by_id("cron_d1").await.unwrap(); + let found = r.get_by_id_system("cron_d1").await.unwrap(); assert!(found.is_none()); } #[tokio::test] async fn cj12_delete_nonexistent() { let (r, _db) = repo().await; - let err = r.delete("cron_nope").await.unwrap_err(); + let err = r.delete_system("cron_nope").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -305,7 +405,7 @@ async fn list_enabled_filters_disabled_jobs() { disabled.enabled = false; r.insert(&disabled).await.unwrap(); - let enabled = r.list_enabled().await.unwrap(); + let enabled = r.list_enabled_system().await.unwrap(); assert_eq!(enabled.len(), 1); assert_eq!(enabled[0].id, "cron_en1"); } @@ -321,9 +421,9 @@ async fn sk1_save_skill_content() { skill_content: Some(Some("---\nname: test\n---\nDo something".into())), ..Default::default() }; - r.update("cron_sk1", ¶ms).await.unwrap(); + r.update_system("cron_sk1", ¶ms).await.unwrap(); - let updated = r.get_by_id("cron_sk1").await.unwrap().unwrap(); + let updated = r.get_by_id_system("cron_sk1").await.unwrap().unwrap(); assert!(updated.skill_content.is_some()); assert!(updated.skill_content.unwrap().contains("Do something")); } @@ -335,7 +435,7 @@ async fn sk2_has_skill_after_save() { job.skill_content = Some("---\nname: s\n---\ncontent".into()); r.insert(&job).await.unwrap(); - let found = r.get_by_id("cron_sk2").await.unwrap().unwrap(); + let found = r.get_by_id_system("cron_sk2").await.unwrap().unwrap(); assert!(found.skill_content.is_some()); } @@ -344,7 +444,7 @@ async fn sk3_no_skill_by_default() { let (r, _db) = repo().await; r.insert(&make_job("cron_sk3")).await.unwrap(); - let found = r.get_by_id("cron_sk3").await.unwrap().unwrap(); + let found = r.get_by_id_system("cron_sk3").await.unwrap().unwrap(); assert!(found.skill_content.is_none()); } @@ -355,8 +455,8 @@ async fn sk7_delete_clears_skill() { job.skill_content = Some("content".into()); r.insert(&job).await.unwrap(); - r.delete("cron_sk7").await.unwrap(); - let found = r.get_by_id("cron_sk7").await.unwrap(); + r.delete_system("cron_sk7").await.unwrap(); + let found = r.get_by_id_system("cron_sk7").await.unwrap(); assert!(found.is_none()); } @@ -368,17 +468,17 @@ async fn cd1_delete_by_conversation_removes_all() { r.insert(&make_job("cron_cd1")).await.unwrap(); r.insert(&make_job("cron_cd2")).await.unwrap(); - let deleted = r.delete_by_conversation("conv_1").await.unwrap(); + let deleted = r.delete_by_conversation_system("conv_1").await.unwrap(); assert_eq!(deleted, 2); - let remaining = r.list_all().await.unwrap(); + let remaining = r.list_all_system().await.unwrap(); assert!(remaining.is_empty()); } #[tokio::test] async fn delete_by_conversation_no_match_returns_zero() { let (r, _db) = repo().await; - let deleted = r.delete_by_conversation("conv_none").await.unwrap(); + let deleted = r.delete_by_conversation_system("conv_none").await.unwrap(); assert_eq!(deleted, 0); } @@ -398,9 +498,9 @@ async fn update_execution_state() { next_run_at: Some(Some(now + 60_000)), ..Default::default() }; - r.update("cron_ex1", ¶ms).await.unwrap(); + r.update_system("cron_ex1", ¶ms).await.unwrap(); - let updated = r.get_by_id("cron_ex1").await.unwrap().unwrap(); + let updated = r.get_by_id_system("cron_ex1").await.unwrap().unwrap(); assert_eq!(updated.last_run_at, Some(now)); assert_eq!(updated.last_status.as_deref(), Some("ok")); assert_eq!(updated.run_count, 1); @@ -418,9 +518,9 @@ async fn update_error_state() { retry_count: Some(1), ..Default::default() }; - r.update("cron_err1", ¶ms).await.unwrap(); + r.update_system("cron_err1", ¶ms).await.unwrap(); - let updated = r.get_by_id("cron_err1").await.unwrap().unwrap(); + let updated = r.get_by_id_system("cron_err1").await.unwrap().unwrap(); assert_eq!(updated.last_status.as_deref(), Some("error")); assert_eq!(updated.last_error.as_deref(), Some("timeout after 30s")); assert_eq!(updated.retry_count, 1); @@ -435,7 +535,7 @@ async fn insert_and_retrieve_agent_config() { job.agent_config = Some(r#"{"backend":"openai","name":"GPT-4","modelId":"gpt-4","workspace":"/home/user"}"#.into()); r.insert(&job).await.unwrap(); - let found = r.get_by_id("cron_ag1").await.unwrap().unwrap(); + let found = r.get_by_id_system("cron_ag1").await.unwrap().unwrap(); let config = found.agent_config.unwrap(); assert!(config.contains("openai")); assert!(config.contains("gpt-4")); @@ -450,6 +550,6 @@ async fn insert_new_conversation_mode() { job.execution_mode = "new_conversation".into(); r.insert(&job).await.unwrap(); - let found = r.get_by_id("cron_nc1").await.unwrap().unwrap(); + let found = r.get_by_id_system("cron_nc1").await.unwrap().unwrap(); assert_eq!(found.execution_mode, "new_conversation"); } diff --git a/crates/aionui-db/tests/feedback_diagnostics_repository.rs b/crates/aionui-db/tests/feedback_diagnostics_repository.rs index 820258187..633760943 100644 --- a/crates/aionui-db/tests/feedback_diagnostics_repository.rs +++ b/crates/aionui-db/tests/feedback_diagnostics_repository.rs @@ -11,12 +11,13 @@ async fn insert_feedback_fixture(db: &aionui_db::Database) { let pool = db.pool(); sqlx::query( "INSERT INTO providers \ - (id, platform, name, base_url, api_key_encrypted, models, enabled, \ + (id, user_id, platform, name, base_url, api_key_encrypted, models, enabled, \ capabilities, context_limit, model_protocols, model_enabled, model_health, \ bedrock_config, is_full_url, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind("prov-secret") + .bind("system_default_user") .bind("openrouter") .bind("OpenRouter") .bind("https://sk-live-secret@example.invalid/v1/chat/completions?api_key=sk-query") @@ -510,9 +511,10 @@ async fn insert_feedback_fixture(db: &aionui_db::Database) { } sqlx::query( - "INSERT INTO client_preferences (key, value, updated_at) VALUES (?, ?, ?) \ - ON CONFLICT(key) DO UPDATE SET value = excluded.value, updated_at = excluded.updated_at", + "INSERT INTO client_preferences (user_id, key, value, updated_at) VALUES (?, ?, ?, ?) \ + ON CONFLICT(user_id, key) DO UPDATE SET value = excluded.value, updated_at = excluded.updated_at", ) + .bind("system_default_user") .bind("appearance.uiScale") .bind("1.25") .bind(ANCHOR_UPDATED_AT + 100) @@ -522,9 +524,9 @@ async fn insert_feedback_fixture(db: &aionui_db::Database) { sqlx::query( "INSERT INTO system_settings \ - (id, language, notification_enabled, cron_notification_enabled, command_queue_enabled, save_upload_to_workspace, updated_at) \ - VALUES (1, ?, ?, ?, ?, ?, ?) \ - ON CONFLICT(id) DO UPDATE SET \ + (user_id, language, notification_enabled, cron_notification_enabled, command_queue_enabled, save_upload_to_workspace, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?) \ + ON CONFLICT(user_id) DO UPDATE SET \ language = excluded.language, \ notification_enabled = excluded.notification_enabled, \ cron_notification_enabled = excluded.cron_notification_enabled, \ @@ -532,6 +534,7 @@ async fn insert_feedback_fixture(db: &aionui_db::Database) { save_upload_to_workspace = excluded.save_upload_to_workspace, \ updated_at = excluded.updated_at", ) + .bind("system_default_user") .bind("zh-CN") .bind(true) .bind(false) @@ -612,11 +615,12 @@ async fn insert_mcp_feedback_fixture(db: &aionui_db::Database) { sqlx::query( "INSERT INTO mcp_servers \ - (id, name, description, enabled, transport_type, transport_config, tools, \ + (id, user_id, name, description, enabled, transport_type, transport_config, tools, \ last_test_status, last_connected, original_json, builtin, deleted_at, created_at, updated_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind("mcp-raw-config") + .bind("system_default_user") .bind("raw-config-mcp") .bind(None::) .bind(true) @@ -635,6 +639,57 @@ async fn insert_mcp_feedback_fixture(db: &aionui_db::Database) { .unwrap(); } +async fn insert_foreign_agent_metadata(db: &aionui_db::Database) { + sqlx::query( + "INSERT INTO users \ + (id, user_type, external_user_id, username, email, password_hash, avatar_path, jwt_secret, \ + status, session_generation, created_at, updated_at, last_login) \ + VALUES (?, 'aionpro', ?, ?, ?, NULL, NULL, NULL, 'active', 0, ?, ?, NULL)", + ) + .bind("user-b") + .bind("external-user-b") + .bind("User B") + .bind("user-b@example.invalid") + .bind(ANCHOR_CREATED_AT) + .bind(ANCHOR_UPDATED_AT) + .execute(db.pool()) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO agent_metadata \ + (id, user_id, name, backend, agent_type, agent_source, enabled, command, args, env, \ + native_skills_dirs, behavior_policy, available_modes, available_models, sort_order, \ + last_check_status, last_check_kind, last_check_error_code, last_check_error_message, \ + created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind("foreign-agent") + .bind("user-b") + .bind("Foreign Agent") + .bind("foreign-backend") + .bind("acp") + .bind("custom") + .bind(true) + .bind("foreign-command") + .bind(json!(["run"]).to_string()) + .bind(json!({}).to_string()) + .bind(None::) + .bind(None::) + .bind(json!([{"id":"foreign-mode"}]).to_string()) + .bind(json!([{"id":"foreign-model"}]).to_string()) + .bind(1_i64) + .bind("ok") + .bind("session") + .bind(None::) + .bind(None::) + .bind(ANCHOR_CREATED_AT) + .bind(ANCHOR_UPDATED_AT + 10_000) + .execute(db.pool()) + .await + .unwrap(); +} + #[tokio::test] async fn collects_conversation_auth_signals_without_sensitive_payloads() { let db = init_database_memory().await.unwrap(); @@ -829,10 +884,40 @@ async fn conversation_profile_is_scoped_to_current_user() { assert!(conversation.data["conversation"].is_null()); } +#[tokio::test] +async fn conversation_profile_rejects_foreign_context_agent_metadata() { + let db = init_database_memory().await.unwrap(); + insert_feedback_fixture(&db).await; + insert_foreign_agent_metadata(&db).await; + let repo = SqliteFeedbackDiagnosticsRepository::new(db.pool().clone()); + + let result = repo + .collect_feedback_diagnostics(&FeedbackDiagnosticsRequest { + user_id: "system_default_user".to_owned(), + profiles: vec![FeedbackDiagnosticsProfile::ConversationSession], + context: FeedbackDiagnosticsDbContext { + conversation_id: Some("conv-auth".to_owned()), + agent_id: Some("foreign-agent".to_owned()), + ..FeedbackDiagnosticsDbContext::default() + }, + }) + .await + .unwrap(); + + let conversation = result + .profiles + .iter() + .find(|profile| profile.name == "conversation-session") + .expect("conversation profile should exist"); + assert_eq!(conversation.mode, "detail"); + assert!(conversation.data["agent_metadata"].is_null()); +} + #[tokio::test] async fn global_summary_includes_recent_diagnostics_without_sensitive_payloads() { let db = init_database_memory().await.unwrap(); insert_feedback_fixture(&db).await; + insert_foreign_agent_metadata(&db).await; let repo = SqliteFeedbackDiagnosticsRepository::new(db.pool().clone()); let result = repo @@ -875,6 +960,9 @@ async fn global_summary_includes_recent_diagnostics_without_sensitive_payloads() ); assert_eq!(direct_conversations[0]["assistant_id"], "assistant-direct"); assert_eq!(direct_conversations[0]["agent_id"], "opencode"); + let global_json = serde_json::to_string(&global.data).expect("global summary should serialize"); + assert!(!global_json.contains("foreign-agent")); + assert!(!global_json.contains("Foreign Agent")); assert_eq!(direct_conversations[0]["current_model_id"], "anthropic/claude-sonnet-4"); assert_eq!( direct_conversations[0]["recent_errors"][0]["content"]["error"]["message"], diff --git a/crates/aionui-db/tests/mcp_server_repository.rs b/crates/aionui-db/tests/mcp_server_repository.rs index cb1d7044c..953cd64bf 100644 --- a/crates/aionui-db/tests/mcp_server_repository.rs +++ b/crates/aionui-db/tests/mcp_server_repository.rs @@ -10,6 +10,8 @@ use aionui_db::{ init_database_memory, }; +const USER_ID: &str = "system_default_user"; + async fn repo() -> (Arc, aionui_db::Database) { let db = init_database_memory().await.unwrap(); let r = Arc::new(SqliteMcpServerRepository::new(db.pool().clone())); @@ -18,6 +20,7 @@ async fn repo() -> (Arc, aionui_db::Database) { fn stdio_params() -> CreateMcpServerParams<'static> { CreateMcpServerParams { + user_id: USER_ID, name: "test-mcp", description: Some("A test MCP server"), enabled: false, @@ -31,6 +34,7 @@ fn stdio_params() -> CreateMcpServerParams<'static> { fn http_params() -> CreateMcpServerParams<'static> { CreateMcpServerParams { + user_id: USER_ID, name: "http-mcp", description: None, enabled: true, @@ -44,6 +48,7 @@ fn http_params() -> CreateMcpServerParams<'static> { fn sse_params() -> CreateMcpServerParams<'static> { CreateMcpServerParams { + user_id: USER_ID, name: "sse-mcp", description: Some("SSE transport server"), enabled: false, @@ -112,7 +117,7 @@ async fn find_by_id_returns_full_record() { let (r, _db) = repo().await; let created = r.create(stdio_params()).await.unwrap(); - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(found.id, created.id); assert_eq!(found.name, "test-mcp"); assert_eq!(found.transport_type, "stdio"); @@ -122,7 +127,7 @@ async fn find_by_id_returns_full_record() { #[tokio::test] async fn find_by_id_nonexistent_returns_none() { let (r, _db) = repo().await; - assert!(r.find_by_id("nonexistent").await.unwrap().is_none()); + assert!(r.find_by_id(USER_ID, "nonexistent").await.unwrap().is_none()); } // -- Find by name -- @@ -132,14 +137,14 @@ async fn find_by_name_returns_matching_record() { let (r, _db) = repo().await; let created = r.create(stdio_params()).await.unwrap(); - let found = r.find_by_name("test-mcp").await.unwrap().unwrap(); + let found = r.find_by_name(USER_ID, "test-mcp").await.unwrap().unwrap(); assert_eq!(found.id, created.id); } #[tokio::test] async fn find_by_name_nonexistent_returns_none() { let (r, _db) = repo().await; - assert!(r.find_by_name("nope").await.unwrap().is_none()); + assert!(r.find_by_name(USER_ID, "nope").await.unwrap().is_none()); } // -- R-3/R-4: List servers -- @@ -147,7 +152,7 @@ async fn find_by_name_nonexistent_returns_none() { #[tokio::test] async fn list_empty_returns_empty_vec() { let (r, _db) = repo().await; - let servers = r.list().await.unwrap(); + let servers = r.list(USER_ID).await.unwrap(); assert!(servers.is_empty()); } @@ -158,7 +163,7 @@ async fn list_returns_all_ordered_by_created_at() { let s2 = r.create(http_params()).await.unwrap(); let s3 = r.create(sse_params()).await.unwrap(); - let all = r.list().await.unwrap(); + let all = r.list(USER_ID).await.unwrap(); assert_eq!(all.len(), 3); assert_eq!(all[0].id, s1.id); assert_eq!(all[1].id, s2.id); @@ -174,6 +179,7 @@ async fn update_name_only_preserves_other_fields() { let updated = r .update( + USER_ID, &created.id, UpdateMcpServerParams { name: Some("renamed-mcp"), @@ -197,6 +203,7 @@ async fn update_transport_type_and_config() { let updated = r .update( + USER_ID, &created.id, UpdateMcpServerParams { transport_type: Some("http"), @@ -218,6 +225,7 @@ async fn update_description() { let updated = r .update( + USER_ID, &created.id, UpdateMcpServerParams { description: Some(Some("new desc")), @@ -236,7 +244,7 @@ async fn update_description() { async fn update_nonexistent_returns_not_found() { let (r, _db) = repo().await; let err = r - .update("nonexistent", UpdateMcpServerParams::default()) + .update(USER_ID, "nonexistent", UpdateMcpServerParams::default()) .await .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); @@ -252,6 +260,7 @@ async fn update_name_to_existing_name_returns_conflict() { let err = r .update( + USER_ID, &s2.id, UpdateMcpServerParams { name: Some("test-mcp"), @@ -274,6 +283,7 @@ async fn update_can_clear_optional_fields() { let updated = r .update( + USER_ID, &created.id, UpdateMcpServerParams { description: Some(None), @@ -296,6 +306,7 @@ async fn update_persists_to_database() { let created = r.create(stdio_params()).await.unwrap(); r.update( + USER_ID, &created.id, UpdateMcpServerParams { enabled: Some(true), @@ -305,7 +316,7 @@ async fn update_persists_to_database() { .await .unwrap(); - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert!(found.enabled); } @@ -316,9 +327,9 @@ async fn delete_existing_removes_record() { let (r, _db) = repo().await; let created = r.create(stdio_params()).await.unwrap(); - r.delete(&created.id).await.unwrap(); - assert!(r.find_by_id(&created.id).await.unwrap().is_none()); - let deleted = r.find_by_id_any(&created.id).await.unwrap().unwrap(); + r.delete(USER_ID, &created.id).await.unwrap(); + assert!(r.find_by_id(USER_ID, &created.id).await.unwrap().is_none()); + let deleted = r.find_by_id_any(USER_ID, &created.id).await.unwrap().unwrap(); assert!(deleted.deleted_at.is_some()); assert!(!deleted.enabled); } @@ -326,7 +337,7 @@ async fn delete_existing_removes_record() { #[tokio::test] async fn delete_nonexistent_returns_not_found() { let (r, _db) = repo().await; - let err = r.delete("nonexistent").await.unwrap_err(); + let err = r.delete(USER_ID, "nonexistent").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -336,9 +347,9 @@ async fn delete_one_does_not_affect_others() { let s1 = r.create(stdio_params()).await.unwrap(); let s2 = r.create(http_params()).await.unwrap(); - r.delete(&s1.id).await.unwrap(); + r.delete(USER_ID, &s1.id).await.unwrap(); - let remaining = r.list().await.unwrap(); + let remaining = r.list(USER_ID).await.unwrap(); assert_eq!(remaining.len(), 1); assert_eq!(remaining[0].id, s2.id); } @@ -348,10 +359,10 @@ async fn list_by_ids_any_includes_soft_deleted_rows() { let (r, _db) = repo().await; let active = r.create(stdio_params()).await.unwrap(); let deleted = r.create(http_params()).await.unwrap(); - r.delete(&deleted.id).await.unwrap(); + r.delete(USER_ID, &deleted.id).await.unwrap(); let rows = r - .list_by_ids_any(&[deleted.id.clone(), active.id.clone()]) + .list_by_ids_any(USER_ID, &[deleted.id.clone(), active.id.clone()]) .await .unwrap(); @@ -369,7 +380,7 @@ async fn batch_upsert_creates_new_servers() { let (r, _db) = repo().await; let results = r - .batch_upsert(&[stdio_params(), http_params(), sse_params()]) + .batch_upsert(USER_ID, &[stdio_params(), http_params(), sse_params()]) .await .unwrap(); @@ -378,7 +389,7 @@ async fn batch_upsert_creates_new_servers() { assert_eq!(results[1].name, "http-mcp"); assert_eq!(results[2].name, "sse-mcp"); - let all = r.list().await.unwrap(); + let all = r.list(USER_ID).await.unwrap(); assert_eq!(all.len(), 3); } @@ -389,14 +400,17 @@ async fn batch_upsert_updates_existing_by_name() { assert!(!existing.enabled); let results = r - .batch_upsert(&[ - CreateMcpServerParams { - enabled: true, - description: Some("Updated via batch"), - ..stdio_params() - }, - http_params(), - ]) + .batch_upsert( + USER_ID, + &[ + CreateMcpServerParams { + enabled: true, + description: Some("Updated via batch"), + ..stdio_params() + }, + http_params(), + ], + ) .await .unwrap(); @@ -412,7 +426,7 @@ async fn batch_upsert_updates_existing_by_name() { #[tokio::test] async fn batch_upsert_empty_list() { let (r, _db) = repo().await; - let results = r.batch_upsert(&[]).await.unwrap(); + let results = r.batch_upsert(USER_ID, &[]).await.unwrap(); assert!(results.is_empty()); } @@ -424,9 +438,11 @@ async fn update_status_with_timestamp() { let created = r.create(stdio_params()).await.unwrap(); let ts = aionui_common::now_ms(); - r.update_status(&created.id, "connected", Some(ts)).await.unwrap(); + r.update_status(USER_ID, &created.id, "connected", Some(ts)) + .await + .unwrap(); - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(found.last_test_status, "connected"); assert_eq!(found.last_connected, Some(ts)); } @@ -437,11 +453,13 @@ async fn update_status_without_timestamp_preserves_existing() { let created = r.create(stdio_params()).await.unwrap(); let ts = aionui_common::now_ms(); - r.update_status(&created.id, "connected", Some(ts)).await.unwrap(); + r.update_status(USER_ID, &created.id, "connected", Some(ts)) + .await + .unwrap(); - r.update_status(&created.id, "error", None).await.unwrap(); + r.update_status(USER_ID, &created.id, "error", None).await.unwrap(); - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(found.last_test_status, "error"); assert_eq!(found.last_connected, Some(ts)); } @@ -449,7 +467,10 @@ async fn update_status_without_timestamp_preserves_existing() { #[tokio::test] async fn update_status_nonexistent_returns_not_found() { let (r, _db) = repo().await; - let err = r.update_status("nonexistent", "connected", None).await.unwrap_err(); + let err = r + .update_status(USER_ID, "nonexistent", "connected", None) + .await + .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -461,9 +482,9 @@ async fn update_tools_sets_json() { let created = r.create(stdio_params()).await.unwrap(); let tools_json = r#"[{"name":"read_file","description":"Read a file"}]"#; - r.update_tools(&created.id, Some(tools_json)).await.unwrap(); + r.update_tools(USER_ID, &created.id, Some(tools_json)).await.unwrap(); - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(found.tools.as_deref(), Some(tools_json)); } @@ -479,16 +500,16 @@ async fn update_tools_clears_to_null() { .unwrap(); assert!(created.tools.is_some()); - r.update_tools(&created.id, None).await.unwrap(); + r.update_tools(USER_ID, &created.id, None).await.unwrap(); - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert!(found.tools.is_none()); } #[tokio::test] async fn update_tools_nonexistent_returns_not_found() { let (r, _db) = repo().await; - let err = r.update_tools("nonexistent", Some("[]")).await.unwrap_err(); + let err = r.update_tools(USER_ID, "nonexistent", Some("[]")).await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -503,12 +524,13 @@ async fn full_crud_lifecycle() { assert_eq!(created.name, "test-mcp"); // Read - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(found.id, created.id); // Update let updated = r .update( + USER_ID, &created.id, UpdateMcpServerParams { name: Some("renamed-mcp"), @@ -522,20 +544,20 @@ async fn full_crud_lifecycle() { assert!(updated.enabled); // Find by new name - let by_name = r.find_by_name("renamed-mcp").await.unwrap().unwrap(); + let by_name = r.find_by_name(USER_ID, "renamed-mcp").await.unwrap().unwrap(); assert_eq!(by_name.id, created.id); // Update status - r.update_status(&created.id, "connected", Some(aionui_common::now_ms())) + r.update_status(USER_ID, &created.id, "connected", Some(aionui_common::now_ms())) .await .unwrap(); - let after_status = r.find_by_id(&created.id).await.unwrap().unwrap(); + let after_status = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(after_status.last_test_status, "connected"); // Delete - r.delete(&created.id).await.unwrap(); - assert!(r.find_by_id(&created.id).await.unwrap().is_none()); - assert!(r.list().await.unwrap().is_empty()); + r.delete(USER_ID, &created.id).await.unwrap(); + assert!(r.find_by_id(USER_ID, &created.id).await.unwrap().is_none()); + assert!(r.list(USER_ID).await.unwrap().is_empty()); } // -- Builtin server -- diff --git a/crates/aionui-db/tests/oauth_token_repository.rs b/crates/aionui-db/tests/oauth_token_repository.rs index 3906909b2..91fb56c1a 100644 --- a/crates/aionui-db/tests/oauth_token_repository.rs +++ b/crates/aionui-db/tests/oauth_token_repository.rs @@ -9,6 +9,8 @@ use aionui_db::{ DbError, IOAuthTokenRepository, SqliteOAuthTokenRepository, UpsertOAuthTokenParams, init_database_memory, }; +const USER_ID: &str = "system_default_user"; + async fn repo() -> (Arc, aionui_db::Database) { let db = init_database_memory().await.unwrap(); let r = Arc::new(SqliteOAuthTokenRepository::new(db.pool().clone())); @@ -17,6 +19,7 @@ async fn repo() -> (Arc, aionui_db::Database) { fn sample_params() -> UpsertOAuthTokenParams<'static> { UpsertOAuthTokenParams { + user_id: USER_ID, server_url: "https://mcp.example.com", access_token: "enc_access_token_123", refresh_token: Some("enc_refresh_token_456"), @@ -30,7 +33,7 @@ fn sample_params() -> UpsertOAuthTokenParams<'static> { #[tokio::test] async fn get_by_url_nonexistent_returns_none() { let (r, _db) = repo().await; - assert!(r.get_by_url("https://nope.com").await.unwrap().is_none()); + assert!(r.get_by_url(USER_ID, "https://nope.com").await.unwrap().is_none()); } // -- OA-2: Insert and retrieve -- @@ -47,7 +50,7 @@ async fn upsert_insert_then_get_returns_token() { assert_eq!(inserted.expires_at, Some(1700000000000)); assert!(inserted.created_at > 0); - let found = r.get_by_url("https://mcp.example.com").await.unwrap().unwrap(); + let found = r.get_by_url(USER_ID, "https://mcp.example.com").await.unwrap().unwrap(); assert_eq!(found.access_token, "enc_access_token_123"); } @@ -60,6 +63,7 @@ async fn upsert_updates_existing_token() { let updated = r .upsert(UpsertOAuthTokenParams { + user_id: USER_ID, server_url: "https://mcp.example.com", access_token: "new_access_token", refresh_token: None, @@ -84,6 +88,7 @@ async fn upsert_without_refresh_token_or_expires_at() { let (r, _db) = repo().await; let token = r .upsert(UpsertOAuthTokenParams { + user_id: USER_ID, server_url: "https://simple.example.com", access_token: "simple_token", refresh_token: None, @@ -104,8 +109,13 @@ async fn delete_existing_token() { let (r, _db) = repo().await; r.upsert(sample_params()).await.unwrap(); - r.delete("https://mcp.example.com").await.unwrap(); - assert!(r.get_by_url("https://mcp.example.com").await.unwrap().is_none()); + r.delete(USER_ID, "https://mcp.example.com").await.unwrap(); + assert!( + r.get_by_url(USER_ID, "https://mcp.example.com") + .await + .unwrap() + .is_none() + ); } // -- OA-7: Delete idempotency (returns NotFound for nonexistent) -- @@ -113,7 +123,7 @@ async fn delete_existing_token() { #[tokio::test] async fn delete_nonexistent_returns_not_found() { let (r, _db) = repo().await; - let err = r.delete("https://nope.com").await.unwrap_err(); + let err = r.delete(USER_ID, "https://nope.com").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -122,7 +132,7 @@ async fn delete_nonexistent_returns_not_found() { #[tokio::test] async fn list_authenticated_urls_empty() { let (r, _db) = repo().await; - let urls = r.list_authenticated_urls().await.unwrap(); + let urls = r.list_authenticated_urls(USER_ID).await.unwrap(); assert!(urls.is_empty()); } @@ -131,6 +141,7 @@ async fn list_authenticated_urls_returns_all() { let (r, _db) = repo().await; r.upsert(sample_params()).await.unwrap(); r.upsert(UpsertOAuthTokenParams { + user_id: USER_ID, server_url: "https://other.example.com", access_token: "token2", refresh_token: None, @@ -140,7 +151,7 @@ async fn list_authenticated_urls_returns_all() { .await .unwrap(); - let urls = r.list_authenticated_urls().await.unwrap(); + let urls = r.list_authenticated_urls(USER_ID).await.unwrap(); assert_eq!(urls.len(), 2); assert!(urls.contains(&"https://mcp.example.com".to_string())); assert!(urls.contains(&"https://other.example.com".to_string())); @@ -153,6 +164,7 @@ async fn delete_one_does_not_affect_others() { let (r, _db) = repo().await; r.upsert(sample_params()).await.unwrap(); r.upsert(UpsertOAuthTokenParams { + user_id: USER_ID, server_url: "https://other.example.com", access_token: "token2", refresh_token: None, @@ -162,9 +174,9 @@ async fn delete_one_does_not_affect_others() { .await .unwrap(); - r.delete("https://mcp.example.com").await.unwrap(); + r.delete(USER_ID, "https://mcp.example.com").await.unwrap(); - let urls = r.list_authenticated_urls().await.unwrap(); + let urls = r.list_authenticated_urls(USER_ID).await.unwrap(); assert_eq!(urls.len(), 1); assert_eq!(urls[0], "https://other.example.com"); } @@ -176,20 +188,26 @@ async fn full_oauth_lifecycle() { let (r, _db) = repo().await; // Initially no tokens - assert!(r.list_authenticated_urls().await.unwrap().is_empty()); - assert!(r.get_by_url("https://mcp.example.com").await.unwrap().is_none()); + assert!(r.list_authenticated_urls(USER_ID).await.unwrap().is_empty()); + assert!( + r.get_by_url(USER_ID, "https://mcp.example.com") + .await + .unwrap() + .is_none() + ); // Store token let token = r.upsert(sample_params()).await.unwrap(); assert_eq!(token.access_token, "enc_access_token_123"); // Verify stored - let urls = r.list_authenticated_urls().await.unwrap(); + let urls = r.list_authenticated_urls(USER_ID).await.unwrap(); assert_eq!(urls.len(), 1); // Update token (refresh) let refreshed = r .upsert(UpsertOAuthTokenParams { + user_id: USER_ID, server_url: "https://mcp.example.com", access_token: "refreshed_token", refresh_token: Some("new_refresh"), @@ -202,6 +220,6 @@ async fn full_oauth_lifecycle() { assert_eq!(refreshed.created_at, token.created_at); // Logout (delete) - r.delete("https://mcp.example.com").await.unwrap(); - assert!(r.list_authenticated_urls().await.unwrap().is_empty()); + r.delete(USER_ID, "https://mcp.example.com").await.unwrap(); + assert!(r.list_authenticated_urls(USER_ID).await.unwrap().is_empty()); } diff --git a/crates/aionui-db/tests/provider_repository.rs b/crates/aionui-db/tests/provider_repository.rs index aa3cc1884..37ae86c81 100644 --- a/crates/aionui-db/tests/provider_repository.rs +++ b/crates/aionui-db/tests/provider_repository.rs @@ -9,6 +9,8 @@ use aionui_db::{ init_database_memory, }; +const USER_ID: &str = "system_default_user"; + async fn repo() -> Arc { let db = init_database_memory().await.unwrap(); Arc::new(SqliteProviderRepository::new(db.pool().clone())) @@ -17,6 +19,7 @@ async fn repo() -> Arc { fn sample_params() -> CreateProviderParams<'static> { CreateProviderParams { id: None, + user_id: USER_ID, platform: "anthropic", name: "Anthropic", base_url: "https://api.anthropic.com", @@ -39,7 +42,7 @@ fn sample_params() -> CreateProviderParams<'static> { #[tokio::test] async fn list_returns_empty_when_no_providers() { let r = repo().await; - assert!(r.list().await.unwrap().is_empty()); + assert!(r.list(USER_ID).await.unwrap().is_empty()); } // -- Create -- @@ -96,7 +99,7 @@ async fn find_by_id_existing_returns_provider() { let r = repo().await; let created = r.create(sample_params()).await.unwrap(); - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(found.id, created.id); assert_eq!(found.name, "Anthropic"); } @@ -104,7 +107,7 @@ async fn find_by_id_existing_returns_provider() { #[tokio::test] async fn find_by_id_nonexistent_returns_none() { let r = repo().await; - assert!(r.find_by_id("no_such_id").await.unwrap().is_none()); + assert!(r.find_by_id(USER_ID, "no_such_id").await.unwrap().is_none()); } // -- List -- @@ -122,7 +125,7 @@ async fn list_returns_all_providers_in_creation_order() { .await .unwrap(); - let all = r.list().await.unwrap(); + let all = r.list(USER_ID).await.unwrap(); assert_eq!(all.len(), 2); assert_eq!(all[0].id, first.id); assert_eq!(all[1].id, second.id); @@ -137,6 +140,7 @@ async fn update_partial_fields_preserves_others() { let updated = r .update( + USER_ID, &created.id, UpdateProviderParams { name: Some("New Name"), @@ -160,6 +164,7 @@ async fn update_api_key_changes_encrypted_value() { let updated = r .update( + USER_ID, &created.id, UpdateProviderParams { api_key_encrypted: Some("new_encrypted"), @@ -179,6 +184,7 @@ async fn update_model_settings_replaces_per_model_overrides() { let updated = r .update( + USER_ID, &created.id, UpdateProviderParams { model_settings: Some(r#"{"gpt-5.6-sol":{"openai_api_mode":"responses"}}"#), @@ -203,6 +209,7 @@ async fn update_optional_fields_can_be_set_and_cleared() { // Set let with_config = r .update( + USER_ID, &created.id, UpdateProviderParams { bedrock_config: Some(Some(r#"{"region":"eu-west-1"}"#)), @@ -216,6 +223,7 @@ async fn update_optional_fields_can_be_set_and_cleared() { // Clear let cleared = r .update( + USER_ID, &created.id, UpdateProviderParams { bedrock_config: Some(None), @@ -231,7 +239,7 @@ async fn update_optional_fields_can_be_set_and_cleared() { async fn update_nonexistent_returns_not_found() { let r = repo().await; let err = r - .update("nonexistent", UpdateProviderParams::default()) + .update(USER_ID, "nonexistent", UpdateProviderParams::default()) .await .unwrap_err(); assert!(matches!(err, DbError::NotFound(_)), "expected NotFound, got: {err:?}"); @@ -244,6 +252,7 @@ async fn update_advances_updated_at() { let updated = r .update( + USER_ID, &created.id, UpdateProviderParams { name: Some("Changed"), @@ -264,14 +273,14 @@ async fn delete_removes_provider() { let r = repo().await; let created = r.create(sample_params()).await.unwrap(); - r.delete(&created.id).await.unwrap(); - assert!(r.find_by_id(&created.id).await.unwrap().is_none()); + r.delete(USER_ID, &created.id).await.unwrap(); + assert!(r.find_by_id(USER_ID, &created.id).await.unwrap().is_none()); } #[tokio::test] async fn delete_nonexistent_returns_not_found() { let r = repo().await; - let err = r.delete("nonexistent").await.unwrap_err(); + let err = r.delete(USER_ID, "nonexistent").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_)), "expected NotFound, got: {err:?}"); } @@ -287,9 +296,9 @@ async fn delete_does_not_affect_other_providers() { .await .unwrap(); - r.delete(&p1.id).await.unwrap(); + r.delete(USER_ID, &p1.id).await.unwrap(); - let all = r.list().await.unwrap(); + let all = r.list(USER_ID).await.unwrap(); assert_eq!(all.len(), 1); assert_eq!(all[0].id, p2.id); } diff --git a/crates/aionui-db/tests/remote_agent_repository.rs b/crates/aionui-db/tests/remote_agent_repository.rs index 708cc7c45..decacf183 100644 --- a/crates/aionui-db/tests/remote_agent_repository.rs +++ b/crates/aionui-db/tests/remote_agent_repository.rs @@ -10,6 +10,8 @@ use aionui_db::{ init_database_memory, }; +const USER_ID: &str = "system_default_user"; + async fn repo() -> (Arc, aionui_db::Database) { let db = init_database_memory().await.unwrap(); let r = Arc::new(SqliteRemoteAgentRepository::new(db.pool().clone())); @@ -18,6 +20,7 @@ async fn repo() -> (Arc, aionui_db::Database) { fn bearer_params() -> CreateRemoteAgentParams<'static> { CreateRemoteAgentParams { + user_id: USER_ID, name: "Remote Server", protocol: "acp", url: "wss://remote.example.com", @@ -35,6 +38,7 @@ fn bearer_params() -> CreateRemoteAgentParams<'static> { fn openclaw_params() -> CreateRemoteAgentParams<'static> { CreateRemoteAgentParams { + user_id: USER_ID, name: "OpenClaw Agent", protocol: "openClaw", url: "wss://openclaw.example.com", @@ -87,7 +91,7 @@ async fn create_openclaw_agent_includes_device_fields() { #[tokio::test] async fn list_empty_returns_empty_vec() { let (r, _db) = repo().await; - let agents = r.list().await.unwrap(); + let agents = r.list(USER_ID).await.unwrap(); assert!(agents.is_empty()); } @@ -97,7 +101,7 @@ async fn list_returns_all_agents_ordered() { let a1 = r.create(bearer_params()).await.unwrap(); let a2 = r.create(openclaw_params()).await.unwrap(); - let all = r.list().await.unwrap(); + let all = r.list(USER_ID).await.unwrap(); assert_eq!(all.len(), 2); assert_eq!(all[0].id, a1.id); assert_eq!(all[1].id, a2.id); @@ -110,7 +114,7 @@ async fn find_by_id_returns_full_record() { let (r, _db) = repo().await; let created = r.create(bearer_params()).await.unwrap(); - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(found.id, created.id); assert_eq!(found.name, "Remote Server"); assert_eq!(found.auth_token.as_deref(), Some("encrypted_bearer_token")); @@ -120,7 +124,7 @@ async fn find_by_id_returns_full_record() { #[tokio::test] async fn find_by_id_nonexistent_returns_none() { let (r, _db) = repo().await; - let result = r.find_by_id("nonexistent-uuid").await.unwrap(); + let result = r.find_by_id(USER_ID, "nonexistent-uuid").await.unwrap(); assert!(result.is_none()); } @@ -133,6 +137,7 @@ async fn update_name_only_preserves_other_fields() { let updated = r .update( + USER_ID, &created.id, UpdateRemoteAgentParams { name: Some("New Name"), @@ -157,6 +162,7 @@ async fn update_multiple_fields() { let updated = r .update( + USER_ID, &created.id, UpdateRemoteAgentParams { name: Some("Updated"), @@ -177,7 +183,7 @@ async fn update_multiple_fields() { async fn update_nonexistent_returns_not_found() { let (r, _db) = repo().await; let err = r - .update("nonexistent-uuid", UpdateRemoteAgentParams::default()) + .update(USER_ID, "nonexistent-uuid", UpdateRemoteAgentParams::default()) .await .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); @@ -192,6 +198,7 @@ async fn update_can_clear_optional_fields() { let updated = r .update( + USER_ID, &created.id, UpdateRemoteAgentParams { description: Some(None), @@ -212,6 +219,7 @@ async fn update_persists_to_database() { let created = r.create(bearer_params()).await.unwrap(); r.update( + USER_ID, &created.id, UpdateRemoteAgentParams { name: Some("Persisted Name"), @@ -221,7 +229,7 @@ async fn update_persists_to_database() { .await .unwrap(); - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(found.name, "Persisted Name"); } @@ -232,16 +240,16 @@ async fn delete_existing_removes_record() { let (r, _db) = repo().await; let created = r.create(bearer_params()).await.unwrap(); - r.delete(&created.id).await.unwrap(); + r.delete(USER_ID, &created.id).await.unwrap(); - let found = r.find_by_id(&created.id).await.unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap(); assert!(found.is_none()); } #[tokio::test] async fn delete_nonexistent_returns_not_found() { let (r, _db) = repo().await; - let err = r.delete("nonexistent-uuid").await.unwrap_err(); + let err = r.delete(USER_ID, "nonexistent-uuid").await.unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); } @@ -251,9 +259,9 @@ async fn delete_one_does_not_affect_others() { let a1 = r.create(bearer_params()).await.unwrap(); let a2 = r.create(openclaw_params()).await.unwrap(); - r.delete(&a1.id).await.unwrap(); + r.delete(USER_ID, &a1.id).await.unwrap(); - let remaining = r.list().await.unwrap(); + let remaining = r.list(USER_ID).await.unwrap(); assert_eq!(remaining.len(), 1); assert_eq!(remaining[0].id, a2.id); } @@ -266,9 +274,11 @@ async fn update_status_to_connected_with_timestamp() { let created = r.create(bearer_params()).await.unwrap(); let ts = aionui_common::now_ms(); - r.update_status(&created.id, "connected", Some(ts)).await.unwrap(); + r.update_status(USER_ID, &created.id, "connected", Some(ts)) + .await + .unwrap(); - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(found.status, "connected"); assert_eq!(found.last_connected_at, Some(ts)); } @@ -278,9 +288,9 @@ async fn update_status_to_error_without_timestamp() { let (r, _db) = repo().await; let created = r.create(bearer_params()).await.unwrap(); - r.update_status(&created.id, "error", None).await.unwrap(); + r.update_status(USER_ID, &created.id, "error", None).await.unwrap(); - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(found.status, "error"); assert!(found.last_connected_at.is_none()); } @@ -289,7 +299,7 @@ async fn update_status_to_error_without_timestamp() { async fn update_status_nonexistent_returns_not_found() { let (r, _db) = repo().await; let err = r - .update_status("nonexistent-uuid", "connected", None) + .update_status(USER_ID, "nonexistent-uuid", "connected", None) .await .unwrap_err(); assert!(matches!(err, DbError::NotFound(_))); @@ -306,12 +316,13 @@ async fn full_crud_lifecycle() { assert_eq!(created.name, "Remote Server"); // Read - let found = r.find_by_id(&created.id).await.unwrap().unwrap(); + let found = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(found.id, created.id); // Update let updated = r .update( + USER_ID, &created.id, UpdateRemoteAgentParams { name: Some("Renamed Server"), @@ -325,16 +336,16 @@ async fn full_crud_lifecycle() { assert_eq!(updated.description.as_deref(), Some("Updated desc")); // Update status - r.update_status(&created.id, "connected", Some(aionui_common::now_ms())) + r.update_status(USER_ID, &created.id, "connected", Some(aionui_common::now_ms())) .await .unwrap(); - let after_status = r.find_by_id(&created.id).await.unwrap().unwrap(); + let after_status = r.find_by_id(USER_ID, &created.id).await.unwrap().unwrap(); assert_eq!(after_status.status, "connected"); // Delete - r.delete(&created.id).await.unwrap(); - assert!(r.find_by_id(&created.id).await.unwrap().is_none()); + r.delete(USER_ID, &created.id).await.unwrap(); + assert!(r.find_by_id(USER_ID, &created.id).await.unwrap().is_none()); // List should be empty - assert!(r.list().await.unwrap().is_empty()); + assert!(r.list(USER_ID).await.unwrap().is_empty()); } diff --git a/crates/aionui-db/tests/settings_repository.rs b/crates/aionui-db/tests/settings_repository.rs index 25ca4839d..75360ff1a 100644 --- a/crates/aionui-db/tests/settings_repository.rs +++ b/crates/aionui-db/tests/settings_repository.rs @@ -6,6 +6,8 @@ use std::sync::Arc; use aionui_db::{ISettingsRepository, SqliteSettingsRepository, init_database_memory}; +const USER_ID: &str = "system_default_user"; + async fn repo() -> Arc { let db = init_database_memory().await.unwrap(); Arc::new(SqliteSettingsRepository::new(db.pool().clone())) @@ -16,7 +18,7 @@ async fn repo() -> Arc { #[tokio::test] async fn get_settings_returns_none_when_no_row_exists() { let r = repo().await; - assert!(r.get_settings().await.unwrap().is_none()); + assert!(r.get_settings(USER_ID).await.unwrap().is_none()); } // -- Upsert creates a row -- @@ -24,7 +26,10 @@ async fn get_settings_returns_none_when_no_row_exists() { #[tokio::test] async fn upsert_creates_settings_with_given_values() { let r = repo().await; - let s = r.upsert_settings("zh-CN", false, true, true, false).await.unwrap(); + let s = r + .upsert_settings(USER_ID, "zh-CN", false, true, true, false) + .await + .unwrap(); assert_eq!(s.language, "zh-CN"); assert!(!s.notification_enabled); @@ -39,9 +44,11 @@ async fn upsert_creates_settings_with_given_values() { #[tokio::test] async fn upsert_then_get_returns_consistent_data() { let r = repo().await; - r.upsert_settings("en-US", true, false, false, true).await.unwrap(); + r.upsert_settings(USER_ID, "en-US", true, false, false, true) + .await + .unwrap(); - let s = r.get_settings().await.unwrap().unwrap(); + let s = r.get_settings(USER_ID).await.unwrap().unwrap(); assert_eq!(s.language, "en-US"); assert!(s.notification_enabled); assert!(!s.cron_notification_enabled); @@ -54,10 +61,14 @@ async fn upsert_then_get_returns_consistent_data() { #[tokio::test] async fn upsert_overwrites_previous_settings() { let r = repo().await; - r.upsert_settings("en-US", true, false, false, false).await.unwrap(); - r.upsert_settings("ja-JP", false, true, true, true).await.unwrap(); - - let s = r.get_settings().await.unwrap().unwrap(); + r.upsert_settings(USER_ID, "en-US", true, false, false, false) + .await + .unwrap(); + r.upsert_settings(USER_ID, "ja-JP", false, true, true, true) + .await + .unwrap(); + + let s = r.get_settings(USER_ID).await.unwrap().unwrap(); assert_eq!(s.language, "ja-JP"); assert!(!s.notification_enabled); assert!(s.cron_notification_enabled); @@ -70,8 +81,14 @@ async fn upsert_overwrites_previous_settings() { #[tokio::test] async fn upsert_advances_updated_at() { let r = repo().await; - let first = r.upsert_settings("en-US", true, false, false, false).await.unwrap(); - let second = r.upsert_settings("en-US", true, false, false, false).await.unwrap(); + let first = r + .upsert_settings(USER_ID, "en-US", true, false, false, false) + .await + .unwrap(); + let second = r + .upsert_settings(USER_ID, "en-US", true, false, false, false) + .await + .unwrap(); assert!(second.updated_at >= first.updated_at); } diff --git a/crates/aionui-db/tests/skill_management_schema.rs b/crates/aionui-db/tests/skill_management_schema.rs index 9e6e74e9f..74826949c 100644 --- a/crates/aionui-db/tests/skill_management_schema.rs +++ b/crates/aionui-db/tests/skill_management_schema.rs @@ -14,6 +14,7 @@ async fn migration_creates_skill_management_tables() { skill_columns, vec![ "id", + "user_id", "name", "description", "path", @@ -49,6 +50,7 @@ async fn migration_creates_skill_management_tables() { "line", "column", "created_at", + "user_id", ] ); } @@ -59,11 +61,17 @@ async fn migration_allows_known_skill_sources_and_rejects_unknown_sources() { let pool = db.pool(); for source in ["user", "builtin", "extension", "cron"] { + let user_id = if source == "builtin" || source == "cron" { + None + } else { + Some("system_default_user") + }; sqlx::query( - "INSERT INTO skills (id, name, description, path, source, enabled, created_at, updated_at) - VALUES (?, ?, NULL, ?, ?, 1, 1, 1)", + "INSERT INTO skills (id, user_id, name, description, path, source, enabled, created_at, updated_at) + VALUES (?, ?, ?, NULL, ?, ?, 1, 1, 1)", ) .bind(format!("skill-{source}")) + .bind(user_id) .bind(format!("skill-{source}")) .bind(format!("/tmp/{source}")) .bind(source) @@ -81,3 +89,41 @@ async fn migration_allows_known_skill_sources_and_rejects_unknown_sources() { assert!(rejected.is_err()); } + +#[tokio::test] +async fn skill_names_are_unique_within_scope() { + let db = init_database_memory().await.unwrap(); + let pool = db.pool(); + + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) + VALUES ('user_b', 'local', 'user_b', 'hash', 'active', 0, 1, 1)", + ) + .execute(pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO skills (id, user_id, name, description, path, source, enabled, created_at, updated_at) + VALUES ('skill-a', 'system_default_user', 'shared-name', NULL, '/tmp/a', 'user', 1, 1, 1)", + ) + .execute(pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO skills (id, user_id, name, description, path, source, enabled, created_at, updated_at) + VALUES ('skill-b', 'user_b', 'shared-name', NULL, '/tmp/b', 'user', 1, 1, 1)", + ) + .execute(pool) + .await + .unwrap(); + + let duplicate = sqlx::query( + "INSERT INTO skills (id, user_id, name, description, path, source, enabled, created_at, updated_at) + VALUES ('skill-c', 'user_b', 'shared-name', NULL, '/tmp/c', 'user', 1, 1, 1)", + ) + .execute(pool) + .await; + assert!(duplicate.is_err()); +} diff --git a/crates/aionui-db/tests/team_repository.rs b/crates/aionui-db/tests/team_repository.rs index f5cf8aa90..91bd14723 100644 --- a/crates/aionui-db/tests/team_repository.rs +++ b/crates/aionui-db/tests/team_repository.rs @@ -17,6 +17,8 @@ use aionui_db::{ DbError, ITeamRepository, SqliteTeamRepository, UpdateTaskParams, UpdateTeamParams, init_database_memory, }; +const DEFAULT_USER_ID: &str = "system_default_user"; + async fn repo() -> (Arc, aionui_db::Database) { let db = init_database_memory().await.unwrap(); let r = Arc::new(SqliteTeamRepository::new(db.pool().clone())); @@ -24,7 +26,7 @@ async fn repo() -> (Arc, aionui_db::Database) { } fn make_team(id: &str, name: &str) -> TeamRow { - make_team_for_user(id, "system_default_user", name) + make_team_for_user(id, DEFAULT_USER_ID, name) } fn make_team_for_user(id: &str, user_id: &str, name: &str) -> TeamRow { @@ -86,7 +88,11 @@ async fn create_and_get_team() { let team = make_team("t1", "Team Alpha"); repo.create_team(&team).await.unwrap(); - let fetched = repo.get_team("t1").await.unwrap().expect("team exists"); + let fetched = repo + .get_team("system_default_user", "t1") + .await + .unwrap() + .expect("team exists"); assert_eq!(fetched.id, "t1"); assert_eq!(fetched.name, "Team Alpha"); assert_eq!(fetched.lead_agent_id, Some("a1".into())); @@ -95,14 +101,14 @@ async fn create_and_get_team() { #[tokio::test] async fn get_nonexistent_team_returns_none() { let (repo, _db) = repo().await; - let result = repo.get_team("nonexistent").await.unwrap(); + let result = repo.get_team("system_default_user", "nonexistent").await.unwrap(); assert!(result.is_none()); } #[tokio::test] async fn list_teams_empty() { let (repo, _db) = repo().await; - let teams = repo.list_teams().await.unwrap(); + let teams = repo.list_teams_for_restore().await.unwrap(); assert!(teams.is_empty()); } @@ -112,7 +118,7 @@ async fn list_teams_multiple() { repo.create_team(&make_team("t1", "Alpha")).await.unwrap(); repo.create_team(&make_team("t2", "Beta")).await.unwrap(); - let teams = repo.list_teams().await.unwrap(); + let teams = repo.list_teams_for_restore().await.unwrap(); assert_eq!(teams.len(), 2); assert_eq!(teams[0].id, "t1"); assert_eq!(teams[1].id, "t2"); @@ -139,12 +145,41 @@ async fn list_teams_by_user_filters_to_owner() { assert!(teams.iter().all(|team| team.user_id == "user-a")); } +#[tokio::test] +async fn scoped_team_crud_rejects_wrong_user() { + let (repo, _db) = repo().await; + repo.create_team(&make_team_for_user("t1", "user-a", "Alpha")) + .await + .unwrap(); + + assert!(repo.get_team("user-b", "t1").await.unwrap().is_none()); + + let update = repo + .update_team( + "user-b", + "t1", + &UpdateTeamParams { + name: Some("Leaked".into()), + ..Default::default() + }, + ) + .await; + assert!(matches!(update, Err(DbError::NotFound(_)))); + + let delete = repo.delete_team("user-b", "t1").await; + assert!(matches!(delete, Err(DbError::NotFound(_)))); + + let team = repo.get_team("user-a", "t1").await.unwrap().unwrap(); + assert_eq!(team.name, "Alpha"); +} + #[tokio::test] async fn update_team_name() { let (repo, _db) = repo().await; repo.create_team(&make_team("t1", "Old Name")).await.unwrap(); repo.update_team( + "system_default_user", "t1", &UpdateTeamParams { name: Some("New Name".into()), @@ -154,7 +189,7 @@ async fn update_team_name() { .await .unwrap(); - let team = repo.get_team("t1").await.unwrap().unwrap(); + let team = repo.get_team("system_default_user", "t1").await.unwrap().unwrap(); assert_eq!(team.name, "New Name"); } @@ -165,6 +200,7 @@ async fn update_team_agents_json() { let new_agents = r#"[{"slotId":"a1"},{"slotId":"a2"}]"#; repo.update_team( + "system_default_user", "t1", &UpdateTeamParams { agents: Some(new_agents.into()), @@ -174,7 +210,7 @@ async fn update_team_agents_json() { .await .unwrap(); - let team = repo.get_team("t1").await.unwrap().unwrap(); + let team = repo.get_team("system_default_user", "t1").await.unwrap().unwrap(); assert_eq!(team.agents, new_agents); } @@ -185,6 +221,7 @@ async fn update_team_can_patch_workspace() { repo.create_team(&make_team("t1", "Team")).await.unwrap(); repo.update_team( + "system_default_user", "t1", &UpdateTeamParams { workspace: Some("/tmp/aionui-team-shared-workspace".into()), @@ -194,7 +231,7 @@ async fn update_team_can_patch_workspace() { .await .unwrap(); - let updated = repo.get_team("t1").await.unwrap().unwrap(); + let updated = repo.get_team("system_default_user", "t1").await.unwrap().unwrap(); assert_eq!(updated.workspace, "/tmp/aionui-team-shared-workspace"); } @@ -204,6 +241,7 @@ async fn update_team_can_patch_session_mode() { repo.create_team(&make_team("t1", "Team")).await.unwrap(); repo.update_team( + "system_default_user", "t1", &UpdateTeamParams { session_mode: Some("full_auto".into()), @@ -213,7 +251,7 @@ async fn update_team_can_patch_session_mode() { .await .unwrap(); - let updated = repo.get_team("t1").await.unwrap().unwrap(); + let updated = repo.get_team("system_default_user", "t1").await.unwrap().unwrap(); assert_eq!(updated.session_mode.as_deref(), Some("full_auto")); } @@ -222,6 +260,7 @@ async fn update_nonexistent_team_returns_not_found() { let (repo, _db) = repo().await; let result = repo .update_team( + "system_default_user", "nonexistent", &UpdateTeamParams { name: Some("X".into()), @@ -236,16 +275,16 @@ async fn update_nonexistent_team_returns_not_found() { async fn delete_team() { let (repo, _db) = repo().await; repo.create_team(&make_team("t1", "Team")).await.unwrap(); - repo.delete_team("t1").await.unwrap(); + repo.delete_team("system_default_user", "t1").await.unwrap(); - let result = repo.get_team("t1").await.unwrap(); + let result = repo.get_team("system_default_user", "t1").await.unwrap(); assert!(result.is_none()); } #[tokio::test] async fn delete_nonexistent_team_returns_not_found() { let (repo, _db) = repo().await; - let result = repo.delete_team("nonexistent").await; + let result = repo.delete_team("system_default_user", "nonexistent").await; assert!(matches!(result, Err(DbError::NotFound(_)))); } @@ -259,17 +298,17 @@ async fn write_and_read_unread_messages() { // Write 3 messages to agent a1 for i in 1..=3 { let msg = make_mailbox_msg(&format!("m{i}"), "t1", "a1", "a2", "message"); - repo.write_message(&msg).await.unwrap(); + repo.write_message(DEFAULT_USER_ID, &msg).await.unwrap(); } // Read unread: should return 3 - let unread = repo.read_unread_and_mark("t1", "a1").await.unwrap(); + let unread = repo.read_unread_and_mark(DEFAULT_USER_ID, "t1", "a1").await.unwrap(); assert_eq!(unread.len(), 3); assert!(!unread[0].read); // returned rows reflect pre-mark state assert_eq!(unread[0].msg_type, "message"); // Read again: should return 0 (all marked read) - let unread2 = repo.read_unread_and_mark("t1", "a1").await.unwrap(); + let unread2 = repo.read_unread_and_mark(DEFAULT_USER_ID, "t1", "a1").await.unwrap(); assert!(unread2.is_empty()); } @@ -278,7 +317,7 @@ async fn read_unread_no_messages() { let (repo, _db) = repo().await; repo.create_team(&make_team("t1", "Team")).await.unwrap(); - let unread = repo.read_unread_and_mark("t1", "a1").await.unwrap(); + let unread = repo.read_unread_and_mark(DEFAULT_USER_ID, "t1", "a1").await.unwrap(); assert!(unread.is_empty()); } @@ -289,9 +328,9 @@ async fn write_idle_notification_with_summary() { let mut msg = make_mailbox_msg("m1", "t1", "a1", "a2", "idle_notification"); msg.summary = Some("Task completed".into()); - repo.write_message(&msg).await.unwrap(); + repo.write_message(DEFAULT_USER_ID, &msg).await.unwrap(); - let history = repo.get_history("t1", "a1", None).await.unwrap(); + let history = repo.get_history(DEFAULT_USER_ID, "t1", "a1", None).await.unwrap(); assert_eq!(history.len(), 1); assert_eq!(history[0].msg_type, "idle_notification"); assert_eq!(history[0].summary.as_deref(), Some("Task completed")); @@ -303,9 +342,9 @@ async fn write_shutdown_request() { repo.create_team(&make_team("t1", "Team")).await.unwrap(); let msg = make_mailbox_msg("m1", "t1", "a1", "a2", "shutdown_request"); - repo.write_message(&msg).await.unwrap(); + repo.write_message(DEFAULT_USER_ID, &msg).await.unwrap(); - let history = repo.get_history("t1", "a1", None).await.unwrap(); + let history = repo.get_history(DEFAULT_USER_ID, "t1", "a1", None).await.unwrap(); assert_eq!(history.len(), 1); assert_eq!(history[0].msg_type, "shutdown_request"); } @@ -317,10 +356,10 @@ async fn get_history_with_limit() { for i in 1..=10 { let msg = make_mailbox_msg(&format!("m{i}"), "t1", "a1", "a2", "message"); - repo.write_message(&msg).await.unwrap(); + repo.write_message(DEFAULT_USER_ID, &msg).await.unwrap(); } - let history = repo.get_history("t1", "a1", Some(5)).await.unwrap(); + let history = repo.get_history(DEFAULT_USER_ID, "t1", "a1", Some(5)).await.unwrap(); assert_eq!(history.len(), 5); } @@ -331,10 +370,10 @@ async fn get_history_no_limit() { for i in 1..=3 { let msg = make_mailbox_msg(&format!("m{i}"), "t1", "a1", "a2", "message"); - repo.write_message(&msg).await.unwrap(); + repo.write_message(DEFAULT_USER_ID, &msg).await.unwrap(); } - let history = repo.get_history("t1", "a1", None).await.unwrap(); + let history = repo.get_history(DEFAULT_USER_ID, "t1", "a1", None).await.unwrap(); assert_eq!(history.len(), 3); } @@ -343,7 +382,7 @@ async fn get_history_empty() { let (repo, _db) = repo().await; repo.create_team(&make_team("t1", "Team")).await.unwrap(); - let history = repo.get_history("t1", "a1", None).await.unwrap(); + let history = repo.get_history(DEFAULT_USER_ID, "t1", "a1", None).await.unwrap(); assert!(history.is_empty()); } @@ -353,13 +392,13 @@ async fn get_history_includes_read_messages() { repo.create_team(&make_team("t1", "Team")).await.unwrap(); let msg = make_mailbox_msg("m1", "t1", "a1", "a2", "message"); - repo.write_message(&msg).await.unwrap(); + repo.write_message(DEFAULT_USER_ID, &msg).await.unwrap(); // Read and mark - repo.read_unread_and_mark("t1", "a1").await.unwrap(); + repo.read_unread_and_mark(DEFAULT_USER_ID, "t1", "a1").await.unwrap(); // History should still return the message - let history = repo.get_history("t1", "a1", None).await.unwrap(); + let history = repo.get_history(DEFAULT_USER_ID, "t1", "a1", None).await.unwrap(); assert_eq!(history.len(), 1); assert!(history[0].read); } @@ -373,21 +412,71 @@ async fn delete_mailbox_by_team() { // Write messages to both teams let msg1 = make_mailbox_msg("m1", "t1", "a1", "a2", "message"); let msg2 = make_mailbox_msg("m2", "t2", "a1", "a2", "message"); - repo.write_message(&msg1).await.unwrap(); - repo.write_message(&msg2).await.unwrap(); + repo.write_message(DEFAULT_USER_ID, &msg1).await.unwrap(); + repo.write_message(DEFAULT_USER_ID, &msg2).await.unwrap(); // Delete team1 mailbox - repo.delete_mailbox_by_team("t1").await.unwrap(); + repo.delete_mailbox_by_team("system_default_user", "t1").await.unwrap(); // Team1 mailbox empty - let h1 = repo.get_history("t1", "a1", None).await.unwrap(); + let h1 = repo.get_history(DEFAULT_USER_ID, "t1", "a1", None).await.unwrap(); assert!(h1.is_empty()); // Team2 mailbox intact - let h2 = repo.get_history("t2", "a1", None).await.unwrap(); + let h2 = repo.get_history(DEFAULT_USER_ID, "t2", "a1", None).await.unwrap(); assert_eq!(h2.len(), 1); } +#[tokio::test] +async fn scoped_mailbox_delete_and_mark_read_do_not_cross_team_owner() { + let (repo, _db) = repo().await; + repo.create_team(&make_team_for_user("t1", "user-a", "Team A")) + .await + .unwrap(); + repo.create_team(&make_team_for_user("t2", "user-b", "Team B")) + .await + .unwrap(); + repo.write_message("user-a", &make_mailbox_msg("m1", "t1", "a1", "a2", "message")) + .await + .unwrap(); + repo.write_message("user-b", &make_mailbox_msg("m2", "t2", "a1", "a2", "message")) + .await + .unwrap(); + + repo.mark_read_batch("user-a", "t1", &["m2".to_owned()]).await.unwrap(); + assert!(!repo.get_history("user-b", "t2", "a1", None).await.unwrap()[0].read); + + repo.delete_mailbox_by_team("user-b", "t1").await.unwrap(); + assert_eq!(repo.get_history("user-a", "t1", "a1", None).await.unwrap().len(), 1); +} + +#[tokio::test] +async fn mailbox_child_methods_do_not_cross_team_owner() { + let (repo, _db) = repo().await; + repo.create_team(&make_team_for_user("t1", "user-a", "Team A")) + .await + .unwrap(); + let msg = make_mailbox_msg("m1", "t1", "a1", "a2", "message"); + repo.write_message("user-a", &msg).await.unwrap(); + + assert!(repo.get_history("user-b", "t1", "a1", None).await.unwrap().is_empty()); + assert!(repo.peek_unread("user-b", "t1", "a1").await.unwrap().is_empty()); + assert!( + repo.read_unread_and_mark("user-b", "t1", "a1") + .await + .unwrap() + .is_empty() + ); + + repo.mark_read_batch("user-b", "t1", &["m1".to_owned()]).await.unwrap(); + assert!(!repo.get_history("user-a", "t1", "a1", None).await.unwrap()[0].read); + + let wrong_user_write = repo + .write_message("user-b", &make_mailbox_msg("m2", "t1", "a1", "a2", "message")) + .await; + assert!(matches!(wrong_user_write, Err(DbError::NotFound(_)))); +} + // ── Task Board Tests ───────────────────────────────────────────────── #[tokio::test] @@ -396,9 +485,9 @@ async fn create_and_list_tasks() { repo.create_team(&make_team("t1", "Team")).await.unwrap(); let task = make_task("tk1", "t1", "Implement feature"); - repo.create_task(&task).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task).await.unwrap(); - let tasks = repo.list_tasks("t1").await.unwrap(); + let tasks = repo.list_tasks(DEFAULT_USER_ID, "t1").await.unwrap(); assert_eq!(tasks.len(), 1); assert_eq!(tasks[0].subject, "Implement feature"); assert_eq!(tasks[0].status, "pending"); @@ -409,7 +498,7 @@ async fn list_tasks_empty() { let (repo, _db) = repo().await; repo.create_team(&make_team("t1", "Team")).await.unwrap(); - let tasks = repo.list_tasks("t1").await.unwrap(); + let tasks = repo.list_tasks(DEFAULT_USER_ID, "t1").await.unwrap(); assert!(tasks.is_empty()); } @@ -419,9 +508,9 @@ async fn find_task_by_id() { repo.create_team(&make_team("t1", "Team")).await.unwrap(); let task = make_task("tk1", "t1", "Task"); - repo.create_task(&task).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task).await.unwrap(); - let found = repo.find_task_by_id("t1", "tk1").await.unwrap(); + let found = repo.find_task_by_id(DEFAULT_USER_ID, "t1", "tk1").await.unwrap(); assert!(found.is_some()); assert_eq!(found.unwrap().id, "tk1"); } @@ -431,7 +520,10 @@ async fn find_task_by_id_not_found() { let (repo, _db) = repo().await; repo.create_team(&make_team("t1", "Team")).await.unwrap(); - let found = repo.find_task_by_id("t1", "nonexistent").await.unwrap(); + let found = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "nonexistent") + .await + .unwrap(); assert!(found.is_none()); } @@ -441,9 +533,11 @@ async fn update_task_status() { repo.create_team(&make_team("t1", "Team")).await.unwrap(); let task = make_task("tk1", "t1", "Task"); - repo.create_task(&task).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task).await.unwrap(); repo.update_task( + DEFAULT_USER_ID, + "t1", "tk1", &UpdateTaskParams { status: Some("in_progress".into()), @@ -453,7 +547,11 @@ async fn update_task_status() { .await .unwrap(); - let updated = repo.find_task_by_id("t1", "tk1").await.unwrap().unwrap(); + let updated = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "tk1") + .await + .unwrap() + .unwrap(); assert_eq!(updated.status, "in_progress"); } @@ -463,9 +561,11 @@ async fn update_task_description_and_owner() { repo.create_team(&make_team("t1", "Team")).await.unwrap(); let task = make_task("tk1", "t1", "Task"); - repo.create_task(&task).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task).await.unwrap(); repo.update_task( + DEFAULT_USER_ID, + "t1", "tk1", &UpdateTaskParams { description: Some("New description".into()), @@ -476,7 +576,11 @@ async fn update_task_description_and_owner() { .await .unwrap(); - let updated = repo.find_task_by_id("t1", "tk1").await.unwrap().unwrap(); + let updated = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "tk1") + .await + .unwrap() + .unwrap(); assert_eq!(updated.description.as_deref(), Some("New description")); assert_eq!(updated.owner.as_deref(), Some("agent-2")); } @@ -486,6 +590,8 @@ async fn update_nonexistent_task_returns_not_found() { let (repo, _db) = repo().await; let result = repo .update_task( + DEFAULT_USER_ID, + "t1", "nonexistent", &UpdateTaskParams { status: Some("completed".into()), @@ -505,20 +611,32 @@ async fn append_to_blocks_and_remove_from_blocked_by() { let task_a = make_task("tkA", "t1", "Task A"); let mut task_b = make_task("tkB", "t1", "Task B"); task_b.blocked_by = r#"["tkA"]"#.into(); - repo.create_task(&task_a).await.unwrap(); - repo.create_task(&task_b).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task_a).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task_b).await.unwrap(); // Append tkB to taskA's blocks - repo.append_to_blocks("tkA", "tkB").await.unwrap(); + repo.append_to_blocks(DEFAULT_USER_ID, "t1", "tkA", "tkB") + .await + .unwrap(); - let a = repo.find_task_by_id("t1", "tkA").await.unwrap().unwrap(); + let a = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "tkA") + .await + .unwrap() + .unwrap(); let blocks: Vec = serde_json::from_str(&a.blocks).unwrap(); assert!(blocks.contains(&"tkB".to_string())); // Now complete taskA: remove tkA from taskB's blocked_by - repo.remove_from_blocked_by("tkB", "tkA").await.unwrap(); + repo.remove_from_blocked_by(DEFAULT_USER_ID, "t1", "tkB", "tkA") + .await + .unwrap(); - let b = repo.find_task_by_id("t1", "tkB").await.unwrap().unwrap(); + let b = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "tkB") + .await + .unwrap() + .unwrap(); let blocked_by: Vec = serde_json::from_str(&b.blocked_by).unwrap(); assert!(!blocked_by.contains(&"tkA".to_string())); } @@ -529,12 +647,20 @@ async fn append_to_blocks_idempotent() { repo.create_team(&make_team("t1", "Team")).await.unwrap(); let task = make_task("tkA", "t1", "Task A"); - repo.create_task(&task).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task).await.unwrap(); - repo.append_to_blocks("tkA", "tkB").await.unwrap(); - repo.append_to_blocks("tkA", "tkB").await.unwrap(); + repo.append_to_blocks(DEFAULT_USER_ID, "t1", "tkA", "tkB") + .await + .unwrap(); + repo.append_to_blocks(DEFAULT_USER_ID, "t1", "tkA", "tkB") + .await + .unwrap(); - let a = repo.find_task_by_id("t1", "tkA").await.unwrap().unwrap(); + let a = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "tkA") + .await + .unwrap() + .unwrap(); let blocks: Vec = serde_json::from_str(&a.blocks).unwrap(); assert_eq!(blocks.len(), 1); // no duplicates } @@ -551,19 +677,35 @@ async fn multi_dependency_unblock() { let mut task_c = make_task("tkC", "t1", "C"); task_c.blocked_by = r#"["tkA"]"#.into(); - repo.create_task(&task_a).await.unwrap(); - repo.create_task(&task_b).await.unwrap(); - repo.create_task(&task_c).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task_a).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task_b).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task_c).await.unwrap(); - repo.append_to_blocks("tkA", "tkB").await.unwrap(); - repo.append_to_blocks("tkA", "tkC").await.unwrap(); + repo.append_to_blocks(DEFAULT_USER_ID, "t1", "tkA", "tkB") + .await + .unwrap(); + repo.append_to_blocks(DEFAULT_USER_ID, "t1", "tkA", "tkC") + .await + .unwrap(); // Complete A: unblock both B and C - repo.remove_from_blocked_by("tkB", "tkA").await.unwrap(); - repo.remove_from_blocked_by("tkC", "tkA").await.unwrap(); + repo.remove_from_blocked_by(DEFAULT_USER_ID, "t1", "tkB", "tkA") + .await + .unwrap(); + repo.remove_from_blocked_by(DEFAULT_USER_ID, "t1", "tkC", "tkA") + .await + .unwrap(); - let b = repo.find_task_by_id("t1", "tkB").await.unwrap().unwrap(); - let c = repo.find_task_by_id("t1", "tkC").await.unwrap().unwrap(); + let b = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "tkB") + .await + .unwrap() + .unwrap(); + let c = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "tkC") + .await + .unwrap() + .unwrap(); let b_blocked: Vec = serde_json::from_str(&b.blocked_by).unwrap(); let c_blocked: Vec = serde_json::from_str(&c.blocked_by).unwrap(); assert!(b_blocked.is_empty()); @@ -580,13 +722,19 @@ async fn partial_unblock_preserves_other_blockers() { let mut task_b = make_task("tkB", "t1", "B"); task_b.blocked_by = r#"["tkA","tkX"]"#.into(); - repo.create_task(&task_a).await.unwrap(); - repo.create_task(&task_b).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task_a).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task_b).await.unwrap(); // Complete A only - repo.remove_from_blocked_by("tkB", "tkA").await.unwrap(); + repo.remove_from_blocked_by(DEFAULT_USER_ID, "t1", "tkB", "tkA") + .await + .unwrap(); - let b = repo.find_task_by_id("t1", "tkB").await.unwrap().unwrap(); + let b = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "tkB") + .await + .unwrap() + .unwrap(); let blocked_by: Vec = serde_json::from_str(&b.blocked_by).unwrap(); assert_eq!(blocked_by, vec!["tkX"]); } @@ -597,10 +745,12 @@ async fn no_blocks_task_completes_cleanly() { repo.create_team(&make_team("t1", "Team")).await.unwrap(); let task = make_task("tkA", "t1", "A"); - repo.create_task(&task).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task).await.unwrap(); // Complete without any blocks to unblock repo.update_task( + DEFAULT_USER_ID, + "t1", "tkA", &UpdateTaskParams { status: Some("completed".into()), @@ -610,7 +760,11 @@ async fn no_blocks_task_completes_cleanly() { .await .unwrap(); - let a = repo.find_task_by_id("t1", "tkA").await.unwrap().unwrap(); + let a = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "tkA") + .await + .unwrap() + .unwrap(); assert_eq!(a.status, "completed"); let blocks: Vec = serde_json::from_str(&a.blocks).unwrap(); assert!(blocks.is_empty()); @@ -622,18 +776,97 @@ async fn delete_tasks_by_team() { repo.create_team(&make_team("t1", "Team1")).await.unwrap(); repo.create_team(&make_team("t2", "Team2")).await.unwrap(); - repo.create_task(&make_task("tk1", "t1", "T1 Task")).await.unwrap(); - repo.create_task(&make_task("tk2", "t2", "T2 Task")).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &make_task("tk1", "t1", "T1 Task")) + .await + .unwrap(); + repo.create_task(DEFAULT_USER_ID, &make_task("tk2", "t2", "T2 Task")) + .await + .unwrap(); - repo.delete_tasks_by_team("t1").await.unwrap(); + repo.delete_tasks_by_team("system_default_user", "t1").await.unwrap(); - let t1_tasks = repo.list_tasks("t1").await.unwrap(); + let t1_tasks = repo.list_tasks(DEFAULT_USER_ID, "t1").await.unwrap(); assert!(t1_tasks.is_empty()); - let t2_tasks = repo.list_tasks("t2").await.unwrap(); + let t2_tasks = repo.list_tasks(DEFAULT_USER_ID, "t2").await.unwrap(); assert_eq!(t2_tasks.len(), 1); } +#[tokio::test] +async fn scoped_task_updates_and_deletes_stay_within_team_owner() { + let (repo, _db) = repo().await; + repo.create_team(&make_team_for_user("t1", "user-a", "Team A")) + .await + .unwrap(); + repo.create_team(&make_team_for_user("t2", "user-b", "Team B")) + .await + .unwrap(); + repo.create_task("user-a", &make_task("tk1", "t1", "A Task")) + .await + .unwrap(); + repo.create_task("user-b", &make_task("tk2", "t2", "B Task")) + .await + .unwrap(); + + let update = repo + .update_task( + "user-b", + "t1", + "tk1", + &UpdateTaskParams { + status: Some("completed".into()), + ..Default::default() + }, + ) + .await; + assert!(matches!(update, Err(DbError::NotFound(_)))); + + repo.delete_tasks_by_team("user-b", "t1").await.unwrap(); + assert_eq!(repo.list_tasks("user-a", "t1").await.unwrap().len(), 1); + assert_eq!(repo.list_tasks("user-b", "t2").await.unwrap().len(), 1); +} + +#[tokio::test] +async fn task_child_methods_do_not_cross_team_owner() { + let (repo, _db) = repo().await; + repo.create_team(&make_team_for_user("t1", "user-a", "Team A")) + .await + .unwrap(); + repo.create_task("user-a", &make_task("tk1", "t1", "A Task")) + .await + .unwrap(); + repo.create_task("user-a", &make_task("tk2", "t1", "B Task")) + .await + .unwrap(); + + assert!(repo.list_tasks("user-b", "t1").await.unwrap().is_empty()); + assert!(repo.find_task_by_id("user-b", "t1", "tk1").await.unwrap().is_none()); + + let wrong_update = repo + .update_task( + "user-b", + "t1", + "tk1", + &UpdateTaskParams { + status: Some("completed".into()), + ..Default::default() + }, + ) + .await; + assert!(matches!(wrong_update, Err(DbError::NotFound(_)))); + + let wrong_append = repo.append_to_blocks("user-b", "t1", "tk1", "tk2").await; + assert!(matches!(wrong_append, Err(DbError::NotFound(_)))); + + let wrong_remove = repo.remove_from_blocked_by("user-b", "t1", "tk2", "tk1").await; + assert!(matches!(wrong_remove, Err(DbError::NotFound(_)))); + + let wrong_create = repo + .create_task("user-b", &make_task("tk3", "t1", "Wrong User Task")) + .await; + assert!(matches!(wrong_create, Err(DbError::NotFound(_)))); +} + #[tokio::test] async fn tasks_contain_dependency_info() { let (repo, _db) = repo().await; @@ -643,11 +876,13 @@ async fn tasks_contain_dependency_info() { let mut task_b = make_task("tkB", "t1", "B"); task_b.blocked_by = r#"["tkA"]"#.into(); - repo.create_task(&task_a).await.unwrap(); - repo.create_task(&task_b).await.unwrap(); - repo.append_to_blocks("tkA", "tkB").await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task_a).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task_b).await.unwrap(); + repo.append_to_blocks(DEFAULT_USER_ID, "t1", "tkA", "tkB") + .await + .unwrap(); - let tasks = repo.list_tasks("t1").await.unwrap(); + let tasks = repo.list_tasks(DEFAULT_USER_ID, "t1").await.unwrap(); assert_eq!(tasks.len(), 2); let a = tasks.iter().find(|t| t.id == "tkA").unwrap(); @@ -669,21 +904,21 @@ async fn delete_team_cascades_mailbox_and_tasks() { // Add mailbox messages and tasks let msg = make_mailbox_msg("m1", "t1", "a1", "a2", "message"); - repo.write_message(&msg).await.unwrap(); + repo.write_message(DEFAULT_USER_ID, &msg).await.unwrap(); let task = make_task("tk1", "t1", "Task"); - repo.create_task(&task).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task).await.unwrap(); // Delete team, then manually clean up related data (as service layer would) - repo.delete_mailbox_by_team("t1").await.unwrap(); - repo.delete_tasks_by_team("t1").await.unwrap(); - repo.delete_team("t1").await.unwrap(); + repo.delete_mailbox_by_team("system_default_user", "t1").await.unwrap(); + repo.delete_tasks_by_team("system_default_user", "t1").await.unwrap(); + repo.delete_team("system_default_user", "t1").await.unwrap(); // Verify all cleaned up - let team = repo.get_team("t1").await.unwrap(); + let team = repo.get_team("system_default_user", "t1").await.unwrap(); assert!(team.is_none()); - let mail = repo.get_history("t1", "a1", None).await.unwrap(); + let mail = repo.get_history(DEFAULT_USER_ID, "t1", "a1", None).await.unwrap(); assert!(mail.is_empty()); - let tasks = repo.list_tasks("t1").await.unwrap(); + let tasks = repo.list_tasks(DEFAULT_USER_ID, "t1").await.unwrap(); assert!(tasks.is_empty()); } @@ -696,13 +931,23 @@ async fn task_blocked_by_blocks_bidirectional_consistency() { let mut task_b = make_task("tkB", "t1", "B"); task_b.blocked_by = r#"["tkA"]"#.into(); - repo.create_task(&task_a).await.unwrap(); - repo.create_task(&task_b).await.unwrap(); - repo.append_to_blocks("tkA", "tkB").await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task_a).await.unwrap(); + repo.create_task(DEFAULT_USER_ID, &task_b).await.unwrap(); + repo.append_to_blocks(DEFAULT_USER_ID, "t1", "tkA", "tkB") + .await + .unwrap(); // Verify bidirectional link - let a = repo.find_task_by_id("t1", "tkA").await.unwrap().unwrap(); - let b = repo.find_task_by_id("t1", "tkB").await.unwrap().unwrap(); + let a = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "tkA") + .await + .unwrap() + .unwrap(); + let b = repo + .find_task_by_id(DEFAULT_USER_ID, "t1", "tkB") + .await + .unwrap() + .unwrap(); let a_blocks: Vec = serde_json::from_str(&a.blocks).unwrap(); let b_blocked_by: Vec = serde_json::from_str(&b.blocked_by).unwrap(); diff --git a/crates/aionui-db/tests/user_repository.rs b/crates/aionui-db/tests/user_repository.rs index 43aecc390..351748679 100644 --- a/crates/aionui-db/tests/user_repository.rs +++ b/crates/aionui-db/tests/user_repository.rs @@ -5,7 +5,9 @@ use std::sync::Arc; -use aionui_db::{DbError, IUserRepository, SqliteUserRepository, init_database_memory}; +use aionui_db::{ + DbError, ExternalUserProjection, IUserRepository, SqliteUserRepository, UserStatus, UserType, init_database_memory, +}; async fn repo() -> Arc { let db = init_database_memory().await.unwrap(); @@ -20,8 +22,11 @@ async fn t2_1_create_user_returns_user_with_populated_fields() { let user = r.create_user("testuser", "$2b$12$fakehash").await.unwrap(); assert!(!user.id.is_empty(), "id should be non-empty"); - assert_eq!(user.username, "testuser"); - assert_eq!(user.password_hash, "$2b$12$fakehash"); + assert_eq!(user.user_type, UserType::Local); + assert_eq!(user.status, UserStatus::Active); + assert_eq!(user.session_generation, 0); + assert_eq!(user.username.as_deref(), Some("testuser")); + assert_eq!(user.password_hash.as_deref(), Some("$2b$12$fakehash")); assert!(user.created_at > 0); assert!(user.updated_at > 0); } @@ -44,7 +49,7 @@ async fn t2_2_find_by_username_existing() { let found = r.find_by_username("findme").await.unwrap(); assert!(found.is_some()); - assert_eq!(found.unwrap().username, "findme"); + assert_eq!(found.unwrap().username.as_deref(), Some("findme")); } #[tokio::test] @@ -153,8 +158,8 @@ async fn t2_9_set_system_user_credentials_updates_username_and_hash() { r.set_system_user_credentials("newadmin", "secure_hash").await.unwrap(); let user = r.get_system_user().await.unwrap().unwrap(); - assert_eq!(user.username, "newadmin"); - assert_eq!(user.password_hash, "secure_hash"); + assert_eq!(user.username.as_deref(), Some("newadmin")); + assert_eq!(user.password_hash.as_deref(), Some("secure_hash")); } #[tokio::test] @@ -176,7 +181,7 @@ async fn t2_10_update_password_changes_hash_and_updated_at() { r.update_password(&user.id, "new_hash").await.unwrap(); let updated = r.find_by_id(&user.id).await.unwrap().unwrap(); - assert_eq!(updated.password_hash, "new_hash"); + assert_eq!(updated.password_hash.as_deref(), Some("new_hash")); assert!(updated.updated_at >= user.updated_at); } @@ -190,7 +195,7 @@ async fn t2_11_update_username_succeeds() { r.update_username(&user.id, "newname").await.unwrap(); let updated = r.find_by_id(&user.id).await.unwrap().unwrap(); - assert_eq!(updated.username, "newname"); + assert_eq!(updated.username.as_deref(), Some("newname")); } #[tokio::test] @@ -231,3 +236,67 @@ async fn t2_13_update_jwt_secret_sets_value() { let updated = r.find_by_id(&user.id).await.unwrap().unwrap(); assert_eq!(updated.jwt_secret.as_deref(), Some("my_secret")); } + +#[tokio::test] +async fn t2_14_external_user_provision_is_idempotent() { + let r = repo().await; + + let first = r + .ensure_external_user( + UserType::Aionpro, + "pro-user-1", + ExternalUserProjection { + username: Some("Pro User".to_string()), + email: Some("pro@example.com".to_string()), + avatar_path: None, + }, + ) + .await + .unwrap(); + let second = r + .ensure_external_user(UserType::Aionpro, "pro-user-1", ExternalUserProjection::default()) + .await + .unwrap(); + + assert_eq!(first.id, second.id); + assert_eq!(first.user_type, UserType::Aionpro); + assert_eq!(first.external_user_id.as_deref(), Some("pro-user-1")); + assert!(first.password_hash.is_none()); + assert!(r.find_by_username("Pro User").await.unwrap().is_none()); +} + +#[tokio::test] +async fn t2_15_disabled_user_is_not_active() { + let r = repo().await; + let user = r + .ensure_external_user( + UserType::Aionpro, + "pro-user-disabled", + ExternalUserProjection::default(), + ) + .await + .unwrap(); + + assert!(r.find_active_by_id(&user.id).await.unwrap().is_some()); + r.set_status(&user.id, UserStatus::Disabled).await.unwrap(); + + let disabled = r.find_by_id(&user.id).await.unwrap().unwrap(); + assert_eq!(disabled.session_generation, 1); + assert!(r.find_active_by_id(&user.id).await.unwrap().is_none()); + + r.set_status(&user.id, UserStatus::Disabled).await.unwrap(); + let disabled_again = r.find_by_id(&user.id).await.unwrap().unwrap(); + assert_eq!(disabled_again.session_generation, 1); +} + +#[tokio::test] +async fn t2_16_increment_session_generation_revokes_old_sessions() { + let r = repo().await; + let user = r.create_user("generation-user", "h").await.unwrap(); + + assert_eq!(r.increment_session_generation(&user.id).await.unwrap(), 1); + assert_eq!(r.increment_session_generation(&user.id).await.unwrap(), 2); + + let updated = r.find_by_id(&user.id).await.unwrap().unwrap(); + assert_eq!(updated.session_generation, 2); +} diff --git a/crates/aionui-db/tests/user_scope_migration.rs b/crates/aionui-db/tests/user_scope_migration.rs new file mode 100644 index 000000000..fc1eedd4d --- /dev/null +++ b/crates/aionui-db/tests/user_scope_migration.rs @@ -0,0 +1,781 @@ +use std::borrow::Cow; +use std::path::Path; + +use aionui_db::init_database_memory; +use sqlx::Row; +use sqlx::migrate::Migrator; +use sqlx::sqlite::SqlitePoolOptions; + +async fn table_columns(pool: &sqlx::SqlitePool, table: &str) -> Vec { + let rows = sqlx::query(&format!("PRAGMA table_info({table})")) + .fetch_all(pool) + .await + .unwrap(); + rows.into_iter().map(|row| row.get::("name")).collect() +} + +async fn table_indexes(pool: &sqlx::SqlitePool, table: &str) -> Vec { + let rows = sqlx::query(&format!("PRAGMA index_list({table})")) + .fetch_all(pool) + .await + .unwrap(); + rows.into_iter().map(|row| row.get::("name")).collect() +} + +async fn run_migrations_through(pool: &sqlx::SqlitePool, max_version: i64) { + sqlx::query("PRAGMA foreign_keys = OFF").execute(pool).await.unwrap(); + let full = Migrator::new(Path::new("migrations")).await.unwrap(); + let migrations = full + .migrations + .iter() + .filter(|migration| migration.version <= max_version) + .cloned() + .collect::>(); + let migrator = Migrator { + migrations: Cow::Owned(migrations), + ignore_missing: false, + locking: true, + no_tx: false, + }; + migrator.run(pool).await.unwrap(); +} + +async fn run_migration(pool: &sqlx::SqlitePool, version: i64) { + run_migration_result(pool, version).await.unwrap(); +} + +async fn run_migration_result(pool: &sqlx::SqlitePool, version: i64) -> Result<(), sqlx::migrate::MigrateError> { + let full = Migrator::new(Path::new("migrations")).await.unwrap(); + let migrations = full + .migrations + .iter() + .filter(|migration| migration.version == version) + .cloned() + .collect::>(); + let migrator = Migrator { + migrations: Cow::Owned(migrations), + ignore_missing: true, + locking: true, + no_tx: false, + }; + migrator.run(pool).await +} + +#[tokio::test] +async fn migration_029_adds_core_user_projection_columns() { + let db = init_database_memory().await.unwrap(); + let columns = table_columns(db.pool(), "users").await; + for column in ["user_type", "external_user_id", "status", "session_generation"] { + assert!( + columns.iter().any(|existing| existing == column), + "missing users.{column}" + ); + } + + let indexes = table_indexes(db.pool(), "users").await; + assert!( + indexes.iter().any(|index| index == "idx_users_external_user"), + "missing external user lookup index" + ); +} + +#[tokio::test] +async fn migration_029_adds_user_scope_to_independent_roots() { + let db = init_database_memory().await.unwrap(); + for (table, column) in [ + ("cron_jobs", "user_id"), + ("providers", "user_id"), + ("remote_agents", "user_id"), + ("mcp_servers", "user_id"), + ("oauth_tokens", "user_id"), + ("system_settings", "user_id"), + ("client_preferences", "user_id"), + ("assistant_plugins", "owner_user_id"), + ("assistant_users", "owner_user_id"), + ("assistant_pairing_codes", "owner_user_id"), + ] { + let columns = table_columns(db.pool(), table).await; + assert!( + columns.iter().any(|existing| existing == column), + "missing {table}.{column}" + ); + } +} + +#[tokio::test] +async fn migration_029_migrates_cron_skills_to_default_user() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 28).await; + + sqlx::query( + "INSERT INTO skills (id, name, description, path, source, enabled, created_at, updated_at) + VALUES ('legacy-cron-skill', 'scheduled-task', NULL, '/tmp/scheduled-task', 'cron', 1, 1, 1)", + ) + .execute(&pool) + .await + .unwrap(); + + run_migration(&pool, 29).await; + + let row = sqlx::query("SELECT user_id, source FROM skills WHERE id = 'legacy-cron-skill'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(row.get::("source"), "cron"); + // Cron-derived skills carry user content (job prompts); legacy rows are + // owned by the pre-migration single user, never exposed as global rows. + assert_eq!( + row.get::, _>("user_id").as_deref(), + Some("system_default_user") + ); +} + +#[tokio::test] +async fn migration_029_backfills_existing_independent_roots_to_default_user() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 28).await; + let now = 1_000_i64; + + sqlx::query( + "INSERT INTO providers ( + id, platform, name, base_url, api_key_encrypted, models, enabled, + capabilities, model_settings, created_at, updated_at + ) VALUES ( + 'provider_legacy', 'openai', 'OpenAI', 'https://api.example.com', + 'encrypted', '[]', 1, '[]', '{}', ?, ? + )", + ) + .bind(now) + .bind(now) + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO remote_agents ( + id, name, protocol, url, auth_type, status, created_at, updated_at + ) VALUES ( + 'remote_legacy', 'Remote', 'http', 'http://127.0.0.1:1', 'none', 'unknown', ?, ? + )", + ) + .bind(now) + .bind(now) + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO mcp_servers ( + id, name, enabled, transport_type, transport_config, last_test_status, + builtin, created_at, updated_at + ) VALUES ( + 'mcp_legacy', 'legacy-mcp', 1, 'stdio', '{}', 'disconnected', 0, ?, ? + )", + ) + .bind(now) + .bind(now) + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO oauth_tokens ( + server_url, access_token, token_type, created_at, updated_at + ) VALUES ( + 'https://mcp.example.com', 'encrypted-token', 'bearer', ?, ? + )", + ) + .bind(now) + .bind(now) + .execute(&pool) + .await + .unwrap(); + + sqlx::query("INSERT INTO client_preferences (key, value, updated_at) VALUES ('theme', 'dark', ?)") + .bind(now) + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO cron_jobs ( + id, name, enabled, schedule_kind, schedule_value, payload_message, + execution_mode, conversation_id, created_by, created_at, + updated_at, run_count, retry_count, max_retries, queue_enabled + ) VALUES ( + 'cron_legacy', 'Legacy Cron', 1, 'every', '60000', 'message', + 'new_conversation', '', 'user', ?, ?, 0, 0, 3, 0 + )", + ) + .bind(now) + .bind(now) + .execute(&pool) + .await + .unwrap(); + + run_migration(&pool, 29).await; + + for (table, id_column, id_value, scope_column) in [ + ("providers", "id", "provider_legacy", "user_id"), + ("remote_agents", "id", "remote_legacy", "user_id"), + ("mcp_servers", "id", "mcp_legacy", "user_id"), + ("client_preferences", "key", "theme", "user_id"), + ("cron_jobs", "id", "cron_legacy", "user_id"), + ] { + let sql = format!("SELECT {scope_column} FROM {table} WHERE {id_column} = ?"); + let owner: String = sqlx::query_scalar(&sql).bind(id_value).fetch_one(&pool).await.unwrap(); + assert_eq!( + owner, "system_default_user", + "{table}.{scope_column} was not backfilled" + ); + } + + let oauth_owner: String = + sqlx::query_scalar("SELECT user_id FROM oauth_tokens WHERE server_url = 'https://mcp.example.com'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(oauth_owner, "system_default_user"); +} + +#[tokio::test] +async fn migration_029_classifies_global_and_user_catalog_rows() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 28).await; + let now = 2_000_i64; + + sqlx::query( + "INSERT INTO agent_metadata ( + id, name, agent_type, agent_source, enabled, created_at, updated_at + ) VALUES + ('agent_builtin_legacy', 'Builtin Agent', 'acp', 'builtin', 1, ?, ?), + ('agent_user_legacy', 'User Agent', 'acp', 'custom', 1, ?, ?)", + ) + .bind(now) + .bind(now) + .bind(now) + .bind(now) + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO assistant_definitions ( + id, assistant_id, source, owner_type, source_ref, name, name_i18n, + description_i18n, avatar_type, agent_id, rule_resource_type, + recommended_prompts, recommended_prompts_i18n, default_model_mode, + default_permission_mode, default_skills_mode, default_skill_ids, + custom_skill_names, default_disabled_builtin_skill_ids, + default_mcps_mode, default_mcp_ids, created_at, updated_at + ) VALUES + ( + 'def_builtin_legacy', 'assistant_builtin', 'builtin', 'system', + 'builtin-ref', 'Builtin Assistant', '{}', '{}', 'none', + 'agent_builtin_legacy', 'none', '[]', '{}', 'auto', 'auto', + 'auto', '[]', '[]', '[]', 'auto', '[]', ?, ? + ), + ( + 'def_user_legacy', 'assistant_user', 'user', 'user', + 'user-ref', 'User Assistant', '{}', '{}', 'none', + 'agent_user_legacy', 'user_file', '[]', '{}', 'auto', 'auto', + 'auto', '[]', '[]', '[]', 'auto', '[]', ?, ? + )", + ) + .bind(now) + .bind(now) + .bind(now) + .bind(now) + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO skills (id, name, description, path, source, enabled, created_at, updated_at) + VALUES + ('skill_builtin_legacy', 'builtin-skill', NULL, '/tmp/builtin', 'builtin', 1, ?, ?), + ('skill_extension_legacy', 'extension-skill', NULL, '/tmp/extension', 'extension', 1, ?, ?)", + ) + .bind(now) + .bind(now) + .bind(now) + .bind(now) + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO skill_import_records ( + id, operation_id, source_label, source_name, status, error_code, created_at + ) VALUES ( + 'import_failed_legacy', 'op_legacy', 'archive.zip', 'broken-skill', + 'failed', 'INVALID_MANIFEST', ? + )", + ) + .bind(now) + .execute(&pool) + .await + .unwrap(); + + run_migration(&pool, 29).await; + + for (table, id_column, id_value, expected_owner) in [ + ("agent_metadata", "id", "agent_builtin_legacy", None), + ("agent_metadata", "id", "agent_user_legacy", Some("system_default_user")), + ("assistant_definitions", "id", "def_builtin_legacy", None), + ( + "assistant_definitions", + "id", + "def_user_legacy", + Some("system_default_user"), + ), + ("skills", "id", "skill_builtin_legacy", None), + ("skills", "id", "skill_extension_legacy", Some("system_default_user")), + ] { + let sql = format!("SELECT user_id FROM {table} WHERE {id_column} = ?"); + let owner: Option = sqlx::query_scalar(&sql).bind(id_value).fetch_one(&pool).await.unwrap(); + assert_eq!( + owner.as_deref(), + expected_owner, + "{table}.{id_column}={id_value} had wrong user_id" + ); + } + + let import_owner: String = + sqlx::query_scalar("SELECT user_id FROM skill_import_records WHERE id = 'import_failed_legacy'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(import_owner, "system_default_user"); +} + +#[tokio::test] +async fn migration_029_keeps_new_conversation_cron_jobs_unanchored_until_run() { + let db = init_database_memory().await.unwrap(); + let now = aionui_common::now_ms(); + + sqlx::query( + "INSERT INTO cron_jobs (\ + id, user_id, name, enabled, schedule_kind, schedule_value, \ + schedule_description, payload_message, execution_mode, agent_config, \ + conversation_id, conversation_title, created_by, created_at, updated_at, \ + next_run_at, run_count, retry_count, max_retries, queue_enabled\ + ) VALUES (\ + 'cron_unanchored', 'system_default_user', 'Unanchored', 1, 'every', '60000', \ + NULL, 'message', 'new_conversation', NULL, '', NULL, 'user', ?, ?, \ + NULL, 0, 0, 3, 0\ + )", + ) + .bind(now) + .bind(now) + .execute(db.pool()) + .await + .unwrap(); + + let row = sqlx::query("SELECT user_id, conversation_id FROM cron_jobs WHERE id = 'cron_unanchored'") + .fetch_one(db.pool()) + .await + .unwrap(); + assert_eq!(row.get::("user_id"), "system_default_user"); + assert_eq!(row.get::("conversation_id"), ""); +} + +async fn insert_valid_aggregate_parents(pool: &sqlx::SqlitePool) { + sqlx::query( + "INSERT INTO conversations (id, user_id, name, type, extra, status, created_at, updated_at) + VALUES ('conv_parent', 'system_default_user', 'Parent Conversation', 'acp', '{}', 'pending', 1, 1)", + ) + .execute(pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO teams (id, user_id, name, workspace, agents, created_at, updated_at) + VALUES ('team_parent', 'system_default_user', 'Parent Team', '', '[]', 1, 1)", + ) + .execute(pool) + .await + .unwrap(); +} + +#[tokio::test] +async fn migration_029_preserves_valid_aggregate_child_rows() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 28).await; + insert_valid_aggregate_parents(&pool).await; + + sqlx::query( + "INSERT INTO messages (id, conversation_id, type, content, created_at) + VALUES ('msg_parent', 'conv_parent', 'user', '{}', 1)", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO conversation_artifacts ( + id, conversation_id, cron_job_id, kind, status, payload, created_at, updated_at + ) VALUES ( + 'artifact_parent', 'conv_parent', 'cron_parent', 'skill_suggest', 'active', '{}', 1, 1 + )", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO conversation_assistant_snapshots ( + conversation_id, assistant_definition_id, assistant_id, assistant_source, + agent_id, rules_content, default_model_mode, default_permission_mode, + default_thought_level_mode, default_skills_mode, resolved_skill_ids, + resolved_disabled_builtin_skill_ids, default_mcps_mode, resolved_mcp_ids, + created_at, updated_at + ) VALUES ( + 'conv_parent', 'definition_parent', 'assistant_parent', 'user', + 'agent_parent', '', 'auto', 'auto', 'auto', 'auto', '[]', '[]', + 'auto', '[]', 1, 1 + )", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO acp_session ( + conversation_id, agent_source, agent_id, session_status, session_config + ) VALUES ( + 'conv_parent', 'builtin', 'agent_parent', 'idle', '{}' + )", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO cron_jobs ( + id, name, enabled, schedule_kind, schedule_value, payload_message, + execution_mode, conversation_id, created_by, created_at, + updated_at, run_count, retry_count, max_retries, queue_enabled + ) VALUES ( + 'cron_parent', 'Parent Cron', 1, 'every', '60000', 'message', + 'existing', 'conv_parent', 'user', 1, 1, 0, 0, 3, 0 + )", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO mailbox (id, team_id, to_agent_id, from_agent_id, type, content, created_at) + VALUES ('mail_parent', 'team_parent', 'a1', 'a2', 'message', 'hello', 1)", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO team_tasks (id, team_id, subject, status, created_at, updated_at) + VALUES ('task_parent', 'team_parent', 'Task', 'pending', 1, 1)", + ) + .execute(&pool) + .await + .unwrap(); + + run_migration(&pool, 29).await; + + for (table, id_column, id_value) in [ + ("messages", "id", "msg_parent"), + ("conversation_artifacts", "id", "artifact_parent"), + ("conversation_assistant_snapshots", "conversation_id", "conv_parent"), + ("acp_session", "conversation_id", "conv_parent"), + ("cron_jobs", "id", "cron_parent"), + ("mailbox", "id", "mail_parent"), + ("team_tasks", "id", "task_parent"), + ] { + let sql = format!("SELECT COUNT(*) FROM {table} WHERE {id_column} = ?"); + let count: i64 = sqlx::query_scalar(&sql).bind(id_value).fetch_one(&pool).await.unwrap(); + assert_eq!(count, 1, "{table} row was not preserved"); + } + + let cron_owner: String = sqlx::query_scalar("SELECT user_id FROM cron_jobs WHERE id = 'cron_parent'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(cron_owner, "system_default_user"); +} + +#[tokio::test] +async fn migration_029_rejects_aggregate_child_orphans() { + for (name, setup_sql) in [ + ( + "messages", + "INSERT INTO messages (id, conversation_id, type, content, created_at) + VALUES ('orphan_msg', 'missing_conv', 'user', '{}', 1)", + ), + ( + "conversation_artifacts", + "INSERT INTO conversation_artifacts ( + id, conversation_id, cron_job_id, kind, status, payload, created_at, updated_at + ) VALUES ( + 'orphan_artifact', 'missing_conv', NULL, 'skill_suggest', 'active', '{}', 1, 1 + )", + ), + ( + "conversation_assistant_snapshots", + "INSERT INTO conversation_assistant_snapshots ( + conversation_id, assistant_definition_id, assistant_id, assistant_source, + agent_id, rules_content, default_model_mode, default_permission_mode, + default_thought_level_mode, default_skills_mode, resolved_skill_ids, + resolved_disabled_builtin_skill_ids, default_mcps_mode, resolved_mcp_ids, + created_at, updated_at + ) VALUES ( + 'missing_conv', 'definition_parent', 'assistant_parent', 'user', + 'agent_parent', '', 'auto', 'auto', 'auto', 'auto', '[]', '[]', + 'auto', '[]', 1, 1 + )", + ), + ( + "acp_session", + "INSERT INTO acp_session ( + conversation_id, agent_source, agent_id, session_status, session_config + ) VALUES ( + 'missing_conv', 'builtin', 'agent_parent', 'idle', '{}' + )", + ), + ( + "cron_jobs", + "INSERT INTO cron_jobs ( + id, name, enabled, schedule_kind, schedule_value, payload_message, + execution_mode, conversation_id, created_by, created_at, + updated_at, run_count, retry_count, max_retries, queue_enabled + ) VALUES ( + 'orphan_cron', 'Orphan Cron', 1, 'every', '60000', 'message', + 'existing', 'missing_conv', 'user', 1, 1, 0, 0, 3, 0 + )", + ), + ( + "mailbox", + "INSERT INTO mailbox (id, team_id, to_agent_id, from_agent_id, type, content, created_at) + VALUES ('orphan_mail', 'missing_team', 'a1', 'a2', 'message', 'hello', 1)", + ), + ( + "team_tasks", + "INSERT INTO team_tasks (id, team_id, subject, status, created_at, updated_at) + VALUES ('orphan_task', 'missing_team', 'Task', 'pending', 1, 1)", + ), + ] { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 28).await; + sqlx::query(setup_sql).execute(&pool).await.unwrap(); + + let err = run_migration_result(&pool, 29).await.unwrap_err(); + assert!( + err.to_string().contains("CHECK constraint failed"), + "unexpected migration error for {name}: {err}" + ); + } +} + +#[tokio::test] +async fn migration_029_rejects_channel_session_without_channel_user() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 28).await; + + sqlx::query( + "INSERT INTO assistant_sessions ( + id, user_id, agent_type, conversation_id, workspace, chat_id, created_at, last_activity + ) VALUES ( + 'session_orphan_channel_user', 'missing_channel_user', 'acp', NULL, NULL, 'chat-1', 1, 1 + )", + ) + .execute(&pool) + .await + .unwrap(); + + let err = run_migration_result(&pool, 29).await.unwrap_err(); + assert!( + err.to_string().contains("CHECK constraint failed"), + "unexpected migration error: {err}" + ); +} + +#[tokio::test] +async fn migration_029_rejects_channel_session_cross_user_conversation() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 28).await; + + sqlx::query( + "INSERT INTO users (id, username, password_hash, created_at, updated_at) + VALUES ('user_b', 'user_b', 'hash', 1, 1)", + ) + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO conversations (id, user_id, name, type, extra, status, created_at, updated_at) + VALUES ('conv_b', 'user_b', 'User B Conversation', 'chat', '{}', 'pending', 1, 1)", + ) + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO assistant_users ( + id, platform_user_id, platform_type, display_name, authorized_at, last_active, session_id + ) VALUES ( + 'channel_user_a', 'platform-user', 'telegram', NULL, 1, NULL, NULL + )", + ) + .execute(&pool) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO assistant_sessions ( + id, user_id, agent_type, conversation_id, workspace, chat_id, created_at, last_activity + ) VALUES ( + 'session_cross_user', 'channel_user_a', 'acp', 'conv_b', NULL, 'chat-1', 1, 1 + )", + ) + .execute(&pool) + .await + .unwrap(); + + let err = run_migration_result(&pool, 29).await.unwrap_err(); + assert!( + err.to_string().contains("CHECK constraint failed"), + "unexpected migration error: {err}" + ); +} + +#[tokio::test] +async fn migration_029_scopes_project_bind_tables_by_user() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 28).await; + run_migration(&pool, 29).await; + + // projects gains a NOT NULL user_id defaulting to the local user so + // phase-1 store code (which never writes the column) keeps working. + sqlx::query( + "INSERT INTO projects (project_id, name, kind, created_at, updated_at) + VALUES ('proj_default', 'Default Owner', 'standard', 1, 1)", + ) + .execute(&pool) + .await + .unwrap(); + let row = sqlx::query("SELECT user_id FROM projects WHERE project_id = 'proj_default'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(row.get::("user_id"), "system_default_user"); + + // The one-workspace-per-folder uniqueness is per owner now: two users may + // use the same folder as their workspace root; the same user may not. + sqlx::query( + "INSERT INTO users (id, username, password_hash, created_at, updated_at) + VALUES ('user_b', 'user_b', 'hash', 0, 0)", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO folders (folder_id, resource_uri, resource_canonical, created_at, updated_at) + VALUES ('folder_shared', 'file:///shared', 'file:///shared', 1, 1)", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO project_explorer + (pe_id, project_id, owner_user_id, folder_id, role, display_name, order_index, created_at, updated_at) + VALUES ('pe_a', 'proj_default', 'system_default_user', 'folder_shared', 'workspace', NULL, 0, 1, 1)", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO project_explorer + (pe_id, project_id, owner_user_id, folder_id, role, display_name, order_index, created_at, updated_at) + VALUES ('pe_b', 'proj_b', 'user_b', 'folder_shared', 'workspace', NULL, 0, 1, 1)", + ) + .execute(&pool) + .await + .unwrap(); + let same_owner_conflict = sqlx::query( + "INSERT INTO project_explorer + (pe_id, project_id, owner_user_id, folder_id, role, display_name, order_index, created_at, updated_at) + VALUES ('pe_a2', 'proj_other', 'system_default_user', 'folder_shared', 'workspace', NULL, 0, 1, 1)", + ) + .execute(&pool) + .await; + assert!( + same_owner_conflict.is_err(), + "same owner must not claim the same folder as workspace twice" + ); +} + +#[tokio::test] +async fn migration_029_backfills_legacy_project_rows_to_default_user() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 28).await; + + // Rows created by phase-1 store code (no owner columns exist in 028) + // must land on system_default_user after the 029 rebuild, mirroring + // every other legacy backfill in this migration. Cross-owner divergence + // cannot exist pre-029 (no owner columns), so the preflight consistency + // checks are trivially green here; their forward-going enforcement is + // covered by migration_029_scopes_project_bind_tables_by_user. + sqlx::query( + "INSERT INTO projects (project_id, name, kind, created_at, updated_at) + VALUES ('proj_legacy', 'Legacy', 'standard', 1, 1)", + ) + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO project_explorer + (pe_id, project_id, folder_id, role, display_name, order_index, created_at, updated_at) + VALUES ('pe_legacy', 'proj_legacy', 'folder_x', 'workspace', NULL, 0, 1, 1)", + ) + .execute(&pool) + .await + .unwrap(); + + // Clean legacy data: both rows backfill to system_default_user → passes. + run_migration(&pool, 29).await; + let row = sqlx::query("SELECT owner_user_id FROM project_explorer WHERE pe_id = 'pe_legacy'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(row.get::("owner_user_id"), "system_default_user"); +} diff --git a/crates/aionui-extension/Cargo.toml b/crates/aionui-extension/Cargo.toml index 8d762320d..e47b3cf24 100644 --- a/crates/aionui-extension/Cargo.toml +++ b/crates/aionui-extension/Cargo.toml @@ -7,6 +7,7 @@ edition.workspace = true aionui-common.workspace = true aionui-db.workspace = true aionui-api-types.workspace = true +aionui-auth.workspace = true aionui-realtime.workspace = true aionui-runtime.workspace = true async-trait.workspace = true diff --git a/crates/aionui-extension/src/classifier.rs b/crates/aionui-extension/src/classifier.rs index 277f540ba..f043a02b3 100644 --- a/crates/aionui-extension/src/classifier.rs +++ b/crates/aionui-extension/src/classifier.rs @@ -38,11 +38,23 @@ impl AssistantClassifier for DefaultUserClassifier { /// `/api/skills/assistant-skill/*` endpoints dispatch per source. #[async_trait::async_trait] pub trait AssistantRuleDispatcher: Send + Sync { - async fn read_rule(&self, id: &str, locale: Option<&str>) -> Result; - async fn write_rule(&self, id: &str, locale: Option<&str>, content: &str) -> Result<(), ExtensionError>; - async fn delete_rule(&self, id: &str) -> Result; + async fn read_rule(&self, user_id: &str, id: &str, locale: Option<&str>) -> Result; + async fn write_rule( + &self, + user_id: &str, + id: &str, + locale: Option<&str>, + content: &str, + ) -> Result<(), ExtensionError>; + async fn delete_rule(&self, user_id: &str, id: &str) -> Result; - async fn read_skill(&self, id: &str, locale: Option<&str>) -> Result; - async fn write_skill(&self, id: &str, locale: Option<&str>, content: &str) -> Result<(), ExtensionError>; - async fn delete_skill(&self, id: &str) -> Result; + async fn read_skill(&self, user_id: &str, id: &str, locale: Option<&str>) -> Result; + async fn write_skill( + &self, + user_id: &str, + id: &str, + locale: Option<&str>, + content: &str, + ) -> Result<(), ExtensionError>; + async fn delete_skill(&self, user_id: &str, id: &str) -> Result; } diff --git a/crates/aionui-extension/src/lib.rs b/crates/aionui-extension/src/lib.rs index 74a6b1fa0..2cd60a11a 100644 --- a/crates/aionui-extension/src/lib.rs +++ b/crates/aionui-extension/src/lib.rs @@ -55,9 +55,10 @@ pub use skill_service::{ SkillPaths, SkillSource, builtin_skills_corpus, builtin_skills_materialize_marker, delete_skill, delete_skill_with_repo, detect_and_count_external_skills, detect_common_skill_paths, export_skill_with_symlink, get_skill_paths, import_skill, import_skills_with_repo, link_workspace_skills, list_available_skills, - list_available_skills_with_repo, materialize_skills_for_agent, materialize_skills_for_agent_with_repo, - read_builtin_rule, read_builtin_skill, read_skill_info, resolve_skill_paths, scan_for_skills, - sync_skill_catalog_into_repo, + list_available_skills_with_repo, list_available_skills_with_repo_for_user, materialize_skills_for_agent, + materialize_skills_for_agent_with_repo, materialize_skills_for_agent_with_repo_for_user, read_builtin_rule, + read_builtin_skill, read_skill_info, resolve_skill_paths, scan_for_skills, sync_skill_catalog_into_repo, + sync_skill_catalog_into_repo_for_user, }; pub use skill_service::{ delete_assistant_rule, delete_assistant_skill, read_assistant_rule, read_assistant_skill, write_assistant_rule, diff --git a/crates/aionui-extension/src/registry.rs b/crates/aionui-extension/src/registry.rs index 2c44b0255..ea04eafcd 100644 --- a/crates/aionui-extension/src/registry.rs +++ b/crates/aionui-extension/src/registry.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::path::Path; use std::sync::Arc; @@ -27,6 +27,8 @@ use crate::types::{ // `registry::{ExtensionRegistry, ExtensionSummary}` continues to work. pub use crate::registry_helpers::ExtensionSummary; +const SYSTEM_DEFAULT_USER_ID: &str = "system_default_user"; + // --------------------------------------------------------------------------- // Public types // --------------------------------------------------------------------------- @@ -118,7 +120,7 @@ impl ExtensionRegistry { let extensions = self.run_activation_hooks(extensions, &persisted).await; // 6. Resolve contributions. - let contributions = resolve_all_contributions(&extensions); + let contributions = resolve_all_installed_contributions(&extensions); // 7. Persist updated states. let states = build_state_map(&extensions); @@ -171,7 +173,7 @@ impl ExtensionRegistry { let extensions = merge_persisted_states(extensions, &persisted); let extensions = self.run_activation_hooks(extensions, &persisted).await; - let contributions = resolve_all_contributions(&extensions); + let contributions = resolve_all_installed_contributions(&extensions); let states = build_state_map(&extensions); self.state_store.set_all(states).await; @@ -198,78 +200,82 @@ impl ExtensionRegistry { impl ExtensionRegistry { /// Enable an extension by name. /// - /// Updates the in-memory state, re-resolves contributions, persists the - /// change, and broadcasts `extensions.state-changed`. + /// Compatibility wrapper for the local system user. pub async fn enable_extension(&self, name: &str) -> Result<(), ExtensionError> { - let state = { - let mut guard = self.inner.write().await; - - let idx = guard - .extensions - .iter() - .position(|e| e.manifest.name == name) - .ok_or_else(|| ExtensionError::NotFound(name.to_owned()))?; - - if guard.extensions[idx].state.enabled { - debug!(name, "extension already enabled"); - return Ok(()); - } - - guard.extensions[idx].state.enabled = true; - guard.extensions[idx].state.last_activated_at = Some(now_ms()); - - // Re-resolve contributions with updated enabled set. - guard.contributions = resolve_all_contributions(&guard.extensions); + self.enable_extension_for_user(SYSTEM_DEFAULT_USER_ID, name).await + } - guard.extensions[idx].state.clone() - }; + /// Enable an installed extension for one application user. + pub async fn enable_extension_for_user(&self, user_id: &str, name: &str) -> Result<(), ExtensionError> { + self.require_installed(name).await?; + if self.extension_enabled_for_user(user_id, name).await { + debug!(user_id, name, "extension already enabled for user"); + return Ok(()); + } - // Persist + broadcast outside the write lock. - self.state_store.set(state).await; - self.broadcast_state_changed(name, true); + self.state_store.set_user_enabled(user_id, name, true).await; + self.broadcast_state_changed(user_id, name, true); - info!(name, "extension enabled"); + info!(user_id, name, "extension enabled for user"); Ok(()) } /// Disable an extension by name. /// - /// Optionally records a reason (logged for auditing). Updates state, - /// re-resolves contributions, persists, and broadcasts - /// `extensions.state-changed`. + /// Compatibility wrapper for the local system user. pub async fn disable_extension(&self, name: &str, reason: Option<&str>) -> Result<(), ExtensionError> { - let state = { - let mut guard = self.inner.write().await; - - let idx = guard - .extensions - .iter() - .position(|e| e.manifest.name == name) - .ok_or_else(|| ExtensionError::NotFound(name.to_owned()))?; - - if !guard.extensions[idx].state.enabled { - debug!(name, "extension already disabled"); - return Ok(()); - } + self.disable_extension_for_user(SYSTEM_DEFAULT_USER_ID, name, reason) + .await + } - guard.extensions[idx].state.enabled = false; + /// Disable an installed extension for one application user. + pub async fn disable_extension_for_user( + &self, + user_id: &str, + name: &str, + reason: Option<&str>, + ) -> Result<(), ExtensionError> { + self.require_installed(name).await?; + if !self.extension_enabled_for_user(user_id, name).await { + debug!(user_id, name, "extension already disabled for user"); + return Ok(()); + } - // Re-resolve contributions with updated enabled set. - guard.contributions = resolve_all_contributions(&guard.extensions); + self.state_store.set_user_enabled(user_id, name, false).await; + self.broadcast_state_changed(user_id, name, false); - guard.extensions[idx].state.clone() - }; + info!( + user_id, + name, + reason_supplied = reason.is_some(), + "extension disabled for user" + ); + Ok(()) + } - // Persist + broadcast outside the write lock. - self.state_store.set(state).await; - self.broadcast_state_changed(name, false); + async fn require_installed(&self, name: &str) -> Result<(), ExtensionError> { + let guard = self.inner.read().await; + guard + .extensions + .iter() + .any(|extension| extension.manifest.name == name) + .then_some(()) + .ok_or_else(|| ExtensionError::NotFound(name.to_owned())) + } - if let Some(r) = reason { - info!(name, reason = r, "extension disabled"); - } else { - info!(name, "extension disabled"); + async fn extension_enabled_for_user(&self, user_id: &str, name: &str) -> bool { + if let Some(enabled) = self.state_store.get_user_enabled(user_id, name).await { + return enabled; } - Ok(()) + if user_id != SYSTEM_DEFAULT_USER_ID { + return true; + } + let guard = self.inner.read().await; + guard + .extensions + .iter() + .find(|extension| extension.manifest.name == name) + .is_none_or(|extension| extension.state.enabled) } } @@ -280,8 +286,12 @@ impl ExtensionRegistry { impl ExtensionRegistry { /// Return summaries of all loaded extensions. pub async fn get_loaded_extensions(&self) -> Vec { - let guard = self.inner.read().await; - guard.extensions.iter().map(to_summary).collect() + self.get_loaded_extensions_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + /// Return summaries with enabled state overlaid for one user. + pub async fn get_loaded_extensions_for_user(&self, user_id: &str) -> Vec { + self.extensions_for_user(user_id).await.iter().map(to_summary).collect() } pub(crate) fn event_broadcaster(&self) -> Arc { @@ -294,78 +304,150 @@ impl ExtensionRegistry { guard.extensions.iter().find(|e| e.manifest.name == name).cloned() } + /// Look up an installed extension with one user's enabled state overlaid. + pub async fn get_extension_by_name_for_user(&self, user_id: &str, name: &str) -> Option { + self.extensions_for_user(user_id) + .await + .into_iter() + .find(|extension| extension.manifest.name == name) + } + /// Snapshot of all resolved contributions. pub async fn get_contributions(&self) -> ResolvedContributions { - let guard = self.inner.read().await; - guard.contributions.clone() + self.get_contributions_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + pub async fn get_contributions_for_user(&self, user_id: &str) -> ResolvedContributions { + let enabled: HashSet = self + .extensions_for_user(user_id) + .await + .into_iter() + .filter(|extension| extension.state.enabled) + .map(|extension| extension.manifest.name) + .collect(); + let mut contributions = { + let guard = self.inner.read().await; + guard.contributions.clone() + }; + retain_contributions(&mut contributions, &enabled); + contributions } pub async fn get_themes(&self) -> Vec { - let guard = self.inner.read().await; - guard.contributions.themes.clone() + self.get_themes_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + pub async fn get_themes_for_user(&self, user_id: &str) -> Vec { + self.get_contributions_for_user(user_id).await.themes } pub async fn get_assistants(&self) -> Vec { - let guard = self.inner.read().await; - guard.contributions.assistants.clone() + self.get_assistants_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + pub async fn get_assistants_for_user(&self, user_id: &str) -> Vec { + self.get_contributions_for_user(user_id).await.assistants } /// Return `true` if any extension contributes an assistant with this id. pub async fn has_assistant(&self, id: &str) -> bool { - let guard = self.inner.read().await; - guard.contributions.assistants.iter().any(|a| a.id == id) + self.has_assistant_for_user(SYSTEM_DEFAULT_USER_ID, id).await + } + + pub async fn has_assistant_for_user(&self, user_id: &str, id: &str) -> bool { + self.get_assistants_for_user(user_id) + .await + .iter() + .any(|assistant| assistant.id == id) } /// Lookup a single extension-contributed assistant by id. pub async fn get_assistant_by_id(&self, id: &str) -> Option { - let guard = self.inner.read().await; - guard.contributions.assistants.iter().find(|a| a.id == id).cloned() + self.get_assistant_by_id_for_user(SYSTEM_DEFAULT_USER_ID, id).await + } + + pub async fn get_assistant_by_id_for_user(&self, user_id: &str, id: &str) -> Option { + self.get_assistants_for_user(user_id) + .await + .into_iter() + .find(|assistant| assistant.id == id) } pub async fn get_acp_adapters(&self) -> Vec { - let guard = self.inner.read().await; - guard.contributions.acp_adapters.clone() + self.get_acp_adapters_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + pub async fn get_acp_adapters_for_user(&self, user_id: &str) -> Vec { + self.get_contributions_for_user(user_id).await.acp_adapters } pub async fn get_agents(&self) -> Vec { - let guard = self.inner.read().await; - guard.contributions.agents.clone() + self.get_agents_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + pub async fn get_agents_for_user(&self, user_id: &str) -> Vec { + self.get_contributions_for_user(user_id).await.agents } pub async fn get_mcp_servers(&self) -> Vec { - let guard = self.inner.read().await; - guard.contributions.mcp_servers.clone() + self.get_mcp_servers_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + pub async fn get_mcp_servers_for_user(&self, user_id: &str) -> Vec { + self.get_contributions_for_user(user_id).await.mcp_servers } pub async fn get_skills(&self) -> Vec { - let guard = self.inner.read().await; - guard.contributions.skills.clone() + self.get_skills_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + pub async fn get_skills_for_user(&self, user_id: &str) -> Vec { + self.get_contributions_for_user(user_id).await.skills } pub async fn get_settings_tabs(&self) -> Vec { - let guard = self.inner.read().await; - guard.contributions.settings_tabs.clone() + self.get_settings_tabs_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + pub async fn get_settings_tabs_for_user(&self, user_id: &str) -> Vec { + self.get_contributions_for_user(user_id).await.settings_tabs } pub async fn get_webui_contributions(&self) -> Vec { - let guard = self.inner.read().await; - guard.contributions.webui.clone() + self.get_webui_contributions_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + pub async fn get_webui_contributions_for_user(&self, user_id: &str) -> Vec { + self.get_contributions_for_user(user_id).await.webui } pub async fn get_channel_plugins(&self) -> Vec { - let guard = self.inner.read().await; - guard.contributions.channel_plugins.clone() + self.get_channel_plugins_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + pub async fn get_channel_plugins_for_user(&self, user_id: &str) -> Vec { + self.get_contributions_for_user(user_id).await.channel_plugins } pub async fn get_model_providers(&self) -> Vec { - let guard = self.inner.read().await; - guard.contributions.model_providers.clone() + self.get_model_providers_for_user(SYSTEM_DEFAULT_USER_ID).await + } + + pub async fn get_model_providers_for_user(&self, user_id: &str) -> Vec { + self.get_contributions_for_user(user_id).await.model_providers } /// Resolve i18n data for a given locale across all enabled extensions. pub async fn get_i18n_for_locale(&self, locale: &str) -> HashMap> { - let guard = self.inner.read().await; - resolve_i18n_for_all(&guard.extensions, locale) + self.get_i18n_for_locale_for_user(SYSTEM_DEFAULT_USER_ID, locale).await + } + + pub async fn get_i18n_for_locale_for_user( + &self, + user_id: &str, + locale: &str, + ) -> HashMap> { + resolve_i18n_for_all(&self.extensions_for_user(user_id).await, locale) } /// Whether the registry has been initialized. @@ -373,6 +455,65 @@ impl ExtensionRegistry { let guard = self.inner.read().await; guard.initialized } + + async fn extensions_for_user(&self, user_id: &str) -> Vec { + let mut extensions = { + let guard = self.inner.read().await; + guard.extensions.clone() + }; + let overrides = self.state_store.get_user_enabled_all(user_id).await; + for extension in &mut extensions { + extension.state.enabled = overrides + .get(&extension.manifest.name) + .copied() + .unwrap_or(user_id != SYSTEM_DEFAULT_USER_ID || extension.state.enabled); + } + extensions + } +} + +fn resolve_all_installed_contributions(extensions: &[LoadedExtension]) -> ResolvedContributions { + let mut installed = extensions.to_vec(); + for extension in &mut installed { + extension.state.enabled = true; + } + resolve_all_contributions(&installed) +} + +fn retain_contributions(contributions: &mut ResolvedContributions, enabled: &HashSet) { + contributions + .acp_adapters + .retain(|item| enabled.contains(&item.extension_name)); + contributions + .mcp_servers + .retain(|item| enabled.contains(&item.extension_name)); + contributions + .assistants + .retain(|item| enabled.contains(&item.extension_name)); + contributions + .agents + .retain(|item| enabled.contains(&item.extension_name)); + contributions + .skills + .retain(|item| enabled.contains(&item.extension_name)); + contributions + .themes + .retain(|item| enabled.contains(&item.extension_name)); + contributions + .channel_plugins + .retain(|item| enabled.contains(&item.extension_name)); + contributions + .webui + .retain(|item| enabled.contains(&item.extension_name)); + contributions + .settings_tabs + .retain(|item| enabled.contains(&item.extension_name)); + contributions + .model_providers + .retain(|item| enabled.contains(&item.extension_name)); + contributions + .i18n + .retain(|extension_name, _| enabled.contains(extension_name)); } // --------------------------------------------------------------------------- @@ -380,8 +521,11 @@ impl ExtensionRegistry { // --------------------------------------------------------------------------- impl ExtensionRegistry { - fn broadcast_state_changed(&self, name: &str, enabled: bool) { - let event = WebSocketMessage::new("extensions.state-changed", json!({ "name": name, "enabled": enabled })); + fn broadcast_state_changed(&self, user_id: &str, name: &str, enabled: bool) { + let event = WebSocketMessage::new( + "extensions.state-changed", + json!({ "user_id": user_id, "name": name, "enabled": enabled }), + ); self.broadcaster.broadcast(event); } @@ -552,13 +696,11 @@ mod tests { // Enable registry.enable_extension("test-ext").await.unwrap(); - { - let guard = registry.inner.read().await; - assert!(guard.extensions[0].state.enabled); - } + assert!(registry.get_loaded_extensions().await[0].enabled); let msg = rx.recv().await.unwrap(); assert_eq!(msg.name, "extensions.state-changed"); + assert_eq!(msg.data["user_id"], SYSTEM_DEFAULT_USER_ID); assert_eq!(msg.data["enabled"], true); // Disable @@ -566,10 +708,7 @@ mod tests { .disable_extension("test-ext", Some("test reason")) .await .unwrap(); - { - let guard = registry.inner.read().await; - assert!(!guard.extensions[0].state.enabled); - } + assert!(!registry.get_loaded_extensions().await[0].enabled); let msg = rx.recv().await.unwrap(); assert_eq!(msg.name, "extensions.state-changed"); diff --git a/crates/aionui-extension/src/routes.rs b/crates/aionui-extension/src/routes.rs index 4d3ef27ad..c53d7a083 100644 --- a/crates/aionui-extension/src/routes.rs +++ b/crates/aionui-extension/src/routes.rs @@ -6,7 +6,7 @@ use std::path::Path as FsPath; use axum::Router; use axum::body::Body; use axum::extract::rejection::JsonRejection; -use axum::extract::{Json, Path, State}; +use axum::extract::{Extension, Json, Path, State}; use axum::http::{HeaderValue, StatusCode, header}; use axum::response::Response; use axum::routing::{get, post}; @@ -15,6 +15,7 @@ use aionui_api_types::{ ApiResponse, DisableExtensionRequest, EnableExtensionRequest, ExtensionSummaryResponse, GetI18nRequest, GetPermissionsRequest, GetRiskLevelRequest, PermissionDetailResponse, PermissionSummaryResponse, }; +use aionui_auth::CurrentUser; use aionui_common::{ApiError, now_ms}; use serde_json::json; @@ -172,8 +173,9 @@ pub fn extension_routes(state: ExtensionRouterState) -> Router { /// `GET /api/extensions` — list all loaded extensions. async fn get_loaded_extensions( State(state): State, + Extension(user): Extension, ) -> Result>>, ApiError> { - let summaries = state.registry.get_loaded_extensions().await; + let summaries = state.registry.get_loaded_extensions_for_user(&user.id).await; let resp: Vec = summaries .into_iter() .map(|s| { @@ -197,8 +199,9 @@ async fn get_loaded_extensions( /// `GET /api/extensions/themes` — get all resolved themes. async fn get_themes( State(state): State, + Extension(user): Extension, ) -> Result>, ApiError> { - let themes = state.registry.get_themes().await; + let themes = state.registry.get_themes_for_user(&user.id).await; let timestamp = now_ms(); let value = serde_json::Value::Array( themes @@ -222,8 +225,9 @@ async fn get_themes( /// `GET /api/extensions/assistants` — get all resolved assistants. async fn get_assistants( State(state): State, + Extension(user): Extension, ) -> Result>, ApiError> { - let assistants = state.registry.get_assistants().await; + let assistants = state.registry.get_assistants_for_user(&user.id).await; let value = serde_json::Value::Array( assistants .into_iter() @@ -254,8 +258,9 @@ async fn get_assistants( /// `GET /api/extensions/acp-adapters` — get all resolved ACP adapters. async fn get_acp_adapters( State(state): State, + Extension(user): Extension, ) -> Result>, ApiError> { - let adapters = state.registry.get_acp_adapters().await; + let adapters = state.registry.get_acp_adapters_for_user(&user.id).await; let value = serde_json::Value::Array( adapters .into_iter() @@ -294,8 +299,9 @@ async fn get_acp_adapters( /// `GET /api/extensions/agents` — get all resolved agents. async fn get_agents( State(state): State, + Extension(user): Extension, ) -> Result>, ApiError> { - let agents = state.registry.get_agents().await; + let agents = state.registry.get_agents_for_user(&user.id).await; let value = serde_json::Value::Array( agents .into_iter() @@ -326,8 +332,9 @@ async fn get_agents( /// `GET /api/extensions/mcp-servers` — get all resolved MCP servers. async fn get_mcp_servers( State(state): State, + Extension(user): Extension, ) -> Result>, ApiError> { - let servers = state.registry.get_mcp_servers().await; + let servers = state.registry.get_mcp_servers_for_user(&user.id).await; let timestamp = now_ms(); let value = serde_json::Value::Array( servers @@ -371,8 +378,9 @@ async fn get_mcp_servers( /// `GET /api/extensions/skills` — get all resolved skills. async fn get_skills( State(state): State, + Extension(user): Extension, ) -> Result>, ApiError> { - let skills = state.registry.get_skills().await; + let skills = state.registry.get_skills_for_user(&user.id).await; let value = serde_json::Value::Array( skills .into_iter() @@ -391,8 +399,9 @@ async fn get_skills( /// `GET /api/extensions/channel-plugins` — get all resolved channel plugins. async fn get_channel_plugins( State(state): State, + Extension(user): Extension, ) -> Result>, ApiError> { - let plugins = state.registry.get_channel_plugins().await; + let plugins = state.registry.get_channel_plugins_for_user(&user.id).await; let value = serde_json::Value::Array( plugins .into_iter() @@ -425,8 +434,9 @@ async fn get_channel_plugins( /// `GET /api/extensions/settings-tabs` — get all resolved settings tabs. async fn get_settings_tabs( State(state): State, + Extension(user): Extension, ) -> Result>, ApiError> { - let tabs = state.registry.get_settings_tabs().await; + let tabs = state.registry.get_settings_tabs_for_user(&user.id).await; let value = serde_json::to_value(&tabs).unwrap_or_default(); Ok(Json(ApiResponse::ok(value))) } @@ -435,13 +445,17 @@ async fn get_settings_tabs( /// extension asset under the trusted extension root. async fn get_extension_asset( State(state): State, + Extension(user): Extension, Path((extension_name, asset_path)): Path<(String, String)>, ) -> Result { let ext = state .registry - .get_extension_by_name(&extension_name) + .get_extension_by_name_for_user(&user.id, &extension_name) .await .ok_or_else(|| ApiError::NotFound(format!("Extension not found: {extension_name}")))?; + if !ext.state.enabled { + return Err(ApiError::NotFound(format!("Extension not found: {extension_name}"))); + } let canonical_root = tokio::fs::canonicalize(&ext.directory) .await @@ -480,8 +494,9 @@ async fn get_extension_asset( /// `GET /api/extensions/webui` — get all WebUI contributions. async fn get_webui( State(state): State, + Extension(user): Extension, ) -> Result>, ApiError> { - let webui = state.registry.get_webui_contributions().await; + let webui = state.registry.get_webui_contributions_for_user(&user.id).await; let value = serde_json::to_value(&webui).unwrap_or_default(); Ok(Json(ApiResponse::ok(value))) } @@ -501,10 +516,11 @@ async fn get_agent_activity( /// `POST /api/extensions/i18n` — get i18n data for a locale. async fn get_i18n( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>>>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let data = state.registry.get_i18n_for_locale(&req.locale).await; + let data = state.registry.get_i18n_for_locale_for_user(&user.id, &req.locale).await; Ok(Json(ApiResponse::ok(data))) } @@ -572,22 +588,24 @@ async fn get_risk_level( /// `POST /api/extensions/enable` — enable an extension. async fn enable_extension( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - state.registry.enable_extension(&req.name).await?; + state.registry.enable_extension_for_user(&user.id, &req.name).await?; Ok(Json(ApiResponse::success())) } /// `POST /api/extensions/disable` — disable an extension. async fn disable_extension( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; state .registry - .disable_extension(&req.name, req.reason.as_deref()) + .disable_extension_for_user(&user.id, &req.name, req.reason.as_deref()) .await?; Ok(Json(ApiResponse::success())) } @@ -700,7 +718,11 @@ mod tests { .await .unwrap(); - (extension_routes(ExtensionRouterState { registry }), tmp, ext_dir) + ( + extension_routes(ExtensionRouterState { registry }).layer(Extension(CurrentUser::local_default())), + tmp, + ext_dir, + ) } #[tokio::test] diff --git a/crates/aionui-extension/src/skill_routes.rs b/crates/aionui-extension/src/skill_routes.rs index 1d05d3d38..a4a900040 100644 --- a/crates/aionui-extension/src/skill_routes.rs +++ b/crates/aionui-extension/src/skill_routes.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use axum::Router; use axum::extract::rejection::JsonRejection; -use axum::extract::{Json, Path as AxumPath, State}; +use axum::extract::{Extension, Json, Path as AxumPath, State}; use axum::routing::{delete, get, post}; use tracing::warn; @@ -17,6 +17,7 @@ use aionui_api_types::{ ScannedSkillResponse, SkillImportLimitsResponse, SkillImportRecordResponse, SkillListItemResponse, SkillPathsResponse, SkillSourceResponse, WriteAssistantRuleRequest, }; +use aionui_auth::CurrentUser; use aionui_common::ApiError; use aionui_db::ISkillRepository; @@ -111,8 +112,14 @@ pub fn skill_routes(state: SkillRouterState) -> Router { /// `GET /api/skills` — list all available skills. async fn list_skills( State(state): State, + Extension(current_user): Extension, ) -> Result>>, ApiError> { - let items = skill_service::list_available_skills_with_repo(&state.skill_paths, state.skill_repo.as_ref()).await?; + let items = skill_service::list_available_skills_with_repo_for_user( + &state.skill_paths, + state.skill_repo.as_ref(), + ¤t_user.id, + ) + .await?; let resp: Vec = items .into_iter() .map(|s| SkillListItemResponse { @@ -164,12 +171,14 @@ async fn get_import_limits() -> Result, + Extension(current_user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let outcome = match skill_service::import_skills_with_repo( + let outcome = match skill_service::import_skills_with_repo_for_user( &state.skill_paths, state.skill_repo.as_ref(), + ¤t_user.id, Path::new(&req.skill_path), ) .await @@ -227,19 +236,27 @@ async fn export_skill_symlink( /// `DELETE /api/skills/:name` — delete a user-custom skill. async fn delete_skill( State(state): State, + Extension(current_user): Extension, AxumPath(name): AxumPath, ) -> Result>, ApiError> { - skill_service::delete_skill_with_repo(&state.skill_paths, state.skill_repo.as_ref(), &name).await?; + skill_service::delete_skill_with_repo_for_user( + &state.skill_paths, + state.skill_repo.as_ref(), + ¤t_user.id, + &name, + ) + .await?; Ok(Json(ApiResponse::success())) } /// `GET /api/skills/import-history` — list recent skill import records. async fn list_import_history( State(state): State, + Extension(current_user): Extension, ) -> Result>>, ApiError> { let records = state .skill_repo - .list_import_records(100) + .list_import_records_for_user(¤t_user.id, 100) .await .map_err(ExtensionError::from)?; let resp = records @@ -358,15 +375,17 @@ async fn read_builtin_skill( /// backend no longer copies any files per-conversation. async fn materialize_for_agent( State(state): State, + Extension(current_user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; if req.conversation_id.trim().is_empty() { return Err(ApiError::BadRequest("conversationId must not be empty".into())); } - let resolved = skill_service::materialize_skills_for_agent_with_repo( + let resolved = skill_service::materialize_skills_for_agent_with_repo_for_user( &state.skill_paths, state.skill_repo.as_ref(), + ¤t_user.id, &req.conversation_id, &req.skills, ) @@ -391,11 +410,14 @@ async fn materialize_for_agent( /// back to user-directory-only legacy behavior otherwise. async fn read_assistant_rule( State(state): State, + Extension(current_user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; if let Some(dispatcher) = &state.assistant_dispatcher { - let content = dispatcher.read_rule(&req.assistant_id, req.locale.as_deref()).await?; + let content = dispatcher + .read_rule(¤t_user.id, &req.assistant_id, req.locale.as_deref()) + .await?; return Ok(Json(ApiResponse::ok(content))); } let content = @@ -408,12 +430,13 @@ async fn read_assistant_rule( /// Dispatches by source: builtin / extension ids reject with 400. async fn write_assistant_rule( State(state): State, + Extension(current_user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; if let Some(dispatcher) = &state.assistant_dispatcher { dispatcher - .write_rule(&req.assistant_id, req.locale.as_deref(), &req.content) + .write_rule(¤t_user.id, &req.assistant_id, req.locale.as_deref(), &req.content) .await?; return Ok(Json(ApiResponse::ok(true))); } @@ -430,10 +453,11 @@ async fn write_assistant_rule( /// `DELETE /api/skills/assistant-rule/:id` — delete all locale versions. async fn delete_assistant_rule( State(state): State, + Extension(current_user): Extension, AxumPath(id): AxumPath, ) -> Result>, ApiError> { if let Some(dispatcher) = &state.assistant_dispatcher { - let ok = dispatcher.delete_rule(&id).await?; + let ok = dispatcher.delete_rule(¤t_user.id, &id).await?; return Ok(Json(ApiResponse::ok(ok))); } let ok = skill_service::delete_assistant_rule(&state.skill_paths, &id).await?; @@ -449,11 +473,14 @@ async fn delete_assistant_rule( /// Dispatches by source via [`AssistantRuleDispatcher`] when wired. async fn read_assistant_skill( State(state): State, + Extension(current_user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; if let Some(dispatcher) = &state.assistant_dispatcher { - let content = dispatcher.read_skill(&req.assistant_id, req.locale.as_deref()).await?; + let content = dispatcher + .read_skill(¤t_user.id, &req.assistant_id, req.locale.as_deref()) + .await?; return Ok(Json(ApiResponse::ok(content))); } let content = @@ -466,12 +493,13 @@ async fn read_assistant_skill( /// Dispatches by source: builtin / extension ids reject with 400. async fn write_assistant_skill( State(state): State, + Extension(current_user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; if let Some(dispatcher) = &state.assistant_dispatcher { dispatcher - .write_skill(&req.assistant_id, req.locale.as_deref(), &req.content) + .write_skill(¤t_user.id, &req.assistant_id, req.locale.as_deref(), &req.content) .await?; return Ok(Json(ApiResponse::ok(true))); } @@ -488,10 +516,11 @@ async fn write_assistant_skill( /// `DELETE /api/skills/assistant-skill/:id` — delete all locale versions. async fn delete_assistant_skill( State(state): State, + Extension(current_user): Extension, AxumPath(id): AxumPath, ) -> Result>, ApiError> { if let Some(dispatcher) = &state.assistant_dispatcher { - let ok = dispatcher.delete_skill(&id).await?; + let ok = dispatcher.delete_skill(¤t_user.id, &id).await?; return Ok(Json(ApiResponse::ok(ok))); } let ok = skill_service::delete_assistant_skill(&state.skill_paths, &id).await?; diff --git a/crates/aionui-extension/src/skill_service.rs b/crates/aionui-extension/src/skill_service.rs index 8694a3fb6..31af1b0ad 100644 --- a/crates/aionui-extension/src/skill_service.rs +++ b/crates/aionui-extension/src/skill_service.rs @@ -29,6 +29,7 @@ pub const BUILTIN_SKILLS_ENV_VAR: &str = "AIONUI_BUILTIN_SKILLS_PATH"; const MAX_SKILL_IMPORT_FILE_BYTES: u64 = 50 * 1024 * 1024; const MAX_SKILL_IMPORT_TOTAL_BYTES: u64 = 200 * 1024 * 1024; const IMPORT_STAGING_PREFIX: &str = ".import-staging-"; +const DEFAULT_USER_ID: &str = "system_default_user"; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct SkillImportLimits { @@ -332,7 +333,16 @@ pub async fn list_available_skills_with_repo( paths: &SkillPaths, repo: &dyn ISkillRepository, ) -> Result, ExtensionError> { - list_skills_from_repo(paths, repo).await + list_available_skills_with_repo_for_user(paths, repo, DEFAULT_USER_ID).await +} + +/// List available global and user-owned skills using the database state source. +pub async fn list_available_skills_with_repo_for_user( + paths: &SkillPaths, + repo: &dyn ISkillRepository, + user_id: &str, +) -> Result, ExtensionError> { + list_skills_from_repo(paths, repo, user_id).await } /// Emit a [`SkillListItem`] for every built-in skill (both auto-inject @@ -513,16 +523,29 @@ pub async fn import_skill_with_repo( repo: &dyn ISkillRepository, skill_path: &Path, ) -> Result { - let copied = copy_skill_into_user_dir(paths, skill_path).await?; - let overwritten = repo.find_by_name(&copied.name).await?.is_some(); + import_skill_with_repo_for_user(paths, repo, DEFAULT_USER_ID, skill_path).await +} + +/// Import a user-owned skill and persist its management metadata in the database. +pub async fn import_skill_with_repo_for_user( + paths: &SkillPaths, + repo: &dyn ISkillRepository, + user_id: &str, + skill_path: &Path, +) -> Result { + let copied = copy_skill_into_user_dir_for_user(paths, user_id, skill_path).await?; + let overwritten = repo.find_by_name_for_user(user_id, &copied.name).await?.is_some(); let row = repo - .upsert(UpsertSkillParams { - name: &copied.name, - description: Some(&copied.description), - path: copied.target_dir.to_string_lossy().as_ref(), - source: "user", - enabled: true, - }) + .upsert_for_user( + user_id, + UpsertSkillParams { + name: &copied.name, + description: Some(&copied.description), + path: copied.target_dir.to_string_lossy().as_ref(), + source: "user", + enabled: true, + }, + ) .await?; debug!( @@ -547,13 +570,22 @@ struct CopiedSkill { } async fn copy_skill_into_user_dir(paths: &SkillPaths, skill_path: &Path) -> Result { + copy_skill_into_user_dir_for_user(paths, DEFAULT_USER_ID, skill_path).await +} + +async fn copy_skill_into_user_dir_for_user( + paths: &SkillPaths, + user_id: &str, + skill_path: &Path, +) -> Result { let (name, description) = read_skill_info(skill_path).await?; validate_filename(&name)?; - let target_dir = paths.user_skills_dir.join(&name); - tokio::fs::create_dir_all(&paths.user_skills_dir).await?; + let user_skill_root = user_skill_root_for_user(paths, user_id); + let target_dir = user_skill_root.join(&name); + tokio::fs::create_dir_all(&user_skill_root).await?; - let staging_dir = import_staging_dir(paths, &name); + let staging_dir = import_staging_dir(&user_skill_root, &name); replace_existing_path(&staging_dir).await?; let mut budget = SkillImportBudget::default(); if let Err(err) = copy_skill_dir_for_import(skill_path, &staging_dir, skill_path, &mut budget).await { @@ -652,6 +684,16 @@ pub async fn import_skills_with_repo( paths: &SkillPaths, repo: &dyn ISkillRepository, source_path: &Path, +) -> Result { + import_skills_with_repo_for_user(paths, repo, DEFAULT_USER_ID, source_path).await +} + +/// Import one or more user-owned skills and persist import history for `user_id`. +pub async fn import_skills_with_repo_for_user( + paths: &SkillPaths, + repo: &dyn ISkillRepository, + user_id: &str, + source_path: &Path, ) -> Result { let operation_id = aionui_common::generate_prefixed_id("skill_import_op"); let source_label = import_source_label(source_path); @@ -692,6 +734,7 @@ pub async fn import_skills_with_repo( import_skill_dirs_batch_with_repo( paths, repo, + user_id, skill_dirs, &operation_id, &source_label, @@ -715,6 +758,7 @@ pub async fn import_skills_with_repo( return import_skill_dirs_batch_with_repo( paths, repo, + user_id, vec![source_path], &operation_id, &source_label, @@ -733,6 +777,7 @@ pub async fn import_skills_with_repo( return import_skill_dirs_batch_with_repo( paths, repo, + user_id, skills.into_iter().map(|skill| PathBuf::from(skill.path)).collect(), &operation_id, &source_label, @@ -826,6 +871,7 @@ async fn import_skill_dirs_batch( async fn import_skill_dirs_batch_with_repo( paths: &SkillPaths, repo: &dyn ISkillRepository, + user_id: &str, skill_dirs: Vec, operation_id: &str, source_label: &str, @@ -836,43 +882,49 @@ async fn import_skill_dirs_batch_with_repo( for skill_dir in skill_dirs { let source_name = import_source_name(&skill_dir); - match import_skill_with_repo(paths, repo, &skill_dir).await { + match import_skill_with_repo_for_user(paths, repo, user_id, &skill_dir).await { Ok(skill) => { - repo.create_import_record(CreateSkillImportRecordParams { - operation_id, - source_label, - source_path, - source_name: &source_name, - skill_id: skill.skill_id.as_deref(), - skill_name: Some(&skill.name), - status: if skill.overwritten { "overwritten" } else { "imported" }, - error_code: None, - error_path: None, - actual_bytes: Some(u64_to_i64(skill.copied_bytes)), - limit_bytes: None, - line: None, - column: None, - }) + repo.create_import_record_for_user( + user_id, + CreateSkillImportRecordParams { + operation_id, + source_label, + source_path, + source_name: &source_name, + skill_id: skill.skill_id.as_deref(), + skill_name: Some(&skill.name), + status: if skill.overwritten { "overwritten" } else { "imported" }, + error_code: None, + error_path: None, + actual_bytes: Some(u64_to_i64(skill.copied_bytes)), + limit_bytes: None, + line: None, + column: None, + }, + ) .await?; imported.push(skill.name); } Err(err) => { let failure = skill_import_failure_from_error(&source_name, &err); - repo.create_import_record(CreateSkillImportRecordParams { - operation_id, - source_label, - source_path, - source_name: &source_name, - skill_id: None, - skill_name: None, - status: "failed", - error_code: Some(&failure.code), - error_path: failure.error_path.as_deref(), - actual_bytes: failure.actual_bytes, - limit_bytes: failure.limit_bytes, - line: failure.line, - column: failure.column, - }) + repo.create_import_record_for_user( + user_id, + CreateSkillImportRecordParams { + operation_id, + source_label, + source_path, + source_name: &source_name, + skill_id: None, + skill_name: None, + status: "failed", + error_code: Some(&failure.code), + error_path: failure.error_path.as_deref(), + actual_bytes: failure.actual_bytes, + limit_bytes: failure.limit_bytes, + line: failure.line, + column: failure.column, + }, + ) .await?; failed.push(failure); } @@ -1007,12 +1059,23 @@ async fn replace_existing_path(path: &Path) -> Result<(), ExtensionError> { Ok(()) } -fn import_staging_dir(paths: &SkillPaths, skill_name: &str) -> PathBuf { +fn user_skill_root_for_user(paths: &SkillPaths, user_id: &str) -> PathBuf { + if user_id == DEFAULT_USER_ID { + return paths.user_skills_dir.clone(); + } + + let mut hasher = Sha256::new(); + hasher.update(user_id.as_bytes()); + let digest = hasher.finalize(); + paths.user_skills_dir.join(".users").join(hex::encode(&digest[..12])) +} + +fn import_staging_dir(user_skill_root: &Path, skill_name: &str) -> PathBuf { let nonce = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_nanos(); - paths.user_skills_dir.join(format!( + user_skill_root.join(format!( "{IMPORT_STAGING_PREFIX}{skill_name}-{}-{nonce}", std::process::id() )) @@ -1153,13 +1216,23 @@ pub async fn delete_skill_with_repo( paths: &SkillPaths, repo: &dyn ISkillRepository, skill_name: &str, +) -> Result<(), ExtensionError> { + delete_skill_with_repo_for_user(paths, repo, DEFAULT_USER_ID, skill_name).await +} + +/// Soft-delete a user-owned skill through the database state source. +pub async fn delete_skill_with_repo_for_user( + paths: &SkillPaths, + repo: &dyn ISkillRepository, + user_id: &str, + skill_name: &str, ) -> Result<(), ExtensionError> { validate_filename(skill_name)?; if builtin_skill_exists(paths, skill_name) { return Err(ExtensionError::BuiltinSkillDeletion(skill_name.to_string())); } - let Some(row) = repo.find_by_name(skill_name).await? else { + let Some(row) = repo.find_by_name_for_user(user_id, skill_name).await? else { return Err(ExtensionError::SkillNotFound(skill_name.to_string())); }; if !PathBuf::from(&row.path).is_dir() { @@ -1170,7 +1243,7 @@ pub async fn delete_skill_with_repo( ); } - repo.delete_by_name(skill_name).await?; + repo.delete_by_name_for_user(user_id, skill_name).await?; debug!(skill = %skill_name, "skill marked deleted"); Ok(()) } @@ -1209,7 +1282,9 @@ pub struct ResolvedAgentSkill { /// 1. `{builtin_skills_dir}/{name}/` — top-level opt-in builtin. /// 2. `{builtin_skills_dir}/auto-inject/{name}/` — auto-inject builtin. /// 3. `{user_skills_dir}/{name}/` — user-created custom skill. -/// 4. `{cron_skills_dir}/{name}/` — per-job cron skill. +/// +/// Per-job cron skill dirs hold user content and are never probed by +/// name; they resolve only through the owning user's repo row. /// /// No files are copied and no per-conversation directory is created — /// the backend just hands the absolute source paths back to the caller, @@ -1260,9 +1335,20 @@ pub async fn materialize_skills_for_agent_with_repo( repo: &dyn ISkillRepository, conversation_id: &str, skills: &[String], +) -> Result, ExtensionError> { + materialize_skills_for_agent_with_repo_for_user(paths, repo, DEFAULT_USER_ID, conversation_id, skills).await +} + +/// Resolve requested skill names using global and user-owned database state. +pub async fn materialize_skills_for_agent_with_repo_for_user( + paths: &SkillPaths, + repo: &dyn ISkillRepository, + user_id: &str, + conversation_id: &str, + skills: &[String], ) -> Result, ExtensionError> { validate_filename(conversation_id)?; - sync_disk_user_skills_into_repo(paths, repo).await?; + sync_disk_user_skills_into_repo_for_user(paths, repo, user_id).await?; let mut resolved = Vec::with_capacity(skills.len()); for name in skills { @@ -1273,7 +1359,7 @@ pub async fn materialize_skills_for_agent_with_repo( warn!(skill = %name, "skipping skill with invalid name"); continue; } - match resolve_skill_source_path_with_repo(paths, repo, name).await? { + match resolve_skill_source_path_with_repo_for_user(paths, repo, user_id, name).await? { Some(source_path) => resolved.push(ResolvedAgentSkill { name: name.clone(), source_path, @@ -1394,27 +1480,18 @@ async fn resolve_skill_source_path(paths: &SkillPaths, name: &str) -> Result Result, ExtensionError> { - let top = paths.builtin_skills_dir.join(name); - if top.is_dir() { - return Ok(Some(top)); - } - let auto = paths.builtin_skills_dir.join(BUILTIN_AUTO_SKILLS_SUBDIR).join(name); - if auto.is_dir() { - return Ok(Some(auto)); - } - if let Some(row) = repo.find_by_name_any(name).await? { + if let Some(row) = repo.find_by_name_any_for_user(user_id, name).await? { let path = PathBuf::from(&row.path); if path.is_dir() { return Ok(Some(path)); @@ -1425,12 +1502,21 @@ async fn resolve_skill_source_path_with_repo( deleted = row.deleted_at.is_some(), "skill row points at a missing directory" ); - return Ok(None); + if row.user_id.is_some() { + return Ok(None); + } } - let cron = paths.cron_skills_dir.join(name); - if cron.is_dir() { - return Ok(Some(cron)); + let top = paths.builtin_skills_dir.join(name); + if top.is_dir() { + return Ok(Some(top)); } + let auto = paths.builtin_skills_dir.join(BUILTIN_AUTO_SKILLS_SUBDIR).join(name); + if auto.is_dir() { + return Ok(Some(auto)); + } + // Per-job cron skill dirs are intentionally NOT probed as a fallback: + // they hold user content (job prompts) and must resolve only through + // a repo row owned by the requesting user. Ok(None) } @@ -1554,15 +1640,32 @@ pub fn get_skill_paths(paths: &SkillPaths) -> (String, String) { async fn list_skills_from_repo( paths: &SkillPaths, repo: &dyn ISkillRepository, + user_id: &str, ) -> Result, ExtensionError> { let mut items = Vec::new(); - for row in repo.list().await? { + for row in dedupe_user_visible_skill_rows(repo.list_for_user(user_id).await?) { let description = row.description.clone().unwrap_or_default(); items.push(skill_row_to_list_item(paths, row, description)); } Ok(items) } +fn dedupe_user_visible_skill_rows(rows: Vec) -> Vec { + let mut merged: Vec = Vec::new(); + + for row in rows { + if let Some(index) = merged.iter().position(|existing| existing.name == row.name) { + if merged[index].user_id.is_none() && row.user_id.is_some() { + merged[index] = row; + } + } else { + merged.push(row); + } + } + + merged +} + /// Synchronize filesystem-backed skill catalogs into the database. /// /// Listing endpoints read the database only. This synchronization is intended @@ -1572,7 +1675,19 @@ pub async fn sync_skill_catalog_into_repo( paths: &SkillPaths, repo: &dyn ISkillRepository, ) -> Result<(), ExtensionError> { - sync_disk_user_skills_into_repo(paths, repo).await?; + sync_skill_catalog_into_repo_for_user(paths, repo, DEFAULT_USER_ID).await +} + +/// Synchronize filesystem-backed skill catalogs for a specific user. +/// +/// Built-in skills are global. Disk-backed custom skills are attributed to +/// `user_id`, matching the user-scoped catalog exposed by `/api/skills`. +pub async fn sync_skill_catalog_into_repo_for_user( + paths: &SkillPaths, + repo: &dyn ISkillRepository, + user_id: &str, +) -> Result<(), ExtensionError> { + sync_disk_user_skills_into_repo_for_user(paths, repo, user_id).await?; sync_builtin_skills_into_repo(paths, repo).await?; Ok(()) } @@ -1583,40 +1698,30 @@ async fn sync_builtin_skills_into_repo(paths: &SkillPaths, repo: &dyn ISkillRepo if skill.name == BUILTIN_AUTO_SKILLS_SUBDIR { continue; } - sync_managed_skill_into_repo(repo, &skill, "builtin").await?; + sync_global_skill_into_repo(repo, &skill, "builtin").await?; } } let auto_inject_dir = paths.builtin_skills_dir.join(BUILTIN_AUTO_SKILLS_SUBDIR); if let Ok(skills) = scan_skill_dirs(&auto_inject_dir).await { for skill in skills { - sync_managed_skill_into_repo(repo, &skill, "builtin").await?; + sync_global_skill_into_repo(repo, &skill, "builtin").await?; } } - if let Ok(skills) = scan_skill_dirs(&paths.cron_skills_dir).await { - for skill in skills { - sync_managed_skill_into_repo(repo, &skill, "cron").await?; - } - } + // Per-job cron skills are user content (job prompts); their rows are + // owned by the job owner and maintained by aionui-cron, never synced + // here as global rows. Ok(()) } -async fn sync_managed_skill_into_repo( +async fn sync_global_skill_into_repo( repo: &dyn ISkillRepository, skill: &ScannedSkill, source: &str, ) -> Result<(), ExtensionError> { - if let Some(existing) = repo.find_by_name_any(&skill.name).await? - && existing.source == "user" - && existing.deleted_at.is_none() - && existing.enabled - { - return Ok(()); - } - - repo.upsert(UpsertSkillParams { + repo.upsert_global(UpsertSkillParams { name: &skill.name, description: Some(&skill.description), path: &skill.path, @@ -1627,14 +1732,22 @@ async fn sync_managed_skill_into_repo( Ok(()) } -async fn sync_disk_user_skills_into_repo( +async fn sync_disk_user_skills_into_repo_for_user( paths: &SkillPaths, repo: &dyn ISkillRepository, + user_id: &str, ) -> Result<(), ExtensionError> { + if user_id != DEFAULT_USER_ID { + return Ok(()); + } + let scanned = scan_skill_dirs(&paths.user_skills_dir).await?; let mut backfilled = 0usize; for skill in scanned { + if is_user_scoped_skill_storage_path(Path::new(&skill.path)) { + continue; + } if let Err(err) = validate_filename(&skill.name) { warn!( skill = %skill.name, @@ -1643,17 +1756,20 @@ async fn sync_disk_user_skills_into_repo( ); continue; } - if repo.find_by_name_any(&skill.name).await?.is_some() { + if repo.find_by_name_any_for_user(user_id, &skill.name).await?.is_some() { continue; } - repo.upsert(UpsertSkillParams { - name: &skill.name, - description: Some(&skill.description), - path: &skill.path, - source: "user", - enabled: true, - }) + repo.upsert_for_user( + user_id, + UpsertSkillParams { + name: &skill.name, + description: Some(&skill.description), + path: &skill.path, + source: "user", + enabled: true, + }, + ) .await?; backfilled += 1; } @@ -1668,6 +1784,15 @@ async fn sync_disk_user_skills_into_repo( Ok(()) } +fn is_user_scoped_skill_storage_path(path: &Path) -> bool { + path.components().any(|component| { + matches!( + component, + Component::Normal(name) if name == std::ffi::OsStr::new(".users") + ) + }) +} + fn skill_row_to_list_item(paths: &SkillPaths, row: SkillRow, description: String) -> SkillListItem { let source = skill_source_from_row(&row.source); let relative_location = skill_relative_location(paths, &row, source); @@ -2863,7 +2988,7 @@ mod tests { } #[tokio::test] - async fn list_available_skills_with_repo_syncs_builtin_auto_inject_and_cron_sources_into_repo() { + async fn list_available_skills_with_repo_syncs_builtin_sources_but_never_cron_dirs() { let tmp = TempDir::new().unwrap(); let paths = make_disk_builtin_paths(tmp.path()); let repo = make_test_skill_repo().await; @@ -2892,11 +3017,10 @@ mod tests { let auto_cron_row = repo.find_by_name("cron").await.unwrap().unwrap(); assert_eq!(auto_cron_row.source, "builtin"); - let scheduled = skills.iter().find(|skill| skill.name == "scheduled-task").unwrap(); - assert_eq!(scheduled.source, SkillSource::Cron); - assert_eq!(scheduled.relative_location, None); - let scheduled_row = repo.find_by_name("scheduled-task").await.unwrap().unwrap(); - assert_eq!(scheduled_row.source, "cron"); + // Per-job cron skills hold user content; the machine-level catalog + // sync must not ingest them as global rows or list them. + assert!(skills.iter().all(|skill| skill.name != "scheduled-task")); + assert!(repo.find_by_name_any("scheduled-task").await.unwrap().is_none()); } #[tokio::test] diff --git a/crates/aionui-extension/src/state.rs b/crates/aionui-extension/src/state.rs index 2c23c05de..40bc65d07 100644 --- a/crates/aionui-extension/src/state.rs +++ b/crates/aionui-extension/src/state.rs @@ -17,9 +17,9 @@ use crate::types::ExtensionState; /// Manages loading and saving extension states to a JSON file with debounced /// writes. /// -/// State is persisted to `extension-states.json`. Writes are debounced by -/// [`STATE_PERSIST_DEBOUNCE_MS`] to avoid excessive disk I/O when multiple -/// state changes happen in quick succession. +/// Machine-level compatibility state is persisted to `extension-states.json`; +/// per-user enablement is persisted alongside it. Writes are debounced by +/// [`STATE_PERSIST_DEBOUNCE_MS`] to avoid excessive disk I/O. #[derive(Clone)] pub struct ExtensionStateStore { inner: Arc, @@ -28,8 +28,12 @@ pub struct ExtensionStateStore { struct Inner { /// Path to the state JSON file. file_path: PathBuf, + /// Path to per-user enablement overrides. + user_file_path: PathBuf, /// In-memory state map protected by a mutex. states: Mutex>, + /// Per-user enabled/disabled overrides keyed by user id and extension name. + user_states: Mutex>>, /// Notifier used to trigger a debounced write. write_notify: Notify, /// Whether the background writer task has been spawned. @@ -38,6 +42,8 @@ struct Inner { const EXTENSION_STATES_FILE_ENV: &str = "AIONUI_EXTENSION_STATES_FILE"; const DEFAULT_STATES_FILE: &str = "extension-states.json"; +const USER_STATES_FILE: &str = "extension-user-states.json"; +const SYSTEM_DEFAULT_USER_ID: &str = "system_default_user"; #[derive(Debug, Deserialize, Serialize)] struct PersistedStates { @@ -55,13 +61,23 @@ struct PersistedExtensionState { last_version: Option, } +#[derive(Debug, Deserialize, Serialize)] +struct PersistedUserStates { + version: u32, + #[serde(default)] + users: BTreeMap>, +} + impl ExtensionStateStore { /// Create a new store backed by the given file path. pub fn new(file_path: PathBuf) -> Self { + let user_file_path = user_state_file_path(&file_path); Self { inner: Arc::new(Inner { file_path, + user_file_path, states: Mutex::new(HashMap::new()), + user_states: Mutex::new(HashMap::new()), write_notify: Notify::new(), writer_spawned: Mutex::new(false), }), @@ -83,8 +99,16 @@ impl ExtensionStateStore { /// default to enabled). Parse errors are propagated as `ExtensionError`. pub async fn load(&self) -> Result, ExtensionError> { let states = load_states_from_file(&self.inner.file_path)?; - let mut guard = self.inner.states.lock().await; - *guard = states.clone(); + let mut user_states = load_user_states_from_file(&self.inner.user_file_path)?; + user_states.entry(SYSTEM_DEFAULT_USER_ID.to_owned()).or_insert_with(|| { + states + .iter() + .map(|(name, state)| (name.clone(), state.enabled)) + .collect() + }); + + *self.inner.states.lock().await = states.clone(); + *self.inner.user_states.lock().await = user_states; Ok(states) } @@ -104,6 +128,18 @@ impl ExtensionStateStore { guard.clone() } + /// Return a user's explicit enablement override, if one exists. + pub async fn get_user_enabled(&self, user_id: &str, name: &str) -> Option { + let guard = self.inner.user_states.lock().await; + guard.get(user_id).and_then(|states| states.get(name)).copied() + } + + /// Snapshot all explicit enablement overrides for one user. + pub async fn get_user_enabled_all(&self, user_id: &str) -> HashMap { + let guard = self.inner.user_states.lock().await; + guard.get(user_id).cloned().unwrap_or_default() + } + // ----------------------------------------------------------------------- // Write (debounced) // ----------------------------------------------------------------------- @@ -127,6 +163,19 @@ impl ExtensionStateStore { self.schedule_write().await; } + /// Persist an enablement override for one user without changing the + /// machine-level installation state. + pub async fn set_user_enabled(&self, user_id: &str, name: &str, enabled: bool) { + { + let mut guard = self.inner.user_states.lock().await; + guard + .entry(user_id.to_owned()) + .or_default() + .insert(name.to_owned(), enabled); + } + self.schedule_write().await; + } + /// Remove the persisted state for an extension and schedule a write. pub async fn remove(&self, name: &str) { { @@ -142,11 +191,16 @@ impl ExtensionStateStore { /// Immediately write the current in-memory states to disk (no debounce). pub async fn flush(&self) -> Result<(), ExtensionError> { - let snapshot = { + let state_snapshot = { let guard = self.inner.states.lock().await; guard.clone() }; - save_states_to_file(&self.inner.file_path, &snapshot) + let user_snapshot = { + let guard = self.inner.user_states.lock().await; + guard.clone() + }; + save_states_to_file(&self.inner.file_path, &state_snapshot)?; + save_user_states_to_file(&self.inner.user_file_path, &user_snapshot) } // ----------------------------------------------------------------------- @@ -177,12 +231,18 @@ impl ExtensionStateStore { // additional notifications. tokio::time::sleep(std::time::Duration::from_millis(STATE_PERSIST_DEBOUNCE_MS)).await; - let snapshot = { + let state_snapshot = { let guard = inner.states.lock().await; guard.clone() }; + let user_snapshot = { + let guard = inner.user_states.lock().await; + guard.clone() + }; - if let Err(e) = save_states_to_file(&inner.file_path, &snapshot) { + if let Err(e) = save_states_to_file(&inner.file_path, &state_snapshot) + .and_then(|()| save_user_states_to_file(&inner.user_file_path, &user_snapshot)) + { error!(error = %e, "failed to persist extension states"); } else { debug!(path = %inner.file_path.display(), "extension states persisted"); @@ -192,6 +252,72 @@ impl ExtensionStateStore { } } +fn user_state_file_path(state_file_path: &Path) -> PathBuf { + if state_file_path.file_name().and_then(|name| name.to_str()) == Some(DEFAULT_STATES_FILE) { + return state_file_path.with_file_name(USER_STATES_FILE); + } + + let stem = state_file_path + .file_stem() + .and_then(|value| value.to_str()) + .unwrap_or("extension-states"); + let file_name = match state_file_path.extension().and_then(|value| value.to_str()) { + Some(extension) => format!("{stem}-user-states.{extension}"), + None => format!("{stem}-user-states.json"), + }; + state_file_path.with_file_name(file_name) +} + +fn load_user_states_from_file(path: &Path) -> Result>, ExtensionError> { + let bytes = match std::fs::read(path) { + Ok(bytes) => bytes, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(HashMap::new()), + Err(error) => return Err(ExtensionError::Io(error)), + }; + let persisted: PersistedUserStates = serde_json::from_slice(&bytes)?; + if persisted.version != 1 { + return Err(ExtensionError::StatePersistence(format!( + "unsupported extension user state file version {} at {}", + persisted.version, + path.display() + ))); + } + Ok(persisted + .users + .into_iter() + .map(|(user_id, states)| (user_id, states.into_iter().collect())) + .collect()) +} + +fn save_user_states_to_file( + path: &Path, + states: &HashMap>, +) -> Result<(), ExtensionError> { + if let Some(parent) = path.parent() + && !parent.exists() + { + std::fs::create_dir_all(parent)?; + } + + let users = states + .iter() + .map(|(user_id, extensions)| { + ( + user_id.clone(), + extensions + .iter() + .map(|(name, enabled)| (name.clone(), *enabled)) + .collect(), + ) + }) + .collect(); + let json = serde_json::to_string_pretty(&PersistedUserStates { version: 1, users })?; + let tmp_path = path.with_extension("json.tmp"); + std::fs::write(&tmp_path, json.as_bytes())?; + std::fs::rename(&tmp_path, path)?; + Ok(()) +} + // --------------------------------------------------------------------------- // File I/O (pure functions, no async needed) // --------------------------------------------------------------------------- @@ -465,6 +591,18 @@ mod tests { assert_eq!(path, override_path); } + #[test] + fn user_state_path_tracks_configured_state_file() { + assert_eq!( + user_state_file_path(Path::new("/data/extension-states.json")), + PathBuf::from("/data/extension-user-states.json") + ); + assert_eq!( + user_state_file_path(Path::new("/tmp/custom.json")), + PathBuf::from("/tmp/custom-user-states.json") + ); + } + #[test] fn save_atomic_write() { let tmp = TempDir::new().unwrap(); diff --git a/crates/aionui-extension/tests/assistant_dispatch_test.rs b/crates/aionui-extension/tests/assistant_dispatch_test.rs index 8653d5a40..5dc0eb4e3 100644 --- a/crates/aionui-extension/tests/assistant_dispatch_test.rs +++ b/crates/aionui-extension/tests/assistant_dispatch_test.rs @@ -9,11 +9,14 @@ use std::path::PathBuf; use std::sync::{Arc, Mutex}; use aionui_api_types::{ApiResponse, AssistantSource, SkillImportLimitsResponse}; +use aionui_auth::CurrentUser; +use aionui_db::{UserStatus, UserType}; use aionui_extension::classifier::{AssistantClassifier, AssistantRuleDispatcher}; use aionui_extension::error::ExtensionError; use aionui_extension::external_paths::ExternalPathsManager; use aionui_extension::skill_routes::{SkillRouterState, skill_routes}; use aionui_extension::skill_service::SkillPaths; +use axum::Extension; use axum::body::Body; use axum::http::{Request, StatusCode}; use http_body_util::BodyExt; @@ -25,12 +28,12 @@ use tower::ServiceExt; #[derive(Default)] struct CallLog { - rule_reads: Vec<(String, Option)>, - rule_writes: Vec<(String, Option, String)>, - rule_deletes: Vec, - skill_reads: Vec<(String, Option)>, - skill_writes: Vec<(String, Option, String)>, - skill_deletes: Vec, + rule_reads: Vec<(String, String, Option)>, + rule_writes: Vec<(String, String, Option, String)>, + rule_deletes: Vec<(String, String)>, + skill_reads: Vec<(String, String, Option)>, + skill_writes: Vec<(String, String, Option, String)>, + skill_deletes: Vec<(String, String)>, } struct FakeDispatcher { @@ -55,69 +58,91 @@ impl AssistantClassifier for FakeDispatcher { #[async_trait::async_trait] impl AssistantRuleDispatcher for FakeDispatcher { - async fn read_rule(&self, id: &str, locale: Option<&str>) -> Result { + async fn read_rule(&self, user_id: &str, id: &str, locale: Option<&str>) -> Result { self.log .lock() .unwrap() .rule_reads - .push((id.to_string(), locale.map(str::to_string))); + .push((user_id.to_string(), id.to_string(), locale.map(str::to_string))); Ok(self.rule_content.get(id).cloned().unwrap_or_default()) } - async fn write_rule(&self, id: &str, locale: Option<&str>, content: &str) -> Result<(), ExtensionError> { + async fn write_rule( + &self, + user_id: &str, + id: &str, + locale: Option<&str>, + content: &str, + ) -> Result<(), ExtensionError> { if self.reject_writes_for.contains(id) { return Err(ExtensionError::InvalidRequest( "Cannot write rule for built-in assistant".into(), )); } - self.log - .lock() - .unwrap() - .rule_writes - .push((id.to_string(), locale.map(str::to_string), content.to_string())); + self.log.lock().unwrap().rule_writes.push(( + user_id.to_string(), + id.to_string(), + locale.map(str::to_string), + content.to_string(), + )); Ok(()) } - async fn delete_rule(&self, id: &str) -> Result { + async fn delete_rule(&self, user_id: &str, id: &str) -> Result { if self.reject_writes_for.contains(id) { return Err(ExtensionError::InvalidRequest( "Cannot delete rule for built-in assistant".into(), )); } - self.log.lock().unwrap().rule_deletes.push(id.to_string()); + self.log + .lock() + .unwrap() + .rule_deletes + .push((user_id.to_string(), id.to_string())); Ok(true) } - async fn read_skill(&self, id: &str, locale: Option<&str>) -> Result { + async fn read_skill(&self, user_id: &str, id: &str, locale: Option<&str>) -> Result { self.log .lock() .unwrap() .skill_reads - .push((id.to_string(), locale.map(str::to_string))); + .push((user_id.to_string(), id.to_string(), locale.map(str::to_string))); Ok(self.skill_content.get(id).cloned().unwrap_or_default()) } - async fn write_skill(&self, id: &str, locale: Option<&str>, content: &str) -> Result<(), ExtensionError> { + async fn write_skill( + &self, + user_id: &str, + id: &str, + locale: Option<&str>, + content: &str, + ) -> Result<(), ExtensionError> { if self.reject_writes_for.contains(id) { return Err(ExtensionError::InvalidRequest( "Cannot write skill for built-in assistant".into(), )); } - self.log - .lock() - .unwrap() - .skill_writes - .push((id.to_string(), locale.map(str::to_string), content.to_string())); + self.log.lock().unwrap().skill_writes.push(( + user_id.to_string(), + id.to_string(), + locale.map(str::to_string), + content.to_string(), + )); Ok(()) } - async fn delete_skill(&self, id: &str) -> Result { + async fn delete_skill(&self, user_id: &str, id: &str) -> Result { if self.reject_writes_for.contains(id) { return Err(ExtensionError::InvalidRequest( "Cannot delete skill for built-in assistant".into(), )); } - self.log.lock().unwrap().skill_deletes.push(id.to_string()); + self.log + .lock() + .unwrap() + .skill_deletes + .push((user_id.to_string(), id.to_string())); Ok(true) } } @@ -149,7 +174,12 @@ async fn router_with_dispatcher(dispatcher: Arc) -> axum::Router external_paths_manager: ext_mgr, assistant_dispatcher: Some(dispatcher), }; - skill_routes(state) + skill_routes(state).layer(Extension(CurrentUser { + id: "user-current".into(), + username: "user-current".into(), + user_type: UserType::Local, + status: UserStatus::Active, + })) } async fn body_json(resp: axum::response::Response) -> T { @@ -238,8 +268,9 @@ async fn read_rule_routes_through_dispatcher_for_builtin() { let log = dispatcher.log.lock().unwrap(); assert_eq!(log.rule_reads.len(), 1); - assert_eq!(log.rule_reads[0].0, "builtin-office"); - assert_eq!(log.rule_reads[0].1.as_deref(), Some("en-US")); + assert_eq!(log.rule_reads[0].0, "user-current"); + assert_eq!(log.rule_reads[0].1, "builtin-office"); + assert_eq!(log.rule_reads[0].2.as_deref(), Some("en-US")); } #[tokio::test] @@ -342,7 +373,8 @@ async fn write_rule_allows_user() { let log = dispatcher.log.lock().unwrap(); assert_eq!(log.rule_writes.len(), 1); - assert_eq!(log.rule_writes[0].2, "rule!"); + assert_eq!(log.rule_writes[0].0, "user-current"); + assert_eq!(log.rule_writes[0].3, "rule!"); } #[tokio::test] @@ -385,7 +417,7 @@ async fn delete_rule_user_dispatches() { assert_eq!(resp.status(), StatusCode::OK); let log = dispatcher.log.lock().unwrap(); - assert_eq!(log.rule_deletes, vec!["u1".to_string()]); + assert_eq!(log.rule_deletes, vec![("user-current".to_string(), "u1".to_string())]); } #[tokio::test] @@ -413,6 +445,7 @@ async fn read_skill_routes_through_dispatcher_for_builtin() { let log = dispatcher.log.lock().unwrap(); assert_eq!(log.skill_reads.len(), 1); + assert_eq!(log.skill_reads[0].0, "user-current"); } #[tokio::test] @@ -457,5 +490,5 @@ async fn delete_skill_user_dispatches() { assert_eq!(resp.status(), StatusCode::OK); let log = dispatcher.log.lock().unwrap(); - assert_eq!(log.skill_deletes, vec!["u1".to_string()]); + assert_eq!(log.skill_deletes, vec![("user-current".to_string(), "u1".to_string())]); } diff --git a/crates/aionui-extension/tests/cron_skill_resolve_test.rs b/crates/aionui-extension/tests/cron_skill_resolve_test.rs index 38b117256..9e735662d 100644 --- a/crates/aionui-extension/tests/cron_skill_resolve_test.rs +++ b/crates/aionui-extension/tests/cron_skill_resolve_test.rs @@ -1,4 +1,6 @@ +use aionui_db::{ISkillRepository, IUserRepository, SqliteSkillRepository, SqliteUserRepository, UpsertSkillParams}; use aionui_extension::{resolve_skill_paths, skill_service}; +use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; fn unique_temp_dir(label: &str) -> std::path::PathBuf { @@ -9,6 +11,27 @@ fn unique_temp_dir(label: &str) -> std::path::PathBuf { std::env::temp_dir().join(format!("aionui-extension-{label}-{}-{nanos}", std::process::id())) } +fn write_cron_skill(base: &std::path::Path, dir_name: &str, skill_name: &str) -> std::path::PathBuf { + let skill_dir = base.join("cron").join("skills").join(dir_name); + std::fs::create_dir_all(&skill_dir).unwrap(); + std::fs::write( + skill_dir.join("SKILL.md"), + format!("---\nname: {skill_name}\ndescription: Saved cron skill\n---\nUse the saved steps."), + ) + .unwrap(); + skill_dir +} + +async fn make_skill_repo_with_users(usernames: &[&str]) -> (Arc, Vec) { + let db = aionui_db::init_database_memory().await.unwrap(); + let user_repo = SqliteUserRepository::new(db.pool().clone()); + let mut ids = Vec::new(); + for username in usernames { + ids.push(user_repo.create_user(username, "hash").await.unwrap().id); + } + (Arc::new(SqliteSkillRepository::new(db.pool().clone())), ids) +} + #[tokio::test] async fn resolve_skill_paths_includes_cron_skills_dir() { let base = unique_temp_dir("cron-paths"); @@ -21,25 +44,70 @@ async fn resolve_skill_paths_includes_cron_skills_dir() { } #[tokio::test] -async fn materialize_resolves_saved_cron_skill() { - let base = unique_temp_dir("cron-materialize"); - let skill_dir = base.join("cron").join("skills").join("cron-job-123"); - std::fs::create_dir_all(&skill_dir).unwrap(); - std::fs::write( - skill_dir.join("SKILL.md"), - "---\nname: cron-job-123\ndescription: Saved cron skill\n---\nUse the saved steps.", - ) - .unwrap(); +async fn materialize_does_not_probe_cron_dirs_without_a_repo_row() { + let base = unique_temp_dir("cron-materialize-unauthorized"); + write_cron_skill(&base, "cron-job-123", "cron-job-123"); let paths = resolve_skill_paths(&base, &base); + // Per-job cron skill content must not resolve by name alone: without an + // owned repo row there is no proof the requester owns the job. let resolved = skill_service::materialize_skills_for_agent(&paths, "conv-1", &["cron-job-123".to_owned()]) .await .unwrap(); + assert!(resolved.is_empty(), "cron dir must not be probed: {resolved:?}"); + std::fs::remove_dir_all(&base).unwrap(); +} + +#[tokio::test] +async fn cron_skill_resolves_only_for_owning_user() { + let base = unique_temp_dir("cron-materialize-owner"); + let skill_dir = write_cron_skill(&base, "cron-job-123", "number-analysis"); + let paths = resolve_skill_paths(&base, &base); + let (repo, user_ids) = make_skill_repo_with_users(&["user-a", "user-b"]).await; + let (user_a, user_b) = (user_ids[0].as_str(), user_ids[1].as_str()); + + repo.upsert_for_user( + user_a, + UpsertSkillParams { + name: "number-analysis", + description: Some("Saved cron skill"), + path: skill_dir.to_str().unwrap(), + source: "cron", + enabled: true, + }, + ) + .await + .unwrap(); + + // The owning user resolves the skill through their repo row. + let resolved = skill_service::materialize_skills_for_agent_with_repo_for_user( + &paths, + repo.as_ref(), + user_a, + "conv-a", + &["number-analysis".to_owned()], + ) + .await + .unwrap(); assert_eq!(resolved.len(), 1); - assert_eq!(resolved[0].name, "cron-job-123"); + assert_eq!(resolved[0].name, "number-analysis"); assert_eq!(resolved[0].source_path, skill_dir); - assert!(resolved[0].source_path.join("SKILL.md").exists()); + + // Another user has no visible row and no disk fallback: no content leak. + let leaked = skill_service::materialize_skills_for_agent_with_repo_for_user( + &paths, + repo.as_ref(), + user_b, + "conv-b", + &["number-analysis".to_owned()], + ) + .await + .unwrap(); + assert!( + leaked.is_empty(), + "user-b must not resolve user-a's cron skill: {leaked:?}" + ); std::fs::remove_dir_all(&base).unwrap(); } diff --git a/crates/aionui-extension/tests/registry_test.rs b/crates/aionui-extension/tests/registry_test.rs index beeff01c6..a6d95447d 100644 --- a/crates/aionui-extension/tests/registry_test.rs +++ b/crates/aionui-extension/tests/registry_test.rs @@ -170,9 +170,8 @@ async fn em1_enable_extension_broadcasts_state_changed() { registry.enable_extension("my-ext").await.unwrap(); // Verify the extension is now enabled. - let ext = registry.get_extension_by_name("my-ext").await.unwrap(); - assert!(ext.state.enabled); - assert!(ext.state.last_activated_at.is_some()); + let extensions = registry.get_loaded_extensions().await; + assert!(extensions[0].enabled); // Verify stateChanged event was broadcast. let msg = rx.recv().await.unwrap(); @@ -198,8 +197,8 @@ async fn em2_disable_extension_broadcasts_state_changed() { .unwrap(); // Verify the extension is now disabled. - let ext = registry.get_extension_by_name("my-ext").await.unwrap(); - assert!(!ext.state.enabled); + let extensions = registry.get_loaded_extensions().await; + assert!(!extensions[0].enabled); // Verify stateChanged event was broadcast. let msg = rx.recv().await.unwrap(); @@ -230,6 +229,32 @@ async fn em4_disable_nonexistent_returns_error() { assert!(result.is_err()); } +#[tokio::test] +async fn extension_enablement_is_isolated_by_user() { + let extensions = vec![make_ext("my-ext", "1.0.0", true)]; + let (registry, bus, _tmp) = seeded_registry(extensions).await; + let mut rx = bus.subscribe(); + + registry + .disable_extension_for_user("user-a", "my-ext", None) + .await + .unwrap(); + + let user_a = registry.get_loaded_extensions_for_user("user-a").await; + let user_b = registry.get_loaded_extensions_for_user("user-b").await; + assert!(!user_a[0].enabled); + assert!(user_b[0].enabled); + + let event = rx.recv().await.unwrap(); + assert_eq!(event.name, "extensions.state-changed"); + assert_eq!(event.data["user_id"], "user-a"); + assert_eq!(event.data["enabled"], false); + + registry.enable_extension_for_user("user-b", "my-ext").await.unwrap(); + let user_a = registry.get_loaded_extensions_for_user("user-a").await; + assert!(!user_a[0].enabled, "user B must not change user A's state"); +} + // --------------------------------------------------------------------------- // HR-3: Hot-reload sequence — deactivate → clear → reload → event // --------------------------------------------------------------------------- @@ -328,9 +353,11 @@ async fn sp_enable_disable_persists_through_reload() { // Flush to disk. store.flush().await.unwrap(); - // Verify persisted state (SP-1). - let loaded = aionui_extension::load_states_from_file(&state_path).unwrap(); - assert!(!loaded["ext-a"].enabled); + // Verify the default user's override, not the machine-level install state. + assert_eq!( + store.get_user_enabled("system_default_user", "ext-a").await, + Some(false) + ); // Re-initialize (simulate restart) — should restore disabled state (SP-2). let store2 = ExtensionStateStore::new(state_path); diff --git a/crates/aionui-extension/tests/skill_integration_test.rs b/crates/aionui-extension/tests/skill_integration_test.rs index 515d04ce3..a9e595e4c 100644 --- a/crates/aionui-extension/tests/skill_integration_test.rs +++ b/crates/aionui-extension/tests/skill_integration_test.rs @@ -6,12 +6,17 @@ use std::path::Path; +use aionui_db::{ + ISkillRepository, IUserRepository, SqliteSkillRepository, SqliteUserRepository, UpsertSkillParams, + init_database_memory, +}; use aionui_extension::external_paths::ExternalPathsManager; use aionui_extension::skill_service::{ NamedPath, SkillPaths, delete_assistant_rule, delete_assistant_skill, delete_skill, - detect_and_count_external_skills, export_skill_with_symlink, import_skill, list_available_skills, - read_assistant_rule, read_assistant_skill, read_builtin_rule, read_builtin_skill, read_skill_info, - resolve_skill_paths, scan_for_skills, write_assistant_rule, write_assistant_skill, + detect_and_count_external_skills, export_skill_with_symlink, import_skill, import_skill_with_repo_for_user, + list_available_skills, materialize_skills_for_agent_with_repo_for_user, read_assistant_rule, read_assistant_skill, + read_builtin_rule, read_builtin_skill, read_skill_info, resolve_skill_paths, scan_for_skills, write_assistant_rule, + write_assistant_skill, }; use tempfile::TempDir; @@ -47,6 +52,24 @@ fn create_skill(base: &Path, name: &str, desc: &str) { .unwrap(); } +fn create_skill_with_body(base: &Path, dir_name: &str, skill_name: &str, desc: &str, body: &str) { + let dir = base.join(dir_name); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write( + dir.join(SKILL_MD), + format!("---\nname: {skill_name}\ndescription: {desc}\n---\n{body}"), + ) + .unwrap(); +} + +async fn create_test_user(db: &aionui_db::Database, username: &str) -> String { + SqliteUserRepository::new(db.pool().clone()) + .create_user(username, "hash") + .await + .unwrap() + .id +} + fn create_builtin_rule(base: &Path, name: &str, content: &str) { std::fs::create_dir_all(base).unwrap(); std::fs::write(base.join(name), content).unwrap(); @@ -172,6 +195,120 @@ async fn sm4_import_skill_replaces_existing_copy() { assert!(content.contains("Updated body")); } +#[tokio::test] +async fn user_scoped_materialize_does_not_backfill_shared_legacy_skill() { + let tmp = TempDir::new().unwrap(); + let paths = make_paths(tmp.path()); + create_skill(&paths.user_skills_dir, "legacy-only", "Legacy default skill"); + + let db = init_database_memory().await.unwrap(); + let user_a = create_test_user(&db, "user_a").await; + let repo = SqliteSkillRepository::new(db.pool().clone()); + + let resolved = + materialize_skills_for_agent_with_repo_for_user(&paths, &repo, &user_a, "conv-1", &["legacy-only".to_owned()]) + .await + .unwrap(); + assert!(resolved.is_empty()); + assert!( + repo.find_by_name_for_user(&user_a, "legacy-only") + .await + .unwrap() + .is_none() + ); + + let default_resolved = materialize_skills_for_agent_with_repo_for_user( + &paths, + &repo, + "system_default_user", + "conv-1", + &["legacy-only".to_owned()], + ) + .await + .unwrap(); + assert_eq!(default_resolved.len(), 1); +} + +#[tokio::test] +async fn user_scoped_imports_with_same_name_use_distinct_storage() { + let tmp = TempDir::new().unwrap(); + let paths = make_paths(tmp.path()); + + let source_a = tmp.path().join("source-a"); + create_skill_with_body(&source_a, "shared-a", "shared", "User A skill", "body-a"); + let source_b = tmp.path().join("source-b"); + create_skill_with_body(&source_b, "shared-b", "shared", "User B skill", "body-b"); + + let db = init_database_memory().await.unwrap(); + let user_a = create_test_user(&db, "user_a").await; + let user_b = create_test_user(&db, "user_b").await; + let repo = SqliteSkillRepository::new(db.pool().clone()); + + import_skill_with_repo_for_user(&paths, &repo, &user_a, &source_a.join("shared-a")) + .await + .unwrap(); + import_skill_with_repo_for_user(&paths, &repo, &user_b, &source_b.join("shared-b")) + .await + .unwrap(); + + let row_a = repo.find_by_name_for_user(&user_a, "shared").await.unwrap().unwrap(); + let row_b = repo.find_by_name_for_user(&user_b, "shared").await.unwrap().unwrap(); + assert_ne!(row_a.path, row_b.path); + + let content_a = std::fs::read_to_string(Path::new(&row_a.path).join(SKILL_MD)).unwrap(); + let content_b = std::fs::read_to_string(Path::new(&row_b.path).join(SKILL_MD)).unwrap(); + assert!(content_a.contains("body-a")); + assert!(content_b.contains("body-b")); + + let resolved_a = + materialize_skills_for_agent_with_repo_for_user(&paths, &repo, &user_a, "conv-1", &["shared".to_owned()]) + .await + .unwrap(); + let resolved_b = + materialize_skills_for_agent_with_repo_for_user(&paths, &repo, &user_b, "conv-1", &["shared".to_owned()]) + .await + .unwrap(); + assert_eq!(resolved_a[0].source_path, Path::new(&row_a.path)); + assert_eq!(resolved_b[0].source_path, Path::new(&row_b.path)); +} + +#[tokio::test] +async fn user_skill_override_wins_over_builtin_during_materialization() { + let tmp = TempDir::new().unwrap(); + let paths = make_paths(tmp.path()); + create_skill(&paths.builtin_skills_dir, "shared", "Builtin skill"); + + let source = tmp.path().join("source-user"); + create_skill_with_body(&source, "shared-user", "shared", "User skill", "user-body"); + + let db = init_database_memory().await.unwrap(); + let user_id = create_test_user(&db, "user_override").await; + let repo = SqliteSkillRepository::new(db.pool().clone()); + let builtin_path = paths.builtin_skills_dir.join("shared"); + repo.upsert_global(UpsertSkillParams { + name: "shared", + description: Some("Builtin skill"), + path: builtin_path.to_string_lossy().as_ref(), + source: "builtin", + enabled: true, + }) + .await + .unwrap(); + + import_skill_with_repo_for_user(&paths, &repo, &user_id, &source.join("shared-user")) + .await + .unwrap(); + let user_row = repo.find_by_name_for_user(&user_id, "shared").await.unwrap().unwrap(); + + let resolved = + materialize_skills_for_agent_with_repo_for_user(&paths, &repo, &user_id, "conv-1", &["shared".to_owned()]) + .await + .unwrap(); + + assert_eq!(resolved.len(), 1); + assert_eq!(resolved[0].source_path, Path::new(&user_row.path)); +} + /// SM-5: Export skill (symlink). #[tokio::test] async fn sm5_export_skill_symlink() { diff --git a/crates/aionui-extension/tests/state_persistence_test.rs b/crates/aionui-extension/tests/state_persistence_test.rs index 95edc1967..66270359d 100644 --- a/crates/aionui-extension/tests/state_persistence_test.rs +++ b/crates/aionui-extension/tests/state_persistence_test.rs @@ -107,6 +107,42 @@ async fn sp2_state_restored_after_restart() { } } +#[tokio::test] +async fn user_enablement_is_persisted_and_isolated() { + let tmp = TempDir::new().unwrap(); + let path = tmp.path().join("extension-states.json"); + + let store = ExtensionStateStore::new(path.clone()); + store.load().await.unwrap(); + store.set_user_enabled("user-a", "ext-a", false).await; + store.set_user_enabled("user-b", "ext-a", true).await; + store.flush().await.unwrap(); + + let restored = ExtensionStateStore::new(path); + restored.load().await.unwrap(); + assert_eq!(restored.get_user_enabled("user-a", "ext-a").await, Some(false)); + assert_eq!(restored.get_user_enabled("user-b", "ext-a").await, Some(true)); + assert_eq!(restored.get_user_enabled("user-c", "ext-a").await, None); +} + +#[tokio::test] +async fn legacy_state_becomes_default_user_enablement() { + let tmp = TempDir::new().unwrap(); + let path = tmp.path().join("extension-states.json"); + let mut states = HashMap::new(); + states.insert("ext-a".to_string(), make_state("ext-a", "1.0.0", false)); + save_states_to_file(&path, &states).unwrap(); + + let store = ExtensionStateStore::new(path); + store.load().await.unwrap(); + + assert_eq!( + store.get_user_enabled("system_default_user", "ext-a").await, + Some(false) + ); + assert_eq!(store.get_user_enabled("user-a", "ext-a").await, None); +} + // --------------------------------------------------------------------------- // SP-3: No state file on first launch → all extensions default to enabled // --------------------------------------------------------------------------- diff --git a/crates/aionui-file/Cargo.toml b/crates/aionui-file/Cargo.toml index fc3c37039..b62876dcd 100644 --- a/crates/aionui-file/Cargo.toml +++ b/crates/aionui-file/Cargo.toml @@ -5,6 +5,7 @@ edition.workspace = true [dependencies] aionui-api-types.workspace = true +aionui-auth.workspace = true aionui-common.workspace = true aionui-realtime.workspace = true async-trait.workspace = true diff --git a/crates/aionui-file/src/routes.rs b/crates/aionui-file/src/routes.rs index ddc0658b8..8c3cf76fe 100644 --- a/crates/aionui-file/src/routes.rs +++ b/crates/aionui-file/src/routes.rs @@ -2,7 +2,7 @@ use axum::Router; use axum::extract::rejection::JsonRejection; -use axum::extract::{DefaultBodyLimit, Json, Multipart, Query, State}; +use axum::extract::{DefaultBodyLimit, Extension, Json, Multipart, Query, State}; use axum::routing::{get, post}; use std::path::{Path, PathBuf}; use std::sync::{Arc, OnceLock}; @@ -16,6 +16,7 @@ use aionui_api_types::{ SnapshotCompareResponse, SnapshotDiscardRequest, SnapshotInfoResponse, SnapshotStageRequest, SnapshotWorkspaceRequest, WorkspaceFlatFileResponse, WorkspaceOfficeWatchRequest, WriteFileRequest, ZipRequest, }; +use aionui_auth::CurrentUser; use aionui_common::ApiError; use aionui_common::constants::UPLOAD_MAX_SIZE; @@ -143,6 +144,7 @@ pub fn file_routes(state: FileRouterState) -> Router { .route("/api/fs/watch/stop-all", post(stop_all_watches)) .route("/api/fs/office-watch/start", post(start_office_watch)) .route("/api/fs/office-watch/stop", post(stop_office_watch)) + .route("/api/fs/office-watch/stop-all", post(stop_all_office_watches)) // E. Workspace snapshot .route("/api/fs/snapshot/init", post(snapshot_init)) .route("/api/fs/snapshot/info", post(snapshot_info)) @@ -256,6 +258,7 @@ async fn read_file_buffer( async fn write_file( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; @@ -267,7 +270,7 @@ async fn write_file( }); let ok = state .file_service - .write_file(&req.path, req.data.as_bytes(), &workspace) + .write_file_for_user(&user.id, &req.path, req.data.as_bytes(), &workspace) .await?; Ok(Json(ApiResponse::ok(ok))) } @@ -286,6 +289,7 @@ async fn copy_files( async fn remove_entry( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; @@ -295,7 +299,10 @@ async fn remove_entry( .map(|p| p.to_string_lossy().into_owned()) .unwrap_or_default() }); - state.file_service.remove_entry(&req.path, &workspace).await?; + state + .file_service + .remove_entry_for_user(&user.id, &req.path, &workspace) + .await?; Ok(Json(ApiResponse::success())) } @@ -478,44 +485,71 @@ async fn cancel_zip( async fn start_watch( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - state.watch_service.start_watch(&req.file_path).await?; + state + .watch_service + .start_watch_for_user(&user.id, &req.file_path) + .await?; Ok(Json(ApiResponse::success())) } async fn stop_watch( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - state.watch_service.stop_watch(&req.file_path).await?; + state + .watch_service + .stop_watch_for_user(&user.id, &req.file_path) + .await?; Ok(Json(ApiResponse::success())) } -async fn stop_all_watches(State(state): State) -> Result>, ApiError> { - state.watch_service.stop_all_watches().await?; +async fn stop_all_watches( + State(state): State, + Extension(user): Extension, +) -> Result>, ApiError> { + state.watch_service.stop_all_watches_for_user(&user.id).await?; Ok(Json(ApiResponse::success())) } async fn start_office_watch( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; let allowed_roots: Vec<&Path> = state.allowed_roots.iter().map(std::path::PathBuf::as_path).collect(); crate::path_safety::validate_path_with_extra_root(&req.workspace, &allowed_roots, Some(Path::new(&req.workspace)))?; - state.watch_service.start_office_watch(&req.workspace).await?; + state + .watch_service + .start_office_watch_for_user(&user.id, &req.workspace) + .await?; Ok(Json(ApiResponse::success())) } async fn stop_office_watch( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - state.watch_service.stop_office_watch(&req.workspace).await?; + state + .watch_service + .stop_office_watch_for_user(&user.id, &req.workspace) + .await?; + Ok(Json(ApiResponse::success())) +} + +async fn stop_all_office_watches( + State(state): State, + Extension(user): Extension, +) -> Result>, ApiError> { + state.watch_service.stop_all_office_watches_for_user(&user.id).await?; Ok(Json(ApiResponse::success())) } diff --git a/crates/aionui-file/src/service.rs b/crates/aionui-file/src/service.rs index 421fc4539..971b161bc 100644 --- a/crates/aionui-file/src/service.rs +++ b/crates/aionui-file/src/service.rs @@ -691,6 +691,17 @@ impl crate::traits::IFileService for FileService { } async fn write_file(&self, path: &str, data: &[u8], workspace: &str) -> Result { + self.write_file_for_user("system_default_user", path, data, workspace) + .await + } + + async fn write_file_for_user( + &self, + user_id: &str, + path: &str, + data: &[u8], + workspace: &str, + ) -> Result { if has_traversal(path) { return Err(FileError::BadRequest(format!( "path '{}' contains invalid traversal patterns", @@ -724,7 +735,8 @@ impl crate::traits::IFileService for FileService { relative_path, operation: ContentUpdateOperation::Write, }; - let payload = serde_json::to_value(&event).unwrap_or_default(); + let mut payload = serde_json::to_value(&event).unwrap_or_default(); + payload["user_id"] = serde_json::Value::String(user_id.to_owned()); let msg = WebSocketMessage::new("fileStream.contentUpdate", payload); self.broadcaster.broadcast(msg); @@ -800,6 +812,10 @@ impl crate::traits::IFileService for FileService { } async fn remove_entry(&self, path: &str, workspace: &str) -> Result<(), FileError> { + self.remove_entry_for_user("system_default_user", path, workspace).await + } + + async fn remove_entry_for_user(&self, user_id: &str, path: &str, workspace: &str) -> Result<(), FileError> { if has_traversal(path) { return Err(FileError::BadRequest(format!( "path '{}' contains invalid traversal patterns", @@ -831,7 +847,8 @@ impl crate::traits::IFileService for FileService { relative_path, operation: ContentUpdateOperation::Delete, }; - let payload = serde_json::to_value(&event).unwrap_or_default(); + let mut payload = serde_json::to_value(&event).unwrap_or_default(); + payload["user_id"] = serde_json::Value::String(user_id.to_owned()); let msg = WebSocketMessage::new("fileStream.contentUpdate", payload); self.broadcaster.broadcast(msg); diff --git a/crates/aionui-file/src/traits.rs b/crates/aionui-file/src/traits.rs index 5f1ccf683..7b16f2e85 100644 --- a/crates/aionui-file/src/traits.rs +++ b/crates/aionui-file/src/traits.rs @@ -52,6 +52,20 @@ pub trait IFileService: Send + Sync { /// `fileStream.contentUpdate` event with `operation = write`. async fn write_file(&self, path: &str, data: &[u8], workspace: &str) -> Result; + /// User-scoped variant of [`write_file`](Self::write_file), used by + /// authenticated routes so WebSocket events are delivered only to the + /// initiating user. + async fn write_file_for_user( + &self, + user_id: &str, + path: &str, + data: &[u8], + workspace: &str, + ) -> Result { + let _ = user_id; + self.write_file(path, data, workspace).await + } + // -- File management -- /// Copy files into `workspace`, preserving directory structure relative to @@ -67,6 +81,12 @@ pub trait IFileService: Send + Sync { /// `fileStream.contentUpdate` event with `operation = delete`. async fn remove_entry(&self, path: &str, workspace: &str) -> Result<(), FileError>; + /// User-scoped variant of [`remove_entry`](Self::remove_entry). + async fn remove_entry_for_user(&self, user_id: &str, path: &str, workspace: &str) -> Result<(), FileError> { + let _ = user_id; + self.remove_entry(path, workspace).await + } + /// Rename a file or directory. Returns the new absolute path. async fn rename_entry(&self, path: &str, new_name: &str) -> Result; @@ -149,19 +169,55 @@ pub trait IFileWatchService: Send + Sync { /// Emits `fileWatch.fileChanged` events on the broadcast channel. async fn start_watch(&self, file_path: &str) -> Result<(), FileError>; + /// Start watching a single file for one WebUI user. + async fn start_watch_for_user(&self, user_id: &str, file_path: &str) -> Result<(), FileError> { + let _ = user_id; + self.start_watch(file_path).await + } + /// Stop watching a previously registered file. async fn stop_watch(&self, file_path: &str) -> Result<(), FileError>; + /// Stop watching a file for one WebUI user. + async fn stop_watch_for_user(&self, user_id: &str, file_path: &str) -> Result<(), FileError> { + let _ = user_id; + self.stop_watch(file_path).await + } + /// Stop all active file watches. async fn stop_all_watches(&self) -> Result<(), FileError>; + /// Stop all active file watches for one WebUI user. + async fn stop_all_watches_for_user(&self, user_id: &str) -> Result<(), FileError> { + let _ = user_id; + self.stop_all_watches().await + } + /// Start watching a workspace directory for new Office files /// (.pptx, .docx, .xlsx). /// Emits `workspaceOfficeWatch.fileAdded` events. async fn start_office_watch(&self, workspace: &str) -> Result<(), FileError>; + /// Start watching a workspace for one WebUI user. + async fn start_office_watch_for_user(&self, user_id: &str, workspace: &str) -> Result<(), FileError> { + let _ = user_id; + self.start_office_watch(workspace).await + } + /// Stop watching a workspace directory for Office files. async fn stop_office_watch(&self, workspace: &str) -> Result<(), FileError>; + + /// Stop watching a workspace for one WebUI user. + async fn stop_office_watch_for_user(&self, user_id: &str, workspace: &str) -> Result<(), FileError> { + let _ = user_id; + self.stop_office_watch(workspace).await + } + + /// Stop all active Office workspace watches for one WebUI user. + async fn stop_all_office_watches_for_user(&self, user_id: &str) -> Result<(), FileError> { + let _ = user_id; + Ok(()) + } } /// Git-based workspace snapshot system for tracking file changes. diff --git a/crates/aionui-file/src/watch_service.rs b/crates/aionui-file/src/watch_service.rs index 0bfa43841..c6f8850f0 100644 --- a/crates/aionui-file/src/watch_service.rs +++ b/crates/aionui-file/src/watch_service.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap}; use std::path::Path; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -73,10 +73,12 @@ pub struct FileWatchService { broadcaster: Arc, /// Shared watcher for all single-file watches. file_watcher: Mutex, - /// Set of canonical paths being watched (shared with the event handler). - watched_files: Arc>, + /// Canonical watched paths and the users subscribed to each path. + watched_files: Arc>>, /// Per-workspace Office watchers, keyed by canonical workspace path. office_watchers: Mutex>, + /// Canonical Office workspaces and the users subscribed to each workspace. + office_watch_users: Arc>>, /// Debounce timestamps shared with watcher callbacks. debounce: Arc>, } @@ -84,7 +86,8 @@ pub struct FileWatchService { impl FileWatchService { /// Create a new watch service backed by the platform's recommended watcher. pub fn new(broadcaster: Arc) -> Result { - let watched_files: Arc> = Arc::new(DashMap::new()); + let watched_files: Arc>> = Arc::new(DashMap::new()); + let office_watch_users: Arc>> = Arc::new(DashMap::new()); let debounce: Arc> = Arc::new(DashMap::new()); let bc = broadcaster.clone(); @@ -107,18 +110,21 @@ impl FileWatchService { for path in &event.paths { let path_str = path.to_string_lossy().into_owned(); - if !wf.contains_key(&path_str) { + let Some(users) = wf.get(&path_str).map(|entry| entry.clone()) else { continue; - } + }; if !should_emit(&db, &path_str) { continue; } - let payload = FileWatchEvent { - file_path: path_str, - event_type: event_type.to_owned(), - }; - let json = serde_json::to_value(&payload).unwrap_or_default(); - bc.broadcast(WebSocketMessage::new("fileWatch.fileChanged", json)); + for user_id in users { + let payload = FileWatchEvent { + file_path: path_str.clone(), + event_type: event_type.to_owned(), + }; + let mut json = serde_json::to_value(&payload).unwrap_or_default(); + json["user_id"] = serde_json::Value::String(user_id); + bc.broadcast(WebSocketMessage::new("fileWatch.fileChanged", json)); + } } }) .map_err(|e| FileError::Internal(format!("failed to create file watcher: {e}")))?; @@ -128,6 +134,7 @@ impl FileWatchService { file_watcher: Mutex::new(file_watcher), watched_files, office_watchers: Mutex::new(HashMap::new()), + office_watch_users, debounce, }) } @@ -136,12 +143,16 @@ impl FileWatchService { #[async_trait::async_trait] impl crate::traits::IFileWatchService for FileWatchService { async fn start_watch(&self, file_path: &str) -> Result<(), FileError> { + self.start_watch_for_user("system_default_user", file_path).await + } + + async fn start_watch_for_user(&self, user_id: &str, file_path: &str) -> Result<(), FileError> { let canonical = std::fs::canonicalize(file_path) .map_err(|e| FileError::NotFound(format!("cannot resolve path {file_path}: {e}")))?; let key = canonical.to_string_lossy().into_owned(); - // Idempotent: already watching → no-op. - if self.watched_files.contains_key(&key) { + if let Some(mut users) = self.watched_files.get_mut(&key) { + users.insert(user_id.to_owned()); return Ok(()); } @@ -152,25 +163,36 @@ impl crate::traits::IFileWatchService for FileWatchService { watcher .watch(&canonical, RecursiveMode::NonRecursive) .map_err(|e| FileError::Internal(format!("failed to watch {file_path}: {e}")))?; - self.watched_files.insert(key, ()); + self.watched_files.insert(key, BTreeSet::from([user_id.to_owned()])); Ok(()) } async fn stop_watch(&self, file_path: &str) -> Result<(), FileError> { + self.stop_watch_for_user("system_default_user", file_path).await + } + + async fn stop_watch_for_user(&self, user_id: &str, file_path: &str) -> Result<(), FileError> { let canonical = std::fs::canonicalize(file_path).unwrap_or_else(|_| file_path.into()); let key = canonical.to_string_lossy().into_owned(); - if self.watched_files.remove(&key).is_none() { - return Ok(()); + let should_unwatch = if let Some(mut users) = self.watched_files.get_mut(&key) { + users.remove(user_id); + users.is_empty() + } else { + false + }; + + if should_unwatch { + self.watched_files.remove(&key); + let mut watcher = self + .file_watcher + .lock() + .map_err(|e| FileError::Internal(format!("file watcher lock poisoned: {e}")))?; + // Ignore unwatch errors because the file may have been deleted. + let _ = watcher.unwatch(&canonical); + self.debounce.remove(&key); } - let mut watcher = self - .file_watcher - .lock() - .map_err(|e| FileError::Internal(format!("file watcher lock poisoned: {e}")))?; - // Ignore unwatch errors — the file may have been deleted. - let _ = watcher.unwatch(&canonical); - self.debounce.remove(&key); Ok(()) } @@ -190,7 +212,19 @@ impl crate::traits::IFileWatchService for FileWatchService { Ok(()) } + async fn stop_all_watches_for_user(&self, user_id: &str) -> Result<(), FileError> { + let keys: Vec = self.watched_files.iter().map(|entry| entry.key().clone()).collect(); + for key in keys { + self.stop_watch_for_user(user_id, &key).await?; + } + Ok(()) + } + async fn start_office_watch(&self, workspace: &str) -> Result<(), FileError> { + self.start_office_watch_for_user("system_default_user", workspace).await + } + + async fn start_office_watch_for_user(&self, user_id: &str, workspace: &str) -> Result<(), FileError> { let canonical = std::fs::canonicalize(workspace) .map_err(|e| FileError::NotFound(format!("cannot resolve workspace {workspace}: {e}")))?; let key = canonical.to_string_lossy().into_owned(); @@ -201,6 +235,10 @@ impl crate::traits::IFileWatchService for FileWatchService { .lock() .map_err(|e| FileError::Internal(format!("office watcher lock poisoned: {e}")))?; if watchers.contains_key(&key) { + self.office_watch_users + .entry(key) + .or_default() + .insert(user_id.to_owned()); return Ok(()); } } @@ -208,6 +246,7 @@ impl crate::traits::IFileWatchService for FileWatchService { let bc = self.broadcaster.clone(); let db = self.debounce.clone(); let ws = key.clone(); + let office_users = self.office_watch_users.clone(); let mut watcher = notify::recommended_watcher(move |res: Result| { let event = match res { @@ -228,15 +267,21 @@ impl crate::traits::IFileWatchService for FileWatchService { } let path_str = path.to_string_lossy().into_owned(); let debounce_key = format!("office:{path_str}"); + let Some(users) = office_users.get(&ws).map(|entry| entry.clone()) else { + continue; + }; if !should_emit(&db, &debounce_key) { continue; } - let payload = OfficeFileAddedEvent { - file_path: path_str, - workspace: ws.clone(), - }; - let json = serde_json::to_value(&payload).unwrap_or_default(); - bc.broadcast(WebSocketMessage::new("workspaceOfficeWatch.fileAdded", json)); + for user_id in users { + let payload = OfficeFileAddedEvent { + file_path: path_str.clone(), + workspace: ws.clone(), + }; + let mut json = serde_json::to_value(&payload).unwrap_or_default(); + json["user_id"] = serde_json::Value::String(user_id); + bc.broadcast(WebSocketMessage::new("workspaceOfficeWatch.fileAdded", json)); + } } }) .map_err(|e| FileError::Internal(format!("failed to create office watcher: {e}")))?; @@ -249,14 +294,32 @@ impl crate::traits::IFileWatchService for FileWatchService { .office_watchers .lock() .map_err(|e| FileError::Internal(format!("office watcher lock poisoned: {e}")))?; - watchers.insert(key, watcher); + watchers.insert(key.clone(), watcher); + self.office_watch_users + .insert(key, BTreeSet::from([user_id.to_owned()])); Ok(()) } async fn stop_office_watch(&self, workspace: &str) -> Result<(), FileError> { + self.stop_office_watch_for_user("system_default_user", workspace).await + } + + async fn stop_office_watch_for_user(&self, user_id: &str, workspace: &str) -> Result<(), FileError> { let canonical = std::fs::canonicalize(workspace).unwrap_or_else(|_| workspace.into()); let key = canonical.to_string_lossy().into_owned(); + let should_stop = if let Some(mut users) = self.office_watch_users.get_mut(&key) { + users.remove(user_id); + users.is_empty() + } else { + false + }; + + if !should_stop { + return Ok(()); + } + + self.office_watch_users.remove(&key); let mut watchers = self .office_watchers .lock() @@ -265,6 +328,18 @@ impl crate::traits::IFileWatchService for FileWatchService { watchers.remove(&key); Ok(()) } + + async fn stop_all_office_watches_for_user(&self, user_id: &str) -> Result<(), FileError> { + let keys: Vec = self + .office_watch_users + .iter() + .map(|entry| entry.key().clone()) + .collect(); + for key in keys { + self.stop_office_watch_for_user(user_id, &key).await?; + } + Ok(()) + } } // --------------------------------------------------------------------------- diff --git a/crates/aionui-file/tests/file_watching.rs b/crates/aionui-file/tests/file_watching.rs index 0a2e3380f..190e1891b 100644 --- a/crates/aionui-file/tests/file_watching.rs +++ b/crates/aionui-file/tests/file_watching.rs @@ -78,6 +78,7 @@ async fn start_watch_and_detect_change() { let ev = events.iter().find(|e| e.name == "fileWatch.fileChanged").unwrap(); assert!(ev.data["file_path"].as_str().is_some()); assert!(ev.data["event_type"].as_str().is_some()); + assert_eq!(ev.data["user_id"].as_str(), Some("system_default_user")); } #[tokio::test] @@ -253,6 +254,57 @@ async fn stop_office_watch_stops_events() { ); } +#[tokio::test] +async fn stop_all_office_watches_for_user_keeps_other_user_subscriptions() { + let dir = tempfile::tempdir().unwrap(); + let (svc, recorder) = make_service(); + let ws = dir.path().to_str().unwrap(); + + svc.start_office_watch_for_user("user_a", ws).await.unwrap(); + svc.start_office_watch_for_user("user_b", ws).await.unwrap(); + settle().await; + + svc.stop_all_office_watches_for_user("user_a").await.unwrap(); + recorder.take_events(); + + std::fs::write(dir.path().join("owned_by_b.docx"), "data").unwrap(); + settle().await; + + let events = recorder.take_events(); + let office_events: Vec<_> = events + .iter() + .filter(|e| e.name == "workspaceOfficeWatch.fileAdded") + .collect(); + assert_eq!(office_events.len(), 1, "expected only user_b event, got: {events:?}"); + assert_eq!(office_events[0].data["user_id"].as_str(), Some("user_b")); +} + +#[tokio::test] +async fn stop_all_office_watches_for_user_stops_owned_subscriptions() { + let dir = tempfile::tempdir().unwrap(); + let (svc, recorder) = make_service(); + let ws = dir.path().to_str().unwrap(); + + svc.start_office_watch_for_user("user_a", ws).await.unwrap(); + settle().await; + + svc.stop_all_office_watches_for_user("user_a").await.unwrap(); + recorder.take_events(); + + std::fs::write(dir.path().join("after_stop_all.docx"), "data").unwrap(); + settle().await; + + let events = recorder.take_events(); + let office_events: Vec<_> = events + .iter() + .filter(|e| e.name == "workspaceOfficeWatch.fileAdded") + .collect(); + assert!( + office_events.is_empty(), + "expected no office events after user stop-all, got: {events:?}" + ); +} + #[tokio::test] async fn idempotent_office_watch() { let dir = tempfile::tempdir().unwrap(); @@ -286,4 +338,5 @@ async fn office_watch_event_has_correct_fields() { data["workspace"].as_str().is_some(), "workspace should be present: {data:?}" ); + assert_eq!(data["user_id"].as_str(), Some("system_default_user")); } diff --git a/crates/aionui-mcp/Cargo.toml b/crates/aionui-mcp/Cargo.toml index 32edd6030..b9381f165 100644 --- a/crates/aionui-mcp/Cargo.toml +++ b/crates/aionui-mcp/Cargo.toml @@ -8,6 +8,7 @@ aionui-runtime.workspace = true aionui-common.workspace = true aionui-db.workspace = true aionui-api-types.workspace = true +aionui-auth.workspace = true aionui-realtime.workspace = true aionui-system.workspace = true async-trait.workspace = true diff --git a/crates/aionui-mcp/src/adapter.rs b/crates/aionui-mcp/src/adapter.rs index 23e37af75..a28e108e8 100644 --- a/crates/aionui-mcp/src/adapter.rs +++ b/crates/aionui-mcp/src/adapter.rs @@ -59,7 +59,7 @@ pub trait McpAgentAdapter: Send + Sync { /// /// Returns an empty vec if the CLI is installed but has no MCP servers. /// Returns `McpError::AgentNotInstalled` if the CLI is not available. - async fn detect_existing(&self) -> Result, McpError>; + async fn detect_existing(&self, _user_id: &str) -> Result, McpError>; /// Installs (or updates) an MCP server configuration in this Agent CLI. /// @@ -80,6 +80,8 @@ mod tests { use std::collections::HashMap; use std::sync::{Arc, Mutex}; + const TEST_USER_ID: &str = "user-1"; + /// In-memory mock adapter for testing the trait interface. struct MockAdapter { source: McpSource, @@ -107,7 +109,7 @@ mod tests { Ok(self.installed) } - async fn detect_existing(&self) -> Result, McpError> { + async fn detect_existing(&self, _user_id: &str) -> Result, McpError> { if !self.installed { return Err(McpError::AgentNotInstalled(format!("{:?}", self.source))); } @@ -166,7 +168,7 @@ mod tests { adapter.install_server("test-mcp", &transport).await.unwrap(); - let detected = adapter.detect_existing().await.unwrap(); + let detected = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert_eq!(detected.len(), 1); assert_eq!(detected[0].name, "test-mcp"); assert_eq!(detected[0].transport, transport); @@ -189,7 +191,7 @@ mod tests { adapter.install_server("test-mcp", &t1).await.unwrap(); adapter.install_server("test-mcp", &t2).await.unwrap(); - let detected = adapter.detect_existing().await.unwrap(); + let detected = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert_eq!(detected.len(), 1); match &detected[0].transport { McpServerTransport::Stdio { command, .. } => assert_eq!(command, "new"), @@ -207,7 +209,7 @@ mod tests { adapter.install_server("srv", &transport).await.unwrap(); adapter.remove_server("srv").await.unwrap(); - let detected = adapter.detect_existing().await.unwrap(); + let detected = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert!(detected.is_empty()); } @@ -221,7 +223,7 @@ mod tests { #[tokio::test] async fn not_installed_detect_fails() { let adapter = MockAdapter::new(McpSource::Qwen, false); - let result = adapter.detect_existing().await; + let result = adapter.detect_existing(TEST_USER_ID).await; assert!(matches!(result, Err(McpError::AgentNotInstalled(_)))); } diff --git a/crates/aionui-mcp/src/adapters/aionrs.rs b/crates/aionui-mcp/src/adapters/aionrs.rs index 094c10ba5..7aead3023 100644 --- a/crates/aionui-mcp/src/adapters/aionrs.rs +++ b/crates/aionui-mcp/src/adapters/aionrs.rs @@ -45,7 +45,7 @@ impl McpAgentAdapter for AionrsAdapter { is_cli_installed(CLI_NAME).await } - async fn detect_existing(&self) -> Result, McpError> { + async fn detect_existing(&self, _user_id: &str) -> Result, McpError> { if !self.is_installed().await? { return Err(McpError::AgentNotInstalled(CLI_NAME.into())); } diff --git a/crates/aionui-mcp/src/adapters/aionui.rs b/crates/aionui-mcp/src/adapters/aionui.rs index 379145d04..c5f191e21 100644 --- a/crates/aionui-mcp/src/adapters/aionui.rs +++ b/crates/aionui-mcp/src/adapters/aionui.rs @@ -39,8 +39,8 @@ impl McpAgentAdapter for AionuiAdapter { Ok(true) } - async fn detect_existing(&self) -> Result, McpError> { - let rows = self.repo.list().await?; + async fn detect_existing(&self, user_id: &str) -> Result, McpError> { + let rows = self.repo.list(user_id).await?; let mut servers = Vec::new(); for row in rows { @@ -82,6 +82,8 @@ mod tests { use crate::types::McpServerTransport; use aionui_db::models::McpServerRow; + const TEST_USER_ID: &str = "user-1"; + /// In-memory mock repository for testing. struct MockRepo { servers: Vec, @@ -95,16 +97,29 @@ mod tests { #[async_trait::async_trait] impl IMcpServerRepository for MockRepo { - async fn list(&self) -> Result, aionui_db::DbError> { - Ok(self.servers.clone()) + async fn list(&self, user_id: &str) -> Result, aionui_db::DbError> { + Ok(self + .servers + .iter() + .filter(|server| server.user_id == user_id) + .cloned() + .collect()) } - async fn find_by_id(&self, id: &str) -> Result, aionui_db::DbError> { - Ok(self.servers.iter().find(|s| s.id == id).cloned()) + async fn find_by_id(&self, user_id: &str, id: &str) -> Result, aionui_db::DbError> { + Ok(self + .servers + .iter() + .find(|s| s.user_id == user_id && s.id == id) + .cloned()) } - async fn find_by_name(&self, name: &str) -> Result, aionui_db::DbError> { - Ok(self.servers.iter().find(|s| s.name == name).cloned()) + async fn find_by_name(&self, user_id: &str, name: &str) -> Result, aionui_db::DbError> { + Ok(self + .servers + .iter() + .find(|s| s.user_id == user_id && s.name == name) + .cloned()) } async fn create( @@ -116,18 +131,20 @@ mod tests { async fn update( &self, + _user_id: &str, _id: &str, _params: aionui_db::UpdateMcpServerParams<'_>, ) -> Result { unimplemented!("not needed for adapter tests") } - async fn delete(&self, _id: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), aionui_db::DbError> { unimplemented!("not needed for adapter tests") } async fn batch_upsert( &self, + _user_id: &str, _servers: &[aionui_db::CreateMcpServerParams<'_>], ) -> Result, aionui_db::DbError> { unimplemented!("not needed for adapter tests") @@ -135,6 +152,7 @@ mod tests { async fn update_status( &self, + _user_id: &str, _id: &str, _status: &str, _last_connected: Option, @@ -142,7 +160,12 @@ mod tests { unimplemented!("not needed for adapter tests") } - async fn update_tools(&self, _id: &str, _tools: Option<&str>) -> Result<(), aionui_db::DbError> { + async fn update_tools( + &self, + _user_id: &str, + _id: &str, + _tools: Option<&str>, + ) -> Result<(), aionui_db::DbError> { unimplemented!("not needed for adapter tests") } } @@ -150,6 +173,7 @@ mod tests { fn make_row(name: &str, transport_type: &str, transport_config: &str) -> McpServerRow { McpServerRow { id: format!("mcp_{name}"), + user_id: "user-1".into(), name: name.to_owned(), description: None, enabled: true, @@ -189,7 +213,7 @@ mod tests { let repo = Arc::new(MockRepo::new(rows)); let adapter = AionuiAdapter::new(repo); - let servers = adapter.detect_existing().await.unwrap(); + let servers = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert_eq!(servers.len(), 2); assert_eq!(servers[0].name, "srv-a"); assert_eq!(servers[1].name, "srv-b"); @@ -202,7 +226,7 @@ mod tests { let repo = Arc::new(MockRepo::new(vec![])); let adapter = AionuiAdapter::new(repo); - let servers = adapter.detect_existing().await.unwrap(); + let servers = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert!(servers.is_empty()); } diff --git a/crates/aionui-mcp/src/adapters/claude.rs b/crates/aionui-mcp/src/adapters/claude.rs index 4db572d34..3a794df98 100644 --- a/crates/aionui-mcp/src/adapters/claude.rs +++ b/crates/aionui-mcp/src/adapters/claude.rs @@ -38,7 +38,7 @@ impl McpAgentAdapter for ClaudeAdapter { is_cli_installed(CLI_NAME).await } - async fn detect_existing(&self) -> Result, McpError> { + async fn detect_existing(&self, _user_id: &str) -> Result, McpError> { if !self.is_installed().await? { return Err(McpError::AgentNotInstalled(CLI_NAME.into())); } diff --git a/crates/aionui-mcp/src/adapters/codebuddy.rs b/crates/aionui-mcp/src/adapters/codebuddy.rs index 25a474d48..e8069b2df 100644 --- a/crates/aionui-mcp/src/adapters/codebuddy.rs +++ b/crates/aionui-mcp/src/adapters/codebuddy.rs @@ -41,7 +41,7 @@ impl McpAgentAdapter for CodeBuddyAdapter { is_cli_installed(CLI_NAME).await } - async fn detect_existing(&self) -> Result, McpError> { + async fn detect_existing(&self, _user_id: &str) -> Result, McpError> { if !self.is_installed().await? { return Err(McpError::AgentNotInstalled(CLI_NAME.into())); } diff --git a/crates/aionui-mcp/src/adapters/codex.rs b/crates/aionui-mcp/src/adapters/codex.rs index 70312e225..d5a2f1ce5 100644 --- a/crates/aionui-mcp/src/adapters/codex.rs +++ b/crates/aionui-mcp/src/adapters/codex.rs @@ -32,7 +32,7 @@ impl McpAgentAdapter for CodexAdapter { is_cli_installed(CLI_NAME).await } - async fn detect_existing(&self) -> Result, McpError> { + async fn detect_existing(&self, _user_id: &str) -> Result, McpError> { if !self.is_installed().await? { return Err(McpError::AgentNotInstalled(CLI_NAME.into())); } diff --git a/crates/aionui-mcp/src/adapters/gemini.rs b/crates/aionui-mcp/src/adapters/gemini.rs index fd7042ecf..fb5a8000c 100644 --- a/crates/aionui-mcp/src/adapters/gemini.rs +++ b/crates/aionui-mcp/src/adapters/gemini.rs @@ -31,7 +31,7 @@ impl McpAgentAdapter for GeminiAdapter { is_cli_installed(CLI_NAME).await } - async fn detect_existing(&self) -> Result, McpError> { + async fn detect_existing(&self, _user_id: &str) -> Result, McpError> { if !self.is_installed().await? { return Err(McpError::AgentNotInstalled(CLI_NAME.into())); } diff --git a/crates/aionui-mcp/src/adapters/opencode.rs b/crates/aionui-mcp/src/adapters/opencode.rs index 8d21a8d98..1f3de3cda 100644 --- a/crates/aionui-mcp/src/adapters/opencode.rs +++ b/crates/aionui-mcp/src/adapters/opencode.rs @@ -48,7 +48,7 @@ impl McpAgentAdapter for OpencodeAdapter { Ok(config_dir().is_some_and(|d| d.exists())) } - async fn detect_existing(&self) -> Result, McpError> { + async fn detect_existing(&self, _user_id: &str) -> Result, McpError> { let path = config_file_path().ok_or_else(|| McpError::AgentNotInstalled("opencode".into()))?; if !path.exists() { diff --git a/crates/aionui-mcp/src/adapters/qwen.rs b/crates/aionui-mcp/src/adapters/qwen.rs index 686bf3d25..5c5353171 100644 --- a/crates/aionui-mcp/src/adapters/qwen.rs +++ b/crates/aionui-mcp/src/adapters/qwen.rs @@ -38,7 +38,7 @@ impl McpAgentAdapter for QwenAdapter { is_cli_installed(CLI_NAME).await } - async fn detect_existing(&self) -> Result, McpError> { + async fn detect_existing(&self, _user_id: &str) -> Result, McpError> { if !self.is_installed().await? { return Err(McpError::AgentNotInstalled(CLI_NAME.into())); } diff --git a/crates/aionui-mcp/src/connection_test/mod.rs b/crates/aionui-mcp/src/connection_test/mod.rs index 78952c39b..8a7c76144 100644 --- a/crates/aionui-mcp/src/connection_test/mod.rs +++ b/crates/aionui-mcp/src/connection_test/mod.rs @@ -68,13 +68,15 @@ impl McpConnectionTestService { /// Dispatches to the appropriate transport handler. Always returns /// a result (never errors) -- failures are encoded in the struct. pub async fn test_connection(&self, name: &str, transport: &McpServerTransport) -> McpConnectionTestResult { - self.test_connection_with_runtime_scope(name, transport, None).await + self.test_connection_with_runtime_scope(name, transport, None, None) + .await } pub async fn test_connection_with_runtime_scope( &self, name: &str, transport: &McpServerTransport, + user_id: Option<&str>, runtime_scope_id: Option<&str>, ) -> McpConnectionTestResult { debug!(name, ?transport, "starting MCP connection test"); @@ -83,7 +85,7 @@ impl McpConnectionTestService { log_mcp_transport_start(mcp_server_id, transport_type); let result = match transport { McpServerTransport::Stdio { command, args, env } => { - self.test_stdio(command, args, env, runtime_scope_id).await + self.test_stdio(command, args, env, user_id, runtime_scope_id).await } McpServerTransport::Http { url, headers } => self.test_http(url, headers).await, McpServerTransport::Sse { url, headers } => self.test_sse(url, headers).await, @@ -99,9 +101,11 @@ impl McpConnectionTestService { command: &str, args: &[String], env: &HashMap, + user_id: Option<&str>, runtime_scope_id: Option<&str>, ) -> McpConnectionTestResult { - self.test_stdio_inner(command, args, env, runtime_scope_id).await + self.test_stdio_inner(command, args, env, user_id, runtime_scope_id) + .await } async fn test_stdio_inner( @@ -109,9 +113,11 @@ impl McpConnectionTestService { command: &str, args: &[String], env: &HashMap, + user_id: Option<&str>, runtime_scope_id: Option<&str>, ) -> McpConnectionTestResult { - let reporter = runtime_scope_id.map(|scope_id| self.runtime_reporter(scope_id.to_owned())); + let reporter = + runtime_scope_id.map(|scope_id| self.runtime_reporter(user_id.map(str::to_owned), scope_id.to_owned())); let mut cmd = match probe_runtime_command(command) { RuntimeCommandProbe::NodeTool { .. } => { let resolved = match ensure_runtime_command_with_reporter(command, reporter.as_deref()).await { @@ -407,10 +413,11 @@ impl McpConnectionTestService { } impl McpConnectionTestService { - fn runtime_reporter(&self, scope_id: String) -> Arc { + fn runtime_reporter(&self, user_id: Option, scope_id: String) -> Arc { let broadcaster = self.broadcaster.clone(); Arc::new(move |update: NodeRuntimeProgress| { let payload = RuntimeStatusPayload { + user_id: user_id.clone(), resource: RuntimeResourceKind::Node, resource_id: None, scope: RuntimeStatusScope { @@ -532,7 +539,10 @@ struct HttpMcpResponse { #[cfg(test)] mod tests { use super::*; + use aionui_api_types::WebSocketMessage; use aionui_realtime::BroadcastEventBus; + use aionui_realtime::EventBroadcaster; + use aionui_runtime::{NodeRuntimeProgress, NodeRuntimeProgressPhase}; use std::io::Write; use std::sync::Mutex; use tracing::Level; @@ -552,6 +562,28 @@ mod tests { } } + struct RecordingBroadcaster { + events: Mutex>>, + } + + impl RecordingBroadcaster { + fn new() -> Self { + Self { + events: Mutex::new(Vec::new()), + } + } + + fn events(&self) -> Vec> { + self.events.lock().unwrap().clone() + } + } + + impl EventBroadcaster for RecordingBroadcaster { + fn broadcast(&self, event: WebSocketMessage) { + self.events.lock().unwrap().push(event); + } + } + fn capture_logs(max_level: Level, f: impl FnOnce()) -> String { let buffer = Arc::new(Mutex::new(Vec::::new())); let make_writer = { @@ -581,6 +613,26 @@ mod tests { assert_eq!(svc.timeout, Duration::from_secs(5)); } + #[test] + fn runtime_reporter_scopes_event_to_user() { + let broadcaster = Arc::new(RecordingBroadcaster::new()); + let svc = McpConnectionTestService::new(reqwest::Client::new(), broadcaster.clone()); + let reporter = svc.runtime_reporter(Some("user-1".to_owned()), "mcp-1".to_owned()); + + reporter.report(NodeRuntimeProgress { + phase: NodeRuntimeProgressPhase::Ready, + failure_kind: None, + message: None, + status_code: None, + }); + + let events = broadcaster.events(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].name, "runtime.statusChanged"); + assert_eq!(events[0].data["user_id"], "user-1"); + assert_eq!(events[0].data["scope"]["id"], "mcp-1"); + } + #[test] fn mcp_transport_diagnostic_logs_use_redacted_contract() { let result = McpConnectionTestResult { diff --git a/crates/aionui-mcp/src/oauth_service.rs b/crates/aionui-mcp/src/oauth_service.rs index a2259598c..dddce6ed0 100644 --- a/crates/aionui-mcp/src/oauth_service.rs +++ b/crates/aionui-mcp/src/oauth_service.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; @@ -69,8 +70,8 @@ struct PendingLogin { pub struct McpOAuthService { token_repo: Arc, http_client: reqwest::Client, - /// Mutex protecting the pending login state (only one login at a time). - pending: Arc>>, + /// Mutex protecting pending login state by (user_id, oauth_state). + pending: Arc>>, } impl McpOAuthService { @@ -78,7 +79,7 @@ impl McpOAuthService { Self { token_repo, http_client, - pending: Arc::new(Mutex::new(None)), + pending: Arc::new(Mutex::new(HashMap::new())), } } @@ -87,8 +88,8 @@ impl McpOAuthService { // ----------------------------------------------------------------------- /// Check whether the given server URL has a valid (non-expired) OAuth token. - pub async fn check_oauth_status(&self, server_url: &str) -> Result { - let authenticated = self.has_valid_token(server_url).await?; + pub async fn check_oauth_status(&self, user_id: &str, server_url: &str) -> Result { + let authenticated = self.has_valid_token(user_id, server_url).await?; Ok(OAuthStatusResponse { authenticated }) } @@ -100,8 +101,8 @@ impl McpOAuthService { /// 4. Build authorization URL and open it in the system browser /// 5. Wait for the redirect with the authorization code /// 6. Exchange code for tokens and persist them - pub async fn login(&self, server_url: &str) -> Result { - let (authorize_url, listener) = self.prepare_login_flow(server_url).await?; + pub async fn login(&self, user_id: &str, server_url: &str) -> Result { + let (authorize_url, listener) = self.prepare_login_flow(user_id, server_url).await?; // Open browser. debug!(url = %authorize_url, "Opening browser for OAuth authorization"); @@ -110,10 +111,10 @@ impl McpOAuthService { } // Wait for callback. - let code = match self.wait_for_callback(listener).await { - Ok(code) => code, + let (code, state) = match self.wait_for_callback(user_id, listener).await { + Ok(callback) => callback, Err(e) => { - self.clear_pending().await; + self.clear_pending_for_user(user_id).await; return Ok(OAuthLoginResponse { success: false, error: Some(e.to_string()), @@ -122,13 +123,13 @@ impl McpOAuthService { }; // Exchange code for tokens. - match self.exchange_code(server_url, code).await { + match self.exchange_code(user_id, server_url, code, state).await { Ok(()) => Ok(OAuthLoginResponse { success: true, error: None, }), Err(e) => { - self.clear_pending().await; + self.clear_pending_for_user(user_id).await; Ok(OAuthLoginResponse { success: false, error: Some(e.to_string()), @@ -140,8 +141,8 @@ impl McpOAuthService { /// Logout from the given MCP server URL (delete stored token). /// /// Idempotent: returns Ok even if no token was stored. - pub async fn logout(&self, server_url: &str) -> Result<(), McpError> { - match self.token_repo.delete(server_url).await { + pub async fn logout(&self, user_id: &str, server_url: &str) -> Result<(), McpError> { + match self.token_repo.delete(user_id, server_url).await { Ok(()) => { debug!(server_url, "OAuth token deleted"); Ok(()) @@ -155,8 +156,8 @@ impl McpOAuthService { } /// Return the list of server URLs that have stored OAuth tokens. - pub async fn get_authenticated_servers(&self) -> Result, McpError> { - let urls = self.token_repo.list_authenticated_urls().await?; + pub async fn get_authenticated_servers(&self, user_id: &str) -> Result, McpError> { + let urls = self.token_repo.list_authenticated_urls(user_id).await?; Ok(urls) } @@ -165,8 +166,8 @@ impl McpOAuthService { /// If the stored token is expired and a refresh token is available, /// automatically refreshes before returning. /// Returns `None` if no token is stored for this URL. - pub async fn get_token(&self, server_url: &str) -> Result, McpError> { - let row = match self.token_repo.get_by_url(server_url).await? { + pub async fn get_token(&self, user_id: &str, server_url: &str) -> Result, McpError> { + let row = match self.token_repo.get_by_url(user_id, server_url).await? { Some(row) => row, None => return Ok(None), }; @@ -177,7 +178,7 @@ impl McpOAuthService { if now >= expires_at - EXPIRY_MARGIN_MS && let Some(ref refresh_token) = row.refresh_token { - match self.refresh_token(server_url, refresh_token).await { + match self.refresh_token(user_id, server_url, refresh_token).await { Ok(new_token) => return Ok(Some(new_token)), Err(e) => { warn!( @@ -199,7 +200,7 @@ impl McpOAuthService { /// Discover endpoints, build OAuth client, generate PKCE, bind callback /// server, store pending state, and return the authorization URL + listener. - async fn prepare_login_flow(&self, server_url: &str) -> Result<(String, TcpListener), McpError> { + async fn prepare_login_flow(&self, user_id: &str, server_url: &str) -> Result<(String, TcpListener), McpError> { let metadata = self.discover_endpoints(server_url).await?; let auth_url_str = metadata.authorization_endpoint.clone(); @@ -236,21 +237,25 @@ impl McpOAuthService { { let mut pending = self.pending.lock().await; - *pending = Some(PendingLogin { - csrf_token, - pkce_verifier, - auth_url: auth_url_str, - token_url: token_url_str, - redirect_url: redirect_url_str, - }); + let state = csrf_token.secret().clone(); + pending.insert( + (user_id.to_string(), state), + PendingLogin { + csrf_token, + pkce_verifier, + auth_url: auth_url_str, + token_url: token_url_str, + redirect_url: redirect_url_str, + }, + ); } Ok((authorize_url.to_string(), listener)) } /// Check if a valid (non-expired) token exists for the URL. - async fn has_valid_token(&self, server_url: &str) -> Result { - let row = match self.token_repo.get_by_url(server_url).await? { + async fn has_valid_token(&self, user_id: &str, server_url: &str) -> Result { + let row = match self.token_repo.get_by_url(user_id, server_url).await? { Some(row) => row, None => return Ok(false), }; @@ -309,12 +314,13 @@ impl McpOAuthService { } /// Wait for the OAuth callback redirect on the given listener. - async fn wait_for_callback(&self, listener: TcpListener) -> Result { - let (code_tx, code_rx) = tokio::sync::oneshot::channel::>(); + async fn wait_for_callback(&self, user_id: &str, listener: TcpListener) -> Result<(String, String), McpError> { + let (code_tx, code_rx) = tokio::sync::oneshot::channel::>(); let pending = self.pending.clone(); + let user_id = user_id.to_string(); tokio::spawn(async move { - let result = Self::handle_callback_connection(listener, pending).await; + let result = Self::handle_callback_connection(&user_id, listener, pending).await; let _ = code_tx.send(result); }); @@ -329,9 +335,10 @@ impl McpOAuthService { /// Handle a single HTTP connection on the callback server. async fn handle_callback_connection( + user_id: &str, listener: TcpListener, - pending: Arc>>, - ) -> Result { + pending: Arc>>, + ) -> Result<(String, String), McpError> { let (mut stream, _) = listener .accept() .await @@ -349,7 +356,7 @@ impl McpOAuthService { // Validate CSRF state. let guard = pending.lock().await; let pending_login = guard - .as_ref() + .get(&(user_id.to_string(), state.clone())) .ok_or_else(|| McpError::OAuth("No pending login state".to_string()))?; if state != *pending_login.csrf_token.secret() { @@ -366,7 +373,7 @@ impl McpOAuthService { let _ = stream.write_all(response.as_bytes()).await; - Ok(code) + Ok((code, state)) } /// Build a no-redirect reqwest client for OAuth token exchange. @@ -378,11 +385,17 @@ impl McpOAuthService { } /// Exchange the authorization code for tokens and persist them. - async fn exchange_code(&self, server_url: &str, code: String) -> Result<(), McpError> { + async fn exchange_code( + &self, + user_id: &str, + server_url: &str, + code: String, + state: String, + ) -> Result<(), McpError> { let (auth_url_str, token_url_str, redirect_url_str, pkce_verifier) = { let mut guard = self.pending.lock().await; let pending = guard - .take() + .remove(&(user_id.to_string(), state)) .ok_or_else(|| McpError::OAuth("No pending login state".to_string()))?; ( pending.auth_url, @@ -411,13 +424,18 @@ impl McpOAuthService { .await .map_err(|e| McpError::OAuth(format!("Token exchange failed: {e}")))?; - self.persist_token(server_url, &token_result).await?; + self.persist_token(user_id, server_url, &token_result).await?; debug!(server_url, "OAuth tokens stored successfully"); Ok(()) } /// Refresh an expired access token using the refresh token. - async fn refresh_token(&self, server_url: &str, refresh_token_value: &str) -> Result { + async fn refresh_token( + &self, + user_id: &str, + server_url: &str, + refresh_token_value: &str, + ) -> Result { let metadata = self.discover_endpoints(server_url).await?; let token_url = TokenUrl::new(metadata.token_endpoint).map_err(|e| McpError::OAuth(format!("Invalid token URL: {e}")))?; @@ -445,6 +463,7 @@ impl McpOAuthService { self.token_repo .upsert(UpsertOAuthTokenParams { + user_id, server_url, access_token: &new_access_token, refresh_token: Some(new_refresh), @@ -458,11 +477,17 @@ impl McpOAuthService { } /// Persist token response to DB. - async fn persist_token(&self, server_url: &str, token_result: &TR) -> Result<(), McpError> { + async fn persist_token( + &self, + user_id: &str, + server_url: &str, + token_result: &TR, + ) -> Result<(), McpError> { let expires_at: Option = token_result.expires_in().map(|d| now_ms() + d.as_millis() as i64); self.token_repo .upsert(UpsertOAuthTokenParams { + user_id, server_url, access_token: token_result.access_token().secret(), refresh_token: token_result.refresh_token().map(|t| t.secret().as_str()), @@ -475,9 +500,9 @@ impl McpOAuthService { } /// Clear the pending login state. - async fn clear_pending(&self) { + async fn clear_pending_for_user(&self, user_id: &str) { let mut guard = self.pending.lock().await; - *guard = None; + guard.retain(|(pending_user_id, _), _| pending_user_id != user_id); } } @@ -563,6 +588,8 @@ fn url_decode(input: &str) -> String { mod tests { use super::*; + const TEST_USER_ID: &str = "user-1"; + // -- parse_callback_query ------------------------------------------------ #[test] @@ -653,13 +680,45 @@ mod tests { let _clone = svc.clone(); } + #[tokio::test] + async fn clear_pending_for_user_keeps_other_user_same_oauth_state() { + let svc = McpOAuthService::new(Arc::new(MockTokenRepo), reqwest::Client::new()); + insert_pending_login(&svc, "user-a", "shared-state").await; + insert_pending_login(&svc, "user-b", "shared-state").await; + + svc.clear_pending_for_user("user-a").await; + + let pending = svc.pending.lock().await; + assert!(!pending.contains_key(&("user-a".to_string(), "shared-state".to_string()))); + assert!(pending.contains_key(&("user-b".to_string(), "shared-state".to_string()))); + } + // -- Mock repositories --------------------------------------------------- + async fn insert_pending_login(svc: &McpOAuthService, user_id: &str, state: &str) { + let (_, pkce_verifier) = PkceCodeChallenge::new_random_sha256(); + let mut pending = svc.pending.lock().await; + pending.insert( + (user_id.to_string(), state.to_string()), + PendingLogin { + csrf_token: CsrfToken::new(state.to_string()), + pkce_verifier, + auth_url: "https://auth.example.com/authorize".to_string(), + token_url: "https://auth.example.com/token".to_string(), + redirect_url: "http://127.0.0.1/callback".to_string(), + }, + ); + } + struct MockTokenRepo; #[async_trait::async_trait] impl IOAuthTokenRepository for MockTokenRepo { - async fn get_by_url(&self, _: &str) -> Result, aionui_db::DbError> { + async fn get_by_url( + &self, + _: &str, + _: &str, + ) -> Result, aionui_db::DbError> { Ok(None) } @@ -670,11 +729,11 @@ mod tests { unimplemented!() } - async fn delete(&self, _: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _: &str, _: &str) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn list_authenticated_urls(&self) -> Result, aionui_db::DbError> { + async fn list_authenticated_urls(&self, _: &str) -> Result, aionui_db::DbError> { Ok(vec![]) } } @@ -683,7 +742,11 @@ mod tests { #[async_trait::async_trait] impl IOAuthTokenRepository for IdempotentDeleteRepo { - async fn get_by_url(&self, _: &str) -> Result, aionui_db::DbError> { + async fn get_by_url( + &self, + _: &str, + _: &str, + ) -> Result, aionui_db::DbError> { Ok(None) } @@ -694,13 +757,13 @@ mod tests { unimplemented!() } - async fn delete(&self, url: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _: &str, url: &str) -> Result<(), aionui_db::DbError> { Err(aionui_db::DbError::NotFound(format!( "OAuth token for '{url}' not found" ))) } - async fn list_authenticated_urls(&self) -> Result, aionui_db::DbError> { + async fn list_authenticated_urls(&self, _: &str) -> Result, aionui_db::DbError> { Ok(vec![]) } } @@ -709,8 +772,13 @@ mod tests { #[async_trait::async_trait] impl IOAuthTokenRepository for ValidTokenRepo { - async fn get_by_url(&self, _: &str) -> Result, aionui_db::DbError> { + async fn get_by_url( + &self, + _: &str, + _: &str, + ) -> Result, aionui_db::DbError> { Ok(Some(aionui_db::models::OAuthTokenRow { + user_id: TEST_USER_ID.to_string(), server_url: "https://example.com".to_string(), access_token: "valid_access_token".to_string(), refresh_token: None, @@ -728,11 +796,11 @@ mod tests { unimplemented!() } - async fn delete(&self, _: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _: &str, _: &str) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn list_authenticated_urls(&self) -> Result, aionui_db::DbError> { + async fn list_authenticated_urls(&self, _: &str) -> Result, aionui_db::DbError> { Ok(vec!["https://example.com".to_string()]) } } @@ -741,8 +809,13 @@ mod tests { #[async_trait::async_trait] impl IOAuthTokenRepository for ExpiredTokenRepo { - async fn get_by_url(&self, _: &str) -> Result, aionui_db::DbError> { + async fn get_by_url( + &self, + _: &str, + _: &str, + ) -> Result, aionui_db::DbError> { Ok(Some(aionui_db::models::OAuthTokenRow { + user_id: TEST_USER_ID.to_string(), server_url: "https://example.com".to_string(), access_token: "expired_token".to_string(), refresh_token: None, @@ -760,11 +833,11 @@ mod tests { unimplemented!() } - async fn delete(&self, _: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _: &str, _: &str) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn list_authenticated_urls(&self) -> Result, aionui_db::DbError> { + async fn list_authenticated_urls(&self, _: &str) -> Result, aionui_db::DbError> { Ok(vec![]) } } @@ -773,8 +846,13 @@ mod tests { #[async_trait::async_trait] impl IOAuthTokenRepository for NoExpiryTokenRepo { - async fn get_by_url(&self, _: &str) -> Result, aionui_db::DbError> { + async fn get_by_url( + &self, + _: &str, + _: &str, + ) -> Result, aionui_db::DbError> { Ok(Some(aionui_db::models::OAuthTokenRow { + user_id: TEST_USER_ID.to_string(), server_url: "https://example.com".to_string(), access_token: "no_expiry_token".to_string(), refresh_token: None, @@ -792,11 +870,11 @@ mod tests { unimplemented!() } - async fn delete(&self, _: &str) -> Result<(), aionui_db::DbError> { + async fn delete(&self, _: &str, _: &str) -> Result<(), aionui_db::DbError> { Ok(()) } - async fn list_authenticated_urls(&self) -> Result, aionui_db::DbError> { + async fn list_authenticated_urls(&self, _: &str) -> Result, aionui_db::DbError> { Ok(vec!["https://example.com".to_string()]) } } @@ -806,62 +884,76 @@ mod tests { #[tokio::test] async fn check_status_no_token_returns_false() { let svc = McpOAuthService::new(Arc::new(MockTokenRepo), reqwest::Client::new()); - let status = svc.check_oauth_status("https://example.com").await.unwrap(); + let status = svc + .check_oauth_status(TEST_USER_ID, "https://example.com") + .await + .unwrap(); assert!(!status.authenticated); } #[tokio::test] async fn check_status_with_valid_token() { let svc = McpOAuthService::new(Arc::new(ValidTokenRepo), reqwest::Client::new()); - let status = svc.check_oauth_status("https://example.com").await.unwrap(); + let status = svc + .check_oauth_status(TEST_USER_ID, "https://example.com") + .await + .unwrap(); assert!(status.authenticated); } #[tokio::test] async fn check_status_with_expired_token() { let svc = McpOAuthService::new(Arc::new(ExpiredTokenRepo), reqwest::Client::new()); - let status = svc.check_oauth_status("https://example.com").await.unwrap(); + let status = svc + .check_oauth_status(TEST_USER_ID, "https://example.com") + .await + .unwrap(); assert!(!status.authenticated); } #[tokio::test] async fn check_status_no_expiry_treated_as_valid() { let svc = McpOAuthService::new(Arc::new(NoExpiryTokenRepo), reqwest::Client::new()); - let status = svc.check_oauth_status("https://example.com").await.unwrap(); + let status = svc + .check_oauth_status(TEST_USER_ID, "https://example.com") + .await + .unwrap(); assert!(status.authenticated); } #[tokio::test] async fn logout_idempotent_for_nonexistent() { let svc = McpOAuthService::new(Arc::new(IdempotentDeleteRepo), reqwest::Client::new()); - svc.logout("https://nonexistent.example.com").await.unwrap(); + svc.logout(TEST_USER_ID, "https://nonexistent.example.com") + .await + .unwrap(); } #[tokio::test] async fn get_authenticated_servers_empty() { let svc = McpOAuthService::new(Arc::new(MockTokenRepo), reqwest::Client::new()); - let urls = svc.get_authenticated_servers().await.unwrap(); + let urls = svc.get_authenticated_servers(TEST_USER_ID).await.unwrap(); assert!(urls.is_empty()); } #[tokio::test] async fn get_authenticated_servers_returns_urls() { let svc = McpOAuthService::new(Arc::new(ValidTokenRepo), reqwest::Client::new()); - let urls = svc.get_authenticated_servers().await.unwrap(); + let urls = svc.get_authenticated_servers(TEST_USER_ID).await.unwrap(); assert_eq!(urls, vec!["https://example.com"]); } #[tokio::test] async fn get_token_returns_none_when_no_token() { let svc = McpOAuthService::new(Arc::new(MockTokenRepo), reqwest::Client::new()); - let token = svc.get_token("https://example.com").await.unwrap(); + let token = svc.get_token(TEST_USER_ID, "https://example.com").await.unwrap(); assert!(token.is_none()); } #[tokio::test] async fn get_token_returns_access_token() { let svc = McpOAuthService::new(Arc::new(ValidTokenRepo), reqwest::Client::new()); - let token = svc.get_token("https://example.com").await.unwrap(); + let token = svc.get_token(TEST_USER_ID, "https://example.com").await.unwrap(); assert_eq!(token.as_deref(), Some("valid_access_token")); } @@ -869,7 +961,7 @@ mod tests { async fn get_token_returns_expired_when_no_refresh() { let svc = McpOAuthService::new(Arc::new(ExpiredTokenRepo), reqwest::Client::new()); // Expired token with no refresh_token: returns the expired token as-is. - let token = svc.get_token("https://example.com").await.unwrap(); + let token = svc.get_token(TEST_USER_ID, "https://example.com").await.unwrap(); assert_eq!(token.as_deref(), Some("expired_token")); } } diff --git a/crates/aionui-mcp/src/routes.rs b/crates/aionui-mcp/src/routes.rs index 93cade5cd..671818ab1 100644 --- a/crates/aionui-mcp/src/routes.rs +++ b/crates/aionui-mcp/src/routes.rs @@ -2,7 +2,7 @@ use axum::Router; use axum::extract::rejection::JsonRejection; -use axum::extract::{Json, Path, State}; +use axum::extract::{Extension, Json, Path, State}; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; @@ -12,6 +12,7 @@ use aionui_api_types::{ McpConnectionTestErrorCode, McpServerResponse, OAuthCheckStatusRequest, OAuthLoginRequest, OAuthLoginResponse, OAuthLogoutRequest, OAuthStatusResponse, TestMcpConnectionRequest, UpdateMcpServerRequest, }; +use aionui_auth::CurrentUser; use aionui_common::ApiError; use crate::connection_test::McpConnectionTestService; @@ -87,40 +88,56 @@ pub fn mcp_routes(state: McpRouterState) -> Router { /// `GET /api/mcp/servers` — list all MCP servers. async fn list_servers( State(state): State, + Extension(user): Extension, ) -> Result>>, ApiError> { - let servers = state.config_service.list_servers().await.map_err(ApiError::from)?; + let servers = state + .config_service + .list_servers(&user.id) + .await + .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(servers))) } /// `GET /api/mcp/servers/:id` — get a single MCP server. async fn get_server( State(state): State, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - let server = state.config_service.get_server(&id).await.map_err(ApiError::from)?; + let server = state + .config_service + .get_server(&user.id, &id) + .await + .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(server))) } /// `POST /api/mcp/servers` — create (or upsert by name) an MCP server. async fn add_server( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result<(StatusCode, Json>), ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let server = state.config_service.add_server(req).await.map_err(ApiError::from)?; + let server = state + .config_service + .add_server(&user.id, req) + .await + .map_err(ApiError::from)?; Ok((StatusCode::CREATED, Json(ApiResponse::ok(server)))) } /// `PUT /api/mcp/servers/:id` — partial update an MCP server. async fn edit_server( State(state): State, + Extension(user): Extension, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; let server = state .config_service - .edit_server(&id, req) + .edit_server(&user.id, &id, req) .await .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(server))) @@ -129,28 +146,43 @@ async fn edit_server( /// `DELETE /api/mcp/servers/:id` — delete an MCP server. async fn delete_server( State(state): State, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - state.config_service.delete_server(&id).await.map_err(ApiError::from)?; + state + .config_service + .delete_server(&user.id, &id) + .await + .map_err(ApiError::from)?; Ok(Json(ApiResponse::success())) } /// `POST /api/mcp/servers/:id/toggle` — toggle enabled state. async fn toggle_server( State(state): State, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - let server = state.config_service.toggle_server(&id).await.map_err(ApiError::from)?; + let server = state + .config_service + .toggle_server(&user.id, &id) + .await + .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(server))) } /// `POST /api/mcp/servers/import` — batch import MCP servers. async fn batch_import( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let servers = state.config_service.batch_import(req).await.map_err(ApiError::from)?; + let servers = state + .config_service + .batch_import(&user.id, req) + .await + .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(servers))) } @@ -163,6 +195,7 @@ async fn batch_import( /// Creates a temporary MCP client, connects, lists tools, and closes. async fn test_connection( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result { let Json(req) = body.map_err(ApiError::from)?; @@ -172,13 +205,14 @@ async fn test_connection( .test_connection_with_runtime_scope( &req.name, &transport, + Some(&user.id), req.runtime_scope_id.as_deref().or(req.id.as_deref()), ) .await; if let Some(server_id) = req.id.as_deref() { state .config_service - .persist_test_result(server_id, &result) + .persist_test_result(&user.id, server_id, &result) .await .map_err(ApiError::from)?; } @@ -227,8 +261,13 @@ fn connection_test_failure_status(code: McpConnectionTestErrorCode) -> StatusCod /// and return their current MCP server configurations. async fn get_agent_configs( State(state): State, + Extension(user): Extension, ) -> Result>>, ApiError> { - let configs = state.sync_service.get_agent_configs().await.map_err(ApiError::from)?; + let configs = state + .sync_service + .get_agent_configs(&user.id) + .await + .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(configs))) } @@ -239,12 +278,13 @@ async fn get_agent_configs( /// `POST /api/mcp/oauth/check-status` — check OAuth authentication status. async fn oauth_check_status( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; let status = state .oauth_service - .check_oauth_status(&req.server_url) + .check_oauth_status(&user.id, &req.server_url) .await .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(status))) @@ -256,12 +296,13 @@ async fn oauth_check_status( /// the callback, and exchanges the code for tokens. async fn oauth_login( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; let result = state .oauth_service - .login(&req.server_url) + .login(&user.id, &req.server_url) .await .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(result))) @@ -270,22 +311,26 @@ async fn oauth_login( /// `POST /api/mcp/oauth/logout` — delete stored OAuth token. async fn oauth_logout( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; state .oauth_service - .logout(&req.server_url) + .logout(&user.id, &req.server_url) .await .map_err(ApiError::from)?; Ok(Json(ApiResponse::success())) } /// `GET /api/mcp/oauth/authenticated` — list server URLs with stored tokens. -async fn oauth_authenticated(State(state): State) -> Result>>, ApiError> { +async fn oauth_authenticated( + State(state): State, + Extension(user): Extension, +) -> Result>>, ApiError> { let urls = state .oauth_service - .get_authenticated_servers() + .get_authenticated_servers(&user.id) .await .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(urls))) diff --git a/crates/aionui-mcp/src/service.rs b/crates/aionui-mcp/src/service.rs index 5bdb5260c..8df26403d 100644 --- a/crates/aionui-mcp/src/service.rs +++ b/crates/aionui-mcp/src/service.rs @@ -31,24 +31,34 @@ pub struct McpConfigService { repo: Arc, } +struct UpsertMcpServer<'a> { + user_id: &'a str, + name: &'a str, + description: Option<&'a str>, + transport: &'a McpServerTransport, + original_json: Option<&'a str>, + builtin: bool, + enabled: bool, +} + impl McpConfigService { pub fn new(repo: Arc) -> Self { Self { repo } } /// List all MCP servers. - pub async fn list_servers(&self) -> Result, McpError> { - let rows = self.repo.list().await?; + pub async fn list_servers(&self, user_id: &str) -> Result, McpError> { + let rows = self.repo.list(user_id).await?; rows.into_iter() .map(|row| McpServer::from_row(row).map(McpServer::into_response)) .collect() } /// Get a single MCP server by ID. - pub async fn get_server(&self, id: &str) -> Result { + pub async fn get_server(&self, user_id: &str, id: &str) -> Result { let row = self .repo - .find_by_id(id) + .find_by_id(user_id, id) .await? .ok_or_else(|| McpError::NotFound(id.to_owned()))?; let server = McpServer::from_row(row)?; @@ -59,25 +69,31 @@ impl McpConfigService { /// /// If a server with the same name already exists, it is updated /// (transport, description, original_json) rather than creating a duplicate. - pub async fn add_server(&self, req: CreateMcpServerRequest) -> Result { + pub async fn add_server(&self, user_id: &str, req: CreateMcpServerRequest) -> Result { let transport = normalize_transport(McpServerTransport::from(req.transport))?; - self.upsert_server( - &req.name, - req.description.as_deref(), - &transport, - req.original_json.as_deref(), - req.builtin, - false, - ) + self.upsert_server(UpsertMcpServer { + user_id, + name: &req.name, + description: req.description.as_deref(), + transport: &transport, + original_json: req.original_json.as_deref(), + builtin: req.builtin, + enabled: false, + }) .await } /// Edit an existing MCP server (partial update). - pub async fn edit_server(&self, id: &str, req: UpdateMcpServerRequest) -> Result { + pub async fn edit_server( + &self, + user_id: &str, + id: &str, + req: UpdateMcpServerRequest, + ) -> Result { // Verify the server exists let existing_server = self .repo - .find_by_id(id) + .find_by_id(user_id, id) .await? .ok_or_else(|| McpError::NotFound(id.to_owned()))?; @@ -92,7 +108,7 @@ impl McpConfigService { // Check name uniqueness if renaming if let Some(ref new_name) = req.name - && let Some(existing) = self.repo.find_by_name_any(new_name).await? + && let Some(existing) = self.repo.find_by_name_any(user_id, new_name).await? && existing.id != id { if existing.builtin { @@ -121,7 +137,7 @@ impl McpConfigService { ..Default::default() }; - let row = self.repo.update(id, params).await?; + let row = self.repo.update(user_id, id, params).await?; let server = McpServer::from_row(row)?; Ok(server.into_response()) } @@ -129,24 +145,24 @@ impl McpConfigService { /// Soft-delete an MCP server by ID. /// /// Returns whether the deleted server was enabled. - pub async fn delete_server(&self, id: &str) -> Result { + pub async fn delete_server(&self, user_id: &str, id: &str) -> Result { let row = self .repo - .find_by_id(id) + .find_by_id(user_id, id) .await? .ok_or_else(|| McpError::NotFound(id.to_owned()))?; let was_enabled = row.enabled; - self.repo.delete(id).await?; + self.repo.delete(user_id, id).await?; Ok(was_enabled) } /// Toggle the enabled state of an MCP server. /// /// Returns the updated server response. - pub async fn toggle_server(&self, id: &str) -> Result { + pub async fn toggle_server(&self, user_id: &str, id: &str) -> Result { let row = self .repo - .find_by_id(id) + .find_by_id(user_id, id) .await? .ok_or_else(|| McpError::NotFound(id.to_owned()))?; @@ -155,7 +171,7 @@ impl McpConfigService { enabled: Some(new_enabled), ..Default::default() }; - let updated = self.repo.update(id, params).await?; + let updated = self.repo.update(user_id, id, params).await?; let server = McpServer::from_row(updated)?; Ok(server.into_response()) } @@ -164,12 +180,16 @@ impl McpConfigService { /// /// Each server is processed individually: existing names are updated, /// new names are created. - pub async fn batch_import(&self, req: BatchImportMcpServersRequest) -> Result, McpError> { + pub async fn batch_import( + &self, + user_id: &str, + req: BatchImportMcpServersRequest, + ) -> Result, McpError> { let requested_count = req.servers.len(); let mut rows = Vec::with_capacity(requested_count); let mut skipped_reserved_count = 0usize; for server_req in req.servers { - if let Some(existing) = self.repo.find_by_name_any(&server_req.name).await? + if let Some(existing) = self.repo.find_by_name_any(user_id, &server_req.name).await? && existing.builtin { skipped_reserved_count += 1; @@ -182,14 +202,15 @@ impl McpConfigService { let transport = normalize_transport(McpServerTransport::from(server_req.transport))?; let server = self - .upsert_server( - &server_req.name, - server_req.description.as_deref(), - &transport, - server_req.original_json.as_deref(), - server_req.builtin, - server_req.enabled.unwrap_or(false), - ) + .upsert_server(UpsertMcpServer { + user_id, + name: &server_req.name, + description: server_req.description.as_deref(), + transport: &transport, + original_json: server_req.original_json.as_deref(), + builtin: server_req.builtin, + enabled: server_req.enabled.unwrap_or(false), + }) .await?; rows.push(server); } @@ -204,28 +225,34 @@ impl McpConfigService { } /// Persist the latest connection test result for an existing MCP server. - pub async fn persist_test_result(&self, id: &str, result: &McpConnectionTestResult) -> Result<(), McpError> { + pub async fn persist_test_result( + &self, + user_id: &str, + id: &str, + result: &McpConnectionTestResult, + ) -> Result<(), McpError> { let status = if result.success { "connected" } else { "error" }; let last_connected = if result.success { Some(now_ms()) } else { None }; let tools_json = result.tools.as_ref().map(serde_json::to_string).transpose()?; - self.repo.update_status(id, status, last_connected).await?; - self.repo.update_tools(id, tools_json.as_deref()).await?; + self.repo.update_status(user_id, id, status, last_connected).await?; + self.repo.update_tools(user_id, id, tools_json.as_deref()).await?; Ok(()) } - async fn upsert_server( - &self, - name: &str, - description: Option<&str>, - transport: &McpServerTransport, - original_json: Option<&str>, - builtin: bool, - enabled: bool, - ) -> Result { + async fn upsert_server(&self, params: UpsertMcpServer<'_>) -> Result { + let UpsertMcpServer { + user_id, + name, + description, + transport, + original_json, + builtin, + enabled, + } = params; let config_json = transport.to_config_json()?; - if let Some(existing) = self.repo.find_by_name_any(name).await? { + if let Some(existing) = self.repo.find_by_name_any(user_id, name).await? { if existing.builtin { return Err(McpError::Conflict(format!( "Builtin MCP server name '{name}' is reserved" @@ -242,12 +269,13 @@ impl McpConfigService { deleted_at: Some(None), ..Default::default() }; - let updated = self.repo.update(&existing.id, params).await?; + let updated = self.repo.update(user_id, &existing.id, params).await?; let server = McpServer::from_row(updated)?; return Ok(server.into_response()); } let params = CreateMcpServerParams { + user_id, name, description, enabled, @@ -361,6 +389,8 @@ mod tests { use std::collections::HashMap; use std::sync::Mutex; + const TEST_USER_ID: &str = "user-1"; + // -- In-memory mock repository ------------------------------------------- #[derive(Debug)] @@ -390,39 +420,46 @@ mod tests { #[async_trait::async_trait] impl IMcpServerRepository for MockMcpServerRepo { - async fn list(&self) -> Result, DbError> { + async fn list(&self, user_id: &str) -> Result, DbError> { let servers = self.servers.lock().unwrap(); - Ok(servers.iter().filter(|s| s.deleted_at.is_none()).cloned().collect()) + Ok(servers + .iter() + .filter(|s| s.user_id == user_id && s.deleted_at.is_none()) + .cloned() + .collect()) } - async fn find_by_id(&self, id: &str) -> Result, DbError> { + async fn find_by_id(&self, user_id: &str, id: &str) -> Result, DbError> { let servers = self.servers.lock().unwrap(); - Ok(servers.iter().find(|s| s.id == id && s.deleted_at.is_none()).cloned()) + Ok(servers + .iter() + .find(|s| s.user_id == user_id && s.id == id && s.deleted_at.is_none()) + .cloned()) } - async fn find_by_name(&self, name: &str) -> Result, DbError> { + async fn find_by_name(&self, user_id: &str, name: &str) -> Result, DbError> { let servers = self.servers.lock().unwrap(); Ok(servers .iter() - .find(|s| s.name == name && s.deleted_at.is_none()) + .find(|s| s.user_id == user_id && s.name == name && s.deleted_at.is_none()) .cloned()) } - async fn find_by_id_any(&self, id: &str) -> Result, DbError> { + async fn find_by_id_any(&self, user_id: &str, id: &str) -> Result, DbError> { let servers = self.servers.lock().unwrap(); - Ok(servers.iter().find(|s| s.id == id).cloned()) + Ok(servers.iter().find(|s| s.user_id == user_id && s.id == id).cloned()) } - async fn find_by_name_any(&self, name: &str) -> Result, DbError> { + async fn find_by_name_any(&self, user_id: &str, name: &str) -> Result, DbError> { let servers = self.servers.lock().unwrap(); - Ok(servers.iter().find(|s| s.name == name).cloned()) + Ok(servers.iter().find(|s| s.user_id == user_id && s.name == name).cloned()) } - async fn list_by_ids_any(&self, ids: &[String]) -> Result, DbError> { + async fn list_by_ids_any(&self, user_id: &str, ids: &[String]) -> Result, DbError> { let servers = self.servers.lock().unwrap(); Ok(servers .iter() - .filter(|server| ids.iter().any(|id| id == &server.id)) + .filter(|server| server.user_id == user_id && ids.iter().any(|id| id == &server.id)) .cloned() .collect()) } @@ -437,6 +474,7 @@ mod tests { } let row = McpServerRow { id: self.next_id(), + user_id: params.user_id.to_owned(), name: params.name.to_owned(), description: params.description.map(String::from), enabled: params.enabled, @@ -455,16 +493,24 @@ mod tests { Ok(row) } - async fn update(&self, id: &str, params: UpdateMcpServerParams<'_>) -> Result { + async fn update( + &self, + user_id: &str, + id: &str, + params: UpdateMcpServerParams<'_>, + ) -> Result { let mut servers = self.servers.lock().unwrap(); let idx = servers .iter() - .position(|s| s.id == id) + .position(|s| s.user_id == user_id && s.id == id) .ok_or_else(|| DbError::NotFound(format!("MCP server {id}")))?; // Check name conflict if let Some(new_name) = params.name { - if servers.iter().any(|s| s.name == new_name && s.id != id) { + if servers + .iter() + .any(|s| s.user_id == user_id && s.name == new_name && s.id != id) + { return Err(DbError::Conflict(format!( "MCP server name '{new_name}' already exists" ))); @@ -499,11 +545,11 @@ mod tests { Ok(servers[idx].clone()) } - async fn delete(&self, id: &str) -> Result<(), DbError> { + async fn delete(&self, user_id: &str, id: &str) -> Result<(), DbError> { let mut servers = self.servers.lock().unwrap(); let idx = servers .iter() - .position(|s| s.id == id && s.deleted_at.is_none()) + .position(|s| s.user_id == user_id && s.id == id && s.deleted_at.is_none()) .ok_or_else(|| DbError::NotFound(format!("MCP server {id}")))?; servers[idx].enabled = false; servers[idx].deleted_at = Some(Self::now()); @@ -511,11 +557,18 @@ mod tests { Ok(()) } - async fn batch_upsert(&self, params_list: &[CreateMcpServerParams<'_>]) -> Result, DbError> { + async fn batch_upsert( + &self, + user_id: &str, + params_list: &[CreateMcpServerParams<'_>], + ) -> Result, DbError> { let mut results = Vec::new(); for params in params_list { let mut servers = self.servers.lock().unwrap(); - if let Some(idx) = servers.iter().position(|s| s.name == params.name) { + if let Some(idx) = servers + .iter() + .position(|s| s.user_id == user_id && s.name == params.name) + { // Update existing servers[idx].description = params.description.map(String::from); servers[idx].transport_type = params.transport_type.to_owned(); @@ -527,6 +580,7 @@ mod tests { // Create new let row = McpServerRow { id: self.next_id(), + user_id: params.user_id.to_owned(), name: params.name.to_owned(), description: params.description.map(String::from), enabled: params.enabled, @@ -550,6 +604,7 @@ mod tests { async fn update_status( &self, + user_id: &str, id: &str, status: &str, last_connected: Option, @@ -557,7 +612,7 @@ mod tests { let mut servers = self.servers.lock().unwrap(); let idx = servers .iter() - .position(|s| s.id == id) + .position(|s| s.user_id == user_id && s.id == id) .ok_or_else(|| DbError::NotFound(format!("MCP server {id}")))?; servers[idx].last_test_status = status.to_owned(); if let Some(lc) = last_connected { @@ -566,11 +621,11 @@ mod tests { Ok(()) } - async fn update_tools(&self, id: &str, tools: Option<&str>) -> Result<(), DbError> { + async fn update_tools(&self, user_id: &str, id: &str, tools: Option<&str>) -> Result<(), DbError> { let mut servers = self.servers.lock().unwrap(); let idx = servers .iter() - .position(|s| s.id == id) + .position(|s| s.user_id == user_id && s.id == id) .ok_or_else(|| DbError::NotFound(format!("MCP server {id}")))?; servers[idx].tools = tools.map(String::from); Ok(()) @@ -642,17 +697,17 @@ mod tests { #[tokio::test] async fn list_servers_empty() { let svc = make_service(); - let result = svc.list_servers().await.unwrap(); + let result = svc.list_servers(TEST_USER_ID).await.unwrap(); assert!(result.is_empty()); } #[tokio::test] async fn list_servers_returns_all() { let svc = make_service(); - svc.add_server(stdio_create_req("a")).await.unwrap(); - svc.add_server(http_create_req("b")).await.unwrap(); + svc.add_server(TEST_USER_ID, stdio_create_req("a")).await.unwrap(); + svc.add_server(TEST_USER_ID, http_create_req("b")).await.unwrap(); - let result = svc.list_servers().await.unwrap(); + let result = svc.list_servers(TEST_USER_ID).await.unwrap(); assert_eq!(result.len(), 2); } @@ -661,8 +716,8 @@ mod tests { #[tokio::test] async fn get_server_found() { let svc = make_service(); - let created = svc.add_server(stdio_create_req("test")).await.unwrap(); - let found = svc.get_server(&created.id).await.unwrap(); + let created = svc.add_server(TEST_USER_ID, stdio_create_req("test")).await.unwrap(); + let found = svc.get_server(TEST_USER_ID, &created.id).await.unwrap(); assert_eq!(found.id, created.id); assert_eq!(found.name, "test"); } @@ -670,7 +725,7 @@ mod tests { #[tokio::test] async fn get_server_not_found() { let svc = make_service(); - let result = svc.get_server("nonexistent").await; + let result = svc.get_server(TEST_USER_ID, "nonexistent").await; assert!(matches!(result, Err(McpError::NotFound(_)))); } @@ -679,7 +734,7 @@ mod tests { #[tokio::test] async fn add_server_creates_new() { let svc = make_service(); - let resp = svc.add_server(stdio_create_req("new-srv")).await.unwrap(); + let resp = svc.add_server(TEST_USER_ID, stdio_create_req("new-srv")).await.unwrap(); assert_eq!(resp.name, "new-srv"); assert!(!resp.enabled); assert_eq!(resp.last_test_status, McpServerStatus::Disconnected); @@ -689,10 +744,16 @@ mod tests { #[tokio::test] async fn add_server_upserts_existing() { let svc = make_service(); - let first = svc.add_server(stdio_create_req("upsert-test")).await.unwrap(); + let first = svc + .add_server(TEST_USER_ID, stdio_create_req("upsert-test")) + .await + .unwrap(); // Second add with same name updates existing - let updated = svc.add_server(http_create_req("upsert-test")).await.unwrap(); + let updated = svc + .add_server(TEST_USER_ID, http_create_req("upsert-test")) + .await + .unwrap(); assert_eq!(updated.id, first.id); // Transport should be updated to http match updated.transport { @@ -707,17 +768,20 @@ mod tests { async fn add_server_stdio_complete() { let svc = make_service(); let resp = svc - .add_server(CreateMcpServerRequest { - name: "stdio-full".into(), - description: Some("full stdio".into()), - transport: McpTransport::Stdio { - command: "node".into(), - args: vec!["index.js".into()], - env: HashMap::from([("KEY".into(), "val".into())]), + .add_server( + TEST_USER_ID, + CreateMcpServerRequest { + name: "stdio-full".into(), + description: Some("full stdio".into()), + transport: McpTransport::Stdio { + command: "node".into(), + args: vec!["index.js".into()], + env: HashMap::from([("KEY".into(), "val".into())]), + }, + original_json: Some(r#"{"name":"stdio-full"}"#.into()), + builtin: true, }, - original_json: Some(r#"{"name":"stdio-full"}"#.into()), - builtin: true, - }) + ) .await .unwrap(); assert_eq!(resp.name, "stdio-full"); @@ -730,9 +794,13 @@ mod tests { #[tokio::test] async fn edit_server_rejects_name_change() { let svc = make_service(); - let created = svc.add_server(stdio_create_req("old-name")).await.unwrap(); + let created = svc + .add_server(TEST_USER_ID, stdio_create_req("old-name")) + .await + .unwrap(); let err = svc .edit_server( + TEST_USER_ID, &created.id, UpdateMcpServerRequest { name: Some("new-name".into()), @@ -750,9 +818,10 @@ mod tests { #[tokio::test] async fn edit_server_updates_transport() { let svc = make_service(); - let created = svc.add_server(stdio_create_req("test")).await.unwrap(); + let created = svc.add_server(TEST_USER_ID, stdio_create_req("test")).await.unwrap(); let updated = svc .edit_server( + TEST_USER_ID, &created.id, UpdateMcpServerRequest { name: None, @@ -776,11 +845,12 @@ mod tests { #[tokio::test] async fn edit_server_clears_description() { let svc = make_service(); - let created = svc.add_server(stdio_create_req("test")).await.unwrap(); + let created = svc.add_server(TEST_USER_ID, stdio_create_req("test")).await.unwrap(); assert!(created.description.is_some()); let updated = svc .edit_server( + TEST_USER_ID, &created.id, UpdateMcpServerRequest { name: None, @@ -800,6 +870,7 @@ mod tests { let svc = make_service(); let result = svc .edit_server( + TEST_USER_ID, "nonexistent", UpdateMcpServerRequest { name: Some("x".into()), @@ -816,11 +887,17 @@ mod tests { #[tokio::test] async fn edit_server_name_conflict() { let svc = make_service(); - svc.add_server(stdio_create_req("server-a")).await.unwrap(); - let b = svc.add_server(stdio_create_req("server-b")).await.unwrap(); + svc.add_server(TEST_USER_ID, stdio_create_req("server-a")) + .await + .unwrap(); + let b = svc + .add_server(TEST_USER_ID, stdio_create_req("server-b")) + .await + .unwrap(); let result = svc .edit_server( + TEST_USER_ID, &b.id, UpdateMcpServerRequest { name: Some("server-a".into()), // conflict @@ -837,11 +914,15 @@ mod tests { #[tokio::test] async fn edit_server_rename_to_same_name() { let svc = make_service(); - let a = svc.add_server(stdio_create_req("server-a")).await.unwrap(); + let a = svc + .add_server(TEST_USER_ID, stdio_create_req("server-a")) + .await + .unwrap(); // Renaming to the same name should succeed let result = svc .edit_server( + TEST_USER_ID, &a.id, UpdateMcpServerRequest { name: Some("server-a".into()), @@ -858,11 +939,15 @@ mod tests { #[tokio::test] async fn edit_server_updates_builtin_flag() { let svc = make_service(); - let created = svc.add_server(stdio_create_req("chrome-devtools")).await.unwrap(); + let created = svc + .add_server(TEST_USER_ID, stdio_create_req("chrome-devtools")) + .await + .unwrap(); assert!(!created.builtin); let updated = svc .edit_server( + TEST_USER_ID, &created.id, UpdateMcpServerRequest { name: None, @@ -882,31 +967,31 @@ mod tests { #[tokio::test] async fn delete_server_removes_and_returns_enabled_status() { let svc = make_service(); - let created = svc.add_server(stdio_create_req("test")).await.unwrap(); + let created = svc.add_server(TEST_USER_ID, stdio_create_req("test")).await.unwrap(); // Not enabled - let was_enabled = svc.delete_server(&created.id).await.unwrap(); + let was_enabled = svc.delete_server(TEST_USER_ID, &created.id).await.unwrap(); assert!(!was_enabled); // Should be hidden from active queries - let result = svc.get_server(&created.id).await; + let result = svc.get_server(TEST_USER_ID, &created.id).await; assert!(matches!(result, Err(McpError::NotFound(_)))); } #[tokio::test] async fn delete_enabled_server_returns_true() { let svc = make_service(); - let created = svc.add_server(stdio_create_req("test")).await.unwrap(); - svc.toggle_server(&created.id).await.unwrap(); // enable + let created = svc.add_server(TEST_USER_ID, stdio_create_req("test")).await.unwrap(); + svc.toggle_server(TEST_USER_ID, &created.id).await.unwrap(); // enable - let was_enabled = svc.delete_server(&created.id).await.unwrap(); + let was_enabled = svc.delete_server(TEST_USER_ID, &created.id).await.unwrap(); assert!(was_enabled); } #[tokio::test] async fn delete_server_not_found() { let svc = make_service(); - let result = svc.delete_server("nonexistent").await; + let result = svc.delete_server(TEST_USER_ID, "nonexistent").await; assert!(matches!(result, Err(McpError::NotFound(_)))); } @@ -915,20 +1000,20 @@ mod tests { #[tokio::test] async fn toggle_server_enables_then_disables() { let svc = make_service(); - let created = svc.add_server(stdio_create_req("toggle")).await.unwrap(); + let created = svc.add_server(TEST_USER_ID, stdio_create_req("toggle")).await.unwrap(); assert!(!created.enabled); - let toggled = svc.toggle_server(&created.id).await.unwrap(); + let toggled = svc.toggle_server(TEST_USER_ID, &created.id).await.unwrap(); assert!(toggled.enabled); - let toggled_back = svc.toggle_server(&created.id).await.unwrap(); + let toggled_back = svc.toggle_server(TEST_USER_ID, &created.id).await.unwrap(); assert!(!toggled_back.enabled); } #[tokio::test] async fn toggle_server_not_found() { let svc = make_service(); - let result = svc.toggle_server("nonexistent").await; + let result = svc.toggle_server(TEST_USER_ID, "nonexistent").await; assert!(matches!(result, Err(McpError::NotFound(_)))); } @@ -940,17 +1025,19 @@ mod tests { let req = BatchImportMcpServersRequest { servers: vec![stdio_import_req("a"), http_import_req("b")], }; - let results = svc.batch_import(req).await.unwrap(); + let results = svc.batch_import(TEST_USER_ID, req).await.unwrap(); assert_eq!(results.len(), 2); - let all = svc.list_servers().await.unwrap(); + let all = svc.list_servers(TEST_USER_ID).await.unwrap(); assert_eq!(all.len(), 2); } #[tokio::test] async fn batch_import_upserts_existing() { let svc = make_service(); - svc.add_server(stdio_create_req("existing")).await.unwrap(); + svc.add_server(TEST_USER_ID, stdio_create_req("existing")) + .await + .unwrap(); let req = BatchImportMcpServersRequest { servers: vec![ @@ -958,20 +1045,23 @@ mod tests { stdio_import_req("brand-new"), // create ], }; - let results = svc.batch_import(req).await.unwrap(); + let results = svc.batch_import(TEST_USER_ID, req).await.unwrap(); assert_eq!(results.len(), 2); - let all = svc.list_servers().await.unwrap(); + let all = svc.list_servers(TEST_USER_ID).await.unwrap(); assert_eq!(all.len(), 2); } #[tokio::test] async fn add_server_restores_soft_deleted_row() { let svc = make_service(); - let created = svc.add_server(stdio_create_req("restored")).await.unwrap(); - svc.delete_server(&created.id).await.unwrap(); + let created = svc + .add_server(TEST_USER_ID, stdio_create_req("restored")) + .await + .unwrap(); + svc.delete_server(TEST_USER_ID, &created.id).await.unwrap(); - let restored = svc.add_server(http_create_req("restored")).await.unwrap(); + let restored = svc.add_server(TEST_USER_ID, http_create_req("restored")).await.unwrap(); assert_eq!(restored.id, created.id); match restored.transport { McpTransport::Http { .. } => {} @@ -982,52 +1072,64 @@ mod tests { #[tokio::test] async fn add_server_rejects_overriding_builtin_name() { let svc = make_service(); - svc.add_server(CreateMcpServerRequest { - name: "chrome-devtools".into(), - description: Some("builtin".into()), - transport: McpTransport::Stdio { - command: "npx".into(), - args: vec!["-y".into(), "chrome-devtools-mcp@latest".into()], - env: HashMap::new(), + svc.add_server( + TEST_USER_ID, + CreateMcpServerRequest { + name: "chrome-devtools".into(), + description: Some("builtin".into()), + transport: McpTransport::Stdio { + command: "npx".into(), + args: vec!["-y".into(), "chrome-devtools-mcp@latest".into()], + env: HashMap::new(), + }, + original_json: None, + builtin: true, }, - original_json: None, - builtin: true, - }) + ) .await .unwrap(); - let err = svc.add_server(stdio_create_req("chrome-devtools")).await.unwrap_err(); + let err = svc + .add_server(TEST_USER_ID, stdio_create_req("chrome-devtools")) + .await + .unwrap_err(); assert!(matches!(err, McpError::Conflict(_))); } #[tokio::test] async fn add_server_rejects_overriding_builtin_name_even_with_builtin_payload() { let svc = make_service(); - svc.add_server(CreateMcpServerRequest { - name: "chrome-devtools".into(), - description: Some("builtin".into()), - transport: McpTransport::Stdio { - command: "npx".into(), - args: vec!["-y".into(), "chrome-devtools-mcp@latest".into()], - env: HashMap::new(), + svc.add_server( + TEST_USER_ID, + CreateMcpServerRequest { + name: "chrome-devtools".into(), + description: Some("builtin".into()), + transport: McpTransport::Stdio { + command: "npx".into(), + args: vec!["-y".into(), "chrome-devtools-mcp@latest".into()], + env: HashMap::new(), + }, + original_json: None, + builtin: true, }, - original_json: None, - builtin: true, - }) + ) .await .unwrap(); let err = svc - .add_server(CreateMcpServerRequest { - name: "chrome-devtools".into(), - description: Some("malicious override".into()), - transport: McpTransport::Http { - url: "https://example.com/mcp".into(), - headers: HashMap::new(), + .add_server( + TEST_USER_ID, + CreateMcpServerRequest { + name: "chrome-devtools".into(), + description: Some("malicious override".into()), + transport: McpTransport::Http { + url: "https://example.com/mcp".into(), + headers: HashMap::new(), + }, + original_json: None, + builtin: true, }, - original_json: None, - builtin: true, - }) + ) .await .unwrap_err(); assert!(matches!(err, McpError::Conflict(_))); @@ -1036,48 +1138,54 @@ mod tests { #[tokio::test] async fn batch_import_skips_reserved_builtin_name() { let svc = make_service(); - svc.add_server(CreateMcpServerRequest { - name: "chrome-devtools".into(), - description: Some("builtin".into()), - transport: McpTransport::Stdio { - command: "npx".into(), - args: vec!["-y".into(), "chrome-devtools-mcp@latest".into()], - env: HashMap::new(), + svc.add_server( + TEST_USER_ID, + CreateMcpServerRequest { + name: "chrome-devtools".into(), + description: Some("builtin".into()), + transport: McpTransport::Stdio { + command: "npx".into(), + args: vec!["-y".into(), "chrome-devtools-mcp@latest".into()], + env: HashMap::new(), + }, + original_json: None, + builtin: true, }, - original_json: None, - builtin: true, - }) + ) .await .unwrap(); let results = svc - .batch_import(BatchImportMcpServersRequest { - servers: vec![ - ImportMcpServerRequest { - name: "chrome-devtools".into(), - description: Some("imported".into()), - transport: McpTransport::Http { - url: "https://example.com/mcp".into(), - headers: HashMap::new(), + .batch_import( + TEST_USER_ID, + BatchImportMcpServersRequest { + servers: vec![ + ImportMcpServerRequest { + name: "chrome-devtools".into(), + description: Some("imported".into()), + transport: McpTransport::Http { + url: "https://example.com/mcp".into(), + headers: HashMap::new(), + }, + original_json: None, + builtin: false, + enabled: Some(false), }, - original_json: None, - builtin: false, - enabled: Some(false), - }, - ImportMcpServerRequest { - name: "playwright".into(), - description: Some("imported".into()), - transport: McpTransport::Stdio { - command: "npx".into(), - args: vec!["@playwright/mcp@latest".into()], - env: HashMap::new(), + ImportMcpServerRequest { + name: "playwright".into(), + description: Some("imported".into()), + transport: McpTransport::Stdio { + command: "npx".into(), + args: vec!["@playwright/mcp@latest".into()], + env: HashMap::new(), + }, + original_json: None, + builtin: false, + enabled: Some(false), }, - original_json: None, - builtin: false, - enabled: Some(false), - }, - ], - }) + ], + }, + ) .await .unwrap(); @@ -1089,17 +1197,20 @@ mod tests { async fn add_server_normalizes_shell_style_stdio_command() { let svc = make_service(); let created = svc - .add_server(CreateMcpServerRequest { - name: "sentry".into(), - description: None, - transport: McpTransport::Stdio { - command: "npx @sentry/mcp-server@latest --organization-slug=demo".into(), - args: vec![], - env: HashMap::new(), + .add_server( + TEST_USER_ID, + CreateMcpServerRequest { + name: "sentry".into(), + description: None, + transport: McpTransport::Stdio { + command: "npx @sentry/mcp-server@latest --organization-slug=demo".into(), + args: vec![], + env: HashMap::new(), + }, + original_json: None, + builtin: false, }, - original_json: None, - builtin: false, - }) + ) .await .unwrap(); @@ -1118,7 +1229,7 @@ mod tests { let mut req = stdio_import_req("enabled-mcp"); req.enabled = Some(true); let result = svc - .batch_import(BatchImportMcpServersRequest { servers: vec![req] }) + .batch_import(TEST_USER_ID, BatchImportMcpServersRequest { servers: vec![req] }) .await .unwrap(); @@ -1130,14 +1241,17 @@ mod tests { async fn batch_import_empty_list() { let svc = make_service(); let req = BatchImportMcpServersRequest { servers: vec![] }; - let results = svc.batch_import(req).await.unwrap(); + let results = svc.batch_import(TEST_USER_ID, req).await.unwrap(); assert!(results.is_empty()); } #[tokio::test] async fn persist_test_result_records_success_status_and_tools() { let svc = make_service(); - let created = svc.add_server(stdio_create_req("persist-success")).await.unwrap(); + let created = svc + .add_server(TEST_USER_ID, stdio_create_req("persist-success")) + .await + .unwrap(); let result = McpConnectionTestResult { success: true, tools: Some(vec![aionui_api_types::McpToolResponse { @@ -1153,9 +1267,11 @@ mod tests { www_authenticate: None, }; - svc.persist_test_result(&created.id, &result).await.unwrap(); + svc.persist_test_result(TEST_USER_ID, &created.id, &result) + .await + .unwrap(); - let updated = svc.get_server(&created.id).await.unwrap(); + let updated = svc.get_server(TEST_USER_ID, &created.id).await.unwrap(); assert_eq!(updated.last_test_status, aionui_common::McpServerStatus::Connected); assert_eq!(updated.tools.unwrap().len(), 1); assert!(updated.last_connected.is_some()); @@ -1164,7 +1280,10 @@ mod tests { #[tokio::test] async fn persist_test_result_records_error_and_clears_tools() { let svc = make_service(); - let created = svc.add_server(stdio_create_req("persist-error")).await.unwrap(); + let created = svc + .add_server(TEST_USER_ID, stdio_create_req("persist-error")) + .await + .unwrap(); let success = McpConnectionTestResult { success: true, @@ -1180,7 +1299,9 @@ mod tests { auth_method: None, www_authenticate: None, }; - svc.persist_test_result(&created.id, &success).await.unwrap(); + svc.persist_test_result(TEST_USER_ID, &created.id, &success) + .await + .unwrap(); let failure = McpConnectionTestResult { success: false, @@ -1192,9 +1313,11 @@ mod tests { auth_method: None, www_authenticate: None, }; - svc.persist_test_result(&created.id, &failure).await.unwrap(); + svc.persist_test_result(TEST_USER_ID, &created.id, &failure) + .await + .unwrap(); - let updated = svc.get_server(&created.id).await.unwrap(); + let updated = svc.get_server(TEST_USER_ID, &created.id).await.unwrap(); assert_eq!(updated.last_test_status, aionui_common::McpServerStatus::Error); assert!(updated.tools.is_none()); assert!(updated.last_connected.is_some()); diff --git a/crates/aionui-mcp/src/sync_service.rs b/crates/aionui-mcp/src/sync_service.rs index 749f6ced2..a2244ceaf 100644 --- a/crates/aionui-mcp/src/sync_service.rs +++ b/crates/aionui-mcp/src/sync_service.rs @@ -34,7 +34,7 @@ impl McpSyncService { /// server configurations. /// /// Agents that are not installed are silently skipped. - pub async fn get_agent_configs(&self) -> Result, McpError> { + pub async fn get_agent_configs(&self, user_id: &str) -> Result, McpError> { let _guard = self.service_lock.lock().await; let mut results = Vec::new(); @@ -46,7 +46,7 @@ impl McpSyncService { continue; } - match adapter.detect_existing().await { + match adapter.detect_existing(user_id).await { Ok(detected) => { let servers = detected.into_iter().map(detected_to_response).collect(); results.push(DetectedMcpServerResponse { @@ -121,6 +121,8 @@ mod tests { use std::collections::HashMap; use std::sync::Mutex as StdMutex; + const TEST_USER_ID: &str = "user-1"; + struct MockAdapter { source: McpSource, installed: bool, @@ -152,7 +154,7 @@ mod tests { Ok(self.installed) } - async fn detect_existing(&self) -> Result, McpError> { + async fn detect_existing(&self, _user_id: &str) -> Result, McpError> { if !self.installed { return Err(McpError::AgentNotInstalled(format!("{:?}", self.source))); } @@ -173,15 +175,15 @@ mod tests { #[async_trait::async_trait] impl IMcpServerRepository for MockRepo { - async fn list(&self) -> Result, DbError> { + async fn list(&self, _user_id: &str) -> Result, DbError> { Ok(Vec::new()) } - async fn find_by_id(&self, _id: &str) -> Result, DbError> { + async fn find_by_id(&self, _user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } - async fn find_by_name(&self, _name: &str) -> Result, DbError> { + async fn find_by_name(&self, _user_id: &str, _name: &str) -> Result, DbError> { Ok(None) } @@ -189,20 +191,30 @@ mod tests { unimplemented!("not needed for detection tests") } - async fn update(&self, _id: &str, _params: UpdateMcpServerParams<'_>) -> Result { + async fn update( + &self, + _user_id: &str, + _id: &str, + _params: UpdateMcpServerParams<'_>, + ) -> Result { unimplemented!("not needed for detection tests") } - async fn delete(&self, _id: &str) -> Result<(), DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), DbError> { unimplemented!("not needed for detection tests") } - async fn batch_upsert(&self, _params_list: &[CreateMcpServerParams<'_>]) -> Result, DbError> { + async fn batch_upsert( + &self, + _user_id: &str, + _params_list: &[CreateMcpServerParams<'_>], + ) -> Result, DbError> { unimplemented!("not needed for detection tests") } async fn update_status( &self, + _user_id: &str, _id: &str, _status: &str, _last_connected: Option, @@ -210,7 +222,7 @@ mod tests { unimplemented!("not needed for detection tests") } - async fn update_tools(&self, _id: &str, _tools: Option<&str>) -> Result<(), DbError> { + async fn update_tools(&self, _user_id: &str, _id: &str, _tools: Option<&str>) -> Result<(), DbError> { unimplemented!("not needed for detection tests") } } @@ -241,7 +253,7 @@ mod tests { let adapter_c = Arc::new(MockAdapter::new(McpSource::Qwen, true).with_existing(vec![])); let svc = make_service(vec![adapter_a, adapter_b, adapter_c]); - let configs = svc.get_agent_configs().await.unwrap(); + let configs = svc.get_agent_configs(TEST_USER_ID).await.unwrap(); assert_eq!(configs.len(), 2); assert_eq!(configs[0].source, McpSource::Claude); @@ -267,7 +279,7 @@ mod tests { #[tokio::test] async fn get_agent_configs_no_adapters() { let svc = make_service(vec![]); - let configs = svc.get_agent_configs().await.unwrap(); + let configs = svc.get_agent_configs(TEST_USER_ID).await.unwrap(); assert!(configs.is_empty()); } diff --git a/crates/aionui-mcp/src/types.rs b/crates/aionui-mcp/src/types.rs index 149bb7a73..13ac082ad 100644 --- a/crates/aionui-mcp/src/types.rs +++ b/crates/aionui-mcp/src/types.rs @@ -442,6 +442,7 @@ mod tests { fn make_test_row(transport_type: &str, transport_config: &str, tools: Option<&str>, status: &str) -> McpServerRow { McpServerRow { id: "mcp_123".into(), + user_id: "user-1".into(), name: "test-server".into(), description: Some("A test server".into()), enabled: true, diff --git a/crates/aionui-mcp/tests/adapter_integration.rs b/crates/aionui-mcp/tests/adapter_integration.rs index 925b72dc2..94ed3e086 100644 --- a/crates/aionui-mcp/tests/adapter_integration.rs +++ b/crates/aionui-mcp/tests/adapter_integration.rs @@ -6,6 +6,8 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex}; +const TEST_USER_ID: &str = "system_default_user"; + use aionui_common::McpSource; use aionui_mcp::{DetectedServer, McpAgentAdapter, McpError, McpServerTransport}; @@ -39,7 +41,7 @@ impl McpAgentAdapter for InMemoryAdapter { Ok(self.installed) } - async fn detect_existing(&self) -> Result, McpError> { + async fn detect_existing(&self, _user_id: &str) -> Result, McpError> { if !self.installed { return Err(McpError::AgentNotInstalled(format!("{:?}", self.source))); } @@ -81,7 +83,7 @@ async fn trait_object_safety_with_arc() { assert_eq!(adapter.source(), McpSource::Claude); assert!(adapter.is_installed().await.unwrap()); - assert!(adapter.detect_existing().await.unwrap().is_empty()); + assert!(adapter.detect_existing(TEST_USER_ID).await.unwrap().is_empty()); } #[tokio::test] @@ -103,7 +105,7 @@ async fn full_lifecycle_install_detect_remove() { adapter.install_server("server-b", &t2).await.unwrap(); // Detect both - let detected = adapter.detect_existing().await.unwrap(); + let detected = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert_eq!(detected.len(), 2); let names: Vec<&str> = detected.iter().map(|s| s.name.as_str()).collect(); @@ -112,13 +114,13 @@ async fn full_lifecycle_install_detect_remove() { // Remove one adapter.remove_server("server-a").await.unwrap(); - let detected = adapter.detect_existing().await.unwrap(); + let detected = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert_eq!(detected.len(), 1); assert_eq!(detected[0].name, "server-b"); // Remove the other adapter.remove_server("server-b").await.unwrap(); - let detected = adapter.detect_existing().await.unwrap(); + let detected = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert!(detected.is_empty()); } @@ -140,7 +142,7 @@ async fn install_replaces_existing_by_name() { adapter.install_server("my-server", &t1).await.unwrap(); adapter.install_server("my-server", &t2).await.unwrap(); - let detected = adapter.detect_existing().await.unwrap(); + let detected = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert_eq!(detected.len(), 1); assert_eq!(detected[0].transport, t2); } @@ -158,7 +160,7 @@ async fn not_installed_errors() { assert!(!adapter.is_installed().await.unwrap()); - let err = adapter.detect_existing().await.unwrap_err(); + let err = adapter.detect_existing(TEST_USER_ID).await.unwrap_err(); assert!(matches!(err, McpError::AgentNotInstalled(_))); let transport = McpServerTransport::Stdio { @@ -187,6 +189,6 @@ async fn multiple_adapters_independent() { claude.install_server("shared-server", &transport).await.unwrap(); // Claude has the server, Gemini does not - assert_eq!(claude.detect_existing().await.unwrap().len(), 1); - assert!(gemini.detect_existing().await.unwrap().is_empty()); + assert_eq!(claude.detect_existing(TEST_USER_ID).await.unwrap().len(), 1); + assert!(gemini.detect_existing(TEST_USER_ID).await.unwrap().is_empty()); } diff --git a/crates/aionui-mcp/tests/file_adapter_integration.rs b/crates/aionui-mcp/tests/file_adapter_integration.rs index 291dd6fbb..5c04300ea 100644 --- a/crates/aionui-mcp/tests/file_adapter_integration.rs +++ b/crates/aionui-mcp/tests/file_adapter_integration.rs @@ -9,6 +9,8 @@ use std::collections::HashMap; use std::sync::Arc; +const TEST_USER_ID: &str = "system_default_user"; + use aionui_common::McpSource; use aionui_mcp::{AionuiAdapter, McpAgentAdapter, McpServerTransport}; @@ -35,36 +37,65 @@ mod aionui { #[async_trait::async_trait] impl IMcpServerRepository for MockRepo { - async fn list(&self) -> Result, DbError> { - Ok(self.servers.lock().await.clone()) + async fn list(&self, user_id: &str) -> Result, DbError> { + Ok(self + .servers + .lock() + .await + .iter() + .filter(|server| server.user_id == user_id) + .cloned() + .collect()) } - async fn find_by_id(&self, id: &str) -> Result, DbError> { - Ok(self.servers.lock().await.iter().find(|s| s.id == id).cloned()) + async fn find_by_id(&self, user_id: &str, id: &str) -> Result, DbError> { + Ok(self + .servers + .lock() + .await + .iter() + .find(|s| s.user_id == user_id && s.id == id) + .cloned()) } - async fn find_by_name(&self, name: &str) -> Result, DbError> { - Ok(self.servers.lock().await.iter().find(|s| s.name == name).cloned()) + async fn find_by_name(&self, user_id: &str, name: &str) -> Result, DbError> { + Ok(self + .servers + .lock() + .await + .iter() + .find(|s| s.user_id == user_id && s.name == name) + .cloned()) } async fn create(&self, _p: CreateMcpServerParams<'_>) -> Result { unimplemented!() } - async fn update(&self, _id: &str, _p: UpdateMcpServerParams<'_>) -> Result { + async fn update( + &self, + _user_id: &str, + _id: &str, + _p: UpdateMcpServerParams<'_>, + ) -> Result { unimplemented!() } - async fn delete(&self, _id: &str) -> Result<(), DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), DbError> { unimplemented!() } - async fn batch_upsert(&self, _s: &[CreateMcpServerParams<'_>]) -> Result, DbError> { + async fn batch_upsert( + &self, + _user_id: &str, + _s: &[CreateMcpServerParams<'_>], + ) -> Result, DbError> { unimplemented!() } async fn update_status( &self, + _user_id: &str, _id: &str, _s: &str, _lc: Option, @@ -72,13 +103,14 @@ mod aionui { unimplemented!() } - async fn update_tools(&self, _id: &str, _t: Option<&str>) -> Result<(), DbError> { + async fn update_tools(&self, _user_id: &str, _id: &str, _t: Option<&str>) -> Result<(), DbError> { unimplemented!() } } fn make_row(name: &str, t_type: &str, t_config: &str) -> McpServerRow { McpServerRow { + user_id: TEST_USER_ID.to_string(), id: format!("mcp_{name}"), name: name.to_owned(), description: None, @@ -120,7 +152,7 @@ mod aionui { let repo = Arc::new(MockRepo::new(rows)); let adapter = AionuiAdapter::new(repo); - let servers = adapter.detect_existing().await.unwrap(); + let servers = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert_eq!(servers.len(), 3); assert_eq!(servers[0].name, "stdio-srv"); assert_eq!(servers[1].name, "http-srv"); @@ -135,7 +167,7 @@ mod aionui { async fn detect_empty_db_returns_empty() { let repo = Arc::new(MockRepo::new(vec![])); let adapter = AionuiAdapter::new(repo); - let servers = adapter.detect_existing().await.unwrap(); + let servers = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert!(servers.is_empty()); } @@ -152,7 +184,7 @@ mod aionui { adapter.install_server("test", &transport).await.unwrap(); // DB should still be empty since install is a no-op - let servers = adapter.detect_existing().await.unwrap(); + let servers = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert!(servers.is_empty()); } @@ -165,7 +197,7 @@ mod aionui { adapter.remove_server("srv").await.unwrap(); // Server should still be in DB since remove is a no-op - let servers = adapter.detect_existing().await.unwrap(); + let servers = adapter.detect_existing(TEST_USER_ID).await.unwrap(); assert_eq!(servers.len(), 1); } diff --git a/crates/aionui-mcp/tests/oauth_integration.rs b/crates/aionui-mcp/tests/oauth_integration.rs index 17bfbe366..d77690568 100644 --- a/crates/aionui-mcp/tests/oauth_integration.rs +++ b/crates/aionui-mcp/tests/oauth_integration.rs @@ -10,6 +10,8 @@ use std::sync::Arc; use aionui_db::{IOAuthTokenRepository, SqliteOAuthTokenRepository, UpsertOAuthTokenParams}; use aionui_mcp::McpOAuthService; +const TEST_USER_ID: &str = "system_default_user"; + async fn make_service() -> (McpOAuthService, Arc) { let db = aionui_db::init_database_memory().await.unwrap(); let repo: Arc = Arc::new(SqliteOAuthTokenRepository::new(db.pool().clone())); @@ -26,7 +28,10 @@ async fn make_service() -> (McpOAuthService, Arc) { #[tokio::test] async fn check_status_unauthenticated_returns_false() { let (svc, _repo) = make_service().await; - let status = svc.check_oauth_status("https://new-server.example.com").await.unwrap(); + let status = svc + .check_oauth_status(TEST_USER_ID, "https://new-server.example.com") + .await + .unwrap(); assert!(!status.authenticated); } @@ -40,6 +45,7 @@ async fn check_status_authenticated_returns_true() { // Seed a valid token. repo.upsert(UpsertOAuthTokenParams { + user_id: TEST_USER_ID, server_url: "https://mcp.example.com", access_token: "access_123", refresh_token: Some("refresh_456"), @@ -50,7 +56,10 @@ async fn check_status_authenticated_returns_true() { .await .unwrap(); - let status = svc.check_oauth_status("https://mcp.example.com").await.unwrap(); + let status = svc + .check_oauth_status(TEST_USER_ID, "https://mcp.example.com") + .await + .unwrap(); assert!(status.authenticated); } @@ -63,6 +72,7 @@ async fn check_status_expired_token_returns_false() { let (svc, repo) = make_service().await; repo.upsert(UpsertOAuthTokenParams { + user_id: TEST_USER_ID, server_url: "https://expired.example.com", access_token: "old_token", refresh_token: None, @@ -73,7 +83,10 @@ async fn check_status_expired_token_returns_false() { .await .unwrap(); - let status = svc.check_oauth_status("https://expired.example.com").await.unwrap(); + let status = svc + .check_oauth_status(TEST_USER_ID, "https://expired.example.com") + .await + .unwrap(); assert!(!status.authenticated); } @@ -86,6 +99,7 @@ async fn check_status_no_expiry_treated_as_valid() { let (svc, repo) = make_service().await; repo.upsert(UpsertOAuthTokenParams { + user_id: TEST_USER_ID, server_url: "https://no-expiry.example.com", access_token: "no_exp_token", refresh_token: None, @@ -95,7 +109,10 @@ async fn check_status_no_expiry_treated_as_valid() { .await .unwrap(); - let status = svc.check_oauth_status("https://no-expiry.example.com").await.unwrap(); + let status = svc + .check_oauth_status(TEST_USER_ID, "https://no-expiry.example.com") + .await + .unwrap(); assert!(status.authenticated); } @@ -108,6 +125,7 @@ async fn get_authenticated_servers_returns_all_urls() { let (svc, repo) = make_service().await; repo.upsert(UpsertOAuthTokenParams { + user_id: TEST_USER_ID, server_url: "https://a.example.com", access_token: "tok_a", refresh_token: None, @@ -118,6 +136,7 @@ async fn get_authenticated_servers_returns_all_urls() { .unwrap(); repo.upsert(UpsertOAuthTokenParams { + user_id: TEST_USER_ID, server_url: "https://b.example.com", access_token: "tok_b", refresh_token: None, @@ -127,7 +146,7 @@ async fn get_authenticated_servers_returns_all_urls() { .await .unwrap(); - let urls = svc.get_authenticated_servers().await.unwrap(); + let urls = svc.get_authenticated_servers(TEST_USER_ID).await.unwrap(); assert_eq!(urls.len(), 2); assert!(urls.contains(&"https://a.example.com".to_string())); assert!(urls.contains(&"https://b.example.com".to_string())); @@ -136,7 +155,7 @@ async fn get_authenticated_servers_returns_all_urls() { #[tokio::test] async fn get_authenticated_servers_empty_when_no_tokens() { let (svc, _repo) = make_service().await; - let urls = svc.get_authenticated_servers().await.unwrap(); + let urls = svc.get_authenticated_servers(TEST_USER_ID).await.unwrap(); assert!(urls.is_empty()); } @@ -148,7 +167,7 @@ async fn get_authenticated_servers_empty_when_no_tokens() { async fn login_invalid_url_returns_error() { let (svc, _repo) = make_service().await; // This URL won't have .well-known endpoints. - let result = svc.login("https://127.0.0.1:1").await; + let result = svc.login(TEST_USER_ID, "https://127.0.0.1:1").await; // Should return an McpError::OAuth about discovery failure. assert!(result.is_err()); } @@ -162,6 +181,7 @@ async fn logout_deletes_stored_token() { let (svc, repo) = make_service().await; repo.upsert(UpsertOAuthTokenParams { + user_id: TEST_USER_ID, server_url: "https://logout.example.com", access_token: "to_delete", refresh_token: None, @@ -172,14 +192,20 @@ async fn logout_deletes_stored_token() { .unwrap(); // Verify token exists. - let status = svc.check_oauth_status("https://logout.example.com").await.unwrap(); + let status = svc + .check_oauth_status(TEST_USER_ID, "https://logout.example.com") + .await + .unwrap(); assert!(status.authenticated); // Logout. - svc.logout("https://logout.example.com").await.unwrap(); + svc.logout(TEST_USER_ID, "https://logout.example.com").await.unwrap(); // Verify token is gone. - let status = svc.check_oauth_status("https://logout.example.com").await.unwrap(); + let status = svc + .check_oauth_status(TEST_USER_ID, "https://logout.example.com") + .await + .unwrap(); assert!(!status.authenticated); } @@ -191,7 +217,9 @@ async fn logout_deletes_stored_token() { async fn logout_idempotent_for_unauthenticated() { let (svc, _repo) = make_service().await; // Should not error. - svc.logout("https://never-authed.example.com").await.unwrap(); + svc.logout(TEST_USER_ID, "https://never-authed.example.com") + .await + .unwrap(); } // --------------------------------------------------------------------------- @@ -201,7 +229,10 @@ async fn logout_idempotent_for_unauthenticated() { #[tokio::test] async fn get_token_returns_none_for_unknown_url() { let (svc, _repo) = make_service().await; - let token = svc.get_token("https://unknown.example.com").await.unwrap(); + let token = svc + .get_token(TEST_USER_ID, "https://unknown.example.com") + .await + .unwrap(); assert!(token.is_none()); } @@ -210,6 +241,7 @@ async fn get_token_returns_access_token_when_valid() { let (svc, repo) = make_service().await; repo.upsert(UpsertOAuthTokenParams { + user_id: TEST_USER_ID, server_url: "https://valid.example.com", access_token: "my_access_token", refresh_token: None, @@ -219,7 +251,7 @@ async fn get_token_returns_access_token_when_valid() { .await .unwrap(); - let token = svc.get_token("https://valid.example.com").await.unwrap(); + let token = svc.get_token(TEST_USER_ID, "https://valid.example.com").await.unwrap(); assert_eq!(token.as_deref(), Some("my_access_token")); } @@ -228,6 +260,7 @@ async fn get_token_returns_expired_token_when_no_refresh_token() { let (svc, repo) = make_service().await; repo.upsert(UpsertOAuthTokenParams { + user_id: TEST_USER_ID, server_url: "https://expired.example.com", access_token: "old_access", refresh_token: None, @@ -238,7 +271,10 @@ async fn get_token_returns_expired_token_when_no_refresh_token() { .unwrap(); // With no refresh_token, returns the expired token as-is. - let token = svc.get_token("https://expired.example.com").await.unwrap(); + let token = svc + .get_token(TEST_USER_ID, "https://expired.example.com") + .await + .unwrap(); assert_eq!(token.as_deref(), Some("old_access")); } @@ -247,6 +283,7 @@ async fn get_token_returns_no_expiry_token() { let (svc, repo) = make_service().await; repo.upsert(UpsertOAuthTokenParams { + user_id: TEST_USER_ID, server_url: "https://noexp.example.com", access_token: "forever_token", refresh_token: None, @@ -256,6 +293,6 @@ async fn get_token_returns_no_expiry_token() { .await .unwrap(); - let token = svc.get_token("https://noexp.example.com").await.unwrap(); + let token = svc.get_token(TEST_USER_ID, "https://noexp.example.com").await.unwrap(); assert_eq!(token.as_deref(), Some("forever_token")); } diff --git a/crates/aionui-mcp/tests/service_integration.rs b/crates/aionui-mcp/tests/service_integration.rs index c1a32330f..0b93f686d 100644 --- a/crates/aionui-mcp/tests/service_integration.rs +++ b/crates/aionui-mcp/tests/service_integration.rs @@ -5,16 +5,25 @@ use std::collections::HashMap; use std::sync::Arc; +const TEST_USER_ID: &str = "system_default_user"; +const OTHER_USERNAME: &str = "user-2"; + use aionui_api_types::{ BatchImportMcpServersRequest, CreateMcpServerRequest, ImportMcpServerRequest, McpTransport, UpdateMcpServerRequest, }; -use aionui_db::SqliteMcpServerRepository; +use aionui_db::{IUserRepository, SqliteMcpServerRepository, SqliteUserRepository}; use aionui_mcp::{McpConfigService, McpError}; async fn make_service() -> McpConfigService { + make_service_with_other_user().await.0 +} + +async fn make_service_with_other_user() -> (McpConfigService, String) { let db = aionui_db::init_database_memory().await.unwrap(); + let user_repo = SqliteUserRepository::new(db.pool().clone()); + let other_user = user_repo.create_user(OTHER_USERNAME, "hash").await.unwrap(); let repo = Arc::new(SqliteMcpServerRepository::new(db.pool().clone())); - McpConfigService::new(repo) + (McpConfigService::new(repo), other_user.id) } fn stdio_req(name: &str) -> CreateMcpServerRequest { @@ -80,21 +89,21 @@ fn http_import_req(name: &str) -> ImportMcpServerRequest { #[tokio::test] async fn create_and_get_stdio_server() { let svc = make_service().await; - let resp = svc.add_server(stdio_req("test-stdio")).await.unwrap(); + let resp = svc.add_server(TEST_USER_ID, stdio_req("test-stdio")).await.unwrap(); assert!(resp.id.starts_with("mcp_")); assert_eq!(resp.name, "test-stdio"); assert!(!resp.enabled); assert_eq!(resp.description.as_deref(), Some("test")); - let found = svc.get_server(&resp.id).await.unwrap(); + let found = svc.get_server(TEST_USER_ID, &resp.id).await.unwrap(); assert_eq!(found.id, resp.id); } #[tokio::test] async fn create_http_with_headers() { let svc = make_service().await; - let resp = svc.add_server(http_req("test-http")).await.unwrap(); + let resp = svc.add_server(TEST_USER_ID, http_req("test-http")).await.unwrap(); match resp.transport { McpTransport::Http { ref url, ref headers } => { @@ -108,8 +117,8 @@ async fn create_http_with_headers() { #[tokio::test] async fn create_same_name_upserts() { let svc = make_service().await; - let first = svc.add_server(stdio_req("dup")).await.unwrap(); - let second = svc.add_server(http_req("dup")).await.unwrap(); + let first = svc.add_server(TEST_USER_ID, stdio_req("dup")).await.unwrap(); + let second = svc.add_server(TEST_USER_ID, http_req("dup")).await.unwrap(); assert_eq!(first.id, second.id); match second.transport { @@ -125,24 +134,66 @@ async fn create_same_name_upserts() { #[tokio::test] async fn list_empty() { let svc = make_service().await; - assert!(svc.list_servers().await.unwrap().is_empty()); + assert!(svc.list_servers(TEST_USER_ID).await.unwrap().is_empty()); } #[tokio::test] async fn list_returns_all() { let svc = make_service().await; - svc.add_server(stdio_req("a")).await.unwrap(); - svc.add_server(http_req("b")).await.unwrap(); - assert_eq!(svc.list_servers().await.unwrap().len(), 2); + svc.add_server(TEST_USER_ID, stdio_req("a")).await.unwrap(); + svc.add_server(TEST_USER_ID, http_req("b")).await.unwrap(); + assert_eq!(svc.list_servers(TEST_USER_ID).await.unwrap().len(), 2); } #[tokio::test] async fn get_not_found() { let svc = make_service().await; - let err = svc.get_server("nonexistent").await.unwrap_err(); + let err = svc.get_server(TEST_USER_ID, "nonexistent").await.unwrap_err(); assert!(matches!(err, McpError::NotFound(_))); } +#[tokio::test] +async fn servers_are_scoped_by_current_user() { + let (svc, other_user_id) = make_service_with_other_user().await; + let owner_server = svc.add_server(TEST_USER_ID, stdio_req("shared")).await.unwrap(); + let other_server = svc.add_server(&other_user_id, http_req("shared")).await.unwrap(); + + assert_ne!(owner_server.id, other_server.id); + assert_eq!(svc.list_servers(TEST_USER_ID).await.unwrap().len(), 1); + assert_eq!(svc.list_servers(&other_user_id).await.unwrap().len(), 1); + + assert!(matches!( + svc.get_server(TEST_USER_ID, &other_server.id).await.unwrap_err(), + McpError::NotFound(_) + )); + assert!(matches!( + svc.edit_server( + TEST_USER_ID, + &other_server.id, + UpdateMcpServerRequest { + name: None, + description: Some(Some("forged".into())), + transport: None, + original_json: None, + builtin: None, + }, + ) + .await + .unwrap_err(), + McpError::NotFound(_) + )); + assert!(matches!( + svc.delete_server(TEST_USER_ID, &other_server.id).await.unwrap_err(), + McpError::NotFound(_) + )); + + svc.delete_server(&other_user_id, &other_server.id).await.unwrap(); + assert_eq!( + svc.get_server(TEST_USER_ID, &owner_server.id).await.unwrap().id, + owner_server.id + ); +} + // --------------------------------------------------------------------------- // Update // --------------------------------------------------------------------------- @@ -150,10 +201,11 @@ async fn get_not_found() { #[tokio::test] async fn edit_name_is_rejected() { let svc = make_service().await; - let created = svc.add_server(stdio_req("old")).await.unwrap(); + let created = svc.add_server(TEST_USER_ID, stdio_req("old")).await.unwrap(); let err = svc .edit_server( + TEST_USER_ID, &created.id, UpdateMcpServerRequest { name: Some("new".into()), @@ -171,10 +223,11 @@ async fn edit_name_is_rejected() { #[tokio::test] async fn edit_transport() { let svc = make_service().await; - let created = svc.add_server(stdio_req("test")).await.unwrap(); + let created = svc.add_server(TEST_USER_ID, stdio_req("test")).await.unwrap(); let updated = svc .edit_server( + TEST_USER_ID, &created.id, UpdateMcpServerRequest { name: None, @@ -198,11 +251,12 @@ async fn edit_transport() { #[tokio::test] async fn edit_clears_description() { let svc = make_service().await; - let created = svc.add_server(stdio_req("test")).await.unwrap(); + let created = svc.add_server(TEST_USER_ID, stdio_req("test")).await.unwrap(); assert!(created.description.is_some()); let updated = svc .edit_server( + TEST_USER_ID, &created.id, UpdateMcpServerRequest { name: None, @@ -222,6 +276,7 @@ async fn edit_not_found() { let svc = make_service().await; let err = svc .edit_server( + TEST_USER_ID, "nonexistent", UpdateMcpServerRequest { name: Some("x".into()), @@ -239,11 +294,12 @@ async fn edit_not_found() { #[tokio::test] async fn edit_name_conflict() { let svc = make_service().await; - svc.add_server(stdio_req("a")).await.unwrap(); - let b = svc.add_server(stdio_req("b")).await.unwrap(); + svc.add_server(TEST_USER_ID, stdio_req("a")).await.unwrap(); + let b = svc.add_server(TEST_USER_ID, stdio_req("b")).await.unwrap(); let err = svc .edit_server( + TEST_USER_ID, &b.id, UpdateMcpServerRequest { name: Some("a".into()), @@ -265,28 +321,28 @@ async fn edit_name_conflict() { #[tokio::test] async fn delete_removes_server() { let svc = make_service().await; - let created = svc.add_server(stdio_req("del")).await.unwrap(); - let was_enabled = svc.delete_server(&created.id).await.unwrap(); + let created = svc.add_server(TEST_USER_ID, stdio_req("del")).await.unwrap(); + let was_enabled = svc.delete_server(TEST_USER_ID, &created.id).await.unwrap(); assert!(!was_enabled); - let err = svc.get_server(&created.id).await.unwrap_err(); + let err = svc.get_server(TEST_USER_ID, &created.id).await.unwrap_err(); assert!(matches!(err, McpError::NotFound(_))); } #[tokio::test] async fn delete_enabled_returns_true() { let svc = make_service().await; - let created = svc.add_server(stdio_req("del-en")).await.unwrap(); - svc.toggle_server(&created.id).await.unwrap(); + let created = svc.add_server(TEST_USER_ID, stdio_req("del-en")).await.unwrap(); + svc.toggle_server(TEST_USER_ID, &created.id).await.unwrap(); - let was_enabled = svc.delete_server(&created.id).await.unwrap(); + let was_enabled = svc.delete_server(TEST_USER_ID, &created.id).await.unwrap(); assert!(was_enabled); } #[tokio::test] async fn delete_not_found() { let svc = make_service().await; - let err = svc.delete_server("nonexistent").await.unwrap_err(); + let err = svc.delete_server(TEST_USER_ID, "nonexistent").await.unwrap_err(); assert!(matches!(err, McpError::NotFound(_))); } @@ -297,13 +353,13 @@ async fn delete_not_found() { #[tokio::test] async fn toggle_enables_then_disables() { let svc = make_service().await; - let created = svc.add_server(stdio_req("tog")).await.unwrap(); + let created = svc.add_server(TEST_USER_ID, stdio_req("tog")).await.unwrap(); assert!(!created.enabled); - let toggled = svc.toggle_server(&created.id).await.unwrap(); + let toggled = svc.toggle_server(TEST_USER_ID, &created.id).await.unwrap(); assert!(toggled.enabled); - let toggled_back = svc.toggle_server(&created.id).await.unwrap(); + let toggled_back = svc.toggle_server(TEST_USER_ID, &created.id).await.unwrap(); assert!(!toggled_back.enabled); } @@ -314,7 +370,7 @@ async fn toggle_enables_then_disables() { #[tokio::test] async fn batch_import_creates_and_upserts() { let svc = make_service().await; - svc.add_server(stdio_req("existing")).await.unwrap(); + svc.add_server(TEST_USER_ID, stdio_req("existing")).await.unwrap(); let req = BatchImportMcpServersRequest { servers: vec![ @@ -322,10 +378,10 @@ async fn batch_import_creates_and_upserts() { stdio_import_req("new"), // create ], }; - let results = svc.batch_import(req).await.unwrap(); + let results = svc.batch_import(TEST_USER_ID, req).await.unwrap(); assert_eq!(results.len(), 2); - let all = svc.list_servers().await.unwrap(); + let all = svc.list_servers(TEST_USER_ID).await.unwrap(); assert_eq!(all.len(), 2); } @@ -336,7 +392,7 @@ async fn batch_import_preserves_enabled_in_database() { req.enabled = Some(true); let result = svc - .batch_import(BatchImportMcpServersRequest { servers: vec![req] }) + .batch_import(TEST_USER_ID, BatchImportMcpServersRequest { servers: vec![req] }) .await .unwrap(); @@ -344,7 +400,7 @@ async fn batch_import_preserves_enabled_in_database() { assert_eq!(result[0].name, "enabled-db-mcp"); assert!(result[0].enabled); - let listed = svc.list_servers().await.unwrap(); + let listed = svc.list_servers(TEST_USER_ID).await.unwrap(); assert_eq!(listed.len(), 1); assert!(listed[0].enabled); } diff --git a/crates/aionui-mcp/tests/sync_integration.rs b/crates/aionui-mcp/tests/sync_integration.rs index fec35a246..d212e966f 100644 --- a/crates/aionui-mcp/tests/sync_integration.rs +++ b/crates/aionui-mcp/tests/sync_integration.rs @@ -3,6 +3,8 @@ use std::collections::HashMap; use std::sync::Arc; +const TEST_USER_ID: &str = "system_default_user"; + use aionui_common::McpSource; use aionui_db::SqliteMcpServerRepository; use aionui_mcp::{DetectedServer, McpAgentAdapter, McpError, McpServerTransport, McpSyncService}; @@ -41,7 +43,7 @@ impl McpAgentAdapter for MockAdapter { Ok(self.installed) } - async fn detect_existing(&self) -> Result, McpError> { + async fn detect_existing(&self, _user_id: &str) -> Result, McpError> { if !self.installed { return Err(McpError::AgentNotInstalled(format!("{:?}", self.source))); } @@ -91,7 +93,7 @@ async fn get_agent_configs_returns_installed_agents() { adapter_qwen, ]) .await; - let configs = sync_svc.get_agent_configs().await.unwrap(); + let configs = sync_svc.get_agent_configs(TEST_USER_ID).await.unwrap(); assert_eq!(configs.len(), 2); assert_eq!(configs[0].source, McpSource::Claude); @@ -106,6 +108,6 @@ async fn get_agent_configs_empty_when_none_installed() { let adapter = Arc::new(MockAdapter::new(McpSource::Claude, false)); let sync_svc = make_service(vec![adapter as Arc]).await; - let configs = sync_svc.get_agent_configs().await.unwrap(); + let configs = sync_svc.get_agent_configs(TEST_USER_ID).await.unwrap(); assert!(configs.is_empty()); } diff --git a/crates/aionui-mcp/tests/types_integration.rs b/crates/aionui-mcp/tests/types_integration.rs index 71c819def..178999d82 100644 --- a/crates/aionui-mcp/tests/types_integration.rs +++ b/crates/aionui-mcp/tests/types_integration.rs @@ -5,6 +5,8 @@ use std::collections::HashMap; +const TEST_USER_ID: &str = "system_default_user"; + use aionui_common::McpServerStatus; use aionui_db::models::McpServerRow; use aionui_mcp::{McpServer, McpServerTransport, McpTool}; @@ -15,6 +17,7 @@ use aionui_mcp::{McpServer, McpServerTransport, McpTool}; fn row(transport_type: &str, transport_config: &str, tools: Option<&str>, status: &str) -> McpServerRow { McpServerRow { + user_id: TEST_USER_ID.to_string(), id: "mcp_integration".into(), name: "integration-test".into(), description: Some("Integration test server".into()), diff --git a/crates/aionui-office/src/proxy.rs b/crates/aionui-office/src/proxy.rs index 62953d247..4ab4719b2 100644 --- a/crates/aionui-office/src/proxy.rs +++ b/crates/aionui-office/src/proxy.rs @@ -63,7 +63,19 @@ impl ProxyService { doc_type: DocType, request_headers: &[(String, String)], ) -> Result { - if !self.watch_manager.is_active_port(port, doc_type) { + self.forward_for_user("system_default_user", port, path, doc_type, request_headers) + .await + } + + pub async fn forward_for_user( + &self, + user_id: &str, + port: u16, + path: &str, + doc_type: DocType, + request_headers: &[(String, String)], + ) -> Result { + if !self.watch_manager.is_active_port_for_user(user_id, port, doc_type) { return Err(ProxyError::PortNotActive(port)); } let proxy_base = format!("/api/{}/{}", doc_type.proxy_prefix(), port); @@ -76,7 +88,18 @@ impl ProxyService { path: &str, request_headers: &[(String, String)], ) -> Result { - if !self.watch_manager.is_active_watch_port(port) { + self.forward_watch_for_user("system_default_user", port, path, request_headers) + .await + } + + pub async fn forward_watch_for_user( + &self, + user_id: &str, + port: u16, + path: &str, + request_headers: &[(String, String)], + ) -> Result { + if !self.watch_manager.is_active_watch_port_for_user(user_id, port) { return Err(ProxyError::PortNotActive(port)); } let proxy_base = format!("/api/office-watch-proxy/{port}"); diff --git a/crates/aionui-office/src/routes.rs b/crates/aionui-office/src/routes.rs index 163cd965f..60ac89154 100644 --- a/crates/aionui-office/src/routes.rs +++ b/crates/aionui-office/src/routes.rs @@ -84,54 +84,55 @@ struct ProxyPortPath { async fn start_word_preview( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { - start_preview(state, body, DocType::Word).await + start_preview(state, &user.id, body, DocType::Word).await } async fn stop_word_preview( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { - stop_preview(state, body, DocType::Word).await + stop_preview(state, &user.id, body, DocType::Word).await } async fn start_excel_preview( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { - start_preview(state, body, DocType::Excel).await + start_preview(state, &user.id, body, DocType::Excel).await } async fn stop_excel_preview( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { - stop_preview(state, body, DocType::Excel).await + stop_preview(state, &user.id, body, DocType::Excel).await } async fn start_ppt_preview( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { - start_preview(state, body, DocType::Ppt).await + start_preview(state, &user.id, body, DocType::Ppt).await } async fn stop_ppt_preview( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { - stop_preview(state, body, DocType::Ppt).await + stop_preview(state, &user.id, body, DocType::Ppt).await } async fn start_preview( state: OfficeRouterState, + user_id: &str, body: Result, JsonRejection>, doc_type: DocType, ) -> Result>, ApiError> { @@ -139,7 +140,10 @@ async fn start_preview( let validated_path = validate_office_path(&state, &req.file_path, req.workspace.as_deref())?; let validated_path = validated_path.to_string_lossy().into_owned(); - let result = state.watch_manager.start(&validated_path, doc_type).await; + let result = state + .watch_manager + .start_for_user(user_id, &validated_path, doc_type) + .await; let resp = match result { Ok(port) => { @@ -157,11 +161,15 @@ async fn start_preview( async fn stop_preview( state: OfficeRouterState, + user_id: &str, body: Result, JsonRejection>, doc_type: DocType, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - state.watch_manager.stop(&req.file_path, doc_type).await; + state + .watch_manager + .stop_for_user(user_id, &req.file_path, doc_type) + .await; Ok(Json(ApiResponse::success())) } @@ -169,33 +177,36 @@ async fn stop_preview( async fn list_snapshots( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let snapshots = state.snapshot_service.list(&req.target).await?; + let snapshots = state.snapshot_service.list_for_user(&user.id, &req.target).await?; Ok(Json(ApiResponse::ok(snapshots))) } async fn save_snapshot( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let info = state.snapshot_service.save(&req.target, &req.content).await?; + let info = state + .snapshot_service + .save_for_user(&user.id, &req.target, &req.content) + .await?; Ok(Json(ApiResponse::ok(info))) } async fn get_snapshot_content( State(state): State, - Extension(_user): Extension, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; let result = state .snapshot_service - .get_content(&req.target, &req.snapshot_id) + .get_content_for_user(&user.id, &req.target, &req.snapshot_id) .await?; Ok(Json(ApiResponse::ok(result))) } @@ -262,15 +273,17 @@ fn preview_error_code(error: &OfficeError) -> &'static str { async fn ppt_proxy( State(state): State, + Extension(user): Extension, Path(params): Path, headers: HeaderMap, ) -> Result { let path = params.path.as_deref().unwrap_or("/"); - proxy_forward(state, params.port, path, DocType::Ppt, &headers).await + proxy_forward(state, &user.id, params.port, path, DocType::Ppt, &headers).await } async fn office_watch_proxy( State(state): State, + Extension(user): Extension, Path(params): Path, headers: HeaderMap, ) -> Result { @@ -282,7 +295,7 @@ async fn office_watch_proxy( let proxy_resp = state .proxy_service - .forward_watch(params.port, path, &request_headers) + .forward_watch_for_user(&user.id, params.port, path, &request_headers) .await?; let status = StatusCode::from_u16(proxy_resp.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); @@ -299,6 +312,7 @@ async fn office_watch_proxy( async fn proxy_forward( state: OfficeRouterState, + user_id: &str, port: u16, path: &str, doc_type: DocType, @@ -311,7 +325,7 @@ async fn proxy_forward( let proxy_resp = state .proxy_service - .forward(port, path, doc_type, &request_headers) + .forward_for_user(user_id, port, path, doc_type, &request_headers) .await?; let status = StatusCode::from_u16(proxy_resp.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); diff --git a/crates/aionui-office/src/snapshot.rs b/crates/aionui-office/src/snapshot.rs index acc1e8790..e74c6713e 100644 --- a/crates/aionui-office/src/snapshot.rs +++ b/crates/aionui-office/src/snapshot.rs @@ -10,6 +10,7 @@ use crate::error::OfficeError; const MAX_SNAPSHOTS: usize = 50; const SNAPSHOT_EXT: &str = ".md"; const INDEX_FILE: &str = "index.json"; +const DEFAULT_USER_ID: &str = "system_default_user"; pub struct SnapshotService { base_dir: PathBuf, @@ -23,7 +24,15 @@ impl SnapshotService { } pub async fn list(&self, target: &PreviewHistoryTargetDto) -> Result, OfficeError> { - let dir = self.target_dir(target); + self.list_for_user(DEFAULT_USER_ID, target).await + } + + pub async fn list_for_user( + &self, + user_id: &str, + target: &PreviewHistoryTargetDto, + ) -> Result, OfficeError> { + let dir = self.target_dir_for_user(user_id, target); let index_path = dir.join(INDEX_FILE); if !index_path.exists() { @@ -41,7 +50,16 @@ impl SnapshotService { target: &PreviewHistoryTargetDto, content: &str, ) -> Result { - let dir = self.target_dir(target); + self.save_for_user(DEFAULT_USER_ID, target, content).await + } + + pub async fn save_for_user( + &self, + user_id: &str, + target: &PreviewHistoryTargetDto, + content: &str, + ) -> Result { + let dir = self.target_dir_for_user(user_id, target); tokio::fs::create_dir_all(&dir).await?; let now_ms = current_timestamp_ms(); @@ -76,7 +94,16 @@ impl SnapshotService { target: &PreviewHistoryTargetDto, snapshot_id: &str, ) -> Result, OfficeError> { - let dir = self.target_dir(target); + self.get_content_for_user(DEFAULT_USER_ID, target, snapshot_id).await + } + + pub async fn get_content_for_user( + &self, + user_id: &str, + target: &PreviewHistoryTargetDto, + snapshot_id: &str, + ) -> Result, OfficeError> { + let dir = self.target_dir_for_user(user_id, target); let snapshots: Vec = self.read_index(&dir).await; let Some(info) = snapshots.into_iter().find(|s| s.id == snapshot_id) else { @@ -95,8 +122,8 @@ impl SnapshotService { })) } - fn target_dir(&self, target: &PreviewHistoryTargetDto) -> PathBuf { - let hash = compute_target_hash(target); + fn target_dir_for_user(&self, user_id: &str, target: &PreviewHistoryTargetDto) -> PathBuf { + let hash = compute_target_hash_for_user(user_id, target); self.base_dir.join(hash) } @@ -132,9 +159,17 @@ impl SnapshotService { } } +#[cfg(test)] fn compute_target_hash(target: &PreviewHistoryTargetDto) -> String { + compute_target_hash_for_user(DEFAULT_USER_ID, target) +} + +fn compute_target_hash_for_user(user_id: &str, target: &PreviewHistoryTargetDto) -> String { let mut hasher = Sha1::new(); + hasher.update(user_id.as_bytes()); + hasher.update(b"\0"); + hasher.update( serde_json::to_value(target.content_type) .unwrap_or_default() @@ -444,6 +479,36 @@ mod tests { assert_eq!(list2.len(), 2); } + #[tokio::test] + async fn service_same_target_isolated_by_user() { + let tmp = tempfile::tempdir().unwrap(); + let svc = SnapshotService::new(tmp.path()); + let target = make_target(PreviewContentType::Markdown, Some("/a.md")); + + let user_a = svc.save_for_user("user-a", &target, "content-a").await.unwrap(); + let user_b = svc.save_for_user("user-b", &target, "content-b").await.unwrap(); + + let list_a = svc.list_for_user("user-a", &target).await.unwrap(); + let list_b = svc.list_for_user("user-b", &target).await.unwrap(); + assert_eq!(list_a.len(), 1); + assert_eq!(list_b.len(), 1); + assert_eq!(list_a[0].id, user_a.id); + assert_eq!(list_b[0].id, user_b.id); + + assert!( + svc.get_content_for_user("user-a", &target, &user_b.id) + .await + .unwrap() + .is_none() + ); + assert!( + svc.get_content_for_user("user-b", &target, &user_a.id) + .await + .unwrap() + .is_none() + ); + } + fn make_target(content_type: PreviewContentType, file_path: Option<&str>) -> PreviewHistoryTargetDto { PreviewHistoryTargetDto { content_type, diff --git a/crates/aionui-office/src/watch_manager.rs b/crates/aionui-office/src/watch_manager.rs index 7af031376..d6b2ce8a5 100644 --- a/crates/aionui-office/src/watch_manager.rs +++ b/crates/aionui-office/src/watch_manager.rs @@ -49,6 +49,7 @@ pub trait ProcessHandle: Send + Sync { // --------------------------------------------------------------------------- struct WatchSession { + user_id: String, port: u16, process: Box, file_path: String, @@ -61,12 +62,29 @@ struct WatchSession { // --------------------------------------------------------------------------- pub struct OfficecliWatchManager { - sessions: DashMap, + sessions: DashMap, spawner: Arc, broadcaster: Arc, last_version_check: Mutex>, } +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct WatchSessionKey { + user_id: String, + doc_type: DocType, + resolved_path: String, +} + +impl WatchSessionKey { + fn new(user_id: &str, resolved_path: &str, doc_type: DocType) -> Self { + Self { + user_id: user_id.to_owned(), + doc_type, + resolved_path: resolved_path.to_owned(), + } + } +} + impl OfficecliWatchManager { pub fn new(spawner: Arc, broadcaster: Arc) -> Self { Self { @@ -78,8 +96,12 @@ impl OfficecliWatchManager { } pub async fn start(&self, file_path: &str, doc_type: DocType) -> Result { + self.start_for_user("system_default_user", file_path, doc_type).await + } + + pub async fn start_for_user(&self, user_id: &str, file_path: &str, doc_type: DocType) -> Result { let resolved = resolve_path(file_path)?; - let key = session_key(&resolved, doc_type); + let key = session_key(user_id, &resolved, doc_type); if let Some(entry) = self.sessions.get(&key) { if !entry.aborted && entry.process.is_alive() { @@ -89,20 +111,25 @@ impl OfficecliWatchManager { self.sessions.remove(&key); } - self.broadcast_status(doc_type, PreviewState::Starting, None); + self.broadcast_status_for_user(user_id, doc_type, PreviewState::Starting, None); - let result = self.try_start(&resolved, doc_type).await; + let result = self.try_start(user_id, &resolved, doc_type).await; match &result { Ok(port) => { - self.broadcast_status(doc_type, PreviewState::Ready, None); + self.broadcast_status_for_user(user_id, doc_type, PreviewState::Ready, None); if doc_type == DocType::Ppt { self.maybe_check_update(doc_type).await; } Ok(*port) } Err(e) => { - self.broadcast_status(doc_type, PreviewState::Error, Some(public_preview_error_message(e))); + self.broadcast_status_for_user( + user_id, + doc_type, + PreviewState::Error, + Some(public_preview_error_message(e)), + ); Err(match e { OfficeError::OfficecliNotFound => OfficeError::OfficecliNotFound, OfficeError::InstallFailed(m) => OfficeError::InstallFailed(m.clone()), @@ -118,17 +145,21 @@ impl OfficecliWatchManager { } } - async fn try_start(&self, resolved: &str, doc_type: DocType) -> Result { + async fn try_start(&self, user_id: &str, resolved: &str, doc_type: DocType) -> Result { for attempt in 1..=START_PORT_MAX_ATTEMPTS { let port = allocate_port()?; - match self.spawn_officecli_with_install(resolved, port, doc_type).await { + match self + .spawn_officecli_with_install(user_id, resolved, port, doc_type) + .await + { Ok(process) => { self.poll_port_ready(port, resolved).await?; - let key = session_key(resolved, doc_type); + let key = session_key(user_id, resolved, doc_type); self.sessions.insert( key, WatchSession { + user_id: user_id.to_owned(), port, process, file_path: resolved.to_owned(), @@ -158,6 +189,7 @@ impl OfficecliWatchManager { async fn spawn_officecli_with_install( &self, + user_id: &str, resolved: &str, port: u16, doc_type: DocType, @@ -165,7 +197,7 @@ impl OfficecliWatchManager { match self.spawner.spawn_officecli(resolved, port, doc_type).await { Ok(process) => Ok(process), Err(OfficeError::OfficecliNotFound) => { - self.broadcast_status(doc_type, PreviewState::Installing, None); + self.broadcast_status_for_user(user_id, doc_type, PreviewState::Installing, None); self.spawner.install_officecli().await?; self.spawner.spawn_officecli(resolved, port, doc_type).await } @@ -184,11 +216,15 @@ impl OfficecliWatchManager { } pub async fn stop(&self, file_path: &str, doc_type: DocType) { + self.stop_for_user("system_default_user", file_path, doc_type).await; + } + + pub async fn stop_for_user(&self, user_id: &str, file_path: &str, doc_type: DocType) { let resolved = match resolve_path(file_path) { Ok(p) => p, Err(_) => return, }; - let key = session_key(&resolved, doc_type); + let key = session_key(user_id, &resolved, doc_type); tokio::time::sleep(Duration::from_millis(STOP_DELAY_MS)).await; @@ -209,18 +245,51 @@ impl OfficecliWatchManager { self.sessions.clear(); } + pub fn stop_all_for_user(&self, user_id: &str) -> usize { + let keys: Vec = self + .sessions + .iter() + .filter(|entry| entry.key().user_id == user_id) + .map(|entry| entry.key().clone()) + .collect(); + let stopped = keys.len(); + + for key in keys { + if let Some((_, session)) = self.sessions.remove(&key) { + session.process.kill(); + } + } + + if stopped > 0 { + tracing::info!(user_id = %user_id, stopped, "stopped office preview sessions for user"); + } + stopped + } + pub fn is_active_port(&self, port: u16, doc_type: DocType) -> bool { self.sessions .iter() .any(|entry| entry.port == port && entry.doc_type == doc_type) } + pub fn is_active_port_for_user(&self, user_id: &str, port: u16, doc_type: DocType) -> bool { + self.sessions + .iter() + .any(|entry| entry.user_id == user_id && entry.port == port && entry.doc_type == doc_type) + } + pub fn is_active_watch_port(&self, port: u16) -> bool { self.sessions .iter() .any(|entry| entry.port == port && matches!(entry.doc_type, DocType::Word | DocType::Excel)) } + pub fn is_active_watch_port_for_user(&self, user_id: &str, port: u16) -> bool { + self.sessions.iter().any(|entry| { + entry.user_id == user_id && entry.port == port && matches!(entry.doc_type, DocType::Word | DocType::Excel) + }) + } + pub fn active_session_count(&self) -> usize { self.sessions.len() } @@ -243,16 +312,23 @@ impl OfficecliWatchManager { } } - fn broadcast_status(&self, doc_type: DocType, state: PreviewState, message: Option) { + fn broadcast_status_for_user( + &self, + user_id: &str, + doc_type: DocType, + state: PreviewState, + message: Option, + ) { let event_name = format!("{}.status", doc_type.event_prefix()); let payload = PreviewStatusEvent { state, message }; - let data = match serde_json::to_value(payload) { + let mut data = match serde_json::to_value(payload) { Ok(v) => v, Err(e) => { tracing::error!("failed to serialize preview status: {e}"); return; } }; + data["user_id"] = serde_json::Value::String(user_id.to_owned()); self.broadcaster.broadcast(WebSocketMessage::new(event_name, data)); } } @@ -424,8 +500,8 @@ fn resolve_path(file_path: &str) -> Result { Ok(resolved.to_string_lossy().into_owned()) } -fn session_key(resolved_path: &str, doc_type: DocType) -> String { - format!("{doc_type}:{resolved_path}") +fn session_key(user_id: &str, resolved_path: &str, doc_type: DocType) -> WatchSessionKey { + WatchSessionKey::new(user_id, resolved_path, doc_type) } fn is_port_in_use_start_failure(error: &OfficeError) -> bool { @@ -604,18 +680,34 @@ mod tests { } #[test] - fn session_key_format() { - let key = session_key("/path/to/doc.docx", DocType::Word); - assert_eq!(key, "word:/path/to/doc.docx"); + fn session_key_preserves_fields() { + let key = session_key("user-1", "/path/to/doc.docx", DocType::Word); + assert_eq!(key.user_id, "user-1"); + assert_eq!(key.doc_type, DocType::Word); + assert_eq!(key.resolved_path, "/path/to/doc.docx"); } #[test] fn session_key_different_doc_types() { - let k1 = session_key("/a.docx", DocType::Word); - let k2 = session_key("/a.docx", DocType::Excel); + let k1 = session_key("user-1", "/a.docx", DocType::Word); + let k2 = session_key("user-1", "/a.docx", DocType::Excel); + assert_ne!(k1, k2); + } + + #[test] + fn session_key_different_users() { + let k1 = session_key("user-1", "/a.docx", DocType::Word); + let k2 = session_key("user-2", "/a.docx", DocType::Word); assert_ne!(k1, k2); } + #[test] + fn session_key_keeps_user_type_and_path_boundaries() { + let first = session_key("user:word", "/a.docx", DocType::Excel); + let second = session_key("user", "word:/a.docx", DocType::Excel); + assert_ne!(first, second); + } + #[test] fn install_failed_public_preview_message_is_sanitized() { let err = OfficeError::InstallFailed("installer stderr".into()); @@ -692,6 +784,36 @@ mod tests { assert_eq!(spawner.spawn_count.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn start_same_file_for_different_users_creates_independent_sessions() { + let spawner = Arc::new(MockSpawner::new()); + let broadcaster = Arc::new(RecordingBroadcaster::new()); + let mgr = make_manager(Arc::clone(&spawner), Arc::clone(&broadcaster)); + + let dir = tempfile::tempdir().unwrap(); + let file = dir.path().join("test.docx"); + std::fs::write(&file, b"test").unwrap(); + let path = file.to_str().unwrap(); + + let user_a_port = mgr.start_for_user("user-a", path, DocType::Word).await.unwrap(); + let user_b_port = mgr.start_for_user("user-b", path, DocType::Word).await.unwrap(); + + assert_ne!(user_a_port, user_b_port); + assert_eq!(mgr.active_session_count(), 2); + assert_eq!(spawner.spawn_count.load(Ordering::SeqCst), 2); + + mgr.stop_for_user("user-a", path, DocType::Word).await; + assert_eq!(mgr.active_session_count(), 1); + assert!(!mgr.is_active_port(user_a_port, DocType::Word)); + assert!(!mgr.is_active_port_for_user("user-a", user_a_port, DocType::Word)); + assert!(mgr.is_active_port(user_b_port, DocType::Word)); + assert!(mgr.is_active_port_for_user("user-b", user_b_port, DocType::Word)); + assert!(!mgr.is_active_port_for_user("user-a", user_b_port, DocType::Word)); + + mgr.stop_for_user("user-b", path, DocType::Word).await; + assert_eq!(mgr.active_session_count(), 0); + } + #[tokio::test] async fn start_retries_when_allocated_port_is_taken() { let spawner = Arc::new(MockSpawner::new()); @@ -766,6 +888,28 @@ mod tests { assert_eq!(mgr.active_session_count(), 0); } + #[tokio::test] + async fn stop_all_for_user_keeps_other_user_sessions() { + let spawner = Arc::new(MockSpawner::new()); + let broadcaster = Arc::new(RecordingBroadcaster::new()); + let mgr = make_manager(Arc::clone(&spawner), Arc::clone(&broadcaster)); + + let dir = tempfile::tempdir().unwrap(); + let file = dir.path().join("shared.docx"); + std::fs::write(&file, b"test").unwrap(); + let path = file.to_str().unwrap(); + + let user_a_port = mgr.start_for_user("user-a", path, DocType::Word).await.unwrap(); + let user_b_port = mgr.start_for_user("user-b", path, DocType::Word).await.unwrap(); + assert_eq!(mgr.active_session_count(), 2); + + assert_eq!(mgr.stop_all_for_user("user-a"), 1); + + assert_eq!(mgr.active_session_count(), 1); + assert!(!mgr.is_active_port_for_user("user-a", user_a_port, DocType::Word)); + assert!(mgr.is_active_port_for_user("user-b", user_b_port, DocType::Word)); + } + #[tokio::test] async fn is_active_port_returns_true_for_active() { let spawner = Arc::new(MockSpawner::new()); @@ -778,6 +922,8 @@ mod tests { let port = mgr.start(file.to_str().unwrap(), DocType::Word).await.unwrap(); assert!(mgr.is_active_port(port, DocType::Word)); + assert!(mgr.is_active_port_for_user("system_default_user", port, DocType::Word)); + assert!(!mgr.is_active_port_for_user("other-user", port, DocType::Word)); assert!(!mgr.is_active_port(port, DocType::Ppt)); assert!(!mgr.is_active_port(12345, DocType::Word)); } @@ -802,6 +948,9 @@ mod tests { assert!(mgr.is_active_watch_port(word_port)); assert!(mgr.is_active_watch_port(excel_port)); + assert!(mgr.is_active_watch_port_for_user("system_default_user", word_port)); + assert!(mgr.is_active_watch_port_for_user("system_default_user", excel_port)); + assert!(!mgr.is_active_watch_port_for_user("other-user", word_port)); assert!(!mgr.is_active_watch_port(ppt_port)); assert!(!mgr.is_active_watch_port(12345)); } diff --git a/crates/aionui-office/tests/proxy_integration.rs b/crates/aionui-office/tests/proxy_integration.rs index be1a40877..cc175e349 100644 --- a/crates/aionui-office/tests/proxy_integration.rs +++ b/crates/aionui-office/tests/proxy_integration.rs @@ -135,6 +135,14 @@ fn build_http_response(status: u16, headers: &[(&str, &str)], body: &str) -> Str } async fn setup_proxy(doc_type: DocType, response_template: &str) -> (ProxyService, u16, tempfile::TempDir) { + setup_proxy_for_user("system_default_user", doc_type, response_template).await +} + +async fn setup_proxy_for_user( + user_id: &str, + doc_type: DocType, + response_template: &str, +) -> (ProxyService, u16, tempfile::TempDir) { let spawner = HttpMockSpawner { response_template: response_template.to_owned(), }; @@ -144,7 +152,10 @@ async fn setup_proxy(doc_type: DocType, response_template: &str) -> (ProxyServic let file = dir.path().join("test.docx"); std::fs::write(&file, b"test").unwrap(); - let port = mgr.start(file.to_str().unwrap(), doc_type).await.unwrap(); + let port = mgr + .start_for_user(user_id, file.to_str().unwrap(), doc_type) + .await + .unwrap(); let proxy = ProxyService::new(mgr); (proxy, port, dir) @@ -206,6 +217,35 @@ async fn ssrf_wrong_doc_type_rejected() { assert!(matches!(result.unwrap_err(), ProxyError::PortNotActive(_))); } +#[tokio::test] +async fn proxy_rejects_active_port_owned_by_another_user() { + let response = build_http_response(200, &[("Content-Type", "text/plain")], "owner preview"); + let (proxy, port, _dir) = setup_proxy_for_user("user-a", DocType::Ppt, &response).await; + + let owner = proxy + .forward_for_user("user-a", port, "/index.html", DocType::Ppt, &[]) + .await + .unwrap(); + assert_eq!(owner.status, 200); + + let other = proxy + .forward_for_user("user-b", port, "/index.html", DocType::Ppt, &[]) + .await; + assert!(matches!(other.unwrap_err(), ProxyError::PortNotActive(_))); +} + +#[tokio::test] +async fn watch_proxy_rejects_active_port_owned_by_another_user() { + let response = build_http_response(200, &[("Content-Type", "text/plain")], "owner preview"); + let (proxy, port, _dir) = setup_proxy_for_user("user-a", DocType::Word, &response).await; + + let owner = proxy.forward_watch_for_user("user-a", port, "/", &[]).await.unwrap(); + assert_eq!(owner.status, 200); + + let other = proxy.forward_watch_for_user("user-b", port, "/", &[]).await; + assert!(matches!(other.unwrap_err(), ProxyError::PortNotActive(_))); +} + // --------------------------------------------------------------------------- // H-1-13.8 fix: forward_watch accepts Excel session ports // --------------------------------------------------------------------------- diff --git a/crates/aionui-office/tests/watch_manager_integration.rs b/crates/aionui-office/tests/watch_manager_integration.rs index fce4e8827..8b1aef972 100644 --- a/crates/aionui-office/tests/watch_manager_integration.rs +++ b/crates/aionui-office/tests/watch_manager_integration.rs @@ -107,6 +107,10 @@ impl TestBroadcaster { .filter_map(|e| e.data["state"].as_str().map(String::from)) .collect() } + + fn events(&self) -> Vec> { + self.events.lock().unwrap().clone() + } } impl EventBroadcaster for TestBroadcaster { @@ -270,6 +274,31 @@ async fn status_events_use_correct_prefix() { assert!(names.contains(&"word-preview.status".to_string())); assert!(names.contains(&"excel-preview.status".to_string())); assert!(names.contains(&"ppt-preview.status".to_string())); + assert!( + broadcaster + .events() + .iter() + .all(|event| event.data["user_id"].as_str() == Some("system_default_user")) + ); +} + +#[tokio::test] +async fn status_events_include_starting_user() { + let spawner = Arc::new(TestSpawner::new(true)); + let broadcaster = Arc::new(TestBroadcaster::new()); + let mgr = OfficecliWatchManager::new(spawner, broadcaster.clone()); + + let dir = tempfile::tempdir().unwrap(); + let file = create_temp_file(&dir, "user.docx"); + + mgr.start_for_user("user-123", &file, DocType::Word).await.unwrap(); + + assert!( + broadcaster + .events() + .iter() + .any(|event| event.name == "word-preview.status" && event.data["user_id"].as_str() == Some("user-123")) + ); } // --------------------------------------------------------------------------- diff --git a/crates/aionui-realtime/src/handler.rs b/crates/aionui-realtime/src/handler.rs index 16ccc5e38..0c0d543bb 100644 --- a/crates/aionui-realtime/src/handler.rs +++ b/crates/aionui-realtime/src/handler.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::{future::Future, pin::Pin}; use aionui_api_types::WebSocketMessage; use axum::extract::WebSocketUpgrade; @@ -20,12 +21,16 @@ use crate::types::{ConnectionId, PER_CONNECTION_BUFFER, RealtimeError, WebSocket /// so that `aionui-realtime` does not depend on `aionui-auth` directly. pub type TokenExtractor = Arc Option + Send + Sync>; +/// Resolves a verified JWT token to the active internal user ID it represents. +pub type TokenUserResolver = Arc Pin> + Send>> + Send + Sync>; + /// Shared state required by the WebSocket upgrade handler. #[derive(Clone)] pub struct WsHandlerState { pub manager: Arc, pub router: Arc, pub token_validator: TokenValidator, + pub token_user_resolver: TokenUserResolver, pub token_extractor: TokenExtractor, } @@ -75,10 +80,15 @@ async fn handle_socket(socket: WebSocket, token: Option, state: WsHandle return; } + let Some(user_id) = (state.token_user_resolver)(token.clone()).await else { + send_realtime_error_and_close(socket, RealtimeError::AuthExpired, "authentication failed").await; + return; + }; + let (tx, rx) = mpsc::channel::(PER_CONNECTION_BUFFER); - let conn_id = state.manager.add_client(token, tx); + let conn_id = state.manager.add_client_for_user(user_id.clone(), token, tx); - info!(%conn_id, "websocket connection established"); + info!(%conn_id, user_id = %user_id, "websocket connection established"); let (ws_sender, ws_receiver) = socket.split(); @@ -88,7 +98,7 @@ async fn handle_socket(socket: WebSocket, token: Option, state: WsHandle // Recv loop exited — client disconnected or errored. send_handle.abort(); state.manager.remove_client(conn_id); - info!(%conn_id, "websocket connection closed"); + info!(%conn_id, user_id = %user_id, "websocket connection closed"); } /// Send a realtime boundary error event, then close with 1008. @@ -260,6 +270,7 @@ mod tests { manager, router: Arc::new(crate::router::NoopMessageRouter), token_validator: Arc::new(|_| true), + token_user_resolver: Arc::new(|_| Box::pin(async { Some("system_default_user".into()) })), token_extractor: Arc::new(|_| None), } } @@ -505,6 +516,7 @@ mod tests { manager, router: router.clone(), token_validator: Arc::new(|_| true), + token_user_resolver: Arc::new(|_| Box::pin(async { Some("system_default_user".into()) })), token_extractor: Arc::new(|_| None), }; diff --git a/crates/aionui-realtime/src/lib.rs b/crates/aionui-realtime/src/lib.rs index 01211dc4f..8ec0dbed6 100644 --- a/crates/aionui-realtime/src/lib.rs +++ b/crates/aionui-realtime/src/lib.rs @@ -8,7 +8,7 @@ pub mod router; pub mod types; pub use broadcaster::{BroadcastEventBus, EventBroadcaster}; -pub use handler::{TokenExtractor, WsHandlerState, ws_upgrade_handler}; +pub use handler::{TokenExtractor, TokenUserResolver, WsHandlerState, ws_upgrade_handler}; pub use manager::{TokenValidator, WebSocketManager}; pub use router::{MessageRouter, NoopMessageRouter}; pub use types::{ diff --git a/crates/aionui-realtime/src/manager.rs b/crates/aionui-realtime/src/manager.rs index 4d5153ba6..2bb25421a 100644 --- a/crates/aionui-realtime/src/manager.rs +++ b/crates/aionui-realtime/src/manager.rs @@ -35,8 +35,14 @@ impl WebSocketManager { /// Register a new client connection and return its assigned ID. pub fn add_client(&self, token: String, tx: mpsc::Sender) -> ConnectionId { + self.add_client_for_user("system_default_user".to_owned(), token, tx) + } + + /// Register a new authenticated client connection for an internal user ID. + pub fn add_client_for_user(&self, user_id: String, token: String, tx: mpsc::Sender) -> ConnectionId { let id = ConnectionId(self.next_id.fetch_add(1, Ordering::Relaxed)); let info = ClientInfo { + user_id, token, last_ping: Instant::now(), tx, @@ -65,6 +71,57 @@ impl WebSocketManager { self.connections.len() } + /// Returns the number of active connections authenticated as `user_id`. + pub fn client_count_for_user(&self, user_id: &str) -> usize { + self.connections + .iter() + .filter(|entry| entry.value().user_id == user_id) + .count() + } + + /// Close and remove every active connection authenticated as `user_id`. + /// + /// Returns the number of connections removed. The close is best-effort: + /// saturated or already-closed outbound queues are still removed from the + /// manager so no future events are delivered to revoked sessions. + pub fn disconnect_user(&self, user_id: &str, reason: &str) -> usize { + let mut to_remove = Vec::new(); + + for entry in self.connections.iter() { + let conn_id = *entry.key(); + let client = entry.value(); + if client.user_id != user_id { + continue; + } + + let outbound = terminal_realtime_error(conn_id, RealtimeError::AuthExpired, reason); + match client.tx.try_send(outbound) { + Ok(()) | Err(mpsc::error::TrySendError::Closed(_)) => { + to_remove.push(conn_id); + } + Err(mpsc::error::TrySendError::Full(_)) => { + warn!( + %conn_id, + user_id = %user_id, + code = RealtimeError::Backpressure.code(), + "outbound channel full, user disconnect close dropped" + ); + to_remove.push(conn_id); + } + } + } + + for conn_id in &to_remove { + self.remove_client(*conn_id); + } + + let removed = to_remove.len(); + if removed > 0 { + info!(user_id = %user_id, removed, "websocket user connections disconnected"); + } + removed + } + /// Send a message to all connected clients. /// /// Uses `try_send` for backpressure. A saturated channel cannot reliably @@ -103,6 +160,45 @@ impl WebSocketManager { } } + /// Send a message to all connections authenticated as `user_id`. + pub fn broadcast_to_user(&self, user_id: &str, msg: WebSocketMessage) { + let text = match serde_json::to_string(&msg) { + Ok(t) => t, + Err(e) => { + warn!(user_id = %user_id, error = %e, "failed to serialize user broadcast message"); + return; + } + }; + + let mut disconnected = Vec::new(); + for entry in self.connections.iter() { + let conn_id = *entry.key(); + let client = entry.value(); + if client.user_id != user_id { + continue; + } + + match client.tx.try_send(WsOutbound::Text(text.clone())) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(_)) => { + warn!( + %conn_id, + user_id = %user_id, + code = RealtimeError::Backpressure.code(), + "outbound channel full, user broadcast message dropped" + ); + } + Err(mpsc::error::TrySendError::Closed(_)) => { + disconnected.push(conn_id); + } + } + } + + for conn_id in disconnected { + self.remove_client(conn_id); + } + } + /// Send a message to a specific connection. pub fn send_to(&self, conn_id: ConnectionId, msg: WebSocketMessage) { let text = match serde_json::to_string(&msg) { @@ -171,7 +267,19 @@ impl Default for WebSocketManager { impl EventBroadcaster for WebSocketManager { fn broadcast(&self, event: WebSocketMessage) { - self.broadcast_all(event); + let Some(user_id) = event + .data + .get("user_id") + .and_then(|value| value.as_str()) + .map(str::to_owned) + else { + warn!( + event_name = %event.name, + "dropping websocket manager event without user_id" + ); + return; + }; + self.broadcast_to_user(&user_id, event); } } @@ -366,6 +474,72 @@ mod tests { } } + #[test] + fn broadcast_to_user_delivers_only_matching_connections() { + let mgr = WebSocketManager::new(); + let (tx1, mut rx1) = new_client_tx(); + let (tx2, mut rx2) = new_client_tx(); + let (tx3, mut rx3) = new_client_tx(); + + mgr.add_client_for_user("user-a".into(), "t1".into(), tx1); + mgr.add_client_for_user("user-b".into(), "t2".into(), tx2); + mgr.add_client_for_user("user-a".into(), "t3".into(), tx3); + + assert_eq!(mgr.client_count_for_user("user-a"), 2); + assert_eq!(mgr.client_count_for_user("user-b"), 1); + + let event = WebSocketMessage::new("scoped-event", json!({"user_id": "user-a"})); + mgr.broadcast_to_user("user-a", event); + + assert!(rx1.try_recv().is_ok()); + assert!(rx2.try_recv().is_err()); + assert!(rx3.try_recv().is_ok()); + } + + #[test] + fn disconnect_user_closes_and_removes_only_matching_connections() { + let mgr = WebSocketManager::new(); + let (tx1, mut rx1) = new_client_tx(); + let (tx2, mut rx2) = new_client_tx(); + let (tx3, mut rx3) = new_client_tx(); + + mgr.add_client_for_user("user-a".into(), "t1".into(), tx1); + mgr.add_client_for_user("user-b".into(), "t2".into(), tx2); + mgr.add_client_for_user("user-a".into(), "t3".into(), tx3); + + let removed = mgr.disconnect_user("user-a", "session revoked"); + + assert_eq!(removed, 2); + assert_eq!(mgr.client_count_for_user("user-a"), 0); + assert_eq!(mgr.client_count_for_user("user-b"), 1); + for msg in [rx1.try_recv().unwrap(), rx3.try_recv().unwrap()] { + match msg { + WsOutbound::TextThenClose(text, code, reason) => { + assert_realtime_auth_expired(&text); + assert_eq!(code, WebSocketCloseCode::PolicyViolation); + assert_eq!(reason, "session revoked"); + } + other => panic!("expected terminal auth close, got {other:?}"), + } + } + assert!(rx2.try_recv().is_err()); + } + + #[test] + fn disconnect_user_removes_matching_connections_when_queue_is_full() { + let mgr = WebSocketManager::new(); + let (tx, mut rx) = mpsc::channel(1); + tx.try_send(WsOutbound::Text("queued".into())).unwrap(); + mgr.add_client_for_user("user-a".into(), "token".into(), tx); + + let removed = mgr.disconnect_user("user-a", "session revoked"); + + assert_eq!(removed, 1); + assert_eq!(mgr.client_count(), 0); + assert_eq!(rx.try_recv().unwrap(), WsOutbound::Text("queued".into())); + assert!(rx.try_recv().is_err()); + } + #[test] fn broadcast_all_removes_closed_channels() { let mgr = WebSocketManager::new(); @@ -442,6 +616,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "valid".into(), last_ping: Instant::now(), tx, @@ -476,6 +651,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "valid".into(), last_ping: old_ping, tx, @@ -507,6 +683,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "expired-token".into(), last_ping: Instant::now(), tx, @@ -540,6 +717,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "expired-token".into(), last_ping: Instant::now(), tx, @@ -563,6 +741,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "expired".into(), last_ping: old_ping, tx, @@ -596,6 +775,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "good".into(), last_ping: Instant::now(), tx: tx1, @@ -607,6 +787,7 @@ mod tests { connections.insert( ConnectionId(2), ClientInfo { + user_id: "user-a".into(), token: "good".into(), last_ping: Instant::now() - (HEARTBEAT_TIMEOUT * 2), tx: tx2, @@ -622,13 +803,15 @@ mod tests { } #[test] - fn event_broadcaster_impl_delegates_to_broadcast_all() { + fn event_broadcaster_impl_routes_to_event_user() { let mgr = WebSocketManager::new(); let (tx, mut rx) = new_client_tx(); - mgr.add_client("tok".into(), tx); + mgr.add_client_for_user("user-a".into(), "tok".into(), tx); + let (tx_other, mut rx_other) = new_client_tx(); + mgr.add_client_for_user("user-b".into(), "tok-other".into(), tx_other); let broadcaster: &dyn EventBroadcaster = &mgr; - broadcaster.broadcast(WebSocketMessage::new("via-trait", json!({}))); + broadcaster.broadcast(WebSocketMessage::new("via-trait", json!({"user_id": "user-a"}))); let msg = rx.try_recv().unwrap(); match msg { @@ -637,6 +820,19 @@ mod tests { } _ => panic!("expected Text"), } + assert!(rx_other.try_recv().is_err()); + } + + #[test] + fn event_broadcaster_impl_drops_unscoped_events() { + let mgr = WebSocketManager::new(); + let (tx, mut rx) = new_client_tx(); + mgr.add_client_for_user("user-a".into(), "tok".into(), tx); + + let broadcaster: &dyn EventBroadcaster = &mgr; + broadcaster.broadcast(WebSocketMessage::new("via-trait", json!({}))); + + assert!(rx.try_recv().is_err()); } #[test] diff --git a/crates/aionui-realtime/src/types.rs b/crates/aionui-realtime/src/types.rs index 231fa8a98..13c33ebbc 100644 --- a/crates/aionui-realtime/src/types.rs +++ b/crates/aionui-realtime/src/types.rs @@ -138,6 +138,8 @@ impl RealtimeError { /// Per-connection client state maintained by the server. pub struct ClientInfo { + /// Internal user ID this connection is authenticated as. + pub user_id: String, /// The JWT token this connection authenticated with. pub token: String, /// Timestamp of last ping/pong activity. diff --git a/crates/aionui-realtime/tests/handler_integration.rs b/crates/aionui-realtime/tests/handler_integration.rs index 1d23288f6..42b8092eb 100644 --- a/crates/aionui-realtime/tests/handler_integration.rs +++ b/crates/aionui-realtime/tests/handler_integration.rs @@ -37,6 +37,9 @@ fn default_state() -> (WsHandlerState, Arc) { manager: manager.clone(), router: Arc::new(NoopMessageRouter), token_validator: Arc::new(|t| t == "valid-token"), + token_user_resolver: Arc::new(|t| { + Box::pin(async move { (t == "valid-token").then(|| "test-user".to_owned()) }) + }), token_extractor: Arc::new(|headers| { headers .get("authorization") @@ -396,6 +399,9 @@ async fn unknown_message_routed_to_message_router() { manager: manager.clone(), router: router.clone(), token_validator: Arc::new(|t| t == "valid-token"), + token_user_resolver: Arc::new(|t| { + Box::pin(async move { (t == "valid-token").then(|| "test-user".to_owned()) }) + }), token_extractor: Arc::new(|headers| { headers .get("authorization") diff --git a/crates/aionui-shell/Cargo.toml b/crates/aionui-shell/Cargo.toml index 7d64ff9aa..d0c5e026f 100644 --- a/crates/aionui-shell/Cargo.toml +++ b/crates/aionui-shell/Cargo.toml @@ -6,6 +6,7 @@ edition.workspace = true [dependencies] aionui-common.workspace = true aionui-api-types.workspace = true +aionui-auth.workspace = true aionui-system.workspace = true aionui-runtime.workspace = true async-trait.workspace = true diff --git a/crates/aionui-shell/src/routes.rs b/crates/aionui-shell/src/routes.rs index d5d73d81e..99024e0fd 100644 --- a/crates/aionui-shell/src/routes.rs +++ b/crates/aionui-shell/src/routes.rs @@ -1,7 +1,7 @@ #![allow(clippy::disallowed_types)] use axum::extract::ws::{Message, WebSocket}; -use axum::extract::{Multipart, State, WebSocketUpgrade}; +use axum::extract::{Extension, Multipart, State, WebSocketUpgrade}; use axum::http::StatusCode; use axum::response::IntoResponse; use axum::routing::{get, post}; @@ -13,6 +13,7 @@ use aionui_api_types::{ ApiResponse, CheckToolInstalledRequest, CheckToolInstalledResponse, OpenExternalRequest, OpenFileRequest, OpenFolderWithRequest, ShowItemInFolderRequest, SpeechToTextConfig, SttStreamServerMessage, }; +use aionui_auth::CurrentUser; use aionui_common::ApiError; use aionui_system::ClientPrefService; @@ -186,6 +187,7 @@ async fn extract_stt_multipart(mut multipart: Multipart) -> Result, + Extension(user): Extension, multipart: Multipart, ) -> Result<(StatusCode, Json), (StatusCode, Json)> { let fields = extract_stt_multipart(multipart).await.map_err(|e| { @@ -198,16 +200,18 @@ async fn speech_to_text( (status, Json(body)) })?; - let config = load_stt_config(&state.client_pref_service).await.map_err(|e| { - let e = ApiError::from(e); - let status = e.status_code(); - let body = serde_json::json!({ - "success": false, - "error": e.to_string(), - "code": e.error_code(), - }); - (status, Json(body)) - })?; + let config = load_stt_config(&state.client_pref_service, &user.id) + .await + .map_err(|e| { + let e = ApiError::from(e); + let status = e.status_code(); + let body = serde_json::json!({ + "success": false, + "error": e.to_string(), + "code": e.error_code(), + }); + (status, Json(body)) + })?; let result = state .stt_service @@ -247,9 +251,10 @@ fn stt_error_response(err: &SttError) -> (StatusCode, Json) { /// `STT_DISABLED` error. pub(crate) async fn load_stt_config( client_pref_service: &ClientPrefService, + user_id: &str, ) -> Result { let prefs = client_pref_service - .get_preferences(Some(&["speechToText", "tools.speechToText"])) + .get_preferences(user_id, Some(&["speechToText", "tools.speechToText"])) .await?; Ok(prefs .get("tools.speechToText") @@ -283,18 +288,22 @@ const STT_STREAM_CHANNEL_CAPACITY: usize = 32; /// reaches the client as a uniform `{"type":"error",...}` WS frame /// (`STT_DISABLED`, `STT_*_NOT_CONFIGURED`, ...) emitted by the session, /// instead of a mix of HTTP handshake statuses and protocol frames. -async fn speech_to_text_stream(State(state): State, ws: WebSocketUpgrade) -> impl IntoResponse { - ws.on_upgrade(move |socket| speech_to_text_stream_socket(socket, state)) +async fn speech_to_text_stream( + State(state): State, + Extension(user): Extension, + ws: WebSocketUpgrade, +) -> impl IntoResponse { + ws.on_upgrade(move |socket| speech_to_text_stream_socket(socket, state, user.id)) } /// Pump WebSocket frames in/out of the transport-agnostic streaming session. /// /// This stays a pure frame adapter: all protocol and business logic lives in /// `stt_stream::run_stream_session`. -async fn speech_to_text_stream_socket(socket: WebSocket, state: ShellRouterState) { +async fn speech_to_text_stream_socket(socket: WebSocket, state: ShellRouterState, user_id: String) { let (mut ws_tx, mut ws_rx) = socket.split(); - let config = match load_stt_config(&state.client_pref_service).await { + let config = match load_stt_config(&state.client_pref_service, &user_id).await { Ok(config) => config, Err(e) => { tracing::error!(error = %e, "stt stream: failed to load config"); diff --git a/crates/aionui-system/src/client_pref.rs b/crates/aionui-system/src/client_pref.rs index a22a6f8c0..f16c13591 100644 --- a/crates/aionui-system/src/client_pref.rs +++ b/crates/aionui-system/src/client_pref.rs @@ -29,20 +29,32 @@ impl ClientPrefService { pub fn with_keep_awake_controller( repo: Arc, keep_awake_controller: DynKeepAwakeController, + keep_awake_restore_user_id: impl Into, ) -> Self { - let service = Self { + let service = Self::with_keep_awake_controller_without_restore(repo, keep_awake_controller); + service.restore_keep_awake_from_preferences(keep_awake_restore_user_id.into()); + service + } + + pub fn with_keep_awake_controller_without_restore( + repo: Arc, + keep_awake_controller: DynKeepAwakeController, + ) -> Self { + Self { repo, keep_awake_controller, - }; - service.restore_keep_awake_from_preferences(); - service + } } /// Get all client preferences, or only the specified keys. - pub async fn get_preferences(&self, keys: Option<&[&str]>) -> Result { + pub async fn get_preferences( + &self, + user_id: &str, + keys: Option<&[&str]>, + ) -> Result { let rows = match keys { - Some(k) if !k.is_empty() => self.repo.get_by_keys(k).await, - _ => self.repo.get_all().await, + Some(k) if !k.is_empty() => self.repo.get_by_keys(user_id, k).await, + _ => self.repo.get_all(user_id).await, } .map_err(|e| SystemError::Internal(format!("Failed to get preferences: {e}")))?; @@ -76,7 +88,11 @@ impl ClientPrefService { } /// Batch update client preferences. Null values delete the key. - pub async fn update_preferences(&self, req: UpdateClientPreferencesRequest) -> Result<(), SystemError> { + pub async fn update_preferences( + &self, + user_id: &str, + req: UpdateClientPreferencesRequest, + ) -> Result<(), SystemError> { let mut upserts: Vec<(String, String)> = Vec::new(); let mut deletes: Vec = Vec::new(); let keep_awake_update = resolve_keep_awake_update(&req)?; @@ -110,7 +126,7 @@ impl ClientPrefService { } let previous_keep_awake = if keep_awake_update.is_some() { - Some(self.get_stored_keep_awake().await?) + Some(self.get_stored_keep_awake(user_id).await?) } else { None }; @@ -123,7 +139,7 @@ impl ClientPrefService { let entries: Vec<(&str, &str)> = upserts.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect(); if let Err(error) = self .repo - .upsert_batch(&entries) + .upsert_batch(user_id, &entries) .await .map_err(|e| SystemError::Internal(format!("Failed to upsert preferences: {e}"))) { @@ -138,7 +154,7 @@ impl ClientPrefService { let keys: Vec<&str> = deletes.iter().map(|k| k.as_str()).collect(); if let Err(error) = self .repo - .delete_keys(&keys) + .delete_keys(user_id, &keys) .await .map_err(|e| SystemError::Internal(format!("Failed to delete preferences: {e}"))) { @@ -161,10 +177,10 @@ impl ClientPrefService { Ok(()) } - async fn get_stored_keep_awake(&self) -> Result { + async fn get_stored_keep_awake(&self, user_id: &str) -> Result { let rows = self .repo - .get_by_keys(&[KEEP_AWAKE_KEY]) + .get_by_keys(user_id, &[KEEP_AWAKE_KEY]) .await .map_err(|e| SystemError::Internal(format!("Failed to get keep-awake preference: {e}")))?; @@ -190,14 +206,14 @@ impl ClientPrefService { Ok(()) } - fn restore_keep_awake_from_preferences(&self) { + fn restore_keep_awake_from_preferences(&self, user_id: String) { let service = self.clone(); let Ok(handle) = tokio::runtime::Handle::try_current() else { warn!("Cannot restore system keep-awake preference without a Tokio runtime"); return; }; handle.spawn(async move { - match service.get_stored_keep_awake().await { + match service.get_stored_keep_awake(&user_id).await { Ok(true) => { if let Err(error) = service.apply_keep_awake(true).await { warn!(error = %error, "Failed to restore system keep-awake assertion"); @@ -258,6 +274,8 @@ mod tests { use tracing::Level; use tracing_subscriber::fmt; + const TEST_USER_ID: &str = "user-1"; + #[derive(Clone)] struct SharedBuf(Arc>>); @@ -299,6 +317,15 @@ mod tests { async fn setup() -> ClientPrefService { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, '', 'active', 0, 1, 1)", + ) + .bind(TEST_USER_ID) + .bind(TEST_USER_ID) + .execute(db.pool()) + .await + .unwrap(); let repo = Arc::new(SqliteClientPreferenceRepository::new(db.pool().clone())); std::mem::forget(db); ClientPrefService::new(repo) @@ -306,6 +333,15 @@ mod tests { async fn setup_with_keep_awake_controller(controller: DynKeepAwakeController) -> ClientPrefService { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, '', 'active', 0, 1, 1)", + ) + .bind(TEST_USER_ID) + .bind(TEST_USER_ID) + .execute(db.pool()) + .await + .unwrap(); let repo = Arc::new(SqliteClientPreferenceRepository::new(db.pool().clone())); std::mem::forget(db); ClientPrefService { @@ -353,7 +389,7 @@ mod tests { #[tokio::test] async fn get_empty_returns_empty_map() { let svc = setup().await; - let prefs = svc.get_preferences(None).await.unwrap(); + let prefs = svc.get_preferences(TEST_USER_ID, None).await.unwrap(); assert!(prefs.is_empty()); } @@ -362,9 +398,9 @@ mod tests { let svc = setup().await; let mut req = UpdateClientPreferencesRequest::new(); req.insert("system.closeToTray".into(), json!(true)); - svc.update_preferences(req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req).await.unwrap(); - let prefs = svc.get_preferences(None).await.unwrap(); + let prefs = svc.get_preferences(TEST_USER_ID, None).await.unwrap(); assert_eq!(prefs["system.closeToTray"], json!(true)); } @@ -373,9 +409,9 @@ mod tests { let svc = setup().await; let mut req = UpdateClientPreferencesRequest::new(); req.insert("pet.size".into(), json!(360)); - svc.update_preferences(req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req).await.unwrap(); - let prefs = svc.get_preferences(None).await.unwrap(); + let prefs = svc.get_preferences(TEST_USER_ID, None).await.unwrap(); assert_eq!(prefs["pet.size"], json!(360)); } @@ -384,9 +420,9 @@ mod tests { let svc = setup().await; let mut req = UpdateClientPreferencesRequest::new(); req.insert("theme".into(), json!("dark")); - svc.update_preferences(req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req).await.unwrap(); - let prefs = svc.get_preferences(None).await.unwrap(); + let prefs = svc.get_preferences(TEST_USER_ID, None).await.unwrap(); assert_eq!(prefs["theme"], json!("dark")); } @@ -396,13 +432,13 @@ mod tests { let mut req = UpdateClientPreferencesRequest::new(); req.insert("theme".into(), json!("dark")); - svc.update_preferences(req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req).await.unwrap(); let mut req2 = UpdateClientPreferencesRequest::new(); req2.insert("theme".into(), json!(null)); - svc.update_preferences(req2).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req2).await.unwrap(); - let prefs = svc.get_preferences(None).await.unwrap(); + let prefs = svc.get_preferences(TEST_USER_ID, None).await.unwrap(); assert!(!prefs.contains_key("theme")); } @@ -414,9 +450,9 @@ mod tests { req.insert("a".into(), json!(1)); req.insert("b".into(), json!(2)); req.insert("c".into(), json!(3)); - svc.update_preferences(req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req).await.unwrap(); - let prefs = svc.get_preferences(Some(&["a", "c"])).await.unwrap(); + let prefs = svc.get_preferences(TEST_USER_ID, Some(&["a", "c"])).await.unwrap(); assert_eq!(prefs.len(), 2); assert_eq!(prefs["a"], json!(1)); assert_eq!(prefs["c"], json!(3)); @@ -434,10 +470,10 @@ mod tests { let mut req = UpdateClientPreferencesRequest::new(); req.insert("appearance.secretTheme".into(), json!("super-secret-value")); req.insert("appearance.deleted".into(), json!(null)); - svc.update_preferences(req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req).await.unwrap(); let _ = svc - .get_preferences(Some(&["appearance.secretTheme", "appearance.missing"])) + .get_preferences(TEST_USER_ID, Some(&["appearance.secretTheme", "appearance.missing"])) .await .unwrap(); }); @@ -467,13 +503,13 @@ mod tests { let mut req1 = UpdateClientPreferencesRequest::new(); req1.insert("k".into(), json!("v1")); - svc.update_preferences(req1).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req1).await.unwrap(); let mut req2 = UpdateClientPreferencesRequest::new(); req2.insert("k".into(), json!("v2")); - svc.update_preferences(req2).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req2).await.unwrap(); - let prefs = svc.get_preferences(None).await.unwrap(); + let prefs = svc.get_preferences(TEST_USER_ID, None).await.unwrap(); assert_eq!(prefs["k"], json!("v2")); } @@ -482,7 +518,7 @@ mod tests { let svc = setup().await; let mut req = UpdateClientPreferencesRequest::new(); req.insert("".into(), json!(true)); - let err = svc.update_preferences(req).await.unwrap_err(); + let err = svc.update_preferences(TEST_USER_ID, req).await.unwrap_err(); assert!(matches!(err, SystemError::BadRequest(_))); } @@ -491,7 +527,7 @@ mod tests { let svc = setup().await; let mut req = UpdateClientPreferencesRequest::new(); req.insert("x".repeat(256), json!(true)); - let err = svc.update_preferences(req).await.unwrap_err(); + let err = svc.update_preferences(TEST_USER_ID, req).await.unwrap_err(); assert!(matches!(err, SystemError::BadRequest(_))); } @@ -502,14 +538,14 @@ mod tests { let mut setup_req = UpdateClientPreferencesRequest::new(); setup_req.insert("keep".into(), json!(1)); setup_req.insert("remove".into(), json!(2)); - svc.update_preferences(setup_req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, setup_req).await.unwrap(); let mut req = UpdateClientPreferencesRequest::new(); req.insert("remove".into(), json!(null)); req.insert("new".into(), json!(3)); - svc.update_preferences(req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req).await.unwrap(); - let prefs = svc.get_preferences(None).await.unwrap(); + let prefs = svc.get_preferences(TEST_USER_ID, None).await.unwrap(); assert_eq!(prefs.len(), 2); assert_eq!(prefs["keep"], json!(1)); assert_eq!(prefs["new"], json!(3)); @@ -522,10 +558,13 @@ mod tests { let mut req = UpdateClientPreferencesRequest::new(); req.insert(KEEP_AWAKE_KEY.into(), json!(true)); - svc.update_preferences(req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req).await.unwrap(); assert_eq!(*controller.calls.lock().unwrap(), vec![true]); - let prefs = svc.get_preferences(Some(&[KEEP_AWAKE_KEY])).await.unwrap(); + let prefs = svc + .get_preferences(TEST_USER_ID, Some(&[KEEP_AWAKE_KEY])) + .await + .unwrap(); assert_eq!(prefs[KEEP_AWAKE_KEY], json!(true)); } @@ -535,14 +574,17 @@ mod tests { let svc = setup_with_keep_awake_controller(controller.clone()).await; let mut setup_req = UpdateClientPreferencesRequest::new(); setup_req.insert(KEEP_AWAKE_KEY.into(), json!(true)); - svc.update_preferences(setup_req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, setup_req).await.unwrap(); let mut req = UpdateClientPreferencesRequest::new(); req.insert(KEEP_AWAKE_KEY.into(), json!(null)); - svc.update_preferences(req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req).await.unwrap(); assert_eq!(*controller.calls.lock().unwrap(), vec![true, false]); - let prefs = svc.get_preferences(Some(&[KEEP_AWAKE_KEY])).await.unwrap(); + let prefs = svc + .get_preferences(TEST_USER_ID, Some(&[KEEP_AWAKE_KEY])) + .await + .unwrap(); assert!(!prefs.contains_key(KEEP_AWAKE_KEY)); } @@ -552,12 +594,15 @@ mod tests { let svc = setup_with_keep_awake_controller(controller.clone()).await; let mut req = UpdateClientPreferencesRequest::new(); req.insert(KEEP_AWAKE_KEY.into(), json!(true)); - svc.update_preferences(req).await.unwrap(); + svc.update_preferences(TEST_USER_ID, req).await.unwrap(); svc.release_keep_awake_for_shutdown().await.unwrap(); assert_eq!(*controller.calls.lock().unwrap(), vec![true, false]); - let prefs = svc.get_preferences(Some(&[KEEP_AWAKE_KEY])).await.unwrap(); + let prefs = svc + .get_preferences(TEST_USER_ID, Some(&[KEEP_AWAKE_KEY])) + .await + .unwrap(); assert_eq!(prefs[KEEP_AWAKE_KEY], json!(true)); } @@ -568,11 +613,14 @@ mod tests { let mut req = UpdateClientPreferencesRequest::new(); req.insert(KEEP_AWAKE_KEY.into(), json!("yes")); - let err = svc.update_preferences(req).await.unwrap_err(); + let err = svc.update_preferences(TEST_USER_ID, req).await.unwrap_err(); assert!(matches!(err, SystemError::BadRequest(_))); assert!(controller.calls.lock().unwrap().is_empty()); - let prefs = svc.get_preferences(Some(&[KEEP_AWAKE_KEY])).await.unwrap(); + let prefs = svc + .get_preferences(TEST_USER_ID, Some(&[KEEP_AWAKE_KEY])) + .await + .unwrap(); assert!(!prefs.contains_key(KEEP_AWAKE_KEY)); } @@ -586,11 +634,14 @@ mod tests { let mut req = UpdateClientPreferencesRequest::new(); req.insert(KEEP_AWAKE_KEY.into(), json!(true)); - let err = svc.update_preferences(req).await.unwrap_err(); + let err = svc.update_preferences(TEST_USER_ID, req).await.unwrap_err(); assert!(matches!(err, SystemError::Internal(_))); assert_eq!(*controller.calls.lock().unwrap(), vec![true]); - let prefs = svc.get_preferences(Some(&[KEEP_AWAKE_KEY])).await.unwrap(); + let prefs = svc + .get_preferences(TEST_USER_ID, Some(&[KEEP_AWAKE_KEY])) + .await + .unwrap(); assert!(!prefs.contains_key(KEEP_AWAKE_KEY)); } @@ -599,10 +650,11 @@ mod tests { let initial = setup().await; let mut req = UpdateClientPreferencesRequest::new(); req.insert(KEEP_AWAKE_KEY.into(), json!(true)); - initial.update_preferences(req).await.unwrap(); + initial.update_preferences(TEST_USER_ID, req).await.unwrap(); let controller = Arc::new(RecordingKeepAwakeController::default()); - let service = ClientPrefService::with_keep_awake_controller(initial.repo.clone(), controller.clone()); + let service = + ClientPrefService::with_keep_awake_controller(initial.repo.clone(), controller.clone(), TEST_USER_ID); for _ in 0..50 { if !controller.calls.lock().unwrap().is_empty() { @@ -614,4 +666,21 @@ mod tests { assert_eq!(*controller.calls.lock().unwrap(), vec![true]); drop(service); } + + #[tokio::test] + async fn keep_awake_without_restore_does_not_read_persisted_default_user_preference() { + let initial = setup().await; + let mut req = UpdateClientPreferencesRequest::new(); + req.insert(KEEP_AWAKE_KEY.into(), json!(true)); + initial.update_preferences(TEST_USER_ID, req).await.unwrap(); + + let controller = Arc::new(RecordingKeepAwakeController::default()); + let service = + ClientPrefService::with_keep_awake_controller_without_restore(initial.repo.clone(), controller.clone()); + + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + + assert!(controller.calls.lock().unwrap().is_empty()); + drop(service); + } } diff --git a/crates/aionui-system/src/model_fetcher/mod.rs b/crates/aionui-system/src/model_fetcher/mod.rs index f315d48fe..ff0c3d11a 100644 --- a/crates/aionui-system/src/model_fetcher/mod.rs +++ b/crates/aionui-system/src/model_fetcher/mod.rs @@ -49,10 +49,11 @@ impl ModelFetchService { /// URL auto-correction with parallel probing. pub async fn fetch_models( &self, + user_id: &str, provider_id: &str, req: &FetchModelsRequest, ) -> Result { - let config = self.load_provider_config(provider_id).await?; + let config = self.load_provider_config(user_id, provider_id).await?; self.fetch_with_config(&config, req.try_fix).await } @@ -89,10 +90,10 @@ impl ModelFetchService { } /// Extract and decrypt provider configuration from DB. - async fn load_provider_config(&self, provider_id: &str) -> Result { + async fn load_provider_config(&self, user_id: &str, provider_id: &str) -> Result { let row = self .repo - .find_by_id(provider_id) + .find_by_id(user_id, provider_id) .await? .ok_or_else(|| SystemError::NotFound(format!("Provider {provider_id} not found")))?; @@ -143,9 +144,19 @@ mod tests { use aionui_db::{CreateProviderParams, SqliteProviderRepository, init_database_memory}; const TEST_KEY: [u8; 32] = [0x42; 32]; + const TEST_USER_ID: &str = "user-1"; async fn setup() -> (ModelFetchService, aionui_db::Database) { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, '', 'active', 0, 1, 1)", + ) + .bind(TEST_USER_ID) + .bind(TEST_USER_ID) + .execute(db.pool()) + .await + .unwrap(); let repo = Arc::new(SqliteProviderRepository::new(db.pool().clone())); let svc = ModelFetchService::new(repo, TEST_KEY, reqwest::Client::new()); (svc, db) @@ -157,6 +168,7 @@ mod tests { let row = repo .create(CreateProviderParams { id: None, + user_id: TEST_USER_ID, platform, name: "Test", base_url, @@ -198,7 +210,7 @@ mod tests { #[tokio::test] async fn load_config_nonexistent_provider_returns_not_found() { let (svc, _db) = setup().await; - let err = svc.load_provider_config("no_such_id").await.unwrap_err(); + let err = svc.load_provider_config(TEST_USER_ID, "no_such_id").await.unwrap_err(); assert!(matches!(err, SystemError::NotFound(_))); } @@ -206,7 +218,7 @@ mod tests { async fn load_config_empty_api_key_returns_bad_request() { let (svc, db) = setup().await; let id = create_provider(&db, "openai", "https://api.openai.com", " ").await; - let err = svc.load_provider_config(&id).await.unwrap_err(); + let err = svc.load_provider_config(TEST_USER_ID, &id).await.unwrap_err(); assert!(matches!(err, SystemError::BadRequest(_))); } @@ -214,7 +226,7 @@ mod tests { async fn load_config_decrypts_api_key() { let (svc, db) = setup().await; let id = create_provider(&db, "openai", "https://api.openai.com", "sk-test-key").await; - let config = svc.load_provider_config(&id).await.unwrap(); + let config = svc.load_provider_config(TEST_USER_ID, &id).await.unwrap(); assert_eq!(config.api_key, "sk-test-key"); assert_eq!(config.platform, "openai"); assert_eq!(config.base_url, "https://api.openai.com"); @@ -226,7 +238,7 @@ mod tests { let (svc, db) = setup().await; let id = create_provider(&db, "vertex-ai", "https://unused", "fake-key").await; let req = FetchModelsRequest { try_fix: false }; - let resp = svc.fetch_models(&id, &req).await.unwrap(); + let resp = svc.fetch_models(TEST_USER_ID, &id, &req).await.unwrap(); assert_eq!(resp.models.len(), 2); assert!(resp.fixed_base_url.is_none()); } @@ -236,7 +248,7 @@ mod tests { let (svc, db) = setup().await; let id = create_provider(&db, "minimax", "https://unused", "fake-key").await; let req = FetchModelsRequest { try_fix: false }; - let resp = svc.fetch_models(&id, &req).await.unwrap(); + let resp = svc.fetch_models(TEST_USER_ID, &id, &req).await.unwrap(); assert_eq!(resp.models.len(), 3); } @@ -244,7 +256,7 @@ mod tests { async fn fetch_models_nonexistent_provider() { let (svc, _db) = setup().await; let req = FetchModelsRequest { try_fix: false }; - let err = svc.fetch_models("no_such_id", &req).await.unwrap_err(); + let err = svc.fetch_models(TEST_USER_ID, "no_such_id", &req).await.unwrap_err(); assert!(matches!(err, SystemError::NotFound(_))); } @@ -354,7 +366,7 @@ mod tests { let (svc, db) = setup().await; // Save a provider with multi-key api_key let id = create_provider(&db, "openai", "https://api.openai.com", "sk-key1\nsk-key2").await; - let config = svc.load_provider_config(&id).await.unwrap(); + let config = svc.load_provider_config(TEST_USER_ID, &id).await.unwrap(); // Must extract only the first key for model fetching assert_eq!(config.api_key, "sk-key1"); } diff --git a/crates/aionui-system/src/provider.rs b/crates/aionui-system/src/provider.rs index e22898df6..37d37d5a2 100644 --- a/crates/aionui-system/src/provider.rs +++ b/crates/aionui-system/src/provider.rs @@ -21,8 +21,8 @@ impl ProviderService { } /// List all providers with masked API keys. - pub async fn list(&self) -> Result, SystemError> { - let rows = self.repo.list().await?; + pub async fn list(&self, user_id: &str) -> Result, SystemError> { + let rows = self.repo.list(user_id).await?; rows.into_iter().map(|row| self.row_to_response(row)).collect() } @@ -32,7 +32,7 @@ impl ProviderService { /// otherwise a fresh id is generated by the repository. This supports the /// frontend-local-store → backend migration path where existing provider /// ids must be preserved. - pub async fn create(&self, req: CreateProviderRequest) -> Result { + pub async fn create(&self, user_id: &str, req: CreateProviderRequest) -> Result { validate_create_request(&req)?; let encrypted_key = encrypt_string(&req.api_key, &self.encryption_key)?; @@ -47,6 +47,7 @@ impl ProviderService { let params = CreateProviderParams { id: trimmed_id, + user_id, platform: &req.platform, name: &req.name, base_url: &req.base_url, @@ -68,7 +69,12 @@ impl ProviderService { } /// Update an existing provider. Only provided fields are changed. - pub async fn update(&self, id: &str, req: UpdateProviderRequest) -> Result { + pub async fn update( + &self, + user_id: &str, + id: &str, + req: UpdateProviderRequest, + ) -> Result { validate_update_request(&req)?; let encrypted_key = req @@ -101,13 +107,13 @@ impl ProviderService { is_full_url: req.is_full_url, }; - let row = self.repo.update(id, params).await?; + let row = self.repo.update(user_id, id, params).await?; self.row_to_response(row) } /// Delete a provider by ID. - pub async fn delete(&self, id: &str) -> Result<(), SystemError> { - self.repo.delete(id).await?; + pub async fn delete(&self, user_id: &str, id: &str) -> Result<(), SystemError> { + self.repo.delete(user_id, id).await?; Ok(()) } @@ -278,6 +284,8 @@ fn validate_base_url(url: &str) -> Result<(), SystemError> { #[cfg(test)] mod tests { use super::*; + + const TEST_USER_ID: &str = "user-1"; use aionui_db::{SqliteProviderRepository, init_database_memory}; // A fixed 32-byte key for testing @@ -285,6 +293,15 @@ mod tests { async fn setup() -> ProviderService { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, '', 'active', 0, 1, 1)", + ) + .bind(TEST_USER_ID) + .bind(TEST_USER_ID) + .execute(db.pool()) + .await + .unwrap(); let repo = Arc::new(SqliteProviderRepository::new(db.pool().clone())); std::mem::forget(db); ProviderService::new(repo, TEST_KEY) @@ -483,14 +500,14 @@ mod tests { #[tokio::test] async fn list_empty() { let svc = setup().await; - let result = svc.list().await.unwrap(); + let result = svc.list(TEST_USER_ID).await.unwrap(); assert!(result.is_empty()); } #[tokio::test] async fn create_and_list() { let svc = setup().await; - let created = svc.create(sample_create_request()).await.unwrap(); + let created = svc.create(TEST_USER_ID, sample_create_request()).await.unwrap(); assert!(created.id.starts_with("prov_")); assert_eq!(created.platform, "anthropic"); @@ -501,7 +518,7 @@ mod tests { assert_eq!(created.models, vec!["claude-sonnet-4-20250514"]); assert!(created.enabled); - let all = svc.list().await.unwrap(); + let all = svc.list(TEST_USER_ID).await.unwrap(); assert_eq!(all.len(), 1); assert_eq!(all[0].id, created.id); assert_eq!(all[0].api_key, "sk-ant-api03-test1234"); @@ -514,7 +531,7 @@ mod tests { id: Some("caller-id-xyz".into()), ..sample_create_request() }; - let created = svc.create(req).await.unwrap(); + let created = svc.create(TEST_USER_ID, req).await.unwrap(); assert_eq!(created.id, "caller-id-xyz"); } @@ -525,7 +542,7 @@ mod tests { id: Some(" ".into()), ..sample_create_request() }; - let err = svc.create(req).await.unwrap_err(); + let err = svc.create(TEST_USER_ID, req).await.unwrap_err(); assert!(matches!(err, SystemError::BadRequest(_))); } @@ -536,13 +553,13 @@ mod tests { id: Some("dup-id".into()), ..sample_create_request() }; - svc.create(req1).await.unwrap(); + svc.create(TEST_USER_ID, req1).await.unwrap(); let req2 = CreateProviderRequest { id: Some("dup-id".into()), ..sample_create_request() }; - let err = svc.create(req2).await.unwrap_err(); + let err = svc.create(TEST_USER_ID, req2).await.unwrap_err(); assert!(matches!(err, SystemError::Conflict(_))); } @@ -555,7 +572,7 @@ mod tests { model_enabled: Some(HashMap::from([("gpt-4".into(), true), ("gpt-3.5".into(), false)])), ..sample_create_request() }; - let created = svc.create(req).await.unwrap(); + let created = svc.create(TEST_USER_ID, req).await.unwrap(); assert_eq!( created.model_protocols.as_ref().and_then(|m| m.get("gpt-4")), @@ -568,7 +585,7 @@ mod tests { ); // And persist through a fresh read. - let all = svc.list().await.unwrap(); + let all = svc.list(TEST_USER_ID).await.unwrap(); assert_eq!(all[0].model_enabled.as_ref().and_then(|m| m.get("gpt-4")), Some(&true)); } @@ -581,7 +598,7 @@ mod tests { api_key: "sk-secret-original-value".into(), ..sample_create_request() }; - let created = svc.create(req).await.unwrap(); + let created = svc.create(TEST_USER_ID, req).await.unwrap(); assert_eq!(created.api_key, "sk-secret-original-value"); assert!(!created.api_key.contains("***")); } @@ -593,17 +610,18 @@ mod tests { platform: "".into(), ..sample_create_request() }; - let err = svc.create(req).await.unwrap_err(); + let err = svc.create(TEST_USER_ID, req).await.unwrap_err(); assert!(matches!(err, SystemError::BadRequest(_))); } #[tokio::test] async fn update_name() { let svc = setup().await; - let created = svc.create(sample_create_request()).await.unwrap(); + let created = svc.create(TEST_USER_ID, sample_create_request()).await.unwrap(); let updated = svc .update( + TEST_USER_ID, &created.id, UpdateProviderRequest { name: Some("New Name".into()), @@ -620,10 +638,11 @@ mod tests { #[tokio::test] async fn update_api_key_re_encrypts() { let svc = setup().await; - let created = svc.create(sample_create_request()).await.unwrap(); + let created = svc.create(TEST_USER_ID, sample_create_request()).await.unwrap(); let updated = svc .update( + TEST_USER_ID, &created.id, UpdateProviderRequest { api_key: Some("new-key-abcdefgh".into()), @@ -641,7 +660,7 @@ mod tests { async fn update_nonexistent_returns_not_found() { let svc = setup().await; let err = svc - .update("no_such_id", UpdateProviderRequest::default()) + .update(TEST_USER_ID, "no_such_id", UpdateProviderRequest::default()) .await .unwrap_err(); assert!(matches!(err, SystemError::NotFound(_))); @@ -650,17 +669,17 @@ mod tests { #[tokio::test] async fn delete_existing() { let svc = setup().await; - let created = svc.create(sample_create_request()).await.unwrap(); + let created = svc.create(TEST_USER_ID, sample_create_request()).await.unwrap(); - svc.delete(&created.id).await.unwrap(); - let all = svc.list().await.unwrap(); + svc.delete(TEST_USER_ID, &created.id).await.unwrap(); + let all = svc.list(TEST_USER_ID).await.unwrap(); assert!(all.is_empty()); } #[tokio::test] async fn delete_nonexistent_returns_not_found() { let svc = setup().await; - let err = svc.delete("no_such_id").await.unwrap_err(); + let err = svc.delete(TEST_USER_ID, "no_such_id").await.unwrap_err(); assert!(matches!(err, SystemError::NotFound(_))); } } diff --git a/crates/aionui-system/src/routes.rs b/crates/aionui-system/src/routes.rs index 347fc1d03..22831d819 100644 --- a/crates/aionui-system/src/routes.rs +++ b/crates/aionui-system/src/routes.rs @@ -106,8 +106,13 @@ pub fn settings_routes(state: SystemRouterState) -> Router { async fn get_settings( State(state): State, + Extension(user): Extension, ) -> Result>, ApiError> { - let settings = state.settings_service.get_settings().await.map_err(ApiError::from)?; + let settings = state + .settings_service + .get_settings(&user.id) + .await + .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(settings))) } @@ -126,12 +131,13 @@ async fn get_feedback_diagnostics( async fn update_settings( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; let settings = state .settings_service - .update_settings(req) + .update_settings(&user.id, req) .await .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(settings))) @@ -148,6 +154,7 @@ struct ClientPrefQuery { async fn get_client_preferences( State(state): State, + Extension(user): Extension, Query(query): Query, ) -> Result>, ApiError> { let keys_filter: Option> = query.keys.map(|k| { @@ -161,7 +168,7 @@ async fn get_client_preferences( let prefs = state .client_pref_service - .get_preferences(key_refs.as_deref()) + .get_preferences(&user.id, key_refs.as_deref()) .await .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(prefs))) @@ -169,12 +176,13 @@ async fn get_client_preferences( async fn update_client_preferences( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; state .client_pref_service - .update_preferences(req) + .update_preferences(&user.id, req) .await .map_err(ApiError::from)?; Ok(Json(ApiResponse::success())) @@ -186,47 +194,64 @@ async fn update_client_preferences( async fn list_providers( State(state): State, + Extension(user): Extension, ) -> Result>>, ApiError> { - let providers = state.provider_service.list().await.map_err(ApiError::from)?; + let providers = state.provider_service.list(&user.id).await.map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(providers))) } async fn create_provider( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result<(StatusCode, Json>), ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let provider = state.provider_service.create(req).await.map_err(ApiError::from)?; + let provider = state + .provider_service + .create(&user.id, req) + .await + .map_err(ApiError::from)?; Ok((StatusCode::CREATED, Json(ApiResponse::ok(provider)))) } async fn update_provider( State(state): State, + Extension(user): Extension, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let provider = state.provider_service.update(&id, req).await.map_err(ApiError::from)?; + let provider = state + .provider_service + .update(&user.id, &id, req) + .await + .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(provider))) } async fn delete_provider( State(state): State, + Extension(user): Extension, Path(id): Path, ) -> Result>, ApiError> { - state.provider_service.delete(&id).await.map_err(ApiError::from)?; + state + .provider_service + .delete(&user.id, &id) + .await + .map_err(ApiError::from)?; Ok(Json(ApiResponse::success())) } async fn fetch_models( State(state): State, + Extension(user): Extension, Path(id): Path, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; let result = state .model_fetch_service - .fetch_models(&id, &req) + .fetch_models(&user.id, &id, &req) .await .map_err(ApiError::from)?; Ok(Json(ApiResponse::ok(result))) @@ -282,9 +307,13 @@ async fn check_update( async fn ensure_node_runtime( State(state): State, + Extension(user): Extension, body: Result, JsonRejection>, ) -> Result>, ApiError> { let Json(req) = body.map_err(ApiError::from)?; - let result = state.runtime_prepare_service.ensure_node_runtime(req.scope).await?; + let result = state + .runtime_prepare_service + .ensure_node_runtime_for_user(&user.id, req.scope) + .await?; Ok(Json(ApiResponse::ok(result))) } diff --git a/crates/aionui-system/src/runtime_prepare.rs b/crates/aionui-system/src/runtime_prepare.rs index cf3c4a14a..5301c40e4 100644 --- a/crates/aionui-system/src/runtime_prepare.rs +++ b/crates/aionui-system/src/runtime_prepare.rs @@ -26,17 +26,39 @@ impl RuntimePrepareService { &self, scope: RuntimeStatusScope, ) -> Result { - let reporter = self.node_runtime_reporter(scope); + self.ensure_node_runtime_with_user(None, scope).await + } + + pub async fn ensure_node_runtime_for_user( + &self, + user_id: &str, + scope: RuntimeStatusScope, + ) -> Result { + self.ensure_node_runtime_with_user(Some(user_id.to_owned()), scope) + .await + } + + async fn ensure_node_runtime_with_user( + &self, + user_id: Option, + scope: RuntimeStatusScope, + ) -> Result { + let reporter = self.node_runtime_reporter(user_id, scope); ensure_node_runtime_with_reporter(Some(reporter.as_ref())) .await .map_err(|error| SystemError::BadRequest(error.to_string()))?; Ok(EnsureNodeRuntimeResponse { ready: true }) } - fn node_runtime_reporter(&self, scope: RuntimeStatusScope) -> SharedNodeRuntimeProgressReporter { + fn node_runtime_reporter( + &self, + user_id: Option, + scope: RuntimeStatusScope, + ) -> SharedNodeRuntimeProgressReporter { let broadcaster = self.broadcaster.clone(); Arc::new(move |update: NodeRuntimeProgress| { let payload = RuntimeStatusPayload { + user_id: user_id.clone(), resource: RuntimeResourceKind::Node, resource_id: None, scope: scope.clone(), @@ -75,3 +97,66 @@ fn map_failure_kind(kind: NodeRuntimeFailureKind) -> RuntimeFailureKind { NodeRuntimeFailureKind::Unknown => RuntimeFailureKind::Unknown, } } + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use aionui_api_types::{RuntimeStatusScopeKind, WebSocketMessage}; + + use super::*; + + #[derive(Default)] + struct RecordingBroadcaster { + events: Mutex>>, + } + + impl RecordingBroadcaster { + fn events(&self) -> Vec> { + self.events.lock().unwrap().clone() + } + } + + impl EventBroadcaster for RecordingBroadcaster { + fn broadcast(&self, event: WebSocketMessage) { + self.events.lock().unwrap().push(event); + } + } + + fn conversation_scope() -> RuntimeStatusScope { + RuntimeStatusScope { + kind: RuntimeStatusScopeKind::Conversation, + id: "conv-1".to_owned(), + } + } + + #[test] + fn node_runtime_reporter_scopes_route_event_to_user() { + let broadcaster = Arc::new(RecordingBroadcaster::default()); + let service = RuntimePrepareService::new(broadcaster.clone()); + let reporter = service.node_runtime_reporter(Some("user-1".to_owned()), conversation_scope()); + + reporter.report(NodeRuntimeProgress::downloading("downloading node")); + + let events = broadcaster.events(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].name, "runtime.statusChanged"); + assert_eq!(events[0].data["user_id"], "user-1"); + assert_eq!(events[0].data["resource"], "node"); + assert_eq!(events[0].data["scope"]["id"], "conv-1"); + } + + #[test] + fn node_runtime_reporter_can_emit_global_startup_event() { + let broadcaster = Arc::new(RecordingBroadcaster::default()); + let service = RuntimePrepareService::new(broadcaster.clone()); + let reporter = service.node_runtime_reporter(None, conversation_scope()); + + reporter.report(NodeRuntimeProgress::validating("validating node")); + + let events = broadcaster.events(); + assert_eq!(events.len(), 1); + assert!(events[0].data.get("user_id").is_none()); + assert_eq!(events[0].data["resource"], "node"); + } +} diff --git a/crates/aionui-system/src/settings.rs b/crates/aionui-system/src/settings.rs index b003813a8..d8b7322e4 100644 --- a/crates/aionui-system/src/settings.rs +++ b/crates/aionui-system/src/settings.rs @@ -23,10 +23,10 @@ impl SettingsService { } /// Get current system settings, falling back to defaults if not yet persisted. - pub async fn get_settings(&self) -> Result { + pub async fn get_settings(&self, user_id: &str) -> Result { let row = self .repo - .get_settings() + .get_settings(user_id) .await .map_err(|e| SystemError::Internal(format!("Failed to get settings: {e}")))?; @@ -42,13 +42,17 @@ impl SettingsService { } /// Partially update system settings. Only fields present in the request are changed. - pub async fn update_settings(&self, req: UpdateSettingsRequest) -> Result { + pub async fn update_settings( + &self, + user_id: &str, + req: UpdateSettingsRequest, + ) -> Result { if let Some(ref lang) = req.language { validate_language(lang)?; } // Merge with current settings (or defaults) - let current = self.get_settings().await?; + let current = self.get_settings(user_id).await?; let language = req.language.unwrap_or(current.language); let notification_enabled = req.notification_enabled.unwrap_or(current.notification_enabled); @@ -61,6 +65,7 @@ impl SettingsService { let row = self .repo .upsert_settings( + user_id, &language, notification_enabled, cron_notification_enabled, @@ -91,10 +96,21 @@ fn validate_language(lang: &str) -> Result<(), SystemError> { #[cfg(test)] mod tests { use super::*; + + const TEST_USER_ID: &str = "user-1"; use aionui_db::{SqliteSettingsRepository, init_database_memory}; async fn setup() -> SettingsService { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, '', 'active', 0, 1, 1)", + ) + .bind(TEST_USER_ID) + .bind(TEST_USER_ID) + .execute(db.pool()) + .await + .unwrap(); let repo = Arc::new(SqliteSettingsRepository::new(db.pool().clone())); // Leak the db handle so the pool stays alive for the test std::mem::forget(db); @@ -118,7 +134,7 @@ mod tests { #[tokio::test] async fn get_settings_returns_defaults_when_empty() { let svc = setup().await; - let settings = svc.get_settings().await.unwrap(); + let settings = svc.get_settings(TEST_USER_ID).await.unwrap(); assert_eq!(settings, SystemSettingsResponse::default()); } @@ -129,7 +145,7 @@ mod tests { language: Some("zh-CN".into()), ..Default::default() }; - let result = svc.update_settings(req).await.unwrap(); + let result = svc.update_settings(TEST_USER_ID, req).await.unwrap(); assert_eq!(result.language, "zh-CN"); // Other fields stay at defaults assert!(result.notification_enabled); @@ -144,7 +160,7 @@ mod tests { command_queue_enabled: Some(true), ..Default::default() }; - let result = svc.update_settings(req).await.unwrap(); + let result = svc.update_settings(TEST_USER_ID, req).await.unwrap(); assert!(!result.notification_enabled); assert!(result.command_queue_enabled); assert_eq!(result.language, "en-US"); @@ -153,7 +169,10 @@ mod tests { #[tokio::test] async fn update_empty_request_returns_current() { let svc = setup().await; - let result = svc.update_settings(UpdateSettingsRequest::default()).await.unwrap(); + let result = svc + .update_settings(TEST_USER_ID, UpdateSettingsRequest::default()) + .await + .unwrap(); assert_eq!(result, SystemSettingsResponse::default()); } @@ -164,22 +183,25 @@ mod tests { language: Some("invalid-lang".into()), ..Default::default() }; - let err = svc.update_settings(req).await.unwrap_err(); + let err = svc.update_settings(TEST_USER_ID, req).await.unwrap_err(); assert!(matches!(err, SystemError::BadRequest(_))); } #[tokio::test] async fn update_then_get_reflects_changes() { let svc = setup().await; - svc.update_settings(UpdateSettingsRequest { - language: Some("ja-JP".into()), - save_upload_to_workspace: Some(true), - ..Default::default() - }) + svc.update_settings( + TEST_USER_ID, + UpdateSettingsRequest { + language: Some("ja-JP".into()), + save_upload_to_workspace: Some(true), + ..Default::default() + }, + ) .await .unwrap(); - let settings = svc.get_settings().await.unwrap(); + let settings = svc.get_settings(TEST_USER_ID).await.unwrap(); assert_eq!(settings.language, "ja-JP"); assert!(settings.save_upload_to_workspace); } diff --git a/crates/aionui-system/tests/feedback_diagnostics_routes.rs b/crates/aionui-system/tests/feedback_diagnostics_routes.rs index 78427d1d8..bc1d915fe 100644 --- a/crates/aionui-system/tests/feedback_diagnostics_routes.rs +++ b/crates/aionui-system/tests/feedback_diagnostics_routes.rs @@ -14,7 +14,7 @@ use tower::ServiceExt; use aionui_db::{ SqliteClientPreferenceRepository, SqliteFeedbackDiagnosticsRepository, SqliteProviderRepository, - SqliteSettingsRepository, init_database_memory, + SqliteSettingsRepository, UserStatus, UserType, init_database_memory, }; const TEST_ENCRYPTION_KEY: [u8; 32] = [0x42; 32]; @@ -132,6 +132,8 @@ fn diagnostics_request(uri: &str) -> Request { req.extensions_mut().insert(CurrentUser { id: "system_default_user".to_owned(), username: "system_default_user".to_owned(), + user_type: UserType::Local, + status: UserStatus::Active, }); req } diff --git a/crates/aionui-system/tests/model_fetch_routes.rs b/crates/aionui-system/tests/model_fetch_routes.rs index ad96cab81..b99291058 100644 --- a/crates/aionui-system/tests/model_fetch_routes.rs +++ b/crates/aionui-system/tests/model_fetch_routes.rs @@ -13,10 +13,11 @@ use tower::ServiceExt; use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +use aionui_auth::CurrentUser; use aionui_common::encrypt_string; use aionui_db::{ CreateProviderParams, IProviderRepository, SqliteClientPreferenceRepository, SqliteFeedbackDiagnosticsRepository, - SqliteProviderRepository, SqliteSettingsRepository, init_database_memory, + SqliteProviderRepository, SqliteSettingsRepository, UserStatus, UserType, init_database_memory, }; use aionui_realtime::BroadcastEventBus; use aionui_system::{ @@ -29,6 +30,7 @@ use aionui_system::{ // --------------------------------------------------------------------------- const TEST_KEY: [u8; 32] = [0x42; 32]; +const TEST_USER_ID: &str = "user-1"; fn build_state(db: &aionui_db::Database) -> SystemRouterState { let provider_repo = Arc::new(SqliteProviderRepository::new(db.pool().clone())); @@ -49,6 +51,15 @@ fn build_state(db: &aionui_db::Database) -> SystemRouterState { async fn setup() -> (axum::Router, aionui_db::Database) { let db = init_database_memory().await.unwrap(); + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, '', 'active', 0, 1, 1)", + ) + .bind(TEST_USER_ID) + .bind(TEST_USER_ID) + .execute(db.pool()) + .await + .unwrap(); let state = build_state(&db); (system_routes(state), db) } @@ -58,6 +69,7 @@ async fn create_provider(db: &aionui_db::Database, platform: &str, base_url: &st let encrypted = encrypt_string(api_key, &TEST_KEY).unwrap(); let row = repo .create(CreateProviderParams { + user_id: TEST_USER_ID, id: None, platform, name: "Test Provider", @@ -85,12 +97,19 @@ async fn body_json(resp: axum::response::Response) -> serde_json::Value { } fn post_request(uri: &str, body: serde_json::Value) -> Request { - Request::builder() + let mut req = Request::builder() .method("POST") .uri(uri) .header("content-type", "application/json") .body(Body::from(body.to_string())) - .unwrap() + .unwrap(); + req.extensions_mut().insert(CurrentUser { + id: TEST_USER_ID.to_owned(), + username: TEST_USER_ID.to_owned(), + user_type: UserType::Local, + status: UserStatus::Active, + }); + req } // --------------------------------------------------------------------------- diff --git a/crates/aionui-system/tests/provider_routes.rs b/crates/aionui-system/tests/provider_routes.rs index 6f15eb46f..4e233346f 100644 --- a/crates/aionui-system/tests/provider_routes.rs +++ b/crates/aionui-system/tests/provider_routes.rs @@ -13,9 +13,10 @@ use http_body_util::BodyExt; use serde_json::json; use tower::ServiceExt; +use aionui_auth::CurrentUser; use aionui_db::{ SqliteClientPreferenceRepository, SqliteFeedbackDiagnosticsRepository, SqliteProviderRepository, - SqliteSettingsRepository, init_database_memory, + SqliteSettingsRepository, UserStatus, UserType, init_database_memory, }; use aionui_system::{ ClientPrefService, FeedbackDiagnosticsService, ModelFetchService, ProtocolDetectionService, ProviderService, @@ -27,6 +28,8 @@ use aionui_system::{ // --------------------------------------------------------------------------- const TEST_ENCRYPTION_KEY: [u8; 32] = [0x42; 32]; +const TEST_USER_ID: &str = "user-1"; +const OTHER_USER_ID: &str = "user-2"; fn build_state(db: &aionui_db::Database) -> SystemRouterState { let provider_repo = Arc::new(SqliteProviderRepository::new(db.pool().clone())); @@ -47,6 +50,17 @@ fn build_state(db: &aionui_db::Database) -> SystemRouterState { async fn setup() -> (axum::Router, aionui_db::Database) { let db = init_database_memory().await.unwrap(); + for user_id in [TEST_USER_ID, OTHER_USER_ID] { + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, '', 'active', 0, 1, 1)", + ) + .bind(user_id) + .bind(user_id) + .execute(db.pool()) + .await + .unwrap(); + } let state = build_state(&db); (system_routes(state), db) } @@ -57,24 +71,57 @@ async fn body_json(resp: axum::response::Response) -> serde_json::Value { } fn get_request(uri: &str) -> Request { - Request::builder().method("GET").uri(uri).body(Body::empty()).unwrap() + get_request_for_user(TEST_USER_ID, uri) +} + +fn get_request_for_user(user_id: &str, uri: &str) -> Request { + let mut req = Request::builder().method("GET").uri(uri).body(Body::empty()).unwrap(); + req.extensions_mut().insert(CurrentUser { + id: user_id.to_owned(), + username: user_id.to_owned(), + user_type: UserType::Local, + status: UserStatus::Active, + }); + req } fn json_request(method: &str, uri: &str, body: serde_json::Value) -> Request { - Request::builder() + json_request_for_user(TEST_USER_ID, method, uri, body) +} + +fn json_request_for_user(user_id: &str, method: &str, uri: &str, body: serde_json::Value) -> Request { + let mut req = Request::builder() .method(method) .uri(uri) .header("content-type", "application/json") .body(Body::from(serde_json::to_vec(&body).unwrap())) - .unwrap() + .unwrap(); + req.extensions_mut().insert(CurrentUser { + id: user_id.to_owned(), + username: user_id.to_owned(), + user_type: UserType::Local, + status: UserStatus::Active, + }); + req } fn delete_request(uri: &str) -> Request { - Request::builder() + delete_request_for_user(TEST_USER_ID, uri) +} + +fn delete_request_for_user(user_id: &str, uri: &str) -> Request { + let mut req = Request::builder() .method("DELETE") .uri(uri) .body(Body::empty()) - .unwrap() + .unwrap(); + req.extensions_mut().insert(CurrentUser { + id: user_id.to_owned(), + username: user_id.to_owned(), + user_type: UserType::Local, + status: UserStatus::Active, + }); + req } fn sample_create_body() -> serde_json::Value { @@ -183,6 +230,33 @@ async fn create_provider_with_supplied_id() { assert_eq!(data["model_enabled"]["gpt-3.5"], false); } +#[tokio::test] +async fn create_provider_ignores_body_user_id() { + let (app, db) = setup().await; + let mut body = sample_create_body(); + body["user_id"] = json!(OTHER_USER_ID); + + let resp = app.oneshot(json_request("POST", "/api/providers", body)).await.unwrap(); + assert_eq!(resp.status(), StatusCode::CREATED); + let json = body_json(resp).await; + let id = json["data"]["id"].as_str().unwrap(); + + let owner: String = sqlx::query_scalar("SELECT user_id FROM providers WHERE id = ?") + .bind(id) + .fetch_one(db.pool()) + .await + .unwrap(); + assert_eq!(owner, TEST_USER_ID); + + let other_app = system_routes(build_state(&db)); + let resp = other_app + .oneshot(get_request_for_user(OTHER_USER_ID, "/api/providers")) + .await + .unwrap(); + let json = body_json(resp).await; + assert_eq!(json["data"], json!([])); +} + #[tokio::test] async fn create_provider_with_duplicate_id_returns_conflict() { let (_app, db) = setup().await; @@ -447,6 +521,37 @@ async fn update_provider_nonexistent() { assert_eq!(resp.status(), StatusCode::NOT_FOUND); } +#[tokio::test] +async fn cross_user_provider_update_delete_are_not_found() { + let (_app, db) = setup().await; + let (_, id) = create_one(&db).await; + + let update_app = system_routes(build_state(&db)); + let resp = update_app + .oneshot(json_request_for_user( + OTHER_USER_ID, + "PUT", + &format!("/api/providers/{id}"), + json!({"name": "Other User Update"}), + )) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + + let delete_app = system_routes(build_state(&db)); + let resp = delete_app + .oneshot(delete_request_for_user(OTHER_USER_ID, &format!("/api/providers/{id}"))) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + + let owner_app = system_routes(build_state(&db)); + let resp = owner_app.oneshot(get_request("/api/providers")).await.unwrap(); + let json = body_json(resp).await; + assert_eq!(json["data"][0]["id"], id); + assert_eq!(json["data"][0]["name"], "Anthropic"); +} + // =========================================================================== // DELETE /api/providers/{id} // =========================================================================== diff --git a/crates/aionui-system/tests/settings_routes.rs b/crates/aionui-system/tests/settings_routes.rs index c89566903..49a4b3d0d 100644 --- a/crates/aionui-system/tests/settings_routes.rs +++ b/crates/aionui-system/tests/settings_routes.rs @@ -12,9 +12,10 @@ use axum::http::{Request, StatusCode}; use http_body_util::BodyExt; use tower::ServiceExt; +use aionui_auth::CurrentUser; use aionui_db::{ SqliteClientPreferenceRepository, SqliteFeedbackDiagnosticsRepository, SqliteProviderRepository, - SqliteSettingsRepository, init_database_memory, + SqliteSettingsRepository, UserStatus, UserType, init_database_memory, }; use aionui_system::{ ClientPrefService, FeedbackDiagnosticsService, ModelFetchService, ProtocolDetectionService, ProviderService, @@ -26,6 +27,8 @@ use aionui_system::{ // --------------------------------------------------------------------------- const TEST_ENCRYPTION_KEY: [u8; 32] = [0x42; 32]; +const TEST_USER_ID: &str = "user-1"; +const OTHER_USER_ID: &str = "user-2"; fn build_state(db: &aionui_db::Database) -> SystemRouterState { let provider_repo = Arc::new(SqliteProviderRepository::new(db.pool().clone())); @@ -46,6 +49,17 @@ fn build_state(db: &aionui_db::Database) -> SystemRouterState { async fn setup() -> (axum::Router, aionui_db::Database) { let db = init_database_memory().await.unwrap(); + for user_id in [TEST_USER_ID, OTHER_USER_ID] { + sqlx::query( + "INSERT INTO users (id, user_type, username, password_hash, status, session_generation, created_at, updated_at) \ + VALUES (?, 'local', ?, '', 'active', 0, 1, 1)", + ) + .bind(user_id) + .bind(user_id) + .execute(db.pool()) + .await + .unwrap(); + } let state = build_state(&db); (settings_routes(state), db) } @@ -56,16 +70,38 @@ async fn body_json(resp: axum::response::Response) -> serde_json::Value { } fn get_request(uri: &str) -> Request { - Request::builder().method("GET").uri(uri).body(Body::empty()).unwrap() + get_request_for_user(TEST_USER_ID, uri) +} + +fn get_request_for_user(user_id: &str, uri: &str) -> Request { + let mut req = Request::builder().method("GET").uri(uri).body(Body::empty()).unwrap(); + req.extensions_mut().insert(CurrentUser { + id: user_id.to_owned(), + username: user_id.to_owned(), + user_type: UserType::Local, + status: UserStatus::Active, + }); + req } fn json_request(method: &str, uri: &str, body: serde_json::Value) -> Request { - Request::builder() + json_request_for_user(TEST_USER_ID, method, uri, body) +} + +fn json_request_for_user(user_id: &str, method: &str, uri: &str, body: serde_json::Value) -> Request { + let mut req = Request::builder() .method(method) .uri(uri) .header("content-type", "application/json") .body(Body::from(serde_json::to_vec(&body).unwrap())) - .unwrap() + .unwrap(); + req.extensions_mut().insert(CurrentUser { + id: user_id.to_owned(), + username: user_id.to_owned(), + user_type: UserType::Local, + status: UserStatus::Active, + }); + req } // =========================================================================== @@ -190,6 +226,44 @@ async fn patch_then_get_reflects_changes() { assert_eq!(json["data"]["save_upload_to_workspace"], true); } +#[tokio::test] +async fn settings_are_scoped_by_current_user() { + let (app, db) = setup().await; + + let resp = app + .oneshot(json_request_for_user( + TEST_USER_ID, + "PATCH", + "/api/settings", + serde_json::json!({ + "language": "zh-CN", + "notification_enabled": false, + "user_id": OTHER_USER_ID + }), + )) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let owner_app = settings_routes(build_state(&db)); + let owner_resp = owner_app + .oneshot(get_request_for_user(TEST_USER_ID, "/api/settings")) + .await + .unwrap(); + let owner_json = body_json(owner_resp).await; + assert_eq!(owner_json["data"]["language"], "zh-CN"); + assert_eq!(owner_json["data"]["notification_enabled"], false); + + let other_app = settings_routes(build_state(&db)); + let other_resp = other_app + .oneshot(get_request_for_user(OTHER_USER_ID, "/api/settings")) + .await + .unwrap(); + let other_json = body_json(other_resp).await; + assert_eq!(other_json["data"]["language"], "en-US"); + assert_eq!(other_json["data"]["notification_enabled"], true); +} + // =========================================================================== // Client Preferences (GET/PUT /api/settings/client) // =========================================================================== @@ -295,6 +369,41 @@ async fn put_batch_write() { assert_eq!(json["data"]["c"], true); } +#[tokio::test] +async fn client_preferences_are_scoped_by_current_user() { + let (app, db) = setup().await; + + let resp = app + .oneshot(json_request_for_user( + TEST_USER_ID, + "PUT", + "/api/settings/client", + serde_json::json!({ + "theme": "dark", + "user_id": OTHER_USER_ID + }), + )) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let owner_app = settings_routes(build_state(&db)); + let owner_resp = owner_app + .oneshot(get_request_for_user(TEST_USER_ID, "/api/settings/client")) + .await + .unwrap(); + let owner_json = body_json(owner_resp).await; + assert_eq!(owner_json["data"]["theme"], "dark"); + + let other_app = settings_routes(build_state(&db)); + let other_resp = other_app + .oneshot(get_request_for_user(OTHER_USER_ID, "/api/settings/client")) + .await + .unwrap(); + let other_json = body_json(other_resp).await; + assert_eq!(other_json["data"], serde_json::json!({})); +} + #[tokio::test] async fn get_client_prefs_with_keys_filter() { let (app, db) = setup().await; diff --git a/crates/aionui-team/src/event_loop.rs b/crates/aionui-team/src/event_loop.rs index d65299722..81c9d056c 100644 --- a/crates/aionui-team/src/event_loop.rs +++ b/crates/aionui-team/src/event_loop.rs @@ -370,7 +370,11 @@ async fn mark_batch_messages_read(ctx: &AgentLoopContext, batch: &WorkBatch) { if batch.mailbox_message_ids.is_empty() { return; } - if let Err(error) = ctx.mailbox.mark_read_batch(&batch.mailbox_message_ids).await { + if let Err(error) = ctx + .mailbox + .mark_read_batch(&ctx.team_id, &batch.mailbox_message_ids) + .await + { warn!( team_id = %ctx.team_id, slot_id = %ctx.slot_id, diff --git a/crates/aionui-team/src/events.rs b/crates/aionui-team/src/events.rs index 5121e2740..3ef28cc91 100644 --- a/crates/aionui-team/src/events.rs +++ b/crates/aionui-team/src/events.rs @@ -6,6 +6,8 @@ use aionui_api_types::{ WebSocketMessage, }; use aionui_realtime::EventBroadcaster; +use serde::Serialize; +use serde_json::Value; use tracing::{debug, info}; use crate::types::{TeamAgent, TeammateStatus}; @@ -36,28 +38,36 @@ pub const TEAM_SLOT_WORK_CHANGED_EVENT: &str = "team.slotWorkChanged"; pub struct TeamEventEmitter { team_id: String, + user_id: String, broadcaster: Arc, } impl TeamEventEmitter { - pub fn new(team_id: String, broadcaster: Arc) -> Self { - Self { team_id, broadcaster } + pub fn new(team_id: String, user_id: String, broadcaster: Arc) -> Self { + Self { + team_id, + user_id, + broadcaster, + } } pub fn team_id(&self) -> &str { &self.team_id } + fn scoped_payload(&self, payload: T) -> Value { + let mut value = serde_json::to_value(payload).expect("serialize team event payload"); + value["user_id"] = Value::String(self.user_id.clone()); + value + } + pub fn broadcast_agent_status(&self, slot_id: &str, status: TeammateStatus) { let payload = TeamAgentStatusPayload { team_id: self.team_id.clone(), slot_id: slot_id.to_owned(), status: status.to_string(), }; - let event = WebSocketMessage::new( - TEAM_AGENT_STATUS_CHANGED_EVENT, - serde_json::to_value(payload).expect("serialize status payload"), - ); + let event = WebSocketMessage::new(TEAM_AGENT_STATUS_CHANGED_EVENT, self.scoped_payload(payload)); self.broadcaster.broadcast(event); } @@ -66,10 +76,7 @@ impl TeamEventEmitter { team_id: self.team_id.clone(), assistant: agent.to_response(), }; - let event = WebSocketMessage::new( - TEAM_AGENT_SPAWNED_EVENT, - serde_json::to_value(payload).expect("serialize spawned payload"), - ); + let event = WebSocketMessage::new(TEAM_AGENT_SPAWNED_EVENT, self.scoped_payload(payload)); self.broadcaster.broadcast(event); } @@ -78,10 +85,7 @@ impl TeamEventEmitter { team_id: self.team_id.clone(), slot_id: slot_id.to_owned(), }; - let event = WebSocketMessage::new( - TEAM_AGENT_REMOVED_EVENT, - serde_json::to_value(payload).expect("serialize removed payload"), - ); + let event = WebSocketMessage::new(TEAM_AGENT_REMOVED_EVENT, self.scoped_payload(payload)); self.broadcaster.broadcast(event); } @@ -91,10 +95,7 @@ impl TeamEventEmitter { slot_id: slot_id.to_owned(), name: name.to_owned(), }; - let event = WebSocketMessage::new( - TEAM_AGENT_RENAMED_EVENT, - serde_json::to_value(payload).expect("serialize renamed payload"), - ); + let event = WebSocketMessage::new(TEAM_AGENT_RENAMED_EVENT, self.scoped_payload(payload)); self.broadcaster.broadcast(event); } @@ -123,10 +124,9 @@ impl TeamEventEmitter { error = payload.error.as_deref().unwrap_or(""), "team member runtime status broadcast" ); - let event = WebSocketMessage::new( - TEAM_AGENT_RUNTIME_STATUS_CHANGED_EVENT, - serde_json::to_value(payload).expect("serialize agent runtime status payload"), - ); + // Keep per-user scoping so the event is delivered only to the owning + // user's WebSocket subscribers. + let event = WebSocketMessage::new(TEAM_AGENT_RUNTIME_STATUS_CHANGED_EVENT, self.scoped_payload(payload)); self.broadcaster.broadcast(event); } @@ -145,10 +145,7 @@ impl TeamEventEmitter { slot_work_count = payload.slot_work.len(), "team websocket event emitted" ); - let event = WebSocketMessage::new( - event_name, - serde_json::to_value(payload).expect("serialize team run payload"), - ); + let event = WebSocketMessage::new(event_name, self.scoped_payload(payload)); self.broadcaster.broadcast(event); } @@ -180,10 +177,7 @@ impl TeamEventEmitter { status = ?payload.status, "team websocket event emitted" ); - let event = WebSocketMessage::new( - event_name, - serde_json::to_value(payload).expect("serialize team child turn payload"), - ); + let event = WebSocketMessage::new(event_name, self.scoped_payload(payload)); self.broadcaster.broadcast(event); } } @@ -221,7 +215,7 @@ mod tests { fn make_emitter() -> (TeamEventEmitter, Arc) { let bc = Arc::new(RecordingBroadcaster::new()); - let emitter = TeamEventEmitter::new("team-1".into(), bc.clone()); + let emitter = TeamEventEmitter::new("team-1".into(), "user-1".into(), bc.clone()); (emitter, bc) } @@ -233,6 +227,7 @@ mod tests { let events = bc.events(); assert_eq!(events.len(), 1); assert_eq!(events[0].name, "team.agentStatusChanged"); + assert_eq!(events[0].data["user_id"], "user-1"); let payload: TeamAgentStatusPayload = serde_json::from_value(events[0].data.clone()).unwrap(); assert_eq!(payload.team_id, "team-1"); diff --git a/crates/aionui-team/src/mailbox.rs b/crates/aionui-team/src/mailbox.rs index 060170144..08566dc8d 100644 --- a/crates/aionui-team/src/mailbox.rs +++ b/crates/aionui-team/src/mailbox.rs @@ -10,11 +10,19 @@ use crate::types::{MailboxMessage, MailboxMessageType}; pub struct Mailbox { repo: Arc, + user_id: String, } impl Mailbox { pub fn new(repo: Arc) -> Self { - Self { repo } + Self::new_for_user(repo, "system_default_user") + } + + pub fn new_for_user(repo: Arc, user_id: impl Into) -> Self { + Self { + repo, + user_id: user_id.into(), + } } pub async fn write( @@ -57,7 +65,7 @@ impl Mailbox { created_at: now_ms(), }; - self.repo.write_message(&row).await?; + self.repo.write_message(&self.user_id, &row).await?; debug!( team_id, @@ -72,7 +80,7 @@ impl Mailbox { } pub async fn read_unread(&self, team_id: &str, agent_id: &str) -> Result, TeamError> { - let rows = self.repo.read_unread_and_mark(team_id, agent_id).await?; + let rows = self.repo.read_unread_and_mark(&self.user_id, team_id, agent_id).await?; debug!(team_id, agent_id, count = rows.len(), "mailbox unread messages read"); @@ -83,15 +91,15 @@ impl Mailbox { /// Reads all unread messages without marking them as read. /// Used by the drain_mailbox pattern: peek → prompt → mark_read on success. pub async fn peek_unread(&self, team_id: &str, agent_id: &str) -> Result, TeamError> { - let rows = self.repo.peek_unread(team_id, agent_id).await?; + let rows = self.repo.peek_unread(&self.user_id, team_id, agent_id).await?; debug!(team_id, agent_id, count = rows.len(), "mailbox peek_unread"); let messages = rows.iter().filter_map(MailboxMessage::from_row).collect(); Ok(messages) } /// Marks the given message IDs as read. Called after successful prompt delivery. - pub async fn mark_read_batch(&self, ids: &[String]) -> Result<(), TeamError> { - self.repo.mark_read_batch(ids).await?; + pub async fn mark_read_batch(&self, team_id: &str, ids: &[String]) -> Result<(), TeamError> { + self.repo.mark_read_batch(&self.user_id, team_id, ids).await?; Ok(()) } @@ -109,7 +117,7 @@ impl Mailbox { } let count = ids.len(); if !ids.is_empty() { - self.mark_read_batch(&ids).await?; + self.mark_read_batch(team_id, &ids).await?; } Ok(count) } @@ -120,18 +128,23 @@ impl Mailbox { agent_id: &str, limit: Option, ) -> Result, TeamError> { - let rows = self.repo.get_history(team_id, agent_id, limit).await?; + let rows = self.repo.get_history(&self.user_id, team_id, agent_id, limit).await?; let messages = rows.iter().filter_map(MailboxMessage::from_row).collect(); Ok(messages) } pub async fn has_unread(&self, team_id: &str, agent_id: &str) -> Result { - let rows = self.repo.get_history(team_id, agent_id, None).await?; + let rows = self.repo.get_history(&self.user_id, team_id, agent_id, None).await?; Ok(rows.iter().any(|r| !r.read)) } pub async fn delete_by_team(&self, team_id: &str) -> Result<(), TeamError> { - self.repo.delete_mailbox_by_team(team_id).await?; + let row = self + .repo + .get_team_for_restore(team_id) + .await? + .ok_or_else(|| TeamError::TeamNotFound(team_id.to_owned()))?; + self.repo.delete_mailbox_by_team(&row.user_id, team_id).await?; debug!(team_id, "mailbox messages deleted for team"); Ok(()) } diff --git a/crates/aionui-team/src/mcp/server.rs b/crates/aionui-team/src/mcp/server.rs index c17f6ff47..bc9427e57 100644 --- a/crates/aionui-team/src/mcp/server.rs +++ b/crates/aionui-team/src/mcp/server.rs @@ -533,13 +533,17 @@ pub(crate) async fn dispatch_tool( "team_shutdown_agent" => { exec_shutdown_agent(arguments, scheduler, service, team_id, caller_slot_id, caller_role).await } - "team_list_assistants" => exec_list_assistants(arguments, service).await, - "team_describe_assistant" => exec_describe_assistant(arguments, service).await, + "team_list_assistants" => exec_list_assistants(arguments, service, team_id).await, + "team_describe_assistant" => exec_describe_assistant(arguments, service, team_id).await, _ => Err(ToolCallError::from_message(format!("Unknown tool: {tool_name}"))), } } -async fn exec_list_assistants(args: &Value, service: &Weak) -> Result { +async fn exec_list_assistants( + args: &Value, + service: &Weak, + team_id: &str, +) -> Result { let props = args.as_object().cloned().unwrap_or_default(); if !props.is_empty() { return Err(ToolCallError::from_message( @@ -549,12 +553,20 @@ async fn exec_list_assistants(args: &Value, service: &Weak) let service = service .upgrade() .ok_or_else(|| ToolCallError::from_message("Team service not available"))?; - let assistants = service.list_team_selectable_assistants().await; + let user_id = service + .team_owner_user_id(team_id) + .await + .map_err(|error| ToolCallError::from_message(error.to_string()))?; + let assistants = service.list_team_selectable_assistants(&user_id).await; let value = json!({ "assistants": assistants }); serde_json::to_string_pretty(&value).map_err(|e| ToolCallError::from_message(format!("Serialization error: {e}"))) } -async fn exec_describe_assistant(args: &Value, service: &Weak) -> Result { +async fn exec_describe_assistant( + args: &Value, + service: &Weak, + team_id: &str, +) -> Result { if args.get("custom_agent_id").is_some() { return Err(ToolCallError::from_message( "custom_agent_id is no longer accepted; use assistant_id", @@ -570,9 +582,13 @@ async fn exec_describe_assistant(args: &Value, service: &Weak = Weak::new(); - let result = exec_list_assistants(&json!({}), &service).await; + let result = exec_list_assistants(&json!({}), &service, "team-1").await; let err = result.expect_err("dead service should be surfaced"); assert!( err.message.contains("Team service not available"), diff --git a/crates/aionui-team/src/message_projection.rs b/crates/aionui-team/src/message_projection.rs index fbe87e316..ad373af5a 100644 --- a/crates/aionui-team/src/message_projection.rs +++ b/crates/aionui-team/src/message_projection.rs @@ -26,6 +26,7 @@ pub enum TeamProjectionSource { #[derive(Debug, Clone)] pub struct TeamProjectionRequest { + pub user_id: String, pub team_id: String, pub slot_id: String, pub conversation_id: String, @@ -38,6 +39,7 @@ pub struct TeamProjectionRequest { impl TeamProjectionRequest { pub fn user_visible( + user_id: impl Into, team_id: impl Into, slot_id: impl Into, conversation_id: impl Into, @@ -45,6 +47,7 @@ impl TeamProjectionRequest { files: Vec, ) -> Self { Self { + user_id: user_id.into(), team_id: team_id.into(), slot_id: slot_id.into(), conversation_id: conversation_id.into(), @@ -56,7 +59,9 @@ impl TeamProjectionRequest { } } + #[allow(clippy::too_many_arguments)] pub fn teammate_visible( + user_id: impl Into, team_id: impl Into, slot_id: impl Into, conversation_id: impl Into, @@ -69,6 +74,7 @@ impl TeamProjectionRequest { let conversation_id = conversation_id.into(); let mailbox_message_id = mailbox_message_id.into(); Self { + user_id: user_id.into(), dedupe_key: Some(teammate_dedupe_key(&team_id, &mailbox_message_id, &conversation_id)), team_id, slot_id: slot_id.into(), @@ -86,6 +92,7 @@ impl TeamProjectionRequest { } pub fn team_system_visible( + user_id: impl Into, team_id: impl Into, slot_id: impl Into, conversation_id: impl Into, @@ -96,6 +103,7 @@ impl TeamProjectionRequest { let conversation_id = conversation_id.into(); let mailbox_message_id = mailbox_message_id.into(); Self { + user_id: user_id.into(), dedupe_key: Some(teammate_dedupe_key(&team_id, &mailbox_message_id, &conversation_id)), team_id, slot_id: slot_id.into(), @@ -202,6 +210,7 @@ where sender_backend, sender_conversation_id, } => Some(serde_json::json!({ + "user_id": request.user_id, "team_id": request.team_id, "slot_id": request.slot_id, "conversation_id": request.conversation_id, @@ -214,6 +223,7 @@ where "sender_conversation_id": sender_conversation_id, })), TeamProjectionSource::TeamSystem => Some(serde_json::json!({ + "user_id": request.user_id, "team_id": request.team_id, "slot_id": request.slot_id, "conversation_id": request.conversation_id, diff --git a/crates/aionui-team/src/ports.rs b/crates/aionui-team/src/ports.rs index 98d625d09..2cb8e26cb 100644 --- a/crates/aionui-team/src/ports.rs +++ b/crates/aionui-team/src/ports.rs @@ -27,14 +27,16 @@ pub trait TeamConversationLookupPort: Send + Sync { #[async_trait] pub trait TeamAssistantCatalogPort: Send + Sync { - async fn list_team_selectable_assistants(&self) -> Result, TeamError>; + async fn list_team_selectable_assistants(&self, user_id: &str) + -> Result, TeamError>; async fn resolve_team_selectable_assistant( &self, + user_id: &str, assistant_id: &str, ) -> Result, TeamError> { Ok(self - .list_team_selectable_assistants() + .list_team_selectable_assistants(user_id) .await? .into_iter() .find(|assistant| assistant.assistant_id == assistant_id)) diff --git a/crates/aionui-team/src/provisioning.rs b/crates/aionui-team/src/provisioning.rs index 75931fee4..5cf80999f 100644 --- a/crates/aionui-team/src/provisioning.rs +++ b/crates/aionui-team/src/provisioning.rs @@ -178,7 +178,7 @@ impl TeamAgentProvisioner { let leader_role = TeammateRole::Lead; let leader_assistant_id = Self::effective_assistant_id(leader_input.assistant_id.as_deref()); let leader_backend = self - .resolve_requested_backend(leader_input.backend.as_deref(), leader_assistant_id.as_deref()) + .resolve_requested_backend(user_id, leader_input.backend.as_deref(), leader_assistant_id.as_deref()) .await?; let leader_conversation = self .create_team_conversation_for_agent( @@ -229,7 +229,7 @@ impl TeamAgentProvisioner { let slot_id = generate_id(); let assistant_id = Self::effective_assistant_id(input.assistant_id.as_deref()); let backend = self - .resolve_requested_backend(input.backend.as_deref(), assistant_id.as_deref()) + .resolve_requested_backend(user_id, input.backend.as_deref(), assistant_id.as_deref()) .await?; let conversation = self .create_team_conversation_for_agent( @@ -294,7 +294,7 @@ impl TeamAgentProvisioner { let workspace = self.workspace_resolver().resolve_for_new_agent(row, team).await?; let assistant_id = Self::effective_assistant_id(req.assistant_id.as_deref()); let backend = self - .resolve_requested_backend(req.backend.as_deref(), assistant_id.as_deref()) + .resolve_requested_backend(user_id, req.backend.as_deref(), assistant_id.as_deref()) .await?; let agent = self .provision_new_agent(NewAgentProvisioning { @@ -317,6 +317,7 @@ impl TeamAgentProvisioner { async fn resolve_requested_backend( &self, + user_id: &str, requested_backend: Option<&str>, assistant_id: Option<&str>, ) -> Result { @@ -324,7 +325,7 @@ impl TeamAgentProvisioner { if let Some(assistant_id) = assistant_id { return self .assistant_catalog - .resolve_team_selectable_assistant(assistant_id) + .resolve_team_selectable_assistant(user_id, assistant_id) .await? .map(|assistant| assistant.backend) .ok_or_else(|| { @@ -343,7 +344,7 @@ impl TeamAgentProvisioner { pub(crate) async fn persist_spawned_agent(&self, req: PersistSpawnedAgentRequest) -> Result { let row = self .repo - .get_team(&req.team_id) + .get_team(&req.user_id, &req.team_id) .await? .ok_or_else(|| TeamError::TeamNotFound(req.team_id.clone()))?; let mut team = Team::from_row(&row)?; @@ -375,10 +376,13 @@ impl TeamAgentProvisioner { task_manager: &Arc, ) -> Result<(), TeamError> { let team_id = mcp_stdio_cfg.team_id.clone(); - let transport = self.team_tool_transport(agent).await?; + let transport = self.team_tool_transport(user_id, agent).await?; match transport { - TeamToolTransport::Mcp => self.write_team_mcp_runtime_config(agent, mcp_stdio_cfg).await?, - TeamToolTransport::CliAssumed => self.write_team_cli_runtime_config(agent).await?, + TeamToolTransport::Mcp => { + self.write_team_mcp_runtime_config(user_id, agent, mcp_stdio_cfg) + .await? + } + TeamToolTransport::CliAssumed => self.write_team_cli_runtime_config(user_id, agent).await?, } task_manager .kill_and_wait(&agent.conversation_id, Some(AgentKillReason::TeamMcpRebuild)) @@ -401,8 +405,12 @@ impl TeamAgentProvisioner { Ok(()) } - pub(crate) async fn team_tool_transport(&self, agent: &TeamAgent) -> Result { - let capabilities = self.agent_capabilities(&agent.backend).await?; + pub(crate) async fn team_tool_transport( + &self, + user_id: &str, + agent: &TeamAgent, + ) -> Result { + let capabilities = self.agent_capabilities(user_id, &agent.backend).await?; if supports_team_mcp_backend(&agent.backend, capabilities.as_ref()) { return Ok(TeamToolTransport::Mcp); } @@ -415,8 +423,8 @@ impl TeamAgentProvisioner { ))) } - async fn agent_capabilities(&self, backend: &str) -> Result, TeamError> { - let Some(metadata) = acp_backend_metadata(&self.agent_metadata_repo, backend).await? else { + async fn agent_capabilities(&self, user_id: &str, backend: &str) -> Result, TeamError> { + let Some(metadata) = acp_backend_metadata(&self.agent_metadata_repo, user_id, backend).await? else { return Ok(None); }; let Some(raw) = metadata @@ -432,10 +440,11 @@ impl TeamAgentProvisioner { pub(crate) async fn write_team_mcp_runtime_config( &self, + user_id: &str, agent: &TeamAgent, mcp_stdio_cfg: TeamMcpStdioConfig, ) -> Result<(), TeamError> { - let acp_metadata = acp_backend_metadata(&self.agent_metadata_repo, &agent.backend).await?; + let acp_metadata = acp_backend_metadata(&self.agent_metadata_repo, user_id, &agent.backend).await?; let agent_type = if acp_metadata.is_some() { AgentType::Acp } else { @@ -457,8 +466,12 @@ impl TeamAgentProvisioner { }) } - pub(crate) async fn write_team_cli_runtime_config(&self, agent: &TeamAgent) -> Result<(), TeamError> { - let acp_metadata = acp_backend_metadata(&self.agent_metadata_repo, &agent.backend).await?; + pub(crate) async fn write_team_cli_runtime_config( + &self, + user_id: &str, + agent: &TeamAgent, + ) -> Result<(), TeamError> { + let acp_metadata = acp_backend_metadata(&self.agent_metadata_repo, user_id, &agent.backend).await?; let agent_type = if acp_metadata.is_some() { AgentType::Acp } else { @@ -539,7 +552,7 @@ impl TeamAgentProvisioner { workspace: Option<&str>, session_mode: Option<&str>, ) -> Result { - let acp_metadata = acp_backend_metadata(&self.agent_metadata_repo, backend).await?; + let acp_metadata = acp_backend_metadata(&self.agent_metadata_repo, user_id, backend).await?; let agent_type = if acp_metadata.is_some() { AgentType::Acp } else { @@ -558,7 +571,7 @@ impl TeamAgentProvisioner { session_mode, ); let provider_id = if agent_type == AgentType::Aionrs { - self.resolve_provider_for_model(model) + self.resolve_provider_for_model(user_id, model) .await .unwrap_or_else(|| backend.to_owned()) } else { @@ -685,8 +698,14 @@ impl TeamAgentProvisioner { async fn persist_agents(&self, team_id: &str, agents: &[TeamAgent]) -> Result<(), TeamError> { let agents_json = serde_json::to_string(agents)?; + let row = self + .repo + .get_team_for_restore(team_id) + .await? + .ok_or_else(|| TeamError::TeamNotFound(team_id.to_owned()))?; self.repo .update_team( + &row.user_id, team_id, &UpdateTeamParams { agents: Some(agents_json), @@ -697,8 +716,8 @@ impl TeamAgentProvisioner { Ok(()) } - async fn resolve_provider_for_model(&self, model: &str) -> Option { - let providers = self.provider_repo.list().await.ok()?; + async fn resolve_provider_for_model(&self, user_id: &str, model: &str) -> Option { + let providers = self.provider_repo.list(user_id).await.ok()?; for provider in providers { if !provider.enabled { continue; @@ -851,6 +870,7 @@ mod tests { impl TeamAssistantCatalogPort for EmptyTeamAssistantCatalog { async fn list_team_selectable_assistants( &self, + _user_id: &str, ) -> Result, TeamError> { Ok(Vec::new()) } @@ -861,9 +881,15 @@ mod tests { async fn list_all(&self) -> Result, DbError> { Ok(Vec::new()) } + async fn list_all_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list_all().await + } async fn get(&self, _id: &str) -> Result, DbError> { Ok(None) } + async fn get_for_user(&self, _user_id: &str, id: &str) -> Result, DbError> { + self.get(id).await + } async fn find_by_source_and_name( &self, _agent_source: &str, @@ -871,12 +897,34 @@ mod tests { ) -> Result, DbError> { Ok(None) } + async fn find_by_source_and_name_for_user( + &self, + _user_id: &str, + agent_source: &str, + name: &str, + ) -> Result, DbError> { + self.find_by_source_and_name(agent_source, name).await + } async fn find_builtin_by_backend(&self, _backend: &str) -> Result, DbError> { Ok(None) } + async fn find_builtin_by_backend_for_user( + &self, + _user_id: &str, + backend: &str, + ) -> Result, DbError> { + self.find_builtin_by_backend(backend).await + } async fn upsert(&self, _params: &UpsertAgentMetadataParams<'_>) -> Result { Err(DbError::Init("unused".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAgentMetadataParams<'_>, + ) -> Result { + self.upsert(params).await + } async fn apply_handshake( &self, _id: &str, @@ -884,6 +932,14 @@ mod tests { ) -> Result, DbError> { Ok(None) } + async fn apply_handshake_for_user( + &self, + _user_id: &str, + id: &str, + params: &UpdateAgentHandshakeParams<'_>, + ) -> Result, DbError> { + self.apply_handshake(id, params).await + } async fn update_availability_snapshot( &self, _id: &str, @@ -891,6 +947,14 @@ mod tests { ) -> Result, DbError> { Ok(None) } + async fn update_availability_snapshot_for_user( + &self, + _user_id: &str, + id: &str, + params: &UpdateAgentAvailabilitySnapshotParams<'_>, + ) -> Result, DbError> { + self.update_availability_snapshot(id, params).await + } async fn update_agent_overrides( &self, _id: &str, @@ -899,31 +963,51 @@ mod tests { ) -> Result<(), DbError> { Ok(()) } + async fn update_agent_overrides_for_user( + &self, + _user_id: &str, + id: &str, + command_override: Option<&str>, + env_override: Option<&str>, + ) -> Result<(), DbError> { + self.update_agent_overrides(id, command_override, env_override).await + } async fn set_enabled(&self, _id: &str, _enabled: bool) -> Result { Ok(false) } + async fn set_enabled_for_user(&self, _user_id: &str, id: &str, enabled: bool) -> Result { + self.set_enabled(id, enabled).await + } async fn delete(&self, _id: &str) -> Result { Ok(false) } + async fn delete_for_user(&self, _user_id: &str, id: &str) -> Result { + self.delete(id).await + } } struct EmptyProviderRepo; #[async_trait] impl IProviderRepository for EmptyProviderRepo { - async fn list(&self) -> Result, DbError> { + async fn list(&self, _user_id: &str) -> Result, DbError> { Ok(Vec::new()) } - async fn find_by_id(&self, _id: &str) -> Result, DbError> { + async fn find_by_id(&self, _user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } async fn create(&self, _params: CreateProviderParams<'_>) -> Result { Err(DbError::Init("unused".into())) } - async fn update(&self, _id: &str, _params: UpdateProviderParams<'_>) -> Result { + async fn update( + &self, + _user_id: &str, + _id: &str, + _params: UpdateProviderParams<'_>, + ) -> Result { Err(DbError::Init("unused".into())) } - async fn delete(&self, _id: &str) -> Result<(), DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), DbError> { Ok(()) } } @@ -976,7 +1060,7 @@ mod tests { let mut agent = test_agent(); agent.backend = "aionrs".into(); - let transport = provisioner.team_tool_transport(&agent).await.unwrap(); + let transport = provisioner.team_tool_transport("user-test", &agent).await.unwrap(); assert_eq!(transport, TeamToolTransport::Mcp); } @@ -987,7 +1071,7 @@ mod tests { let mut agent = test_agent(); agent.backend = "custom-acp".into(); - let transport = provisioner.team_tool_transport(&agent).await.unwrap(); + let transport = provisioner.team_tool_transport("user-test", &agent).await.unwrap(); assert_eq!(transport, TeamToolTransport::CliAssumed); } @@ -998,7 +1082,10 @@ mod tests { let patches = Arc::new(Mutex::new(Vec::new())); let provisioner = test_provisioner_with_patches(events, Arc::clone(&patches)); - provisioner.write_team_cli_runtime_config(&test_agent()).await.unwrap(); + provisioner + .write_team_cli_runtime_config("user-test", &test_agent()) + .await + .unwrap(); let patches = patches.lock().unwrap(); assert_eq!(patches[0]["team_mcp_stdio_config"], serde_json::Value::Null); diff --git a/crates/aionui-team/src/scheduler/mod.rs b/crates/aionui-team/src/scheduler/mod.rs index 4d50fc89d..e31b39f81 100644 --- a/crates/aionui-team/src/scheduler/mod.rs +++ b/crates/aionui-team/src/scheduler/mod.rs @@ -130,6 +130,7 @@ pub struct TeammateManager { impl TeammateManager { pub fn new( team_id: String, + user_id: String, agents: &[TeamAgent], mailbox: Arc, task_board: Arc, @@ -148,7 +149,7 @@ impl TeammateManager { }, ); } - let events = TeamEventEmitter::new(team_id.clone(), broadcaster); + let events = TeamEventEmitter::new(team_id.clone(), user_id, broadcaster); Self { team_id, slots: Mutex::new(slots), diff --git a/crates/aionui-team/src/scheduler/tests.rs b/crates/aionui-team/src/scheduler/tests.rs index 66f6f0346..50d9710f6 100644 --- a/crates/aionui-team/src/scheduler/tests.rs +++ b/crates/aionui-team/src/scheduler/tests.rs @@ -94,7 +94,14 @@ fn make_manager(agents: &[TeamAgent]) -> (TeammateManager, Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); let action = SchedulerAction::SendMessage { to: "worker-1".into(), @@ -486,7 +500,14 @@ async fn execute_broadcast_message_writes_to_all_others() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); let action = SchedulerAction::SendMessage { to: "*".into(), @@ -512,7 +533,14 @@ async fn execute_task_create() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox, task_board.clone(), broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox, + task_board.clone(), + broadcaster, + ); let action = SchedulerAction::TaskCreate { subject: "Implement feature".into(), @@ -537,7 +565,14 @@ async fn execute_task_update() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox, task_board.clone(), broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox, + task_board.clone(), + broadcaster, + ); let task = task_board.create_task("t1", "Work", None, None, &[]).await.unwrap(); @@ -563,7 +598,14 @@ async fn execute_idle_notification_writes_to_lead_mailbox() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.set_status("worker-1", TeammateStatus::Working).await.unwrap(); @@ -587,7 +629,14 @@ async fn lead_idle_notification_does_not_write_to_self() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.set_status("lead-1", TeammateStatus::Working).await.unwrap(); @@ -609,7 +658,14 @@ async fn execute_shutdown_agent_writes_shutdown_request() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); let action = SchedulerAction::ShutdownAgent { slot_id: "worker-1".into(), @@ -643,7 +699,14 @@ async fn lead_cannot_shutdown_lead() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); let action = SchedulerAction::ShutdownAgent { slot_id: "lead-1".into(), @@ -669,7 +732,14 @@ async fn lead_can_shutdown_worker() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); let action = SchedulerAction::ShutdownAgent { slot_id: "worker-1".into(), @@ -715,7 +785,14 @@ async fn finalize_turn_executes_actions_and_marks_idle() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board.clone(), broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board.clone(), + broadcaster, + ); mgr.set_status("worker-1", TeammateStatus::Working).await.unwrap(); mgr.set_status("worker-2", TeammateStatus::Working).await.unwrap(); @@ -767,7 +844,14 @@ async fn finalize_turn_with_idle_notification_skips_double_idle() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox, task_board, broadcaster.clone()); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox, + task_board, + broadcaster.clone(), + ); mgr.set_status("worker-1", TeammateStatus::Working).await.unwrap(); @@ -794,7 +878,7 @@ async fn finalize_turn_all_teammates_done_signals_leader_wake() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox, task_board, broadcaster); + let mgr = TeammateManager::new("t1".into(), "user-1".into(), &agents, mailbox, task_board, broadcaster); mgr.set_status("worker-1", TeammateStatus::Working).await.unwrap(); mgr.set_status("worker-2", TeammateStatus::Working).await.unwrap(); @@ -814,7 +898,14 @@ async fn wake_payload_includes_tasks_and_unread() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board.clone(), broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board.clone(), + broadcaster, + ); task_board.create_task("t1", "Task A", None, None, &[]).await.unwrap(); @@ -845,7 +936,14 @@ async fn mark_idle_with_summary_writes_idle_notification_to_lead() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.set_status("worker-1", TeammateStatus::Working).await.unwrap(); mgr.mark_idle("worker-1", Some("sub-task done")).await.unwrap(); @@ -865,7 +963,14 @@ async fn mark_idle_without_summary_still_writes_fallback_content() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.set_status("worker-1", TeammateStatus::Working).await.unwrap(); mgr.mark_idle("worker-1", None).await.unwrap(); @@ -883,7 +988,14 @@ async fn mark_idle_from_lead_does_not_write_notification() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.set_status("lead-1", TeammateStatus::Working).await.unwrap(); mgr.mark_idle("lead-1", Some("done")).await.unwrap(); @@ -1304,7 +1416,14 @@ async fn write_crash_testament_delivers_to_lead_mailbox() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.write_crash_testament("worker-1", "Worker1", &CrashReason::ProcessExited, None) .await @@ -1330,7 +1449,14 @@ async fn write_crash_testament_noop_when_no_lead() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); // Should not panic or error mgr.write_crash_testament("worker-1", "Worker1", &CrashReason::SessionNotFound, Some("last words")) @@ -1353,7 +1479,14 @@ async fn write_crash_testament_noop_when_lead_crashes() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); // Lead crashing should not write to itself mgr.write_crash_testament("lead-1", "Lead", &CrashReason::ProcessExited, None) @@ -1410,7 +1543,14 @@ async fn handle_agent_crash_writes_testament_to_lead() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.handle_agent_crash("worker-1", CrashReason::Unknown("segfault".into()), Some("cleaning up")) .await @@ -1454,7 +1594,14 @@ async fn handle_agent_crash_leader_branch_returns_none() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.set_status("lead-1", TeammateStatus::Working).await.unwrap(); assert!(mgr.acquire_wake_lock("lead-1")); @@ -1564,7 +1711,14 @@ async fn handle_inactivity_timeout_teammate_marks_error_and_wakes_lead() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.set_status("worker-1", TeammateStatus::Working).await.unwrap(); @@ -1592,7 +1746,14 @@ async fn handle_inactivity_timeout_leader_returns_none_no_mailbox_write() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.set_status("lead-1", TeammateStatus::Working).await.unwrap(); assert!(mgr.acquire_wake_lock("lead-1")); @@ -1669,7 +1830,14 @@ async fn handle_inactivity_timeout_no_lead_returns_none() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); let wake_target = mgr.handle_inactivity_timeout("worker-1").await.unwrap(); @@ -1689,7 +1857,14 @@ async fn notify_shutdown_rejected_delivers_to_lead_mailbox() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.notify_shutdown_rejected("worker-1", "still working on task X") .await @@ -1716,7 +1891,14 @@ async fn notify_shutdown_rejected_noop_when_no_lead() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.notify_shutdown_rejected("worker-1", "busy").await.unwrap(); @@ -1733,7 +1915,14 @@ async fn notify_shutdown_rejected_noop_when_sender_is_lead() { let mailbox = Arc::new(Mailbox::new(repo.clone())); let task_board = Arc::new(TaskBoard::new(repo)); let broadcaster: Arc = Arc::new(RecordingBroadcaster::new()); - let mgr = TeammateManager::new("t1".into(), &agents, mailbox.clone(), task_board, broadcaster); + let mgr = TeammateManager::new( + "t1".into(), + "user-1".into(), + &agents, + mailbox.clone(), + task_board, + broadcaster, + ); mgr.notify_shutdown_rejected("lead-1", "irrelevant").await.unwrap(); diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 137b1174b..2d8465335 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -201,17 +201,24 @@ impl TeamSessionService { } async fn load_owned_team(&self, user_id: &str, team_id: &str) -> Result { + let row = self.load_owned_team_row(user_id, team_id).await?; + Ok(Team::from_row(&row)?) + } + + async fn load_owned_team_row(&self, user_id: &str, team_id: &str) -> Result { + self.repo + .get_team(user_id, team_id) + .await? + .ok_or_else(|| TeamError::TeamNotFound(team_id.into())) + } + + pub(crate) async fn team_owner_user_id(&self, team_id: &str) -> Result { let row = self .repo - .get_team(team_id) + .get_team_for_restore(team_id) .await? .ok_or_else(|| TeamError::TeamNotFound(team_id.into()))?; - if row.user_id != user_id { - return Err(TeamError::Forbidden(format!( - "team {team_id} is not owned by current user" - ))); - } - Ok(Team::from_row(&row)?) + Ok(row.user_id) } pub async fn renew_active_lease( @@ -220,9 +227,14 @@ impl TeamSessionService { team_id: &str, active_leases: &ActiveLeaseRegistry, ) -> Result<(), TeamError> { - let team = match self.load_owned_team(user_id, team_id).await { + let team = match self.repo.get_team(user_id, team_id).await { + Ok(Some(row)) => Team::from_row(&row).map_err(TeamError::from), + Ok(None) => Err(TeamError::TeamNotFound(team_id.to_owned())), + Err(error) => Err(TeamError::Database(error)), + }; + let team = match team { Ok(team) => team, - Err(error @ (TeamError::TeamNotFound(_) | TeamError::Forbidden(_))) => { + Err(error @ TeamError::TeamNotFound(_)) => { debug!( kind = "team", team_id, @@ -261,7 +273,7 @@ impl TeamSessionService { /// Restore sessions for all existing teams. Called once at app startup /// so that MCP servers are available before any user sends a message. pub async fn restore_all_sessions(&self) { - let teams = match self.repo.list_teams().await { + let teams = match self.repo.list_teams_for_restore().await { Ok(t) => t, Err(e) => { tracing::warn!(error = %e, "failed to list teams for session restore"); @@ -269,7 +281,7 @@ impl TeamSessionService { } }; for team in &teams { - if let Err(e) = self.ensure_session_inner(&team.id).await { + if let Err(e) = self.ensure_session_inner(&team.id, None).await { tracing::warn!(team_id = %team.id, error = %e, "failed to restore session on startup"); continue; } @@ -347,9 +359,9 @@ impl TeamSessionService { "Team created" ); - self.broadcast_team_created(&team.id, &team.name); + self.broadcast_team_created(user_id, &team.id, &team.name); - self.build_team_response(&team).await + self.build_team_response(user_id, &team).await } pub async fn list_teams(&self, user_id: &str) -> Result, TeamError> { @@ -357,7 +369,7 @@ impl TeamSessionService { let mut teams = Vec::with_capacity(rows.len()); for row in &rows { match Team::from_row(row) { - Ok(team) => match self.build_team_response(&team).await { + Ok(team) => match self.build_team_response(user_id, &team).await { Ok(resp) => teams.push(resp), Err(e) => { tracing::warn!(team_id = %row.id, error = %e, "skipping team with build error"); @@ -373,7 +385,7 @@ impl TeamSessionService { pub async fn get_team(&self, user_id: &str, team_id: &str) -> Result { let team = self.load_owned_team(user_id, team_id).await?; - self.build_team_response(&team).await + self.build_team_response(user_id, &team).await } pub async fn remove_team(&self, user_id: &str, team_id: &str) -> Result<(), TeamError> { @@ -403,14 +415,14 @@ impl TeamSessionService { .await; } - self.repo.delete_mailbox_by_team(team_id).await?; - self.repo.delete_tasks_by_team(team_id).await?; - self.repo.delete_team(team_id).await?; + self.repo.delete_mailbox_by_team(user_id, team_id).await?; + self.repo.delete_tasks_by_team(user_id, team_id).await?; + self.repo.delete_team(user_id, team_id).await?; self.add_agent_locks.remove(team_id); info!(team_id = %team_id, "Team removed"); - self.broadcast_team_removed(team_id); + self.broadcast_team_removed(user_id, team_id); Ok(()) } @@ -419,6 +431,7 @@ impl TeamSessionService { self.repo .update_team( + user_id, team_id, &UpdateTeamParams { name: Some(name.to_owned()), @@ -426,7 +439,7 @@ impl TeamSessionService { }, ) .await?; - self.broadcast_team_renamed(team_id, name); + self.broadcast_team_renamed(user_id, team_id, name); Ok(()) } @@ -443,16 +456,7 @@ impl TeamSessionService { .clone(); let _guard = lock.lock().await; - let row = self - .repo - .get_team(team_id) - .await? - .ok_or_else(|| TeamError::TeamNotFound(team_id.into()))?; - if row.user_id != user_id { - return Err(TeamError::Forbidden(format!( - "team {team_id} is not owned by current user" - ))); - } + let row = self.load_owned_team_row(user_id, team_id).await?; let mut team = Team::from_row(&row)?; let agent = self.provisioner().add_agent(user_id, &row, &mut team, req).await?; @@ -463,7 +467,7 @@ impl TeamSessionService { .self_ref .upgrade() .ok_or_else(|| TeamError::InvalidRequest("add_agent requires a live TeamSessionService".into()))?; - self.broadcast_agent_runtime_status(team_id, &agent, TeamAgentRuntimeStatus::Pending, None); + self.broadcast_agent_runtime_status(user_id, team_id, &agent, TeamAgentRuntimeStatus::Pending, None); spawn_attach_agent_process_bg( service, session, @@ -484,7 +488,8 @@ impl TeamSessionService { "manual teammate added" ); } else { - TeamEventEmitter::new(team_id.to_owned(), self.broadcaster.clone()).broadcast_agent_spawned(&agent); + TeamEventEmitter::new(team_id.to_owned(), user_id.to_owned(), self.broadcaster.clone()) + .broadcast_agent_spawned(&agent); info!( team_id = %team_id, slot_id = %agent.slot_id, @@ -496,7 +501,7 @@ impl TeamSessionService { ); } - self.build_agent_response(&agent).await + self.build_agent_response(user_id, &agent).await } pub async fn remove_agent(&self, user_id: &str, team_id: &str, slot_id: &str) -> Result<(), TeamError> { @@ -560,6 +565,7 @@ impl TeamSessionService { let agents_json = serde_json::to_string(¤t.agents)?; self.repo .update_team( + user_id, team_id, &UpdateTeamParams { agents: Some(agents_json), @@ -642,7 +648,8 @@ impl TeamSessionService { "manual teammate removed" ); } else { - TeamEventEmitter::new(team_id.to_owned(), self.broadcaster.clone()).broadcast_agent_removed(slot_id); + TeamEventEmitter::new(team_id.to_owned(), user_id.to_owned(), self.broadcaster.clone()) + .broadcast_agent_removed(slot_id); info!( team_id = %team_id, slot_id = %removed.slot_id, @@ -693,6 +700,7 @@ impl TeamSessionService { let agents_json = serde_json::to_string(&team.agents)?; self.repo .update_team( + user_id, team_id, &UpdateTeamParams { agents: Some(agents_json), @@ -723,19 +731,11 @@ impl TeamSessionService { /// any failure, stop the session and leave the map untouched so a /// retry can start cleanly. pub async fn ensure_session(&self, user_id: &str, team_id: &str) -> Result<(), TeamError> { - let row = match self.repo.get_team(team_id).await { - Ok(Some(row)) => row, - Ok(None) | Err(_) => return self.ensure_session_inner(team_id).await, - }; - if row.user_id != user_id { - return Err(TeamError::Forbidden(format!( - "team {team_id} is not owned by current user" - ))); - } - self.ensure_session_inner(team_id).await + self.load_owned_team_row(user_id, team_id).await?; + self.ensure_session_inner(team_id, Some(user_id)).await } - async fn ensure_session_inner(&self, team_id: &str) -> Result<(), TeamError> { + async fn ensure_session_inner(&self, team_id: &str, requested_user_id: Option<&str>) -> Result<(), TeamError> { let membership_lock = self .add_agent_locks .entry(team_id.to_owned()) @@ -743,28 +743,34 @@ impl TeamSessionService { .clone(); let membership_guard = membership_lock.lock().await; - let row = match self.repo.get_team(team_id).await { + let row = match self.repo.get_team_for_restore(team_id).await { Ok(Some(row)) => row, Ok(None) => { - self.broadcast_session_status( - team_id, - TeamSessionStatus::Failed, - Some(TeamSessionPhase::LoadingTeam), - |p| { - p.error = Some(format!("team not found: {team_id}")); - }, - ); + if let Some(user_id) = requested_user_id { + self.broadcast_session_status( + user_id, + team_id, + TeamSessionStatus::Failed, + Some(TeamSessionPhase::LoadingTeam), + |p| { + p.error = Some(format!("team not found: {team_id}")); + }, + ); + } return Err(TeamError::TeamNotFound(team_id.into())); } Err(e) => { - self.broadcast_session_status( - team_id, - TeamSessionStatus::Failed, - Some(TeamSessionPhase::LoadingTeam), - |p| { - p.error = Some(e.to_string()); - }, - ); + if let Some(user_id) = requested_user_id { + self.broadcast_session_status( + user_id, + team_id, + TeamSessionStatus::Failed, + Some(TeamSessionPhase::LoadingTeam), + |p| { + p.error = Some(e.to_string()); + }, + ); + } return Err(e.into()); } }; @@ -801,6 +807,7 @@ impl TeamSessionService { } self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Starting, Some(TeamSessionPhase::LoadingTeam), @@ -808,6 +815,7 @@ impl TeamSessionService { ); self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Starting, Some(TeamSessionPhase::StartingBridge), @@ -832,6 +840,7 @@ impl TeamSessionService { Ok(session) => Arc::new(session), Err(e) => { self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Failed, Some(TeamSessionPhase::StartingBridge), @@ -844,6 +853,7 @@ impl TeamSessionService { }; self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Starting, Some(TeamSessionPhase::AttachingAgents), @@ -865,6 +875,7 @@ impl TeamSessionService { else { let error = TeamError::InvalidRequest("team has no lead agent".to_owned()); self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Failed, Some(TeamSessionPhase::AttachingAgents), @@ -891,7 +902,7 @@ impl TeamSessionService { drop(membership_guard); drop(ensure_guard); - self.broadcast_agent_runtime_status(team_id, &leader, TeamAgentRuntimeStatus::Pending, None); + self.broadcast_agent_runtime_status(&user_id, team_id, &leader, TeamAgentRuntimeStatus::Pending, None); let leader_outcome = match session.member_runtimes().reserve_attach(&leader.slot_id, false) { ReserveAttach::Start(lease) => { attach_member_runtime( @@ -916,6 +927,7 @@ impl TeamSessionService { AttachOutcome::Ready | AttachOutcome::Removed => {} AttachOutcome::Failed(failure) => { self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Failed, Some(TeamSessionPhase::AttachingAgents), @@ -942,10 +954,11 @@ impl TeamSessionService { // Teammates start dormant; the leader's Ready was already broadcast by // its successful attach. for agent in agents_snapshot.iter().filter(|a| a.role != TeammateRole::Lead) { - self.broadcast_agent_runtime_status(team_id, agent, TeamAgentRuntimeStatus::Dormant, None); + self.broadcast_agent_runtime_status(&user_id, team_id, agent, TeamAgentRuntimeStatus::Dormant, None); } self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Starting, Some(TeamSessionPhase::Recovering), @@ -960,7 +973,7 @@ impl TeamSessionService { ); } - self.broadcast_session_status(team_id, TeamSessionStatus::Ready, None, |p| { + self.broadcast_session_status(&user_id, team_id, TeamSessionStatus::Ready, None, |p| { p.server_count = Some(agents_snapshot.len()); }); @@ -1005,6 +1018,7 @@ impl TeamSessionService { match reservation { ReserveAttach::Start(owner) => { self.broadcast_agent_runtime_status( + session.user_id(), session.team_id(), agent, TeamAgentRuntimeStatus::Pending, @@ -1092,7 +1106,7 @@ impl TeamSessionService { .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) .clone(); let _membership_guard = membership_lock.lock().await; - let current_agents = match self.repo.get_team(team_id).await? { + let current_agents = match self.repo.get_team(user_id, team_id).await? { Some(row) => Team::from_row(&row)?.agents, None => return Err(TeamError::TeamNotFound(team_id.to_owned())), }; @@ -1125,6 +1139,7 @@ impl TeamSessionService { // must not fail just because an unrelated teammate is broken. if agent.role == TeammateRole::Lead { self.broadcast_session_status( + user_id, team_id, TeamSessionStatus::Failed, Some(TeamSessionPhase::AttachingAgents), @@ -1146,7 +1161,7 @@ impl TeamSessionService { } } - self.broadcast_session_status(team_id, TeamSessionStatus::Ready, None, |payload| { + self.broadcast_session_status(user_id, team_id, TeamSessionStatus::Ready, None, |payload| { payload.server_count = Some(current_agents.len()); }); Ok(()) @@ -1181,16 +1196,7 @@ impl TeamSessionService { team_id: &str, conversation_id: &str, ) -> Result { - let row = self - .repo - .get_team(team_id) - .await? - .ok_or_else(|| TeamError::TeamNotFound(team_id.to_owned()))?; - if row.user_id != user_id { - return Err(TeamError::Forbidden(format!( - "team {team_id} is not owned by current user" - ))); - } + let row = self.load_owned_team_row(user_id, team_id).await?; let team = Team::from_row(&row)?; let member = team.agents.iter().any(|agent| agent.conversation_id == conversation_id); @@ -1203,6 +1209,7 @@ impl TeamSessionService { fn broadcast_session_status( &self, + user_id: &str, team_id: &str, status: TeamSessionStatus, phase: Option, @@ -1230,10 +1237,11 @@ impl TeamSessionService { error = payload.error.as_deref().unwrap_or(""), "team session status broadcast" ); - let event = WebSocketMessage::new( - TEAM_SESSION_STATUS_CHANGED_EVENT, - serde_json::to_value(payload).expect("serialize team session status payload"), - ); + // Keep per-user scoping so the overlay event reaches only the owning + // user's WebSocket subscribers. + let mut value = serde_json::to_value(payload).expect("serialize team session status payload"); + value["user_id"] = serde_json::Value::String(user_id.to_owned()); + let event = WebSocketMessage::new(TEAM_SESSION_STATUS_CHANGED_EVENT, value); self.broadcaster.broadcast(event); } @@ -1249,49 +1257,50 @@ impl TeamSessionService { }) } - fn broadcast_team_created(&self, team_id: &str, team_name: &str) { + fn broadcast_team_created(&self, user_id: &str, team_id: &str, team_name: &str) { info!(team_id = %team_id, event_name = TEAM_CREATED_EVENT, "team event broadcast"); self.broadcaster.broadcast(WebSocketMessage::new( TEAM_CREATED_EVENT, - serde_json::json!({ "team_id": team_id, "team_name": team_name }), + serde_json::json!({ "user_id": user_id, "team_id": team_id, "team_name": team_name }), )); - self.broadcast_team_list_changed(team_id, "created"); + self.broadcast_team_list_changed(user_id, team_id, "created"); } - fn broadcast_team_removed(&self, team_id: &str) { + fn broadcast_team_removed(&self, user_id: &str, team_id: &str) { info!(team_id = %team_id, event_name = TEAM_REMOVED_EVENT, "team event broadcast"); self.broadcaster.broadcast(WebSocketMessage::new( TEAM_REMOVED_EVENT, - serde_json::json!({ "team_id": team_id }), + serde_json::json!({ "user_id": user_id, "team_id": team_id }), )); - self.broadcast_team_list_changed(team_id, "removed"); + self.broadcast_team_list_changed(user_id, team_id, "removed"); } - fn broadcast_team_renamed(&self, team_id: &str, team_name: &str) { + fn broadcast_team_renamed(&self, user_id: &str, team_id: &str, team_name: &str) { info!(team_id = %team_id, event_name = TEAM_RENAMED_EVENT, "team event broadcast"); self.broadcaster.broadcast(WebSocketMessage::new( TEAM_RENAMED_EVENT, - serde_json::json!({ "team_id": team_id, "team_name": team_name }), + serde_json::json!({ "user_id": user_id, "team_id": team_id, "team_name": team_name }), )); - self.broadcast_team_list_changed(team_id, "renamed"); + self.broadcast_team_list_changed(user_id, team_id, "renamed"); } - fn broadcast_team_list_changed(&self, team_id: &str, action: &str) { + fn broadcast_team_list_changed(&self, user_id: &str, team_id: &str, action: &str) { info!(team_id = %team_id, event_name = crate::events::TEAM_LIST_CHANGED_EVENT, action, "team event broadcast"); self.broadcaster.broadcast(WebSocketMessage::new( crate::events::TEAM_LIST_CHANGED_EVENT, - serde_json::json!({ "team_id": team_id, "action": action }), + serde_json::json!({ "user_id": user_id, "team_id": team_id, "action": action }), )); } pub(crate) fn broadcast_agent_runtime_status( &self, + user_id: &str, team_id: &str, agent: &TeamAgent, status: TeamAgentRuntimeStatus, error: Option, ) { - TeamEventEmitter::new(team_id.to_owned(), self.broadcaster.clone()) + TeamEventEmitter::new(team_id.to_owned(), user_id.to_owned(), self.broadcaster.clone()) .broadcast_agent_runtime_status(agent, status, error); } @@ -1356,7 +1365,13 @@ impl TeamSessionService { pub(crate) fn publish_member_runtime_ready_if_current(&self, expected: &TeamSession, agent: &TeamAgent) -> bool { self.with_published_session(expected, |_| { - self.broadcast_agent_runtime_status(expected.team_id(), agent, TeamAgentRuntimeStatus::Ready, None); + self.broadcast_agent_runtime_status( + expected.user_id(), + expected.team_id(), + agent, + TeamAgentRuntimeStatus::Ready, + None, + ); }) .is_some() } @@ -1364,6 +1379,7 @@ impl TeamSessionService { pub(crate) fn publish_member_runtime_starting_if_current(&self, expected: &TeamSession) -> bool { self.with_published_session(expected, |_| { self.broadcast_session_status( + expected.user_id(), expected.team_id(), TeamSessionStatus::Starting, Some(TeamSessionPhase::AttachingAgents), @@ -1376,6 +1392,7 @@ impl TeamSessionService { pub(crate) fn publish_member_runtime_failed_if_current(&self, expected: &TeamSession, reason: &str) -> bool { self.with_published_session(expected, |_| { self.broadcast_session_status( + expected.user_id(), expected.team_id(), TeamSessionStatus::Failed, Some(TeamSessionPhase::AttachingAgents), @@ -1392,7 +1409,7 @@ impl TeamSessionService { .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) .clone(); let _membership_guard = membership_lock.lock().await; - let Ok(Some(row)) = self.repo.get_team(expected.team_id()).await else { + let Ok(Some(row)) = self.repo.get_team(expected.user_id(), expected.team_id()).await else { return; }; let Ok(team) = Team::from_row(&row) else { @@ -1415,9 +1432,15 @@ impl TeamSessionService { match expected.member_runtimes().snapshot(&leader.slot_id) { MemberRuntimeSnapshot::Ready => { let _ = self.with_published_session(expected, |_| { - self.broadcast_session_status(expected.team_id(), TeamSessionStatus::Ready, None, |payload| { - payload.server_count = Some(team.agents.len()); - }); + self.broadcast_session_status( + expected.user_id(), + expected.team_id(), + TeamSessionStatus::Ready, + None, + |payload| { + payload.server_count = Some(team.agents.len()); + }, + ); }); } MemberRuntimeSnapshot::Failed { failure, .. } => { @@ -1513,16 +1536,10 @@ impl TeamSessionService { let team_row = self .repo - .get_team(&team_id) + .get_team(user_id, &team_id) .await .map_err(|error| error_payload(TeamToolErrorCode::RuntimeContextMissing, error.to_string()))? .ok_or_else(|| error_payload(TeamToolErrorCode::TeamNotFound, "team not found"))?; - if team_row.user_id != user_id { - return Err(error_payload( - TeamToolErrorCode::PermissionDenied, - "team does not belong to user", - )); - } let binding = TeamSessionBinding { team_id: team_id.clone(), @@ -1592,6 +1609,20 @@ impl TeamSessionService { Ok(()) } + pub fn stop_sessions_for_user(&self, user_id: &str) -> usize { + let team_ids: Vec = self + .sessions + .iter() + .filter(|entry| entry.session.user_id() == user_id) + .map(|entry| entry.key().clone()) + .collect(); + let stopped = team_ids.len(); + for team_id in team_ids { + self.stop_session_unchecked(&team_id); + } + stopped + } + fn stop_session_unchecked(&self, team_id: &str) { if let Some((_, entry)) = self.sessions.remove(team_id) { entry.slow_monitor_handle.abort(); @@ -1673,7 +1704,15 @@ impl TeamSessionService { "team idle cleanup stopping idle team session" ); info!(team_id, reason = "idle_cleanup", "broadcasting team session stopped"); - self.broadcast_session_status(&team_id, TeamSessionStatus::Stopped, None, |_| {}); + if let Some(entry) = self.sessions.get(&team_id) { + self.broadcast_session_status( + entry.session.user_id(), + &team_id, + TeamSessionStatus::Stopped, + None, + |_| {}, + ); + } self.stop_session_unchecked(&team_id); for agent in agents { self.task_manager @@ -1696,7 +1735,7 @@ impl TeamSessionService { files: Option>, ) -> Result { self.load_owned_team(user_id, team_id).await?; - self.ensure_session_inner(team_id).await?; + self.ensure_session_inner(team_id, Some(user_id)).await?; let session = { let entry = self .sessions @@ -1716,7 +1755,7 @@ impl TeamSessionService { files: Option>, ) -> Result { self.load_owned_team(user_id, team_id).await?; - self.ensure_session_inner(team_id).await?; + self.ensure_session_inner(team_id, Some(user_id)).await?; let session = { let entry = self .sessions @@ -1735,7 +1774,7 @@ impl TeamSessionService { /// the member's event loop via `reconcile_mailbox`. pub async fn attach_agent_runtime(&self, user_id: &str, team_id: &str, slot_id: &str) -> Result<(), TeamError> { self.load_owned_team(user_id, team_id).await?; - self.ensure_session_inner(team_id).await?; + self.ensure_session_inner(team_id, Some(user_id)).await?; let session = { let entry = self .sessions @@ -1749,7 +1788,7 @@ impl TeamSessionService { .upgrade() .ok_or_else(|| TeamError::InvalidRequest("team service is shutting down".to_owned()))?; let reservation = session.member_runtimes().reserve_attach(slot_id, true); - self.broadcast_agent_runtime_status(team_id, &agent, TeamAgentRuntimeStatus::Pending, None); + self.broadcast_agent_runtime_status(user_id, team_id, &agent, TeamAgentRuntimeStatus::Pending, None); spawn_attach_agent_process_bg( service, Arc::clone(&session), @@ -1772,7 +1811,7 @@ impl TeamSessionService { reason: Option, ) -> Result<(), TeamError> { self.load_owned_team(user_id, team_id).await?; - self.ensure_session_inner(team_id).await?; + self.ensure_session_inner(team_id, Some(user_id)).await?; let session = { let entry = self .sessions @@ -1792,7 +1831,7 @@ impl TeamSessionService { reason: Option, ) -> Result<(), TeamError> { self.load_owned_team(user_id, team_id).await?; - self.ensure_session_inner(team_id).await?; + self.ensure_session_inner(team_id, Some(user_id)).await?; let session = { let entry = self .sessions @@ -1812,7 +1851,7 @@ impl TeamSessionService { reason: Option, ) -> Result<(), TeamError> { self.load_owned_team(user_id, team_id).await?; - self.ensure_session_inner(team_id).await?; + self.ensure_session_inner(team_id, Some(user_id)).await?; let session = { let entry = self .sessions @@ -1828,6 +1867,7 @@ impl TeamSessionService { let provisioner = self.provisioner(); self.repo .update_team( + user_id, team_id, &UpdateTeamParams { session_mode: Some(mode.to_owned()), @@ -2282,6 +2322,29 @@ mod tests { svc.stop_session("user-test", &created.id).await.unwrap(); } + #[tokio::test] + async fn stop_sessions_for_user_keeps_other_user_sessions() { + let (svc, _repo, _task_manager, _conv_repo) = setup_with_factory_metadata_team_repo_and_conversation_repo(); + let owned = svc + .create_team("user-test", single_agent_team_request("Owned Session")) + .await + .unwrap(); + let other = svc + .create_team("user-other", single_agent_team_request("Other Session")) + .await + .unwrap(); + + svc.ensure_session("user-test", &owned.id).await.unwrap(); + svc.ensure_session("user-other", &other.id).await.unwrap(); + + assert_eq!(svc.stop_sessions_for_user("user-test"), 1); + assert_eq!(svc.session_count_for_test(), 1); + assert!(!svc.session_has_slow_monitor(&owned.id)); + assert!(svc.session_has_slow_monitor(&other.id)); + + svc.stop_session("user-other", &other.id).await.unwrap(); + } + #[tokio::test] async fn ensure_session_emits_agent_runtime_ready_after_member_warmup() { let (svc, _repo, _task_manager, _conv_repo, broadcaster) = @@ -2737,7 +2800,11 @@ mod tests { .await .unwrap(); - let row = repo.get_team(&created.id).await.unwrap().expect("team row"); + let row = repo + .get_team("user-test", &created.id) + .await + .unwrap() + .expect("team row"); assert_eq!(row.session_mode.as_deref(), Some("full_auto")); let added = svc @@ -2787,7 +2854,11 @@ mod tests { .create_team("user-test", single_agent_team_request("Partial Mode Seed")) .await .unwrap(); - let mut row = repo.get_team(&created.id).await.unwrap().expect("team row"); + let mut row = repo + .get_team("user-test", &created.id) + .await + .unwrap() + .expect("team row"); row.agents = serde_json::json!([ { "slot_id": "slot-accepts", @@ -2810,6 +2881,7 @@ mod tests { ]) .to_string(); repo.update_team( + "user-test", &created.id, &aionui_db::UpdateTeamParams { agents: Some(row.agents), @@ -2867,7 +2939,11 @@ mod tests { .await .unwrap(); - let team = repo.get_team(&created.id).await.unwrap().expect("team row"); + let team = repo + .get_team("user-test", &created.id) + .await + .unwrap() + .expect("team row"); assert_eq!(team.session_mode.as_deref(), Some("read-only")); let accepting_extra = conv_repo.get_extra(accepting_conversation_id).unwrap(); @@ -3017,6 +3093,6 @@ mod tests { .await .expect_err("team config options must reject cross-user access"); - assert!(matches!(err, crate::error::TeamError::Forbidden(_))); + assert!(matches!(err, crate::error::TeamError::TeamNotFound(_))); } } diff --git a/crates/aionui-team/src/service/describe_support.rs b/crates/aionui-team/src/service/describe_support.rs index 00b12c229..ebdf9b095 100644 --- a/crates/aionui-team/src/service/describe_support.rs +++ b/crates/aionui-team/src/service/describe_support.rs @@ -9,21 +9,25 @@ use aionui_db::models::AssistantDefinitionRow; impl TeamSessionService { pub(crate) async fn describe_assistant( &self, + user_id: &str, assistant_id: &str, locale: Option<&str>, ) -> Result { let definition = self .assistant_definition_repo - .get_by_assistant_id(assistant_id) + .get_by_assistant_id_for_user(user_id, assistant_id) .await? .ok_or_else(|| TeamError::InvalidRequest(format!("Preset assistant not found: {assistant_id}")))?; - let overlay = self.assistant_overlay_repo.get(&definition.id).await?; + let overlay = self + .assistant_overlay_repo + .get_for_user(user_id, &definition.id) + .await?; let effective_agent_id = overlay .as_ref() .and_then(|row| row.agent_id_override.as_deref()) .filter(|value| !value.trim().is_empty()) .unwrap_or(definition.agent_id.as_str()); - let effective_backend = resolve_runtime_backend(&self.agent_metadata_repo, effective_agent_id).await?; + let effective_backend = resolve_runtime_backend(&self.agent_metadata_repo, user_id, effective_agent_id).await?; Ok(render_assistant_description_json( &definition, diff --git a/crates/aionui-team/src/service/response_builder.rs b/crates/aionui-team/src/service/response_builder.rs index 1cc387553..93c3a9521 100644 --- a/crates/aionui-team/src/service/response_builder.rs +++ b/crates/aionui-team/src/service/response_builder.rs @@ -1,10 +1,10 @@ use super::*; impl TeamSessionService { - pub(super) async fn build_team_response(&self, team: &Team) -> Result { + pub(super) async fn build_team_response(&self, user_id: &str, team: &Team) -> Result { let mut agents = Vec::with_capacity(team.agents.len()); for agent in &team.agents { - agents.push(self.build_agent_response(agent).await?); + agents.push(self.build_agent_response(user_id, agent).await?); } Ok(TeamResponse { @@ -20,9 +20,10 @@ impl TeamSessionService { pub(super) async fn build_agent_response( &self, + user_id: &str, agent: &TeamAgent, ) -> Result { - let icon = self.resolve_agent_icon(agent).await?; + let icon = self.resolve_agent_icon(user_id, agent).await?; let mut response = agent.to_response_with_icon(icon); response.pending_confirmations = self.pending_confirmation_count(&agent.conversation_id); Ok(response) @@ -35,9 +36,12 @@ impl TeamSessionService { .unwrap_or(0) } - async fn resolve_agent_icon(&self, agent: &TeamAgent) -> Result, TeamError> { + async fn resolve_agent_icon(&self, user_id: &str, agent: &TeamAgent) -> Result, TeamError> { if let Some(assistant_id) = agent.assistant_id.as_deref() - && let Some(definition) = self.assistant_definition_repo.get_by_assistant_id(assistant_id).await? + && let Some(definition) = self + .assistant_definition_repo + .get_by_assistant_id_for_user(user_id, assistant_id) + .await? && let Some(icon) = assistant_icon( definition.assistant_id.as_str(), &definition.avatar_type, @@ -49,7 +53,7 @@ impl TeamSessionService { if let Some(row) = self .agent_metadata_repo - .find_builtin_by_backend(agent.backend.as_str()) + .find_builtin_by_backend_for_user(user_id, agent.backend.as_str()) .await? && row.icon.is_some() { @@ -59,7 +63,7 @@ impl TeamSessionService { if agent.backend == "acp" && let Some(row) = self .agent_metadata_repo - .find_builtin_by_backend(agent.model.as_str()) + .find_builtin_by_backend_for_user(user_id, agent.model.as_str()) .await? { return Ok(row.icon); diff --git a/crates/aionui-team/src/service/spawn_support.rs b/crates/aionui-team/src/service/spawn_support.rs index 987d8fa3c..6ac91ae47 100644 --- a/crates/aionui-team/src/service/spawn_support.rs +++ b/crates/aionui-team/src/service/spawn_support.rs @@ -30,9 +30,10 @@ fn find_acp_backend_metadata(rows: &[AgentMetadataRow], backend: &str) -> Option pub(crate) async fn acp_backend_metadata( agent_metadata_repo: &Arc, + user_id: &str, backend: &str, ) -> Result, TeamError> { - let rows = agent_metadata_repo.list_all().await?; + let rows = agent_metadata_repo.list_all_for_user(user_id).await?; Ok(find_acp_backend_metadata(&rows, backend)) } @@ -51,9 +52,10 @@ pub(crate) fn session_mode_for_backend( pub(crate) async fn resolve_runtime_backend( agent_metadata_repo: &Arc, + user_id: &str, agent_id: &str, ) -> Result { - let rows = agent_metadata_repo.list_all().await?; + let rows = agent_metadata_repo.list_all_for_user(user_id).await?; Ok(resolve_agent_binding_from_rows(&rows, agent_id) .map(|binding| binding.runtime_backend) .unwrap_or_else(|| agent_id.to_owned())) @@ -62,10 +64,11 @@ pub(crate) async fn resolve_runtime_backend( impl TeamSessionService { pub(crate) async fn resolve_team_selectable_assistant( &self, + user_id: &str, assistant_id: &str, ) -> Result { self.assistant_catalog - .resolve_team_selectable_assistant(assistant_id) + .resolve_team_selectable_assistant(user_id, assistant_id) .await? .ok_or_else(|| { TeamError::InvalidRequest(format!("Assistant is not available for team mode: {assistant_id}")) @@ -74,16 +77,17 @@ impl TeamSessionService { pub(crate) async fn resolve_spawn_backend_and_model( &self, + user_id: &str, assistant_id: Option<&str>, requested_model: Option<&str>, fallback_backend: &str, fallback_model: &str, ) -> Result<(String, String), TeamError> { if let Some(assistant_id) = assistant_id.map(str::trim).filter(|value| !value.is_empty()) { - let selectable = self.resolve_team_selectable_assistant(assistant_id).await?; + let selectable = self.resolve_team_selectable_assistant(user_id, assistant_id).await?; let definition = self .assistant_definition_repo - .get_by_assistant_id(assistant_id) + .get_by_assistant_id_for_user(user_id, assistant_id) .await? .ok_or_else(|| TeamError::InvalidRequest(format!("Preset assistant not found: {assistant_id}")))?; let backend = selectable.backend; @@ -96,7 +100,7 @@ impl TeamSessionService { .flatten() .map(|value| value.trim().to_owned()) .filter(|value| !value.is_empty()); - let backend_default_model = self.default_model_for_backend(&backend).await; + let backend_default_model = self.default_model_for_backend(user_id, &backend).await; let model = requested_model .or(fixed_model) .or(backend_default_model) @@ -109,7 +113,7 @@ impl TeamSessionService { .map(str::trim) .filter(|value| !value.is_empty()) .map(str::to_owned); - let backend_default_model = self.default_model_for_backend(&backend).await; + let backend_default_model = self.default_model_for_backend(user_id, &backend).await; let model = requested_model .or(backend_default_model) .unwrap_or_else(|| fallback_model.to_owned()); @@ -119,8 +123,8 @@ impl TeamSessionService { /// Return all enabled assistants that can currently participate in team mode. /// This consumes the same assistant projection as the Team creation UI, so /// `team_selectable` has a single source of truth. - pub(crate) async fn list_team_selectable_assistants(&self) -> Vec { - let Ok(assistants) = self.assistant_catalog.list_team_selectable_assistants().await else { + pub(crate) async fn list_team_selectable_assistants(&self, user_id: &str) -> Vec { + let Ok(assistants) = self.assistant_catalog.list_team_selectable_assistants(user_id).await else { return Vec::new(); }; @@ -138,8 +142,8 @@ impl TeamSessionService { /// Collect all enabled provider model IDs grouped by provider name. /// Returns a flat list of model IDs for use by internal agents (aionrs). - async fn collect_provider_models(&self) -> Vec { - let Ok(providers) = self.provider_repo.list().await else { + async fn collect_provider_models(&self, user_id: &str) -> Vec { + let Ok(providers) = self.provider_repo.list(user_id).await else { return vec![]; }; providers @@ -149,11 +153,15 @@ impl TeamSessionService { .collect() } - pub(crate) async fn default_model_for_backend(&self, backend: &str) -> Option { + pub(crate) async fn default_model_for_backend(&self, user_id: &str, backend: &str) -> Option { if backend == "aionrs" { - return self.collect_provider_models().await.into_iter().next(); + return self.collect_provider_models(user_id).await.into_iter().next(); } - let row = self.agent_metadata_repo.find_builtin_by_backend(backend).await.ok()??; + let row = self + .agent_metadata_repo + .find_builtin_by_backend_for_user(user_id, backend) + .await + .ok()??; let json: serde_json::Value = serde_json::from_str(row.available_models.as_deref()?).ok()?; if let Some(id) = json.get("current_model_id").and_then(|v| v.as_str()) && !id.is_empty() @@ -228,6 +236,7 @@ mod tests { fn agent_metadata_row(backend: &str, yolo_id: Option<&str>) -> AgentMetadataRow { AgentMetadataRow { id: format!("agent-{backend}"), + user_id: None, icon: None, name: format!("{backend} agent"), name_i18n: None, @@ -278,14 +287,49 @@ mod tests { Ok(vec![self.row.clone()]) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + + async fn list_including_deleted_for_user( + &self, + _user_id: &str, + ) -> Result, DbError> { + self.list().await + } + async fn get_by_assistant_id(&self, assistant_id: &str) -> Result, DbError> { Ok((self.row.assistant_id == assistant_id).then_some(self.row.clone())) } + async fn get_by_assistant_id_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + + async fn get_by_assistant_id_including_deleted_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + async fn get_by_id(&self, definition_id: &str) -> Result, DbError> { Ok((self.row.id == definition_id).then_some(self.row.clone())) } + async fn get_by_id_for_user( + &self, + _user_id: &str, + definition_id: &str, + ) -> Result, DbError> { + self.get_by_id(definition_id).await + } + async fn get_by_source_ref( &self, _source: &str, @@ -294,6 +338,24 @@ mod tests { Ok(None) } + async fn get_by_source_ref_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + + async fn get_by_source_ref_including_deleted_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + async fn upsert( &self, _params: &UpsertAssistantDefinitionParams<'_>, @@ -301,9 +363,26 @@ mod tests { Err(DbError::Init("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn soft_delete(&self, _definition_id: &str, _deleted_at: i64) -> Result { Ok(false) } + + async fn soft_delete_for_user( + &self, + _user_id: &str, + definition_id: &str, + deleted_at: i64, + ) -> Result { + self.soft_delete(definition_id, deleted_at).await + } } #[derive(Clone)] @@ -317,14 +396,49 @@ mod tests { Ok(self.rows.clone()) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + + async fn list_including_deleted_for_user( + &self, + _user_id: &str, + ) -> Result, DbError> { + self.list().await + } + async fn get_by_assistant_id(&self, assistant_id: &str) -> Result, DbError> { Ok(self.rows.iter().find(|row| row.assistant_id == assistant_id).cloned()) } + async fn get_by_assistant_id_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + + async fn get_by_assistant_id_including_deleted_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + async fn get_by_id(&self, definition_id: &str) -> Result, DbError> { Ok(self.rows.iter().find(|row| row.id == definition_id).cloned()) } + async fn get_by_id_for_user( + &self, + _user_id: &str, + definition_id: &str, + ) -> Result, DbError> { + self.get_by_id(definition_id).await + } + async fn get_by_source_ref( &self, _source: &str, @@ -333,6 +447,24 @@ mod tests { Ok(None) } + async fn get_by_source_ref_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + + async fn get_by_source_ref_including_deleted_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + async fn upsert( &self, _params: &UpsertAssistantDefinitionParams<'_>, @@ -340,9 +472,26 @@ mod tests { Err(DbError::Init("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn soft_delete(&self, _definition_id: &str, _deleted_at: i64) -> Result { Ok(false) } + + async fn soft_delete_for_user( + &self, + _user_id: &str, + definition_id: &str, + deleted_at: i64, + ) -> Result { + self.soft_delete(definition_id, deleted_at).await + } } #[derive(Clone)] @@ -356,17 +505,41 @@ mod tests { Ok((self.row.assistant_definition_id == definition_id).then_some(self.row.clone())) } + async fn get_for_user( + &self, + _user_id: &str, + definition_id: &str, + ) -> Result, DbError> { + self.get(definition_id).await + } + async fn list(&self) -> Result, DbError> { Ok(vec![self.row.clone()]) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + async fn upsert(&self, _params: &UpsertAssistantOverlayParams<'_>) -> Result { Err(DbError::Init("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantOverlayParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn delete(&self, _definition_id: &str) -> Result { Ok(false) } + + async fn delete_for_user(&self, _user_id: &str, definition_id: &str) -> Result { + self.delete(definition_id).await + } } #[derive(Clone)] @@ -384,17 +557,41 @@ mod tests { .cloned()) } + async fn get_for_user( + &self, + _user_id: &str, + definition_id: &str, + ) -> Result, DbError> { + self.get(definition_id).await + } + async fn list(&self) -> Result, DbError> { Ok(self.rows.clone()) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + async fn upsert(&self, _params: &UpsertAssistantOverlayParams<'_>) -> Result { Err(DbError::Init("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantOverlayParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn delete(&self, _definition_id: &str) -> Result { Ok(false) } + + async fn delete_for_user(&self, _user_id: &str, definition_id: &str) -> Result { + self.delete(definition_id).await + } } struct SingleProviderRepo { @@ -403,11 +600,11 @@ mod tests { #[async_trait::async_trait] impl IProviderRepository for SingleProviderRepo { - async fn list(&self) -> Result, DbError> { + async fn list(&self, _user_id: &str) -> Result, DbError> { Ok(self.rows.clone()) } - async fn find_by_id(&self, _id: &str) -> Result, DbError> { + async fn find_by_id(&self, _user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } @@ -415,11 +612,16 @@ mod tests { Err(DbError::NotFound("not implemented".into())) } - async fn update(&self, _id: &str, _params: aionui_db::UpdateProviderParams<'_>) -> Result { + async fn update( + &self, + _user_id: &str, + _id: &str, + _params: aionui_db::UpdateProviderParams<'_>, + ) -> Result { Err(DbError::NotFound("not implemented".into())) } - async fn delete(&self, _id: &str) -> Result<(), DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), DbError> { Err(DbError::NotFound("not implemented".into())) } } @@ -427,6 +629,7 @@ mod tests { fn provider_row(id: &str, models: &[&str]) -> Provider { Provider { id: id.into(), + user_id: "user1".into(), platform: "openai".into(), name: id.into(), base_url: "https://example.com".into(), @@ -456,10 +659,18 @@ mod tests { Ok(self.rows.clone()) } + async fn list_all_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list_all().await + } + async fn get(&self, id: &str) -> Result, DbError> { Ok(self.rows.iter().find(|row| row.id == id).cloned()) } + async fn get_for_user(&self, _user_id: &str, id: &str) -> Result, DbError> { + self.get(id).await + } + async fn find_by_source_and_name( &self, agent_source: &str, @@ -472,6 +683,15 @@ mod tests { .cloned()) } + async fn find_by_source_and_name_for_user( + &self, + _user_id: &str, + agent_source: &str, + name: &str, + ) -> Result, DbError> { + self.find_by_source_and_name(agent_source, name).await + } + async fn find_builtin_by_backend(&self, backend: &str) -> Result, DbError> { Ok(self .rows @@ -480,10 +700,26 @@ mod tests { .cloned()) } + async fn find_builtin_by_backend_for_user( + &self, + _user_id: &str, + backend: &str, + ) -> Result, DbError> { + self.find_builtin_by_backend(backend).await + } + async fn upsert(&self, _params: &UpsertAgentMetadataParams<'_>) -> Result { Err(DbError::Init("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAgentMetadataParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn apply_handshake( &self, _id: &str, @@ -492,6 +728,15 @@ mod tests { Ok(None) } + async fn apply_handshake_for_user( + &self, + _user_id: &str, + id: &str, + params: &UpdateAgentHandshakeParams<'_>, + ) -> Result, DbError> { + self.apply_handshake(id, params).await + } + async fn update_availability_snapshot( &self, _id: &str, @@ -500,6 +745,15 @@ mod tests { Ok(None) } + async fn update_availability_snapshot_for_user( + &self, + _user_id: &str, + id: &str, + params: &aionui_db::models::UpdateAgentAvailabilitySnapshotParams<'_>, + ) -> Result, DbError> { + self.update_availability_snapshot(id, params).await + } + async fn update_agent_overrides( &self, _id: &str, @@ -509,13 +763,31 @@ mod tests { Ok(()) } + async fn update_agent_overrides_for_user( + &self, + _user_id: &str, + id: &str, + command_override: Option<&str>, + env_override: Option<&str>, + ) -> Result<(), DbError> { + self.update_agent_overrides(id, command_override, env_override).await + } + async fn set_enabled(&self, _id: &str, _enabled: bool) -> Result { Ok(false) } + async fn set_enabled_for_user(&self, _user_id: &str, id: &str, enabled: bool) -> Result { + self.set_enabled(id, enabled).await + } + async fn delete(&self, _id: &str) -> Result { Ok(false) } + + async fn delete_for_user(&self, _user_id: &str, id: &str) -> Result { + self.delete(id).await + } } struct RowsTeamAssistantCatalog { @@ -526,6 +798,7 @@ mod tests { impl crate::ports::TeamAssistantCatalogPort for RowsTeamAssistantCatalog { async fn list_team_selectable_assistants( &self, + _user_id: &str, ) -> Result, TeamError> { Ok(self.rows.clone()) } @@ -669,7 +942,7 @@ mod tests { base.backend_binary_path.clone(), ); - let assistants = svc.list_team_selectable_assistants().await; + let assistants = svc.list_team_selectable_assistants("user1").await; let ids: Vec<&str> = assistants .iter() .map(|assistant| assistant.assistant_id.as_str()) @@ -731,7 +1004,7 @@ mod tests { let svc = service_with_selectable_catalog(vec![], vec![assistant_definition("word-creator", "aionrs")]); let err = svc - .resolve_spawn_backend_and_model(Some("word-creator"), None, "gemini", "gemini-2.5-pro") + .resolve_spawn_backend_and_model("user1", Some("word-creator"), None, "gemini", "gemini-2.5-pro") .await .expect_err("spawn must reject assistants outside the team-selectable catalog"); @@ -845,7 +1118,7 @@ mod tests { ); let (backend, model) = svc - .resolve_spawn_backend_and_model(Some("word-creator"), None, "gemini", "gemini-2.5-pro") + .resolve_spawn_backend_and_model("user1", Some("word-creator"), None, "gemini", "gemini-2.5-pro") .await .unwrap(); diff --git a/crates/aionui-team/src/session.rs b/crates/aionui-team/src/session.rs index 95017031d..496e1503b 100644 --- a/crates/aionui-team/src/session.rs +++ b/crates/aionui-team/src/session.rs @@ -166,12 +166,16 @@ impl TeamSession { service: Weak, prompt_dump: TeamPromptDumpConfig, ) -> Result { - let mailbox = Arc::new(Mailbox::new(repo.clone())); - let task_board = Arc::new(TaskBoard::new(repo)); + let mailbox = Arc::new(Mailbox::new_for_user(repo.clone(), user_id.clone())); + let task_board = Arc::new(TaskBoard::new_for_user(repo, user_id.clone())); let member_runtimes = Arc::new(MemberRuntimeRegistry::new(generate_id())); let team_run_manager = Arc::new(TeamRunManager::new( team.id.clone(), - Arc::new(TeamEventEmitter::new(team.id.clone(), broadcaster.clone())), + Arc::new(TeamEventEmitter::new( + team.id.clone(), + user_id.clone(), + broadcaster.clone(), + )), )); let work_coordinator = Arc::new(SlotWorkCoordinator::new( team.id.clone(), @@ -181,6 +185,7 @@ impl TeamSession { let scheduler = Arc::new(TeammateManager::new( team.id.clone(), + user_id.clone(), &team.agents, mailbox.clone(), task_board.clone(), @@ -262,7 +267,11 @@ impl TeamSession { } pub(crate) fn team_event_emitter(&self) -> Arc { - Arc::new(TeamEventEmitter::new(self.team.id.clone(), self.broadcaster.clone())) + Arc::new(TeamEventEmitter::new( + self.team.id.clone(), + self.user_id.clone(), + self.broadcaster.clone(), + )) } pub fn team_run_manager(&self) -> &Arc { @@ -295,7 +304,7 @@ impl TeamSession { let Some(service) = self.service.upgrade() else { return Ok(TeamToolTransport::Mcp); }; - service.provisioner().team_tool_transport(agent).await + service.provisioner().team_tool_transport(&self.user_id, agent).await } pub(crate) async fn prepare_next_batch(&self, slot_id: &str) -> Result { @@ -316,7 +325,9 @@ impl TeamSession { .work_coordinator .set_runtime_constraint(slot_id, runtime_constraint); if !update.terminal_message_ids.is_empty() { - self.mailbox.mark_read_batch(&update.terminal_message_ids).await?; + self.mailbox + .mark_read_batch(&self.team.id, &update.terminal_message_ids) + .await?; } let unread = self @@ -538,7 +549,13 @@ impl TeamSession { let reservation = self.member_runtimes.reserve_attach(slot_id, false); if matches!(reservation, ReserveAttach::Start(_)) { let agent = self.scheduler.get_agent(slot_id).await?; - service.broadcast_agent_runtime_status(&self.team.id, &agent, TeamAgentRuntimeStatus::Pending, None); + service.broadcast_agent_runtime_status( + self.user_id(), + &self.team.id, + &agent, + TeamAgentRuntimeStatus::Pending, + None, + ); info!( team_id = %self.team.id, slot_id, @@ -603,6 +620,7 @@ impl TeamSession { let projection = TeamMessageProjection::new(self.projection_store.clone(), self.broadcaster.clone()); let request = TeamProjectionRequest::user_visible( + &self.user_id, &self.team.id, slot_id, &agent.conversation_id, @@ -654,7 +672,9 @@ impl TeamSession { }; let update = self.work_coordinator.set_runtime_constraint(slot_id, constraint); if !update.terminal_message_ids.is_empty() { - self.mailbox.mark_read_batch(&update.terminal_message_ids).await?; + self.mailbox + .mark_read_batch(&self.team.id, &update.terminal_message_ids) + .await?; } Ok(()) } @@ -702,6 +722,7 @@ impl TeamSession { let projection = TeamMessageProjection::new(self.projection_store.clone(), self.broadcaster.clone()); let request = TeamProjectionRequest { + user_id: self.user_id.clone(), team_id: self.team.id.clone(), slot_id: to_slot_id.to_owned(), conversation_id: to_agent.conversation_id.clone(), @@ -830,7 +851,7 @@ impl TeamSession { self.work_coordinator.abort_enqueue(lease, "enqueue_commit_failed"); if let Err(mark_read_error) = self .mailbox - .mark_read_batch(std::slice::from_ref(&mailbox_message_id)) + .mark_read_batch(&self.team.id, std::slice::from_ref(&mailbox_message_id)) .await { warn!( @@ -948,6 +969,7 @@ impl TeamSession { msg.content.clone() }; let request = TeamProjectionRequest { + user_id: self.user_id.clone(), team_id: self.team.id.clone(), slot_id: msg.to_agent_id.clone(), conversation_id: input.conversation_id.clone(), @@ -988,7 +1010,9 @@ impl TeamSession { self.team_run_manager.begin_cancel(team_run_id, reason)?; let result = self.work_coordinator.cancel_run(team_run_id); if !result.terminal_message_ids.is_empty() { - self.mailbox.mark_read_batch(&result.terminal_message_ids).await?; + self.mailbox + .mark_read_batch(&self.team.id, &result.terminal_message_ids) + .await?; } for target in result.cancel_targets { let Some(turn_id) = target.turn_id else { @@ -1304,6 +1328,7 @@ impl TeamSession { ) { let projection = TeamMessageProjection::new(self.projection_store.clone(), self.broadcaster.clone()); let request = TeamProjectionRequest::team_system_visible( + &self.user_id, &self.team.id, slot_id, conversation_id, @@ -1340,7 +1365,9 @@ impl TeamSession { } let work = self.work_coordinator.remove_slot(slot_id); if !work.terminal_message_ids.is_empty() { - self.mailbox.mark_read_batch(&work.terminal_message_ids).await?; + self.mailbox + .mark_read_batch(&self.team.id, &work.terminal_message_ids) + .await?; } if let Some(target) = work.cancel_target && let Some(turn_id) = target.turn_id @@ -1421,7 +1448,13 @@ impl TeamSession { // capability checks. Assistant spawns derive backend from the preset // identity rather than inheriting the caller backend. let (backend, model) = service - .resolve_spawn_backend_and_model(Some(assistant_id), None, caller.backend.as_str(), caller.model.as_str()) + .resolve_spawn_backend_and_model( + &self.user_id, + Some(assistant_id), + None, + caller.backend.as_str(), + caller.model.as_str(), + ) .await?; // Step 4: DB side-effects (new conversation + persisted agent slot). @@ -1645,6 +1678,7 @@ pub(crate) async fn attach_member_runtime( let failure = sanitize_member_runtime_failure(&error); if session.member_runtimes.commit_failed(&lease, failure.clone()) { service.broadcast_agent_runtime_status( + session.user_id(), session.team_id(), &agent, TeamAgentRuntimeStatus::Failed, @@ -2845,7 +2879,7 @@ mod tests { #[test] fn session_replacement_rejects_old_generation_batch_and_attach_completion() { let old_broadcaster: Arc = Arc::new(NullBroadcaster); - let old_emitter = Arc::new(TeamEventEmitter::new("t1".into(), old_broadcaster)); + let old_emitter = Arc::new(TeamEventEmitter::new("t1".into(), "user-1".into(), old_broadcaster)); let old_runs = Arc::new(TeamRunManager::new("t1".into(), old_emitter)); let old_coordinator = SlotWorkCoordinator::new("t1".into(), "old-generation".into(), old_runs); old_coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); @@ -2863,7 +2897,7 @@ mod tests { }; let new_broadcaster: Arc = Arc::new(NullBroadcaster); - let new_emitter = Arc::new(TeamEventEmitter::new("t1".into(), new_broadcaster)); + let new_emitter = Arc::new(TeamEventEmitter::new("t1".into(), "user-1".into(), new_broadcaster)); let new_runs = Arc::new(TeamRunManager::new("t1".into(), new_emitter)); let new_coordinator = SlotWorkCoordinator::new("t1".into(), "new-generation".into(), new_runs); new_coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); diff --git a/crates/aionui-team/src/task_board.rs b/crates/aionui-team/src/task_board.rs index 54b8f83a4..0af9fd5c1 100644 --- a/crates/aionui-team/src/task_board.rs +++ b/crates/aionui-team/src/task_board.rs @@ -11,6 +11,7 @@ use crate::types::{TaskStatus, TeamTask}; pub struct TaskBoard { repo: Arc, + user_id: String, } /// Optional fields for task update. @@ -25,7 +26,14 @@ pub struct TaskUpdate { impl TaskBoard { pub fn new(repo: Arc) -> Self { - Self { repo } + Self::new_for_user(repo, "system_default_user") + } + + pub fn new_for_user(repo: Arc, user_id: impl Into) -> Self { + Self { + repo, + user_id: user_id.into(), + } } pub async fn create_task( @@ -37,7 +45,7 @@ impl TaskBoard { blocked_by: &[String], ) -> Result { for dep_id in blocked_by { - let dep = self.repo.find_task_by_id(team_id, dep_id).await?; + let dep = self.repo.find_task_by_id(&self.user_id, team_id, dep_id).await?; if dep.is_none() { return Err(TeamError::BlockedTaskNotFound(dep_id.clone())); } @@ -61,10 +69,12 @@ impl TaskBoard { updated_at: now, }; - self.repo.create_task(&row).await?; + self.repo.create_task(&self.user_id, &row).await?; for dep_id in blocked_by { - self.repo.append_to_blocks(dep_id, &task_id).await?; + self.repo + .append_to_blocks(&self.user_id, team_id, dep_id, &task_id) + .await?; } debug!(team_id, task_id = %task_id, subject, "task created"); @@ -75,7 +85,7 @@ impl TaskBoard { pub async fn update_task(&self, team_id: &str, task_id: &str, update: &TaskUpdate) -> Result { let existing = self .repo - .find_task_by_id(team_id, task_id) + .find_task_by_id(&self.user_id, team_id, task_id) .await? .ok_or_else(|| TeamError::TaskNotFound(task_id.to_owned()))?; @@ -87,15 +97,15 @@ impl TaskBoard { metadata: update.metadata.as_ref().map(serde_json::to_string).transpose()?, }; - self.repo.update_task(task_id, ¶ms).await?; + self.repo.update_task(&self.user_id, team_id, task_id, ¶ms).await?; if update.status == Some(TaskStatus::Completed) { - self.check_unblocks(task_id, &existing).await?; + self.check_unblocks(team_id, task_id, &existing).await?; } let updated = self .repo - .find_task_by_id(team_id, task_id) + .find_task_by_id(&self.user_id, team_id, task_id) .await? .ok_or_else(|| TeamError::TaskNotFound(task_id.to_owned()))?; @@ -105,16 +115,21 @@ impl TaskBoard { } pub async fn list_tasks(&self, team_id: &str) -> Result, TeamError> { - let rows = self.repo.list_tasks(team_id).await?; + let rows = self.repo.list_tasks(&self.user_id, team_id).await?; let tasks = rows.iter().filter_map(|r| TeamTask::from_row(r).ok()).collect(); Ok(tasks) } - async fn check_unblocks(&self, completed_task_id: &str, completed_row: &TeamTaskRow) -> Result<(), TeamError> { + async fn check_unblocks( + &self, + team_id: &str, + completed_task_id: &str, + completed_row: &TeamTaskRow, + ) -> Result<(), TeamError> { let blocks: Vec = serde_json::from_str(&completed_row.blocks)?; for downstream_id in &blocks { self.repo - .remove_from_blocked_by(downstream_id, completed_task_id) + .remove_from_blocked_by(&self.user_id, team_id, downstream_id, completed_task_id) .await?; debug!( completed = completed_task_id, @@ -177,7 +192,11 @@ mod tests { assert_eq!(task_b.blocked_by, vec![task_a.id.clone()]); - let updated_a = repo.find_task_by_id("t1", &task_a.id).await.unwrap().unwrap(); + let updated_a = repo + .find_task_by_id("system_default_user", "t1", &task_a.id) + .await + .unwrap() + .unwrap(); let blocks_a: Vec = serde_json::from_str(&updated_a.blocks).unwrap(); assert_eq!(blocks_a, vec![task_b.id]); } diff --git a/crates/aionui-team/src/team_run/tests.rs b/crates/aionui-team/src/team_run/tests.rs index d54851b1d..e8acf2ab7 100644 --- a/crates/aionui-team/src/team_run/tests.rs +++ b/crates/aionui-team/src/team_run/tests.rs @@ -12,7 +12,7 @@ use crate::work_source::WorkSource; fn coordinator_and_manager() -> (Arc, Arc) { let broadcaster = Arc::new(RecordingBroadcaster::new()); - let emitter = Arc::new(TeamEventEmitter::new("team-1".into(), broadcaster)); + let emitter = Arc::new(TeamEventEmitter::new("team-1".into(), "user-1".into(), broadcaster)); let manager = Arc::new(TeamRunManager::new("team-1".into(), emitter)); let coordinator = Arc::new(SlotWorkCoordinator::new( "team-1".into(), @@ -133,7 +133,11 @@ fn aborting_the_creator_lease_preserves_a_concurrent_committed_enqueue() { #[test] fn accepted_event_is_emitted_only_after_the_durable_enqueue_commits() { let broadcaster = Arc::new(RecordingBroadcaster::new()); - let emitter = Arc::new(TeamEventEmitter::new("team-1".into(), broadcaster.clone())); + let emitter = Arc::new(TeamEventEmitter::new( + "team-1".into(), + "user-1".into(), + broadcaster.clone(), + )); let manager = Arc::new(TeamRunManager::new("team-1".into(), emitter)); let coordinator = SlotWorkCoordinator::new("team-1".into(), "generation-1".into(), manager); coordinator.set_runtime_constraint("lead-1", RuntimeConstraint::Ready); diff --git a/crates/aionui-team/src/test_utils.rs b/crates/aionui-team/src/test_utils.rs index ff409483b..badaa5dc8 100644 --- a/crates/aionui-team/src/test_utils.rs +++ b/crates/aionui-team/src/test_utils.rs @@ -30,25 +30,40 @@ impl ITeamRepository for MockTeamRepo { async fn create_team(&self, _row: &TeamRow) -> Result<(), DbError> { Ok(()) } - async fn list_teams(&self) -> Result, DbError> { + async fn list_teams_for_restore(&self) -> Result, DbError> { Ok(vec![]) } async fn list_teams_by_user(&self, _user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn get_team(&self, _id: &str) -> Result, DbError> { + async fn get_team(&self, _user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } - async fn update_team(&self, _id: &str, _p: &UpdateTeamParams) -> Result<(), DbError> { + async fn get_team_for_restore(&self, _id: &str) -> Result, DbError> { + Ok(Some(TeamRow { + id: _id.to_owned(), + user_id: "system_default_user".to_owned(), + name: "mock".to_owned(), + workspace: String::new(), + workspace_mode: "shared".to_owned(), + agents: "[]".to_owned(), + lead_agent_id: None, + session_mode: None, + agents_version: "1.0.1".to_owned(), + created_at: now_ms(), + updated_at: now_ms(), + })) + } + async fn update_team(&self, _user_id: &str, _id: &str, _p: &UpdateTeamParams) -> Result<(), DbError> { Ok(()) } - async fn delete_team(&self, _id: &str) -> Result<(), DbError> { + async fn delete_team(&self, _user_id: &str, _id: &str) -> Result<(), DbError> { Ok(()) } // ── Mailbox ───────────────────────────────────────────────────── - async fn write_message(&self, row: &MailboxMessageRow) -> Result<(), DbError> { + async fn write_message(&self, _user_id: &str, row: &MailboxMessageRow) -> Result<(), DbError> { let mut state = self.state.lock().unwrap(); if state.fail_message_writes { return Err(DbError::Init("forced mailbox write failure".into())); @@ -57,7 +72,12 @@ impl ITeamRepository for MockTeamRepo { Ok(()) } - async fn read_unread_and_mark(&self, team_id: &str, to_agent_id: &str) -> Result, DbError> { + async fn read_unread_and_mark( + &self, + _user_id: &str, + team_id: &str, + to_agent_id: &str, + ) -> Result, DbError> { let mut state = self.state.lock().unwrap(); let mut result = vec![]; for msg in &mut state.messages { @@ -69,7 +89,12 @@ impl ITeamRepository for MockTeamRepo { Ok(result) } - async fn peek_unread(&self, team_id: &str, to_agent_id: &str) -> Result, DbError> { + async fn peek_unread( + &self, + _user_id: &str, + team_id: &str, + to_agent_id: &str, + ) -> Result, DbError> { let state = self.state.lock().unwrap(); let result = state .messages @@ -80,10 +105,10 @@ impl ITeamRepository for MockTeamRepo { Ok(result) } - async fn mark_read_batch(&self, ids: &[String]) -> Result<(), DbError> { + async fn mark_read_batch(&self, _user_id: &str, team_id: &str, ids: &[String]) -> Result<(), DbError> { let mut state = self.state.lock().unwrap(); for msg in &mut state.messages { - if ids.contains(&msg.id) { + if msg.team_id == team_id && ids.contains(&msg.id) { msg.read = true; } } @@ -92,6 +117,7 @@ impl ITeamRepository for MockTeamRepo { async fn get_history( &self, + _user_id: &str, team_id: &str, to_agent_id: &str, limit: Option, @@ -108,19 +134,24 @@ impl ITeamRepository for MockTeamRepo { Ok(msgs) } - async fn delete_mailbox_by_team(&self, team_id: &str) -> Result<(), DbError> { + async fn delete_mailbox_by_team(&self, _user_id: &str, team_id: &str) -> Result<(), DbError> { self.state.lock().unwrap().messages.retain(|m| m.team_id != team_id); Ok(()) } // ── TaskBoard ─────────────────────────────────────────────────── - async fn create_task(&self, row: &TeamTaskRow) -> Result<(), DbError> { + async fn create_task(&self, _user_id: &str, row: &TeamTaskRow) -> Result<(), DbError> { self.state.lock().unwrap().tasks.push(row.clone()); Ok(()) } - async fn find_task_by_id(&self, team_id: &str, task_id: &str) -> Result, DbError> { + async fn find_task_by_id( + &self, + _user_id: &str, + team_id: &str, + task_id: &str, + ) -> Result, DbError> { let state = self.state.lock().unwrap(); let found = state .tasks @@ -130,12 +161,18 @@ impl ITeamRepository for MockTeamRepo { Ok(found) } - async fn update_task(&self, task_id: &str, params: &UpdateTaskParams) -> Result<(), DbError> { + async fn update_task( + &self, + _user_id: &str, + team_id: &str, + task_id: &str, + params: &UpdateTaskParams, + ) -> Result<(), DbError> { let mut state = self.state.lock().unwrap(); let task = state .tasks .iter_mut() - .find(|t| t.id == task_id) + .find(|t| t.team_id == team_id && t.id == task_id) .ok_or_else(|| DbError::NotFound(task_id.to_owned()))?; if let Some(ref s) = params.status { task.status = s.clone(); @@ -156,7 +193,7 @@ impl ITeamRepository for MockTeamRepo { Ok(()) } - async fn list_tasks(&self, team_id: &str) -> Result, DbError> { + async fn list_tasks(&self, _user_id: &str, team_id: &str) -> Result, DbError> { let state = self.state.lock().unwrap(); if state.fail_task_lists { return Err(DbError::Init("forced task list failure".into())); @@ -165,12 +202,18 @@ impl ITeamRepository for MockTeamRepo { Ok(tasks) } - async fn append_to_blocks(&self, task_id: &str, blocked_task_id: &str) -> Result<(), DbError> { + async fn append_to_blocks( + &self, + _user_id: &str, + team_id: &str, + task_id: &str, + blocked_task_id: &str, + ) -> Result<(), DbError> { let mut state = self.state.lock().unwrap(); let task = state .tasks .iter_mut() - .find(|t| t.id == task_id) + .find(|t| t.team_id == team_id && t.id == task_id) .ok_or_else(|| DbError::NotFound(task_id.to_owned()))?; let mut blocks: Vec = serde_json::from_str(&task.blocks).unwrap_or_default(); blocks.push(blocked_task_id.to_owned()); @@ -178,12 +221,18 @@ impl ITeamRepository for MockTeamRepo { Ok(()) } - async fn remove_from_blocked_by(&self, task_id: &str, unblocked_task_id: &str) -> Result<(), DbError> { + async fn remove_from_blocked_by( + &self, + _user_id: &str, + team_id: &str, + task_id: &str, + unblocked_task_id: &str, + ) -> Result<(), DbError> { let mut state = self.state.lock().unwrap(); let task = state .tasks .iter_mut() - .find(|t| t.id == task_id) + .find(|t| t.team_id == team_id && t.id == task_id) .ok_or_else(|| DbError::NotFound(task_id.to_owned()))?; let mut blocked_by: Vec = serde_json::from_str(&task.blocked_by).unwrap_or_default(); blocked_by.retain(|id| id != unblocked_task_id); @@ -191,7 +240,7 @@ impl ITeamRepository for MockTeamRepo { Ok(()) } - async fn delete_tasks_by_team(&self, team_id: &str) -> Result<(), DbError> { + async fn delete_tasks_by_team(&self, _user_id: &str, team_id: &str) -> Result<(), DbError> { self.state.lock().unwrap().tasks.retain(|t| t.team_id != team_id); Ok(()) } @@ -264,8 +313,24 @@ pub(crate) mod workspace_harness { #[async_trait] impl IConversationRepository for MockConversationRepo { - async fn get(&self, id: &str) -> Result, DbError> { - Ok(self.conversations.lock().unwrap().iter().find(|c| c.id == id).cloned()) + async fn get(&self, user_id: &str, id: &str) -> Result, DbError> { + Ok(self + .conversations + .lock() + .unwrap() + .iter() + .find(|c| c.user_id == user_id && c.id == id) + .cloned()) + } + + async fn owner_user_id(&self, id: &str) -> Result, DbError> { + Ok(self + .conversations + .lock() + .unwrap() + .iter() + .find(|c| c.id == id) + .map(|c| c.user_id.clone())) } async fn create(&self, row: &ConversationRow) -> Result<(), DbError> { @@ -273,11 +338,11 @@ pub(crate) mod workspace_harness { Ok(()) } - async fn update(&self, id: &str, updates: &ConversationRowUpdate) -> Result<(), DbError> { + async fn update(&self, user_id: &str, id: &str, updates: &ConversationRowUpdate) -> Result<(), DbError> { let mut conversations = self.conversations.lock().unwrap(); let conversation = conversations .iter_mut() - .find(|c| c.id == id) + .find(|c| c.user_id == user_id && c.id == id) .ok_or_else(|| DbError::NotFound(id.to_owned()))?; if let Some(ref extra) = updates.extra { conversation.extra = extra.clone(); @@ -297,8 +362,11 @@ pub(crate) mod workspace_harness { Ok(()) } - async fn delete(&self, id: &str) -> Result<(), DbError> { - self.conversations.lock().unwrap().retain(|c| c.id != id); + async fn delete(&self, user_id: &str, id: &str) -> Result<(), DbError> { + self.conversations + .lock() + .unwrap() + .retain(|c| c.user_id != user_id || c.id != id); Ok(()) } @@ -338,6 +406,7 @@ pub(crate) mod workspace_harness { async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &MessagePageParams, ) -> Result { @@ -348,20 +417,27 @@ pub(crate) mod workspace_harness { }) } - async fn insert_message(&self, _message: &MessageRow) -> Result<(), DbError> { + async fn insert_message(&self, _user_id: &str, _message: &MessageRow) -> Result<(), DbError> { Ok(()) } - async fn update_message(&self, _id: &str, _updates: &MessageRowUpdate) -> Result<(), DbError> { + async fn update_message( + &self, + _user_id: &str, + _conversation_id: &str, + _id: &str, + _updates: &MessageRowUpdate, + ) -> Result<(), DbError> { Ok(()) } - async fn delete_messages_by_conversation(&self, _conv_id: &str) -> Result<(), DbError> { + async fn delete_messages_by_conversation(&self, _user_id: &str, _conv_id: &str) -> Result<(), DbError> { Ok(()) } async fn get_message_by_msg_id( &self, + _user_id: &str, _conv_id: &str, _msg_id: &str, _msg_type: &str, @@ -403,7 +479,7 @@ pub(crate) mod workspace_harness { Ok(()) } - async fn list_teams(&self) -> Result, DbError> { + async fn list_teams_for_restore(&self) -> Result, DbError> { Ok(self.teams.lock().unwrap().clone()) } @@ -418,15 +494,25 @@ pub(crate) mod workspace_harness { .collect()) } - async fn get_team(&self, id: &str) -> Result, DbError> { + async fn get_team(&self, user_id: &str, id: &str) -> Result, DbError> { + Ok(self + .teams + .lock() + .unwrap() + .iter() + .find(|t| t.user_id == user_id && t.id == id) + .cloned()) + } + + async fn get_team_for_restore(&self, id: &str) -> Result, DbError> { Ok(self.teams.lock().unwrap().iter().find(|t| t.id == id).cloned()) } - async fn update_team(&self, id: &str, params: &UpdateTeamParams) -> Result<(), DbError> { + async fn update_team(&self, user_id: &str, id: &str, params: &UpdateTeamParams) -> Result<(), DbError> { let mut teams = self.teams.lock().unwrap(); let team = teams .iter_mut() - .find(|t| t.id == id) + .find(|t| t.user_id == user_id && t.id == id) .ok_or_else(|| DbError::NotFound(id.to_owned()))?; if let Some(ref name) = params.name { team.name = name.clone(); @@ -447,17 +533,25 @@ pub(crate) mod workspace_harness { Ok(()) } - async fn delete_team(&self, id: &str) -> Result<(), DbError> { - self.teams.lock().unwrap().retain(|t| t.id != id); + async fn delete_team(&self, user_id: &str, id: &str) -> Result<(), DbError> { + self.teams + .lock() + .unwrap() + .retain(|t| t.user_id != user_id || t.id != id); Ok(()) } - async fn write_message(&self, _row: &aionui_db::models::MailboxMessageRow) -> Result<(), DbError> { + async fn write_message( + &self, + _user_id: &str, + _row: &aionui_db::models::MailboxMessageRow, + ) -> Result<(), DbError> { Ok(()) } async fn read_unread_and_mark( &self, + _user_id: &str, _team_id: &str, _to_agent_id: &str, ) -> Result, DbError> { @@ -466,18 +560,20 @@ pub(crate) mod workspace_harness { async fn peek_unread( &self, + _user_id: &str, _team_id: &str, _to_agent_id: &str, ) -> Result, DbError> { Ok(vec![]) } - async fn mark_read_batch(&self, _ids: &[String]) -> Result<(), DbError> { + async fn mark_read_batch(&self, _user_id: &str, _team_id: &str, _ids: &[String]) -> Result<(), DbError> { Ok(()) } async fn get_history( &self, + _user_id: &str, _team_id: &str, _to_agent_id: &str, _limit: Option, @@ -485,35 +581,58 @@ pub(crate) mod workspace_harness { Ok(vec![]) } - async fn delete_mailbox_by_team(&self, _team_id: &str) -> Result<(), DbError> { + async fn delete_mailbox_by_team(&self, _user_id: &str, _team_id: &str) -> Result<(), DbError> { Ok(()) } - async fn create_task(&self, _row: &TeamTaskRow) -> Result<(), DbError> { + async fn create_task(&self, _user_id: &str, _row: &TeamTaskRow) -> Result<(), DbError> { Ok(()) } - async fn find_task_by_id(&self, _team_id: &str, _task_id: &str) -> Result, DbError> { + async fn find_task_by_id( + &self, + _user_id: &str, + _team_id: &str, + _task_id: &str, + ) -> Result, DbError> { Ok(None) } - async fn update_task(&self, _task_id: &str, _params: &aionui_db::UpdateTaskParams) -> Result<(), DbError> { + async fn update_task( + &self, + _user_id: &str, + _team_id: &str, + _task_id: &str, + _params: &aionui_db::UpdateTaskParams, + ) -> Result<(), DbError> { Ok(()) } - async fn list_tasks(&self, _team_id: &str) -> Result, DbError> { + async fn list_tasks(&self, _user_id: &str, _team_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn append_to_blocks(&self, _task_id: &str, _blocked_task_id: &str) -> Result<(), DbError> { + async fn append_to_blocks( + &self, + _user_id: &str, + _team_id: &str, + _task_id: &str, + _blocked_task_id: &str, + ) -> Result<(), DbError> { Ok(()) } - async fn remove_from_blocked_by(&self, _task_id: &str, _unblocked_task_id: &str) -> Result<(), DbError> { + async fn remove_from_blocked_by( + &self, + _user_id: &str, + _team_id: &str, + _task_id: &str, + _unblocked_task_id: &str, + ) -> Result<(), DbError> { Ok(()) } - async fn delete_tasks_by_team(&self, _team_id: &str) -> Result<(), DbError> { + async fn delete_tasks_by_team(&self, _user_id: &str, _team_id: &str) -> Result<(), DbError> { Ok(()) } } @@ -620,8 +739,14 @@ pub(crate) mod workspace_harness { target.insert(key.clone(), value.clone()); } } + let user_id = self + .repo + .owner_user_id(conversation_id) + .await? + .ok_or_else(|| TeamError::AgentNotFound(conversation_id.to_owned()))?; self.repo .update( + &user_id, conversation_id, &ConversationRowUpdate { name: None, @@ -689,7 +814,8 @@ pub(crate) mod workspace_harness { Ok(()) } - async fn delete_team_conversation(&self, _user_id: &str, _conversation_id: &str) -> Result<(), TeamError> { + async fn delete_team_conversation(&self, user_id: &str, conversation_id: &str) -> Result<(), TeamError> { + self.repo.delete(user_id, conversation_id).await?; Ok(()) } } @@ -809,7 +935,10 @@ pub(crate) mod workspace_harness { #[async_trait] impl TeamAssistantCatalogPort for EmptyTeamAssistantCatalog { - async fn list_team_selectable_assistants(&self) -> Result, TeamError> { + async fn list_team_selectable_assistants( + &self, + _user_id: &str, + ) -> Result, TeamError> { Ok(Vec::new()) } } @@ -820,10 +949,18 @@ pub(crate) mod workspace_harness { Ok(vec![]) } + async fn list_all_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list_all().await + } + async fn get(&self, _id: &str) -> Result, DbError> { Ok(None) } + async fn get_for_user(&self, _user_id: &str, id: &str) -> Result, DbError> { + self.get(id).await + } + async fn find_by_source_and_name( &self, _agent_source: &str, @@ -832,14 +969,39 @@ pub(crate) mod workspace_harness { Ok(None) } + async fn find_by_source_and_name_for_user( + &self, + _user_id: &str, + agent_source: &str, + name: &str, + ) -> Result, DbError> { + self.find_by_source_and_name(agent_source, name).await + } + async fn find_builtin_by_backend(&self, _backend: &str) -> Result, DbError> { Ok(None) } + async fn find_builtin_by_backend_for_user( + &self, + _user_id: &str, + backend: &str, + ) -> Result, DbError> { + self.find_builtin_by_backend(backend).await + } + async fn upsert(&self, _params: &UpsertAgentMetadataParams<'_>) -> Result { Err(DbError::NotFound("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAgentMetadataParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn apply_handshake( &self, _id: &str, @@ -848,6 +1010,15 @@ pub(crate) mod workspace_harness { Ok(None) } + async fn apply_handshake_for_user( + &self, + _user_id: &str, + id: &str, + params: &UpdateAgentHandshakeParams<'_>, + ) -> Result, DbError> { + self.apply_handshake(id, params).await + } + async fn update_availability_snapshot( &self, _id: &str, @@ -856,6 +1027,15 @@ pub(crate) mod workspace_harness { Ok(None) } + async fn update_availability_snapshot_for_user( + &self, + _user_id: &str, + id: &str, + params: &aionui_db::models::UpdateAgentAvailabilitySnapshotParams<'_>, + ) -> Result, DbError> { + self.update_availability_snapshot(id, params).await + } + async fn update_agent_overrides( &self, _id: &str, @@ -865,13 +1045,31 @@ pub(crate) mod workspace_harness { Ok(()) } + async fn update_agent_overrides_for_user( + &self, + _user_id: &str, + id: &str, + command_override: Option<&str>, + env_override: Option<&str>, + ) -> Result<(), DbError> { + self.update_agent_overrides(id, command_override, env_override).await + } + async fn set_enabled(&self, _id: &str, _enabled: bool) -> Result { Ok(false) } + async fn set_enabled_for_user(&self, _user_id: &str, id: &str, enabled: bool) -> Result { + self.set_enabled(id, enabled).await + } + async fn delete(&self, _id: &str) -> Result { Ok(false) } + + async fn delete_for_user(&self, _user_id: &str, id: &str) -> Result { + self.delete(id).await + } } struct EmptyAssistantDefinitionRepo; @@ -882,14 +1080,49 @@ pub(crate) mod workspace_harness { Ok(vec![]) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + + async fn list_including_deleted_for_user( + &self, + _user_id: &str, + ) -> Result, DbError> { + self.list().await + } + async fn get_by_assistant_id(&self, _assistant_id: &str) -> Result, DbError> { Ok(None) } + async fn get_by_assistant_id_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + + async fn get_by_assistant_id_including_deleted_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + async fn get_by_id(&self, _definition_id: &str) -> Result, DbError> { Ok(None) } + async fn get_by_id_for_user( + &self, + _user_id: &str, + definition_id: &str, + ) -> Result, DbError> { + self.get_by_id(definition_id).await + } + async fn get_by_source_ref( &self, _source: &str, @@ -898,6 +1131,24 @@ pub(crate) mod workspace_harness { Ok(None) } + async fn get_by_source_ref_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + + async fn get_by_source_ref_including_deleted_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + async fn upsert( &self, _params: &UpsertAssistantDefinitionParams<'_>, @@ -905,9 +1156,26 @@ pub(crate) mod workspace_harness { Err(DbError::Init("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn soft_delete(&self, _definition_id: &str, _deleted_at: i64) -> Result { Ok(false) } + + async fn soft_delete_for_user( + &self, + _user_id: &str, + definition_id: &str, + deleted_at: i64, + ) -> Result { + self.soft_delete(definition_id, deleted_at).await + } } struct EmptyAssistantOverlayRepo; @@ -918,28 +1186,52 @@ pub(crate) mod workspace_harness { Ok(None) } + async fn get_for_user( + &self, + _user_id: &str, + definition_id: &str, + ) -> Result, DbError> { + self.get(definition_id).await + } + async fn list(&self) -> Result, DbError> { Ok(vec![]) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + async fn upsert(&self, _params: &UpsertAssistantOverlayParams<'_>) -> Result { Err(DbError::Init("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantOverlayParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn delete(&self, _definition_id: &str) -> Result { Ok(false) } + + async fn delete_for_user(&self, _user_id: &str, definition_id: &str) -> Result { + self.delete(definition_id).await + } } struct EmptyProviderRepo; #[async_trait] impl IProviderRepository for EmptyProviderRepo { - async fn list(&self) -> Result, DbError> { + async fn list(&self, _user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn find_by_id(&self, _id: &str) -> Result, DbError> { + async fn find_by_id(&self, _user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } @@ -952,13 +1244,14 @@ pub(crate) mod workspace_harness { async fn update( &self, + _user_id: &str, _id: &str, _params: aionui_db::UpdateProviderParams<'_>, ) -> Result { Err(DbError::NotFound("not implemented".into())) } - async fn delete(&self, _id: &str) -> Result<(), DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), DbError> { Ok(()) } } @@ -1056,6 +1349,7 @@ pub(crate) mod workspace_harness { pub(crate) async fn force_team_workspace(repo: &Arc, team_id: &str, workspace: &str) { repo.update_team( + "user1", team_id, &UpdateTeamParams { workspace: Some(workspace.to_owned()), diff --git a/crates/aionui-team/src/workspace.rs b/crates/aionui-team/src/workspace.rs index 438ed2941..1c33c31e0 100644 --- a/crates/aionui-team/src/workspace.rs +++ b/crates/aionui-team/src/workspace.rs @@ -94,8 +94,14 @@ impl TeamWorkspaceResolver { } async fn write_team_workspace(&self, team_id: &str, workspace: &str) -> Result<(), TeamError> { + let row = self + .repo + .get_team_for_restore(team_id) + .await? + .ok_or_else(|| TeamError::TeamNotFound(team_id.to_owned()))?; self.repo .update_team( + &row.user_id, team_id, &UpdateTeamParams { workspace: Some(workspace.to_owned()), diff --git a/crates/aionui-team/tests/common/mod.rs b/crates/aionui-team/tests/common/mod.rs index d6aa820df..d6d3d821f 100644 --- a/crates/aionui-team/tests/common/mod.rs +++ b/crates/aionui-team/tests/common/mod.rs @@ -26,28 +26,36 @@ impl ITeamRepository for MockTeamRepo { async fn create_team(&self, _row: &TeamRow) -> Result<(), DbError> { Ok(()) } - async fn list_teams(&self) -> Result, DbError> { + async fn list_teams_for_restore(&self) -> Result, DbError> { Ok(vec![]) } async fn list_teams_by_user(&self, _user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn get_team(&self, _id: &str) -> Result, DbError> { + async fn get_team(&self, _user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } - async fn update_team(&self, _id: &str, _p: &UpdateTeamParams) -> Result<(), DbError> { + async fn get_team_for_restore(&self, _id: &str) -> Result, DbError> { + Ok(None) + } + async fn update_team(&self, _user_id: &str, _id: &str, _p: &UpdateTeamParams) -> Result<(), DbError> { Ok(()) } - async fn delete_team(&self, _id: &str) -> Result<(), DbError> { + async fn delete_team(&self, _user_id: &str, _id: &str) -> Result<(), DbError> { Ok(()) } - async fn write_message(&self, row: &MailboxMessageRow) -> Result<(), DbError> { + async fn write_message(&self, _user_id: &str, row: &MailboxMessageRow) -> Result<(), DbError> { self.state.lock().unwrap().messages.push(row.clone()); Ok(()) } - async fn read_unread_and_mark(&self, team_id: &str, to_agent_id: &str) -> Result, DbError> { + async fn read_unread_and_mark( + &self, + _user_id: &str, + team_id: &str, + to_agent_id: &str, + ) -> Result, DbError> { let mut state = self.state.lock().unwrap(); let mut result = vec![]; for msg in &mut state.messages { @@ -59,7 +67,12 @@ impl ITeamRepository for MockTeamRepo { Ok(result) } - async fn peek_unread(&self, team_id: &str, to_agent_id: &str) -> Result, DbError> { + async fn peek_unread( + &self, + _user_id: &str, + team_id: &str, + to_agent_id: &str, + ) -> Result, DbError> { let state = self.state.lock().unwrap(); let result = state .messages @@ -70,10 +83,10 @@ impl ITeamRepository for MockTeamRepo { Ok(result) } - async fn mark_read_batch(&self, ids: &[String]) -> Result<(), DbError> { + async fn mark_read_batch(&self, _user_id: &str, team_id: &str, ids: &[String]) -> Result<(), DbError> { let mut state = self.state.lock().unwrap(); for msg in &mut state.messages { - if ids.contains(&msg.id) { + if msg.team_id == team_id && ids.contains(&msg.id) { msg.read = true; } } @@ -82,6 +95,7 @@ impl ITeamRepository for MockTeamRepo { async fn get_history( &self, + _user_id: &str, team_id: &str, to_agent_id: &str, limit: Option, @@ -98,17 +112,22 @@ impl ITeamRepository for MockTeamRepo { Ok(msgs) } - async fn delete_mailbox_by_team(&self, team_id: &str) -> Result<(), DbError> { + async fn delete_mailbox_by_team(&self, _user_id: &str, team_id: &str) -> Result<(), DbError> { self.state.lock().unwrap().messages.retain(|m| m.team_id != team_id); Ok(()) } - async fn create_task(&self, row: &TeamTaskRow) -> Result<(), DbError> { + async fn create_task(&self, _user_id: &str, row: &TeamTaskRow) -> Result<(), DbError> { self.state.lock().unwrap().tasks.push(row.clone()); Ok(()) } - async fn find_task_by_id(&self, team_id: &str, task_id: &str) -> Result, DbError> { + async fn find_task_by_id( + &self, + _user_id: &str, + team_id: &str, + task_id: &str, + ) -> Result, DbError> { let state = self.state.lock().unwrap(); let found = state .tasks @@ -118,12 +137,18 @@ impl ITeamRepository for MockTeamRepo { Ok(found) } - async fn update_task(&self, task_id: &str, params: &UpdateTaskParams) -> Result<(), DbError> { + async fn update_task( + &self, + _user_id: &str, + team_id: &str, + task_id: &str, + params: &UpdateTaskParams, + ) -> Result<(), DbError> { let mut state = self.state.lock().unwrap(); let task = state .tasks .iter_mut() - .find(|t| t.id == task_id) + .find(|t| t.team_id == team_id && t.id == task_id) .ok_or_else(|| DbError::NotFound(task_id.to_owned()))?; if let Some(ref s) = params.status { task.status = s.clone(); @@ -144,18 +169,24 @@ impl ITeamRepository for MockTeamRepo { Ok(()) } - async fn list_tasks(&self, team_id: &str) -> Result, DbError> { + async fn list_tasks(&self, _user_id: &str, team_id: &str) -> Result, DbError> { let state = self.state.lock().unwrap(); let tasks = state.tasks.iter().filter(|t| t.team_id == team_id).cloned().collect(); Ok(tasks) } - async fn append_to_blocks(&self, task_id: &str, blocked_task_id: &str) -> Result<(), DbError> { + async fn append_to_blocks( + &self, + _user_id: &str, + team_id: &str, + task_id: &str, + blocked_task_id: &str, + ) -> Result<(), DbError> { let mut state = self.state.lock().unwrap(); let task = state .tasks .iter_mut() - .find(|t| t.id == task_id) + .find(|t| t.team_id == team_id && t.id == task_id) .ok_or_else(|| DbError::NotFound(task_id.to_owned()))?; let mut blocks: Vec = serde_json::from_str(&task.blocks).unwrap_or_default(); blocks.push(blocked_task_id.to_owned()); @@ -163,12 +194,18 @@ impl ITeamRepository for MockTeamRepo { Ok(()) } - async fn remove_from_blocked_by(&self, task_id: &str, unblocked_task_id: &str) -> Result<(), DbError> { + async fn remove_from_blocked_by( + &self, + _user_id: &str, + team_id: &str, + task_id: &str, + unblocked_task_id: &str, + ) -> Result<(), DbError> { let mut state = self.state.lock().unwrap(); let task = state .tasks .iter_mut() - .find(|t| t.id == task_id) + .find(|t| t.team_id == team_id && t.id == task_id) .ok_or_else(|| DbError::NotFound(task_id.to_owned()))?; let mut blocked_by: Vec = serde_json::from_str(&task.blocked_by).unwrap_or_default(); blocked_by.retain(|id| id != unblocked_task_id); @@ -176,7 +213,7 @@ impl ITeamRepository for MockTeamRepo { Ok(()) } - async fn delete_tasks_by_team(&self, team_id: &str) -> Result<(), DbError> { + async fn delete_tasks_by_team(&self, _user_id: &str, team_id: &str) -> Result<(), DbError> { self.state.lock().unwrap().tasks.retain(|t| t.team_id != team_id); Ok(()) } diff --git a/crates/aionui-team/tests/e2e_smoke.rs b/crates/aionui-team/tests/e2e_smoke.rs index ddfc18a8e..a18b3f933 100644 --- a/crates/aionui-team/tests/e2e_smoke.rs +++ b/crates/aionui-team/tests/e2e_smoke.rs @@ -101,6 +101,7 @@ async fn setup_team_with_lead() -> SmokeEnv { ]; let scheduler = Arc::new(TeammateManager::new( team_id.clone(), + "user-1".into(), &agents, mailbox.clone(), task_board.clone(), diff --git a/crates/aionui-team/tests/mailbox_integration.rs b/crates/aionui-team/tests/mailbox_integration.rs index 967f45d33..a4290e216 100644 --- a/crates/aionui-team/tests/mailbox_integration.rs +++ b/crates/aionui-team/tests/mailbox_integration.rs @@ -10,13 +10,31 @@ use std::sync::Arc; +use aionui_common::now_ms; use aionui_db::{ITeamRepository, SqliteTeamRepository, init_database_memory}; use aionui_team::{Mailbox, MailboxMessageType}; async fn setup() -> (Mailbox, aionui_db::Database) { let db = init_database_memory().await.unwrap(); - let repo = Arc::new(SqliteTeamRepository::new(db.pool().clone())) as Arc; - (Mailbox::new(repo), db) + let repo = Arc::new(SqliteTeamRepository::new(db.pool().clone())); + for team_id in ["t1", "t2"] { + repo.create_team(&aionui_db::models::TeamRow { + id: team_id.to_owned(), + user_id: "system_default_user".to_owned(), + name: team_id.to_owned(), + workspace: String::new(), + workspace_mode: "shared".to_owned(), + agents: "[]".to_owned(), + lead_agent_id: None, + session_mode: None, + agents_version: "1.0.1".to_owned(), + created_at: now_ms(), + updated_at: now_ms(), + }) + .await + .unwrap(); + } + (Mailbox::new(repo as Arc), db) } // -- MW: Write messages ------------------------------------------------------- diff --git a/crates/aionui-team/tests/mcp_server_integration.rs b/crates/aionui-team/tests/mcp_server_integration.rs index 2e6892ef7..d210542b4 100644 --- a/crates/aionui-team/tests/mcp_server_integration.rs +++ b/crates/aionui-team/tests/mcp_server_integration.rs @@ -83,6 +83,7 @@ async fn setup_with_prompt_dump(prompt_dump: Option) -> Te let agents = make_agents(); let scheduler = Arc::new(TeammateManager::new( "team-1".into(), + "user-1".into(), &agents, mailbox, task_board, diff --git a/crates/aionui-team/tests/prompts_events_integration.rs b/crates/aionui-team/tests/prompts_events_integration.rs index d3f68a61b..59671df9f 100644 --- a/crates/aionui-team/tests/prompts_events_integration.rs +++ b/crates/aionui-team/tests/prompts_events_integration.rs @@ -127,6 +127,7 @@ fn visibility_policy_user_message_has_explicit_decisions() { #[test] fn projection_request_for_teammate_mirror_uses_stable_mailbox_dedupe_key() { let req = TeamProjectionRequest::teammate_visible( + "user-1", "team-1", "lead-1", "conv-lead", @@ -147,6 +148,7 @@ fn projection_request_for_teammate_mirror_uses_stable_mailbox_dedupe_key() { #[test] fn projection_request_for_team_system_message_uses_stable_mailbox_dedupe_key() { let req = TeamProjectionRequest::team_system_visible( + "user-1", "team-1", "lead-1", "conv-lead", @@ -168,6 +170,7 @@ async fn projection_inserts_user_visible_bubble_with_stripped_system_notes() { let bc = Arc::new(RecordingBroadcaster::new()); let projection = TeamMessageProjection::new(store.clone(), bc.clone()); let req = TeamProjectionRequest { + user_id: "user-1".into(), team_id: "team-1".into(), slot_id: "lead-1".into(), conversation_id: "conv-lead".into(), @@ -197,6 +200,7 @@ async fn projection_dedupes_teammate_mirror_and_broadcasts_persisted_msg_id() { let bc = Arc::new(RecordingBroadcaster::new()); let projection = TeamMessageProjection::new(store.clone(), bc.clone()); let req = TeamProjectionRequest::teammate_visible( + "user-1", "team-1", "lead-1", "conv-lead", @@ -226,6 +230,7 @@ async fn projection_dedupes_teammate_mirror_and_broadcasts_persisted_msg_id() { let events = bc.events(); assert_eq!(events.len(), 1, "duplicate projection must not re-broadcast"); assert_eq!(events[0].name, "team.teammateMessage"); + assert_eq!(events[0].data["user_id"], "user-1"); assert_eq!(events[0].data["conversation_id"], "conv-lead"); assert_eq!(events[0].data["msg_id"], first_msg_id); assert_eq!(events[0].data["from_slot_id"], "worker-1"); @@ -238,6 +243,7 @@ async fn projection_inserts_team_system_bubble_and_broadcasts_it_like_existing_m let bc = Arc::new(RecordingBroadcaster::new()); let projection = TeamMessageProjection::new(store.clone(), bc.clone()); let req = TeamProjectionRequest::team_system_visible( + "user-1", "team-1", "lead-1", "conv-lead", @@ -522,7 +528,7 @@ async fn we1_agent_status_change_event() { make_agent("lead-1", "Lead", TeammateRole::Lead), make_agent("w1", "Worker", TeammateRole::Teammate), ]; - let mgr = TeammateManager::new("t1".into(), &agents, mailbox, task_board, bc.clone()); + let mgr = TeammateManager::new("t1".into(), "user-1".into(), &agents, mailbox, task_board, bc.clone()); mgr.set_status("w1", TeammateStatus::Working).await.unwrap(); @@ -545,7 +551,7 @@ async fn we2_agent_spawned_event() { let task_board = Arc::new(TaskBoard::new(repo)); let bc = Arc::new(RecordingBroadcaster::new()); let agents = vec![make_agent("lead-1", "Lead", TeammateRole::Lead)]; - let mgr = TeammateManager::new("t1".into(), &agents, mailbox, task_board, bc.clone()); + let mgr = TeammateManager::new("t1".into(), "user-1".into(), &agents, mailbox, task_board, bc.clone()); let new_agent = make_agent("w2", "NewWorker", TeammateRole::Teammate); mgr.add_agent(&new_agent).await; @@ -575,7 +581,7 @@ async fn we3_agent_removed_event() { make_agent("lead-1", "Lead", TeammateRole::Lead), make_agent("w1", "Worker", TeammateRole::Teammate), ]; - let mgr = TeammateManager::new("t1".into(), &agents, mailbox, task_board, bc.clone()); + let mgr = TeammateManager::new("t1".into(), "user-1".into(), &agents, mailbox, task_board, bc.clone()); mgr.remove_agent("w1").await.unwrap(); @@ -603,7 +609,7 @@ async fn we4_agent_renamed_event() { make_agent("lead-1", "Lead", TeammateRole::Lead), make_agent("w1", "Worker", TeammateRole::Teammate), ]; - let mgr = TeammateManager::new("t1".into(), &agents, mailbox, task_board, bc.clone()); + let mgr = TeammateManager::new("t1".into(), "user-1".into(), &agents, mailbox, task_board, bc.clone()); mgr.rename_agent("w1", "SuperWorker").await.unwrap(); @@ -625,7 +631,7 @@ async fn we4_agent_renamed_event() { #[test] fn event_emitter_uses_typed_payloads() { let bc = Arc::new(RecordingBroadcaster::new()); - let emitter = TeamEventEmitter::new("team-x".into(), bc.clone()); + let emitter = TeamEventEmitter::new("team-x".into(), "user-1".into(), bc.clone()); let agent = make_agent("s1", "A", TeammateRole::Teammate); emitter.broadcast_agent_status("s1", TeammateStatus::Thinking); diff --git a/crates/aionui-team/tests/scheduler_integration.rs b/crates/aionui-team/tests/scheduler_integration.rs index 16ebec258..be298e069 100644 --- a/crates/aionui-team/tests/scheduler_integration.rs +++ b/crates/aionui-team/tests/scheduler_integration.rs @@ -71,6 +71,7 @@ fn setup_team(agents: &[TeamAgent]) -> TestHarness { let broadcaster = Arc::new(RecordingBroadcaster::new()); let mgr = TeammateManager::new( "team-1".into(), + "user-1".into(), agents, mailbox.clone(), task_board.clone(), diff --git a/crates/aionui-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index 3e33201da..71d990fbc 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -99,19 +99,23 @@ impl MockConversationRepo { #[async_trait::async_trait] impl IConversationRepository for MockConversationRepo { - async fn get(&self, id: &str) -> Result, DbError> { + async fn get(&self, user_id: &str, id: &str) -> Result, DbError> { let convs = self.conversations.lock().unwrap(); - Ok(convs.iter().find(|c| c.id == id).cloned()) + Ok(convs.iter().find(|c| c.user_id == user_id && c.id == id).cloned()) + } + async fn owner_user_id(&self, id: &str) -> Result, DbError> { + let convs = self.conversations.lock().unwrap(); + Ok(convs.iter().find(|c| c.id == id).map(|c| c.user_id.clone())) } async fn create(&self, row: &ConversationRow) -> Result<(), DbError> { self.conversations.lock().unwrap().push(row.clone()); Ok(()) } - async fn update(&self, id: &str, updates: &ConversationRowUpdate) -> Result<(), DbError> { + async fn update(&self, user_id: &str, id: &str, updates: &ConversationRowUpdate) -> Result<(), DbError> { let mut convs = self.conversations.lock().unwrap(); let conv = convs .iter_mut() - .find(|c| c.id == id) + .find(|c| c.user_id == user_id && c.id == id) .ok_or_else(|| DbError::NotFound(id.to_owned()))?; if let Some(ref extra) = updates.extra { conv.extra = extra.clone(); @@ -130,8 +134,11 @@ impl IConversationRepository for MockConversationRepo { } Ok(()) } - async fn delete(&self, id: &str) -> Result<(), DbError> { - self.conversations.lock().unwrap().retain(|c| c.id != id); + async fn delete(&self, user_id: &str, id: &str) -> Result<(), DbError> { + self.conversations + .lock() + .unwrap() + .retain(|c| c.user_id != user_id || c.id != id); Ok(()) } async fn list_paginated( @@ -162,6 +169,7 @@ impl IConversationRepository for MockConversationRepo { } async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &MessagePageParams, ) -> Result { @@ -171,18 +179,25 @@ impl IConversationRepository for MockConversationRepo { has_more_after: false, }) } - async fn insert_message(&self, message: &MessageRow) -> Result<(), DbError> { + async fn insert_message(&self, _user_id: &str, message: &MessageRow) -> Result<(), DbError> { self.messages.lock().unwrap().push(message.clone()); Ok(()) } - async fn update_message(&self, _id: &str, _updates: &MessageRowUpdate) -> Result<(), DbError> { + async fn update_message( + &self, + _user_id: &str, + _conversation_id: &str, + _id: &str, + _updates: &MessageRowUpdate, + ) -> Result<(), DbError> { Ok(()) } - async fn delete_messages_by_conversation(&self, _conv_id: &str) -> Result<(), DbError> { + async fn delete_messages_by_conversation(&self, _user_id: &str, _conv_id: &str) -> Result<(), DbError> { Ok(()) } async fn get_message_by_msg_id( &self, + _user_id: &str, conv_id: &str, msg_id: &str, msg_type: &str, @@ -466,8 +481,14 @@ impl TeamConversationProvisioningPort for FakeConversationPorts { target.insert(key.clone(), value.clone()); } } + let user_id = self + .repo + .owner_user_id(conversation_id) + .await? + .ok_or_else(|| aionui_team::TeamError::AgentNotFound(conversation_id.to_owned()))?; self.repo .update( + &user_id, conversation_id, &ConversationRowUpdate { name: None, @@ -526,14 +547,9 @@ impl TeamConversationProvisioningPort for FakeConversationPorts { conversation_id: &str, task_manager: &Arc, ) -> Result<(), aionui_team::TeamError> { - let row = self - .repo - .get(conversation_id) - .await? - .filter(|row| row.user_id == user_id) - .ok_or_else(|| { - aionui_team::TeamError::InvalidRequest(format!("conversation not found: {conversation_id}")) - })?; + let row = self.repo.get(user_id, conversation_id).await?.ok_or_else(|| { + aionui_team::TeamError::InvalidRequest(format!("conversation not found: {conversation_id}")) + })?; let extra: serde_json::Value = serde_json::from_str(&row.extra)?; let team = aionui_api_types::TeamSessionBinding::from_extra_value(&extra)?; let config: AcpBuildExtra = serde_json::from_value(extra.clone()).unwrap_or_default(); @@ -594,7 +610,7 @@ impl TeamConversationProvisioningPort for FakeConversationPorts { _user_id: &str, conversation_id: &str, ) -> Result<(), aionui_team::TeamError> { - self.repo.delete(conversation_id).await?; + self.repo.delete(_user_id, conversation_id).await?; Ok(()) } } @@ -611,14 +627,22 @@ impl TeamProjectionMessageStore for FakeConversationPorts { msg_id: &str, msg_type: &str, ) -> Result, aionui_team::TeamError> { + let Some(user_id) = self.repo.owner_user_id(conversation_id).await? else { + return Ok(None); + }; Ok(self .repo - .get_message_by_msg_id(conversation_id, msg_id, msg_type) + .get_message_by_msg_id(&user_id, conversation_id, msg_id, msg_type) .await?) } async fn insert_projected_message(&self, row: &MessageRow) -> Result<(), aionui_team::TeamError> { - self.repo.insert_message(row).await?; + let user_id = self + .repo + .owner_user_id(&row.conversation_id) + .await? + .ok_or_else(|| aionui_team::TeamError::InvalidRequest("conversation not found".into()))?; + self.repo.insert_message(&user_id, row).await?; Ok(()) } } @@ -629,7 +653,10 @@ impl TeamConversationLookupPort for FakeConversationPorts { &self, conversation_id: &str, ) -> Result, aionui_team::TeamError> { - let Some(row) = self.repo.get(conversation_id).await? else { + let Some(user_id) = self.repo.owner_user_id(conversation_id).await? else { + return Ok(None); + }; + let Some(row) = self.repo.get(&user_id, conversation_id).await? else { return Ok(None); }; let extra: serde_json::Value = serde_json::from_str(&row.extra).unwrap_or(serde_json::Value::Null); @@ -722,7 +749,7 @@ impl ITeamRepository for FullMockTeamRepo { self.teams.lock().unwrap().push(row.clone()); Ok(()) } - async fn list_teams(&self) -> Result, DbError> { + async fn list_teams_for_restore(&self) -> Result, DbError> { Ok(self.teams.lock().unwrap().clone()) } async fn list_teams_by_user(&self, user_id: &str) -> Result, DbError> { @@ -735,10 +762,19 @@ impl ITeamRepository for FullMockTeamRepo { .cloned() .collect()) } - async fn get_team(&self, id: &str) -> Result, DbError> { + async fn get_team(&self, user_id: &str, id: &str) -> Result, DbError> { + Ok(self + .teams + .lock() + .unwrap() + .iter() + .find(|t| t.user_id == user_id && t.id == id) + .cloned()) + } + async fn get_team_for_restore(&self, id: &str) -> Result, DbError> { Ok(self.teams.lock().unwrap().iter().find(|t| t.id == id).cloned()) } - async fn update_team(&self, id: &str, params: &aionui_db::UpdateTeamParams) -> Result<(), DbError> { + async fn update_team(&self, user_id: &str, id: &str, params: &aionui_db::UpdateTeamParams) -> Result<(), DbError> { if params.workspace.is_some() && *self.fail_workspace_update.lock().unwrap() { return Err(DbError::Init("forced workspace writeback failure".into())); } @@ -748,7 +784,7 @@ impl ITeamRepository for FullMockTeamRepo { let mut teams = self.teams.lock().unwrap(); let team = teams .iter_mut() - .find(|t| t.id == id) + .find(|t| t.user_id == user_id && t.id == id) .ok_or_else(|| DbError::NotFound(id.to_owned()))?; if let Some(ref name) = params.name { team.name = name.clone(); @@ -765,70 +801,99 @@ impl ITeamRepository for FullMockTeamRepo { team.updated_at = aionui_common::now_ms(); Ok(()) } - async fn delete_team(&self, id: &str) -> Result<(), DbError> { - self.teams.lock().unwrap().retain(|t| t.id != id); + async fn delete_team(&self, user_id: &str, id: &str) -> Result<(), DbError> { + self.teams + .lock() + .unwrap() + .retain(|t| t.user_id != user_id || t.id != id); Ok(()) } - async fn write_message(&self, row: &aionui_db::models::MailboxMessageRow) -> Result<(), DbError> { + async fn write_message(&self, user_id: &str, row: &aionui_db::models::MailboxMessageRow) -> Result<(), DbError> { if *self.fail_message_writes.lock().unwrap() { return Err(DbError::Init("forced mailbox write failure".into())); } - self.inner.write_message(row).await + self.inner.write_message(user_id, row).await } async fn read_unread_and_mark( &self, + user_id: &str, team_id: &str, to_agent_id: &str, ) -> Result, DbError> { - self.inner.read_unread_and_mark(team_id, to_agent_id).await + self.inner.read_unread_and_mark(user_id, team_id, to_agent_id).await } async fn peek_unread( &self, + user_id: &str, team_id: &str, to_agent_id: &str, ) -> Result, DbError> { - self.inner.peek_unread(team_id, to_agent_id).await + self.inner.peek_unread(user_id, team_id, to_agent_id).await } - async fn mark_read_batch(&self, ids: &[String]) -> Result<(), DbError> { - self.inner.mark_read_batch(ids).await + async fn mark_read_batch(&self, user_id: &str, team_id: &str, ids: &[String]) -> Result<(), DbError> { + self.inner.mark_read_batch(user_id, team_id, ids).await } async fn get_history( &self, + user_id: &str, team_id: &str, to_agent_id: &str, limit: Option, ) -> Result, DbError> { - self.inner.get_history(team_id, to_agent_id, limit).await + self.inner.get_history(user_id, team_id, to_agent_id, limit).await } - async fn delete_mailbox_by_team(&self, team_id: &str) -> Result<(), DbError> { - self.inner.delete_mailbox_by_team(team_id).await + async fn delete_mailbox_by_team(&self, user_id: &str, team_id: &str) -> Result<(), DbError> { + self.inner.delete_mailbox_by_team(user_id, team_id).await } - async fn create_task(&self, row: &aionui_db::models::TeamTaskRow) -> Result<(), DbError> { - self.inner.create_task(row).await + async fn create_task(&self, user_id: &str, row: &aionui_db::models::TeamTaskRow) -> Result<(), DbError> { + self.inner.create_task(user_id, row).await } async fn find_task_by_id( &self, + user_id: &str, team_id: &str, task_id: &str, ) -> Result, DbError> { - self.inner.find_task_by_id(team_id, task_id).await + self.inner.find_task_by_id(user_id, team_id, task_id).await } - async fn update_task(&self, task_id: &str, params: &aionui_db::UpdateTaskParams) -> Result<(), DbError> { - self.inner.update_task(task_id, params).await + async fn update_task( + &self, + user_id: &str, + team_id: &str, + task_id: &str, + params: &aionui_db::UpdateTaskParams, + ) -> Result<(), DbError> { + self.inner.update_task(user_id, team_id, task_id, params).await } - async fn list_tasks(&self, team_id: &str) -> Result, DbError> { - self.inner.list_tasks(team_id).await + async fn list_tasks(&self, user_id: &str, team_id: &str) -> Result, DbError> { + self.inner.list_tasks(user_id, team_id).await } - async fn append_to_blocks(&self, task_id: &str, blocked_task_id: &str) -> Result<(), DbError> { - self.inner.append_to_blocks(task_id, blocked_task_id).await + async fn append_to_blocks( + &self, + user_id: &str, + team_id: &str, + task_id: &str, + blocked_task_id: &str, + ) -> Result<(), DbError> { + self.inner + .append_to_blocks(user_id, team_id, task_id, blocked_task_id) + .await } - async fn remove_from_blocked_by(&self, task_id: &str, unblocked_task_id: &str) -> Result<(), DbError> { - self.inner.remove_from_blocked_by(task_id, unblocked_task_id).await + async fn remove_from_blocked_by( + &self, + user_id: &str, + team_id: &str, + task_id: &str, + unblocked_task_id: &str, + ) -> Result<(), DbError> { + self.inner + .remove_from_blocked_by(user_id, team_id, task_id, unblocked_task_id) + .await } - async fn delete_tasks_by_team(&self, team_id: &str) -> Result<(), DbError> { - self.inner.delete_tasks_by_team(team_id).await + async fn delete_tasks_by_team(&self, user_id: &str, team_id: &str) -> Result<(), DbError> { + self.inner.delete_tasks_by_team(user_id, team_id).await } } @@ -866,9 +931,15 @@ impl IAgentMetadataRepository for StubAgentMetadataRepo { async fn list_all(&self) -> Result, DbError> { Ok(self.rows_by_id.values().cloned().collect()) } + async fn list_all_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list_all().await + } async fn get(&self, id: &str) -> Result, DbError> { Ok(self.rows_by_id.get(id).cloned()) } + async fn get_for_user(&self, _user_id: &str, id: &str) -> Result, DbError> { + self.get(id).await + } async fn find_by_source_and_name( &self, agent_source: &str, @@ -880,12 +951,34 @@ impl IAgentMetadataRepository for StubAgentMetadataRepo { .find(|row| row.agent_source == agent_source && row.name == name) .cloned()) } + async fn find_by_source_and_name_for_user( + &self, + _user_id: &str, + agent_source: &str, + name: &str, + ) -> Result, DbError> { + self.find_by_source_and_name(agent_source, name).await + } async fn find_builtin_by_backend(&self, backend: &str) -> Result, DbError> { Ok(self.builtin_by_backend.get(backend).cloned()) } + async fn find_builtin_by_backend_for_user( + &self, + _user_id: &str, + backend: &str, + ) -> Result, DbError> { + self.find_builtin_by_backend(backend).await + } async fn upsert(&self, _params: &UpsertAgentMetadataParams<'_>) -> Result { Err(DbError::Init("stub".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAgentMetadataParams<'_>, + ) -> Result { + self.upsert(params).await + } async fn apply_handshake( &self, _id: &str, @@ -893,6 +986,14 @@ impl IAgentMetadataRepository for StubAgentMetadataRepo { ) -> Result, DbError> { Ok(None) } + async fn apply_handshake_for_user( + &self, + _user_id: &str, + id: &str, + params: &UpdateAgentHandshakeParams<'_>, + ) -> Result, DbError> { + self.apply_handshake(id, params).await + } async fn update_availability_snapshot( &self, _id: &str, @@ -900,6 +1001,14 @@ impl IAgentMetadataRepository for StubAgentMetadataRepo { ) -> Result, DbError> { Ok(None) } + async fn update_availability_snapshot_for_user( + &self, + _user_id: &str, + id: &str, + params: &UpdateAgentAvailabilitySnapshotParams<'_>, + ) -> Result, DbError> { + self.update_availability_snapshot(id, params).await + } async fn update_agent_overrides( &self, _id: &str, @@ -908,12 +1017,27 @@ impl IAgentMetadataRepository for StubAgentMetadataRepo { ) -> Result<(), DbError> { Ok(()) } + async fn update_agent_overrides_for_user( + &self, + _user_id: &str, + id: &str, + command_override: Option<&str>, + env_override: Option<&str>, + ) -> Result<(), DbError> { + self.update_agent_overrides(id, command_override, env_override).await + } async fn set_enabled(&self, _id: &str, _enabled: bool) -> Result { Ok(false) } + async fn set_enabled_for_user(&self, _user_id: &str, id: &str, enabled: bool) -> Result { + self.set_enabled(id, enabled).await + } async fn delete(&self, _id: &str) -> Result { Ok(false) } + async fn delete_for_user(&self, _user_id: &str, id: &str) -> Result { + self.delete(id).await + } } // --------------------------------------------------------------------------- @@ -1251,7 +1375,10 @@ struct EmptyTeamAssistantCatalog; #[async_trait::async_trait] impl TeamAssistantCatalogPort for EmptyTeamAssistantCatalog { - async fn list_team_selectable_assistants(&self) -> Result, TeamError> { + async fn list_team_selectable_assistants( + &self, + _user_id: &str, + ) -> Result, TeamError> { Ok(Vec::new()) } } @@ -1264,13 +1391,19 @@ struct TestTeamAssistantCatalog { #[async_trait::async_trait] impl TeamAssistantCatalogPort for TestTeamAssistantCatalog { - async fn list_team_selectable_assistants(&self) -> Result, TeamError> { - let agent_rows = self.agent_metadata_repo.list_all().await?; - let definitions = self.assistant_definition_repo.list().await?; + async fn list_team_selectable_assistants( + &self, + user_id: &str, + ) -> Result, TeamError> { + let agent_rows = self.agent_metadata_repo.list_all_for_user(user_id).await?; + let definitions = self.assistant_definition_repo.list_for_user(user_id).await?; let mut result = Vec::new(); for definition in definitions { - let overlay = self.assistant_overlay_repo.get(&definition.id).await?; + let overlay = self + .assistant_overlay_repo + .get_for_user(user_id, &definition.id) + .await?; if overlay.as_ref().is_some_and(|row| !row.enabled) { continue; } @@ -1299,10 +1432,10 @@ struct EmptyProviderRepo; #[async_trait::async_trait] impl IProviderRepository for EmptyProviderRepo { - async fn list(&self) -> Result, DbError> { + async fn list(&self, _user_id: &str) -> Result, DbError> { Ok(vec![]) } - async fn find_by_id(&self, _id: &str) -> Result, DbError> { + async fn find_by_id(&self, _user_id: &str, _id: &str) -> Result, DbError> { Ok(None) } async fn create( @@ -1313,12 +1446,13 @@ impl IProviderRepository for EmptyProviderRepo { } async fn update( &self, + _user_id: &str, _id: &str, _params: aionui_db::UpdateProviderParams<'_>, ) -> Result { Err(DbError::NotFound("not implemented".into())) } - async fn delete(&self, _id: &str) -> Result<(), DbError> { + async fn delete(&self, _user_id: &str, _id: &str) -> Result<(), DbError> { Err(DbError::NotFound("not implemented".into())) } } @@ -1331,14 +1465,46 @@ impl IAssistantDefinitionRepository for EmptyAssistantDefinitionRepo { Ok(vec![]) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + + async fn list_including_deleted_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + async fn get_by_assistant_id(&self, _assistant_id: &str) -> Result, DbError> { Ok(None) } + async fn get_by_assistant_id_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + + async fn get_by_assistant_id_including_deleted_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + async fn get_by_id(&self, _definition_id: &str) -> Result, DbError> { Ok(None) } + async fn get_by_id_for_user( + &self, + _user_id: &str, + definition_id: &str, + ) -> Result, DbError> { + self.get_by_id(definition_id).await + } + async fn get_by_source_ref( &self, _source: &str, @@ -1347,13 +1513,48 @@ impl IAssistantDefinitionRepository for EmptyAssistantDefinitionRepo { Ok(None) } + async fn get_by_source_ref_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + + async fn get_by_source_ref_including_deleted_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + async fn upsert(&self, _params: &UpsertAssistantDefinitionParams<'_>) -> Result { Err(DbError::Init("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn soft_delete(&self, _definition_id: &str, _deleted_at: i64) -> Result { Ok(false) } + + async fn soft_delete_for_user( + &self, + _user_id: &str, + definition_id: &str, + deleted_at: i64, + ) -> Result { + self.soft_delete(definition_id, deleted_at).await + } } struct EmptyAssistantOverlayRepo; @@ -1364,17 +1565,37 @@ impl IAssistantOverlayRepository for EmptyAssistantOverlayRepo { Ok(None) } + async fn get_for_user(&self, _user_id: &str, definition_id: &str) -> Result, DbError> { + self.get(definition_id).await + } + async fn list(&self) -> Result, DbError> { Ok(vec![]) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + async fn upsert(&self, _params: &UpsertAssistantOverlayParams<'_>) -> Result { Err(DbError::Init("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantOverlayParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn delete(&self, _definition_id: &str) -> Result { Ok(false) } + + async fn delete_for_user(&self, _user_id: &str, definition_id: &str) -> Result { + self.delete(definition_id).await + } } struct SingleAssistantDefinitionRepo { @@ -1387,14 +1608,46 @@ impl IAssistantDefinitionRepository for SingleAssistantDefinitionRepo { Ok(vec![self.row.clone()]) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + + async fn list_including_deleted_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + async fn get_by_assistant_id(&self, assistant_id: &str) -> Result, DbError> { Ok((self.row.assistant_id == assistant_id).then_some(self.row.clone())) } + async fn get_by_assistant_id_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + + async fn get_by_assistant_id_including_deleted_for_user( + &self, + _user_id: &str, + assistant_id: &str, + ) -> Result, DbError> { + self.get_by_assistant_id(assistant_id).await + } + async fn get_by_id(&self, definition_id: &str) -> Result, DbError> { Ok((self.row.id == definition_id).then_some(self.row.clone())) } + async fn get_by_id_for_user( + &self, + _user_id: &str, + definition_id: &str, + ) -> Result, DbError> { + self.get_by_id(definition_id).await + } + async fn get_by_source_ref( &self, _source: &str, @@ -1403,13 +1656,48 @@ impl IAssistantDefinitionRepository for SingleAssistantDefinitionRepo { Ok(None) } + async fn get_by_source_ref_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + + async fn get_by_source_ref_including_deleted_for_user( + &self, + _user_id: &str, + source: &str, + source_ref: &str, + ) -> Result, DbError> { + self.get_by_source_ref(source, source_ref).await + } + async fn upsert(&self, _params: &UpsertAssistantDefinitionParams<'_>) -> Result { Err(DbError::Init("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn soft_delete(&self, _definition_id: &str, _deleted_at: i64) -> Result { Ok(false) } + + async fn soft_delete_for_user( + &self, + _user_id: &str, + definition_id: &str, + deleted_at: i64, + ) -> Result { + self.soft_delete(definition_id, deleted_at).await + } } struct SingleAssistantOverlayRepo { @@ -1422,17 +1710,37 @@ impl IAssistantOverlayRepository for SingleAssistantOverlayRepo { Ok((self.row.assistant_definition_id == definition_id).then_some(self.row.clone())) } + async fn get_for_user(&self, _user_id: &str, definition_id: &str) -> Result, DbError> { + self.get(definition_id).await + } + async fn list(&self) -> Result, DbError> { Ok(vec![self.row.clone()]) } + async fn list_for_user(&self, _user_id: &str) -> Result, DbError> { + self.list().await + } + async fn upsert(&self, _params: &UpsertAssistantOverlayParams<'_>) -> Result { Err(DbError::Init("not implemented".into())) } + async fn upsert_for_user( + &self, + _user_id: &str, + params: &UpsertAssistantOverlayParams<'_>, + ) -> Result { + self.upsert(params).await + } + async fn delete(&self, _definition_id: &str) -> Result { Ok(false) } + + async fn delete_for_user(&self, _user_id: &str, definition_id: &str) -> Result { + self.delete(definition_id).await + } } fn setup_with_factory(factory: AgentFactory) -> (Arc, Arc) { @@ -1663,18 +1971,21 @@ async fn recovery_creates_background_intents_without_restoring_old_memory_run() .expect("clear existing session"); team_repo - .write_message(&aionui_db::models::MailboxMessageRow { - id: "mailbox-orphan-1".into(), - team_id: created.id.clone(), - to_agent_id: lead_slot_id.clone(), - from_agent_id: "worker-or-user".into(), - msg_type: "message".into(), - content: "orphan backlog".into(), - summary: None, - files: None, - read: false, - created_at: aionui_common::now_ms(), - }) + .write_message( + "user1", + &aionui_db::models::MailboxMessageRow { + id: "mailbox-orphan-1".into(), + team_id: created.id.clone(), + to_agent_id: lead_slot_id.clone(), + from_agent_id: "worker-or-user".into(), + msg_type: "message".into(), + content: "orphan backlog".into(), + summary: None, + files: None, + read: false, + created_at: aionui_common::now_ms(), + }, + ) .await .expect("seed orphan mailbox"); @@ -1781,18 +2092,21 @@ async fn ensure_session_does_not_run_self_message_only_recovery_turn() { .expect("clear existing session"); team_repo - .write_message(&aionui_db::models::MailboxMessageRow { - id: "mailbox-self-1".into(), - team_id: created.id.clone(), - to_agent_id: lead_slot_id.clone(), - from_agent_id: lead_slot_id, - msg_type: "message".into(), - content: "self backlog".into(), - summary: None, - files: None, - read: false, - created_at: aionui_common::now_ms(), - }) + .write_message( + "user1", + &aionui_db::models::MailboxMessageRow { + id: "mailbox-self-1".into(), + team_id: created.id.clone(), + to_agent_id: lead_slot_id.clone(), + from_agent_id: lead_slot_id, + msg_type: "message".into(), + content: "self backlog".into(), + summary: None, + files: None, + read: false, + created_at: aionui_common::now_ms(), + }, + ) .await .expect("seed self mailbox"); @@ -1891,6 +2205,7 @@ fn setup_with_factory_recording_broadcaster_and_conversation_repo(factory: Agent fn make_agent_metadata_row(id: &str, backend: &str, icon: &str) -> AgentMetadataRow { AgentMetadataRow { id: id.to_owned(), + user_id: None, icon: Some(icon.to_owned()), name: backend.to_owned(), name_i18n: None, @@ -2183,7 +2498,7 @@ async fn renew_active_lease_rejects_team_owned_by_other_user() { .await .unwrap_err(); - assert!(matches!(err, TeamError::Forbidden(_))); + assert!(matches!(err, TeamError::TeamNotFound(_))); for agent in &created.assistants { assert!(!active_leases.is_active(&agent.conversation_id)); } @@ -2191,6 +2506,7 @@ async fn renew_active_lease_rejects_team_owned_by_other_user() { async fn force_team_workspace(repo: &Arc, team_id: &str, workspace: &str) { repo.update_team( + "user1", team_id, &aionui_db::UpdateTeamParams { workspace: Some(workspace.to_owned()), @@ -2465,7 +2781,7 @@ async fn tc_create_team_carries_assistant_identity_into_lead_conversation_extra( .unwrap(); let row = conv_repo - .get(&resp.assistants[0].conversation_id) + .get("user1", &resp.assistants[0].conversation_id) .await .unwrap() .expect("lead conversation row"); @@ -3333,7 +3649,7 @@ async fn tg3_get_team_rejects_cross_user_access() { let result = svc.get_team("user2", &created.id).await; - assert!(matches!(result, Err(aionui_team::TeamError::Forbidden(_)))); + assert!(matches!(result, Err(aionui_team::TeamError::TeamNotFound(_)))); } // -- Delete team -------------------------------------------------------------- @@ -3411,7 +3727,7 @@ async fn tr5_rename_team_rejects_cross_user_access() { let result = svc.rename_team("user2", &created.id, "Nope").await; - assert!(matches!(result, Err(aionui_team::TeamError::Forbidden(_)))); + assert!(matches!(result, Err(aionui_team::TeamError::TeamNotFound(_)))); } // =========================================================================== @@ -3804,7 +4120,10 @@ async fn manual_add_agent_attach_failure_marks_slot_error_without_leader_notice( ); let lead_slot_id = created.leader_assistant_id.as_deref().expect("leader slot"); - let leader_messages = team_repo.get_history(&created.id, lead_slot_id, None).await.unwrap(); + let leader_messages = team_repo + .get_history("user1", &created.id, lead_slot_id, None) + .await + .unwrap(); assert!( !leader_messages .iter() @@ -4139,7 +4458,7 @@ async fn remove_during_attach_cancels_work_and_rejects_late_ready() { assert!(task_manager.get_task(&added.conversation_id).is_none()); assert!( - conv_repo.get(&added.conversation_id).await.unwrap().is_none(), + conv_repo.get("user1", &added.conversation_id).await.unwrap().is_none(), "removed attaching member conversation must be deleted" ); assert!( @@ -4539,7 +4858,7 @@ async fn provisioning_resolves_acp_backend_from_agent_metadata() { .unwrap(); let row = conv_repo - .get(&created.assistants[0].conversation_id) + .get("user1", &created.assistants[0].conversation_id) .await .unwrap() .expect("conversation row"); @@ -4631,7 +4950,7 @@ async fn membership_persist_failure_does_not_delete_the_conversation() { assert!(error.to_string().contains("forced agent update failure")); assert!( - conv_repo.get(&worker.conversation_id).await.unwrap().is_some(), + conv_repo.get("user1", &worker.conversation_id).await.unwrap().is_some(), "conversation deletion must happen only after membership persistence" ); assert!( @@ -4678,7 +4997,7 @@ async fn remove_tolerates_current_session_already_missing_the_slot() { .await .expect("post-persistence runtime cleanup must be idempotent"); - assert!(conv_repo.get(&worker.conversation_id).await.unwrap().is_none()); + assert!(conv_repo.get("user1", &worker.conversation_id).await.unwrap().is_none()); assert!( svc.get_team("user1", &created.id) .await @@ -5177,7 +5496,7 @@ async fn es4_ensure_session_rejects_cross_user_access() { let result = svc.ensure_session("user2", &created.id).await; - assert!(matches!(result, Err(aionui_team::TeamError::Forbidden(_)))); + assert!(matches!(result, Err(aionui_team::TeamError::TeamNotFound(_)))); } // -- W5-D31b-2: team.sessionStatusChanged service-layer broadcasts ----------- @@ -5191,19 +5510,10 @@ async fn d31b2_ensure_session_broadcasts_failed_loading_team_for_missing_team() let err = svc.ensure_session("user1", "nonexistent-team-xyz").await.unwrap_err(); assert!(matches!(err, aionui_team::TeamError::TeamNotFound(_))); - let failed = recorder - .events_by_name("team.sessionStatusChanged") - .into_iter() - .find(|e| { - e.data.get("status").and_then(|v| v.as_str()) == Some("failed") - && e.data.get("phase").and_then(|v| v.as_str()) == Some("loading_team") - }) - .expect("failed/loading_team broadcast expected"); - assert_eq!( - failed.data.get("team_id").and_then(|v| v.as_str()), - Some("nonexistent-team-xyz") + assert!( + recorder.events_by_name("team.sessionStatusChanged").is_empty(), + "missing teams should not emit session status events" ); - assert!(failed.data.get("error").is_some()); } #[tokio::test] @@ -5334,7 +5644,7 @@ async fn ss4_stop_session_rejects_cross_user_access() { let result = svc.stop_session("user2", &created.id).await; - assert!(matches!(result, Err(aionui_team::TeamError::Forbidden(_)))); + assert!(matches!(result, Err(aionui_team::TeamError::TeamNotFound(_)))); } // =========================================================================== @@ -5386,7 +5696,7 @@ async fn sm2_send_message_rejects_cross_user_access() { let result = svc.send_message("user2", &created.id, "Hello", None).await; - assert!(matches!(result, Err(aionui_team::TeamError::Forbidden(_)))); + assert!(matches!(result, Err(aionui_team::TeamError::TeamNotFound(_)))); } #[tokio::test] @@ -5431,7 +5741,7 @@ async fn sa2_send_message_to_agent_rejects_cross_user_access() { .send_message_to_agent("user2", &created.id, &worker_slot, "Do this", None) .await; - assert!(matches!(result, Err(aionui_team::TeamError::Forbidden(_)))); + assert!(matches!(result, Err(aionui_team::TeamError::TeamNotFound(_)))); } #[tokio::test] @@ -6471,9 +6781,12 @@ async fn attach_agent_runtime_rejects_cross_user() { .attach_agent_runtime("intruder", &created.id, &worker.slot_id) .await .expect_err("cross-user directed attach must be rejected"); + // User-scope convention: cross-user access returns TeamNotFound (404) to + // avoid leaking team existence, consistent with every other cross-user + // team operation on this branch. assert!( - matches!(error, TeamError::Forbidden(_)), - "expected Forbidden, got {error:?}" + matches!(error, TeamError::TeamNotFound(_)), + "expected TeamNotFound, got {error:?}" ); } @@ -6720,7 +7033,10 @@ async fn lazy_attach_failure_preserves_unread_and_skips_leader_on_human_delivery // Preserve-unread (spec 5.4b): the delivered message must remain unread so a // retry re-drains it via reconcile_mailbox. - let worker_unread = team_repo.peek_unread(&created.id, &worker.slot_id).await.unwrap(); + let worker_unread = team_repo + .peek_unread("user1", &created.id, &worker.slot_id) + .await + .unwrap(); assert!( worker_unread.iter().any(|message| message.content == "please do X"), "a failed lazy attach must not mark the pending delivery read" @@ -6728,7 +7044,10 @@ async fn lazy_attach_failure_preserves_unread_and_skips_leader_on_human_delivery // Human-direct failure must NOT wake the leader (spec 5.4c): the failure is // surfaced inline to the user instead. - let lead_unread = team_repo.peek_unread(&created.id, &lead.slot_id).await.unwrap(); + let lead_unread = team_repo + .peek_unread("user1", &created.id, &lead.slot_id) + .await + .unwrap(); assert!( !lead_unread .iter() @@ -6804,7 +7123,10 @@ async fn agent_triggered_attach_failure_notifies_leader() { // re-delegate. tokio::time::timeout(std::time::Duration::from_secs(2), async { loop { - let leader_messages = team_repo.get_history(&created.id, &lead_slot_id, None).await.unwrap(); + let leader_messages = team_repo + .get_history("user1", &created.id, &lead_slot_id, None) + .await + .unwrap(); if leader_messages.iter().any(|message| { message.from_agent_id == spawned.slot_id && message.content.contains("failed to start its runtime") }) { diff --git a/crates/aionui-team/tests/task_board_integration.rs b/crates/aionui-team/tests/task_board_integration.rs index 0302cedde..8f996d283 100644 --- a/crates/aionui-team/tests/task_board_integration.rs +++ b/crates/aionui-team/tests/task_board_integration.rs @@ -11,13 +11,29 @@ use std::sync::Arc; +use aionui_common::now_ms; use aionui_db::{ITeamRepository, SqliteTeamRepository, init_database_memory}; use aionui_team::{TaskBoard, TaskStatus, TaskUpdate}; async fn setup() -> (TaskBoard, aionui_db::Database) { let db = init_database_memory().await.unwrap(); - let repo = Arc::new(SqliteTeamRepository::new(db.pool().clone())) as Arc; - (TaskBoard::new(repo), db) + let repo = Arc::new(SqliteTeamRepository::new(db.pool().clone())); + repo.create_team(&aionui_db::models::TeamRow { + id: "t1".to_owned(), + user_id: "system_default_user".to_owned(), + name: "t1".to_owned(), + workspace: String::new(), + workspace_mode: "shared".to_owned(), + agents: "[]".to_owned(), + lead_agent_id: None, + session_mode: None, + agents_version: "1.0.1".to_owned(), + created_at: now_ms(), + updated_at: now_ms(), + }) + .await + .unwrap(); + (TaskBoard::new(repo as Arc), db) } // -- TK: Create tasks ---------------------------------------------------------