From a8af8792ea3980a932b4331cd5c518238f7e487f Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 11:38:13 +0800 Subject: [PATCH 001/145] db: extend users for external identity projection --- crates/aionui-api-types/src/auth.rs | 87 +++++- crates/aionui-api-types/src/lib.rs | 8 +- .../aionui-db/migrations/026_user_scope.sql | 92 ++++++ crates/aionui-db/src/lib.rs | 7 +- crates/aionui-db/src/models/mod.rs | 2 +- crates/aionui-db/src/models/user.rs | 49 ++- .../aionui-db/src/repository/sqlite_user.rs | 280 +++++++++++++++--- crates/aionui-db/src/repository/user.rs | 28 +- crates/aionui-db/tests/user_repository.rs | 80 ++++- 9 files changed, 573 insertions(+), 60 deletions(-) create mode 100644 crates/aionui-db/migrations/026_user_scope.sql diff --git a/crates/aionui-api-types/src/auth.rs b/crates/aionui-api-types/src/auth.rs index 6d069ac12..584e19585 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,61 @@ 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 token: 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 +247,36 @@ 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_change_password_request_snake_case() { let raw = r#"{"current_password":"old123","new_password":"new456"}"#; diff --git a/crates/aionui-api-types/src/lib.rs b/crates/aionui-api-types/src/lib.rs index 795667e22..74cb78c6a 100644 --- a/crates/aionui-api-types/src/lib.rs +++ b/crates/aionui-api-types/src/lib.rs @@ -61,9 +61,11 @@ pub use assistant::{ is_local_avatar_value, }; pub use auth::{ - AuthStatusResponse, ChangePasswordRequest, LoginRequest, LoginResponse, PublicUser, QrLoginRequest, - RefreshResponse, RefreshTokenRequest, UserInfoResponse, WebuiChangePasswordRequest, WebuiChangeUsernameRequest, - WebuiChangeUsernameResponse, WebuiGenerateQrTokenResponse, WebuiResetPasswordResponse, WsTokenResponse, + AuthStatusResponse, ChangePasswordRequest, EnsureExternalSessionRequest, EnsureExternalSessionResponse, + EnsureExternalUserRequest, EnsureExternalUserResponse, ExternalUserType, InternalAuthErrorCode, LoginRequest, + LoginResponse, PublicUser, QrLoginRequest, RefreshResponse, RefreshTokenRequest, UserInfoResponse, + WebuiChangePasswordRequest, WebuiChangeUsernameRequest, WebuiChangeUsernameResponse, WebuiGenerateQrTokenResponse, + WebuiResetPasswordResponse, WsTokenResponse, }; pub use channel::{ ApprovePairingRequest, BridgeResponse, ChannelAssistantSettingRequest, ChannelAssistantSettingResponse, diff --git a/crates/aionui-db/migrations/026_user_scope.sql b/crates/aionui-db/migrations/026_user_scope.sql new file mode 100644 index 000000000..488c98c2e --- /dev/null +++ b/crates/aionui-db/migrations/026_user_scope.sql @@ -0,0 +1,92 @@ +-- Migration 026: extend users for local and AionPro identity projection. + +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; + +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); + +PRAGMA foreign_keys = ON; diff --git a/crates/aionui-db/src/lib.rs b/crates/aionui-db/src/lib.rs index dae365284..3da6e9096 100644 --- a/crates/aionui-db/src/lib.rs +++ b/crates/aionui-db/src/lib.rs @@ -25,9 +25,10 @@ pub use instance_lock::{DataDirInstanceGuard, instance_lock_path}; pub use models::{ AgentMetadataRow, AssistantDefinitionRow, AssistantOverlayRow, AssistantOverrideRow, AssistantPreferenceRow, AssistantRow, ConversationArtifactRow, ConversationAssistantSnapshotRow, CreateAssistantParams, - SkillImportRecordRow, SkillRow, UpdateAgentAvailabilitySnapshotParams, UpdateAgentHandshakeParams, - UpdateAssistantParams, UpsertAgentMetadataParams, UpsertAssistantDefinitionParams, UpsertAssistantOverlayParams, - UpsertAssistantPreferenceParams, UpsertConversationAssistantSnapshotParams, UpsertOverrideParams, + ExternalUserProjection, SkillImportRecordRow, SkillRow, UpdateAgentAvailabilitySnapshotParams, + UpdateAgentHandshakeParams, UpdateAssistantParams, UpsertAgentMetadataParams, UpsertAssistantDefinitionParams, + UpsertAssistantOverlayParams, UpsertAssistantPreferenceParams, UpsertConversationAssistantSnapshotParams, + UpsertOverrideParams, UserStatus, UserType, }; pub use repository::channel::UpdatePluginStatusParams; pub use repository::conversation::{ diff --git a/crates/aionui-db/src/models/mod.rs b/crates/aionui-db/src/models/mod.rs index 076836557..054a3bb2e 100644 --- a/crates/aionui-db/src/models/mod.rs +++ b/crates/aionui-db/src/models/mod.rs @@ -38,4 +38,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/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/sqlite_user.rs b/crates/aionui-db/src/repository/sqlite_user.rs index cf55bcf0f..1ade5b323 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,8 +116,74 @@ 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?; @@ -126,6 +199,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 +226,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 +245,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 +299,52 @@ 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 = ?, updated_at = ? WHERE id = ?") + .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 +368,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 +417,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 +457,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 +480,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 +533,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 +552,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 +606,74 @@ 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!(repo.find_active_by_id(&user.id).await.unwrap().is_none()); + } + + #[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); } } diff --git a/crates/aionui-db/src/repository/user.rs b/crates/aionui-db/src/repository/user.rs index 690aa5a85..2bf771b65 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,30 @@ 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>; + /// 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 +75,10 @@ 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. + 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/user_repository.rs b/crates/aionui-db/tests/user_repository.rs index 43aecc390..d48332540 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,62 @@ 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(); + + assert!(r.find_by_id(&user.id).await.unwrap().is_some()); + assert!(r.find_active_by_id(&user.id).await.unwrap().is_none()); +} + +#[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); +} From 24a46c2182fef17460910f4c87126dd9344302dd Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 11:38:57 +0800 Subject: [PATCH 002/145] auth: add trusted external user provisioning --- .../aionui-app/src/bootstrap/environment.rs | 25 ++- crates/aionui-app/src/cli.rs | 21 ++ crates/aionui-app/src/config.rs | 35 +++ crates/aionui-app/src/lib.rs | 2 +- crates/aionui-app/src/router/routes.rs | 19 +- crates/aionui-app/src/services.rs | 9 +- crates/aionui-app/tests/local_mode.rs | 39 ++++ crates/aionui-auth/src/csrf.rs | 7 +- crates/aionui-auth/src/lib.rs | 5 +- crates/aionui-auth/src/middleware.rs | 15 +- crates/aionui-auth/src/routes.rs | 159 ++++++++++++-- crates/aionui-auth/src/service.rs | 102 +++++++++ crates/aionui-auth/tests/middleware_tests.rs | 8 +- crates/aionui-auth/tests/route_tests.rs | 201 +++++++++++++++++- 14 files changed, 607 insertions(+), 40 deletions(-) create mode 100644 crates/aionui-auth/src/service.rs diff --git a/crates/aionui-app/src/bootstrap/environment.rs b/crates/aionui-app/src/bootstrap/environment.rs index c9817e3e7..56c79d80b 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,35 @@ pub fn init_environment(cli: &Cli, merged_path: &str) -> 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/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 bb50d63f6..3ddb0fa7e 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -18,7 +18,8 @@ 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")] @@ -140,13 +141,15 @@ 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(), local: services.local, }; 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 @@ -258,7 +261,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( @@ -284,7 +287,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([ @@ -302,6 +305,14 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates } } +fn auth_identity_mode(identity_mode: crate::config::IdentityMode) -> AuthIdentityMode { + if identity_mode.is_local() { + AuthIdentityMode::Local + } else { + AuthIdentityMode::UserSession + } +} + 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) { diff --git a/crates/aionui-app/src/services.rs b/crates/aionui-app/src/services.rs index 56cdcc91c..bf7d88a11 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())); @@ -228,6 +231,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/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-auth/src/csrf.rs b/crates/aionui-auth/src/csrf.rs index 27e9f54a8..3577d9bcb 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 == "/api/auth/internal/external-sessions"; if needs_validation && !is_exempt { let header_token = request diff --git a/crates/aionui-auth/src/lib.rs b/crates/aionui-auth/src/lib.rs index fdefc280d..bbe6d78d5 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 service::{AuthProvisionService, ProvisionError}; diff --git a/crates/aionui-auth/src/middleware.rs b/crates/aionui-auth/src/middleware.rs index d9b8520a5..9ee0adb26 100644 --- a/crates/aionui-auth/src/middleware.rs +++ b/crates/aionui-auth/src/middleware.rs @@ -12,6 +12,12 @@ use aionui_db::IUserRepository; use crate::JwtService; use crate::extract::extract_token_from_headers; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AuthIdentityMode { + Local, + UserSession, +} + /// Authenticated user injected into request extensions by the auth middleware. /// /// Route handlers extract this from `request.extensions()` to identify @@ -29,8 +35,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,7 +55,7 @@ pub async fn auth_middleware( next: Next, ) -> Result { // In local mode, skip JWT verification and inject a fixed default user. - if state.local { + if state.identity_mode == AuthIdentityMode::Local { request.extensions_mut().insert(CurrentUser { id: "system_default_user".to_string(), username: "system_default_user".to_string(), @@ -68,7 +73,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"); @@ -78,7 +83,7 @@ pub async fn auth_middleware( request.extensions_mut().insert(CurrentUser { id: user.id, - username: user.username, + username: user.username.unwrap_or_else(|| "external_user".to_string()), }); Ok(next.run(request).await) diff --git a/crates/aionui-auth/src/routes.rs b/crates/aionui-auth/src/routes.rs index 670e5b545..b0fc7137f 100644 --- a/crates/aionui-auth/src/routes.rs +++ b/crates/aionui-auth/src/routes.rs @@ -5,16 +5,17 @@ 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 aionui_api_types::{ - ApiResponse, AuthStatusResponse, ChangePasswordRequest, LoginRequest, LoginResponse, PublicUser, QrLoginRequest, - RefreshResponse, RefreshTokenRequest, UserInfoResponse, WebuiChangePasswordRequest, WebuiChangeUsernameRequest, + ApiResponse, AuthStatusResponse, ChangePasswordRequest, EnsureExternalSessionRequest, EnsureExternalUserRequest, + EnsureExternalUserResponse, LoginRequest, LoginResponse, PublicUser, QrLoginRequest, RefreshResponse, + RefreshTokenRequest, UserInfoResponse, WebuiChangePasswordRequest, WebuiChangeUsernameRequest, WebuiChangeUsernameResponse, WebuiGenerateQrTokenResponse, WebuiResetPasswordResponse, WsTokenResponse, }; use aionui_common::ApiError; @@ -23,15 +24,18 @@ use aionui_db::{DbError, IUserRepository, 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"; + impl From for ApiError { fn from(err: AuthError) -> Self { match err { @@ -64,6 +68,8 @@ 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 local: bool, } @@ -103,6 +109,67 @@ 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}")), + } +} + /// Build the auth router with all endpoints and middleware layers. /// /// Returns a `Router` with these endpoints: @@ -133,7 +200,7 @@ 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: AuthIdentityMode::UserSession, }; // Auth rate limited routes (login, qr-login) @@ -146,6 +213,14 @@ 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/users", get(list_internal_users_handler).post(create_internal_user_handler), @@ -220,6 +295,55 @@ 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 response = service + .create_external_session(req) + .await + .map_err(provision_error_to_api_error)?; + tracing::info!( + user_id = %response.user.id, + "external core session exchange succeeded" + ); + let cookie = state.cookie_config.build_session_cookie(&response.token); + Ok(([(header::SET_COOKIE, cookie)], Json(ApiResponse::ok(response))).into_response()) +} + // --------------------------------------------------------------------------- // POST /login // --------------------------------------------------------------------------- @@ -246,7 +370,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 +378,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 +400,7 @@ async fn login_handler( let token = state .jwt_service - .sign(&user.id, &user.username) + .sign(&user.id, user.username.as_deref().unwrap_or("external_user")) .map_err(|e| ApiError::Internal(format!("Token signing error: {e}")))?; // Update last login (best-effort) @@ -284,7 +412,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, ); @@ -505,7 +633,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())); } @@ -618,7 +749,7 @@ async fn qr_login_handler( let token = state .jwt_service - .sign(&user.id, &user.username) + .sign(&user.id, user.username.as_deref().unwrap_or("external_user")) .map_err(|e| ApiError::Internal(format!("Token signing error: {e}")))?; // Update last login (best-effort) @@ -630,7 +761,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 +894,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..e49632472 --- /dev/null +++ b/crates/aionui-auth/src/service.rs @@ -0,0 +1,102 @@ +use std::sync::Arc; + +use aionui_api_types::{ + EnsureExternalSessionRequest, EnsureExternalSessionResponse, EnsureExternalUserRequest, EnsureExternalUserResponse, + ExternalUserType, PublicUser, +}; +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, +} + +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); + } + + 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(&user.id, &username)?; + self.user_repo.update_last_login(&user.id).await?; + + Ok(EnsureExternalSessionResponse { + user: PublicUser { id: user.id, username }, + token, + session_generation: user.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..0d87f3df0 100644 --- a/crates/aionui-auth/tests/middleware_tests.rs +++ b/crates/aionui-auth/tests/middleware_tests.rs @@ -10,9 +10,9 @@ 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}; @@ -134,7 +134,7 @@ fn protected_auth_app(jwt_service: Arc, user_repo: Arc (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) { 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,6 +45,12 @@ 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 { + AuthIdentityMode::UserSession + }, + bootstrap_secret: bootstrap_secret.map(Arc::::from), local, }; @@ -88,6 +100,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() @@ -713,6 +754,162 @@ 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_and_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 json = body_json(session).await; + let token = json["data"]["token"].as_str().unwrap(); + 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_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 // =========================================================================== From ada0ab2131d142d7832c73480339e0fcd118e2df Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 11:40:10 +0800 Subject: [PATCH 003/145] conversation: enforce repository-level user scope --- .../src/acp_error_recovery.rs | 8 +- .../src/message_persistence.rs | 3 +- crates/aionui-conversation/src/routes_aux.rs | 8 +- .../src/runtime_completion.rs | 7 +- crates/aionui-conversation/src/service.rs | 296 ++++++----- crates/aionui-conversation/src/service_ops.rs | 11 +- .../aionui-conversation/src/service_test.rs | 311 ++++++++---- .../service_test/acp_error_recovery_test.rs | 3 +- .../src/startup_recovery.rs | 9 +- .../src/stream_persistence.rs | 72 ++- .../aionui-conversation/src/stream_relay.rs | 36 +- .../src/turn_orchestrator.rs | 23 +- .../tests/acp_tool_call_persistence.rs | 12 +- .../tests/conversation_crud.rs | 4 +- .../tests/conversation_extended.rs | 48 +- .../tests/stream_relay_tips.rs | 1 + .../tests/stream_relay_tool_call.rs | 3 + crates/aionui-db/src/lib.rs | 2 +- .../aionui-db/src/repository/conversation.rs | 84 +++- .../src/repository/sqlite_conversation.rs | 475 +++++++++++++----- .../tests/conversation_repository.rs | 210 ++++---- 21 files changed, 1065 insertions(+), 561 deletions(-) diff --git a/crates/aionui-conversation/src/acp_error_recovery.rs b/crates/aionui-conversation/src/acp_error_recovery.rs index e03b68e55..23b84a9a6 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, @@ -202,6 +203,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, @@ -227,7 +229,7 @@ impl ConversationService { .await; self.clear_persisted_acp_model_after_model_not_found(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_aux.rs b/crates/aionui-conversation/src/routes_aux.rs index 5a56e27dd..8ada037fe 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)?, ))) @@ -82,14 +82,14 @@ async fn get_slash_commands( 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..9c6413d3d 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, diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index d8d1ea0af..a1a600079 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -553,8 +553,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(), @@ -596,14 +597,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,19 +619,27 @@ 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? { + if let Some(snapshot) = self + .conversation_repo + .get_assistant_snapshot(user_id, &response.id) + .await? + { response.assistant = Some(self.assistant_identity_from_snapshot(&snapshot).await?); } @@ -1130,25 +1145,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"))?; } @@ -1527,12 +1545,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 +1563,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 +1595,7 @@ impl ConversationService { pub(crate) async fn persist_runtime_assistant_preferences( &self, + user_id: &str, conversation_id: &str, updates: AssistantRuntimePreferenceUpdate<'_>, ) -> Result<(), ConversationError> { @@ -1584,7 +1607,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 +1616,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(()); }; @@ -1737,16 +1764,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 +1812,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 +1848,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)?; @@ -1923,11 +1948,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 +1962,7 @@ impl ConversationService { ) .await?; self.persist_runtime_assistant_preferences( + user_id, id, AssistantRuntimePreferenceUpdate { model: Some(selected_model), @@ -1958,7 +1985,7 @@ 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"))?; @@ -1975,15 +2002,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,7 +2032,9 @@ 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(()) } @@ -2038,9 +2072,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 @@ -2060,7 +2093,7 @@ impl ConversationService { hook.on_conversation_deleted(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()); } @@ -2129,14 +2162,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 +2181,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 +2194,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 +2218,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 +2253,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 +2283,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 +2353,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 +2389,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 +2405,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 +2430,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,7 +2443,7 @@ 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(), @@ -2473,9 +2504,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 +2531,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(), })?; @@ -2544,9 +2573,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 +2615,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 +2661,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()); } @@ -2670,6 +2697,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 +2706,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 +2755,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 +2781,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 +2790,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 +2809,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 +2818,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 +2866,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 +2886,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; @@ -2894,8 +2925,8 @@ 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 = @@ -2926,9 +2957,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 +3029,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 +3043,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 +3090,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 +3131,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 +3174,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"); @@ -3371,6 +3399,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 +3417,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 +3432,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, @@ -3437,7 +3466,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 +3488,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), @@ -3953,6 +3982,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 +3990,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 +4016,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..d696912e2 100644 --- a/crates/aionui-conversation/src/service_ops.rs +++ b/crates/aionui-conversation/src/service_ops.rs @@ -37,6 +37,7 @@ impl ConversationService { pub async fn set_config_option( &self, + user_id: &str, conversation_id: &str, option_id: &str, req: SetConfigOptionRequest, @@ -103,7 +104,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 +116,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!( @@ -167,6 +171,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 +183,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 a3148c3e1..9e3c13212 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -232,6 +232,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 +246,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 +261,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 +292,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 +355,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 +364,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 +394,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 +407,7 @@ impl IConversationRepository for MockRepo { async fn list_messages_page( &self, + _user_id: &str, conv_id: &str, params: &MessagePageParams, ) -> Result { @@ -474,17 +487,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 +518,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 +528,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 +544,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 +554,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 +581,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 +598,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 +613,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 +627,7 @@ impl IConversationRepository for MockRepo { async fn update_artifact_status( &self, + _user_id: &str, conversation_id: &str, artifact_id: &str, status: &str, @@ -610,6 +647,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 +664,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 +678,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 @@ -1594,7 +1637,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] @@ -1640,7 +1683,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] @@ -1688,7 +1731,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] @@ -1737,7 +1780,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] @@ -1858,6 +1901,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()), @@ -2129,7 +2173,7 @@ 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(); + let exists = self.repo.get("user_1", conversation_id).await.unwrap().is_some(); self.observations.lock().unwrap().push(exists); } } @@ -2148,7 +2192,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] @@ -2347,22 +2391,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()); } @@ -2370,22 +2417,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(); @@ -3634,6 +3684,7 @@ async fn set_config_option_returns_observed_confirmation() { let result = svc .set_config_option( + "user_1", &conv.id, "reasoning_effort", SetConfigOptionRequest { @@ -3745,6 +3796,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 { @@ -3779,6 +3831,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 { @@ -3838,6 +3891,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 { @@ -3851,10 +3905,11 @@ async fn set_config_option_persists_runtime_model_into_assistant_preference_when let pref_after_model = preference_repo.get("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 { @@ -3866,12 +3921,13 @@ async fn set_config_option_persists_runtime_model_into_assistant_preference_when let pref_after_mode = preference_repo.get("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,7 +3940,7 @@ async fn set_config_option_persists_runtime_model_into_assistant_preference_when 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") @@ -3943,6 +3999,7 @@ async fn set_config_option_does_not_persist_preference_on_error() { let result = svc .set_config_option( + "user_1", &conv.id, "model", SetConfigOptionRequest { @@ -4004,6 +4061,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 { @@ -4013,6 +4071,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 { @@ -4022,6 +4081,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 { @@ -4037,7 +4097,7 @@ async fn set_config_option_skips_preference_write_back_when_default_mode_is_fixe 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")); @@ -4091,6 +4151,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 { @@ -4103,7 +4164,7 @@ async fn set_config_option_command_ack_does_not_persist_assistant_preference() { let pref = preference_repo.get("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")); } @@ -4172,7 +4233,11 @@ async fn update_aionrs_model_updates_assistant_preference_only_when_snapshot_mod ); let auto_pref = preference_repo.get("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( @@ -4230,11 +4295,15 @@ async fn update_aionrs_model_updates_assistant_preference_only_when_snapshot_mod let fixed_pref = preference_repo.get("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 @@ -4252,6 +4321,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()), @@ -4274,6 +4344,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, @@ -4379,7 +4450,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), @@ -4403,6 +4474,7 @@ async fn send_message_persists_hidden_user_message_when_requested() { let messages = repo .list_messages_page( + "user_1", &conv.id, &MessagePageParams { limit: 20, @@ -4443,6 +4515,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, @@ -4483,7 +4556,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), @@ -4543,32 +4616,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(), @@ -4612,6 +4688,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, @@ -4721,7 +4798,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"); @@ -4761,6 +4838,7 @@ async fn send_message_rejects_when_runtime_is_shutting_down() { let messages = repo .list_messages_page( + "user_1", &conv.id, &MessagePageParams { limit: 20, @@ -4798,6 +4876,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, @@ -4842,7 +4921,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 @@ -4868,30 +4947,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(); @@ -4899,6 +4984,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, @@ -5000,6 +5086,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, @@ -5038,6 +5125,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, @@ -5807,6 +5895,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()), @@ -6324,7 +6413,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"); @@ -6606,7 +6695,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 { @@ -6648,7 +6737,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 { @@ -6752,7 +6841,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")); @@ -7233,7 +7322,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()); @@ -7356,7 +7445,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"); 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/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 7bcad7a9b..b0568f53a 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,7 +119,7 @@ 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"); } @@ -151,7 +159,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 +178,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 +201,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 +221,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 +252,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 +271,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 +291,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 +312,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 +338,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 +371,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 +401,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"); } } @@ -411,7 +439,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, @@ -461,7 +489,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"); } } @@ -485,7 +513,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); @@ -495,7 +523,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 { @@ -510,7 +542,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 b3bdfc4f9..0b3ce7fbb 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, @@ -2048,7 +2049,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; } @@ -2252,19 +2254,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( @@ -2303,6 +2313,7 @@ mod tests { } async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &aionui_db::MessagePageParams, ) -> Result { @@ -2312,7 +2323,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))); } @@ -2322,7 +2333,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))); } @@ -2346,18 +2357,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..bc56578c7 100644 --- a/crates/aionui-conversation/src/turn_orchestrator.rs +++ b/crates/aionui-conversation/src/turn_orchestrator.rs @@ -136,6 +136,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 +152,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 +172,7 @@ impl ConversationTurnOrchestrator { ); self.service .persist_and_broadcast_send_failure_tip( + &input.user_id, &input.conv_id, &input.turn_id, &send_error, @@ -244,6 +251,7 @@ impl ConversationTurnOrchestrator { ); self.service .persist_and_broadcast_send_failure_tip( + &input.user_id, &input.conv_id, &input.turn_id, &send_error, @@ -311,6 +319,7 @@ impl ConversationTurnOrchestrator { persist_session_key( self.service.conversation_repo(), &persistence, + &input.user_id, &input.conv_id, &session_key, ) @@ -451,6 +460,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 +476,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 +491,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, @@ -507,7 +524,7 @@ impl ConversationTurnOrchestrator { 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 { 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..0a26e3ff5 100644 --- a/crates/aionui-conversation/tests/conversation_crud.rs +++ b/crates/aionui-conversation/tests/conversation_crud.rs @@ -790,9 +790,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")); } 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 834a62dd2..81e02766e 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, @@ -333,6 +335,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-db/src/lib.rs b/crates/aionui-db/src/lib.rs index 3da6e9096..7ee9a5b9e 100644 --- a/crates/aionui-db/src/lib.rs +++ b/crates/aionui-db/src/lib.rs @@ -33,7 +33,7 @@ pub use models::{ 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/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/sqlite_conversation.rs b/crates/aionui-db/src/repository/sqlite_conversation.rs index c27aae2a0..1906314c1 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, }; /// SQLite-backed implementation of [`IConversationRepository`]. @@ -23,7 +23,27 @@ impl SqliteConversationRepository { Self { pool } } - async fn insert_message_once(&self, message: &MessageRow) -> Result<(), sqlx::Error> { + 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?; + + Ok(exists != 0) + } + + 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 insert_message_once(&self, user_id: &str, message: &MessageRow) -> Result<(), DbError> { + self.ensure_conversation_for_user(user_id, &message.conversation_id) + .await?; sqlx::query( "INSERT INTO messages \ (id, conversation_id, msg_id, type, content, position, \ @@ -45,7 +65,9 @@ impl SqliteConversationRepository { Ok(()) } - async fn upsert_message_once(&self, message: &MessageRow) -> Result<(), sqlx::Error> { + async fn upsert_message_once(&self, user_id: &str, message: &MessageRow) -> Result<(), DbError> { + self.ensure_conversation_for_user(user_id, &message.conversation_id) + .await?; sqlx::query( "INSERT INTO messages \ (id, conversation_id, msg_id, type, content, position, \ @@ -91,15 +113,23 @@ impl SqliteConversationRepository { Ok(()) } - 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) @@ -110,15 +140,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) @@ -129,7 +167,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, @@ -142,8 +185,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, @@ -157,8 +204,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?; @@ -166,6 +214,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 \ @@ -192,7 +249,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(); @@ -230,12 +287,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?; @@ -247,8 +308,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?; @@ -403,11 +465,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?; @@ -417,8 +483,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 ( @@ -483,14 +552,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) } @@ -499,21 +577,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) @@ -524,13 +607,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) @@ -544,13 +630,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) @@ -563,11 +652,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) @@ -576,13 +668,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) @@ -593,13 +688,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) @@ -630,16 +728,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) @@ -648,15 +749,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(); @@ -677,13 +784,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?; @@ -694,25 +811,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) @@ -722,9 +850,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') \ @@ -734,7 +862,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( @@ -803,12 +949,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?; @@ -818,12 +970,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) @@ -832,7 +988,13 @@ impl IConversationRepository for SqliteConversationRepository { Ok(row) } - async fn upsert_artifact(&self, artifact: &ConversationArtifactRow) -> Result { + async fn upsert_artifact( + &self, + user_id: &str, + artifact: &ConversationArtifactRow, + ) -> Result { + self.ensure_conversation_for_user(user_id, &artifact.conversation_id) + .await?; sqlx::query( "INSERT INTO conversation_artifacts \ (id, conversation_id, cron_job_id, kind, status, payload, created_at, updated_at) \ @@ -856,13 +1018,14 @@ impl IConversationRepository for SqliteConversationRepository { .execute(&self.pool) .await?; - self.get_artifact(&artifact.conversation_id, &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, @@ -871,12 +1034,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?; @@ -884,29 +1052,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?; @@ -914,21 +1090,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?; @@ -1078,7 +1268,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"); @@ -1090,7 +1280,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()); } #[tokio::test] @@ -1101,6 +1291,7 @@ mod tests { let now = aionui_common::now_ms(); repo.update( + &conv.user_id, &conv.id, &ConversationRowUpdate { name: Some("Updated Name".to_string()), @@ -1111,7 +1302,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); } @@ -1124,6 +1315,7 @@ mod tests { let pin_time = aionui_common::now_ms(); repo.update( + &conv.user_id, &conv.id, &ConversationRowUpdate { pinned: Some(true), @@ -1135,7 +1327,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)); } @@ -1145,6 +1337,7 @@ mod tests { let (repo, _db) = setup().await; let err = repo .update( + "user_1", "no_id", &ConversationRowUpdate { name: Some("x".to_string()), @@ -1163,7 +1356,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] @@ -1172,39 +1367,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(_))); } @@ -1518,10 +1707,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, @@ -1545,11 +1735,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, @@ -1577,11 +1768,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, @@ -1592,6 +1784,7 @@ mod tests { .unwrap(); let older = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 3, @@ -1618,9 +1811,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()), @@ -1632,6 +1827,7 @@ mod tests { let result = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 50, @@ -1648,6 +1844,8 @@ mod tests { let (repo, _db) = setup().await; let err = repo .update_message( + "user_1", + "conv_1", "no_id", &MessageRowUpdate { hidden: Some(true), @@ -1668,13 +1866,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, @@ -1693,10 +1894,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()); @@ -1704,7 +1905,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()); @@ -1718,12 +1919,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); @@ -1738,7 +1939,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) @@ -1759,7 +1960,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/tests/conversation_repository.rs b/crates/aionui-db/tests/conversation_repository.rs index 47c2eaab7..c9c57f3cd 100644 --- a/crates/aionui-db/tests/conversation_repository.rs +++ b/crates/aionui-db/tests/conversation_repository.rs @@ -76,13 +76,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 +95,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 +131,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 +411,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 +444,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 +460,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 +499,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 +538,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 +563,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 +596,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 +613,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 +640,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 +671,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 +713,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 +744,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 +761,7 @@ async fn update_message_fields() { let msgs = repo .list_messages_page( + &conv.user_id, &conv.id, &MessagePageParams { limit: 50, @@ -770,13 +784,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 +814,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 +850,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 +875,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 +895,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 +923,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 +935,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 +954,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 +966,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 +981,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 +997,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 +1020,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 +1031,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 +1041,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 +1085,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 +1121,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 +1154,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()); } From 7287c238cbd1fd034c1ee4d3bb174a7458dd9058 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 11:41:00 +0800 Subject: [PATCH 004/145] team-cron: adapt aggregate adapters to conversation user scope --- .../src/router/team_conversation_adapters.rs | 46 ++++- crates/aionui-cron/src/executor.rs | 190 ++++++++++++++---- crates/aionui-cron/src/skill_suggest.rs | 48 +++-- .../aionui-cron/tests/service_integration.rs | 70 +++++-- 4 files changed, 281 insertions(+), 73 deletions(-) diff --git a/crates/aionui-app/src/router/team_conversation_adapters.rs b/crates/aionui-app/src/router/team_conversation_adapters.rs index 4e6ea5662..0e582ba3b 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,7 +271,11 @@ 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) } @@ -288,7 +317,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-cron/src/executor.rs b/crates/aionui-cron/src/executor.rs index f899ea43d..a0908a5ec 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -50,6 +50,7 @@ pub(crate) enum PreparedRunNow { } struct SkillSuggestContext { + user_id: String, conversation_id: String, job_id: String, workspace: String, @@ -212,7 +213,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,8 +227,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(conversation_id) + .get(&user_id, conversation_id) .await .map_err(CronError::Database) } @@ -248,8 +257,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 +315,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,7 +366,7 @@ impl JobExecutor { ..Default::default() }; self.conversation_repo - .update(conversation_id, &update) + .update(&row.user_id, conversation_id, &update) .await .map_err(CronError::Database)?; } @@ -395,7 +413,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 +446,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 +472,7 @@ impl JobExecutor { ..Default::default() }; self.conversation_repo - .update(conversation_id, &update) + .update(&row.user_id, conversation_id, &update) .await .map_err(CronError::Database) } @@ -595,10 +613,14 @@ impl JobExecutor { async fn resolve_conversation_owner_user_id(&self, job: &CronJob) -> Result { 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()) @@ -653,16 +675,18 @@ 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) { warn!( @@ -688,7 +712,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 +727,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 +746,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 +784,7 @@ impl JobExecutor { async fn conversation_turn_failed_message( &self, + user_id: &str, conversation_id: &str, outcome: ConversationAgentTurnOutcome, ) -> String { @@ -771,7 +799,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, @@ -811,7 +839,7 @@ impl JobExecutor { pub async fn mark_skill_suggest_artifacts_saved(&self, job_id: &str) -> Result<(), CronError> { let rows = self .conversation_repo - .mark_skill_suggest_artifacts_saved(job_id, now_ms()) + .mark_skill_suggest_artifacts_saved(SYSTEM_DEFAULT_USER_ID, job_id, now_ms()) .await .map_err(CronError::Database)?; @@ -849,7 +877,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 +935,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 { @@ -2360,16 +2391,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 +2451,7 @@ mod tests { } async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &MessagePageParams, ) -> Result { @@ -2417,17 +2461,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,7 +2856,11 @@ 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(), @@ -2818,15 +2881,24 @@ mod tests { })) } + async fn owner_user_id(&self, _id: &str) -> Result, aionui_db::DbError> { + Ok(Some("cron".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 +2942,7 @@ mod tests { async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &MessagePageParams, ) -> Result { @@ -2880,20 +2953,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 +3089,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 +3159,7 @@ mod tests { async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &MessagePageParams, ) -> Result { @@ -3068,7 +3170,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 +3184,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 +3228,7 @@ mod tests { async fn upsert_artifact( &self, + _user_id: &str, artifact: &ConversationArtifactRow, ) -> Result { self.operations diff --git a/crates/aionui-cron/src/skill_suggest.rs b/crates/aionui-cron/src/skill_suggest.rs index 029f59bd5..1c9489082 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!( @@ -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/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index 61cb7b928..0a62e28a7 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -173,7 +173,11 @@ impl StubConvRepo { #[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> { let mut rows = self.rows.lock().unwrap(); if let Some(existing) = rows.get(id) { @@ -481,8 +485,19 @@ impl IConversationRepository for StubConvRepo { 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 +505,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 +541,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 +577,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 +630,7 @@ impl IConversationRepository for StubConvRepo { } async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &MessagePageParams, ) -> Result { @@ -618,18 +640,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 +684,7 @@ impl IConversationRepository for StubConvRepo { } async fn list_artifacts( &self, + _user_id: &str, conversation_id: &str, ) -> Result, aionui_db::DbError> { Ok(self @@ -664,6 +698,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 +712,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 +720,7 @@ impl IConversationRepository for StubConvRepo { } async fn update_artifact_status( &self, + _user_id: &str, conversation_id: &str, artifact_id: &str, status: &str, @@ -702,6 +739,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> { @@ -1432,7 +1470,7 @@ async fn cj7b_add_job_binds_existing_conversation_to_job() { let job = svc.add_job(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); @@ -1559,7 +1597,7 @@ async fn update_existing_job_to_new_conversation_removes_previous_conversation_b create_req.execution_mode = Some("existing".into()); let created = svc.add_job(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); @@ -1589,7 +1627,7 @@ async fn update_existing_job_to_new_conversation_removes_previous_conversation_b assert_eq!(row.conversation_id, ""); 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()); @@ -1623,7 +1661,7 @@ async fn update_existing_job_to_new_conversation_clears_previous_auto_workspace( create_req.agent_config.as_mut().unwrap().workspace = Some(auto_workspace); let created = svc.add_job(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) @@ -2233,7 +2271,7 @@ async fn existing_job_with_missing_conversation_run_now_creates_replacement_conv assert_ne!(response.conversation_id, "missing-conv-run-now"); assert!( - conv_repo.get(&response.conversation_id).await.unwrap().is_some(), + conv_repo.get("u1", &response.conversation_id).await.unwrap().is_some(), "run-now should create a replacement conversation for an existing job whose previous conversation was deleted" ); @@ -2328,7 +2366,7 @@ async fn create_for_conversation_helper_creates_claimed_conversation_job_with_mu 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); @@ -2359,7 +2397,7 @@ async fn create_for_conversation_helper_keeps_conversation_extra_mode_unchanged( 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); @@ -2713,6 +2751,7 @@ async fn update_for_conversation_helper_updates_claimed_conversation_job() { conv_repo .update( + "u1", "conv_1", &ConversationRowUpdate { extra: Some("{}".into()), @@ -2737,7 +2776,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 +2797,7 @@ async fn update_for_conversation_helper_fails_when_conversation_binding_fails() conv_repo .update( + "u1", "conv_1", &ConversationRowUpdate { extra: Some("{}".into()), @@ -3068,7 +3108,7 @@ async fn cd3b_on_conversation_delete_clears_deleted_workspace_from_jobs() { let job = svc.add_job(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) From 6ad4cd7a99403c855f6c8b530e60e02b77d2ca5f Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 11:42:43 +0800 Subject: [PATCH 005/145] test: update app checks for scoped conversation repository --- crates/aionui-app/src/router/state.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index a2461348b..7277ae94b 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -1073,13 +1073,18 @@ mod tests { .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"); From 01032c4e2bf001b0a2645dea820114561b8d547c Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 11:44:25 +0800 Subject: [PATCH 006/145] test: update app e2e tests for scoped conversation repository --- crates/aionui-app/tests/message_e2e.rs | 18 ++++++++++++++---- crates/aionui-app/tests/team_e2e.rs | 17 +++++++++++++++-- 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/crates/aionui-app/tests/message_e2e.rs b/crates/aionui-app/tests/message_e2e.rs index 9eb7e75e7..ef678410b 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,19 @@ 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 +433,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/team_e2e.rs b/crates/aionui-app/tests/team_e2e.rs index 423e87eda..4c87d0abf 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"]); @@ -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,14 @@ 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, From 126f6d3bf91191864a858f9fb1dba1175a74df1f Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 11:48:17 +0800 Subject: [PATCH 007/145] test: update cron e2e checks for scoped conversation repository --- crates/aionui-app/tests/cron_e2e.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/crates/aionui-app/tests/cron_e2e.rs b/crates/aionui-app/tests/cron_e2e.rs index 62f152766..fabdcbfff 100644 --- a/crates/aionui-app/tests/cron_e2e.rs +++ b/crates/aionui-app/tests/cron_e2e.rs @@ -784,8 +784,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 +810,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"); From 54fbcd27c8c6e487d7d679143e9dd3b8160a4e63 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 11:49:19 +0800 Subject: [PATCH 008/145] test: update conversation e2e checks for scoped repository --- crates/aionui-app/tests/conversation_e2e.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) 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(); From c6f76894ef829da7e02e022b8ca66ffc34928790 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 12:54:07 +0800 Subject: [PATCH 009/145] team-cron: enforce aggregate root user scope --- Cargo.lock | 1 + crates/aionui-app/tests/message_e2e.rs | 6 +- crates/aionui-app/tests/team_e2e.rs | 6 +- crates/aionui-cron/Cargo.toml | 1 + crates/aionui-cron/src/routes.rs | 36 +- crates/aionui-cron/src/service.rs | 88 ++-- .../aionui-cron/tests/service_integration.rs | 413 ++++++++++-------- .../migrations/027_aggregate_user_scope.sql | 79 ++++ crates/aionui-db/src/repository/cron.rs | 19 + .../aionui-db/src/repository/sqlite_cron.rs | 325 ++++++++++---- .../aionui-db/src/repository/sqlite_team.rs | 90 ++-- crates/aionui-db/src/repository/team.rs | 32 +- crates/aionui-db/tests/cron_repository.rs | 78 ++++ crates/aionui-db/tests/team_repository.rs | 149 +++++-- crates/aionui-team/src/event_loop.rs | 6 +- crates/aionui-team/src/mailbox.rs | 13 +- crates/aionui-team/src/provisioning.rs | 8 +- crates/aionui-team/src/service.rs | 81 ++-- crates/aionui-team/src/session.rs | 23 +- crates/aionui-team/src/task_board.rs | 15 +- crates/aionui-team/src/test_utils.rs | 155 +++++-- crates/aionui-team/src/workspace.rs | 6 + crates/aionui-team/tests/common/mod.rs | 36 +- .../aionui-team/tests/mailbox_integration.rs | 22 +- .../tests/session_service_integration.rs | 169 ++++--- .../tests/task_board_integration.rs | 20 +- 26 files changed, 1287 insertions(+), 590 deletions(-) create mode 100644 crates/aionui-db/migrations/027_aggregate_user_scope.sql diff --git a/Cargo.lock b/Cargo.lock index 1469baa3c..9d2556db2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -604,6 +604,7 @@ dependencies = [ "dashmap", "serde", "serde_json", + "sqlx", "tempfile", "thiserror 2.0.18", "tokio", diff --git a/crates/aionui-app/tests/message_e2e.rs b/crates/aionui-app/tests/message_e2e.rs index ef678410b..7f3d47384 100644 --- a/crates/aionui-app/tests/message_e2e.rs +++ b/crates/aionui-app/tests/message_e2e.rs @@ -110,11 +110,7 @@ async fn insert_acp_tool_message( async fn upsert_artifact(services: &aionui_app::AppServices, artifact: aionui_db::ConversationArtifactRow) { let repo = aionui_db::SqliteConversationRepository::new(services.database.pool().clone()); - let user_id = repo - .owner_user_id(&artifact.conversation_id) - .await - .unwrap() - .unwrap(); + let user_id = repo.owner_user_id(&artifact.conversation_id).await.unwrap().unwrap(); aionui_db::IConversationRepository::upsert_artifact(&repo, &user_id, &artifact) .await .unwrap(); diff --git a/crates/aionui-app/tests/team_e2e.rs b/crates/aionui-app/tests/team_e2e.rs index 4c87d0abf..ab426daa1 100644 --- a/crates/aionui-app/tests/team_e2e.rs +++ b/crates/aionui-app/tests/team_e2e.rs @@ -1262,11 +1262,7 @@ 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 user_id = repo.owner_user_id(lead_conversation_id).await.unwrap().unwrap(); let messages = repo .list_messages_page( &user_id, diff --git a/crates/aionui-cron/Cargo.toml b/crates/aionui-cron/Cargo.toml index 040398696..0edd90d52 100644 --- a/crates/aionui-cron/Cargo.toml +++ b/crates/aionui-cron/Cargo.toml @@ -28,6 +28,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/routes.rs b/crates/aionui-cron/src/routes.rs index 5b7ff0f9f..68aa865d1 100644 --- a/crates/aionui-cron/src/routes.rs +++ b/crates/aionui-cron/src/routes.rs @@ -73,60 +73,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))) } @@ -200,12 +200,12 @@ fn header_value(headers: &HeaderMap, name: &'static str) -> Result, - 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 +220,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())) } diff --git a/crates/aionui-cron/src/service.rs b/crates/aionui-cron/src/service.rs index 43d7fe580..60771efcc 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -87,8 +87,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( @@ -125,7 +125,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 @@ -136,7 +136,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, @@ -166,9 +166,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 } @@ -182,13 +185,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), @@ -244,6 +248,7 @@ impl CronService { async fn add_job_internal( &self, + user_id: &str, req: CreateCronJobRequest, runtime_agent_type: Option, assistant_backend_override: Option, @@ -259,6 +264,21 @@ impl CronService { 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 conversation_id.is_empty() + || self + .executor + .get_conversation_row(conversation_id) + .await? + .filter(|row| row.user_id == user_id) + .is_none() + { + return Err(CronError::Conversation( + aionui_conversation::ConversationError::NotFound { + id: req.conversation_id, + }, + )); + } let agent_config = match req.agent_config { Some(config) => Some( @@ -283,7 +303,7 @@ impl CronService { 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, @@ -313,10 +333,15 @@ impl CronService { 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)?; @@ -352,9 +377,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() @@ -399,13 +422,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 @@ -429,21 +451,21 @@ impl CronService { 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.repo.delete_for_user(user_id, job_id).await?; self.emitter.emit_job_removed(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)?; @@ -451,11 +473,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()); @@ -686,10 +708,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)?; @@ -718,10 +740,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()))?; @@ -733,17 +755,17 @@ impl CronService { skill_content: Some(Some(req.content)), ..Default::default() }; - self.repo.update(job_id, ¶ms).await?; + self.repo.update_for_user(user_id, job_id, ¶ms).await?; self.executor.mark_skill_suggest_artifacts_saved(job_id).await?; info!(job_id, "Skill content saved"); Ok(()) } - pub async fn has_skill(&self, job_id: &str) -> Result { + 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()))?; @@ -753,9 +775,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()))?; @@ -765,7 +787,7 @@ 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?; info!(job_id, "Skill content deleted"); Ok(()) diff --git a/crates/aionui-cron/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index 0a62e28a7..907cab1c0 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -26,8 +26,8 @@ use aionui_db::{ IAssistantDefinitionRepository, IAssistantOverlayRepository, IAssistantPreferenceRepository, IConversationRepository, ICronRepository, MessagePageParams, MessagePageResult, MessageRowUpdate, MessageSearchRow, SqliteAcpSessionRepository, SqliteAgentMetadataRepository, SqliteAssistantDefinitionRepository, - SqliteAssistantOverlayRepository, SqliteAssistantPreferenceRepository, SqliteCronRepository, - UpsertAgentMetadataParams, UpsertAssistantDefinitionParams, UpsertAssistantOverlayParams, + SqliteAssistantOverlayRepository, SqliteAssistantPreferenceRepository, SqliteConversationRepository, + SqliteCronRepository, UpsertAgentMetadataParams, UpsertAssistantDefinitionParams, UpsertAssistantOverlayParams, UpsertConversationAssistantSnapshotParams, init_database_memory, models::{ConversationAssistantSnapshotRow, CronJobRow, MessageRow}, }; @@ -45,6 +45,37 @@ use tower::ServiceExt; // ── Test infrastructure ──────────────────────────────────────────── +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,19 +144,33 @@ 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 { let mut guard = self.messages.lock().unwrap(); std::mem::take(&mut *guard) @@ -178,10 +223,9 @@ impl IConversationRepository for StubConvRepo { _user_id: &str, 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())); + if let Some(existing) = { self.rows.lock().unwrap().get(id).cloned() } { + self.seed_sqlite_row(&existing).await?; + return Ok(Some(existing)); } if id.starts_with("missing") { return Ok(None); @@ -481,7 +525,8 @@ impl IConversationRepository for StubConvRepo { } }; - rows.insert(id.to_owned(), row.clone()); + self.rows.lock().unwrap().insert(id.to_owned(), row.clone()); + self.seed_sqlite_row(&row).await?; Ok(Some(row)) } @@ -792,6 +837,7 @@ async fn setup_with_conv_runtime_and_agent_metadata() -> ( ) { 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())); @@ -801,7 +847,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(); @@ -830,7 +876,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( @@ -894,6 +940,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())); @@ -903,7 +950,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(); @@ -932,7 +979,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( @@ -1124,7 +1171,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"); @@ -1143,7 +1190,7 @@ 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(); @@ -1166,7 +1213,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")); @@ -1179,7 +1226,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")); @@ -1202,7 +1249,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"); } @@ -1213,7 +1260,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")); } @@ -1239,7 +1286,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"); } @@ -1272,7 +1319,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"); } @@ -1294,7 +1341,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"); @@ -1318,7 +1365,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"); @@ -1334,17 +1381,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); @@ -1355,9 +1405,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"); } @@ -1367,7 +1420,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(_))); } @@ -1377,12 +1430,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); } @@ -1408,7 +1461,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, @@ -1427,7 +1480,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() @@ -1444,20 +1497,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); } @@ -1468,7 +1521,7 @@ 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("u1", "conv_existing_bind").await.unwrap().unwrap(); let extra: serde_json::Value = serde_json::from_str(&bound.extra).unwrap(); @@ -1485,7 +1538,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 { @@ -1501,7 +1557,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); @@ -1516,7 +1572,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(); @@ -1542,7 +1601,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")) ); @@ -1552,10 +1611,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(); @@ -1581,21 +1640,19 @@ 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("u1", "conv_mode_switch").await.unwrap().unwrap(); let extra_before: serde_json::Value = serde_json::from_str(&bound_before.extra).unwrap(); @@ -1604,6 +1661,7 @@ async fn update_existing_job_to_new_conversation_removes_previous_conversation_b let updated = svc .update_job( + "u1", &created.id, UpdateCronJobRequest { name: None, @@ -1624,19 +1682,13 @@ async fn update_existing_job_to_new_conversation_removes_previous_conversation_b let row = cron_repo.get_by_id(&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("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] @@ -1659,7 +1711,7 @@ 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("u1", &conversation_id).await.unwrap().unwrap(); assert!( @@ -1670,6 +1722,7 @@ async fn update_existing_job_to_new_conversation_clears_previous_auto_workspace( ); svc.update_job( + "u1", &created.id, UpdateCronJobRequest { name: None, @@ -1711,9 +1764,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, @@ -1748,7 +1802,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, @@ -1763,7 +1817,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"))); } @@ -1773,7 +1827,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, @@ -1797,7 +1851,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")); @@ -1810,7 +1864,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, @@ -1835,7 +1889,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"); @@ -1852,7 +1906,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(); @@ -1869,7 +1923,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 { .. } @@ -1894,7 +1948,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(_))); } @@ -1903,12 +1957,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(); @@ -1921,7 +1978,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(_)) @@ -1933,19 +1990,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(); @@ -1967,6 +2027,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(), @@ -1996,9 +2057,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(), @@ -2007,7 +2072,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); } @@ -2016,9 +2081,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); } @@ -2027,10 +2095,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(_))); @@ -2042,12 +2113,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(), @@ -2065,6 +2137,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(), @@ -2081,10 +2154,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(), @@ -2093,9 +2167,9 @@ 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(_))); } @@ -2105,7 +2179,10 @@ async fn sk7_delete_cleans_skill() { 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(); @@ -2125,7 +2202,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(_))); } @@ -2143,7 +2220,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); } @@ -2159,7 +2236,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(_))); } @@ -2175,111 +2252,59 @@ 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_rejects_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_err = svc.add_job("u1", empty_req).await.unwrap_err(); + assert!(matches!(empty_err, aionui_cron::error::CronError::Conversation(_))); 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_err = svc.add_job("u1", stale_req).await.unwrap_err(); + assert!(matches!(stale_err, aionui_cron::error::CronError::Conversation(_))); } #[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(); - - svc.init().await; - - 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 found = svc.get_job(&normal.id).await; - assert!(found.is_ok()); + let err = svc.add_job("u1", missing_req).await.unwrap_err(); + assert!(matches!(err, aionui_cron::error::CronError::Conversation(_))); } #[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("u1", &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] @@ -2289,7 +2314,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 @@ -2297,7 +2322,7 @@ 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 { @@ -2320,9 +2345,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(), @@ -2330,10 +2359,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 { @@ -2868,7 +2897,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 { @@ -2883,7 +2915,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); } @@ -2892,12 +2924,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), @@ -2933,7 +2966,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)); } @@ -2950,7 +2983,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)); } @@ -2961,7 +2994,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; @@ -2973,7 +3006,7 @@ async fn sr1_system_resume_missed_job() { 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" @@ -3017,24 +3050,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; - 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(); @@ -3051,7 +3084,9 @@ 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; @@ -3059,7 +3094,7 @@ async fn cd2_delete_by_conversation_no_matching_jobs() { 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"); } @@ -3073,12 +3108,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; - 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()); @@ -3105,7 +3140,7 @@ 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("u1", &conversation_id).await.unwrap().unwrap(); @@ -3124,7 +3159,7 @@ async fn cd3b_on_conversation_delete_clears_deleted_workspace_from_jobs() { svc.on_conversation_deleted(&conversation_id).await; - assert!(svc.get_job(&job.id).await.is_ok()); + assert!(svc.get_job("u1", &job.id).await.is_ok()); let row = cron_repo.get_by_id(&job.id).await.unwrap().unwrap(); let config: CronAgentConfig = serde_json::from_str(row.agent_config.as_deref().unwrap()).unwrap(); assert!( @@ -3153,7 +3188,7 @@ 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; @@ -3175,23 +3210,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; 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/027_aggregate_user_scope.sql b/crates/aionui-db/migrations/027_aggregate_user_scope.sql new file mode 100644 index 000000000..93faf92cd --- /dev/null +++ b/crates/aionui-db/migrations/027_aggregate_user_scope.sql @@ -0,0 +1,79 @@ +-- Enforce aggregate parent chains for new writes without rewriting existing tables. + +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; + +CREATE TRIGGER IF NOT EXISTS trg_cron_jobs_conversation_parent_insert +BEFORE INSERT ON cron_jobs +FOR EACH ROW +WHEN NEW.conversation_id IS NULL + OR NEW.conversation_id = '' + OR NOT EXISTS (SELECT 1 FROM conversations WHERE id = NEW.conversation_id) +BEGIN + SELECT RAISE(ABORT, 'cron_jobs.conversation_id must reference conversations.id'); +END; + +CREATE TRIGGER IF NOT EXISTS trg_cron_jobs_conversation_parent_update +BEFORE UPDATE OF conversation_id ON cron_jobs +FOR EACH ROW +WHEN NEW.conversation_id IS NULL + OR NEW.conversation_id = '' + OR NOT EXISTS (SELECT 1 FROM conversations WHERE id = NEW.conversation_id) +BEGIN + SELECT RAISE(ABORT, 'cron_jobs.conversation_id must reference conversations.id'); +END; + +CREATE TRIGGER IF NOT EXISTS trg_cron_job_runs_job_parent_insert +BEFORE INSERT ON cron_job_runs +FOR EACH ROW +WHEN NOT EXISTS ( + SELECT 1 + FROM cron_jobs j + JOIN conversations c ON c.id = j.conversation_id + WHERE j.id = NEW.job_id +) +BEGIN + SELECT RAISE(ABORT, 'cron_job_runs.job_id must reference cron_jobs.id with conversation parent'); +END; + +CREATE TRIGGER IF NOT EXISTS trg_cron_job_runs_job_parent_update +BEFORE UPDATE OF job_id ON cron_job_runs +FOR EACH ROW +WHEN NOT EXISTS ( + SELECT 1 + FROM cron_jobs j + JOIN conversations c ON c.id = j.conversation_id + WHERE j.id = NEW.job_id +) +BEGIN + SELECT RAISE(ABORT, 'cron_job_runs.job_id must reference cron_jobs.id with conversation parent'); +END; diff --git a/crates/aionui-db/src/repository/cron.rs b/crates/aionui-db/src/repository/cron.rs index d2c27359c..ff17d58a6 100644 --- a/crates/aionui-db/src/repository/cron.rs +++ b/crates/aionui-db/src/repository/cron.rs @@ -74,21 +74,40 @@ pub trait ICronRepository: Send + Sync { /// Returns `DbError::NotFound` if absent. async fn update(&self, id: &str, params: &UpdateCronJobParams) -> 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>; + /// Deletes a cron job by ID. Returns `DbError::NotFound` if absent. async fn delete(&self, id: &str) -> 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 a single cron job by ID, or `None` if not found. async fn get_by_id(&self, 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>; + /// Returns all cron jobs ordered by creation time ascending. async fn list_all(&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>; + /// Returns all enabled cron jobs. async fn list_enabled(&self) -> Result, DbError>; /// Returns all cron jobs for a given conversation. async fn list_by_conversation(&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>; + /// Deletes all cron jobs associated with a conversation. /// Returns the number of deleted rows. async fn delete_by_conversation(&self, conversation_id: &str) -> Result; diff --git a/crates/aionui-db/src/repository/sqlite_cron.rs b/crates/aionui-db/src/repository/sqlite_cron.rs index 916e133da..e8e2ff161 100644 --- a/crates/aionui-db/src/repository/sqlite_cron.rs +++ b/crates/aionui-db/src/repository/sqlite_cron.rs @@ -64,93 +64,11 @@ impl ICronRepository for SqliteCronRepository { } 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(", ")); - - let mut query = sqlx::query(&sql); - for bind in &binds { - query = bind_value(query, bind); - } - query = query.bind(id); + self.update_inner(None, id, params).await + } - let result = query.execute(&self.pool).await?; - if result.rows_affected() == 0 { - return Err(DbError::NotFound(format!("cron job '{id}'"))); - } - Ok(()) + async fn update_for_user(&self, user_id: &str, id: &str, params: &UpdateCronJobParams) -> Result<(), DbError> { + self.update_inner(Some(user_id), id, params).await } async fn delete(&self, id: &str) -> Result<(), DbError> { @@ -164,6 +82,25 @@ impl ICronRepository for SqliteCronRepository { Ok(()) } + async fn delete_for_user(&self, user_id: &str, id: &str) -> Result<(), DbError> { + let result = sqlx::query( + "DELETE FROM cron_jobs \ + WHERE id = ? \ + AND EXISTS (\ + SELECT 1 FROM conversations c \ + WHERE c.id = cron_jobs.conversation_id AND c.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> { let row = sqlx::query_as::<_, CronJobRow>("SELECT * FROM cron_jobs WHERE id = ?") .bind(id) @@ -172,6 +109,19 @@ impl ICronRepository for SqliteCronRepository { Ok(row) } + async fn get_by_id_for_user(&self, user_id: &str, id: &str) -> Result, DbError> { + let row = sqlx::query_as::<_, CronJobRow>( + "SELECT j.* FROM cron_jobs j \ + JOIN conversations c ON c.id = j.conversation_id \ + WHERE c.user_id = ? AND j.id = ?", + ) + .bind(user_id) + .bind(id) + .fetch_optional(&self.pool) + .await?; + Ok(row) + } + async fn list_all(&self) -> Result, DbError> { let rows = sqlx::query_as::<_, CronJobRow>("SELECT * FROM cron_jobs ORDER BY created_at ASC") .fetch_all(&self.pool) @@ -179,10 +129,28 @@ impl ICronRepository for SqliteCronRepository { Ok(rows) } + async fn list_all_for_user(&self, user_id: &str) -> Result, DbError> { + let rows = sqlx::query_as::<_, CronJobRow>( + "SELECT j.* FROM cron_jobs j \ + JOIN conversations c ON c.id = j.conversation_id \ + WHERE c.user_id = ? \ + ORDER BY j.created_at ASC", + ) + .bind(user_id) + .fetch_all(&self.pool) + .await?; + Ok(rows) + } + async fn list_enabled(&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?; + let rows = sqlx::query_as::<_, CronJobRow>( + "SELECT j.* FROM cron_jobs j \ + JOIN conversations c ON c.id = j.conversation_id \ + WHERE j.enabled = 1 \ + ORDER BY j.created_at ASC", + ) + .fetch_all(&self.pool) + .await?; Ok(rows) } @@ -196,6 +164,24 @@ impl ICronRepository for SqliteCronRepository { Ok(rows) } + async fn list_by_conversation_for_user( + &self, + user_id: &str, + conversation_id: &str, + ) -> Result, DbError> { + let rows = sqlx::query_as::<_, CronJobRow>( + "SELECT j.* FROM cron_jobs j \ + JOIN conversations c ON c.id = j.conversation_id \ + WHERE c.user_id = ? AND j.conversation_id = ? \ + ORDER BY j.created_at ASC", + ) + .bind(user_id) + .bind(conversation_id) + .fetch_all(&self.pool) + .await?; + Ok(rows) + } + async fn delete_by_conversation(&self, conversation_id: &str) -> Result { let result = sqlx::query("DELETE FROM cron_jobs WHERE conversation_id = ?") .bind(conversation_id) @@ -209,6 +195,20 @@ 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 j \ + JOIN conversations c ON c.id = j.conversation_id \ + WHERE j.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 +312,12 @@ 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 \ + JOIN conversations c ON c.id = j.conversation_id \ + WHERE j.id = cron_job_runs.job_id\ + )", ) .bind(lease_until) .bind(updated_at) @@ -334,7 +339,12 @@ 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 \ + JOIN conversations c ON c.id = j.conversation_id \ + WHERE j.id = cron_job_runs.job_id\ + )", ) .bind(retry_at) .bind(updated_at) @@ -350,7 +360,12 @@ 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 \ + JOIN conversations c ON c.id = j.conversation_id \ + WHERE j.id = cron_job_runs.job_id\ + )", ) .bind(params.status) .bind(params.conversation_id) @@ -376,9 +391,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 \ + JOIN conversations c ON c.id = j.conversation_id \ + WHERE r.job_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 +406,124 @@ impl ICronRepository for SqliteCronRepository { } } +impl SqliteCronRepository { + 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()) { + let owns_new_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_new_conversation { + return Err(DbError::NotFound(format!("conversation '{conversation_id}'"))); + } + } + + let sql = if user_id.is_some() { + format!( + "UPDATE cron_jobs SET {} WHERE id = ? \ + AND EXISTS (\ + SELECT 1 FROM conversations c \ + WHERE c.id = cron_jobs.conversation_id AND c.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)] diff --git a/crates/aionui-db/src/repository/sqlite_team.rs b/crates/aionui-db/src/repository/sqlite_team.rs index 63990280c..4277fb134 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?; @@ -201,15 +215,16 @@ impl ITeamRepository for SqliteTeamRepository { Ok(rows) } - async fn mark_read_batch(&self, ids: &[String]) -> Result<(), DbError> { + async fn mark_read_batch(&self, 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})"); let mut query = sqlx::query(&sql); + query = query.bind(team_id); for id in chunk { query = query.bind(id); } @@ -254,11 +269,16 @@ impl ITeamRepository for SqliteTeamRepository { 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(()) } @@ -296,7 +316,7 @@ impl ITeamRepository for SqliteTeamRepository { Ok(row) } - async fn update_task(&self, task_id: &str, params: &UpdateTaskParams) -> Result<(), DbError> { + async fn update_task(&self, 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 +339,10 @@ 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 = ?", + set_clauses.join(", ") + ); let mut query = sqlx::query(&sql); if let Some(ref status) = params.status { @@ -338,6 +361,7 @@ impl ITeamRepository for SqliteTeamRepository { query = query.bind(metadata); } query = query.bind(now_ms()); + query = query.bind(team_id); query = query.bind(task_id); let result = query.execute(&self.pool).await?; @@ -356,11 +380,12 @@ impl ITeamRepository for SqliteTeamRepository { Ok(rows) } - async fn append_to_blocks(&self, task_id: &str, blocked_task_id: &str) -> Result<(), DbError> { + async fn append_to_blocks(&self, 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 = ?") + let row = sqlx::query_as::<_, TeamTaskRow>("SELECT * FROM team_tasks WHERE team_id = ? AND id = ?") + .bind(team_id) .bind(task_id) .fetch_optional(&mut *tx) .await? @@ -372,9 +397,10 @@ 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 = ?") + sqlx::query("UPDATE team_tasks SET blocks = ?, updated_at = ? WHERE team_id = ? AND id = ?") .bind(&new_blocks) .bind(now_ms()) + .bind(team_id) .bind(task_id) .execute(&mut *tx) .await?; @@ -383,11 +409,17 @@ impl ITeamRepository for SqliteTeamRepository { Ok(()) } - async fn remove_from_blocked_by(&self, task_id: &str, unblocked_task_id: &str) -> Result<(), DbError> { + async fn remove_from_blocked_by( + &self, + 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 = ?") + let row = sqlx::query_as::<_, TeamTaskRow>("SELECT * FROM team_tasks WHERE team_id = ? AND id = ?") + .bind(team_id) .bind(task_id) .fetch_optional(&mut *tx) .await? @@ -397,9 +429,10 @@ impl ITeamRepository for SqliteTeamRepository { 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 = ?") + sqlx::query("UPDATE team_tasks SET blocked_by = ?, updated_at = ? WHERE team_id = ? AND id = ?") .bind(&new_blocked_by) .bind(now_ms()) + .bind(team_id) .bind(task_id) .execute(&mut *tx) .await?; @@ -408,11 +441,16 @@ impl ITeamRepository for SqliteTeamRepository { 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/team.rs b/crates/aionui-db/src/repository/team.rs index e7b1c2095..da0d4f694 100644 --- a/crates/aionui-db/src/repository/team.rs +++ b/crates/aionui-db/src/repository/team.rs @@ -33,21 +33,24 @@ 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 ────────────────────────────────────────────────────── @@ -62,7 +65,7 @@ pub trait ITeamRepository: Send + Sync { async fn peek_unread(&self, 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, team_id: &str, ids: &[String]) -> Result<(), DbError>; /// Returns message history for an agent, optionally limited. /// Messages are ordered by `created_at` ascending. @@ -74,7 +77,7 @@ pub trait ITeamRepository: Send + Sync { ) -> 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 ──────────────────────────────────────────────────────── @@ -86,19 +89,24 @@ pub trait ITeamRepository: Send + Sync { /// 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, 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>; /// 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, 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, + 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/tests/cron_repository.rs b/crates/aionui-db/tests/cron_repository.rs index 0c90c1e2b..bbf92992f 100644 --- a/crates/aionui-db/tests/cron_repository.rs +++ b/crates/aionui-db/tests/cron_repository.rs @@ -71,6 +71,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] @@ -177,6 +199,62 @@ async fn cj7_list_by_conversation() { assert_eq!(conv2.len(), 1); } +#[tokio::test] +async fn scoped_crud_filters_by_conversation_owner() { + 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.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 scheduler_enabled_scan_requires_conversation_owner() { + let (r, _db) = repo().await; + r.insert(&make_job("cron_owned")).await.unwrap(); + let mut orphan = make_job("cron_orphan"); + orphan.conversation_id = "missing_conversation".into(); + let orphan_insert = r.insert(&orphan).await; + assert!(orphan_insert.is_err()); + + let enabled = r.list_enabled().await.unwrap(); + + assert!(enabled.iter().any(|job| job.id == "cron_owned")); +} + #[tokio::test] async fn cj8_update_name_and_enabled() { let (r, _db) = repo().await; diff --git a/crates/aionui-db/tests/team_repository.rs b/crates/aionui-db/tests/team_repository.rs index f5cf8aa90..1cbfc6fa9 100644 --- a/crates/aionui-db/tests/team_repository.rs +++ b/crates/aionui-db/tests/team_repository.rs @@ -86,7 +86,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 +99,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 +116,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 +143,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 +187,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 +198,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 +208,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 +219,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 +229,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 +239,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 +249,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 +258,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 +273,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(_)))); } @@ -377,7 +414,7 @@ async fn delete_mailbox_by_team() { repo.write_message(&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(); @@ -388,6 +425,29 @@ async fn delete_mailbox_by_team() { 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(&make_mailbox_msg("m1", "t1", "a1", "a2", "message")) + .await + .unwrap(); + repo.write_message(&make_mailbox_msg("m2", "t2", "a1", "a2", "message")) + .await + .unwrap(); + + repo.mark_read_batch("t1", &["m2".to_owned()]).await.unwrap(); + assert!(!repo.get_history("t2", "a1", None).await.unwrap()[0].read); + + repo.delete_mailbox_by_team("user-b", "t1").await.unwrap(); + assert_eq!(repo.get_history("t1", "a1", None).await.unwrap().len(), 1); +} + // ── Task Board Tests ───────────────────────────────────────────────── #[tokio::test] @@ -444,6 +504,7 @@ async fn update_task_status() { repo.create_task(&task).await.unwrap(); repo.update_task( + "t1", "tk1", &UpdateTaskParams { status: Some("in_progress".into()), @@ -466,6 +527,7 @@ async fn update_task_description_and_owner() { repo.create_task(&task).await.unwrap(); repo.update_task( + "t1", "tk1", &UpdateTaskParams { description: Some("New description".into()), @@ -486,6 +548,7 @@ async fn update_nonexistent_task_returns_not_found() { let (repo, _db) = repo().await; let result = repo .update_task( + "t1", "nonexistent", &UpdateTaskParams { status: Some("completed".into()), @@ -509,14 +572,14 @@ async fn append_to_blocks_and_remove_from_blocked_by() { repo.create_task(&task_b).await.unwrap(); // Append tkB to taskA's blocks - repo.append_to_blocks("tkA", "tkB").await.unwrap(); + repo.append_to_blocks("t1", "tkA", "tkB").await.unwrap(); let a = repo.find_task_by_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("t1", "tkB", "tkA").await.unwrap(); let b = repo.find_task_by_id("t1", "tkB").await.unwrap().unwrap(); let blocked_by: Vec = serde_json::from_str(&b.blocked_by).unwrap(); @@ -531,8 +594,8 @@ async fn append_to_blocks_idempotent() { let task = make_task("tkA", "t1", "Task A"); repo.create_task(&task).await.unwrap(); - repo.append_to_blocks("tkA", "tkB").await.unwrap(); - repo.append_to_blocks("tkA", "tkB").await.unwrap(); + repo.append_to_blocks("t1", "tkA", "tkB").await.unwrap(); + repo.append_to_blocks("t1", "tkA", "tkB").await.unwrap(); let a = repo.find_task_by_id("t1", "tkA").await.unwrap().unwrap(); let blocks: Vec = serde_json::from_str(&a.blocks).unwrap(); @@ -555,12 +618,12 @@ async fn multi_dependency_unblock() { repo.create_task(&task_b).await.unwrap(); repo.create_task(&task_c).await.unwrap(); - repo.append_to_blocks("tkA", "tkB").await.unwrap(); - repo.append_to_blocks("tkA", "tkC").await.unwrap(); + repo.append_to_blocks("t1", "tkA", "tkB").await.unwrap(); + repo.append_to_blocks("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("t1", "tkB", "tkA").await.unwrap(); + repo.remove_from_blocked_by("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(); @@ -584,7 +647,7 @@ async fn partial_unblock_preserves_other_blockers() { repo.create_task(&task_b).await.unwrap(); // Complete A only - repo.remove_from_blocked_by("tkB", "tkA").await.unwrap(); + repo.remove_from_blocked_by("t1", "tkB", "tkA").await.unwrap(); let b = repo.find_task_by_id("t1", "tkB").await.unwrap().unwrap(); let blocked_by: Vec = serde_json::from_str(&b.blocked_by).unwrap(); @@ -601,6 +664,7 @@ async fn no_blocks_task_completes_cleanly() { // Complete without any blocks to unblock repo.update_task( + "t1", "tkA", &UpdateTaskParams { status: Some("completed".into()), @@ -625,7 +689,7 @@ async fn delete_tasks_by_team() { repo.create_task(&make_task("tk1", "t1", "T1 Task")).await.unwrap(); repo.create_task(&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(); assert!(t1_tasks.is_empty()); @@ -634,6 +698,35 @@ async fn delete_tasks_by_team() { 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(&make_task("tk1", "t1", "A Task")).await.unwrap(); + repo.create_task(&make_task("tk2", "t2", "B Task")).await.unwrap(); + + let update = repo + .update_task( + "t1", + "tk2", + &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("t1").await.unwrap().len(), 1); + assert_eq!(repo.list_tasks("t2").await.unwrap().len(), 1); +} + #[tokio::test] async fn tasks_contain_dependency_info() { let (repo, _db) = repo().await; @@ -645,7 +738,7 @@ async fn tasks_contain_dependency_info() { repo.create_task(&task_a).await.unwrap(); repo.create_task(&task_b).await.unwrap(); - repo.append_to_blocks("tkA", "tkB").await.unwrap(); + repo.append_to_blocks("t1", "tkA", "tkB").await.unwrap(); let tasks = repo.list_tasks("t1").await.unwrap(); assert_eq!(tasks.len(), 2); @@ -674,12 +767,12 @@ async fn delete_team_cascades_mailbox_and_tasks() { repo.create_task(&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(); assert!(mail.is_empty()); @@ -698,7 +791,7 @@ async fn task_blocked_by_blocks_bidirectional_consistency() { repo.create_task(&task_a).await.unwrap(); repo.create_task(&task_b).await.unwrap(); - repo.append_to_blocks("tkA", "tkB").await.unwrap(); + repo.append_to_blocks("t1", "tkA", "tkB").await.unwrap(); // Verify bidirectional link let a = repo.find_task_by_id("t1", "tkA").await.unwrap().unwrap(); 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/mailbox.rs b/crates/aionui-team/src/mailbox.rs index 060170144..dc4d14298 100644 --- a/crates/aionui-team/src/mailbox.rs +++ b/crates/aionui-team/src/mailbox.rs @@ -90,8 +90,8 @@ impl Mailbox { } /// 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(team_id, ids).await?; Ok(()) } @@ -109,7 +109,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) } @@ -131,7 +131,12 @@ impl Mailbox { } 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/provisioning.rs b/crates/aionui-team/src/provisioning.rs index 75931fee4..4dad1196e 100644 --- a/crates/aionui-team/src/provisioning.rs +++ b/crates/aionui-team/src/provisioning.rs @@ -343,7 +343,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)?; @@ -685,8 +685,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), diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 68daf8a04..7504a6129 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -293,14 +293,9 @@ impl TeamSessionService { async fn load_owned_team(&self, user_id: &str, team_id: &str) -> Result { let row = self .repo - .get_team(team_id) + .get_team(user_id, 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)?) } @@ -351,7 +346,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"); @@ -493,9 +488,9 @@ 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); @@ -509,6 +504,7 @@ impl TeamSessionService { self.repo .update_team( + user_id, team_id, &UpdateTeamParams { name: Some(name.to_owned()), @@ -535,14 +531,9 @@ impl TeamSessionService { let row = self .repo - .get_team(team_id) + .get_team(user_id, 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 mut team = Team::from_row(&row)?; let agent = self.provisioner().add_agent(user_id, &row, &mut team, req).await?; @@ -648,6 +639,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), @@ -781,6 +773,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), @@ -811,15 +804,10 @@ 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.repo + .get_team(user_id, team_id) + .await? + .ok_or_else(|| TeamError::TeamNotFound(team_id.to_owned()))?; self.ensure_session_inner(team_id).await } @@ -831,7 +819,7 @@ 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( @@ -1131,7 +1119,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())), }; @@ -1220,14 +1208,9 @@ impl TeamSessionService { ) -> Result { let row = self .repo - .get_team(team_id) + .get_team(user_id, 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 team = Team::from_row(&row)?; let member = team.agents.iter().any(|agent| agent.conversation_id == conversation_id); @@ -1613,7 +1596,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 { @@ -1722,16 +1705,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(), @@ -2001,6 +1978,7 @@ impl TeamSessionService { let provisioner = self.provisioner(); self.repo .update_team( + user_id, team_id, &UpdateTeamParams { session_mode: Some(mode.to_owned()), @@ -2899,7 +2877,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 @@ -2949,7 +2931,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", @@ -2972,6 +2958,7 @@ mod tests { ]) .to_string(); repo.update_team( + "user-test", &created.id, &aionui_db::UpdateTeamParams { agents: Some(row.agents), @@ -3029,7 +3016,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(); @@ -3179,6 +3170,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/session.rs b/crates/aionui-team/src/session.rs index 48453cf52..07ebba499 100644 --- a/crates/aionui-team/src/session.rs +++ b/crates/aionui-team/src/session.rs @@ -316,7 +316,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 @@ -604,7 +606,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(()) } @@ -777,7 +781,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!( @@ -928,7 +932,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 { @@ -1280,7 +1286,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 @@ -1589,7 +1597,10 @@ pub(crate) async fn attach_member_runtime( }, ); if !update.terminal_message_ids.is_empty() { - let _ = session.mailbox.mark_read_batch(&update.terminal_message_ids).await; + let _ = session + .mailbox + .mark_read_batch(&session.team.id, &update.terminal_message_ids) + .await; } service.refresh_member_runtime_status(&session).await; if let Err(notify_error) = session diff --git a/crates/aionui-team/src/task_board.rs b/crates/aionui-team/src/task_board.rs index 54b8f83a4..35cc43bca 100644 --- a/crates/aionui-team/src/task_board.rs +++ b/crates/aionui-team/src/task_board.rs @@ -64,7 +64,7 @@ impl TaskBoard { self.repo.create_task(&row).await?; for dep_id in blocked_by { - self.repo.append_to_blocks(dep_id, &task_id).await?; + self.repo.append_to_blocks(team_id, dep_id, &task_id).await?; } debug!(team_id, task_id = %task_id, subject, "task created"); @@ -87,10 +87,10 @@ 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(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 @@ -110,11 +110,16 @@ impl TaskBoard { 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(team_id, downstream_id, completed_task_id) .await?; debug!( completed = completed_task_id, diff --git a/crates/aionui-team/src/test_utils.rs b/crates/aionui-team/src/test_utils.rs index ff409483b..b8da93626 100644 --- a/crates/aionui-team/src/test_utils.rs +++ b/crates/aionui-team/src/test_utils.rs @@ -30,19 +30,34 @@ 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(()) } @@ -80,10 +95,10 @@ impl ITeamRepository for MockTeamRepo { Ok(result) } - async fn mark_read_batch(&self, ids: &[String]) -> Result<(), DbError> { + async fn mark_read_batch(&self, 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; } } @@ -108,7 +123,7 @@ 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(()) } @@ -130,12 +145,12 @@ impl ITeamRepository for MockTeamRepo { Ok(found) } - async fn update_task(&self, task_id: &str, params: &UpdateTaskParams) -> Result<(), DbError> { + async fn update_task(&self, 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(); @@ -165,12 +180,12 @@ 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, 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 +193,17 @@ 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, + 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 +211,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 +284,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 +309,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 +333,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 +377,7 @@ pub(crate) mod workspace_harness { async fn list_messages_page( &self, + _user_id: &str, _conv_id: &str, _params: &MessagePageParams, ) -> Result { @@ -348,20 +388,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 +450,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 +465,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,8 +504,11 @@ 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(()) } @@ -472,7 +532,7 @@ pub(crate) mod workspace_harness { Ok(vec![]) } - async fn mark_read_batch(&self, _ids: &[String]) -> Result<(), DbError> { + async fn mark_read_batch(&self, _team_id: &str, _ids: &[String]) -> Result<(), DbError> { Ok(()) } @@ -485,7 +545,7 @@ 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(()) } @@ -497,7 +557,12 @@ pub(crate) mod workspace_harness { Ok(None) } - async fn update_task(&self, _task_id: &str, _params: &aionui_db::UpdateTaskParams) -> Result<(), DbError> { + async fn update_task( + &self, + _team_id: &str, + _task_id: &str, + _params: &aionui_db::UpdateTaskParams, + ) -> Result<(), DbError> { Ok(()) } @@ -505,15 +570,25 @@ pub(crate) mod workspace_harness { Ok(vec![]) } - async fn append_to_blocks(&self, _task_id: &str, _blocked_task_id: &str) -> Result<(), DbError> { + async fn append_to_blocks( + &self, + _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, + _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 +695,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 +770,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(()) } } @@ -1056,6 +1138,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..27d64927d 100644 --- a/crates/aionui-team/tests/common/mod.rs +++ b/crates/aionui-team/tests/common/mod.rs @@ -26,19 +26,22 @@ 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(()) } @@ -70,10 +73,10 @@ impl ITeamRepository for MockTeamRepo { Ok(result) } - async fn mark_read_batch(&self, ids: &[String]) -> Result<(), DbError> { + async fn mark_read_batch(&self, 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; } } @@ -98,7 +101,7 @@ 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(()) } @@ -118,12 +121,12 @@ impl ITeamRepository for MockTeamRepo { Ok(found) } - async fn update_task(&self, task_id: &str, params: &UpdateTaskParams) -> Result<(), DbError> { + async fn update_task(&self, 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(); @@ -150,12 +153,12 @@ 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, 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 +166,17 @@ 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, + 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 +184,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/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/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index 6c68bce43..f8c381bbd 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,8 +801,11 @@ 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(()) } @@ -790,8 +829,8 @@ impl ITeamRepository for FullMockTeamRepo { ) -> Result, DbError> { self.inner.peek_unread(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, team_id: &str, ids: &[String]) -> Result<(), DbError> { + self.inner.mark_read_batch(team_id, ids).await } async fn get_history( &self, @@ -801,8 +840,8 @@ impl ITeamRepository for FullMockTeamRepo { ) -> Result, DbError> { self.inner.get_history(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> { @@ -815,20 +854,32 @@ impl ITeamRepository for FullMockTeamRepo { ) -> Result, DbError> { self.inner.find_task_by_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, + team_id: &str, + task_id: &str, + params: &aionui_db::UpdateTaskParams, + ) -> Result<(), DbError> { + self.inner.update_task(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 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, team_id: &str, task_id: &str, blocked_task_id: &str) -> Result<(), DbError> { + self.inner.append_to_blocks(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, + team_id: &str, + task_id: &str, + unblocked_task_id: &str, + ) -> Result<(), DbError> { + self.inner + .remove_from_blocked_by(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 } } @@ -2206,7 +2257,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)); } @@ -2214,6 +2265,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()), @@ -2488,7 +2540,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"); @@ -3356,7 +3408,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 -------------------------------------------------------------- @@ -3434,7 +3486,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(_)))); } // =========================================================================== @@ -4030,7 +4082,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!( @@ -4430,7 +4482,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"); @@ -4522,7 +4574,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!( @@ -4569,7 +4621,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 @@ -5068,7 +5120,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 ----------- @@ -5082,19 +5134,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] @@ -5225,7 +5268,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(_)))); } // =========================================================================== @@ -5277,7 +5320,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] @@ -5322,7 +5365,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] 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 --------------------------------------------------------- From c0336e484ee30d4221ca313a9309c2fe14106448 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 13:17:28 +0800 Subject: [PATCH 010/145] db: scope configuration repositories by user --- .../028_configuration_user_scope.sql | 177 +++++++++++++ .../aionui-db/src/models/client_preference.rs | 1 + crates/aionui-db/src/models/mcp_server.rs | 1 + crates/aionui-db/src/models/oauth_token.rs | 1 + crates/aionui-db/src/models/provider.rs | 1 + crates/aionui-db/src/models/remote_agent.rs | 1 + .../aionui-db/src/models/system_settings.rs | 4 +- .../src/repository/client_preference.rs | 8 +- crates/aionui-db/src/repository/mcp_server.rs | 33 ++- .../aionui-db/src/repository/oauth_token.rs | 7 +- crates/aionui-db/src/repository/provider.rs | 9 +- .../aionui-db/src/repository/remote_agent.rs | 15 +- crates/aionui-db/src/repository/settings.rs | 7 +- .../repository/sqlite_client_preference.rs | 102 +++++--- .../src/repository/sqlite_diagnostics.rs | 55 ++-- .../src/repository/sqlite_mcp_server.rs | 234 +++++++++++++----- .../src/repository/sqlite_oauth_token.rs | 112 +++++++-- .../src/repository/sqlite_provider.rs | 107 ++++++-- .../src/repository/sqlite_remote_agent.rs | 126 +++++++--- .../src/repository/sqlite_settings.rs | 71 ++++-- .../tests/client_preference_repository.rs | 69 +++--- .../tests/feedback_diagnostics_repository.rs | 22 +- .../aionui-db/tests/mcp_server_repository.rs | 112 +++++---- .../aionui-db/tests/oauth_token_repository.rs | 46 ++-- crates/aionui-db/tests/provider_repository.rs | 28 ++- .../tests/remote_agent_repository.rs | 55 ++-- crates/aionui-db/tests/settings_repository.rs | 37 ++- 27 files changed, 1064 insertions(+), 377 deletions(-) create mode 100644 crates/aionui-db/migrations/028_configuration_user_scope.sql diff --git a/crates/aionui-db/migrations/028_configuration_user_scope.sql b/crates/aionui-db/migrations/028_configuration_user_scope.sql new file mode 100644 index 000000000..23553942c --- /dev/null +++ b/crates/aionui-db/migrations/028_configuration_user_scope.sql @@ -0,0 +1,177 @@ +-- Migration 028: scope independent configuration roots by user. + +PRAGMA foreign_keys = OFF; + +CREATE TABLE providers_new ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + platform TEXT NOT NULL, + name TEXT NOT NULL, + base_url TEXT NOT NULL, + api_key_encrypted TEXT NOT NULL, + models TEXT NOT NULL DEFAULT '[]', + enabled INTEGER NOT NULL DEFAULT 1, + capabilities TEXT NOT NULL DEFAULT '[]', + context_limit INTEGER, + model_protocols TEXT, + model_enabled TEXT, + model_health TEXT, + bedrock_config TEXT, + is_full_url INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); + +INSERT INTO providers_new ( + 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 +) +SELECT + id, 'system_default_user', 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 +FROM providers; + +DROP TABLE providers; +ALTER TABLE providers_new RENAME TO providers; +CREATE INDEX IF NOT EXISTS idx_providers_user_platform ON providers(user_id, platform); + +CREATE TABLE remote_agents_new ( + id TEXT PRIMARY KEY NOT NULL, + user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), + name TEXT NOT NULL, + protocol TEXT NOT NULL, + url TEXT NOT NULL, + auth_type TEXT NOT NULL, + auth_token TEXT, + allow_insecure INTEGER NOT NULL DEFAULT 0, + avatar TEXT, + description TEXT, + device_id TEXT, + device_public_key TEXT, + device_private_key TEXT, + device_token TEXT, + status TEXT NOT NULL DEFAULT 'unknown', + last_connected_at INTEGER, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL +); + +INSERT INTO remote_agents_new ( + 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 +) +SELECT + id, 'system_default_user', 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 +FROM remote_agents; + +DROP TABLE remote_agents; +ALTER TABLE remote_agents_new RENAME TO remote_agents; +CREATE INDEX IF NOT EXISTS idx_remote_agents_user_status ON remote_agents(user_id, status); + +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; + +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; + +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; + +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; + +DROP TABLE client_preferences; +ALTER TABLE client_preferences_new RENAME TO client_preferences; + +PRAGMA foreign_keys = ON; 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/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/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/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/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/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/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_diagnostics.rs b/crates/aionui-db/src/repository/sqlite_diagnostics.rs index 8793717cf..4970ce45b 100644 --- a/crates/aionui-db/src/repository/sqlite_diagnostics.rs +++ b/crates/aionui-db/src/repository/sqlite_diagnostics.rs @@ -528,17 +528,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 +731,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 +739,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 +783,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 +816,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 +974,18 @@ impl SqliteFeedbackDiagnosticsRepository { .bind(&request.user_id) .fetch_one(&self.pool) .await?; - let provider_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM providers") + 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") .fetch_one(&self.pool) .await?; - let active_mcp_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM mcp_servers WHERE deleted_at IS NULL") - .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, @@ -985,7 +1000,7 @@ impl SqliteFeedbackDiagnosticsRepository { "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?, + "provider_health": self.collect_global_provider_health(&request.user_id).await?, }), )) } @@ -1451,14 +1466,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 +1531,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..511dd4997 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", ) + .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,12 +77,14 @@ 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); @@ -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/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/feedback_diagnostics_repository.rs b/crates/aionui-db/tests/feedback_diagnostics_repository.rs index 820258187..c8df0eca3 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) 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..ca3dadfc4 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"), @@ -203,6 +208,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 +222,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 +238,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 +251,7 @@ async fn update_advances_updated_at() { let updated = r .update( + USER_ID, &created.id, UpdateProviderParams { name: Some("Changed"), @@ -264,14 +272,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 +295,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); } From 5f6a9ae651f00a1f593be806d76ab57ff388eafb Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 14:11:25 +0800 Subject: [PATCH 011/145] system: scope user configuration APIs --- crates/aionui-app/src/router/state.rs | 1 + crates/aionui-app/tests/stt_stream_e2e.rs | 2 +- crates/aionui-system/src/client_pref.rs | 132 ++++++++++++------ crates/aionui-system/src/model_fetcher/mod.rs | 32 +++-- crates/aionui-system/src/provider.rs | 69 +++++---- crates/aionui-system/src/routes.rs | 43 ++++-- crates/aionui-system/src/settings.rs | 52 +++++-- .../aionui-system/tests/model_fetch_routes.rs | 21 ++- crates/aionui-system/tests/provider_routes.rs | 36 ++++- crates/aionui-system/tests/settings_routes.rs | 27 +++- 10 files changed, 300 insertions(+), 115 deletions(-) diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index 7277ae94b..9ac3f5db1 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -384,6 +384,7 @@ pub fn build_system_state(services: &AppServices) -> SystemRouterState { client_pref_service: ClientPrefService::with_keep_awake_controller( Arc::new(SqliteClientPreferenceRepository::new(pool.clone())), Arc::new(aionui_system::SystemKeepAwakeController::new()), + "system_default_user", ), provider_service: ProviderService::new(provider_repo.clone(), encryption_key), model_fetch_service: ModelFetchService::new(provider_repo, encryption_key, http_client.clone()), 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-system/src/client_pref.rs b/crates/aionui-system/src/client_pref.rs index 659d6129c..07c4c0fd8 100644 --- a/crates/aionui-system/src/client_pref.rs +++ b/crates/aionui-system/src/client_pref.rs @@ -29,20 +29,25 @@ 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 { repo, keep_awake_controller, }; - service.restore_keep_awake_from_preferences(); + service.restore_keep_awake_from_preferences(keep_awake_restore_user_id.into()); 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 +81,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 +119,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 +132,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 +147,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}"))) { @@ -152,10 +161,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}")))?; @@ -181,14 +190,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"); @@ -249,6 +258,8 @@ mod tests { use tracing::Level; use tracing_subscriber::fmt; + const TEST_USER_ID: &str = "user-1"; + #[derive(Clone)] struct SharedBuf(Arc>>); @@ -290,6 +301,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) @@ -297,6 +317,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 { @@ -344,7 +373,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()); } @@ -353,9 +382,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)); } @@ -364,9 +393,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)); } @@ -375,9 +404,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")); } @@ -387,13 +416,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")); } @@ -405,9 +434,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)); @@ -425,10 +454,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(); }); @@ -458,13 +487,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")); } @@ -473,7 +502,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(_))); } @@ -482,7 +511,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(_))); } @@ -493,14 +522,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)); @@ -513,10 +542,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)); } @@ -526,14 +558,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)); } @@ -544,11 +579,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)); } @@ -562,11 +600,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)); } @@ -575,10 +616,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() { 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 86f031740..42cb70a30 100644 --- a/crates/aionui-system/src/routes.rs +++ b/crates/aionui-system/src/routes.rs @@ -108,8 +108,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))) } @@ -128,12 +133,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))) @@ -150,6 +156,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| { @@ -163,7 +170,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))) @@ -171,12 +178,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())) @@ -188,47 +196,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))) 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/model_fetch_routes.rs b/crates/aionui-system/tests/model_fetch_routes.rs index ad96cab81..a3405fcdb 100644 --- a/crates/aionui-system/tests/model_fetch_routes.rs +++ b/crates/aionui-system/tests/model_fetch_routes.rs @@ -13,6 +13,7 @@ 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, @@ -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,17 @@ 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(), + }); + req } // --------------------------------------------------------------------------- diff --git a/crates/aionui-system/tests/provider_routes.rs b/crates/aionui-system/tests/provider_routes.rs index 6f15eb46f..24b897a3e 100644 --- a/crates/aionui-system/tests/provider_routes.rs +++ b/crates/aionui-system/tests/provider_routes.rs @@ -13,6 +13,7 @@ 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, @@ -27,6 +28,7 @@ use aionui_system::{ // --------------------------------------------------------------------------- const TEST_ENCRYPTION_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())); @@ -47,6 +49,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) } @@ -57,24 +68,39 @@ 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() + let mut req = Request::builder().method("GET").uri(uri).body(Body::empty()).unwrap(); + req.extensions_mut().insert(CurrentUser { + id: TEST_USER_ID.to_owned(), + username: TEST_USER_ID.to_owned(), + }); + req } fn json_request(method: &str, uri: &str, body: serde_json::Value) -> Request { - Request::builder() + 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: TEST_USER_ID.to_owned(), + username: TEST_USER_ID.to_owned(), + }); + req } fn delete_request(uri: &str) -> Request { - Request::builder() + let mut req = Request::builder() .method("DELETE") .uri(uri) .body(Body::empty()) - .unwrap() + .unwrap(); + req.extensions_mut().insert(CurrentUser { + id: TEST_USER_ID.to_owned(), + username: TEST_USER_ID.to_owned(), + }); + req } fn sample_create_body() -> serde_json::Value { diff --git a/crates/aionui-system/tests/settings_routes.rs b/crates/aionui-system/tests/settings_routes.rs index c89566903..a62a666da 100644 --- a/crates/aionui-system/tests/settings_routes.rs +++ b/crates/aionui-system/tests/settings_routes.rs @@ -12,6 +12,7 @@ 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, @@ -26,6 +27,7 @@ use aionui_system::{ // --------------------------------------------------------------------------- const TEST_ENCRYPTION_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())); @@ -46,6 +48,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); (settings_routes(state), db) } @@ -56,16 +67,26 @@ 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() + let mut req = Request::builder().method("GET").uri(uri).body(Body::empty()).unwrap(); + req.extensions_mut().insert(CurrentUser { + id: TEST_USER_ID.to_owned(), + username: TEST_USER_ID.to_owned(), + }); + req } fn json_request(method: &str, uri: &str, body: serde_json::Value) -> Request { - Request::builder() + 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: TEST_USER_ID.to_owned(), + username: TEST_USER_ID.to_owned(), + }); + req } // =========================================================================== From 40cf7660c4b1d514bcac60e38d619e231c0c7287 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 14:12:32 +0800 Subject: [PATCH 012/145] mcp: scope server state by current user --- Cargo.lock | 1 + crates/aionui-mcp/Cargo.toml | 1 + crates/aionui-mcp/src/adapter.rs | 14 +- crates/aionui-mcp/src/adapters/aionrs.rs | 2 +- crates/aionui-mcp/src/adapters/aionui.rs | 48 +- crates/aionui-mcp/src/adapters/claude.rs | 2 +- crates/aionui-mcp/src/adapters/codebuddy.rs | 2 +- crates/aionui-mcp/src/adapters/codex.rs | 2 +- crates/aionui-mcp/src/adapters/gemini.rs | 2 +- crates/aionui-mcp/src/adapters/opencode.rs | 2 +- crates/aionui-mcp/src/adapters/qwen.rs | 2 +- crates/aionui-mcp/src/oauth_service.rs | 200 ++++--- crates/aionui-mcp/src/routes.rs | 74 ++- crates/aionui-mcp/src/service.rs | 521 +++++++++++------- crates/aionui-mcp/src/sync_service.rs | 36 +- crates/aionui-mcp/src/types.rs | 1 + .../aionui-mcp/tests/adapter_integration.rs | 20 +- .../tests/file_adapter_integration.rs | 60 +- crates/aionui-mcp/tests/oauth_integration.rs | 67 ++- .../aionui-mcp/tests/service_integration.rs | 67 ++- crates/aionui-mcp/tests/sync_integration.rs | 8 +- crates/aionui-mcp/tests/types_integration.rs | 3 + 22 files changed, 745 insertions(+), 390 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9d2556db2..57ac1d3f4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -693,6 +693,7 @@ name = "aionui-mcp" version = "0.1.50" dependencies = [ "aionui-api-types", + "aionui-auth", "aionui-common", "aionui-db", "aionui-realtime", 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/oauth_service.rs b/crates/aionui-mcp/src/oauth_service.rs index a2259598c..911bf56cf 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] @@ -659,7 +686,11 @@ mod tests { #[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 +701,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 +714,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 +729,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 +744,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 +768,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 +781,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 +805,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 +818,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 +842,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 +856,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 +933,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..3d315c730 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)?; @@ -178,7 +211,7 @@ async fn test_connection( 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 +260,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 +277,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 +295,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 +310,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..33db49ee2 100644 --- a/crates/aionui-mcp/tests/service_integration.rs +++ b/crates/aionui-mcp/tests/service_integration.rs @@ -5,6 +5,8 @@ use std::collections::HashMap; use std::sync::Arc; +const TEST_USER_ID: &str = "system_default_user"; + use aionui_api_types::{ BatchImportMcpServersRequest, CreateMcpServerRequest, ImportMcpServerRequest, McpTransport, UpdateMcpServerRequest, }; @@ -80,21 +82,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 +110,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,21 +127,21 @@ 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(_))); } @@ -150,10 +152,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 +174,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 +202,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 +227,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 +245,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 +272,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 +304,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 +321,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 +329,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 +343,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 +351,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()), From ab70abb3685086bdf7a8486721e5532f9d82f014 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 14:12:52 +0800 Subject: [PATCH 013/145] ai-agent: resolve runtime config per user --- crates/aionui-ai-agent/src/factory/acp.rs | 52 ++++++++++++------ crates/aionui-ai-agent/src/factory/aionrs.rs | 53 ++++++++++++++----- crates/aionui-ai-agent/src/factory/context.rs | 2 + crates/aionui-ai-agent/src/routes/agent.rs | 12 ++--- crates/aionui-ai-agent/src/routes/remote.rs | 44 ++++++++++----- crates/aionui-ai-agent/src/services/agent.rs | 10 ++-- .../src/services/availability/mod.rs | 23 +++++--- .../src/services/provider_health.rs | 18 +++++-- crates/aionui-ai-agent/src/services/remote.rs | 38 ++++++++----- .../tests/agent_availability_integration.rs | 7 ++- .../tests/factory_provider_integration.rs | 5 +- 11 files changed, 187 insertions(+), 77 deletions(-) diff --git a/crates/aionui-ai-agent/src/factory/acp.rs b/crates/aionui-ai-agent/src/factory/acp.rs index 399c1cc77..ab99916c6 100644 --- a/crates/aionui-ai-agent/src/factory/acp.rs +++ b/crates/aionui-ai-agent/src/factory/acp.rs @@ -79,6 +79,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, ) @@ -294,12 +295,13 @@ async fn resolve_builtin_managed_acp_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, @@ -522,6 +524,8 @@ mod tests { path::{Path, PathBuf}, }; + const TEST_USER_ID: &str = "user-1"; + fn make_row( name: &str, transport_type: &str, @@ -531,6 +535,7 @@ mod tests { ) -> McpServerRow { McpServerRow { id: format!("mcp_{name}"), + user_id: TEST_USER_ID.to_owned(), name: name.to_owned(), description: None, enabled, @@ -849,26 +854,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( @@ -879,29 +893,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!() } } @@ -928,7 +950,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"), @@ -947,7 +969,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()); } @@ -966,7 +988,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"), @@ -991,7 +1013,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] { @@ -1018,7 +1040,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/aionrs.rs b/crates/aionui-ai-agent/src/factory/aionrs.rs index 07d568f46..cf33b7377 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(), ) @@ -81,7 +82,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 +443,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, @@ -761,6 +763,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 +864,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 +887,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 +916,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 +937,7 @@ mod tests { async fn update_status( &self, + _user_id: &str, _id: &str, _status: &str, _last_connected: Option, @@ -929,7 +945,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 +972,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); 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/routes/agent.rs b/crates/aionui-ai-agent/src/routes/agent.rs index ad6614cd4..f8c42fe45 100644 --- a/crates/aionui-ai-agent/src/routes/agent.rs +++ b/crates/aionui-ai-agent/src/routes/agent.rs @@ -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)?, ))) @@ -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/services/agent.rs b/crates/aionui-ai-agent/src/services/agent.rs index c0789d81b..b07d956f3 100644 --- a/crates/aionui-ai-agent/src/services/agent.rs +++ b/crates/aionui-ai-agent/src/services/agent.rs @@ -102,19 +102,21 @@ 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 { @@ -163,7 +165,7 @@ impl AgentService { .await .map_err(|e| AgentError::internal(format!("repo.update_agent_overrides: {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 { diff --git a/crates/aionui-ai-agent/src/services/availability/mod.rs b/crates/aionui-ai-agent/src/services/availability/mod.rs index c47acdf9b..4803169bb 100644 --- a/crates/aionui-ai-agent/src/services/availability/mod.rs +++ b/crates/aionui-ai-agent/src/services/availability/mod.rs @@ -53,7 +53,7 @@ impl AgentAvailabilityService { self.registry.list_management_rows().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) @@ -71,6 +71,7 @@ impl AgentAvailabilityService { &self.registry, &self.provider_repo, &meta, + user_id, AgentSnapshotCheckKind::Manual, ) .await; @@ -155,6 +156,7 @@ async fn run_probe( registry: &Arc, provider_repo: &Arc, meta: &AgentMetadata, + user_id: &str, kind: AgentSnapshotCheckKind, ) -> AvailabilitySnapshot { let started_at = now_ms(); @@ -233,7 +235,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) }; @@ -279,8 +281,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, @@ -393,6 +396,7 @@ mod tests { fn enabled_provider_params() -> CreateProviderParams<'static> { CreateProviderParams { id: None, + user_id: "system_default_user", platform: "openai", name: "OpenAI", base_url: "https://api.openai.com", @@ -415,7 +419,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, "system_default_user").await; assert_eq!(status, AgentSnapshotCheckStatus::Offline); assert_eq!(code.as_deref(), Some("no_provider")); @@ -427,7 +431,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, "system_default_user").await; assert_eq!(status, AgentSnapshotCheckStatus::Online); assert!(code.is_none()); @@ -622,7 +626,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/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/tests/agent_availability_integration.rs b/crates/aionui-ai-agent/tests/agent_availability_integration.rs index e799e0d8e..03b1a8483 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 TEST_USER_ID: &str = "user-1"; + struct NoopBroadcaster; impl EventBroadcaster for NoopBroadcaster { @@ -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(TEST_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(); diff --git a/crates/aionui-ai-agent/tests/factory_provider_integration.rs b/crates/aionui-ai-agent/tests/factory_provider_integration.rs index 2bbb52df8..f2e8eb562 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", @@ -91,7 +94,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, }, From bdfd8e54fbfa0f9e96d0ac2f3bc5687cc0993de1 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 14:13:06 +0800 Subject: [PATCH 014/145] channel: use user scoped assistant settings --- Cargo.lock | 1 + crates/aionui-channel/Cargo.toml | 1 + crates/aionui-channel/src/action.rs | 20 ++- crates/aionui-channel/src/channel_settings.rs | 123 +++++++++++++----- crates/aionui-channel/src/message_service.rs | 9 +- crates/aionui-channel/src/routes.rs | 18 ++- .../tests/message_service_integration.rs | 51 ++++++-- crates/aionui-conversation/src/service.rs | 4 +- 8 files changed, 164 insertions(+), 63 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 57ac1d3f4..697cbc5d8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -515,6 +515,7 @@ version = "0.1.50" dependencies = [ "aionui-ai-agent", "aionui-api-types", + "aionui-auth", "aionui-common", "aionui-conversation", "aionui-db", 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..e953a9fd8 100644 --- a/crates/aionui-channel/src/action.rs +++ b/crates/aionui-channel/src/action.rs @@ -86,7 +86,7 @@ impl ActionExecutor { } // 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(&internal_user_id, msg.platform).await?; let session = self .session_mgr .get_or_create_session(&internal_user_id, chat_id, &agent_config.agent_type, None) @@ -224,7 +224,10 @@ 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(internal_user_id, action.context.platform) + .await?; let session = self .session_mgr .reset_session(user_id, chat_id, &agent_config.agent_type, None) @@ -251,7 +254,10 @@ 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(internal_user_id, action.context.platform) + .await?; let session = self .session_mgr .get_or_create_session(user_id, chat_id, &agent_config.agent_type, None) @@ -673,16 +679,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(()) } } diff --git a/crates/aionui-channel/src/channel_settings.rs b/crates/aionui-channel/src/channel_settings.rs index be223a69b..d5ede8d7a 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()); @@ -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; @@ -197,10 +206,11 @@ 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; @@ -217,24 +227,30 @@ 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(()) } @@ -506,6 +522,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 +544,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 +557,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 +571,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 +583,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(()) @@ -731,7 +751,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 +764,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 +777,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 +792,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 +805,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 +818,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 +835,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 +855,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 +871,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 +889,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 +901,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 +918,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 +928,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 +938,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 +957,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 +967,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 +985,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 +1010,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 +1031,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 +1056,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 +1077,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 +1102,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 +1127,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/message_service.rs b/crates/aionui-channel/src/message_service.rs index 37b627fbf..7a4608210 100644 --- a/crates/aionui-channel/src/message_service.rs +++ b/crates/aionui-channel/src/message_service.rs @@ -127,10 +127,13 @@ impl ChannelMessageService { let source = platform_to_source(platform); let agent_config = self .settings - .get_agent_config(platform) + .get_agent_config(&self.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(&self.owner_user_id, platform) + .await?; let assistant_id = assistant_setting .as_ref() .and_then(|setting| setting.assistant_id.as_deref()) @@ -141,7 +144,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(&self.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() { diff --git a/crates/aionui-channel/src/routes.rs b/crates/aionui-channel/src/routes.rs index cc190d172..93127dbaf 100644 --- a/crates/aionui-channel/src/routes.rs +++ b/crates/aionui-channel/src/routes.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use axum::Router; use axum::extract::rejection::JsonRejection; -use axum::extract::{Json, Path, State}; +use axum::extract::{Extension, Json, Path, State}; use axum::routing::{get, post, put}; use tracing::warn; @@ -14,6 +14,7 @@ use aionui_api_types::{ EnablePluginRequest, PairingRequestResponse, PluginStatusResponse, RejectPairingRequest, RevokeUserRequest, SyncChannelSettingsRequest, TestPluginRequest, TestPluginResponse, }; +use aionui_auth::CurrentUser; use aionui_common::ApiError; use aionui_db::{DbError, IChannelRepository}; use aionui_extension::{ExtensionRegistry, ResolvedChannelPlugin}; @@ -568,18 +569,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,7 +590,10 @@ 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 + .settings_service + .set_assistant_setting(&user.id, platform, &req) + .await?; state.session_manager.clear_all_sessions().await?; Ok(Json(ApiResponse::ok(BridgeResponse { @@ -600,6 +606,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,7 +614,10 @@ 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 + .settings_service + .set_model_setting(&user.id, platform, &req) + .await?; state.session_manager.clear_all_sessions().await?; Ok(Json(ApiResponse::ok(BridgeResponse { diff --git a/crates/aionui-channel/tests/message_service_integration.rs b/crates/aionui-channel/tests/message_service_integration.rs index 085cedc2d..ed68a7e13 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>>, } @@ -306,10 +308,13 @@ 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(); @@ -338,7 +343,7 @@ async fn send_to_agent_persists_assistant_snapshot_for_channel_bound_assistant() .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,7 +351,11 @@ 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) @@ -380,10 +389,13 @@ 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())); @@ -478,11 +490,15 @@ async fn send_to_agent_without_saved_binding_defaults_to_bare_aionrs_assistant() .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,7 +541,10 @@ 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(); @@ -553,6 +572,10 @@ async fn send_to_agent_without_assistant_name_falls_back_to_legacy_channel_name( .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-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index a1a600079..7f3db2979 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -1040,11 +1040,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}")))?, }; From 404c19d57dbc268fbb10f24d0479f312d8a68cab Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 14:31:20 +0800 Subject: [PATCH 015/145] db: consolidate user scope migration --- .../aionui-db/migrations/026_user_scope.sql | 187 +++++++++++++++++- .../migrations/027_aggregate_user_scope.sql | 79 -------- .../028_configuration_user_scope.sql | 177 ----------------- .../src/repository/sqlite_mcp_server.rs | 4 +- 4 files changed, 188 insertions(+), 259 deletions(-) delete mode 100644 crates/aionui-db/migrations/027_aggregate_user_scope.sql delete mode 100644 crates/aionui-db/migrations/028_configuration_user_scope.sql diff --git a/crates/aionui-db/migrations/026_user_scope.sql b/crates/aionui-db/migrations/026_user_scope.sql index 488c98c2e..b9bd2e3f6 100644 --- a/crates/aionui-db/migrations/026_user_scope.sql +++ b/crates/aionui-db/migrations/026_user_scope.sql @@ -1,4 +1,4 @@ --- Migration 026: extend users for local and AionPro identity projection. +-- Migration 026: add user scope for local, aggregate, and configuration data. PRAGMA foreign_keys = OFF; @@ -89,4 +89,189 @@ CREATE UNIQUE INDEX idx_users_external_user 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; + +CREATE TRIGGER IF NOT EXISTS trg_cron_jobs_conversation_parent_insert +BEFORE INSERT ON cron_jobs +FOR EACH ROW +WHEN NEW.conversation_id IS NULL + OR NEW.conversation_id = '' + OR NOT EXISTS (SELECT 1 FROM conversations WHERE id = NEW.conversation_id) +BEGIN + SELECT RAISE(ABORT, 'cron_jobs.conversation_id must reference conversations.id'); +END; + +CREATE TRIGGER IF NOT EXISTS trg_cron_jobs_conversation_parent_update +BEFORE UPDATE OF conversation_id ON cron_jobs +FOR EACH ROW +WHEN NEW.conversation_id IS NULL + OR NEW.conversation_id = '' + OR NOT EXISTS (SELECT 1 FROM conversations WHERE id = NEW.conversation_id) +BEGIN + SELECT RAISE(ABORT, 'cron_jobs.conversation_id must reference conversations.id'); +END; + +CREATE TRIGGER IF NOT EXISTS trg_cron_job_runs_job_parent_insert +BEFORE INSERT ON cron_job_runs +FOR EACH ROW +WHEN NOT EXISTS ( + SELECT 1 + FROM cron_jobs j + JOIN conversations c ON c.id = j.conversation_id + WHERE j.id = NEW.job_id +) +BEGIN + SELECT RAISE(ABORT, 'cron_job_runs.job_id must reference cron_jobs.id with conversation parent'); +END; + +CREATE TRIGGER IF NOT EXISTS trg_cron_job_runs_job_parent_update +BEFORE UPDATE OF job_id ON cron_job_runs +FOR EACH ROW +WHEN NOT EXISTS ( + SELECT 1 + FROM cron_jobs j + JOIN conversations c ON c.id = j.conversation_id + WHERE j.id = NEW.job_id +) +BEGIN + SELECT RAISE(ABORT, 'cron_job_runs.job_id must reference cron_jobs.id with conversation parent'); +END; + +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 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; + +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; + +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; + +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; + +DROP TABLE client_preferences; +ALTER TABLE client_preferences_new RENAME TO client_preferences; + PRAGMA foreign_keys = ON; diff --git a/crates/aionui-db/migrations/027_aggregate_user_scope.sql b/crates/aionui-db/migrations/027_aggregate_user_scope.sql deleted file mode 100644 index 93faf92cd..000000000 --- a/crates/aionui-db/migrations/027_aggregate_user_scope.sql +++ /dev/null @@ -1,79 +0,0 @@ --- Enforce aggregate parent chains for new writes without rewriting existing tables. - -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; - -CREATE TRIGGER IF NOT EXISTS trg_cron_jobs_conversation_parent_insert -BEFORE INSERT ON cron_jobs -FOR EACH ROW -WHEN NEW.conversation_id IS NULL - OR NEW.conversation_id = '' - OR NOT EXISTS (SELECT 1 FROM conversations WHERE id = NEW.conversation_id) -BEGIN - SELECT RAISE(ABORT, 'cron_jobs.conversation_id must reference conversations.id'); -END; - -CREATE TRIGGER IF NOT EXISTS trg_cron_jobs_conversation_parent_update -BEFORE UPDATE OF conversation_id ON cron_jobs -FOR EACH ROW -WHEN NEW.conversation_id IS NULL - OR NEW.conversation_id = '' - OR NOT EXISTS (SELECT 1 FROM conversations WHERE id = NEW.conversation_id) -BEGIN - SELECT RAISE(ABORT, 'cron_jobs.conversation_id must reference conversations.id'); -END; - -CREATE TRIGGER IF NOT EXISTS trg_cron_job_runs_job_parent_insert -BEFORE INSERT ON cron_job_runs -FOR EACH ROW -WHEN NOT EXISTS ( - SELECT 1 - FROM cron_jobs j - JOIN conversations c ON c.id = j.conversation_id - WHERE j.id = NEW.job_id -) -BEGIN - SELECT RAISE(ABORT, 'cron_job_runs.job_id must reference cron_jobs.id with conversation parent'); -END; - -CREATE TRIGGER IF NOT EXISTS trg_cron_job_runs_job_parent_update -BEFORE UPDATE OF job_id ON cron_job_runs -FOR EACH ROW -WHEN NOT EXISTS ( - SELECT 1 - FROM cron_jobs j - JOIN conversations c ON c.id = j.conversation_id - WHERE j.id = NEW.job_id -) -BEGIN - SELECT RAISE(ABORT, 'cron_job_runs.job_id must reference cron_jobs.id with conversation parent'); -END; diff --git a/crates/aionui-db/migrations/028_configuration_user_scope.sql b/crates/aionui-db/migrations/028_configuration_user_scope.sql deleted file mode 100644 index 23553942c..000000000 --- a/crates/aionui-db/migrations/028_configuration_user_scope.sql +++ /dev/null @@ -1,177 +0,0 @@ --- Migration 028: scope independent configuration roots by user. - -PRAGMA foreign_keys = OFF; - -CREATE TABLE providers_new ( - id TEXT PRIMARY KEY NOT NULL, - user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), - platform TEXT NOT NULL, - name TEXT NOT NULL, - base_url TEXT NOT NULL, - api_key_encrypted TEXT NOT NULL, - models TEXT NOT NULL DEFAULT '[]', - enabled INTEGER NOT NULL DEFAULT 1, - capabilities TEXT NOT NULL DEFAULT '[]', - context_limit INTEGER, - model_protocols TEXT, - model_enabled TEXT, - model_health TEXT, - bedrock_config TEXT, - is_full_url INTEGER NOT NULL DEFAULT 0, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -INSERT INTO providers_new ( - 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 -) -SELECT - id, 'system_default_user', 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 -FROM providers; - -DROP TABLE providers; -ALTER TABLE providers_new RENAME TO providers; -CREATE INDEX IF NOT EXISTS idx_providers_user_platform ON providers(user_id, platform); - -CREATE TABLE remote_agents_new ( - id TEXT PRIMARY KEY NOT NULL, - user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id), - name TEXT NOT NULL, - protocol TEXT NOT NULL, - url TEXT NOT NULL, - auth_type TEXT NOT NULL, - auth_token TEXT, - allow_insecure INTEGER NOT NULL DEFAULT 0, - avatar TEXT, - description TEXT, - device_id TEXT, - device_public_key TEXT, - device_private_key TEXT, - device_token TEXT, - status TEXT NOT NULL DEFAULT 'unknown', - last_connected_at INTEGER, - created_at INTEGER NOT NULL, - updated_at INTEGER NOT NULL -); - -INSERT INTO remote_agents_new ( - 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 -) -SELECT - id, 'system_default_user', 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 -FROM remote_agents; - -DROP TABLE remote_agents; -ALTER TABLE remote_agents_new RENAME TO remote_agents; -CREATE INDEX IF NOT EXISTS idx_remote_agents_user_status ON remote_agents(user_id, status); - -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; - -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; - -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; - -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; - -DROP TABLE client_preferences; -ALTER TABLE client_preferences_new RENAME TO client_preferences; - -PRAGMA foreign_keys = ON; diff --git a/crates/aionui-db/src/repository/sqlite_mcp_server.rs b/crates/aionui-db/src/repository/sqlite_mcp_server.rs index 511dd4997..6dcbd529a 100644 --- a/crates/aionui-db/src/repository/sqlite_mcp_server.rs +++ b/crates/aionui-db/src/repository/sqlite_mcp_server.rs @@ -24,7 +24,7 @@ impl SqliteMcpServerRepository { impl IMcpServerRepository for SqliteMcpServerRepository { async fn list(&self, user_id: &str) -> Result, DbError> { let rows = sqlx::query_as::<_, McpServerRow>( - "SELECT * FROM mcp_servers WHERE user_id = ? AND 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) @@ -89,7 +89,7 @@ impl IMcpServerRepository for SqliteMcpServerRepository { 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(); From 6b876fb61daecb7baeb62ca9cb787d295a878976 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 14:49:37 +0800 Subject: [PATCH 016/145] db: add mixed user scope schema --- .../aionui-db/migrations/026_user_scope.sql | 195 ++++++++++++++++++ .../src/repository/sqlite_assistant.rs | 63 +++--- .../aionui-db/src/repository/sqlite_skill.rs | 47 +++-- .../tests/skill_management_schema.rs | 50 ++++- 4 files changed, 316 insertions(+), 39 deletions(-) diff --git a/crates/aionui-db/migrations/026_user_scope.sql b/crates/aionui-db/migrations/026_user_scope.sql index b9bd2e3f6..80b5709b3 100644 --- a/crates/aionui-db/migrations/026_user_scope.sql +++ b/crates/aionui-db/migrations/026_user_scope.sql @@ -274,4 +274,199 @@ FROM client_preferences; DROP TABLE client_preferences; ALTER TABLE client_preferences_new RENAME TO client_preferences; +ALTER TABLE agent_metadata + ADD COLUMN user_id TEXT REFERENCES users(id); +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; + +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); + +ALTER TABLE assistant_definitions + ADD COLUMN user_id TEXT REFERENCES users(id); +UPDATE assistant_definitions +SET user_id = 'system_default_user' +WHERE user_id IS NULL + AND NOT (source = 'builtin' AND owner_type = 'system'); +DROP INDEX IF EXISTS idx_assistant_definitions_source_ref; +DROP INDEX IF EXISTS idx_assistant_definitions_assistant_id; +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_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; + +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; + +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; + +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); + +ALTER TABLE assistant_plugins + ADD COLUMN owner_user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id); +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; + +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); + +ALTER TABLE assistant_pairing_codes + ADD COLUMN owner_user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id); +CREATE INDEX IF NOT EXISTS idx_pairing_codes_owner_status + ON assistant_pairing_codes(owner_user_id, status); + PRAGMA foreign_keys = ON; diff --git a/crates/aionui-db/src/repository/sqlite_assistant.rs b/crates/aionui-db/src/repository/sqlite_assistant.rs index afd11b38c..073beca9b 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 { @@ -269,15 +271,19 @@ 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?; + let row = sqlx::query_as::<_, AssistantOverrideRow>( + "SELECT * FROM assistant_overrides WHERE user_id = ? AND assistant_id = ?", + ) + .bind(DEFAULT_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") + let rows = sqlx::query_as::<_, AssistantOverrideRow>("SELECT * FROM assistant_overrides WHERE user_id = ?") + .bind(DEFAULT_USER_ID) .fetch_all(&self.pool) .await?; Ok(rows) @@ -289,14 +295,15 @@ impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { 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(DEFAULT_USER_ID) .bind(params.assistant_id) .bind(params.enabled) .bind(params.sort_order) @@ -315,7 +322,8 @@ impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { } async fn delete(&self, assistant_id: &str) -> Result { - let result = sqlx::query("DELETE FROM assistant_overrides WHERE assistant_id = ?") + let result = sqlx::query("DELETE FROM assistant_overrides WHERE user_id = ? AND assistant_id = ?") + .bind(DEFAULT_USER_ID) .bind(assistant_id) .execute(&self.pool) .await?; @@ -324,15 +332,16 @@ impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { async fn delete_orphans(&self, 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(DEFAULT_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(DEFAULT_USER_ID); for id in valid_ids { q = q.bind(*id); } @@ -550,8 +559,9 @@ impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { impl IAssistantOverlayRepository for SqliteAssistantOverlayRepository { async fn get(&self, 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(DEFAULT_USER_ID) .bind(assistant_definition_id) .fetch_optional(&self.pool) .await?; @@ -560,8 +570,9 @@ impl IAssistantOverlayRepository for SqliteAssistantOverlayRepository { async fn list(&self) -> 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(DEFAULT_USER_ID) .fetch_all(&self.pool) .await?; Ok(rows) @@ -571,15 +582,17 @@ impl IAssistantOverlayRepository for SqliteAssistantOverlayRepository { 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(DEFAULT_USER_ID) .bind(params.assistant_definition_id) .bind(params.enabled) .bind(params.sort_order) @@ -599,7 +612,8 @@ impl IAssistantOverlayRepository for SqliteAssistantOverlayRepository { } async fn delete(&self, assistant_definition_id: &str) -> Result { - let result = sqlx::query("DELETE FROM assistant_overlays WHERE assistant_definition_id = ?") + let result = sqlx::query("DELETE FROM assistant_overlays WHERE user_id = ? AND assistant_definition_id = ?") + .bind(DEFAULT_USER_ID) .bind(assistant_definition_id) .execute(&self.pool) .await?; @@ -611,8 +625,9 @@ impl IAssistantOverlayRepository for SqliteAssistantOverlayRepository { impl IAssistantPreferenceRepository for SqliteAssistantPreferenceRepository { async fn get(&self, 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(DEFAULT_USER_ID) .bind(assistant_definition_id) .fetch_optional(&self.pool) .await?; @@ -623,10 +638,10 @@ impl IAssistantPreferenceRepository for SqliteAssistantPreferenceRepository { 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 +650,7 @@ impl IAssistantPreferenceRepository for SqliteAssistantPreferenceRepository { last_mcp_ids = excluded.last_mcp_ids, updated_at = excluded.updated_at", ) + .bind(DEFAULT_USER_ID) .bind(params.assistant_definition_id) .bind(params.last_model_id) .bind(params.last_permission_value) @@ -656,7 +672,8 @@ impl IAssistantPreferenceRepository for SqliteAssistantPreferenceRepository { } async fn delete(&self, assistant_definition_id: &str) -> Result { - let result = sqlx::query("DELETE FROM assistant_preferences WHERE assistant_definition_id = ?") + let result = sqlx::query("DELETE FROM assistant_preferences WHERE user_id = ? AND assistant_definition_id = ?") + .bind(DEFAULT_USER_ID) .bind(assistant_definition_id) .execute(&self.pool) .await?; diff --git a/crates/aionui-db/src/repository/sqlite_skill.rs b/crates/aionui-db/src/repository/sqlite_skill.rs index 72281d8b3..ee8c6c889 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 { @@ -20,27 +22,41 @@ impl SqliteSkillRepository { impl ISkillRepository for SqliteSkillRepository { async fn list(&self) -> 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(DEFAULT_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?; + 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(DEFAULT_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?; + 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(DEFAULT_USER_ID) + .bind(name) + .fetch_optional(&self.pool) + .await?; Ok(row) } @@ -55,9 +71,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 +82,7 @@ impl ISkillRepository for SqliteSkillRepository { updated_at = excluded.updated_at", ) .bind(&id) + .bind(DEFAULT_USER_ID) .bind(params.name) .bind(params.description) .bind(params.path) @@ -84,10 +101,12 @@ impl ISkillRepository for SqliteSkillRepository { async fn delete_by_name(&self, 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(DEFAULT_USER_ID) .bind(name) .execute(&self.pool) .await?; diff --git a/crates/aionui-db/tests/skill_management_schema.rs b/crates/aionui-db/tests/skill_management_schema.rs index 9e6e74e9f..39c6e7dcd 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" { + 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()); +} From 91a2f6ceb697fcff2a6f4bf027cc46e1fe485690 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 14:52:41 +0800 Subject: [PATCH 017/145] db: add user scoped skill repository methods --- crates/aionui-db/src/repository/skill.rs | 50 ++++++ .../aionui-db/src/repository/sqlite_skill.rs | 166 ++++++++++++++++-- 2 files changed, 205 insertions(+), 11 deletions(-) diff --git a/crates/aionui-db/src/repository/skill.rs b/crates/aionui-db/src/repository/skill.rs index 1a2c80931..ad504af33 100644 --- a/crates/aionui-db/src/repository/skill.rs +++ b/crates/aionui-db/src/repository/skill.rs @@ -7,26 +7,76 @@ 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> { + let _ = user_id; + self.list().await + } + /// 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> { + let _ = user_id; + self.find_by_name(name).await + } + /// 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> { + let _ = user_id; + self.find_by_name_any(name).await + } + /// 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 { + let _ = user_id; + 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 { + let _ = user_id; + self.delete_by_name(name).await + } + /// 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 { + let _ = user_id; + self.create_import_record(params).await + } + /// 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> { + let _ = user_id; + self.list_import_records(limit).await + } } /// Parameters for creating or updating a skill row. diff --git a/crates/aionui-db/src/repository/sqlite_skill.rs b/crates/aionui-db/src/repository/sqlite_skill.rs index ee8c6c889..d74eb3a89 100644 --- a/crates/aionui-db/src/repository/sqlite_skill.rs +++ b/crates/aionui-db/src/repository/sqlite_skill.rs @@ -21,25 +21,33 @@ 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 (user_id IS NULL OR user_id = ?) AND deleted_at IS NULL AND enabled = 1 \ ORDER BY updated_at DESC, name ASC", ) - .bind(DEFAULT_USER_ID) + .bind(user_id) .fetch_all(&self.pool) .await?; Ok(rows) } async fn find_by_name(&self, name: &str) -> Result, DbError> { + 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(DEFAULT_USER_ID) + .bind(user_id) .bind(name) .fetch_optional(&self.pool) .await?; @@ -47,13 +55,17 @@ impl ISkillRepository for SqliteSkillRepository { } async fn find_by_name_any(&self, name: &str) -> Result, DbError> { + 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(DEFAULT_USER_ID) + .bind(user_id) .bind(name) .fetch_optional(&self.pool) .await?; @@ -61,8 +73,12 @@ impl ISkillRepository for SqliteSkillRepository { } 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 = self.find_by_name_any_for_user(user_id, params.name).await?; let id = existing .as_ref() .map(|row| row.id.clone()) @@ -82,7 +98,7 @@ impl ISkillRepository for SqliteSkillRepository { updated_at = excluded.updated_at", ) .bind(&id) - .bind(DEFAULT_USER_ID) + .bind(user_id) .bind(params.name) .bind(params.description) .bind(params.path) @@ -93,12 +109,16 @@ 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 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 = ? \ @@ -106,7 +126,7 @@ impl ISkillRepository for SqliteSkillRepository { ) .bind(now) .bind(now) - .bind(DEFAULT_USER_ID) + .bind(user_id) .bind(name) .execute(&self.pool) .await?; @@ -115,7 +135,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}'"))) } @@ -123,6 +143,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(); @@ -130,8 +158,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) @@ -148,6 +176,7 @@ impl ISkillRepository for SqliteSkillRepository { .bind(params.line) .bind(params.column) .bind(now) + .bind(user_id) .execute(&self.pool) .await?; @@ -159,9 +188,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?; @@ -174,12 +212,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; @@ -273,4 +327,94 @@ 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 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); + } } From 00c9dc6683285cd330ee175da7454c356aec36fd Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 15:00:13 +0800 Subject: [PATCH 018/145] extension: scope skill APIs by current user --- Cargo.lock | 1 + crates/aionui-db/src/repository/skill.rs | 5 + .../aionui-db/src/repository/sqlite_skill.rs | 102 ++++++++- crates/aionui-extension/Cargo.toml | 1 + crates/aionui-extension/src/skill_routes.rs | 31 ++- crates/aionui-extension/src/skill_service.rs | 204 +++++++++++++----- 6 files changed, 278 insertions(+), 66 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 697cbc5d8..f8cfa5444 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -634,6 +634,7 @@ name = "aionui-extension" version = "0.1.50" dependencies = [ "aionui-api-types", + "aionui-auth", "aionui-common", "aionui-db", "aionui-realtime", diff --git a/crates/aionui-db/src/repository/skill.rs b/crates/aionui-db/src/repository/skill.rs index ad504af33..bcde126f5 100644 --- a/crates/aionui-db/src/repository/skill.rs +++ b/crates/aionui-db/src/repository/skill.rs @@ -40,6 +40,11 @@ pub trait ISkillRepository: Send + Sync { self.upsert(params).await } + /// 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; diff --git a/crates/aionui-db/src/repository/sqlite_skill.rs b/crates/aionui-db/src/repository/sqlite_skill.rs index d74eb3a89..a40452e68 100644 --- a/crates/aionui-db/src/repository/sqlite_skill.rs +++ b/crates/aionui-db/src/repository/sqlite_skill.rs @@ -78,7 +78,11 @@ impl ISkillRepository for SqliteSkillRepository { 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_for_user(user_id, 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()) @@ -114,6 +118,48 @@ impl ISkillRepository for SqliteSkillRepository { .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 } @@ -383,6 +429,60 @@ mod tests { 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; 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/skill_routes.rs b/crates/aionui-extension/src/skill_routes.rs index 1d05d3d38..d762b02bc 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, ) diff --git a/crates/aionui-extension/src/skill_service.rs b/crates/aionui-extension/src/skill_service.rs index 8694a3fb6..b36a46118 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 @@ -512,17 +522,30 @@ pub async fn import_skill_with_repo( paths: &SkillPaths, repo: &dyn ISkillRepository, skill_path: &Path, +) -> Result { + 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(paths, skill_path).await?; - let overwritten = repo.find_by_name(&copied.name).await?.is_some(); + 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!( @@ -652,6 +675,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 +725,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 +749,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 +768,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 +862,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 +873,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); } @@ -1153,13 +1196,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 +1223,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(()) } @@ -1260,9 +1313,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 +1337,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, @@ -1401,9 +1465,10 @@ async fn resolve_skill_source_path(paths: &SkillPaths, name: &str) -> Result Result, ExtensionError> { let top = paths.builtin_skills_dir.join(name); @@ -1414,7 +1479,7 @@ async fn resolve_skill_source_path_with_repo( 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)); @@ -1554,9 +1619,10 @@ 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 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)); } @@ -1572,7 +1638,7 @@ 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_disk_user_skills_into_repo_for_user(paths, repo, DEFAULT_USER_ID).await?; sync_builtin_skills_into_repo(paths, repo).await?; Ok(()) } @@ -1583,14 +1649,14 @@ 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?; } } @@ -1627,9 +1693,26 @@ async fn sync_managed_skill_into_repo( Ok(()) } -async fn sync_disk_user_skills_into_repo( +async fn sync_global_skill_into_repo( + repo: &dyn ISkillRepository, + skill: &ScannedSkill, + source: &str, +) -> Result<(), ExtensionError> { + repo.upsert_global(UpsertSkillParams { + name: &skill.name, + description: Some(&skill.description), + path: &skill.path, + source, + enabled: true, + }) + .await?; + Ok(()) +} + +async fn sync_disk_user_skills_into_repo_for_user( paths: &SkillPaths, repo: &dyn ISkillRepository, + user_id: &str, ) -> Result<(), ExtensionError> { let scanned = scan_skill_dirs(&paths.user_skills_dir).await?; let mut backfilled = 0usize; @@ -1643,17 +1726,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; } From 8b836087deb8149c7424155e196a8c0c8f293deb Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 15:08:22 +0800 Subject: [PATCH 019/145] db: add user scoped assistant repository methods --- crates/aionui-db/src/repository/assistant.rs | 180 ++++ .../src/repository/sqlite_assistant.rs | 836 +++++++++++++++--- 2 files changed, 877 insertions(+), 139 deletions(-) diff --git a/crates/aionui-db/src/repository/assistant.rs b/crates/aionui-db/src/repository/assistant.rs index 5f6085d5e..8a12789c3 100644 --- a/crates/aionui-db/src/repository/assistant.rs +++ b/crates/aionui-db/src/repository/assistant.rs @@ -15,24 +15,67 @@ 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> { + let _ = user_id; + self.list().await + } + /// 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> { + let _ = user_id; + self.get(id).await + } + /// 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 { + let _ = user_id; + self.create(params).await + } + /// 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> { + let _ = user_id; + self.update(id, params).await + } + /// 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 { + let _ = user_id; + self.delete(id).await + } + /// 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 { + let _ = user_id; + self.upsert(params).await + } } /// Per-assistant user state (enabled flag, sort order, last-used timestamp). @@ -41,41 +84,107 @@ 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> { + let _ = user_id; + self.get(assistant_id).await + } + /// Fetch all override rows. async fn get_all(&self) -> Result, DbError>; + async fn get_all_for_user(&self, user_id: &str) -> Result, DbError> { + let _ = user_id; + self.get_all().await + } + /// 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 { + let _ = user_id; + self.upsert(params).await + } + /// 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 { + let _ = user_id; + self.delete(assistant_id).await + } + /// 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 { + let _ = user_id; + self.delete_orphans(valid_ids).await + } } /// 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> { + let _ = user_id; + self.list().await + } async fn list_including_deleted(&self) -> Result, DbError> { self.list().await } + async fn list_including_deleted_for_user(&self, user_id: &str) -> Result, DbError> { + let _ = user_id; + self.list_including_deleted().await + } 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> { + let _ = user_id; + self.get_by_assistant_id(assistant_id).await + } 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> { + let _ = user_id; + self.get_by_assistant_id_including_deleted(assistant_id).await + } 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> { + let _ = user_id; + self.get_by_id(id).await + } 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> { + let _ = user_id; + self.get_by_source_ref(source, source_ref).await + } async fn get_by_source_ref_including_deleted( &self, source: &str, @@ -83,7 +192,30 @@ 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> { + let _ = user_id; + self.get_by_source_ref_including_deleted(source, source_ref).await + } async fn upsert(&self, params: &UpsertAssistantDefinitionParams<'_>) -> Result; + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + let _ = user_id; + self.upsert(params).await + } + async fn upsert_global( + &self, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert(params).await + } async fn update_avatar_fields_preserving_deleted( &self, id: &str, @@ -96,21 +228,69 @@ 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 { + let _ = user_id; + self.soft_delete(id, deleted_at).await + } } /// 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> { + let _ = user_id; + self.get(assistant_definition_id).await + } async fn list(&self) -> Result, DbError>; + async fn list_for_user(&self, user_id: &str) -> Result, DbError> { + let _ = user_id; + self.list().await + } async fn upsert(&self, params: &UpsertAssistantOverlayParams<'_>) -> Result; + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAssistantOverlayParams<'_>, + ) -> Result { + let _ = user_id; + self.upsert(params).await + } async fn delete(&self, assistant_definition_id: &str) -> Result; + async fn delete_for_user(&self, user_id: &str, assistant_definition_id: &str) -> Result { + let _ = user_id; + self.delete(assistant_definition_id).await + } } /// 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> { + let _ = user_id; + self.get(assistant_definition_id).await + } async fn upsert(&self, params: &UpsertAssistantPreferenceParams<'_>) -> Result; + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAssistantPreferenceParams<'_>, + ) -> Result { + let _ = user_id; + self.upsert(params).await + } async fn delete(&self, assistant_definition_id: &str) -> Result; + async fn delete_for_user(&self, user_id: &str, assistant_definition_id: &str) -> Result { + let _ = user_id; + self.delete(assistant_definition_id).await + } } diff --git a/crates/aionui-db/src/repository/sqlite_assistant.rs b/crates/aionui-db/src/repository/sqlite_assistant.rs index 073beca9b..0d09b547f 100644 --- a/crates/aionui-db/src/repository/sqlite_assistant.rs +++ b/crates/aionui-db/src/repository/sqlite_assistant.rs @@ -35,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?; @@ -50,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) @@ -101,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); }; @@ -113,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) @@ -127,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?; @@ -135,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?; @@ -143,15 +178,24 @@ 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( "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 \ + user_id = excluded.user_id, \ name = excluded.name, \ description = excluded.description, \ avatar = excluded.avatar, \ @@ -166,6 +210,7 @@ impl IAssistantRepository for SqliteAssistantRepository { updated_at = excluded.updated_at", ) .bind(params.id) + .bind(user_id) .bind(params.name) .bind(params.description) .bind(params.avatar) @@ -183,7 +228,7 @@ impl IAssistantRepository for SqliteAssistantRepository { .await?; 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) @@ -236,6 +281,105 @@ 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(); + + 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 + user_id = excluded.user_id, + 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", + ) + .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?; + + 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`]. @@ -271,10 +415,14 @@ impl SqliteAssistantOverrideRepository { #[async_trait::async_trait] impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { async fn get(&self, assistant_id: &str) -> Result, DbError> { + 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(DEFAULT_USER_ID) + .bind(user_id) .bind(assistant_id) .fetch_optional(&self.pool) .await?; @@ -282,14 +430,26 @@ impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { } async fn get_all(&self) -> Result, DbError> { + 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(DEFAULT_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; @@ -303,7 +463,7 @@ impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { last_used_at = COALESCE(excluded.last_used_at, assistant_overrides.last_used_at), \ updated_at = excluded.updated_at", ) - .bind(DEFAULT_USER_ID) + .bind(user_id) .bind(params.assistant_id) .bind(params.enabled) .bind(params.sort_order) @@ -312,7 +472,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 @@ -322,8 +482,12 @@ impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { } async fn delete(&self, assistant_id: &str) -> Result { + 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(DEFAULT_USER_ID) + .bind(user_id) .bind(assistant_id) .execute(&self.pool) .await?; @@ -331,9 +495,13 @@ 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 WHERE user_id = ?") - .bind(DEFAULT_USER_ID) + .bind(user_id) .execute(&self.pool) .await?; return Ok(result.rows_affected()); @@ -341,7 +509,7 @@ impl IAssistantOverrideRepository for SqliteAssistantOverrideRepository { let placeholders = std::iter::repeat_n("?", valid_ids.len()).collect::>().join(","); let sql = format!("DELETE FROM assistant_overrides WHERE user_id = ? AND assistant_id NOT IN ({placeholders})"); - let mut q = sqlx::query(&sql).bind(DEFAULT_USER_ID); + let mut q = sqlx::query(&sql).bind(user_id); for id in valid_ids { q = q.bind(*id); } @@ -353,26 +521,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?; @@ -383,18 +619,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?; @@ -405,10 +661,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) @@ -420,10 +690,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) @@ -432,91 +716,26 @@ impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { } async fn upsert(&self, params: &UpsertAssistantDefinitionParams<'_>) -> Result { - let now = now_ms(); + if params.source == "builtin" && params.owner_type == "system" { + self.upsert_global(params).await + } else { + self.upsert_for_user(DEFAULT_USER_ID, params).await + } + } - 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", - ) - .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) - .await?; + async fn upsert_for_user( + &self, + user_id: &str, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert_with_user_id(Some(user_id), params).await + } - 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_global( + &self, + params: &UpsertAssistantDefinitionParams<'_>, + ) -> Result { + self.upsert_with_user_id(None, params).await } async fn update_avatar_fields_preserving_deleted( @@ -541,13 +760,18 @@ impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { } async fn soft_delete(&self, id: &str, deleted_at: i64) -> Result { + self.soft_delete_for_user(DEFAULT_USER_ID, id, deleted_at).await + } + + 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 id = ? AND deleted_at IS NULL", + 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?; @@ -558,10 +782,18 @@ impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { #[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 user_id = ? AND assistant_definition_id = ?", ) - .bind(DEFAULT_USER_ID) + .bind(user_id) .bind(assistant_definition_id) .fetch_optional(&self.pool) .await?; @@ -569,16 +801,28 @@ 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 WHERE user_id = ? ORDER BY sort_order, updated_at", ) - .bind(DEFAULT_USER_ID) + .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 ( @@ -592,7 +836,7 @@ impl IAssistantOverlayRepository for SqliteAssistantOverlayRepository { last_used_at = excluded.last_used_at, updated_at = excluded.updated_at", ) - .bind(DEFAULT_USER_ID) + .bind(user_id) .bind(params.assistant_definition_id) .bind(params.enabled) .bind(params.sort_order) @@ -603,17 +847,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 { + 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(DEFAULT_USER_ID) + .bind(user_id) .bind(assistant_definition_id) .execute(&self.pool) .await?; @@ -624,10 +874,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 user_id = ? AND assistant_definition_id = ?", ) - .bind(DEFAULT_USER_ID) + .bind(user_id) .bind(assistant_definition_id) .fetch_optional(&self.pool) .await?; @@ -635,6 +893,14 @@ 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 ( @@ -650,7 +916,7 @@ impl IAssistantPreferenceRepository for SqliteAssistantPreferenceRepository { last_mcp_ids = excluded.last_mcp_ids, updated_at = excluded.updated_at", ) - .bind(DEFAULT_USER_ID) + .bind(user_id) .bind(params.assistant_definition_id) .bind(params.last_model_id) .bind(params.last_permission_value) @@ -663,17 +929,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 { + 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(DEFAULT_USER_ID) + .bind(user_id) .bind(assistant_definition_id) .execute(&self.pool) .await?; @@ -690,6 +962,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, @@ -698,6 +973,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) } @@ -711,9 +987,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, @@ -732,12 +1021,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"), @@ -764,6 +1062,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; @@ -886,6 +1197,39 @@ 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 override_get_missing_returns_none() { let (_a, o, _db) = setup().await; @@ -1014,6 +1358,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; @@ -1030,6 +1410,104 @@ 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 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; @@ -1052,6 +1530,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; @@ -1075,4 +1592,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"]"#); + } } From 57151b64113efcb0b000e93b1e228cb23692f2dc Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 15:32:49 +0800 Subject: [PATCH 020/145] assistant: scope APIs by current user --- crates/aionui-assistant/src/routes.rs | 35 +- crates/aionui-assistant/src/service.rs | 510 +++++++++++++----- .../aionui-db/migrations/026_user_scope.sql | 88 ++- crates/aionui-db/src/repository/assistant.rs | 13 + .../src/repository/sqlite_assistant.rs | 44 +- 5 files changed, 529 insertions(+), 161 deletions(-) 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..cad6de2b1 100644 --- a/crates/aionui-assistant/src/service.rs +++ b/crates/aionui-assistant/src/service.rs @@ -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 @@ -228,7 +229,7 @@ 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 @@ -265,7 +266,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", @@ -457,12 +458,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 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 +481,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,6 +510,7 @@ impl AssistantService { for row in generated_rows { if let Err(error) = self .reconcile_generated_assistant( + user_id, row, &definitions, has_existing_generated, @@ -524,6 +533,7 @@ impl AssistantService { async fn reconcile_generated_assistant( &self, + user_id: &str, row: &AgentManagementRow, definitions: &[AssistantDefinitionRow], has_existing_generated: bool, @@ -540,7 +550,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 +622,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,7 +633,7 @@ 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() @@ -636,13 +646,16 @@ impl AssistantService { 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}")))?; } @@ -654,6 +667,16 @@ impl AssistantService { &self, row: &AssistantRow, requested_agent_id: Option<&str>, + ) -> Result<(), AssistantError> { + self.upsert_definition_from_legacy_user_row_for_user(DEFAULT_USER_ID, row, requested_agent_id) + .await + } + + async fn upsert_definition_from_legacy_user_row_for_user( + &self, + user_id: &str, + row: &AssistantRow, + requested_agent_id: Option<&str>, ) -> Result<(), AssistantError> { // User-defined assistants do not expose locale-aware editing in the // current product. Keep the unified definition canonical fields as the @@ -666,58 +689,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.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 +760,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 +768,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 +788,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 +814,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 @@ -811,7 +856,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,12 +864,16 @@ 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) .await?; @@ -835,13 +884,22 @@ 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 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(id, locale).await?; let projection = self .project_definition(&definition, state.as_ref(), &projections) @@ -879,9 +937,13 @@ 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}")))?; @@ -945,6 +1007,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,7 +1040,7 @@ 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?; let avatar = self.normalize_user_avatar_input(&id, req.avatar.as_deref())?; @@ -989,19 +1059,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(&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()) + self.apply_detail_overrides_for_user(user_id, &row.id, detail_overrides, false) + .await?; + 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 +1120,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) @@ -1057,27 +1141,37 @@ impl AssistantService { .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.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 +1194,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 +1222,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 +1241,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() { @@ -1177,22 +1276,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 +1308,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 +1432,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 +1447,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 +1466,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,23 +1480,27 @@ 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}"); } } @@ -1400,11 +1516,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 +1538,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 +1562,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 +1586,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 +1612,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 +1664,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 @@ -1580,9 +1716,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; } @@ -1746,10 +1882,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; } @@ -2024,8 +2164,9 @@ impl AssistantService { assistant_md_path(&self.user_skills_dir(), 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 +2174,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 +2183,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,6 +2193,34 @@ impl AssistantService { Ok((generate_prefixed_id("asstdef"), assistant_id.to_string())) } + 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(&self, id: &str) { remove_assistant_md_files(&self.user_rules_dir(), id); remove_assistant_md_files(&self.user_skills_dir(), id); @@ -3053,7 +3222,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, }; @@ -3302,8 +3471,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 +3497,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 +3879,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 @@ -6196,6 +6417,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-db/migrations/026_user_scope.sql b/crates/aionui-db/migrations/026_user_scope.sql index 80b5709b3..dee21e2b1 100644 --- a/crates/aionui-db/migrations/026_user_scope.sql +++ b/crates/aionui-db/migrations/026_user_scope.sql @@ -309,14 +309,86 @@ 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); -ALTER TABLE assistant_definitions - ADD COLUMN user_id TEXT REFERENCES users(id); -UPDATE assistant_definitions -SET user_id = 'system_default_user' -WHERE user_id IS NULL - AND NOT (source = 'builtin' AND owner_type = 'system'); 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; + +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; @@ -329,6 +401,10 @@ CREATE UNIQUE INDEX IF NOT EXISTS idx_assistant_definitions_global_assistant_id 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); diff --git a/crates/aionui-db/src/repository/assistant.rs b/crates/aionui-db/src/repository/assistant.rs index 8a12789c3..6b108e273 100644 --- a/crates/aionui-db/src/repository/assistant.rs +++ b/crates/aionui-db/src/repository/assistant.rs @@ -201,6 +201,19 @@ pub trait IAssistantDefinitionRepository: Send + Sync { let _ = user_id; self.get_by_source_ref_including_deleted(source, source_ref).await } + 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, diff --git a/crates/aionui-db/src/repository/sqlite_assistant.rs b/crates/aionui-db/src/repository/sqlite_assistant.rs index 0d09b547f..9c7d7379e 100644 --- a/crates/aionui-db/src/repository/sqlite_assistant.rs +++ b/crates/aionui-db/src/repository/sqlite_assistant.rs @@ -715,6 +715,38 @@ impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { Ok(row) } + 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) + } + + 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(assistant_id) + .fetch_optional(&self.pool) + .await?; + Ok(row) + } + async fn upsert(&self, params: &UpsertAssistantDefinitionParams<'_>) -> Result { if params.source == "builtin" && params.owner_type == "system" { self.upsert_global(params).await @@ -760,7 +792,17 @@ impl IAssistantDefinitionRepository for SqliteAssistantDefinitionRepository { } async fn soft_delete(&self, id: &str, deleted_at: i64) -> Result { - self.soft_delete_for_user(DEFAULT_USER_ID, id, deleted_at).await + let result = sqlx::query( + "UPDATE assistant_definitions + SET deleted_at = ?, updated_at = ? + WHERE id = ? AND deleted_at IS NULL", + ) + .bind(deleted_at) + .bind(now_ms()) + .bind(id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) } async fn soft_delete_for_user(&self, user_id: &str, id: &str, deleted_at: i64) -> Result { From fbec97ba1627bc33b66b55f2623bded998266562 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 15:48:06 +0800 Subject: [PATCH 021/145] db: add user scoped agent metadata methods --- .../aionui-db/migrations/026_user_scope.sql | 81 ++- .../src/repository/agent_metadata.rs | 83 +++ .../src/repository/sqlite_agent_metadata.rs | 607 +++++++++++++++--- 3 files changed, 670 insertions(+), 101 deletions(-) diff --git a/crates/aionui-db/migrations/026_user_scope.sql b/crates/aionui-db/migrations/026_user_scope.sql index dee21e2b1..f87607c10 100644 --- a/crates/aionui-db/migrations/026_user_scope.sql +++ b/crates/aionui-db/migrations/026_user_scope.sql @@ -274,8 +274,85 @@ FROM client_preferences; DROP TABLE client_preferences; ALTER TABLE client_preferences_new RENAME TO client_preferences; -ALTER TABLE agent_metadata - ADD COLUMN user_id TEXT REFERENCES users(id); +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; + +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); diff --git a/crates/aionui-db/src/repository/agent_metadata.rs b/crates/aionui-db/src/repository/agent_metadata.rs index 33df906ad..055aba11b 100644 --- a/crates/aionui-db/src/repository/agent_metadata.rs +++ b/crates/aionui-db/src/repository/agent_metadata.rs @@ -17,9 +17,19 @@ 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> { + let _ = user_id; + self.list_all().await + } + /// 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> { + let _ = user_id; + self.get(id).await + } + /// Look up by the unique `(agent_source, name)` pair. async fn find_by_source_and_name( &self, @@ -27,14 +37,46 @@ 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> { + let _ = user_id; + self.find_by_source_and_name(agent_source, name).await + } + /// 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> { + let _ = user_id; + self.find_builtin_by_backend(backend).await + } + /// 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 { + let _ = user_id; + self.upsert(params).await + } + + 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 +85,16 @@ 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> { + let _ = user_id; + self.apply_handshake(id, params).await + } + /// Persist the latest availability snapshot for an existing row. /// Returns `Ok(None)` if no row matches `id`. async fn update_availability_snapshot( @@ -51,6 +103,16 @@ 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> { + let _ = user_id; + self.update_availability_snapshot(id, params).await + } + /// 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 +124,30 @@ 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> { + let _ = user_id; + self.update_agent_overrides(id, command_override, env_override).await + } + /// 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 { + let _ = user_id; + self.set_enabled(id, enabled).await + } + /// 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 { + let _ = user_id; + self.delete(id).await + } } diff --git a/crates/aionui-db/src/repository/sqlite_agent_metadata.rs b/crates/aionui-db/src/repository/sqlite_agent_metadata.rs index e3eb7bf03..5a20814fa 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")?, @@ -236,15 +240,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 +265,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 +365,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 +393,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 +447,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 +628,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 +637,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 +650,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 +677,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 +689,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 +698,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 +707,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 +738,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 +758,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 +771,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 +1359,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()); + } } From 50b4ff6b55a33cc832d9f1ada69eb7386d2e0869 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 16:19:24 +0800 Subject: [PATCH 022/145] ai-agent: scope agent management APIs by current user --- crates/aionui-ai-agent/src/registry.rs | 227 +++++++++++------- crates/aionui-ai-agent/src/routes/agent.rs | 24 +- crates/aionui-ai-agent/src/services/agent.rs | 22 +- .../src/services/availability/mod.rs | 63 +++-- crates/aionui-ai-agent/src/services/custom.rs | 64 +++-- .../tests/agent_availability_integration.rs | 16 +- 6 files changed, 260 insertions(+), 156 deletions(-) diff --git a/crates/aionui-ai-agent/src/registry.rs b/crates/aionui-ai-agent/src/registry.rs index 0f133e5de..97a5aad52 100644 --- a/crates/aionui-ai-agent/src/registry.rs +++ b/crates/aionui-ai-agent/src/registry.rs @@ -355,100 +355,72 @@ 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| decode_row(row, AvailabilityProjection::Cached)) + .map(|(meta, reason)| { + let cached_meta = cached_rows.get(&meta.id); + let cached_reason = cached_reasons.get(&meta.id); + let (meta, reason) = overlay_hydrated_availability(meta, reason, cached_meta, cached_reason); + 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((meta, reason)) = row.and_then(|row| decode_row(row, AvailabilityProjection::Cached)) else { + return Ok(None); + }; + let cached_meta = self.by_id.read().await.get(&meta.id).cloned(); + let cached_reason = self.unavailable_reasons.read().await.get(&meta.id).cloned(); + let (meta, reason) = overlay_hydrated_availability(meta, reason, cached_meta.as_ref(), cached_reason.as_ref()); + Ok(Some(agent_management_row(meta, reason.as_ref()))) } /// Like [`Self::list_all_including_hidden`] but pairs every row @@ -496,6 +468,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() diff --git a/crates/aionui-ai-agent/src/routes/agent.rs b/crates/aionui-ai-agent/src/routes/agent.rs index f8c42fe45..f30a59ae1 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)?, ))) @@ -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)?, ))) diff --git a/crates/aionui-ai-agent/src/services/agent.rs b/crates/aionui-ai-agent/src/services/agent.rs index b07d956f3..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. @@ -122,9 +122,9 @@ impl AgentService { ) -> 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 @@ -161,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(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 4803169bb..bcab97712 100644 --- a/crates/aionui-ai-agent/src/services/availability/mod.rs +++ b/crates/aionui-ai-agent/src/services/availability/mod.rs @@ -18,6 +18,8 @@ use crate::error::AgentError; use crate::protocol::{cli_detect, custom_agent_probe}; use crate::registry::{AgentRegistry, guidance_for_snapshot_error_code}; +const SYSTEM_DEFAULT_USER_ID: &str = "system_default_user"; + #[async_trait::async_trait] pub trait AgentAvailabilityFeedbackPort: Send + Sync { async fn record_session_success(&self, agent_id: &str) -> Result<(), AgentError>; @@ -49,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, 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"))); } @@ -75,9 +77,9 @@ impl AgentAvailabilityService { 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"))) } @@ -107,17 +109,31 @@ impl AgentAvailabilityService { self.persist_snapshot(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> { + self.persist_snapshot_for_user(SYSTEM_DEFAULT_USER_ID, id, snapshot) + .await + } + + 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 { @@ -144,10 +160,12 @@ 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}")))?; + if user_id == SYSTEM_DEFAULT_USER_ID { + self.registry.reload_one(id).await?; + } Ok(()) } } @@ -381,7 +399,10 @@ impl AgentAvailabilityFeedbackPort for AgentAvailabilityService { mod tests { use std::sync::Arc; - use super::{AgentAvailabilityService, explicit_probe_args, probe_aionrs_provider_readiness, run_probe}; + use super::{ + AgentAvailabilityService, SYSTEM_DEFAULT_USER_ID, explicit_probe_args, probe_aionrs_provider_readiness, + run_probe, + }; use crate::registry::AgentRegistry; use aionui_api_types::{ AgentHandshake, AgentManagementStatus, AgentMetadata, AgentSnapshotCheckKind, AgentSnapshotCheckStatus, @@ -486,8 +507,9 @@ mod tests { .unwrap(); let row = service - .list_management_rows() + .list_management_rows(SYSTEM_DEFAULT_USER_ID) .await + .unwrap() .into_iter() .find(|item| item.id == "agent-session-failure") .unwrap(); @@ -560,8 +582,9 @@ mod tests { service.record_session_success("agent-session-success").await.unwrap(); let row = service - .list_management_rows() + .list_management_rows(SYSTEM_DEFAULT_USER_ID) .await + .unwrap() .into_iter() .find(|item| item.id == "agent-session-success") .unwrap(); diff --git a/crates/aionui-ai-agent/src/services/custom.rs b/crates/aionui-ai-agent/src/services/custom.rs index 5b3bc2a62..dbd76df48 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 @@ -45,16 +46,22 @@ impl AgentService { 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/tests/agent_availability_integration.rs b/crates/aionui-ai-agent/tests/agent_availability_integration.rs index 03b1a8483..639a5f217 100644 --- a/crates/aionui-ai-agent/tests/agent_availability_integration.rs +++ b/crates/aionui-ai-agent/tests/agent_availability_integration.rs @@ -8,7 +8,7 @@ use aionui_db::{ }; use aionui_realtime::EventBroadcaster; -const TEST_USER_ID: &str = "user-1"; +const SYSTEM_DEFAULT_USER_ID: &str = "system_default_user"; struct NoopBroadcaster; @@ -283,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); @@ -324,7 +324,7 @@ async fn manual_health_check_does_not_refresh_unrelated_agents() { let service = agent_service(registry.clone(), provider_repo, temp.path().to_path_buf()); service - .health_check_agent_by_id(TEST_USER_ID, "agent-target-missing") + .health_check_agent_by_id(SYSTEM_DEFAULT_USER_ID, "agent-target-missing") .await .unwrap(); @@ -371,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(); @@ -416,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(); From e90033dcbfc029b06837a16e2b49197585dfc9d5 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 16:20:04 +0800 Subject: [PATCH 023/145] app: pass user scope to remaining service reads --- Cargo.lock | 1 + crates/aionui-shell/Cargo.toml | 1 + crates/aionui-shell/src/routes.rs | 41 +++++++++++-------- crates/aionui-team/src/provisioning.rs | 6 +-- .../aionui-team/src/service/spawn_support.rs | 17 ++++---- crates/aionui-team/src/session.rs | 8 +++- 6 files changed, 46 insertions(+), 28 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f8cfa5444..77e68ef8a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -790,6 +790,7 @@ name = "aionui-shell" version = "0.1.50" dependencies = [ "aionui-api-types", + "aionui-auth", "aionui-common", "aionui-db", "aionui-runtime", 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-team/src/provisioning.rs b/crates/aionui-team/src/provisioning.rs index 4dad1196e..42751226e 100644 --- a/crates/aionui-team/src/provisioning.rs +++ b/crates/aionui-team/src/provisioning.rs @@ -558,7 +558,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 { @@ -703,8 +703,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; diff --git a/crates/aionui-team/src/service/spawn_support.rs b/crates/aionui-team/src/service/spawn_support.rs index 987d8fa3c..47ce0a7d0 100644 --- a/crates/aionui-team/src/service/spawn_support.rs +++ b/crates/aionui-team/src/service/spawn_support.rs @@ -74,6 +74,7 @@ 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, @@ -96,7 +97,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 +110,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()); @@ -138,8 +139,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,9 +150,9 @@ 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 json: serde_json::Value = serde_json::from_str(row.available_models.as_deref()?).ok()?; @@ -731,7 +732,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 +846,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 07ebba499..72e69430e 100644 --- a/crates/aionui-team/src/session.rs +++ b/crates/aionui-team/src/session.rs @@ -1369,7 +1369,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). From 50c27d883e1b34545669269922c6b433ca3bb23c Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 16:23:16 +0800 Subject: [PATCH 024/145] team: update provider test doubles for user scope --- crates/aionui-team/src/provisioning.rs | 13 +++++++++---- crates/aionui-team/src/service/spawn_support.rs | 14 ++++++++++---- crates/aionui-team/src/test_utils.rs | 7 ++++--- .../tests/session_service_integration.rs | 7 ++++--- 4 files changed, 27 insertions(+), 14 deletions(-) diff --git a/crates/aionui-team/src/provisioning.rs b/crates/aionui-team/src/provisioning.rs index 42751226e..e1da9f1c1 100644 --- a/crates/aionui-team/src/provisioning.rs +++ b/crates/aionui-team/src/provisioning.rs @@ -917,19 +917,24 @@ mod tests { #[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(()) } } diff --git a/crates/aionui-team/src/service/spawn_support.rs b/crates/aionui-team/src/service/spawn_support.rs index 47ce0a7d0..8a92eac73 100644 --- a/crates/aionui-team/src/service/spawn_support.rs +++ b/crates/aionui-team/src/service/spawn_support.rs @@ -404,11 +404,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) } @@ -416,11 +416,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())) } } @@ -428,6 +433,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(), diff --git a/crates/aionui-team/src/test_utils.rs b/crates/aionui-team/src/test_utils.rs index b8da93626..fdf44fafb 100644 --- a/crates/aionui-team/src/test_utils.rs +++ b/crates/aionui-team/src/test_utils.rs @@ -1017,11 +1017,11 @@ pub(crate) mod workspace_harness { #[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) } @@ -1034,13 +1034,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(()) } } diff --git a/crates/aionui-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index f8c381bbd..2f9bd772a 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -1350,10 +1350,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( @@ -1364,12 +1364,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())) } } From 333999c8b46cdd5cf062228c7fe9fe979533c0ae Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 17:15:34 +0800 Subject: [PATCH 025/145] channel: scope repository operations by owner user --- crates/aionui-app/src/router/routes.rs | 4 +- crates/aionui-app/src/router/state.rs | 25 +- crates/aionui-app/tests/channel_e2e.rs | 71 +- crates/aionui-channel/src/action.rs | 126 +++- crates/aionui-channel/src/manager.rs | 243 ++++--- crates/aionui-channel/src/orchestrator.rs | 7 +- crates/aionui-channel/src/pairing.rs | 232 ++++--- crates/aionui-channel/src/routes.rs | 47 +- crates/aionui-channel/src/session.rs | 220 ++++-- .../tests/manager_integration.rs | 98 +-- .../aionui-channel/tests/orchestrator_test.rs | 9 +- .../tests/pairing_integration.rs | 146 ++-- .../tests/session_action_integration.rs | 42 +- .../aionui-db/migrations/026_user_scope.sql | 25 +- crates/aionui-db/src/models/channel.rs | 3 + crates/aionui-db/src/repository/channel.rs | 77 ++- .../src/repository/sqlite_channel.rs | 635 ++++++++++++------ crates/aionui-db/tests/channel_repository.rs | 120 ++-- 18 files changed, 1417 insertions(+), 713 deletions(-) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index 3ddb0fa7e..6d95ef1a0 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -92,11 +92,13 @@ 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 { + 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" ); diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index 9ac3f5db1..cd3b7720d 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -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: String, } /// Build all default `ModuleStates` from application services. @@ -507,7 +508,17 @@ fn build_channel_settings_service( async fn build_channel_message_service( services: &AppServices, channel_settings: Arc, + owner_user_id: String, ) -> Arc { + Arc::new(aionui_channel::message_service::ChannelMessageService::new( + Arc::new(services.conversation_service.clone()), + services.worker_task_manager.clone(), + channel_settings, + owner_user_id, + )) +} + +async fn primary_channel_owner_user_id(services: &AppServices) -> String { let owner_user_id = services .user_repo .get_primary_webui_user() @@ -516,13 +527,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, - )) + owner_user_id } /// Build the default `ChannelRouterState` and orchestrator components. @@ -557,15 +562,18 @@ 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 runtime_owner_user_id = primary_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), + runtime_owner_user_id.clone(), )); - let message_service = build_channel_message_service(services, Arc::clone(&channel_settings)).await; + let message_service = + build_channel_message_service(services, Arc::clone(&channel_settings), runtime_owner_user_id.clone()).await; let orchestrator = aionui_channel::orchestrator::ChannelOrchestrator::new( action_executor, @@ -590,6 +598,7 @@ pub async fn build_channel_state( confirm_rx, manager, plugin_factory, + owner_user_id: runtime_owner_user_id, }; (state, components) 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-channel/src/action.rs b/crates/aionui-channel/src/action.rs index e953a9fd8..e82c186d2 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: String, } impl ActionExecutor { @@ -47,11 +49,13 @@ impl ActionExecutor { pairing: Arc, session_mgr: Arc, settings: Arc, + owner_user_id: String, ) -> Self { Self { pairing, session_mgr, settings, + owner_user_id, } } @@ -67,7 +71,10 @@ impl ActionExecutor { let chat_id = &msg.chat_id; // 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(&self.owner_user_id, user_id, &platform_type) + .await?; let internal_user_id = match internal_user_id { Some(id) => id, @@ -86,10 +93,19 @@ impl ActionExecutor { } // 3. Text message → session resolution → AI dispatch - let agent_config = self.settings.get_agent_config(&internal_user_id, msg.platform).await?; + let agent_config = self + .settings + .get_agent_config(&self.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( + &self.owner_user_id, + &internal_user_id, + chat_id, + &agent_config.agent_type, + None, + ) .await?; info!( @@ -101,6 +117,7 @@ impl ActionExecutor { ); Ok(MessageResult::Dispatched { + owner_user_id: self.owner_user_id.clone(), session_id: session.id, conversation_id: session.conversation_id, }) @@ -116,7 +133,7 @@ impl ActionExecutor { ) -> Result { let code = self .pairing - .request_pairing(platform_user_id, platform_type, Some(display_name)) + .request_pairing(&self.owner_user_id, platform_user_id, platform_type, Some(display_name)) .await?; debug!( @@ -148,14 +165,23 @@ impl ActionExecutor { "pairing.show" | "pairing.refresh" => { let code = self .pairing - .request_pairing(&action.context.user_id, &action.context.platform.to_string(), None) + .request_pairing( + &self.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( + &self.owner_user_id, + &action.context.user_id, + &action.context.platform.to_string(), + ) .await?; if authorized { Ok(ActionResponse { @@ -230,7 +256,7 @@ impl ActionExecutor { .await?; let session = self .session_mgr - .reset_session(user_id, chat_id, &agent_config.agent_type, None) + .reset_session(&self.owner_user_id, user_id, chat_id, &agent_config.agent_type, None) .await?; Ok(ActionResponse { @@ -256,11 +282,11 @@ impl ActionExecutor { let chat_id = &action.context.chat_id; let agent_config = self .settings - .get_agent_config(internal_user_id, action.context.platform) + .get_agent_config(&self.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(&self.owner_user_id, user_id, chat_id, &agent_config.agent_type, None) .await?; Ok(ActionResponse { @@ -519,6 +545,7 @@ mod tests { } // ── Mock IChannelRepository ──────────────────────────────────────── + const OWNER_ID: &str = "owner-test"; struct MockRepo { users: Mutex>, @@ -538,6 +565,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()), @@ -551,27 +579,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> { @@ -581,26 +615,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, @@ -616,10 +656,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()); @@ -628,7 +678,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(); @@ -637,29 +692,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(); @@ -668,7 +732,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) } } @@ -702,7 +766,7 @@ 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, OWNER_ID.to_owned()); (executor, repo) } diff --git a/crates/aionui-channel/src/manager.rs b/crates/aionui-channel/src/manager.rs index c673c76c1..53fdf311c 100644 --- a/crates/aionui-channel/src/manager.rs +++ b/crates/aionui-channel/src/manager.rs @@ -77,8 +77,8 @@ 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| { @@ -102,6 +102,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,7 +116,7 @@ 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 @@ -132,6 +133,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,7 +143,7 @@ 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) @@ -153,14 +155,14 @@ impl ChannelManager { }; 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 +172,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); - 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,6 +191,7 @@ 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, @@ -200,9 +205,10 @@ impl ChannelManager { .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 +218,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,7 +229,7 @@ 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; @@ -233,10 +239,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 +287,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 +296,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 +305,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; } } @@ -403,10 +412,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(|| { @@ -436,7 +445,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()))?; @@ -462,23 +476,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; + 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 +503,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" @@ -616,6 +632,7 @@ mod tests { } // ── Mock IChannelRepository ──────────────────────────────────────── + const OWNER_ID: &str = "owner-test"; struct MockRepo { plugins: Mutex>, @@ -635,16 +652,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 +671,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 +695,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 +707,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 +986,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 +996,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 +1007,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 +1021,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 +1033,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 +1046,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,7 +1067,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(); @@ -1027,7 +1080,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 +1096,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 +1114,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 +1126,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 +1139,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 +1150,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 +1168,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 +1185,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")); - 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 +1203,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 +1221,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 +1233,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 +1298,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 +1309,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 +1323,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,7 +1334,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(), 1); assert!(mgr.is_plugin_running("telegram")); } @@ -1288,6 +1351,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 +1363,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,7 +1376,7 @@ 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); @@ -1328,7 +1393,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 +1404,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 +1414,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); @@ -1364,7 +1431,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(); @@ -1390,7 +1457,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(); @@ -1418,12 +1485,12 @@ 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); } diff --git a/crates/aionui-channel/src/orchestrator.rs b/crates/aionui-channel/src/orchestrator.rs index db64e6148..9689adc91 100644 --- a/crates/aionui-channel/src/orchestrator.rs +++ b/crates/aionui-channel/src/orchestrator.rs @@ -79,6 +79,7 @@ impl ChannelOrchestrator { send_action_response(&sender, &plugin_id, &chat_id, &response).await; } Ok(MessageResult::Dispatched { + owner_user_id, session_id, conversation_id, }) => { @@ -86,6 +87,7 @@ impl ChannelOrchestrator { &msg_svc, &session_mgr, &sender, + &owner_user_id, &session_id, conversation_id.as_deref(), &text, @@ -145,6 +147,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 +155,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"); @@ -189,7 +192,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"); diff --git a/crates/aionui-channel/src/pairing.rs b/crates/aionui-channel/src/pairing.rs index a82bc3b24..beb2e7d59 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" @@ -101,14 +104,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 +120,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" @@ -148,20 +152,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 +173,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 +191,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 +225,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 +241,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 +253,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 +335,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 +375,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 +387,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 +402,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 +415,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 +477,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 +501,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 +515,7 @@ mod tests { } // ── Helpers ──────────────────────────────────────────────────────── + const OWNER_ID: &str = "owner-test"; fn make_service() -> (PairingService, Arc, Arc) { let repo = Arc::new(MockRepo::new()); @@ -503,7 +561,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 +579,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 +595,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 +608,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 +629,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 +640,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 +663,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 +683,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 +703,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 +713,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 +722,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 +734,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 +755,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 +770,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 +778,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 +787,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 +810,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 +821,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/routes.rs b/crates/aionui-channel/src/routes.rs index 93127dbaf..bdb6c6c72 100644 --- a/crates/aionui-channel/src/routes.rs +++ b/crates/aionui-channel/src/routes.rs @@ -133,8 +133,9 @@ 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 statuses = state.manager.get_plugin_status(&user.id).await?; let extension_plugins = state.extension_registry.get_channel_plugins().await; let extension_map: HashMap = extension_plugins @@ -308,6 +309,7 @@ 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)?; @@ -316,7 +318,7 @@ async fn enable_plugin( 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(()) => { @@ -339,7 +341,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 { @@ -361,6 +363,7 @@ 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)?; @@ -368,7 +371,7 @@ async fn disable_plugin( if resolve_extension_channel_plugin(&state, &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() @@ -380,7 +383,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()), @@ -440,8 +443,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 { @@ -459,11 +463,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, @@ -475,11 +480,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, @@ -495,8 +501,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 { @@ -516,17 +527,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)?; @@ -544,8 +559,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 { @@ -594,7 +610,7 @@ async fn set_channel_assistant_setting( .settings_service .set_assistant_setting(&user.id, platform, &req) .await?; - state.session_manager.clear_all_sessions().await?; + state.session_manager.clear_all_sessions(&user.id).await?; Ok(Json(ApiResponse::ok(BridgeResponse { success: true, @@ -618,7 +634,7 @@ async fn set_channel_default_model_setting( .settings_service .set_model_setting(&user.id, platform, &req) .await?; - state.session_manager.clear_all_sessions().await?; + state.session_manager.clear_all_sessions(&user.id).await?; Ok(Json(ApiResponse::ok(BridgeResponse { success: true, @@ -638,6 +654,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)?; @@ -645,7 +662,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, 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/tests/manager_integration.rs b/crates/aionui-channel/tests/manager_integration.rs index 5504949ae..a7880dec2 100644 --- a/crates/aionui-channel/tests/manager_integration.rs +++ b/crates/aionui-channel/tests/manager_integration.rs @@ -21,6 +21,7 @@ use aionui_realtime::EventBroadcaster; use tokio::sync::mpsc; // ── Test infrastructure ───────────────────────────────────────────── +const OWNER_ID: &str = "system_default_user"; struct MockBroadcaster { events: Mutex>>, @@ -234,7 +235,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 +246,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,12 +266,12 @@ 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"); @@ -288,7 +289,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 +298,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 +320,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")); - 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 +348,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 +363,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 +377,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")); - 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 +397,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 +475,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 +487,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 +513,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 +536,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 +554,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 +577,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,7 +586,7 @@ 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")); } @@ -592,12 +598,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,16 +614,18 @@ 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")); - let statuses = mgr.get_plugin_status().await.unwrap(); + let statuses = mgr.get_plugin_status(OWNER_ID).await.unwrap(); assert_eq!(statuses.len(), 2); } @@ -628,10 +636,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,7 +654,7 @@ 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(); @@ -671,7 +681,7 @@ 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(); @@ -687,11 +697,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 +716,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/orchestrator_test.rs b/crates/aionui-channel/tests/orchestrator_test.rs index 3b7c710aa..2e2bd2897 100644 --- a/crates/aionui-channel/tests/orchestrator_test.rs +++ b/crates/aionui-channel/tests/orchestrator_test.rs @@ -8,6 +8,8 @@ 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 { id: "msg-1".into(), @@ -45,7 +47,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, + 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..4ed5ef350 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, 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,7 +78,7 @@ 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 } @@ -145,10 +147,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 +158,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 +173,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 +203,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 +228,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 +248,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 +269,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 +411,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-db/migrations/026_user_scope.sql b/crates/aionui-db/migrations/026_user_scope.sql index f87607c10..74fb17229 100644 --- a/crates/aionui-db/migrations/026_user_scope.sql +++ b/crates/aionui-db/migrations/026_user_scope.sql @@ -591,29 +591,8 @@ ALTER TABLE 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; - -DROP TABLE assistant_users; -ALTER TABLE assistant_users_new RENAME TO assistant_users; +ALTER TABLE assistant_users + ADD COLUMN owner_user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id); CREATE INDEX IF NOT EXISTS idx_assistant_users_owner_authorized_at ON assistant_users(owner_user_id, authorized_at DESC); diff --git a/crates/aionui-db/src/models/channel.rs b/crates/aionui-db/src/models/channel.rs index 51ee8a320..086f94bc9 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, @@ -28,6 +29,7 @@ pub struct ChannelPluginRow { #[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/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/sqlite_channel.rs b/crates/aionui-db/src/repository/sqlite_channel.rs index 1052df95e..4a3f62736 100644 --- a/crates/aionui-db/src/repository/sqlite_channel.rs +++ b/crates/aionui-db/src/repository/sqlite_channel.rs @@ -20,27 +20,33 @@ 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 (?, ?, ?, ?, ?, ?, ?, ?, ?) \ + (id, owner_user_id, type, name, enabled, config, status, last_connected, created_at, updated_at) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \ ON CONFLICT(id) DO UPDATE SET \ + owner_user_id = excluded.owner_user_id, \ type = excluded.type, \ name = excluded.name, \ enabled = excluded.enabled, \ @@ -50,6 +56,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 +70,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 +92,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 +110,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 +120,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 +134,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 +162,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 +192,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 +210,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 +224,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,48 +289,80 @@ 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 { @@ -296,16 +371,20 @@ impl IChannelRepository for SqliteChannelRepository { 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 +393,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 +431,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 +458,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 +494,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 +525,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 +570,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 +598,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 +613,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 +621,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 +634,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 +642,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 +652,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 +667,31 @@ 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 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 +702,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 +713,7 @@ mod tests { let (repo, _db) = setup().await; let err = repo .update_plugin_status( + OWNER_A, "nope", &UpdatePluginStatusParams { status: Some("error".into()), @@ -586,9 +728,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 +738,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 +755,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 +763,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 +777,72 @@ 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 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 +852,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 +872,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 +900,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 +921,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 +946,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 +1058,116 @@ 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::NotFound(_))); + 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 +1176,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 +1187,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 +1197,36 @@ 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 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 +1234,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 +1269,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 +1277,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 +1301,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/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"); From a12a1c2b1f9ccbc14141b03597bd754104c557bc Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 17:31:27 +0800 Subject: [PATCH 026/145] realtime: bind websocket events to user scope --- crates/aionui-app/src/router/routes.rs | 11 ++- crates/aionui-app/src/router/state.rs | 5 ++ crates/aionui-app/tests/websocket_e2e.rs | 29 ++++++- crates/aionui-realtime/src/handler.rs | 17 +++- crates/aionui-realtime/src/manager.rs | 82 +++++++++++++++++++ crates/aionui-realtime/src/types.rs | 2 + .../tests/handler_integration.rs | 2 + 7 files changed, 143 insertions(+), 5 deletions(-) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index 6d95ef1a0..9f243d725 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -70,7 +70,16 @@ 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 { + ws_manager.broadcast_all(event); + } } }); diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index cd3b7720d..10f7fb955 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -859,12 +859,16 @@ 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(|_| 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 token_user_resolver = + Arc::new(move |token: &str| jwt_service.verify(token).ok().map(|payload| payload.user_id)); let token_extractor = Arc::new(|headers: &axum::http::HeaderMap| extract_token_from_ws_headers(headers)); @@ -872,6 +876,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, } } diff --git a/crates/aionui-app/tests/websocket_e2e.rs b/crates/aionui-app/tests/websocket_e2e.rs index 481190926..f9bdc788c 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; @@ -349,6 +349,33 @@ 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"); + let token_b = sign_token(&app, "user-b"); + + 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; + + assert_eq!(ws_manager(&app).client_count(), 2); + 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_2_unicast_reaches_only_target() { use aionui_realtime::ConnectionId; diff --git a/crates/aionui-realtime/src/handler.rs b/crates/aionui-realtime/src/handler.rs index 16ccc5e38..3f1091949 100644 --- a/crates/aionui-realtime/src/handler.rs +++ b/crates/aionui-realtime/src/handler.rs @@ -20,12 +20,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 internal user ID carried by it. +pub type TokenUserResolver = Arc Option + 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 +79,15 @@ async fn handle_socket(socket: WebSocket, token: Option, state: WsHandle return; } + let Some(user_id) = (state.token_user_resolver)(&token) 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 +97,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 +269,7 @@ mod tests { manager, router: Arc::new(crate::router::NoopMessageRouter), token_validator: Arc::new(|_| true), + token_user_resolver: Arc::new(|_| Some("system_default_user".into())), token_extractor: Arc::new(|_| None), } } @@ -505,6 +515,7 @@ mod tests { manager, router: router.clone(), token_validator: Arc::new(|_| true), + token_user_resolver: Arc::new(|_| Some("system_default_user".into())), token_extractor: Arc::new(|_| None), }; diff --git a/crates/aionui-realtime/src/manager.rs b/crates/aionui-realtime/src/manager.rs index 4d5153ba6..a70c7c0e7 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,14 @@ 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() + } + /// Send a message to all connected clients. /// /// Uses `try_send` for backpressure. A saturated channel cannot reliably @@ -103,6 +117,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) { @@ -366,6 +419,28 @@ 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 broadcast_all_removes_closed_channels() { let mgr = WebSocketManager::new(); @@ -442,6 +517,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "valid".into(), last_ping: Instant::now(), tx, @@ -476,6 +552,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "valid".into(), last_ping: old_ping, tx, @@ -507,6 +584,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "expired-token".into(), last_ping: Instant::now(), tx, @@ -540,6 +618,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "expired-token".into(), last_ping: Instant::now(), tx, @@ -563,6 +642,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "expired".into(), last_ping: old_ping, tx, @@ -596,6 +676,7 @@ mod tests { connections.insert( ConnectionId(1), ClientInfo { + user_id: "user-a".into(), token: "good".into(), last_ping: Instant::now(), tx: tx1, @@ -607,6 +688,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, 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..53b0eae5f 100644 --- a/crates/aionui-realtime/tests/handler_integration.rs +++ b/crates/aionui-realtime/tests/handler_integration.rs @@ -37,6 +37,7 @@ 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| (t == "valid-token").then(|| "test-user".to_owned())), token_extractor: Arc::new(|headers| { headers .get("authorization") @@ -396,6 +397,7 @@ 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| (t == "valid-token").then(|| "test-user".to_owned())), token_extractor: Arc::new(|headers| { headers .get("authorization") From 718eb5b135587a97d129942adcfac266c6cd99a3 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 17:57:49 +0800 Subject: [PATCH 027/145] channel: isolate plugin runtime by owner --- crates/aionui-api-types/src/channel.rs | 2 + crates/aionui-channel/src/action.rs | 6 + crates/aionui-channel/src/manager.rs | 115 +++++++++++------- crates/aionui-channel/src/orchestrator.rs | 14 ++- .../src/plugins/dingtalk/plugin.rs | 2 + .../aionui-channel/src/plugins/lark/plugin.rs | 3 + .../src/plugins/telegram/plugin.rs | 2 + .../src/plugins/weixin/plugin.rs | 1 + crates/aionui-channel/src/stream_relay.rs | 65 ++++++++-- crates/aionui-channel/src/types.rs | 4 + .../tests/dingtalk_integration.rs | 2 +- .../aionui-channel/tests/lark_integration.rs | 2 +- .../tests/manager_integration.rs | 82 +++++++++++-- .../aionui-channel/tests/orchestrator_test.rs | 1 + .../tests/session_action_integration.rs | 2 + .../aionui-channel/tests/stream_relay_test.rs | 7 ++ .../tests/telegram_integration.rs | 2 +- .../tests/weixin_integration.rs | 2 +- .../src/repository/sqlite_channel.rs | 94 ++++++++++++-- 19 files changed, 329 insertions(+), 79 deletions(-) diff --git a/crates/aionui-api-types/src/channel.rs b/crates/aionui-api-types/src/channel.rs index df84a53fd..666be2cf4 100644 --- a/crates/aionui-api-types/src/channel.rs +++ b/crates/aionui-api-types/src/channel.rs @@ -267,6 +267,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, } @@ -714,6 +715,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(), diff --git a/crates/aionui-channel/src/action.rs b/crates/aionui-channel/src/action.rs index e82c186d2..cb37363d8 100644 --- a/crates/aionui-channel/src/action.rs +++ b/crates/aionui-channel/src/action.rs @@ -59,6 +59,10 @@ impl ActionExecutor { } } + pub fn owner_user_id(&self) -> &str { + &self.owner_user_id + } + /// Main entry point: handle an incoming message from any platform. /// /// Flow: @@ -772,6 +776,7 @@ mod tests { 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(), @@ -802,6 +807,7 @@ mod tests { params: Option>, ) -> UnifiedIncomingMessage { UnifiedIncomingMessage { + owner_user_id: None, id: "msg_1".into(), platform, chat_id: chat_id.into(), diff --git a/crates/aionui-channel/src/manager.rs b/crates/aionui-channel/src/manager.rs index 53fdf311c..8608c98d9 100644 --- a/crates/aionui-channel/src/manager.rs +++ b/crates/aionui-channel/src/manager.rs @@ -28,7 +28,7 @@ pub struct ChannelManager { repo: Arc, broadcaster: Arc, encryption_key: [u8; 32], - /// Active plugin instances keyed by plugin ID. + /// 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. @@ -82,7 +82,10 @@ impl ChannelManager { 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(); @@ -119,10 +122,7 @@ impl ChannelManager { 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)?; @@ -149,10 +149,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); if let Err(e) = plugin.initialize(config, callbacks).await { self.update_plugin_error(owner_user_id, plugin_id, &e.to_string()).await; @@ -177,7 +174,7 @@ impl ChannelManager { .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!(owner_user_id = %owner_user_id, plugin_id = %plugin_id, "plugin enabled and started"); self.broadcast_status_change(owner_user_id, plugin_id).await; @@ -196,9 +193,7 @@ impl ChannelManager { 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) @@ -231,7 +226,7 @@ impl ChannelManager { /// Idempotent — disabling an already-disabled plugin is a no-op. 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 { @@ -330,7 +325,7 @@ impl ChannelManager { 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"); } @@ -341,9 +336,9 @@ impl ChannelManager { } /// 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) } @@ -354,13 +349,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 } @@ -368,6 +364,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, @@ -375,7 +372,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 } @@ -431,8 +428,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) -> String { + format!("{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: &str) { + 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, @@ -462,10 +489,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?; @@ -478,7 +502,7 @@ impl ChannelManager { }; self.repo.update_plugin_status(owner_user_id, &row.id, ¶ms).await?; - self.plugins.insert(row.id.clone(), plugin); + 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(()) @@ -521,10 +545,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, }; @@ -541,7 +569,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(), @@ -576,21 +606,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 } } @@ -1072,7 +1105,7 @@ mod tests { .unwrap(); assert_eq!(mgr.active_plugin_count(), 1); - assert!(mgr.is_plugin_running("telegram")); + assert!(mgr.is_plugin_running(OWNER_ID, "telegram")); } #[tokio::test] @@ -1188,7 +1221,7 @@ mod tests { 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(OWNER_ID, "telegram").await.unwrap(); @@ -1336,7 +1369,7 @@ mod tests { 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] @@ -1380,8 +1413,8 @@ mod tests { // 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(); @@ -1436,7 +1469,7 @@ mod tests { .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"); @@ -1446,7 +1479,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(_))); @@ -1461,7 +1494,7 @@ mod tests { .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(); } @@ -1470,7 +1503,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(_))); @@ -1497,7 +1530,7 @@ mod tests { #[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/orchestrator.rs b/crates/aionui-channel/src/orchestrator.rs index 9689adc91..4a8da3d62 100644 --- a/crates/aionui-channel/src/orchestrator.rs +++ b/crates/aionui-channel/src/orchestrator.rs @@ -67,6 +67,8 @@ impl ChannelOrchestrator { let chat_id = msg.chat_id.clone(); let plugin_id = platform.to_string(); let text = msg.content.text.clone(); + let fallback_owner_user_id = self.action_executor.owner_user_id().to_owned(); + let message_owner_user_id = msg.owner_user_id.clone().unwrap_or(fallback_owner_user_id); let executor = Arc::clone(&self.action_executor); let msg_svc = Arc::clone(&self.message_service); @@ -76,7 +78,7 @@ 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; + send_action_response(&sender, &message_owner_user_id, &plugin_id, &chat_id, &response).await; } Ok(MessageResult::Dispatched { owner_user_id, @@ -110,6 +112,7 @@ impl ChannelOrchestrator { async fn send_action_response( sender: &Arc, + owner_user_id: &str, plugin_id: &str, chat_id: &str, response: &crate::types::ActionResponse, @@ -132,11 +135,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; } } } @@ -184,7 +189,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; } }; @@ -201,6 +206,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/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 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..bf26d96fd 100644 --- a/crates/aionui-channel/tests/dingtalk_integration.rs +++ b/crates/aionui-channel/tests/dingtalk_integration.rs @@ -273,6 +273,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("system_default_user", "dingtalk")); } } diff --git a/crates/aionui-channel/tests/lark_integration.rs b/crates/aionui-channel/tests/lark_integration.rs index 1656ce0c0..75f861565 100644 --- a/crates/aionui-channel/tests/lark_integration.rs +++ b/crates/aionui-channel/tests/lark_integration.rs @@ -272,6 +272,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("system_default_user", "lark")); } } diff --git a/crates/aionui-channel/tests/manager_integration.rs b/crates/aionui-channel/tests/manager_integration.rs index a7880dec2..ab207c62e 100644 --- a/crates/aionui-channel/tests/manager_integration.rs +++ b/crates/aionui-channel/tests/manager_integration.rs @@ -16,7 +16,9 @@ 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; @@ -146,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); @@ -154,7 +184,7 @@ async fn setup() -> (ChannelManager, Arc, Arc PluginFactory { @@ -278,7 +308,7 @@ async fn ep1_enable_telegram_plugin() { 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); } @@ -331,7 +361,7 @@ async fn ep6_re_enable_empty_config_reuses_stored_credentials() { .await .unwrap(); - assert!(mgr.is_plugin_running("telegram")); + assert!(mgr.is_plugin_running(OWNER_ID, "telegram")); let row = repo.get_plugin(OWNER_ID, "telegram").await.unwrap().unwrap(); assert!(row.enabled); @@ -383,7 +413,7 @@ async fn dp1_disable_enabled_plugin() { 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(OWNER_ID, "telegram").await.unwrap().unwrap(); assert!(!row.enabled); @@ -588,7 +618,7 @@ async fn restore_starts_enabled_plugins() { // Restore should bring it back 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 ───────────────────────── @@ -622,13 +652,43 @@ async fn enable_multiple_plugins() { .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(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] @@ -659,7 +719,7 @@ async fn send_message_routes_to_plugin() { .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"); @@ -670,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(_))); @@ -685,7 +745,7 @@ async fn edit_message_routes_to_plugin() { .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(); } diff --git a/crates/aionui-channel/tests/orchestrator_test.rs b/crates/aionui-channel/tests/orchestrator_test.rs index 2e2bd2897..6e2a1b30a 100644 --- a/crates/aionui-channel/tests/orchestrator_test.rs +++ b/crates/aionui-channel/tests/orchestrator_test.rs @@ -12,6 +12,7 @@ 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(), diff --git a/crates/aionui-channel/tests/session_action_integration.rs b/crates/aionui-channel/tests/session_action_integration.rs index 4ed5ef350..a4c748aba 100644 --- a/crates/aionui-channel/tests/session_action_integration.rs +++ b/crates/aionui-channel/tests/session_action_integration.rs @@ -84,6 +84,7 @@ async fn create_user(repo: &Arc, platform_user_id: &str, 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(), @@ -112,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(), 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..34fef1232 100644 --- a/crates/aionui-channel/tests/telegram_integration.rs +++ b/crates/aionui-channel/tests/telegram_integration.rs @@ -254,6 +254,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("system_default_user", "telegram")); } } diff --git a/crates/aionui-channel/tests/weixin_integration.rs b/crates/aionui-channel/tests/weixin_integration.rs index c97f82f35..955b6dcef 100644 --- a/crates/aionui-channel/tests/weixin_integration.rs +++ b/crates/aionui-channel/tests/weixin_integration.rs @@ -256,7 +256,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("system_default_user", "weixin")); } // -- Login event serialization ------------------------------------------ diff --git a/crates/aionui-db/src/repository/sqlite_channel.rs b/crates/aionui-db/src/repository/sqlite_channel.rs index 4a3f62736..0c6c093e5 100644 --- a/crates/aionui-db/src/repository/sqlite_channel.rs +++ b/crates/aionui-db/src/repository/sqlite_channel.rs @@ -16,6 +16,27 @@ impl SqliteChannelRepository { } } +fn plugin_storage_id(owner_user_id: &str, plugin_id: &str) -> String { + if owner_user_id == "system_default_user" { + plugin_id.to_owned() + } else { + format!("{owner_user_id}:{plugin_id}") + } +} + +fn plugin_logical_id(owner_user_id: &str, stored_id: &str) -> String { + let prefix = format!("{owner_user_id}:"); + stored_id + .strip_prefix(&prefix) + .map(str::to_owned) + .unwrap_or_else(|| stored_id.to_owned()) +} + +fn plugin_row_with_logical_id(owner_user_id: &str, mut row: ChannelPluginRow) -> ChannelPluginRow { + row.id = plugin_logical_id(owner_user_id, &row.id); + row +} + #[async_trait::async_trait] impl IChannelRepository for SqliteChannelRepository { // ── Plugin CRUD ────────────────────────────────────────────────── @@ -27,20 +48,28 @@ impl IChannelRepository for SqliteChannelRepository { .bind(owner_user_id) .fetch_all(&self.pool) .await?; - Ok(rows) + Ok(rows + .into_iter() + .map(|row| plugin_row_with_logical_id(owner_user_id, row)) + .collect()) } 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) + let storage_id = plugin_storage_id(owner_user_id, id); + let row = sqlx::query_as::<_, ChannelPluginRow>( + "SELECT * FROM assistant_plugins WHERE owner_user_id = ? AND id IN (?, ?) ORDER BY id = ? DESC LIMIT 1", + ) + .bind(owner_user_id) + .bind(&storage_id) + .bind(id) + .bind(&storage_id) + .fetch_optional(&self.pool) + .await?; + Ok(row.map(|row| plugin_row_with_logical_id(owner_user_id, row))) } async fn upsert_plugin(&self, owner_user_id: &str, row: &ChannelPluginRow) -> Result<(), DbError> { + let storage_id = plugin_storage_id(owner_user_id, &row.id); sqlx::query( "INSERT INTO assistant_plugins \ (id, owner_user_id, type, name, enabled, config, status, last_connected, created_at, updated_at) \ @@ -55,7 +84,7 @@ impl IChannelRepository for SqliteChannelRepository { last_connected = excluded.last_connected, \ updated_at = excluded.updated_at", ) - .bind(&row.id) + .bind(&storage_id) .bind(owner_user_id) .bind(&row.r#type) .bind(&row.name) @@ -76,6 +105,7 @@ impl IChannelRepository for SqliteChannelRepository { id: &str, params: &UpdatePluginStatusParams, ) -> Result<(), DbError> { + let storage_id = plugin_storage_id(owner_user_id, id); let mut set_clauses = Vec::new(); if params.status.is_some() { set_clauses.push("status = ?"); @@ -111,7 +141,7 @@ impl IChannelRepository for SqliteChannelRepository { } query = query.bind(now); query = query.bind(owner_user_id); - query = query.bind(id); + query = query.bind(&storage_id); let result = query.execute(&self.pool).await?; if result.rows_affected() == 0 { @@ -121,9 +151,10 @@ impl IChannelRepository for SqliteChannelRepository { } async fn delete_plugin(&self, owner_user_id: &str, id: &str) -> Result<(), DbError> { + let storage_id = plugin_storage_id(owner_user_id, id); let result = sqlx::query("DELETE FROM assistant_plugins WHERE owner_user_id = ? AND id = ?") .bind(owner_user_id) - .bind(id) + .bind(&storage_id) .execute(&self.pool) .await?; if result.rows_affected() == 0 { @@ -684,6 +715,47 @@ mod tests { 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; From 43c79b7a6a21465d43711259c7b872eacca55288 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 18:03:20 +0800 Subject: [PATCH 028/145] channel: scope pairing websocket events --- crates/aionui-api-types/src/channel.rs | 9 +++++++++ crates/aionui-channel/src/pairing.rs | 2 ++ 2 files changed, 11 insertions(+) diff --git a/crates/aionui-api-types/src/channel.rs b/crates/aionui-api-types/src/channel.rs index 666be2cf4..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, @@ -277,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, @@ -685,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(), @@ -692,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"); @@ -702,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(), @@ -733,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"); @@ -742,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"); @@ -757,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-channel/src/pairing.rs b/crates/aionui-channel/src/pairing.rs index beb2e7d59..545f5ac59 100644 --- a/crates/aionui-channel/src/pairing.rs +++ b/crates/aionui-channel/src/pairing.rs @@ -85,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(), @@ -136,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, From 8e588a3a010f77ec002855c4a264873f50bd5ee5 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 18:10:31 +0800 Subject: [PATCH 029/145] app: drop unscoped websocket business events --- crates/aionui-app/src/router/routes.rs | 11 ++++++- crates/aionui-app/tests/websocket_e2e.rs | 38 ++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index 9f243d725..8526924be 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -77,8 +77,13 @@ pub async fn create_router_with_runtime(services: &AppServices) -> Result<(Route .map(str::to_owned) { ws_manager.broadcast_to_user(&user_id, event); - } else { + } 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" + ); } } }); @@ -324,6 +329,10 @@ fn auth_identity_mode(identity_mode: crate::config::IdentityMode) -> AuthIdentit } } +fn is_global_websocket_event(event_name: &str) -> bool { + matches!(event_name, "runtime.statusChanged") +} + 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) { diff --git a/crates/aionui-app/tests/websocket_e2e.rs b/crates/aionui-app/tests/websocket_e2e.rs index f9bdc788c..acdc0af88 100644 --- a/crates/aionui-app/tests/websocket_e2e.rs +++ b/crates/aionui-app/tests/websocket_e2e.rs @@ -376,6 +376,44 @@ async fn t4_1_scoped_event_reaches_only_matching_user() { 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"); + + 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"); + let token_b = sign_token(&app, "user-b"); + + 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"); +} + #[tokio::test] async fn t4_2_unicast_reaches_only_target() { use aionui_realtime::ConnectionId; From 8d45033f73bb4231d81a8a1b7d099e8e3c599999 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 18:16:04 +0800 Subject: [PATCH 030/145] conversation: include user scope in websocket events --- .../src/acp_error_recovery.rs | 2 +- .../src/runtime_completion.rs | 1 + crates/aionui-conversation/src/service.rs | 24 ++++++++++++------- .../src/stream_persistence.rs | 1 + .../aionui-conversation/src/stream_relay.rs | 2 ++ 5 files changed, 21 insertions(+), 9 deletions(-) diff --git a/crates/aionui-conversation/src/acp_error_recovery.rs b/crates/aionui-conversation/src/acp_error_recovery.rs index 23b84a9a6..9e0f3f05f 100644 --- a/crates/aionui-conversation/src/acp_error_recovery.rs +++ b/crates/aionui-conversation/src/acp_error_recovery.rs @@ -123,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, diff --git a/crates/aionui-conversation/src/runtime_completion.rs b/crates/aionui-conversation/src/runtime_completion.rs index 9c6413d3d..6b7161550 100644 --- a/crates/aionui-conversation/src/runtime_completion.rs +++ b/crates/aionui-conversation/src/runtime_completion.rs @@ -81,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/service.rs b/crates/aionui-conversation/src/service.rs index 7f3db2979..74b4b6999 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -1199,7 +1199,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); @@ -1992,7 +1992,7 @@ impl ConversationService { 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) } @@ -2136,7 +2136,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(()) } @@ -2450,11 +2450,13 @@ impl ConversationService { })?; 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) } @@ -2553,6 +2555,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, }); @@ -2675,6 +2678,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, @@ -2905,6 +2909,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, @@ -2932,6 +2937,7 @@ impl ConversationService { 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, @@ -3445,11 +3451,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, diff --git a/crates/aionui-conversation/src/stream_persistence.rs b/crates/aionui-conversation/src/stream_persistence.rs index b0568f53a..dc9acfac0 100644 --- a/crates/aionui-conversation/src/stream_persistence.rs +++ b/crates/aionui-conversation/src/stream_persistence.rs @@ -124,6 +124,7 @@ impl StreamPersistenceAdapter { } let payload = json!({ + "user_id": self.user_id, "conversation_id": self.conversation_id, "session_id": self.conversation_id, "turn_id": turn_id, diff --git a/crates/aionui-conversation/src/stream_relay.rs b/crates/aionui-conversation/src/stream_relay.rs index 0b3ce7fbb..a38f6a1a4 100644 --- a/crates/aionui-conversation/src/stream_relay.rs +++ b/crates/aionui-conversation/src/stream_relay.rs @@ -768,6 +768,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())); } From 7c76351b5f0415c5a4a5862d23793731240b417e Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 18:24:47 +0800 Subject: [PATCH 031/145] db: use owner-scoped channel constraints --- .../aionui-db/migrations/026_user_scope.sql | 78 ++++++++++++-- crates/aionui-db/src/models/channel.rs | 2 +- .../src/repository/sqlite_channel.rs | 102 ++++++++++-------- 3 files changed, 131 insertions(+), 51 deletions(-) diff --git a/crates/aionui-db/migrations/026_user_scope.sql b/crates/aionui-db/migrations/026_user_scope.sql index 74fb17229..b0a29e6a3 100644 --- a/crates/aionui-db/migrations/026_user_scope.sql +++ b/crates/aionui-db/migrations/026_user_scope.sql @@ -586,18 +586,84 @@ ALTER TABLE skill_import_records CREATE INDEX IF NOT EXISTS idx_skill_import_records_user_created_at ON skill_import_records(user_id, created_at DESC); -ALTER TABLE assistant_plugins - ADD COLUMN owner_user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id); +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; + +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); -ALTER TABLE assistant_users - ADD COLUMN owner_user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id); +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; + +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); -ALTER TABLE assistant_pairing_codes - ADD COLUMN owner_user_id TEXT NOT NULL DEFAULT 'system_default_user' REFERENCES users(id); +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; + +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); diff --git a/crates/aionui-db/src/models/channel.rs b/crates/aionui-db/src/models/channel.rs index 086f94bc9..cbadca203 100644 --- a/crates/aionui-db/src/models/channel.rs +++ b/crates/aionui-db/src/models/channel.rs @@ -25,7 +25,7 @@ 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, diff --git a/crates/aionui-db/src/repository/sqlite_channel.rs b/crates/aionui-db/src/repository/sqlite_channel.rs index 0c6c093e5..0376f3e00 100644 --- a/crates/aionui-db/src/repository/sqlite_channel.rs +++ b/crates/aionui-db/src/repository/sqlite_channel.rs @@ -16,27 +16,6 @@ impl SqliteChannelRepository { } } -fn plugin_storage_id(owner_user_id: &str, plugin_id: &str) -> String { - if owner_user_id == "system_default_user" { - plugin_id.to_owned() - } else { - format!("{owner_user_id}:{plugin_id}") - } -} - -fn plugin_logical_id(owner_user_id: &str, stored_id: &str) -> String { - let prefix = format!("{owner_user_id}:"); - stored_id - .strip_prefix(&prefix) - .map(str::to_owned) - .unwrap_or_else(|| stored_id.to_owned()) -} - -fn plugin_row_with_logical_id(owner_user_id: &str, mut row: ChannelPluginRow) -> ChannelPluginRow { - row.id = plugin_logical_id(owner_user_id, &row.id); - row -} - #[async_trait::async_trait] impl IChannelRepository for SqliteChannelRepository { // ── Plugin CRUD ────────────────────────────────────────────────── @@ -48,34 +27,25 @@ impl IChannelRepository for SqliteChannelRepository { .bind(owner_user_id) .fetch_all(&self.pool) .await?; - Ok(rows - .into_iter() - .map(|row| plugin_row_with_logical_id(owner_user_id, row)) - .collect()) + Ok(rows) } async fn get_plugin(&self, owner_user_id: &str, id: &str) -> Result, DbError> { - let storage_id = plugin_storage_id(owner_user_id, id); - let row = sqlx::query_as::<_, ChannelPluginRow>( - "SELECT * FROM assistant_plugins WHERE owner_user_id = ? AND id IN (?, ?) ORDER BY id = ? DESC LIMIT 1", - ) - .bind(owner_user_id) - .bind(&storage_id) - .bind(id) - .bind(&storage_id) - .fetch_optional(&self.pool) - .await?; - Ok(row.map(|row| plugin_row_with_logical_id(owner_user_id, row))) + 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, owner_user_id: &str, row: &ChannelPluginRow) -> Result<(), DbError> { - let storage_id = plugin_storage_id(owner_user_id, &row.id); sqlx::query( "INSERT INTO assistant_plugins \ (id, owner_user_id, type, name, enabled, config, status, last_connected, created_at, updated_at) \ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \ - ON CONFLICT(id) DO UPDATE SET \ - owner_user_id = excluded.owner_user_id, \ + ON CONFLICT(owner_user_id, id) DO UPDATE SET \ type = excluded.type, \ name = excluded.name, \ enabled = excluded.enabled, \ @@ -84,7 +54,7 @@ impl IChannelRepository for SqliteChannelRepository { last_connected = excluded.last_connected, \ updated_at = excluded.updated_at", ) - .bind(&storage_id) + .bind(&row.id) .bind(owner_user_id) .bind(&row.r#type) .bind(&row.name) @@ -105,7 +75,6 @@ impl IChannelRepository for SqliteChannelRepository { id: &str, params: &UpdatePluginStatusParams, ) -> Result<(), DbError> { - let storage_id = plugin_storage_id(owner_user_id, id); let mut set_clauses = Vec::new(); if params.status.is_some() { set_clauses.push("status = ?"); @@ -141,7 +110,7 @@ impl IChannelRepository for SqliteChannelRepository { } query = query.bind(now); query = query.bind(owner_user_id); - query = query.bind(&storage_id); + query = query.bind(id); let result = query.execute(&self.pool).await?; if result.rows_affected() == 0 { @@ -151,10 +120,9 @@ impl IChannelRepository for SqliteChannelRepository { } async fn delete_plugin(&self, owner_user_id: &str, id: &str) -> Result<(), DbError> { - let storage_id = plugin_storage_id(owner_user_id, id); let result = sqlx::query("DELETE FROM assistant_plugins WHERE owner_user_id = ? AND id = ?") .bind(owner_user_id) - .bind(&storage_id) + .bind(id) .execute(&self.pool) .await?; if result.rows_affected() == 0 { @@ -894,6 +862,33 @@ mod tests { ); } + #[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; @@ -1285,6 +1280,25 @@ mod tests { 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; From 4e6c61583071b14141158c76ea33c412e88ae75b Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 18:36:17 +0800 Subject: [PATCH 032/145] team: include user scope in websocket events --- crates/aionui-team/src/events.rs | 55 ++-- crates/aionui-team/src/message_projection.rs | 10 + crates/aionui-team/src/scheduler/mod.rs | 3 +- crates/aionui-team/src/scheduler/tests.rs | 245 ++++++++++++++++-- crates/aionui-team/src/service.rs | 151 +++++++---- crates/aionui-team/src/session.rs | 22 +- crates/aionui-team/src/team_run/tests.rs | 8 +- crates/aionui-team/tests/e2e_smoke.rs | 1 + .../tests/mcp_server_integration.rs | 1 + .../tests/prompts_events_integration.rs | 16 +- .../tests/scheduler_integration.rs | 1 + 11 files changed, 387 insertions(+), 126 deletions(-) diff --git a/crates/aionui-team/src/events.rs b/crates/aionui-team/src/events.rs index c33457c40..2d2922793 100644 --- a/crates/aionui-team/src/events.rs +++ b/crates/aionui-team/src/events.rs @@ -5,6 +5,8 @@ use aionui_api_types::{ TeamAgentSpawnedPayload, TeamAgentStatusPayload, TeamChildTurnPayload, TeamRunPayload, WebSocketMessage, }; use aionui_realtime::EventBroadcaster; +use serde::Serialize; +use serde_json::Value; use tracing::debug; use crate::types::{TeamAgent, TeammateStatus}; @@ -34,28 +36,36 @@ pub const TEAM_CHILD_TURN_CANCELLED_EVENT: &str = "team.childTurnCancelled"; 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); } @@ -64,10 +74,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); } @@ -76,10 +83,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); } @@ -89,10 +93,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); } @@ -109,10 +110,7 @@ impl TeamEventEmitter { status, error, }; - let event = WebSocketMessage::new( - TEAM_AGENT_RUNTIME_STATUS_CHANGED_EVENT, - serde_json::to_value(payload).expect("serialize agent runtime status payload"), - ); + let event = WebSocketMessage::new(TEAM_AGENT_RUNTIME_STATUS_CHANGED_EVENT, self.scoped_payload(payload)); self.broadcaster.broadcast(event); } @@ -131,10 +129,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); } @@ -150,10 +145,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); } } @@ -191,7 +183,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) } @@ -203,6 +195,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/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/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 7504a6129..b3c41446b 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -354,7 +354,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; } @@ -432,7 +432,7 @@ 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 } @@ -495,7 +495,7 @@ impl TeamSessionService { 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(()) } @@ -512,7 +512,7 @@ impl TeamSessionService { }, ) .await?; - self.broadcast_team_renamed(team_id, name); + self.broadcast_team_renamed(user_id, team_id, name); Ok(()) } @@ -544,7 +544,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, @@ -563,7 +563,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, @@ -722,7 +723,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, @@ -808,10 +810,10 @@ impl TeamSessionService { .get_team(user_id, team_id) .await? .ok_or_else(|| TeamError::TeamNotFound(team_id.to_owned()))?; - self.ensure_session_inner(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()) @@ -822,25 +824,31 @@ impl TeamSessionService { 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()); } }; @@ -877,6 +885,7 @@ impl TeamSessionService { } self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Starting, Some(TeamSessionPhase::LoadingTeam), @@ -884,6 +893,7 @@ impl TeamSessionService { ); self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Starting, Some(TeamSessionPhase::StartingBridge), @@ -908,6 +918,7 @@ impl TeamSessionService { Ok(session) => Arc::new(session), Err(e) => { self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Failed, Some(TeamSessionPhase::StartingBridge), @@ -920,6 +931,7 @@ impl TeamSessionService { }; self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Starting, Some(TeamSessionPhase::AttachingAgents), @@ -931,6 +943,7 @@ impl TeamSessionService { .await { self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Failed, Some(TeamSessionPhase::AttachingAgents), @@ -953,6 +966,7 @@ impl TeamSessionService { session.stop(); self.cleanup_bootstrap_runtime_tasks(&agents_snapshot).await; self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Failed, Some(TeamSessionPhase::AttachingAgents), @@ -985,10 +999,11 @@ impl TeamSessionService { drop(membership_guard); for agent in &agents_snapshot { - self.broadcast_agent_runtime_status(team_id, agent, TeamAgentRuntimeStatus::Ready, None); + self.broadcast_agent_runtime_status(&user_id, team_id, agent, TeamAgentRuntimeStatus::Ready, None); } self.broadcast_session_status( + &user_id, team_id, TeamSessionStatus::Starting, Some(TeamSessionPhase::Recovering), @@ -1003,7 +1018,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()); }); @@ -1039,6 +1054,7 @@ impl TeamSessionService { match reservation { ReserveAttach::Start(owner) => { self.broadcast_agent_runtime_status( + session.user_id(), session.team_id(), agent, TeamAgentRuntimeStatus::Pending, @@ -1143,6 +1159,7 @@ impl TeamSessionService { AttachOutcome::Ready | AttachOutcome::Removed => {} AttachOutcome::Failed(failure) => { self.broadcast_session_status( + user_id, team_id, TeamSessionStatus::Failed, Some(TeamSessionPhase::AttachingAgents), @@ -1163,7 +1180,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(()) @@ -1223,6 +1240,7 @@ impl TeamSessionService { fn broadcast_session_status( &self, + user_id: &str, team_id: &str, status: TeamSessionStatus, phase: Option, @@ -1238,10 +1256,9 @@ impl TeamSessionService { error: None, }; customize(&mut payload); - let event = WebSocketMessage::new( - TEAM_SESSION_STATUS_CHANGED_EVENT, - serde_json::to_value(payload).expect("serialize team session status payload"), - ); + 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); } @@ -1257,49 +1274,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); } @@ -1368,7 +1386,7 @@ impl TeamSessionService { } let cfg = session.mcp_stdio_config(&agent.slot_id); - 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_rebuild_agent_process( &mut jobs, provisioner.clone(), @@ -1428,6 +1446,7 @@ impl TeamSessionService { .map(ToString::to_string) .unwrap_or_else(|| "unknown rebuild failure".to_owned()); self.broadcast_agent_runtime_status( + user_id, team_id, &failure.agent, TeamAgentRuntimeStatus::Failed, @@ -1560,7 +1579,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() } @@ -1568,6 +1593,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), @@ -1580,6 +1606,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), @@ -1624,9 +1651,15 @@ impl TeamSessionService { self.publish_member_runtime_starting_if_current(expected); } else { 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()); + }, + ); }); } } @@ -1859,7 +1892,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 @@ -1882,7 +1923,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 @@ -1902,7 +1943,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 @@ -1922,7 +1963,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 @@ -1942,7 +1983,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 @@ -1962,7 +2003,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 diff --git a/crates/aionui-team/src/session.rs b/crates/aionui-team/src/session.rs index 72e69430e..f73c2dd0a 100644 --- a/crates/aionui-team/src/session.rs +++ b/crates/aionui-team/src/session.rs @@ -171,7 +171,11 @@ impl TeamSession { 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 { @@ -555,6 +564,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, @@ -653,6 +663,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(), @@ -892,6 +903,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(), @@ -1250,6 +1262,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, @@ -1586,6 +1599,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, @@ -2785,7 +2799,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); @@ -2803,7 +2817,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/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/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/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(), From 687409137dc80142b973fc8ea122dadaf462ebbd Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 18:46:35 +0800 Subject: [PATCH 033/145] cron: include user scope in websocket events --- crates/aionui-cron/src/artifacts.rs | 4 +- crates/aionui-cron/src/events.rs | 46 +++++++------ crates/aionui-cron/src/executor.rs | 4 +- crates/aionui-cron/src/service.rs | 89 ++++++++++++++++++++----- crates/aionui-cron/src/skill_suggest.rs | 2 +- 5 files changed, 105 insertions(+), 40 deletions(-) diff --git a/crates/aionui-cron/src/artifacts.rs b/crates/aionui-cron/src/artifacts.rs index f0097f016..11886d196 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(()) } 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 a0908a5ec..26279fb74 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -688,7 +688,7 @@ impl JobExecutor { let row = build_cron_trigger_artifact(&started.conversation_id, &job, created_at); 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, @@ -844,7 +844,7 @@ impl JobExecutor { .map_err(CronError::Database)?; for row in rows { - broadcast_artifact(&self.broadcaster, &row)?; + broadcast_artifact(&self.broadcaster, SYSTEM_DEFAULT_USER_ID, &row)?; } Ok(()) diff --git a/crates/aionui-cron/src/service.rs b/crates/aionui-cron/src/service.rs index 60771efcc..80bfbbc54 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -327,7 +327,7 @@ 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) @@ -445,7 +445,7 @@ 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) @@ -457,7 +457,7 @@ impl CronService { warn!(job_id, error = %err, "Failed to delete cron skill file during job removal"); } self.repo.delete_for_user(user_id, job_id).await?; - self.emitter.emit_job_removed(job_id); + self.emitter.emit_job_removed(user_id, job_id); info!(job_id, "Cron job removed"); Ok(()) } @@ -605,7 +605,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) => { @@ -698,7 +700,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; } @@ -727,10 +731,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 }) @@ -833,6 +838,43 @@ impl CronService { self.resolve_agent_type_for_assistant_id(assistant_id).await } + async fn owner_user_id_for_job(&self, job: &CronJob) -> Option { + 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 is_team_conversation_job(&self, job: &CronJob) -> Result { let conversation_id = job.conversation_id.trim(); if conversation_id.is_empty() { @@ -967,6 +1009,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 } => { @@ -978,7 +1021,9 @@ impl CronService { } self.update_job_after_success(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 { @@ -1008,7 +1053,9 @@ impl CronService { 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 @@ -1019,20 +1066,22 @@ impl CronService { } self.update_job_after_error(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.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.emitter.emit_job_executed(user_id, job_id, "error", Some(&message)); } ExecutionResult::Retrying { attempt } => { let params = UpdateCronJobParams { @@ -1060,7 +1109,7 @@ impl CronService { "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); } } } @@ -1164,7 +1213,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; @@ -1216,9 +1268,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, diff --git a/crates/aionui-cron/src/skill_suggest.rs b/crates/aionui-cron/src/skill_suggest.rs index 1c9489082..0c3d479a7 100644 --- a/crates/aionui-cron/src/skill_suggest.rs +++ b/crates/aionui-cron/src/skill_suggest.rs @@ -155,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, From 4bfffedbffb4264d0568da1f6b5ca2f33584b80c Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 18:49:53 +0800 Subject: [PATCH 034/145] app: whitelist global extension websocket events --- crates/aionui-app/src/router/routes.rs | 5 ++++- crates/aionui-app/tests/websocket_e2e.rs | 10 ++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index 8526924be..a27691a6f 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -330,7 +330,10 @@ fn auth_identity_mode(identity_mode: crate::config::IdentityMode) -> AuthIdentit } fn is_global_websocket_event(event_name: &str) -> bool { - matches!(event_name, "runtime.statusChanged") + matches!( + event_name, + "runtime.statusChanged" | "extensions.state-changed" | "extensions.lifecycle" | "hub.state-changed" + ) } async fn normalize_boundary_error_response(request: Request, next: Next) -> Response { diff --git a/crates/aionui-app/tests/websocket_e2e.rs b/crates/aionui-app/tests/websocket_e2e.rs index acdc0af88..01ebf2215 100644 --- a/crates/aionui-app/tests/websocket_e2e.rs +++ b/crates/aionui-app/tests/websocket_e2e.rs @@ -412,6 +412,16 @@ async fn t4_1_whitelisted_global_event_reaches_all_users() { 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] From 0e5079c0a75c02d16eb4306f5afa3d8ff35581cd Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 19:11:56 +0800 Subject: [PATCH 035/145] file: scope filesystem websocket events by user --- Cargo.lock | 1 + crates/aionui-file/Cargo.toml | 1 + crates/aionui-file/src/routes.rs | 43 ++++-- crates/aionui-file/src/service.rs | 21 ++- crates/aionui-file/src/traits.rs | 50 +++++++ crates/aionui-file/src/watch_service.rs | 125 +++++++++++++----- crates/aionui-file/tests/file_watching.rs | 2 + crates/aionui-office/src/routes.rs | 18 ++- crates/aionui-office/src/watch_manager.rs | 38 ++++-- .../tests/watch_manager_integration.rs | 29 ++++ 10 files changed, 270 insertions(+), 58 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 77e68ef8a..49df9f92e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -667,6 +667,7 @@ name = "aionui-file" version = "0.1.50" dependencies = [ "aionui-api-types", + "aionui-auth", "aionui-common", "aionui-realtime", "async-trait", 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..954b460e5 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; @@ -256,6 +257,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 +269,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 +288,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 +298,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 +484,63 @@ 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())) } 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..410163260 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,49 @@ 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 + } } /// 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..7ec04c2be 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() diff --git a/crates/aionui-file/tests/file_watching.rs b/crates/aionui-file/tests/file_watching.rs index 0a2e3380f..fb6f1165d 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] @@ -286,4 +287,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-office/src/routes.rs b/crates/aionui-office/src/routes.rs index 163cd965f..ed4975de0 100644 --- a/crates/aionui-office/src/routes.rs +++ b/crates/aionui-office/src/routes.rs @@ -84,10 +84,10 @@ 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( @@ -100,10 +100,10 @@ async fn stop_word_preview( 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( @@ -116,10 +116,10 @@ async fn stop_excel_preview( 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( @@ -132,6 +132,7 @@ async fn stop_ppt_preview( 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) => { diff --git a/crates/aionui-office/src/watch_manager.rs b/crates/aionui-office/src/watch_manager.rs index 7af031376..dc1ee8667 100644 --- a/crates/aionui-office/src/watch_manager.rs +++ b/crates/aionui-office/src/watch_manager.rs @@ -78,6 +78,10 @@ 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); @@ -89,20 +93,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,10 +127,13 @@ 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?; @@ -158,6 +170,7 @@ impl OfficecliWatchManager { async fn spawn_officecli_with_install( &self, + user_id: &str, resolved: &str, port: u16, doc_type: DocType, @@ -165,7 +178,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 } @@ -243,16 +256,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)); } } 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")) + ); } // --------------------------------------------------------------------------- From 4c5f95bacdab8ad19d193d7838454e3b8086011a Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 19:24:56 +0800 Subject: [PATCH 036/145] system: avoid default user keep-awake restore outside local mode --- crates/aionui-app/src/router/state.rs | 28 +++++++++++------- crates/aionui-app/tests/assistants_e2e.rs | 1 + crates/aionui-app/tests/cron_e2e.rs | 2 ++ crates/aionui-app/tests/extension_e2e.rs | 36 +++++++++++++++-------- crates/aionui-system/src/client_pref.rs | 32 +++++++++++++++++--- 5 files changed, 72 insertions(+), 27 deletions(-) diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index 10f7fb955..8a11b3648 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -380,13 +380,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()), - "system_default_user", - ), + 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()), @@ -1062,15 +1066,19 @@ 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(); let settings = build_channel_settings_service(&services); - let message_service = build_channel_message_service(&services, settings).await; + let message_service = + build_channel_message_service(&services, settings, "system_default_user".to_owned()).await; let session = AssistantSessionRow { id: "session-channel-state".to_owned(), user_id: "channel-user-state".to_owned(), diff --git a/crates/aionui-app/tests/assistants_e2e.rs b/crates/aionui-app/tests/assistants_e2e.rs index ecdd20c47..7d02b1e6e 100644 --- a/crates/aionui-app/tests/assistants_e2e.rs +++ b/crates/aionui-app/tests/assistants_e2e.rs @@ -289,6 +289,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", diff --git a/crates/aionui-app/tests/cron_e2e.rs b/crates/aionui-app/tests/cron_e2e.rs index fabdcbfff..a9026fcd4 100644 --- a/crates/aionui-app/tests/cron_e2e.rs +++ b/crates/aionui-app/tests/cron_e2e.rs @@ -686,6 +686,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 +700,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, diff --git a/crates/aionui-app/tests/extension_e2e.rs b/crates/aionui-app/tests/extension_e2e.rs index ea2e7d30a..ae46be1c7 100644 --- a/crates/aionui-app/tests/extension_e2e.rs +++ b/crates/aionui-app/tests/extension_e2e.rs @@ -617,22 +617,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 +669,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 +692,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")); diff --git a/crates/aionui-system/src/client_pref.rs b/crates/aionui-system/src/client_pref.rs index 07c4c0fd8..3030383f3 100644 --- a/crates/aionui-system/src/client_pref.rs +++ b/crates/aionui-system/src/client_pref.rs @@ -31,14 +31,21 @@ impl ClientPrefService { keep_awake_controller: DynKeepAwakeController, keep_awake_restore_user_id: impl Into, ) -> Self { - let service = Self { - repo, - keep_awake_controller, - }; + 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, + } + } + /// Get all client preferences, or only the specified keys. pub async fn get_preferences( &self, @@ -632,4 +639,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); + } } From 2fe6fddd11364293eb6b55760c47c97be7be26fb Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 19:33:30 +0800 Subject: [PATCH 037/145] cron: require conversation owner for background execution --- crates/aionui-cron/src/executor.rs | 23 +++++++++++++++++------ crates/aionui-cron/src/service.rs | 4 +++- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/crates/aionui-cron/src/executor.rs b/crates/aionui-cron/src/executor.rs index 26279fb74..5da16ab37 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 { @@ -612,6 +611,13 @@ impl JobExecutor { } async fn resolve_conversation_owner_user_id(&self, job: &CronJob) -> Result { + 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(user_id) = self .conversation_repo @@ -623,7 +629,10 @@ impl JobExecutor { 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( @@ -830,21 +839,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(SYSTEM_DEFAULT_USER_ID, 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, SYSTEM_DEFAULT_USER_ID, &row)?; + broadcast_artifact(&self.broadcaster, user_id, &row)?; } Ok(()) diff --git a/crates/aionui-cron/src/service.rs b/crates/aionui-cron/src/service.rs index 80bfbbc54..31dce391f 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -761,7 +761,9 @@ impl CronService { ..Default::default() }; self.repo.update_for_user(user_id, job_id, ¶ms).await?; - self.executor.mark_skill_suggest_artifacts_saved(job_id).await?; + self.executor + .mark_skill_suggest_artifacts_saved(user_id, job_id) + .await?; info!(job_id, "Skill content saved"); Ok(()) From 97e50930164696a2904a387aade99b61caea8924 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 19:48:01 +0800 Subject: [PATCH 038/145] channel: require explicit owner in aionpro mode --- crates/aionui-app/src/router/routes.rs | 19 +++-- crates/aionui-app/src/router/state.rs | 30 +++---- crates/aionui-channel/src/action.rs | 78 +++++++++++++------ crates/aionui-channel/src/message_service.rs | 23 +++--- crates/aionui-channel/src/orchestrator.rs | 14 +++- .../tests/message_service_integration.rs | 48 ++++-------- .../aionui-channel/tests/orchestrator_test.rs | 2 +- .../tests/session_action_integration.rs | 2 +- 8 files changed, 118 insertions(+), 98 deletions(-) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index a27691a6f..bd7db5a66 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -108,13 +108,20 @@ pub async fn create_router_with_runtime(services: &AppServices) -> Result<(Route 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_owner_user_id, &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", - owner_user_id = %chan_owner_user_id, - error = %e, - "failed to restore channel plugins" + "skipping channel plugin restore until an owner user is available" ); } }); diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index 8a11b3648..2018ede46 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -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,7 +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: String, + pub owner_user_id: Option, } /// Build all default `ModuleStates` from application services. @@ -512,17 +512,19 @@ fn build_channel_settings_service( async fn build_channel_message_service( services: &AppServices, channel_settings: Arc, - owner_user_id: String, ) -> Arc { Arc::new(aionui_channel::message_service::ChannelMessageService::new( Arc::new(services.conversation_service.clone()), services.worker_task_manager.clone(), channel_settings, - owner_user_id, )) } -async fn primary_channel_owner_user_id(services: &AppServices) -> String { +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() @@ -531,7 +533,7 @@ async fn primary_channel_owner_user_id(services: &AppServices) -> String { .flatten() .map(|u| u.id) .unwrap_or_else(|| "system_default_user".to_string()); - owner_user_id + Some(owner_user_id) } /// Build the default `ChannelRouterState` and orchestrator components. @@ -566,18 +568,17 @@ 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 runtime_owner_user_id = primary_channel_owner_user_id(services).await; + 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), - runtime_owner_user_id.clone(), + startup_owner_user_id.clone(), )); - let message_service = - build_channel_message_service(services, Arc::clone(&channel_settings), runtime_owner_user_id.clone()).await; + let message_service = build_channel_message_service(services, Arc::clone(&channel_settings)).await; let orchestrator = aionui_channel::orchestrator::ChannelOrchestrator::new( action_executor, @@ -602,7 +603,7 @@ pub async fn build_channel_state( confirm_rx, manager, plugin_factory, - owner_user_id: runtime_owner_user_id, + owner_user_id: startup_owner_user_id, }; (state, components) @@ -1077,8 +1078,7 @@ mod tests { .unwrap(); let settings = build_channel_settings_service(&services); - let message_service = - build_channel_message_service(&services, settings, "system_default_user".to_owned()).await; + let message_service = build_channel_message_service(&services, settings).await; let session = AssistantSessionRow { id: "session-channel-state".to_owned(), user_id: "channel-user-state".to_owned(), @@ -1091,7 +1091,7 @@ 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(); @@ -1122,7 +1122,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-channel/src/action.rs b/crates/aionui-channel/src/action.rs index cb37363d8..88dd85b4f 100644 --- a/crates/aionui-channel/src/action.rs +++ b/crates/aionui-channel/src/action.rs @@ -41,7 +41,7 @@ pub struct ActionExecutor { pairing: Arc, session_mgr: Arc, settings: Arc, - owner_user_id: String, + owner_user_id: Option, } impl ActionExecutor { @@ -49,7 +49,7 @@ impl ActionExecutor { pairing: Arc, session_mgr: Arc, settings: Arc, - owner_user_id: String, + owner_user_id: Option, ) -> Self { Self { pairing, @@ -59,8 +59,8 @@ impl ActionExecutor { } } - pub fn owner_user_id(&self) -> &str { - &self.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. @@ -73,18 +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(&self.owner_user_id, user_id, &platform_type) + .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)); } @@ -92,19 +97,16 @@ 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(&self.owner_user_id, 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( - &self.owner_user_id, + owner_user_id, &internal_user_id, chat_id, &agent_config.agent_type, @@ -121,7 +123,7 @@ impl ActionExecutor { ); Ok(MessageResult::Dispatched { - owner_user_id: self.owner_user_id.clone(), + owner_user_id: owner_user_id.to_owned(), session_id: session.id, conversation_id: session.conversation_id, }) @@ -131,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(&self.owner_user_id, platform_user_id, platform_type, Some(display_name)) + .request_pairing(owner_user_id, platform_user_id, platform_type, Some(display_name)) .await?; debug!( @@ -152,25 +155,30 @@ 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( - &self.owner_user_id, + owner_user_id, &action.context.user_id, &action.context.platform.to_string(), None, @@ -182,7 +190,7 @@ impl ActionExecutor { let authorized = self .pairing .is_user_authorized( - &self.owner_user_id, + owner_user_id, &action.context.user_id, &action.context.platform.to_string(), ) @@ -247,6 +255,7 @@ impl ActionExecutor { async fn handle_system_action( &self, + owner_user_id: &str, action: &UnifiedAction, internal_user_id: &str, ) -> Result { @@ -256,11 +265,11 @@ impl ActionExecutor { let chat_id = &action.context.chat_id; let agent_config = self .settings - .get_agent_config(internal_user_id, action.context.platform) + .get_agent_config(owner_user_id, action.context.platform) .await?; let session = self .session_mgr - .reset_session(&self.owner_user_id, 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 { @@ -286,11 +295,11 @@ impl ActionExecutor { let chat_id = &action.context.chat_id; let agent_config = self .settings - .get_agent_config(&self.owner_user_id, action.context.platform) + .get_agent_config(owner_user_id, action.context.platform) .await?; let session = self .session_mgr - .get_or_create_session(&self.owner_user_id, 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 { @@ -770,7 +779,18 @@ 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, OWNER_ID.to_owned()); + 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) } @@ -859,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/message_service.rs b/crates/aionui-channel/src/message_service.rs index 7a4608210..2f28b7129 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,19 +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(&self.owner_user_id, 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(&self.owner_user_id, 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()) @@ -144,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(&self.owner_user_id, 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() { @@ -185,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 4a8da3d62..092e025e1 100644 --- a/crates/aionui-channel/src/orchestrator.rs +++ b/crates/aionui-channel/src/orchestrator.rs @@ -67,8 +67,10 @@ impl ChannelOrchestrator { let chat_id = msg.chat_id.clone(); let plugin_id = platform.to_string(); let text = msg.content.text.clone(); - let fallback_owner_user_id = self.action_executor.owner_user_id().to_owned(); - let message_owner_user_id = msg.owner_user_id.clone().unwrap_or(fallback_owner_user_id); + 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); @@ -78,7 +80,11 @@ impl ChannelOrchestrator { tokio::spawn(async move { match executor.handle_incoming_message(&msg).await { Ok(MessageResult::Action(response)) => { - send_action_response(&sender, &message_owner_user_id, &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, @@ -172,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"); diff --git a/crates/aionui-channel/tests/message_service_integration.rs b/crates/aionui-channel/tests/message_service_integration.rs index ed68a7e13..7ec82fc80 100644 --- a/crates/aionui-channel/tests/message_service_integration.rs +++ b/crates/aionui-channel/tests/message_service_integration.rs @@ -239,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(), @@ -263,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(), @@ -319,12 +317,7 @@ async fn send_to_agent_persists_assistant_snapshot_for_channel_bound_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".to_owned(), @@ -338,7 +331,7 @@ 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(); @@ -401,12 +394,7 @@ async fn send_to_agent_rejects_unresolvable_channel_assistant_binding() { 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(), @@ -420,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(_))); @@ -466,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(), @@ -485,7 +468,7 @@ 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(); @@ -549,12 +532,7 @@ async fn send_to_agent_without_assistant_name_falls_back_to_legacy_channel_name( .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(), @@ -568,7 +546,7 @@ 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(); diff --git a/crates/aionui-channel/tests/orchestrator_test.rs b/crates/aionui-channel/tests/orchestrator_test.rs index 6e2a1b30a..a1527ab34 100644 --- a/crates/aionui-channel/tests/orchestrator_test.rs +++ b/crates/aionui-channel/tests/orchestrator_test.rs @@ -52,7 +52,7 @@ async fn unauthorized_user_gets_pairing_response() { pairing, Arc::clone(&session_mgr), settings, - OWNER_ID.to_owned(), + Some(OWNER_ID.to_owned()), )); let msg = make_text_message("unknown_user", "chat_1", "hello"); diff --git a/crates/aionui-channel/tests/session_action_integration.rs b/crates/aionui-channel/tests/session_action_integration.rs index a4c748aba..b84292cba 100644 --- a/crates/aionui-channel/tests/session_action_integration.rs +++ b/crates/aionui-channel/tests/session_action_integration.rs @@ -58,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, OWNER_ID.to_owned()); + let executor = ActionExecutor::new(pairing_arc, session_mgr_arc, settings, Some(OWNER_ID.to_owned())); // Keep db alive std::mem::forget(db); From f02826589a959e437c19995a64a7e0a82f689ab7 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 19:54:09 +0800 Subject: [PATCH 039/145] agent: scope availability feedback by conversation user --- .../src/services/availability/mod.rs | 72 +++++++++++-------- .../aionui-conversation/src/service_test.rs | 26 +++++-- .../src/turn_orchestrator.rs | 20 ++++-- 3 files changed, 79 insertions(+), 39 deletions(-) diff --git a/crates/aionui-ai-agent/src/services/availability/mod.rs b/crates/aionui-ai-agent/src/services/availability/mod.rs index bcab97712..1df808a0e 100644 --- a/crates/aionui-ai-agent/src/services/availability/mod.rs +++ b/crates/aionui-ai-agent/src/services/availability/mod.rs @@ -18,12 +18,16 @@ use crate::error::AgentError; use crate::protocol::{cli_detect, custom_agent_probe}; use crate::registry::{AgentRegistry, guidance_for_snapshot_error_code}; -const SYSTEM_DEFAULT_USER_ID: &str = "system_default_user"; - #[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 { @@ -83,7 +87,13 @@ impl AgentAvailabilityService { .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", @@ -93,10 +103,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", @@ -106,7 +116,7 @@ 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( @@ -117,11 +127,6 @@ impl AgentAvailabilityService { self.registry.management_row_by_id_for_user(user_id, id).await } - async fn persist_snapshot(&self, id: &str, snapshot: &AvailabilitySnapshot) -> Result<(), AgentError> { - self.persist_snapshot_for_user(SYSTEM_DEFAULT_USER_ID, id, snapshot) - .await - } - async fn persist_snapshot_for_user( &self, user_id: &str, @@ -163,9 +168,6 @@ impl AgentAvailabilityService { .update_availability_snapshot_for_user(user_id, id, ¶ms) .await .map_err(|error| AgentError::internal(format!("repo.update_availability_snapshot_for_user: {error}")))?; - if user_id == SYSTEM_DEFAULT_USER_ID { - self.registry.reload_one(id).await?; - } Ok(()) } } @@ -386,12 +388,18 @@ async fn try_connect_builtin_managed_agent( #[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 } } @@ -399,10 +407,7 @@ impl AgentAvailabilityFeedbackPort for AgentAvailabilityService { mod tests { use std::sync::Arc; - use super::{ - AgentAvailabilityService, SYSTEM_DEFAULT_USER_ID, explicit_probe_args, probe_aionrs_provider_readiness, - run_probe, - }; + use super::{AgentAvailabilityService, explicit_probe_args, probe_aionrs_provider_readiness, run_probe}; use crate::registry::AgentRegistry; use aionui_api_types::{ AgentHandshake, AgentManagementStatus, AgentMetadata, AgentSnapshotCheckKind, AgentSnapshotCheckStatus, @@ -414,10 +419,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: "system_default_user", + user_id: TEST_USER_ID, platform: "openai", name: "OpenAI", base_url: "https://api.openai.com", @@ -440,7 +447,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, "system_default_user").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")); @@ -452,7 +459,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, "system_default_user").await; + let (status, code, _msg) = probe_aionrs_provider_readiness(&provider_repo, TEST_USER_ID).await; assert_eq!(status, AgentSnapshotCheckStatus::Online); assert!(code.is_none()); @@ -499,6 +506,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", @@ -507,7 +515,7 @@ mod tests { .unwrap(); let row = service - .list_management_rows(SYSTEM_DEFAULT_USER_ID) + .list_management_rows(TEST_USER_ID) .await .unwrap() .into_iter() @@ -572,6 +580,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", @@ -579,10 +588,13 @@ 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(SYSTEM_DEFAULT_USER_ID) + .list_management_rows(TEST_USER_ID) .await .unwrap() .into_iter() diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index 9e3c13212..8c2bb8d5c 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -174,6 +174,7 @@ impl EventBroadcaster for MockBroadcaster { #[derive(Debug, Clone, PartialEq, Eq)] struct RecordedAvailabilityFailure { + user_id: String, agent_id: String, code: String, message: String, @@ -181,19 +182,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(), @@ -5181,11 +5192,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(), @@ -5226,7 +5239,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()); } diff --git a/crates/aionui-conversation/src/turn_orchestrator.rs b/crates/aionui-conversation/src/turn_orchestrator.rs index bc56578c7..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, @@ -268,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(); @@ -277,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, @@ -511,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 @@ -519,7 +523,12 @@ 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); @@ -598,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, @@ -608,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), @@ -618,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" From eb16f5a30fc82c42b21693c5a31eac0ed17902bf Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 20:00:48 +0800 Subject: [PATCH 040/145] auth: bind sessions to user generation --- crates/aionui-auth/src/jwt.rs | 26 ++++++++++++++++ crates/aionui-auth/src/middleware.rs | 4 +++ crates/aionui-auth/src/routes.rs | 32 ++++++++++++++++++-- crates/aionui-auth/src/service.rs | 4 ++- crates/aionui-auth/tests/middleware_tests.rs | 26 ++++++++++++++++ 5 files changed, 88 insertions(+), 4 deletions(-) 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/middleware.rs b/crates/aionui-auth/src/middleware.rs index 9ee0adb26..0a0cbb5d1 100644 --- a/crates/aionui-auth/src/middleware.rs +++ b/crates/aionui-auth/src/middleware.rs @@ -81,6 +81,10 @@ pub async fn auth_middleware( })? .ok_or_else(|| ApiError::Unauthorized("Invalid authentication subject".into()))?; + 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.unwrap_or_else(|| "external_user".to_string()), diff --git a/crates/aionui-auth/src/routes.rs b/crates/aionui-auth/src/routes.rs index b0fc7137f..71773f500 100644 --- a/crates/aionui-auth/src/routes.rs +++ b/crates/aionui-auth/src/routes.rs @@ -400,7 +400,11 @@ async fn login_handler( let token = state .jwt_service - .sign(&user.id, user.username.as_deref().unwrap_or("external_user")) + .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) @@ -685,9 +689,27 @@ 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 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 { @@ -749,7 +771,11 @@ async fn qr_login_handler( let token = state .jwt_service - .sign(&user.id, user.username.as_deref().unwrap_or("external_user")) + .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) diff --git a/crates/aionui-auth/src/service.rs b/crates/aionui-auth/src/service.rs index e49632472..9569ad1ce 100644 --- a/crates/aionui-auth/src/service.rs +++ b/crates/aionui-auth/src/service.rs @@ -75,7 +75,9 @@ impl AuthProvisionService { } let username = user.username.clone().unwrap_or_else(|| "external_user".to_string()); - let token = self.jwt_service.sign(&user.id, &username)?; + 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(EnsureExternalSessionResponse { diff --git a/crates/aionui-auth/tests/middleware_tests.rs b/crates/aionui-auth/tests/middleware_tests.rs index 0d87f3df0..348bf5441 100644 --- a/crates/aionui-auth/tests/middleware_tests.rs +++ b/crates/aionui-auth/tests/middleware_tests.rs @@ -239,6 +239,32 @@ 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_database_error_returns_internal_error_code() { let jwt_service = Arc::new(JwtService::new("middleware_test_secret".into())); From ec8b8f3817dcfd71f0f018c327177944e7b1213c Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 20:06:08 +0800 Subject: [PATCH 041/145] channel: update integration tests for owner scope --- .../tests/dingtalk_integration.rs | 18 ++++++++++++------ .../aionui-channel/tests/lark_integration.rs | 16 ++++++++++------ .../tests/telegram_integration.rs | 18 ++++++++++++------ .../aionui-channel/tests/weixin_integration.rs | 14 +++++++++----- 4 files changed, 43 insertions(+), 23 deletions(-) diff --git a/crates/aionui-channel/tests/dingtalk_integration.rs b/crates/aionui-channel/tests/dingtalk_integration.rs index bf26d96fd..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("system_default_user", "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 75f861565..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("system_default_user", "lark")); + assert!(!manager.is_plugin_running(OWNER_USER_ID, "lark")); } } diff --git a/crates/aionui-channel/tests/telegram_integration.rs b/crates/aionui-channel/tests/telegram_integration.rs index 34fef1232..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("system_default_user", "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 955b6dcef..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("system_default_user", "weixin")); + assert!(!manager.is_plugin_running(OWNER_USER_ID, "weixin")); } // -- Login event serialization ------------------------------------------ From a7d985b1458edad9e013611726758d6cc0c3f3b3 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 20:16:51 +0800 Subject: [PATCH 042/145] agent: project user scoped management availability --- crates/aionui-ai-agent/src/registry.rs | 16 ++- crates/aionui-app/tests/acp_e2e.rs | 181 ++++++++++++++----------- 2 files changed, 114 insertions(+), 83 deletions(-) diff --git a/crates/aionui-ai-agent/src/registry.rs b/crates/aionui-ai-agent/src/registry.rs index 97a5aad52..e9b45403b 100644 --- a/crates/aionui-ai-agent/src/registry.rs +++ b/crates/aionui-ai-agent/src/registry.rs @@ -374,10 +374,14 @@ impl AgentRegistry { let mut rows: Vec = rows .into_iter() .filter_map(|row| decode_row(row, AvailabilityProjection::Cached)) - .map(|(meta, reason)| { + .map(|(mut meta, mut reason)| { let cached_meta = cached_rows.get(&meta.id); let cached_reason = cached_reasons.get(&meta.id); - let (meta, reason) = overlay_hydrated_availability(meta, reason, cached_meta, cached_reason); + 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(); @@ -414,12 +418,16 @@ impl AgentRegistry { .get_for_user(user_id, id) .await .map_err(|e| AgentError::internal(format!("load agent_metadata '{id}' for user: {e}")))?; - let Some((meta, reason)) = row.and_then(|row| decode_row(row, AvailabilityProjection::Cached)) else { + let Some((mut meta, mut reason)) = row.and_then(|row| decode_row(row, AvailabilityProjection::Cached)) else { return Ok(None); }; let cached_meta = self.by_id.read().await.get(&meta.id).cloned(); let cached_reason = self.unavailable_reasons.read().await.get(&meta.id).cloned(); - let (meta, reason) = overlay_hydrated_availability(meta, reason, cached_meta.as_ref(), cached_reason.as_ref()); + 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()))) } 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(); From 8c68df914ddc99766861542aee60bde0c980f74d Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 20:22:42 +0800 Subject: [PATCH 043/145] test: align team active lease cross-user response --- crates/aionui-app/tests/active_lease_e2e.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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()); } From f8f3093718ea33309b466a58afa61c105b446667 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 20:47:41 +0800 Subject: [PATCH 044/145] cron: scope jobs by user --- crates/aionui-app/tests/cron_e2e.rs | 99 ++++++++++++++----- crates/aionui-cron/src/artifacts.rs | 1 + crates/aionui-cron/src/executor.rs | 5 + crates/aionui-cron/src/scheduler.rs | 1 + crates/aionui-cron/src/service.rs | 17 ++-- crates/aionui-cron/src/types.rs | 5 + .../aionui-cron/tests/service_integration.rs | 1 + .../aionui-db/migrations/026_user_scope.sql | 48 +-------- crates/aionui-db/src/models/cron_job.rs | 3 + .../aionui-db/src/repository/sqlite_cron.rs | 84 ++++++---------- crates/aionui-db/tests/cron_repository.rs | 1 + 11 files changed, 138 insertions(+), 127 deletions(-) diff --git a/crates/aionui-app/tests/cron_e2e.rs b/crates/aionui-app/tests/cron_e2e.rs index a9026fcd4..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 @@ -837,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"}); @@ -859,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"}); @@ -887,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); @@ -905,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": ""}); @@ -927,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"}); @@ -968,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( @@ -1035,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-cron/src/artifacts.rs b/crates/aionui-cron/src/artifacts.rs index 11886d196..e688cf1d7 100644 --- a/crates/aionui-cron/src/artifacts.rs +++ b/crates/aionui-cron/src/artifacts.rs @@ -103,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/executor.rs b/crates/aionui-cron/src/executor.rs index 5da16ab37..ac8c4a7e1 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -611,6 +611,10 @@ 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", @@ -1282,6 +1286,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/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 31dce391f..e49c49739 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -265,13 +265,14 @@ impl CronService { 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 conversation_id.is_empty() - || self - .executor - .get_conversation_row(conversation_id) - .await? - .filter(|row| row.user_id == user_id) - .is_none() + if matches!(execution_mode, ExecutionMode::Existing) + && (conversation_id.is_empty() + || self + .executor + .get_conversation_row(conversation_id) + .await? + .filter(|row| row.user_id == user_id) + .is_none()) { return Err(CronError::Conversation( aionui_conversation::ConversationError::NotFound { @@ -297,6 +298,7 @@ impl CronService { let job = CronJob { id: generate_prefixed_id("cron"), + user_id: user_id.to_owned(), name: req.name, enabled: true, schedule, @@ -2314,6 +2316,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/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 907cab1c0..123e63019 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -1445,6 +1445,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: "user1".into(), name: "Legacy custom agent job".into(), enabled: true, schedule_kind: "every".into(), diff --git a/crates/aionui-db/migrations/026_user_scope.sql b/crates/aionui-db/migrations/026_user_scope.sql index b0a29e6a3..838541369 100644 --- a/crates/aionui-db/migrations/026_user_scope.sql +++ b/crates/aionui-db/migrations/026_user_scope.sql @@ -121,51 +121,9 @@ BEGIN SELECT RAISE(ABORT, 'team_tasks.team_id must reference teams.id'); END; -CREATE TRIGGER IF NOT EXISTS trg_cron_jobs_conversation_parent_insert -BEFORE INSERT ON cron_jobs -FOR EACH ROW -WHEN NEW.conversation_id IS NULL - OR NEW.conversation_id = '' - OR NOT EXISTS (SELECT 1 FROM conversations WHERE id = NEW.conversation_id) -BEGIN - SELECT RAISE(ABORT, 'cron_jobs.conversation_id must reference conversations.id'); -END; - -CREATE TRIGGER IF NOT EXISTS trg_cron_jobs_conversation_parent_update -BEFORE UPDATE OF conversation_id ON cron_jobs -FOR EACH ROW -WHEN NEW.conversation_id IS NULL - OR NEW.conversation_id = '' - OR NOT EXISTS (SELECT 1 FROM conversations WHERE id = NEW.conversation_id) -BEGIN - SELECT RAISE(ABORT, 'cron_jobs.conversation_id must reference conversations.id'); -END; - -CREATE TRIGGER IF NOT EXISTS trg_cron_job_runs_job_parent_insert -BEFORE INSERT ON cron_job_runs -FOR EACH ROW -WHEN NOT EXISTS ( - SELECT 1 - FROM cron_jobs j - JOIN conversations c ON c.id = j.conversation_id - WHERE j.id = NEW.job_id -) -BEGIN - SELECT RAISE(ABORT, 'cron_job_runs.job_id must reference cron_jobs.id with conversation parent'); -END; - -CREATE TRIGGER IF NOT EXISTS trg_cron_job_runs_job_parent_update -BEFORE UPDATE OF job_id ON cron_job_runs -FOR EACH ROW -WHEN NOT EXISTS ( - SELECT 1 - FROM cron_jobs j - JOIN conversations c ON c.id = j.conversation_id - WHERE j.id = NEW.job_id -) -BEGIN - SELECT RAISE(ABORT, 'cron_job_runs.job_id must reference cron_jobs.id with conversation parent'); -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'; 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/repository/sqlite_cron.rs b/crates/aionui-db/src/repository/sqlite_cron.rs index e8e2ff161..e9df2747f 100644 --- a/crates/aionui-db/src/repository/sqlite_cron.rs +++ b/crates/aionui-db/src/repository/sqlite_cron.rs @@ -24,16 +24,17 @@ impl ICronRepository for SqliteCronRepository { async fn insert(&self, row: &CronJobRow) -> Result<(), DbError> { 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) @@ -86,10 +87,7 @@ impl ICronRepository for SqliteCronRepository { let result = sqlx::query( "DELETE FROM cron_jobs \ WHERE id = ? \ - AND EXISTS (\ - SELECT 1 FROM conversations c \ - WHERE c.id = cron_jobs.conversation_id AND c.user_id = ?\ - )", + AND user_id = ?", ) .bind(id) .bind(user_id) @@ -110,15 +108,11 @@ impl ICronRepository for SqliteCronRepository { } async fn get_by_id_for_user(&self, user_id: &str, id: &str) -> Result, DbError> { - let row = sqlx::query_as::<_, CronJobRow>( - "SELECT j.* FROM cron_jobs j \ - JOIN conversations c ON c.id = j.conversation_id \ - WHERE c.user_id = ? AND j.id = ?", - ) - .bind(user_id) - .bind(id) - .fetch_optional(&self.pool) - .await?; + 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) } @@ -130,27 +124,17 @@ impl ICronRepository for SqliteCronRepository { } async fn list_all_for_user(&self, user_id: &str) -> Result, DbError> { - let rows = sqlx::query_as::<_, CronJobRow>( - "SELECT j.* FROM cron_jobs j \ - JOIN conversations c ON c.id = j.conversation_id \ - WHERE c.user_id = ? \ - ORDER BY j.created_at ASC", - ) - .bind(user_id) - .fetch_all(&self.pool) - .await?; + 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(&self) -> Result, DbError> { - let rows = sqlx::query_as::<_, CronJobRow>( - "SELECT j.* FROM cron_jobs j \ - JOIN conversations c ON c.id = j.conversation_id \ - WHERE j.enabled = 1 \ - ORDER BY j.created_at ASC", - ) - .fetch_all(&self.pool) - .await?; + 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) } @@ -170,10 +154,9 @@ impl ICronRepository for SqliteCronRepository { conversation_id: &str, ) -> Result, DbError> { let rows = sqlx::query_as::<_, CronJobRow>( - "SELECT j.* FROM cron_jobs j \ - JOIN conversations c ON c.id = j.conversation_id \ - WHERE c.user_id = ? AND j.conversation_id = ? \ - ORDER BY j.created_at ASC", + "SELECT * FROM cron_jobs \ + WHERE user_id = ? AND conversation_id = ? \ + ORDER BY created_at ASC", ) .bind(user_id) .bind(conversation_id) @@ -197,9 +180,8 @@ impl ICronRepository for SqliteCronRepository { let result = async { let job_has_owner: bool = sqlx::query_scalar( "SELECT EXISTS(\ - SELECT 1 FROM cron_jobs j \ - JOIN conversations c ON c.id = j.conversation_id \ - WHERE j.id = ?\ + SELECT 1 FROM cron_jobs \ + WHERE id = ? AND user_id IS NOT NULL AND user_id != ''\ )", ) .bind(params.job_id) @@ -315,8 +297,7 @@ impl ICronRepository for SqliteCronRepository { WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ? \ AND EXISTS (\ SELECT 1 FROM cron_jobs j \ - JOIN conversations c ON c.id = j.conversation_id \ - WHERE j.id = cron_job_runs.job_id\ + WHERE j.id = cron_job_runs.job_id AND j.user_id IS NOT NULL AND j.user_id != ''\ )", ) .bind(lease_until) @@ -342,8 +323,7 @@ impl ICronRepository for SqliteCronRepository { WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ? \ AND EXISTS (\ SELECT 1 FROM cron_jobs j \ - JOIN conversations c ON c.id = j.conversation_id \ - WHERE j.id = cron_job_runs.job_id\ + WHERE j.id = cron_job_runs.job_id AND j.user_id IS NOT NULL AND j.user_id != ''\ )", ) .bind(retry_at) @@ -363,8 +343,7 @@ impl ICronRepository for SqliteCronRepository { WHERE job_id = ? AND scheduled_at = ? AND status = 'running' AND owner_id = ? \ AND EXISTS (\ SELECT 1 FROM cron_jobs j \ - JOIN conversations c ON c.id = j.conversation_id \ - WHERE j.id = cron_job_runs.job_id\ + WHERE j.id = cron_job_runs.job_id AND j.user_id IS NOT NULL AND j.user_id != ''\ )", ) .bind(params.status) @@ -393,8 +372,8 @@ impl ICronRepository for SqliteCronRepository { let row = sqlx::query_as::<_, (TimestampMs, TimestampMs)>( "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 \ - JOIN conversations c ON c.id = j.conversation_id \ - WHERE r.job_id = ? AND r.status IN ('running', 'retrying') AND r.lease_until IS NOT NULL \ + 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) @@ -482,7 +461,9 @@ impl SqliteCronRepository { 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()) { + if let (Some(user_id), Some(conversation_id)) = (user_id, params.conversation_id.as_deref()) + && !conversation_id.trim().is_empty() + { let owns_new_conversation: bool = sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM conversations WHERE user_id = ? AND id = ?)") .bind(user_id) @@ -496,11 +477,7 @@ impl SqliteCronRepository { let sql = if user_id.is_some() { format!( - "UPDATE cron_jobs SET {} WHERE id = ? \ - AND EXISTS (\ - SELECT 1 FROM conversations c \ - WHERE c.id = cron_jobs.conversation_id AND c.user_id = ?\ - )", + "UPDATE cron_jobs SET {} WHERE id = ? AND user_id = ?", set_parts.join(", ") ) } else { @@ -580,6 +557,7 @@ mod tests { let now = now_ms(); CronJobRow { id: id.into(), + user_id: "user1".into(), name: "Test Job".into(), enabled: true, schedule_kind: "every".into(), diff --git a/crates/aionui-db/tests/cron_repository.rs b/crates/aionui-db/tests/cron_repository.rs index bbf92992f..b18a1ea3b 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: "user1".into(), name: "Test Job".into(), enabled: true, schedule_kind: "every".into(), From 2bf1396ccb23e66943781535c50e0a0b1ef71be3 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 20:49:57 +0800 Subject: [PATCH 045/145] test: cover user scope migration shape --- .../aionui-db/tests/user_scope_migration.rs | 90 +++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 crates/aionui-db/tests/user_scope_migration.rs 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..36eff9d9b --- /dev/null +++ b/crates/aionui-db/tests/user_scope_migration.rs @@ -0,0 +1,90 @@ +use aionui_db::init_database_memory; +use sqlx::Row; + +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() +} + +#[tokio::test] +async fn migration_026_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_026_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_026_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"), ""); +} From ceb869a7371862aad8bac8d92f21519b400d506e Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 20:56:45 +0800 Subject: [PATCH 046/145] cron: resolve assistant metadata by job user --- crates/aionui-cron/src/service.rs | 124 +++++++++++++++++++++--------- 1 file changed, 89 insertions(+), 35 deletions(-) diff --git a/crates/aionui-cron/src/service.rs b/crates/aionui-cron/src/service.rs index e49c49739..d0d6f5bbf 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -257,7 +257,10 @@ 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())?; @@ -284,6 +287,7 @@ impl CronService { 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(), @@ -392,9 +396,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 @@ -812,6 +819,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 { @@ -820,7 +828,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 { @@ -839,10 +847,15 @@ 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"); @@ -896,18 +909,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()) } @@ -1572,7 +1592,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 ( @@ -1580,7 +1600,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(); @@ -1589,7 +1609,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!( @@ -1604,7 +1624,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!( @@ -1621,7 +1644,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)) }; @@ -1652,6 +1675,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()), @@ -1708,6 +1732,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>, @@ -1719,7 +1744,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!( @@ -1728,6 +1753,7 @@ impl CronService { })?; let full_auto_mode = self .resolve_cron_full_auto_mode( + user_id, runtime_agent_type, Some(assistant_id), None, @@ -1749,41 +1775,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() @@ -1800,9 +1843,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() @@ -1819,8 +1862,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())) @@ -1828,20 +1871,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)); } @@ -1850,33 +1894,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); }; From 1e74e1022b98c9a54350e238cf9350fb23c671d4 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 20:58:13 +0800 Subject: [PATCH 047/145] channel: resolve assistant settings by owner --- crates/aionui-channel/src/channel_settings.rs | 55 ++++++++++++------- 1 file changed, 35 insertions(+), 20 deletions(-) diff --git a/crates/aionui-channel/src/channel_settings.rs b/crates/aionui-channel/src/channel_settings.rs index d5ede8d7a..b278a9273 100644 --- a/crates/aionui-channel/src/channel_settings.rs +++ b/crates/aionui-channel/src/channel_settings.rs @@ -92,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, @@ -186,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); @@ -194,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 { @@ -213,11 +216,14 @@ impl ChannelSettingsService { 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 }; @@ -256,6 +262,7 @@ impl ChannelSettingsService { async fn resolve_assistant_agent_config( &self, + user_id: &str, assistant_id: &str, ) -> Result, ChannelError> { let (Some(definition_repo), Some(overlay_repo)) = @@ -264,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 }; @@ -282,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)) = @@ -300,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)); } @@ -317,6 +328,7 @@ impl ChannelSettingsService { async fn normalize_channel_assistant_setting_for_response( &self, + user_id: &str, assistant: ChannelAssistantSettingResponse, ) -> Result { let assistant_id = assistant @@ -335,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() { @@ -353,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); }; @@ -367,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; } @@ -399,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 { @@ -407,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())) From decbc73cf7c24fb46273a8d85131b2e884d56d88 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 21:05:58 +0800 Subject: [PATCH 048/145] team: resolve assistant catalog by owner --- crates/aionui-app/src/router/state.rs | 3 +- crates/aionui-team/src/mcp/server.rs | 30 ++++++++--- crates/aionui-team/src/ports.rs | 6 ++- crates/aionui-team/src/provisioning.rs | 53 ++++++++++++------- crates/aionui-team/src/service.rs | 17 ++++-- .../src/service/describe_support.rs | 10 ++-- .../src/service/response_builder.rs | 18 ++++--- .../aionui-team/src/service/spawn_support.rs | 26 +++++---- crates/aionui-team/src/session.rs | 2 +- crates/aionui-team/src/test_utils.rs | 5 +- .../tests/session_service_integration.rs | 19 +++++-- 11 files changed, 131 insertions(+), 58 deletions(-) diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index 2018ede46..f70d35b4c 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -629,8 +629,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}")) })?; 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/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 e1da9f1c1..d682f99f8 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(|| { @@ -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 { @@ -857,6 +870,7 @@ mod tests { impl TeamAssistantCatalogPort for EmptyTeamAssistantCatalog { async fn list_team_selectable_assistants( &self, + _user_id: &str, ) -> Result, TeamError> { Ok(Vec::new()) } @@ -987,7 +1001,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); } @@ -998,7 +1012,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); } @@ -1009,7 +1023,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/service.rs b/crates/aionui-team/src/service.rs index b3c41446b..cda3d239c 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -299,6 +299,15 @@ impl TeamSessionService { Ok(Team::from_row(&row)?) } + pub(crate) async fn team_owner_user_id(&self, team_id: &str) -> Result { + let row = self + .repo + .get_team_for_restore(team_id) + .await? + .ok_or_else(|| TeamError::TeamNotFound(team_id.into()))?; + Ok(row.user_id) + } + pub async fn renew_active_lease( &self, user_id: &str, @@ -434,7 +443,7 @@ impl TeamSessionService { 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> { @@ -442,7 +451,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"); @@ -458,7 +467,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> { @@ -576,7 +585,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> { 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 8a92eac73..e72a629ac 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}")) @@ -81,10 +84,10 @@ impl TeamSessionService { 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; @@ -120,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(); }; @@ -154,7 +157,11 @@ impl TeamSessionService { if backend == "aionrs" { 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() @@ -533,6 +540,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()) } @@ -676,7 +684,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()) diff --git a/crates/aionui-team/src/session.rs b/crates/aionui-team/src/session.rs index f73c2dd0a..534702fe6 100644 --- a/crates/aionui-team/src/session.rs +++ b/crates/aionui-team/src/session.rs @@ -304,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 { diff --git a/crates/aionui-team/src/test_utils.rs b/crates/aionui-team/src/test_utils.rs index fdf44fafb..07508b3c1 100644 --- a/crates/aionui-team/src/test_utils.rs +++ b/crates/aionui-team/src/test_utils.rs @@ -891,7 +891,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()) } } diff --git a/crates/aionui-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index 2f9bd772a..5848aaaa0 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -1302,7 +1302,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()) } } @@ -1315,13 +1318,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; } From 12b59b3e1b8bdb3e39c7260526cd1d82664c6594 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 21:10:01 +0800 Subject: [PATCH 049/145] conversation: resolve runtime catalog by owner --- crates/aionui-conversation/src/service.rs | 54 ++++++++++++------- .../src/session_context.rs | 4 +- 2 files changed, 38 insertions(+), 20 deletions(-) diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index 74b4b6999..ac4aa1d8b 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -640,7 +640,7 @@ impl ConversationService { .get_assistant_snapshot(user_id, &response.id) .await? { - response.assistant = Some(self.assistant_identity_from_snapshot(&snapshot).await?); + response.assistant = Some(self.assistant_identity_from_snapshot(user_id, &snapshot).await?); } Ok(()) @@ -648,14 +648,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, @@ -686,13 +689,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}"))) } @@ -732,7 +736,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, @@ -968,6 +972,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() @@ -1026,7 +1031,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(); @@ -1175,7 +1182,7 @@ 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?; } @@ -1209,6 +1216,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>, @@ -1251,7 +1259,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) @@ -1298,11 +1306,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)) @@ -1310,6 +1319,7 @@ impl ConversationService { async fn resolve_assistant_snapshot( &self, + user_id: &str, assistant_id: &str, locale: Option<&str>, overrides: &AssistantConversationOverrides, @@ -1324,7 +1334,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 { @@ -1332,11 +1342,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}")))?; @@ -1424,7 +1434,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"), @@ -1658,7 +1668,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 { @@ -3368,6 +3378,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(), ) @@ -3760,6 +3771,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> { @@ -3767,7 +3779,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(); } @@ -3779,11 +3795,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), } @@ -3792,6 +3809,7 @@ impl ConversationService { async fn resolve_acp_mcp_support_policy( repo: &Arc, + user_id: &str, extra: &serde_json::Value, ) -> Result { let agent_id = extra @@ -3809,12 +3827,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, diff --git a/crates/aionui-conversation/src/session_context.rs b/crates/aionui-conversation/src/session_context.rs index d11e7a23a..35dcdcf63 100644 --- a/crates/aionui-conversation/src/session_context.rs +++ b/crates/aionui-conversation/src/session_context.rs @@ -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 { From a8a73c74752d6f1042a9bac0d5f8f49598d7fa72 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 21:14:17 +0800 Subject: [PATCH 050/145] cron: resolve executor registry by job user --- crates/aionui-ai-agent/src/registry.rs | 11 +++++++ crates/aionui-cron/src/executor.rs | 43 ++++++++++++++++++++------ 2 files changed, 44 insertions(+), 10 deletions(-) diff --git a/crates/aionui-ai-agent/src/registry.rs b/crates/aionui-ai-agent/src/registry.rs index e9b45403b..a890cf291 100644 --- a/crates/aionui-ai-agent/src/registry.rs +++ b/crates/aionui-ai-agent/src/registry.rs @@ -301,6 +301,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. diff --git a/crates/aionui-cron/src/executor.rs b/crates/aionui-cron/src/executor.rs index ac8c4a7e1..09aab2f4f 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -559,7 +559,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?; @@ -987,7 +987,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())); @@ -995,7 +999,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); } @@ -1030,7 +1038,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)); @@ -1555,15 +1566,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() @@ -1576,15 +1593,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 ); } From b84849b7000d26d2af4a0b79403750ddca290533 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 21:17:19 +0800 Subject: [PATCH 051/145] agent: resolve acp runtime metadata by user --- crates/aionui-ai-agent/src/factory/acp.rs | 4 +++- crates/aionui-ai-agent/src/protocol/cli_detect.rs | 4 ++-- crates/aionui-ai-agent/src/services/availability/mod.rs | 2 +- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/crates/aionui-ai-agent/src/factory/acp.rs b/crates/aionui-ai-agent/src/factory/acp.rs index ab99916c6..3ae0716f6 100644 --- a/crates/aionui-ai-agent/src/factory/acp.rs +++ b/crates/aionui-ai-agent/src/factory/acp.rs @@ -34,7 +34,9 @@ pub(super) async fn build( 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 + deps.agent_registry + .find_builtin_by_backend_for_user(&ctx.user_id, vendor) + .await } else { None } 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/services/availability/mod.rs b/crates/aionui-ai-agent/src/services/availability/mod.rs index 1df808a0e..40b76a406 100644 --- a/crates/aionui-ai-agent/src/services/availability/mod.rs +++ b/crates/aionui-ai-agent/src/services/availability/mod.rs @@ -240,7 +240,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 { From 598e07bf919d7521c93948a4b47dd77be6179543 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 21:22:19 +0800 Subject: [PATCH 052/145] assistant: project generated agents by user --- crates/aionui-app/src/router/state.rs | 10 ++++++++-- crates/aionui-app/tests/assistants_e2e.rs | 5 ++++- crates/aionui-assistant/src/agent_catalog.rs | 2 +- crates/aionui-assistant/src/service.rs | 7 +++++-- 4 files changed, 18 insertions(+), 6 deletions(-) diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index f70d35b4c..e462fa186 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -329,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}"))) } } diff --git a/crates/aionui-app/tests/assistants_e2e.rs b/crates/aionui-app/tests/assistants_e2e.rs index 7d02b1e6e..67fe09db3 100644 --- a/crates/aionui-app/tests/assistants_e2e.rs +++ b/crates/aionui-app/tests/assistants_e2e.rs @@ -65,7 +65,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()) } } 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/service.rs b/crates/aionui-assistant/src/service.rs index cad6de2b1..bacefea32 100644 --- a/crates/aionui-assistant/src/service.rs +++ b/crates/aionui-assistant/src/service.rs @@ -469,7 +469,7 @@ impl AssistantService { return Ok(Vec::new()); }; - let rows = agent_catalog.list_management_agents().await?; + 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}")) })?; @@ -3346,7 +3346,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()) } } From fb343f13f51f6126437a5ad01cf978f10a1c8f89 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 21:43:34 +0800 Subject: [PATCH 053/145] extension: scope skill catalog sync by user --- crates/aionui-app/tests/common/mod.rs | 13 ++++++-- crates/aionui-app/tests/extension_e2e.rs | 6 ++-- crates/aionui-db/src/models/skill.rs | 1 + crates/aionui-extension/src/lib.rs | 2 +- crates/aionui-extension/src/skill_service.rs | 32 ++++++++++++++++++-- 5 files changed, 45 insertions(+), 9 deletions(-) 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/extension_e2e.rs b/crates/aionui-app/tests/extension_e2e.rs index ae46be1c7..c620ab23a 100644 --- a/crates/aionui-app/tests/extension_e2e.rs +++ b/crates/aionui-app/tests/extension_e2e.rs @@ -1356,7 +1356,7 @@ async fn sl1_list_skills_tags_builtin_and_custom_with_source_field() { 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; let resp = app.oneshot(get_with_token("/api/skills", &token)).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); @@ -1395,7 +1395,7 @@ async fn sl2_list_skills_user_custom_overrides_builtin() { 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; let resp = app.oneshot(get_with_token("/api/skills", &token)).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); @@ -1437,7 +1437,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-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-extension/src/lib.rs b/crates/aionui-extension/src/lib.rs index 74a6b1fa0..5aa63b043 100644 --- a/crates/aionui-extension/src/lib.rs +++ b/crates/aionui-extension/src/lib.rs @@ -57,7 +57,7 @@ pub use skill_service::{ 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, + 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/skill_service.rs b/crates/aionui-extension/src/skill_service.rs index b36a46118..2f0f4d400 100644 --- a/crates/aionui-extension/src/skill_service.rs +++ b/crates/aionui-extension/src/skill_service.rs @@ -1622,13 +1622,29 @@ async fn list_skills_from_repo( user_id: &str, ) -> Result, ExtensionError> { let mut items = Vec::new(); - for row in repo.list_for_user(user_id).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 @@ -1638,7 +1654,19 @@ pub async fn sync_skill_catalog_into_repo( paths: &SkillPaths, repo: &dyn ISkillRepository, ) -> Result<(), ExtensionError> { - sync_disk_user_skills_into_repo_for_user(paths, repo, DEFAULT_USER_ID).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(()) } From 68e96b61857bf735a50fc52afad7b4ed1d89f157 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 22:10:23 +0800 Subject: [PATCH 054/145] team: distinguish cross-user access from missing teams --- crates/aionui-team/src/service.rs | 29 +++++++++++++---------------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index cda3d239c..4b1d71ac3 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -291,12 +291,20 @@ 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 { let row = self .repo - .get_team(user_id, team_id) + .get_team_for_restore(team_id) .await? .ok_or_else(|| TeamError::TeamNotFound(team_id.into()))?; - Ok(Team::from_row(&row)?) + if row.user_id != user_id { + return Err(TeamError::Forbidden("team belongs to another user".into())); + } + Ok(row) } pub(crate) async fn team_owner_user_id(&self, team_id: &str) -> Result { @@ -538,11 +546,7 @@ impl TeamSessionService { .clone(); let _guard = lock.lock().await; - let row = self - .repo - .get_team(user_id, team_id) - .await? - .ok_or_else(|| TeamError::TeamNotFound(team_id.into()))?; + 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?; @@ -815,10 +819,7 @@ 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> { - self.repo - .get_team(user_id, team_id) - .await? - .ok_or_else(|| TeamError::TeamNotFound(team_id.to_owned()))?; + self.load_owned_team_row(user_id, team_id).await?; self.ensure_session_inner(team_id, Some(user_id)).await } @@ -1232,11 +1233,7 @@ impl TeamSessionService { team_id: &str, conversation_id: &str, ) -> Result { - let row = self - .repo - .get_team(user_id, team_id) - .await? - .ok_or_else(|| TeamError::TeamNotFound(team_id.to_owned()))?; + 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); From b106c77411134a3ddc93234247649bd8489072da Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 22:16:21 +0800 Subject: [PATCH 055/145] team: preserve active lease cross-user hiding --- crates/aionui-team/src/service.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 4b1d71ac3..012821022 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -322,9 +322,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, From 919e2afb7d0d8d3bd5e021c1fbb89c6cc7fd7e52 Mon Sep 17 00:00:00 2001 From: zk <> Date: Mon, 20 Jul 2026 22:45:03 +0800 Subject: [PATCH 056/145] conversation: scope assistant preferences by user --- Cargo.lock | 1 + crates/aionui-conversation/Cargo.toml | 1 + crates/aionui-conversation/src/service.rs | 68 +- .../aionui-conversation/src/service_test.rs | 606 +++++++++++------- 4 files changed, 406 insertions(+), 270 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 49df9f92e..dd8c563b6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -578,6 +578,7 @@ dependencies = [ "regex", "serde", "serde_json", + "sqlx", "tempfile", "thiserror 2.0.18", "tokio", 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/service.rs b/crates/aionui-conversation/src/service.rs index ac4aa1d8b..195be2d20 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -1187,7 +1187,8 @@ impl ConversationService { } 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)?; @@ -1480,6 +1481,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 { @@ -1487,7 +1489,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" { @@ -1538,15 +1540,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}")))?; @@ -1706,7 +1711,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}")))?; @@ -1742,24 +1747,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}")))?; diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index 8c2bb8d5c..6b25e3f2d 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -1155,6 +1155,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 = @@ -1180,6 +1182,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 = @@ -1209,6 +1213,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 = @@ -1230,6 +1236,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!({ @@ -3883,15 +3905,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(); @@ -3913,7 +3938,11 @@ 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("user_1", &conv.id).await.unwrap().unwrap(); @@ -3929,7 +3958,11 @@ 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("user_1", &conv.id).await.unwrap().unwrap(); @@ -3947,7 +3980,11 @@ 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")); @@ -3984,15 +4021,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(); @@ -4020,7 +4060,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"), @@ -4055,15 +4099,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(); @@ -4102,7 +4149,11 @@ 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); @@ -4140,15 +4191,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(); @@ -4172,7 +4226,11 @@ 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("user_1", &conv.id).await.unwrap().unwrap(); @@ -4205,15 +4263,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(); @@ -4242,7 +4303,11 @@ 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("user_1", &auto_conv.id) @@ -4271,15 +4336,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(); @@ -4304,7 +4372,11 @@ 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("user_1", &fixed_conv.id) @@ -6365,25 +6437,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(); @@ -6439,7 +6517,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"]"#); @@ -6492,46 +6574,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(); @@ -6548,36 +6636,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(); @@ -6619,46 +6710,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(); @@ -6810,25 +6907,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(); @@ -6883,13 +6986,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(); @@ -6971,25 +7077,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(); @@ -7015,7 +7127,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"]"#); @@ -7069,25 +7185,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(); @@ -7110,7 +7232,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, "[]"); From 319be2f943240e43dc154e45c0f8bb1d125f2513 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 10:08:56 +0800 Subject: [PATCH 057/145] cron: seed assistant fixtures for scoped user --- .../aionui-cron/tests/service_integration.rs | 80 ++++++++++--------- 1 file changed, 43 insertions(+), 37 deletions(-) diff --git a/crates/aionui-cron/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index 123e63019..6cb749177 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -1071,36 +1071,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(); } @@ -1121,13 +1124,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(); } From 2f33aed937b07219e787197883fb135f3a4d010d Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 10:17:32 +0800 Subject: [PATCH 058/145] cron: align owner anchor tests with user scope --- crates/aionui-cron/tests/service_integration.rs | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/crates/aionui-cron/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index 6cb749177..805ac9ec2 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -1451,7 +1451,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: "user1".into(), + user_id: "u1".into(), name: "Legacy custom agent job".into(), enabled: true, schedule_kind: "every".into(), @@ -2277,20 +2277,22 @@ async fn oc1_rejects_lazy_existing_jobs() { } #[tokio::test] -async fn oc1b_rejects_new_conversation_jobs_without_owner_anchor() { +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_err = svc.add_job("u1", empty_req).await.unwrap_err(); - assert!(matches!(empty_err, aionui_cron::error::CronError::Conversation(_))); + 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 = "missing-conv-that-no-longer-exists".into(); stale_req.execution_mode = Some("new_conversation".into()); - let stale_err = svc.add_job("u1", stale_req).await.unwrap_err(); - assert!(matches!(stale_err, aionui_cron::error::CronError::Conversation(_))); + 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] From eb908a9bc48531a4640de6deb75d0ec7a275acfb Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 10:23:11 +0800 Subject: [PATCH 059/145] db: align cron repository tests with user scope --- crates/aionui-db/tests/cron_repository.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/crates/aionui-db/tests/cron_repository.rs b/crates/aionui-db/tests/cron_repository.rs index b18a1ea3b..72635ecf6 100644 --- a/crates/aionui-db/tests/cron_repository.rs +++ b/crates/aionui-db/tests/cron_repository.rs @@ -44,7 +44,7 @@ fn make_job(id: &str) -> CronJobRow { let now = now_ms(); CronJobRow { id: id.into(), - user_id: "user1".into(), + user_id: "user_1".into(), name: "Test Job".into(), enabled: true, schedule_kind: "every".into(), @@ -201,12 +201,13 @@ async fn cj7_list_by_conversation() { } #[tokio::test] -async fn scoped_crud_filters_by_conversation_owner() { +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(); @@ -243,17 +244,17 @@ async fn scoped_crud_filters_by_conversation_owner() { } #[tokio::test] -async fn scheduler_enabled_scan_requires_conversation_owner() { +async fn scheduler_enabled_scan_uses_job_user_id() { let (r, _db) = repo().await; r.insert(&make_job("cron_owned")).await.unwrap(); let mut orphan = make_job("cron_orphan"); orphan.conversation_id = "missing_conversation".into(); - let orphan_insert = r.insert(&orphan).await; - assert!(orphan_insert.is_err()); + r.insert(&orphan).await.unwrap(); let enabled = r.list_enabled().await.unwrap(); assert!(enabled.iter().any(|job| job.id == "cron_owned")); + assert!(enabled.iter().any(|job| job.id == "cron_orphan")); } #[tokio::test] From 76bc1e99b412d9c06891a8acc70e2c13c7b7aa5d Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 10:28:23 +0800 Subject: [PATCH 060/145] team: hide cross-user team lookups --- crates/aionui-team/src/service.rs | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 012821022..141b17d2e 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -296,15 +296,11 @@ impl TeamSessionService { } async fn load_owned_team_row(&self, user_id: &str, team_id: &str) -> Result { - let row = self + self .repo - .get_team_for_restore(team_id) + .get_team(user_id, team_id) .await? - .ok_or_else(|| TeamError::TeamNotFound(team_id.into()))?; - if row.user_id != user_id { - return Err(TeamError::Forbidden("team belongs to another user".into())); - } - Ok(row) + .ok_or_else(|| TeamError::TeamNotFound(team_id.into())) } pub(crate) async fn team_owner_user_id(&self, team_id: &str) -> Result { From e2a322a2bc2600832acb9df89a7ecb7828a00c3f Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 10:55:57 +0800 Subject: [PATCH 061/145] app: expect hidden cross-user team access --- crates/aionui-app/tests/team_e2e.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/aionui-app/tests/team_e2e.rs b/crates/aionui-app/tests/team_e2e.rs index ab426daa1..36a2ac376 100644 --- a/crates/aionui-app/tests/team_e2e.rs +++ b/crates/aionui-app/tests/team_e2e.rs @@ -489,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"), @@ -536,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); } } @@ -658,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 From 952a4f129d6879e55765cc1ee889af46d2df246e Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 11:17:15 +0800 Subject: [PATCH 062/145] style: apply rustfmt to team service --- crates/aionui-team/src/service.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 141b17d2e..56baa58f1 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -296,8 +296,7 @@ impl TeamSessionService { } async fn load_owned_team_row(&self, user_id: &str, team_id: &str) -> Result { - self - .repo + self.repo .get_team(user_id, team_id) .await? .ok_or_else(|| TeamError::TeamNotFound(team_id.into())) From 1542c01f7664cd2ade1c645c8ff04d5d4312c4f7 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 11:21:16 +0800 Subject: [PATCH 063/145] assistant: reduce generated reconcile arguments --- crates/aionui-assistant/src/service.rs | 39 ++++++++++++++++---------- 1 file changed, 24 insertions(+), 15 deletions(-) diff --git a/crates/aionui-assistant/src/service.rs b/crates/aionui-assistant/src/service.rs index bacefea32..1a38ca30d 100644 --- a/crates/aionui-assistant/src/service.rs +++ b/crates/aionui-assistant/src/service.rs @@ -134,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. /// @@ -512,11 +520,13 @@ impl AssistantService { .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 { @@ -535,13 +545,10 @@ impl AssistantService { &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()) @@ -638,10 +645,12 @@ impl AssistantService { .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 }; From 1d27099dce47e2ed5d78834739196b1e23b2b6f0 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 11:26:21 +0800 Subject: [PATCH 064/145] auth: expose user context identity fields --- crates/aionui-auth/src/middleware.rs | 29 +++++++++++++------ crates/aionui-auth/tests/middleware_tests.rs | 4 ++- .../tests/feedback_diagnostics_routes.rs | 4 ++- .../aionui-system/tests/model_fetch_routes.rs | 4 ++- crates/aionui-system/tests/provider_routes.rs | 8 ++++- crates/aionui-system/tests/settings_routes.rs | 6 +++- 6 files changed, 41 insertions(+), 14 deletions(-) diff --git a/crates/aionui-auth/src/middleware.rs b/crates/aionui-auth/src/middleware.rs index 0a0cbb5d1..40a3ca318 100644 --- a/crates/aionui-auth/src/middleware.rs +++ b/crates/aionui-auth/src/middleware.rs @@ -7,7 +7,7 @@ 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; @@ -28,6 +28,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. @@ -56,10 +71,7 @@ pub async fn auth_middleware( ) -> Result { // In local mode, skip JWT verification and inject a fixed default user. if state.identity_mode == AuthIdentityMode::Local { - request.extensions_mut().insert(CurrentUser { - id: "system_default_user".to_string(), - username: "system_default_user".to_string(), - }); + request.extensions_mut().insert(CurrentUser::local_default()); return Ok(next.run(request).await); } @@ -88,6 +100,8 @@ pub async fn auth_middleware( request.extensions_mut().insert(CurrentUser { id: user.id, username: user.username.unwrap_or_else(|| "external_user".to_string()), + user_type: user.user_type, + status: user.status, }); Ok(next.run(request).await) @@ -98,10 +112,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/tests/middleware_tests.rs b/crates/aionui-auth/tests/middleware_tests.rs index 348bf5441..792485d59 100644 --- a/crates/aionui-auth/tests/middleware_tests.rs +++ b/crates/aionui-auth/tests/middleware_tests.rs @@ -14,7 +14,7 @@ use aionui_auth::{ 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(); @@ -436,6 +436,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-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 a3405fcdb..b99291058 100644 --- a/crates/aionui-system/tests/model_fetch_routes.rs +++ b/crates/aionui-system/tests/model_fetch_routes.rs @@ -17,7 +17,7 @@ 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::{ @@ -106,6 +106,8 @@ fn post_request(uri: &str, body: serde_json::Value) -> Request { 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 24b897a3e..dd0ded32b 100644 --- a/crates/aionui-system/tests/provider_routes.rs +++ b/crates/aionui-system/tests/provider_routes.rs @@ -16,7 +16,7 @@ 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, @@ -72,6 +72,8 @@ fn get_request(uri: &str) -> Request { 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 } @@ -86,6 +88,8 @@ fn json_request(method: &str, uri: &str, body: serde_json::Value) -> Request Request { 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/settings_routes.rs b/crates/aionui-system/tests/settings_routes.rs index a62a666da..51c17540d 100644 --- a/crates/aionui-system/tests/settings_routes.rs +++ b/crates/aionui-system/tests/settings_routes.rs @@ -15,7 +15,7 @@ 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, @@ -71,6 +71,8 @@ fn get_request(uri: &str) -> Request { 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 } @@ -85,6 +87,8 @@ fn json_request(method: &str, uri: &str, body: serde_json::Value) -> Request Date: Tue, 21 Jul 2026 11:32:49 +0800 Subject: [PATCH 065/145] cron: scope internal job updates by user --- crates/aionui-cron/src/service.rs | 45 ++++++++++++++++--------------- 1 file changed, 23 insertions(+), 22 deletions(-) diff --git a/crates/aionui-cron/src/service.rs b/crates/aionui-cron/src/service.rs index d0d6f5bbf..118825de4 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -972,7 +972,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, @@ -1043,7 +1043,8 @@ 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; if let Some(user_id) = owner_user_id.as_deref() { self.emitter.emit_job_executed(user_id, job_id, "ok", None); @@ -1057,7 +1058,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"); } } @@ -1073,7 +1074,7 @@ 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; @@ -1088,7 +1089,7 @@ 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; if let Some(user_id) = owner_user_id.as_deref() { self.emitter.emit_job_executed(user_id, job_id, "error", Some(&message)); @@ -1100,11 +1101,11 @@ impl CronService { 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.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.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 } => { @@ -1112,7 +1113,7 @@ impl CronService { 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, @@ -1126,7 +1127,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(user_id, job_id, ¶ms).await { error!( job_id, error = %err, @@ -1138,8 +1139,8 @@ impl CronService { } } - 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) => { @@ -1165,7 +1166,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; } @@ -1185,8 +1186,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) => { @@ -1203,7 +1204,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"); } } @@ -1227,7 +1228,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); @@ -1255,7 +1256,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); @@ -1268,7 +1269,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, @@ -1317,7 +1318,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, @@ -1333,7 +1334,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, @@ -1435,7 +1436,7 @@ 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"); } } @@ -1522,7 +1523,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!( From 00bdf81e0cb49f3d30ff65ebf120a6308cd91c50 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 12:01:34 +0800 Subject: [PATCH 066/145] office: scope preview sessions by user --- crates/aionui-office/src/routes.rs | 18 +++++--- crates/aionui-office/src/watch_manager.rs | 56 +++++++++++++++++++---- 2 files changed, 58 insertions(+), 16 deletions(-) diff --git a/crates/aionui-office/src/routes.rs b/crates/aionui-office/src/routes.rs index ed4975de0..fa25be108 100644 --- a/crates/aionui-office/src/routes.rs +++ b/crates/aionui-office/src/routes.rs @@ -92,10 +92,10 @@ async fn start_word_preview( 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( @@ -108,10 +108,10 @@ async fn start_excel_preview( 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( @@ -124,10 +124,10 @@ async fn start_ppt_preview( 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( @@ -161,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())) } diff --git a/crates/aionui-office/src/watch_manager.rs b/crates/aionui-office/src/watch_manager.rs index dc1ee8667..690b4b3a9 100644 --- a/crates/aionui-office/src/watch_manager.rs +++ b/crates/aionui-office/src/watch_manager.rs @@ -83,7 +83,7 @@ impl OfficecliWatchManager { 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() { @@ -137,7 +137,7 @@ impl OfficecliWatchManager { 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 { @@ -197,11 +197,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; @@ -444,8 +448,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) -> String { + format!("{user_id}:{doc_type}:{resolved_path}") } fn is_port_in_use_start_failure(error: &OfficeError) -> bool { @@ -625,14 +629,21 @@ 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"); + let key = session_key("user-1", "/path/to/doc.docx", DocType::Word); + assert_eq!(key, "user-1:word:/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); } @@ -712,6 +723,33 @@ 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(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()); From 4009fd0a6d24a67edf8c432ef53246fa7dfb33ab Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 12:07:32 +0800 Subject: [PATCH 067/145] conversation: guard runtime ops by owner --- crates/aionui-conversation/src/routes_aux.rs | 16 +++++---- crates/aionui-conversation/src/service_ops.rs | 35 +++++++++++++++++-- .../aionui-conversation/src/service_test.rs | 15 +++++++- 3 files changed, 57 insertions(+), 9 deletions(-) diff --git a/crates/aionui-conversation/src/routes_aux.rs b/crates/aionui-conversation/src/routes_aux.rs index 8ada037fe..ec5d646d9 100644 --- a/crates/aionui-conversation/src/routes_aux.rs +++ b/crates/aionui-conversation/src/routes_aux.rs @@ -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,11 +72,15 @@ 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)?, ))) } diff --git a/crates/aionui-conversation/src/service_ops.rs b/crates/aionui-conversation/src/service_ops.rs index d696912e2..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 @@ -52,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, @@ -134,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 @@ -152,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)? @@ -163,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 diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index 6b25e3f2d..fe52ba238 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -3614,12 +3614,25 @@ 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 ensure_runtime_recovers_missing_agent_and_returns_snapshot() { let task_mgr = Arc::new(MockTaskManager::new()); From 96e7dc4cc609334c4c2e5aaca34e6f54e09e5aaa Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 12:11:51 +0800 Subject: [PATCH 068/145] conversation: scope active runtime count --- crates/aionui-ai-agent/src/task_manager.rs | 25 +++++++++++++++++ crates/aionui-conversation/src/routes.rs | 4 +-- crates/aionui-conversation/src/service.rs | 16 +++++++++++ .../aionui-conversation/src/service_test.rs | 28 +++++++++++++++++++ 4 files changed, 71 insertions(+), 2 deletions(-) diff --git a/crates/aionui-ai-agent/src/task_manager.rs b/crates/aionui-ai-agent/src/task_manager.rs index b51b2778f..2cc87ee63 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 @@ -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-conversation/src/routes.rs b/crates/aionui-conversation/src/routes.rs index bbb06d2b3..dbebe2993 100644 --- a/crates/aionui-conversation/src/routes.rs +++ b/crates/aionui-conversation/src/routes.rs @@ -419,9 +419,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 }))) } diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index 195be2d20..247875be6 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -584,6 +584,22 @@ 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) + } + async fn send_message_response( &self, conversation_id: &str, diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index fe52ba238..ae1a71580 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -3159,6 +3159,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![] } @@ -3281,6 +3285,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![] } @@ -3633,6 +3641,26 @@ async fn get_config_options_rejects_cross_user_active_task() { 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 ensure_runtime_recovers_missing_agent_and_returns_snapshot() { let task_mgr = Arc::new(MockTaskManager::new()); From 08b5d89f166bcb26d5a3326f6d3538efb9d5ba80 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 12:26:15 +0800 Subject: [PATCH 069/145] runtime: scope status events by user --- crates/aionui-ai-agent/src/factory/acp.rs | 25 ++++++-- crates/aionui-ai-agent/src/factory/aionrs.rs | 26 +++++--- crates/aionui-ai-agent/src/routes/agent.rs | 4 +- crates/aionui-ai-agent/src/runtime_status.rs | 60 ++++++++++++++++++ crates/aionui-ai-agent/src/services/custom.rs | 8 +-- crates/aionui-api-types/src/runtime.rs | 4 ++ crates/aionui-mcp/src/connection_test/mod.rs | 62 +++++++++++++++++-- crates/aionui-mcp/src/routes.rs | 1 + crates/aionui-system/src/runtime_prepare.rs | 2 + 9 files changed, 168 insertions(+), 24 deletions(-) diff --git a/crates/aionui-ai-agent/src/factory/acp.rs b/crates/aionui-ai-agent/src/factory/acp.rs index 3ae0716f6..caf94fa95 100644 --- a/crates/aionui-ai-agent/src/factory/acp.rs +++ b/crates/aionui-ai-agent/src/factory/acp.rs @@ -52,8 +52,14 @@ pub(super) async fn build( config.backend.clone_from(&meta.backend); } - 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 { @@ -172,6 +178,7 @@ pub(super) async fn build( async fn resolve_agent_command_spec( meta: &aionui_api_types::AgentMetadata, + user_id: &str, workspace: &str, conversation_id: &str, broadcaster: Arc, @@ -180,7 +187,8 @@ async fn resolve_agent_command_spec( && let Some(backend) = meta.backend.as_deref() && let Some(tool) = ManagedAcpToolId::from_backend(backend) { - return resolve_builtin_managed_acp_command_spec(meta, workspace, conversation_id, broadcaster, tool).await; + return resolve_builtin_managed_acp_command_spec(meta, workspace, user_id, conversation_id, broadcaster, tool) + .await; } let command = meta @@ -188,7 +196,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)))?; @@ -233,6 +241,7 @@ async fn resolve_agent_command_spec( async fn resolve_builtin_managed_acp_command_spec( meta: &aionui_api_types::AgentMetadata, workspace: &str, + user_id: &str, conversation_id: &str, broadcaster: Arc, tool: ManagedAcpToolId, @@ -246,12 +255,14 @@ async fn resolve_builtin_managed_acp_command_spec( ))); } - let node_reporter = conversation_runtime_reporter(broadcaster.clone(), conversation_id.to_owned()); + let node_reporter = + conversation_runtime_reporter(broadcaster.clone(), user_id.to_owned(), conversation_id.to_owned()); let node_runtime = ensure_node_runtime_with_reporter(Some(node_reporter.as_ref())) .await .map_err(|error| AgentError::bad_request(format!("Agent '{}' CLI unavailable: {error}", meta.name)))?; - let tool_reporter = conversation_acp_tool_runtime_reporter(broadcaster, conversation_id.to_owned(), tool); + let tool_reporter = + conversation_acp_tool_runtime_reporter(broadcaster, user_id.to_owned(), conversation_id.to_owned(), tool); let managed_tool = ensure_managed_acp_tool_with_reporter(tool, Some(tool_reporter.as_ref())) .await .map_err(|error| AgentError::bad_request(format!("Agent '{}' CLI unavailable: {error}", meta.name)))?; @@ -740,6 +751,7 @@ mod tests { let spec = resolve_agent_command_spec( &meta, + "user-acp", "/tmp/workspace", "conv-acp", Arc::new(BroadcastEventBus::new(16)), @@ -761,6 +773,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)), diff --git a/crates/aionui-ai-agent/src/factory/aionrs.rs b/crates/aionui-ai-agent/src/factory/aionrs.rs index cf33b7377..da9fba085 100644 --- a/crates/aionui-ai-agent/src/factory/aionrs.rs +++ b/crates/aionui-ai-agent/src/factory/aionrs.rs @@ -65,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(), ) @@ -472,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); } @@ -493,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 { @@ -520,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, @@ -591,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 { @@ -601,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), @@ -664,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!( @@ -695,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())?; @@ -1001,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"); @@ -1834,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/routes/agent.rs b/crates/aionui-ai-agent/src/routes/agent.rs index f30a59ae1..b50c5ddb4 100644 --- a/crates/aionui-ai-agent/src/routes/agent.rs +++ b/crates/aionui-ai-agent/src/routes/agent.rs @@ -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)?, ))) diff --git a/crates/aionui-ai-agent/src/runtime_status.rs b/crates/aionui-ai-agent/src/runtime_status.rs index f6d797a11..08f03bcb7 100644 --- a/crates/aionui-ai-agent/src/runtime_status.rs +++ b/crates/aionui-ai-agent/src/runtime_status.rs @@ -12,10 +12,12 @@ use aionui_runtime::{ 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(), @@ -25,10 +27,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(), @@ -38,11 +42,13 @@ pub(crate) fn custom_agent_runtime_reporter( pub(crate) fn conversation_acp_tool_runtime_reporter( broadcaster: Arc, + user_id: impl Into, conversation_id: impl Into, tool: ManagedAcpToolId, ) -> SharedManagedAcpToolProgressReporter { acp_tool_runtime_reporter( broadcaster, + Some(user_id.into()), RuntimeStatusScope { kind: RuntimeStatusScopeKind::Conversation, id: conversation_id.into(), @@ -53,10 +59,12 @@ pub(crate) fn conversation_acp_tool_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(), @@ -72,11 +80,13 @@ fn node_runtime_reporter( fn acp_tool_runtime_reporter( broadcaster: Arc, + user_id: Option, scope: RuntimeStatusScope, tool: ManagedAcpToolId, ) -> SharedManagedAcpToolProgressReporter { Arc::new(move |update: ManagedAcpToolProgress| { let payload = RuntimeStatusPayload { + user_id: user_id.clone(), resource: RuntimeResourceKind::AcpTool, resource_id: Some(tool.slug().to_owned()), scope: scope.clone(), @@ -139,3 +149,53 @@ fn map_acp_failure_kind(kind: ManagedAcpToolFailureKind) -> RuntimeFailureKind { ManagedAcpToolFailureKind::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/custom.rs b/crates/aionui-ai-agent/src/services/custom.rs index dbd76df48..d25b7f119 100644 --- a/crates/aionui-ai-agent/src/services/custom.rs +++ b/crates/aionui-ai-agent/src/services/custom.rs @@ -34,15 +34,15 @@ 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) } diff --git a/crates/aionui-api-types/src/runtime.rs b/crates/aionui-api-types/src/runtime.rs index c7ee7c730..30e5ca5fd 100644 --- a/crates/aionui-api-types/src/runtime.rs +++ b/crates/aionui-api-types/src/runtime.rs @@ -48,6 +48,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, @@ -89,6 +91,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 { @@ -103,6 +106,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-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/routes.rs b/crates/aionui-mcp/src/routes.rs index 3d315c730..671818ab1 100644 --- a/crates/aionui-mcp/src/routes.rs +++ b/crates/aionui-mcp/src/routes.rs @@ -205,6 +205,7 @@ 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; diff --git a/crates/aionui-system/src/runtime_prepare.rs b/crates/aionui-system/src/runtime_prepare.rs index 3433a73e6..416ccad43 100644 --- a/crates/aionui-system/src/runtime_prepare.rs +++ b/crates/aionui-system/src/runtime_prepare.rs @@ -57,6 +57,7 @@ impl RuntimePrepareService { let broadcaster = self.broadcaster.clone(); Arc::new(move |update: NodeRuntimeProgress| { let payload = RuntimeStatusPayload { + user_id: None, resource: RuntimeResourceKind::Node, resource_id: None, scope: scope.clone(), @@ -78,6 +79,7 @@ impl RuntimePrepareService { let broadcaster = self.broadcaster.clone(); Arc::new(move |update: ManagedAcpToolProgress| { let payload = RuntimeStatusPayload { + user_id: None, resource: RuntimeResourceKind::AcpTool, resource_id: Some(tool.slug().to_owned()), scope: scope.clone(), From f965a4be75f6ede94a33d1f9e9b9051732bd48ec Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 12:31:41 +0800 Subject: [PATCH 070/145] system: scope runtime prepare events by user --- crates/aionui-system/src/routes.rs | 9 +- crates/aionui-system/src/runtime_prepare.rs | 136 +++++++++++++++++++- 2 files changed, 137 insertions(+), 8 deletions(-) diff --git a/crates/aionui-system/src/routes.rs b/crates/aionui-system/src/routes.rs index 42cb70a30..9ba61cb5b 100644 --- a/crates/aionui-system/src/routes.rs +++ b/crates/aionui-system/src/routes.rs @@ -309,21 +309,26 @@ 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))) } async fn ensure_managed_acp_tool( 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_managed_acp_tool(req.scope, &req.tool_id) + .ensure_managed_acp_tool_for_user(&user.id, req.scope, &req.tool_id) .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 416ccad43..667efb0ee 100644 --- a/crates/aionui-system/src/runtime_prepare.rs +++ b/crates/aionui-system/src/runtime_prepare.rs @@ -27,7 +27,24 @@ 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()))?; @@ -38,26 +55,49 @@ impl RuntimePrepareService { &self, scope: RuntimeStatusScope, tool_id: &str, + ) -> Result { + self.ensure_managed_acp_tool_with_user(None, scope, tool_id).await + } + + pub async fn ensure_managed_acp_tool_for_user( + &self, + user_id: &str, + scope: RuntimeStatusScope, + tool_id: &str, + ) -> Result { + self.ensure_managed_acp_tool_with_user(Some(user_id.to_owned()), scope, tool_id) + .await + } + + async fn ensure_managed_acp_tool_with_user( + &self, + user_id: Option, + scope: RuntimeStatusScope, + tool_id: &str, ) -> Result { let tool = ManagedAcpToolId::from_slug(tool_id) .ok_or_else(|| SystemError::BadRequest(format!("Unsupported managed ACP tool '{tool_id}'")))?; - let node_reporter = self.node_runtime_reporter(scope.clone()); + let node_reporter = self.node_runtime_reporter(user_id.clone(), scope.clone()); ensure_node_runtime_with_reporter(Some(node_reporter.as_ref())) .await .map_err(|error| SystemError::BadRequest(error.to_string()))?; - let tool_reporter = self.acp_tool_runtime_reporter(scope, tool); + let tool_reporter = self.acp_tool_runtime_reporter(user_id, scope, tool); ensure_managed_acp_tool_with_reporter(tool, Some(tool_reporter.as_ref())) .await .map_err(|error| SystemError::BadRequest(error.to_string()))?; Ok(EnsureManagedAcpToolResponse { 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: None, + user_id: user_id.clone(), resource: RuntimeResourceKind::Node, resource_id: None, scope: scope.clone(), @@ -73,13 +113,14 @@ impl RuntimePrepareService { fn acp_tool_runtime_reporter( &self, + user_id: Option, scope: RuntimeStatusScope, tool: ManagedAcpToolId, ) -> SharedManagedAcpToolProgressReporter { let broadcaster = self.broadcaster.clone(); Arc::new(move |update: ManagedAcpToolProgress| { let payload = RuntimeStatusPayload { - user_id: None, + user_id: user_id.clone(), resource: RuntimeResourceKind::AcpTool, resource_id: Some(tool.slug().to_owned()), scope: scope.clone(), @@ -143,3 +184,86 @@ fn map_acp_failure_kind(kind: ManagedAcpToolFailureKind) -> RuntimeFailureKind { ManagedAcpToolFailureKind::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 managed_acp_tool_reporter_scopes_route_event_to_user() { + let broadcaster = Arc::new(RecordingBroadcaster::default()); + let service = RuntimePrepareService::new(broadcaster.clone()); + let reporter = service.acp_tool_runtime_reporter( + Some("user-1".to_owned()), + conversation_scope(), + ManagedAcpToolId::CodexAcp, + ); + + reporter.report(ManagedAcpToolProgress::validating("validating tool")); + + 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"], "acp_tool"); + assert_eq!(events[0].data["resource_id"], "codex-acp"); + } + + #[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"); + } +} From 9362476a7a7a110384f28cf1f1f77f6462ac6597 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 12:40:20 +0800 Subject: [PATCH 071/145] realtime: validate websocket user sessions --- crates/aionui-app/src/router/state.rs | 41 +++++++++++++++++-- crates/aionui-realtime/src/handler.rs | 11 ++--- crates/aionui-realtime/src/lib.rs | 2 +- .../tests/handler_integration.rs | 8 +++- 4 files changed, 50 insertions(+), 12 deletions(-) diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index e462fa186..b5b00b841 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, @@ -871,7 +871,7 @@ 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(|_| Some("system_default_user".to_owned())), + token_user_resolver: Arc::new(|_| Box::pin(async { Some("system_default_user".to_owned()) })), token_extractor: Arc::new(|_| Some("local".into())), }; } @@ -879,8 +879,16 @@ pub fn build_ws_state(services: &AppServices) -> WsHandlerState { 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 token_user_resolver = - Arc::new(move |token: &str| jwt_service.verify(token).ok().map(|payload| payload.user_id)); + let user_repo = services.user_repo.clone(); + 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()??; + (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)); @@ -1057,6 +1065,31 @@ 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_channel_message_service_uses_app_conversation_service_for_assistant_bindings() { let db = aionui_db::init_database_memory().await.unwrap(); diff --git a/crates/aionui-realtime/src/handler.rs b/crates/aionui-realtime/src/handler.rs index 3f1091949..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,8 +21,8 @@ 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 internal user ID carried by it. -pub type TokenUserResolver = 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)] @@ -79,7 +80,7 @@ async fn handle_socket(socket: WebSocket, token: Option, state: WsHandle return; } - let Some(user_id) = (state.token_user_resolver)(&token) else { + let Some(user_id) = (state.token_user_resolver)(token.clone()).await else { send_realtime_error_and_close(socket, RealtimeError::AuthExpired, "authentication failed").await; return; }; @@ -269,7 +270,7 @@ mod tests { manager, router: Arc::new(crate::router::NoopMessageRouter), token_validator: Arc::new(|_| true), - token_user_resolver: Arc::new(|_| Some("system_default_user".into())), + token_user_resolver: Arc::new(|_| Box::pin(async { Some("system_default_user".into()) })), token_extractor: Arc::new(|_| None), } } @@ -515,7 +516,7 @@ mod tests { manager, router: router.clone(), token_validator: Arc::new(|_| true), - token_user_resolver: Arc::new(|_| Some("system_default_user".into())), + 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/tests/handler_integration.rs b/crates/aionui-realtime/tests/handler_integration.rs index 53b0eae5f..42b8092eb 100644 --- a/crates/aionui-realtime/tests/handler_integration.rs +++ b/crates/aionui-realtime/tests/handler_integration.rs @@ -37,7 +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| (t == "valid-token").then(|| "test-user".to_owned())), + 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") @@ -397,7 +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| (t == "valid-token").then(|| "test-user".to_owned())), + 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") From 193ad419716f252cd00ee38b8dd174ceb6fbcb3c Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 12:42:31 +0800 Subject: [PATCH 072/145] auth: revoke sessions when disabling users --- .../aionui-db/src/repository/sqlite_user.rs | 27 ++++++++++++++----- crates/aionui-db/src/repository/user.rs | 3 ++- crates/aionui-db/tests/user_repository.rs | 7 ++++- 3 files changed, 29 insertions(+), 8 deletions(-) diff --git a/crates/aionui-db/src/repository/sqlite_user.rs b/crates/aionui-db/src/repository/sqlite_user.rs index 1ade5b323..4e6026cf9 100644 --- a/crates/aionui-db/src/repository/sqlite_user.rs +++ b/crates/aionui-db/src/repository/sqlite_user.rs @@ -302,12 +302,22 @@ impl IUserRepository for SqliteUserRepository { 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 = ?, updated_at = ? WHERE id = ?") - .bind(status.as_str()) - .bind(now) - .bind(user_id) - .execute(&self.pool) - .await?; + 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"))); @@ -662,7 +672,12 @@ mod tests { 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] diff --git a/crates/aionui-db/src/repository/user.rs b/crates/aionui-db/src/repository/user.rs index 2bf771b65..0ed780bd3 100644 --- a/crates/aionui-db/src/repository/user.rs +++ b/crates/aionui-db/src/repository/user.rs @@ -76,7 +76,8 @@ 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. + /// 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. diff --git a/crates/aionui-db/tests/user_repository.rs b/crates/aionui-db/tests/user_repository.rs index d48332540..351748679 100644 --- a/crates/aionui-db/tests/user_repository.rs +++ b/crates/aionui-db/tests/user_repository.rs @@ -280,8 +280,13 @@ async fn t2_15_disabled_user_is_not_active() { assert!(r.find_active_by_id(&user.id).await.unwrap().is_some()); r.set_status(&user.id, UserStatus::Disabled).await.unwrap(); - assert!(r.find_by_id(&user.id).await.unwrap().is_some()); + 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] From e7e4e7dae57686d6330f4ca2458e58223dbc337f Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 12:45:40 +0800 Subject: [PATCH 073/145] auth: add external session revocation --- crates/aionui-api-types/src/auth.rs | 33 ++++++++++++++++++ crates/aionui-api-types/src/lib.rs | 6 ++-- crates/aionui-auth/src/csrf.rs | 2 +- crates/aionui-auth/src/routes.rs | 33 ++++++++++++++++-- crates/aionui-auth/src/service.rs | 19 +++++++++- crates/aionui-auth/tests/route_tests.rs | 46 +++++++++++++++++++++++++ 6 files changed, 132 insertions(+), 7 deletions(-) diff --git a/crates/aionui-api-types/src/auth.rs b/crates/aionui-api-types/src/auth.rs index 584e19585..889d60acb 100644 --- a/crates/aionui-api-types/src/auth.rs +++ b/crates/aionui-api-types/src/auth.rs @@ -130,6 +130,20 @@ pub struct EnsureExternalSessionResponse { 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")] @@ -277,6 +291,25 @@ mod tests { 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/lib.rs b/crates/aionui-api-types/src/lib.rs index 74cb78c6a..e5b5bd25f 100644 --- a/crates/aionui-api-types/src/lib.rs +++ b/crates/aionui-api-types/src/lib.rs @@ -63,9 +63,9 @@ pub use assistant::{ pub use auth::{ AuthStatusResponse, ChangePasswordRequest, EnsureExternalSessionRequest, EnsureExternalSessionResponse, EnsureExternalUserRequest, EnsureExternalUserResponse, ExternalUserType, InternalAuthErrorCode, LoginRequest, - LoginResponse, PublicUser, QrLoginRequest, RefreshResponse, RefreshTokenRequest, UserInfoResponse, - WebuiChangePasswordRequest, WebuiChangeUsernameRequest, WebuiChangeUsernameResponse, WebuiGenerateQrTokenResponse, - WebuiResetPasswordResponse, WsTokenResponse, + LoginResponse, PublicUser, QrLoginRequest, RefreshResponse, RefreshTokenRequest, RevokeExternalSessionRequest, + RevokeExternalSessionResponse, UserInfoResponse, WebuiChangePasswordRequest, WebuiChangeUsernameRequest, + WebuiChangeUsernameResponse, WebuiGenerateQrTokenResponse, WebuiResetPasswordResponse, WsTokenResponse, }; pub use channel::{ ApprovePairingRequest, BridgeResponse, ChannelAssistantSettingRequest, ChannelAssistantSettingResponse, diff --git a/crates/aionui-auth/src/csrf.rs b/crates/aionui-auth/src/csrf.rs index 3577d9bcb..35b5504bf 100644 --- a/crates/aionui-auth/src/csrf.rs +++ b/crates/aionui-auth/src/csrf.rs @@ -38,7 +38,7 @@ pub async fn csrf_middleware( let is_exempt = path == "/login" || path == "/api/auth/qr-login" || path.starts_with("/api/auth/internal/external-users/") - || path == "/api/auth/internal/external-sessions"; + || path.starts_with("/api/auth/internal/external-sessions"); if needs_validation && !is_exempt { let header_token = request diff --git a/crates/aionui-auth/src/routes.rs b/crates/aionui-auth/src/routes.rs index 71773f500..ca928a9d4 100644 --- a/crates/aionui-auth/src/routes.rs +++ b/crates/aionui-auth/src/routes.rs @@ -15,8 +15,9 @@ use serde::Deserialize; use aionui_api_types::{ ApiResponse, AuthStatusResponse, ChangePasswordRequest, EnsureExternalSessionRequest, EnsureExternalUserRequest, EnsureExternalUserResponse, LoginRequest, LoginResponse, PublicUser, QrLoginRequest, RefreshResponse, - RefreshTokenRequest, UserInfoResponse, WebuiChangePasswordRequest, WebuiChangeUsernameRequest, - WebuiChangeUsernameResponse, WebuiGenerateQrTokenResponse, WebuiResetPasswordResponse, WsTokenResponse, + RefreshTokenRequest, RevokeExternalSessionRequest, RevokeExternalSessionResponse, UserInfoResponse, + WebuiChangePasswordRequest, WebuiChangeUsernameRequest, WebuiChangeUsernameResponse, WebuiGenerateQrTokenResponse, + WebuiResetPasswordResponse, WsTokenResponse, }; use aionui_common::ApiError; use aionui_common::constants::COOKIE_MAX_AGE_DAYS; @@ -221,6 +222,10 @@ pub fn auth_routes(state: AuthRouterState) -> Router { "/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), @@ -344,6 +349,30 @@ async fn create_external_session_handler( Ok(([(header::SET_COOKIE, cookie)], Json(ApiResponse::ok(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" + ); + Ok(Json(ApiResponse::ok(response))) +} + // --------------------------------------------------------------------------- // POST /login // --------------------------------------------------------------------------- diff --git a/crates/aionui-auth/src/service.rs b/crates/aionui-auth/src/service.rs index 9569ad1ce..59b3614e9 100644 --- a/crates/aionui-auth/src/service.rs +++ b/crates/aionui-auth/src/service.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use aionui_api_types::{ EnsureExternalSessionRequest, EnsureExternalSessionResponse, EnsureExternalUserRequest, EnsureExternalUserResponse, - ExternalUserType, PublicUser, + ExternalUserType, PublicUser, RevokeExternalSessionRequest, RevokeExternalSessionResponse, }; use aionui_db::{ExternalUserProjection, IUserRepository, UserStatus, UserType, models::User}; @@ -86,6 +86,23 @@ impl AuthProvisionService { session_generation: user.session_generation, }) } + + 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 { diff --git a/crates/aionui-auth/tests/route_tests.rs b/crates/aionui-auth/tests/route_tests.rs index d48dec208..1c9b0fbeb 100644 --- a/crates/aionui-auth/tests/route_tests.rs +++ b/crates/aionui-auth/tests/route_tests.rs @@ -880,6 +880,52 @@ async fn external_session_exchange_rejects_unprovisioned_user() { assert_eq!(json["code"], "USER_CONTEXT_REQUIRED"); } +#[tokio::test] +async fn external_session_revoke_invalidates_existing_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-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 session_json = body_json(session).await; + let token = session_json["data"]["token"].as_str().unwrap().to_owned(); + assert_eq!(session_json["data"]["session_generation"], 0); + + 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); + + 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; From 535f67b8cedf02bd51fc97f2e3ae2be8b3276f90 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 12:50:05 +0800 Subject: [PATCH 074/145] file: add user scoped office watch cleanup --- crates/aionui-file/src/routes.rs | 9 ++++ crates/aionui-file/src/traits.rs | 6 +++ crates/aionui-file/src/watch_service.rs | 12 ++++++ crates/aionui-file/tests/file_watching.rs | 51 +++++++++++++++++++++++ 4 files changed, 78 insertions(+) diff --git a/crates/aionui-file/src/routes.rs b/crates/aionui-file/src/routes.rs index 954b460e5..8c3cf76fe 100644 --- a/crates/aionui-file/src/routes.rs +++ b/crates/aionui-file/src/routes.rs @@ -144,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)) @@ -544,6 +545,14 @@ async fn stop_office_watch( 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())) +} + // --------------------------------------------------------------------------- // E. Workspace snapshot — handlers // --------------------------------------------------------------------------- diff --git a/crates/aionui-file/src/traits.rs b/crates/aionui-file/src/traits.rs index 410163260..7b16f2e85 100644 --- a/crates/aionui-file/src/traits.rs +++ b/crates/aionui-file/src/traits.rs @@ -212,6 +212,12 @@ pub trait IFileWatchService: Send + Sync { 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 7ec04c2be..c6f8850f0 100644 --- a/crates/aionui-file/src/watch_service.rs +++ b/crates/aionui-file/src/watch_service.rs @@ -328,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 fb6f1165d..190e1891b 100644 --- a/crates/aionui-file/tests/file_watching.rs +++ b/crates/aionui-file/tests/file_watching.rs @@ -254,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(); From ab21983d97dcb15be6c4d24dff643c73f396de0f Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 12:51:24 +0800 Subject: [PATCH 075/145] realtime: add user websocket disconnect --- crates/aionui-realtime/src/manager.rs | 87 +++++++++++++++++++++++++++ 1 file changed, 87 insertions(+) diff --git a/crates/aionui-realtime/src/manager.rs b/crates/aionui-realtime/src/manager.rs index a70c7c0e7..b0b39ccb1 100644 --- a/crates/aionui-realtime/src/manager.rs +++ b/crates/aionui-realtime/src/manager.rs @@ -79,6 +79,49 @@ impl WebSocketManager { .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 @@ -441,6 +484,50 @@ mod tests { 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(); From 92f8f66f487a34a7a8e3f8b20e51fc1b9577e943 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 12:53:14 +0800 Subject: [PATCH 076/145] auth: disconnect websockets on external revoke --- crates/aionui-app/src/router/routes.rs | 6 +++++ crates/aionui-auth/src/lib.rs | 2 +- crates/aionui-auth/src/routes.rs | 6 +++++ crates/aionui-auth/tests/route_tests.rs | 29 ++++++++++++++++++++++--- 4 files changed, 39 insertions(+), 4 deletions(-) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index bd7db5a66..b4d9d268a 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -166,6 +166,12 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates 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(); + Some(Arc::new(move |user_id: &str| { + ws_manager.disconnect_user(user_id, "session revoked"); + })) + }, local: services.local, }; diff --git a/crates/aionui-auth/src/lib.rs b/crates/aionui-auth/src/lib.rs index bbe6d78d5..8d9c410c4 100644 --- a/crates/aionui-auth/src/lib.rs +++ b/crates/aionui-auth/src/lib.rs @@ -57,6 +57,6 @@ pub use middleware::{AuthIdentityMode, AuthState, CurrentUser, auth_middleware, 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/routes.rs b/crates/aionui-auth/src/routes.rs index ca928a9d4..b256eaaca 100644 --- a/crates/aionui-auth/src/routes.rs +++ b/crates/aionui-auth/src/routes.rs @@ -37,6 +37,8 @@ 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 { @@ -71,6 +73,7 @@ pub struct AuthRouterState { pub qr_token_store: Arc, pub identity_mode: AuthIdentityMode, pub bootstrap_secret: Option>, + pub session_revoked_hook: Option>, pub local: bool, } @@ -370,6 +373,9 @@ async fn revoke_external_session_handler( 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))) } diff --git a/crates/aionui-auth/tests/route_tests.rs b/crates/aionui-auth/tests/route_tests.rs index 1c9b0fbeb..ddcf0404e 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; @@ -13,7 +13,8 @@ use http_body_util::BodyExt; use tower::ServiceExt; use aionui_auth::{ - AuthIdentityMode, AuthRouterState, CookieConfig, JwtService, QrTokenStore, auth_routes, hash_password, + AuthIdentityMode, AuthRouterState, CookieConfig, JwtService, QrTokenStore, SessionRevokedHook, auth_routes, + hash_password, }; use aionui_db::{IUserRepository, SqliteUserRepository, UserStatus, init_database_memory}; @@ -31,6 +32,14 @@ async fn test_app_with_local(local: bool) -> (Router, TestContext) { } async fn test_app_with_options(local: bool, bootstrap_secret: Option<&str>) -> (Router, TestContext) { + test_app_with_options_and_hook(local, bootstrap_secret, None).await +} + +async fn test_app_with_options_and_hook( + local: bool, + bootstrap_secret: Option<&str>, + 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())); @@ -51,6 +60,7 @@ async fn test_app_with_options(local: bool, bootstrap_secret: Option<&str>) -> ( AuthIdentityMode::UserSession }, bootstrap_secret: bootstrap_secret.map(Arc::::from), + session_revoked_hook, local, }; @@ -882,7 +892,16 @@ async fn external_session_exchange_rejects_unprovisioned_user() { #[tokio::test] async fn external_session_revoke_invalidates_existing_token() { - let (app, _ctx) = test_app_with_options(false, Some("bootstrap-secret")).await; + 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"), + Some(Arc::new(move |user_id: &str| { + hook_users.lock().unwrap().push(user_id.to_owned()); + })), + ) + .await; let provision = app .clone() @@ -921,6 +940,10 @@ async fn external_session_revoke_invalidates_existing_token() { 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); From 034658aa801bf680eb8f8e8676af0aa7e5a3b1fb Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 13:01:47 +0800 Subject: [PATCH 077/145] conversation: scope runtime skill resolution --- crates/aionui-conversation/src/service.rs | 10 ++- .../aionui-conversation/src/skill_resolver.rs | 81 ++++++++++++++++++- .../aionui-conversation/src/stream_relay.rs | 5 +- crates/aionui-extension/src/lib.rs | 4 +- 4 files changed, 93 insertions(+), 7 deletions(-) diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index 247875be6..48d7a38c6 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -997,7 +997,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 @@ -3414,7 +3417,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; } 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/stream_relay.rs b/crates/aionui-conversation/src/stream_relay.rs index a38f6a1a4..25b8a7256 100644 --- a/crates/aionui-conversation/src/stream_relay.rs +++ b/crates/aionui-conversation/src/stream_relay.rs @@ -734,6 +734,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 })); @@ -780,6 +781,7 @@ impl StreamRelay { struct SharedSkillResolver { resolver: Arc, + user_id: String, allowed_skill_names: Vec, } @@ -794,7 +796,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 } } @@ -886,6 +888,7 @@ mod tests { let resolver: Arc = concrete.clone(); let loader = SharedSkillResolver { resolver, + user_id: "system_default_user".into(), allowed_skill_names: vec!["cron".into()], }; diff --git a/crates/aionui-extension/src/lib.rs b/crates/aionui-extension/src/lib.rs index 5aa63b043..b46376b84 100644 --- a/crates/aionui-extension/src/lib.rs +++ b/crates/aionui-extension/src/lib.rs @@ -56,8 +56,8 @@ pub use skill_service::{ 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, sync_skill_catalog_into_repo_for_user, + 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, From b4f18626455a53b264a7000607f78dfd485fb057 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 13:04:01 +0800 Subject: [PATCH 078/145] extension: keep managed cron skills global --- crates/aionui-db/migrations/026_user_scope.sql | 2 +- crates/aionui-db/tests/skill_management_schema.rs | 2 +- crates/aionui-extension/src/skill_service.rs | 5 +++-- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/crates/aionui-db/migrations/026_user_scope.sql b/crates/aionui-db/migrations/026_user_scope.sql index 838541369..710b3bc80 100644 --- a/crates/aionui-db/migrations/026_user_scope.sql +++ b/crates/aionui-db/migrations/026_user_scope.sql @@ -523,7 +523,7 @@ INSERT INTO skills_new ( ) SELECT id, - CASE WHEN source = 'builtin' THEN NULL ELSE 'system_default_user' END, + CASE WHEN source IN ('builtin', 'cron') THEN NULL ELSE 'system_default_user' END, name, description, path, source, enabled, deleted_at, created_at, updated_at FROM skills; diff --git a/crates/aionui-db/tests/skill_management_schema.rs b/crates/aionui-db/tests/skill_management_schema.rs index 39c6e7dcd..74826949c 100644 --- a/crates/aionui-db/tests/skill_management_schema.rs +++ b/crates/aionui-db/tests/skill_management_schema.rs @@ -61,7 +61,7 @@ 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" { + let user_id = if source == "builtin" || source == "cron" { None } else { Some("system_default_user") diff --git a/crates/aionui-extension/src/skill_service.rs b/crates/aionui-extension/src/skill_service.rs index 2f0f4d400..edcd6713f 100644 --- a/crates/aionui-extension/src/skill_service.rs +++ b/crates/aionui-extension/src/skill_service.rs @@ -1702,7 +1702,7 @@ async fn sync_managed_skill_into_repo( skill: &ScannedSkill, source: &str, ) -> Result<(), ExtensionError> { - if let Some(existing) = repo.find_by_name_any(&skill.name).await? + if let Some(existing) = repo.find_by_name_any_for_user(DEFAULT_USER_ID, &skill.name).await? && existing.source == "user" && existing.deleted_at.is_none() && existing.enabled @@ -1710,7 +1710,7 @@ async fn sync_managed_skill_into_repo( return Ok(()); } - repo.upsert(UpsertSkillParams { + repo.upsert_global(UpsertSkillParams { name: &skill.name, description: Some(&skill.description), path: &skill.path, @@ -3011,6 +3011,7 @@ mod tests { assert_eq!(scheduled.relative_location, None); let scheduled_row = repo.find_by_name("scheduled-task").await.unwrap().unwrap(); assert_eq!(scheduled_row.source, "cron"); + assert_eq!(scheduled_row.user_id, None); } #[tokio::test] From 02d85bfed0f542256e726e0c7b34c098115330b7 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 13:06:27 +0800 Subject: [PATCH 079/145] team: pass owner to config option lookup --- crates/aionui-app/src/router/team_conversation_adapters.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/aionui-app/src/router/team_conversation_adapters.rs b/crates/aionui-app/src/router/team_conversation_adapters.rs index 0e582ba3b..8542c5e9d 100644 --- a/crates/aionui-app/src/router/team_conversation_adapters.rs +++ b/crates/aionui-app/src/router/team_conversation_adapters.rs @@ -288,8 +288,9 @@ impl TeamConversationProvisioningPort for TeamConversationAdapters { } 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) } From 5a30d94520f3ab877d5c48393b7aa68af5442b9a Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 13:12:01 +0800 Subject: [PATCH 080/145] db: cover cron skill user scope migration --- .../aionui-db/tests/user_scope_migration.rs | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/crates/aionui-db/tests/user_scope_migration.rs b/crates/aionui-db/tests/user_scope_migration.rs index 36eff9d9b..8660781c0 100644 --- a/crates/aionui-db/tests/user_scope_migration.rs +++ b/crates/aionui-db/tests/user_scope_migration.rs @@ -1,5 +1,10 @@ +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})")) @@ -17,6 +22,41 @@ async fn table_indexes(pool: &sqlx::SqlitePool, table: &str) -> Vec { 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) { + 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.unwrap(); +} + #[tokio::test] async fn migration_026_adds_core_user_projection_columns() { let db = init_database_memory().await.unwrap(); @@ -58,6 +98,33 @@ async fn migration_026_adds_user_scope_to_independent_roots() { } } +#[tokio::test] +async fn migration_026_migrates_cron_skills_as_global_rows() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 25).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, 26).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"); + assert_eq!(row.get::, _>("user_id"), None); +} + #[tokio::test] async fn migration_026_keeps_new_conversation_cron_jobs_unanchored_until_run() { let db = init_database_memory().await.unwrap(); From f7804004d2b82106200c76220480f765db0b7827 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 13:15:50 +0800 Subject: [PATCH 081/145] app: require bootstrap secret for aionpro mode --- .../aionui-app/src/bootstrap/environment.rs | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/crates/aionui-app/src/bootstrap/environment.rs b/crates/aionui-app/src/bootstrap/environment.rs index 56c79d80b..720705a79 100644 --- a/crates/aionui-app/src/bootstrap/environment.rs +++ b/crates/aionui-app/src/bootstrap/environment.rs @@ -49,6 +49,7 @@ 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 @@ -135,6 +151,8 @@ pub async fn init_data_layer(config: &AppConfig) -> Result Date: Tue, 21 Jul 2026 13:17:37 +0800 Subject: [PATCH 082/145] mcp: cover user scoped oauth pending state --- crates/aionui-mcp/src/oauth_service.rs | 28 ++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/crates/aionui-mcp/src/oauth_service.rs b/crates/aionui-mcp/src/oauth_service.rs index 911bf56cf..dddce6ed0 100644 --- a/crates/aionui-mcp/src/oauth_service.rs +++ b/crates/aionui-mcp/src/oauth_service.rs @@ -680,8 +680,36 @@ 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] From 963fd0443cdc44ca0bf4be40b8a66315c3cdfe06 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 13:24:07 +0800 Subject: [PATCH 083/145] cron: validate helper user header --- crates/aionui-cron/src/routes.rs | 56 ++++++++++++++++++- .../aionui-cron/tests/service_integration.rs | 21 +++++-- 2 files changed, 70 insertions(+), 7 deletions(-) diff --git a/crates/aionui-cron/src/routes.rs b/crates/aionui-cron/src/routes.rs index 68aa865d1..979b4439e 100644 --- a/crates/aionui-cron/src/routes.rs +++ b/crates/aionui-cron/src/routes.rs @@ -145,10 +145,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 +161,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 +176,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,6 +201,19 @@ 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, @@ -265,6 +281,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(); diff --git a/crates/aionui-cron/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index 805ac9ec2..6d33ab3be 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}; @@ -32,6 +33,7 @@ use aionui_db::{ 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,15 @@ 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) \ @@ -2553,7 +2564,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() @@ -2637,7 +2649,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() @@ -2698,9 +2711,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] From 18ff21033bb4c2ce921be70697f9d4a1ac4a7cbd Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 13:29:56 +0800 Subject: [PATCH 084/145] app: fix test clippy warnings --- crates/aionui-app/src/router/state.rs | 8 ++++---- crates/aionui-app/tests/team_e2e.rs | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index b5b00b841..7b88e57e4 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -983,6 +983,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 { @@ -996,9 +998,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 { @@ -1015,7 +1015,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; diff --git a/crates/aionui-app/tests/team_e2e.rs b/crates/aionui-app/tests/team_e2e.rs index 36a2ac376..809385a87 100644 --- a/crates/aionui-app/tests/team_e2e.rs +++ b/crates/aionui-app/tests/team_e2e.rs @@ -622,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] From 563a620639bfc364e807be6aed9fe61d010ff6ca Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 13:35:27 +0800 Subject: [PATCH 085/145] db: scope feedback diagnostics agents --- .../src/repository/sqlite_diagnostics.rs | 21 +++-- .../tests/feedback_diagnostics_repository.rs | 84 +++++++++++++++++++ 2 files changed, 97 insertions(+), 8 deletions(-) diff --git a/crates/aionui-db/src/repository/sqlite_diagnostics.rs b/crates/aionui-db/src/repository/sqlite_diagnostics.rs index 4970ce45b..69133a34a 100644 --- a/crates/aionui-db/src/repository/sqlite_diagnostics.rs +++ b/crates/aionui-db/src/repository/sqlite_diagnostics.rs @@ -92,7 +92,7 @@ impl SqliteFeedbackDiagnosticsRepository { 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?, + "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?, }); @@ -428,7 +428,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 +441,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?; @@ -978,9 +979,11 @@ impl SqliteFeedbackDiagnosticsRepository { .bind(&request.user_id) .fetch_one(&self.pool) .await?; - let agent_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM agent_metadata") - .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) @@ -999,7 +1002,7 @@ 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?, + "agent_health": self.collect_global_agent_health(&request.user_id).await?, "provider_health": self.collect_global_provider_health(&request.user_id).await?, }), )) @@ -1419,7 +1422,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, \ @@ -1427,9 +1430,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?; diff --git a/crates/aionui-db/tests/feedback_diagnostics_repository.rs b/crates/aionui-db/tests/feedback_diagnostics_repository.rs index c8df0eca3..633760943 100644 --- a/crates/aionui-db/tests/feedback_diagnostics_repository.rs +++ b/crates/aionui-db/tests/feedback_diagnostics_repository.rs @@ -639,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(); @@ -833,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 @@ -879,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"], From 2c7f0d2e965a30242f912f3f7dadfc1fe7b7b111 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 13:45:55 +0800 Subject: [PATCH 086/145] agent: scope prompt skill lookup by user --- .../src/capability/first_message_injector.rs | 10 +- .../src/capability/skill_manager/mod.rs | 23 +++- crates/aionui-ai-agent/src/factory/acp.rs | 1 + .../src/factory/acp_assembler.rs | 4 + .../aionui-ai-agent/src/manager/acp/hooks.rs | 1 + .../tests/acp_agent_integration.rs | 1 + .../tests/prompt_pipeline_integration.rs | 1 + .../tests/skill_manager_integration.rs | 102 ++++++++++++++++++ crates/aionui-extension/src/lib.rs | 7 +- 9 files changed, 142 insertions(+), 8 deletions(-) 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 caf94fa95..032e5c7aa 100644 --- a/crates/aionui-ai-agent/src/factory/acp.rs +++ b/crates/aionui-ai-agent/src/factory/acp.rs @@ -123,6 +123,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, 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/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/tests/acp_agent_integration.rs b/crates/aionui-ai-agent/tests/acp_agent_integration.rs index 286d014da..aaa34c7a6 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/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-extension/src/lib.rs b/crates/aionui-extension/src/lib.rs index b46376b84..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, - 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, + 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, From 00cb3611822bb53e68b015dcfe2a772afa57e304 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 13:54:04 +0800 Subject: [PATCH 087/145] auth: disable qr login in aionpro mode --- crates/aionui-app/src/router/routes.rs | 1 + crates/aionui-auth/src/routes.rs | 10 ++++++++++ crates/aionui-auth/tests/route_tests.rs | 18 +++++++++++++++++- 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index b4d9d268a..90d7d030b 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -173,6 +173,7 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates })) }, local: services.local, + aionpro_mode: services.identity_mode == crate::config::IdentityMode::AionPro, }; let auth_mw_state = AuthState { diff --git a/crates/aionui-auth/src/routes.rs b/crates/aionui-auth/src/routes.rs index b256eaaca..b3644312d 100644 --- a/crates/aionui-auth/src/routes.rs +++ b/crates/aionui-auth/src/routes.rs @@ -75,6 +75,7 @@ pub struct AuthRouterState { pub bootstrap_secret: Option>, pub session_revoked_hook: Option>, pub local: bool, + pub aionpro_mode: bool, } #[derive(Debug, Deserialize)] @@ -791,6 +792,15 @@ async fn qr_login_handler( State(state): State, body: Result, JsonRejection>, ) -> Result { + if state.aionpro_mode { + return Err(ApiError::coded( + StatusCode::UNAUTHORIZED, + "USER_CONTEXT_REQUIRED", + "User context required.", + None, + )); + } + let Json(req) = body.map_err(ApiError::from)?; // Validate and consume QR token (one-time use) diff --git a/crates/aionui-auth/tests/route_tests.rs b/crates/aionui-auth/tests/route_tests.rs index ddcf0404e..b1c040c0c 100644 --- a/crates/aionui-auth/tests/route_tests.rs +++ b/crates/aionui-auth/tests/route_tests.rs @@ -32,12 +32,13 @@ async fn test_app_with_local(local: bool) -> (Router, TestContext) { } async fn test_app_with_options(local: bool, bootstrap_secret: Option<&str>) -> (Router, TestContext) { - test_app_with_options_and_hook(local, bootstrap_secret, None).await + 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(); @@ -62,6 +63,7 @@ async fn test_app_with_options_and_hook( bootstrap_secret: bootstrap_secret.map(Arc::::from), session_revoked_hook, local, + aionpro_mode, }; let app = auth_routes(state); @@ -748,6 +750,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) // =========================================================================== @@ -897,6 +912,7 @@ async fn external_session_revoke_invalidates_existing_token() { 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()); })), From 9c73cff7818081a155dc8e5c13f713228f9449b4 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 14:02:02 +0800 Subject: [PATCH 088/145] auth: reject local sessions in aionpro mode --- crates/aionui-app/src/router/routes.rs | 8 ++--- crates/aionui-auth/src/middleware.rs | 11 ++++++ crates/aionui-auth/src/routes.rs | 30 +++++++++++++---- crates/aionui-auth/tests/middleware_tests.rs | 35 +++++++++++++++++++- crates/aionui-auth/tests/route_tests.rs | 33 ++++++++++++++++++ 5 files changed, 105 insertions(+), 12 deletions(-) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index 90d7d030b..65e887789 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -336,10 +336,10 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates } fn auth_identity_mode(identity_mode: crate::config::IdentityMode) -> AuthIdentityMode { - if identity_mode.is_local() { - AuthIdentityMode::Local - } else { - AuthIdentityMode::UserSession + match identity_mode { + crate::config::IdentityMode::Local => AuthIdentityMode::Local, + crate::config::IdentityMode::WebUi => AuthIdentityMode::UserSession, + crate::config::IdentityMode::AionPro => AuthIdentityMode::AionPro, } } diff --git a/crates/aionui-auth/src/middleware.rs b/crates/aionui-auth/src/middleware.rs index 40a3ca318..05ccaa89e 100644 --- a/crates/aionui-auth/src/middleware.rs +++ b/crates/aionui-auth/src/middleware.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use axum::extract::{Request, State}; +use axum::http::StatusCode; use axum::middleware::Next; use axum::response::Response; @@ -16,6 +17,7 @@ use crate::extract::extract_token_from_headers; pub enum AuthIdentityMode { Local, UserSession, + AionPro, } /// Authenticated user injected into request extensions by the auth middleware. @@ -93,6 +95,15 @@ 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())); } diff --git a/crates/aionui-auth/src/routes.rs b/crates/aionui-auth/src/routes.rs index b3644312d..118a967e4 100644 --- a/crates/aionui-auth/src/routes.rs +++ b/crates/aionui-auth/src/routes.rs @@ -175,6 +175,15 @@ fn provision_error_to_api_error(err: ProvisionError) -> ApiError { } } +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: @@ -205,7 +214,11 @@ pub fn auth_routes(state: AuthRouterState) -> Router { let auth_state = AuthState { jwt_service: state.jwt_service.clone(), user_repo: state.user_repo.clone(), - identity_mode: AuthIdentityMode::UserSession, + identity_mode: if state.aionpro_mode { + AuthIdentityMode::AionPro + } else { + AuthIdentityMode::UserSession + }, }; // Auth rate limited routes (login, qr-login) @@ -388,6 +401,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) @@ -735,6 +752,10 @@ async fn refresh_handler( })? .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())); } @@ -793,12 +814,7 @@ async fn qr_login_handler( body: Result, JsonRejection>, ) -> Result { if state.aionpro_mode { - return Err(ApiError::coded( - StatusCode::UNAUTHORIZED, - "USER_CONTEXT_REQUIRED", - "User context required.", - None, - )); + return Err(user_context_required()); } let Json(req) = body.map_err(ApiError::from)?; diff --git a/crates/aionui-auth/tests/middleware_tests.rs b/crates/aionui-auth/tests/middleware_tests.rs index 792485d59..20bfab36a 100644 --- a/crates/aionui-auth/tests/middleware_tests.rs +++ b/crates/aionui-auth/tests/middleware_tests.rs @@ -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, - identity_mode: AuthIdentityMode::UserSession, + identity_mode, }; Router::new() @@ -265,6 +273,31 @@ async fn auth_middleware_session_generation_mismatch_returns_unauthorized_code() 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())); diff --git a/crates/aionui-auth/tests/route_tests.rs b/crates/aionui-auth/tests/route_tests.rs index b1c040c0c..1c088d4f7 100644 --- a/crates/aionui-auth/tests/route_tests.rs +++ b/crates/aionui-auth/tests/route_tests.rs @@ -57,6 +57,8 @@ async fn test_app_with_options_and_hook( qr_token_store: qr_token_store.clone(), identity_mode: if local { AuthIdentityMode::Local + } else if aionpro_mode { + AuthIdentityMode::AionPro } else { AuthIdentityMode::UserSession }, @@ -268,6 +270,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 @@ -632,6 +649,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) // =========================================================================== From 5d41afa147f55ac7f776a8b9b458ce4127ee6074 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 14:03:32 +0800 Subject: [PATCH 089/145] realtime: reject local websocket sessions in aionpro mode --- crates/aionui-app/src/router/state.rs | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index 7b88e57e4..7c5bdc01a 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -880,12 +880,16 @@ pub fn build_ws_state(services: &AppServices) -> WsHandlerState { 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) }) }); @@ -1090,6 +1094,27 @@ mod tests { 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(); From 7ae2cc4bd142d52d89bbe21272d647438ce30e9e Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 14:10:45 +0800 Subject: [PATCH 090/145] conversation: scope acp runtime mode persistence --- .../src/router/team_conversation_adapters.rs | 3 +- crates/aionui-conversation/src/service.rs | 14 ++++++++- .../aionui-conversation/src/service_test.rs | 29 +++++++++++++++++-- crates/aionui-cron/src/executor.rs | 2 +- 4 files changed, 43 insertions(+), 5 deletions(-) diff --git a/crates/aionui-app/src/router/team_conversation_adapters.rs b/crates/aionui-app/src/router/team_conversation_adapters.rs index 8542c5e9d..4a5317b75 100644 --- a/crates/aionui-app/src/router/team_conversation_adapters.rs +++ b/crates/aionui-app/src/router/team_conversation_adapters.rs @@ -281,8 +281,9 @@ impl TeamConversationProvisioningPort for TeamConversationAdapters { } 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) } diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index 48d7a38c6..a172e5804 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -2076,7 +2076,19 @@ impl ConversationService { 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) diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index ae1a71580..18495fc6b 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -3790,11 +3790,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); @@ -3806,6 +3810,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()); diff --git a/crates/aionui-cron/src/executor.rs b/crates/aionui-cron/src/executor.rs index 09aab2f4f..9520442d2 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -372,7 +372,7 @@ impl JobExecutor { 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?; } From 45fd436332465652692c5de92ad633f86b63feea Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 14:16:37 +0800 Subject: [PATCH 091/145] agent: scope acp factory catalog lookup --- crates/aionui-ai-agent/src/factory/acp.rs | 120 ++++++++++++++++++++-- 1 file changed, 109 insertions(+), 11 deletions(-) diff --git a/crates/aionui-ai-agent/src/factory/acp.rs b/crates/aionui-ai-agent/src/factory/acp.rs index 032e5c7aa..23162f3c4 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; @@ -31,16 +32,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_for_user(&ctx.user_id, 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 @@ -177,6 +169,30 @@ pub(super) async fn build( 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, @@ -530,6 +546,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; @@ -588,6 +608,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(); From 68e0a78963d66170073ae8e222416fb079fd014a Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 14:29:50 +0800 Subject: [PATCH 092/145] agent: scope acp catalog sync writes --- crates/aionui-ai-agent/src/factory/acp.rs | 1 + .../aionui-ai-agent/src/manager/acp/agent.rs | 6 +- .../src/manager/acp/catalog_forwarder.rs | 5 +- crates/aionui-ai-agent/src/registry.rs | 193 ++++++++++++++++-- .../src/registry_config_option_tests.rs | 10 +- 5 files changed, 195 insertions(+), 20 deletions(-) diff --git a/crates/aionui-ai-agent/src/factory/acp.rs b/crates/aionui-ai-agent/src/factory/acp.rs index 23162f3c4..69ed262d3 100644 --- a/crates/aionui-ai-agent/src/factory/acp.rs +++ b/crates/aionui-ai-agent/src/factory/acp.rs @@ -140,6 +140,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, 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/registry.rs b/crates/aionui-ai-agent/src/registry.rs index a890cf291..5f8ff2b94 100644 --- a/crates/aionui-ai-agent/src/registry.rs +++ b/crates/aionui-ai-agent/src/registry.rs @@ -42,10 +42,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, } @@ -89,8 +91,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" @@ -108,13 +114,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) { @@ -139,12 +147,12 @@ 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(()); }; @@ -172,6 +180,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 { @@ -1225,8 +1245,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, }; @@ -1385,7 +1406,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(); @@ -1395,6 +1418,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 @@ -1544,13 +1596,118 @@ 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""#)) + ); + } + /// Partial updates must leave unrelated columns untouched. /// /// Three consecutive writes target three different columns — each @@ -1567,6 +1724,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})), @@ -1578,6 +1736,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"}])), @@ -1589,6 +1748,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"}])), @@ -1656,6 +1816,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})), @@ -1665,7 +1826,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(); 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!({ From 72912461bdb7e6b8def8cb11409408e58980a659 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 14:35:58 +0800 Subject: [PATCH 093/145] conversation: override runtime extra user scope --- .../src/session_context.rs | 31 +++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/crates/aionui-conversation/src/session_context.rs b/crates/aionui-conversation/src/session_context.rs index 35dcdcf63..b611ddd9d 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); @@ -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 @@ -744,6 +744,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; @@ -1009,6 +1026,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; From e041d8ac64eae1c3d6f099d6dd06cd95d6edff7c Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 14:44:49 +0800 Subject: [PATCH 094/145] config: bind CLI user selector to environment --- crates/aionui-app/src/commands/cmd_config.rs | 24 ++++++++++++++++++- .../src/commands/config_capabilities.rs | 2 +- 2 files changed, 24 insertions(+), 2 deletions(-) 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" } } }, From 3c1bcbec25352a04c19365059419c6ab72597e97 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 14:48:15 +0800 Subject: [PATCH 095/145] realtime: avoid unscoped manager broadcasts --- crates/aionui-realtime/src/manager.rs | 35 ++++++++++++++++++++++++--- 1 file changed, 31 insertions(+), 4 deletions(-) diff --git a/crates/aionui-realtime/src/manager.rs b/crates/aionui-realtime/src/manager.rs index b0b39ccb1..2bb25421a 100644 --- a/crates/aionui-realtime/src/manager.rs +++ b/crates/aionui-realtime/src/manager.rs @@ -267,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); } } @@ -791,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 { @@ -806,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] From 50df71a25cc4ade4a0ecfe83b3da704add3e20a7 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 14:54:51 +0800 Subject: [PATCH 096/145] cron: use scoped conversation lookup for helpers --- crates/aionui-cron/src/executor.rs | 11 +++++++++++ crates/aionui-cron/src/service.rs | 6 ++---- crates/aionui-cron/tests/service_integration.rs | 8 +++++++- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/crates/aionui-cron/src/executor.rs b/crates/aionui-cron/src/executor.rs index 9520442d2..e992f9e25 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -240,6 +240,17 @@ impl JobExecutor { .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, conversation_id: &str, diff --git a/crates/aionui-cron/src/service.rs b/crates/aionui-cron/src/service.rs index 118825de4..270126203 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -236,9 +236,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(), @@ -272,9 +271,8 @@ impl CronService { && (conversation_id.is_empty() || self .executor - .get_conversation_row(conversation_id) + .get_conversation_row_for_user(user_id, conversation_id) .await? - .filter(|row| row.user_id == user_id) .is_none()) { return Err(CronError::Conversation( diff --git a/crates/aionui-cron/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index 6d33ab3be..ca217456a 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -231,10 +231,13 @@ impl StubConvRepo { impl IConversationRepository for StubConvRepo { async fn get( &self, - _user_id: &str, + 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)); } @@ -536,6 +539,9 @@ impl IConversationRepository for StubConvRepo { } }; + 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)) From d25ba1939ad0bb3514628381688ec6fda292f399 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 15:00:38 +0800 Subject: [PATCH 097/145] db: scope diagnostics message samples by user --- .../src/repository/sqlite_diagnostics.rs | 52 +++++++++++-------- 1 file changed, 30 insertions(+), 22 deletions(-) diff --git a/crates/aionui-db/src/repository/sqlite_diagnostics.rs b/crates/aionui-db/src/repository/sqlite_diagnostics.rs index 69133a34a..0beed89d7 100644 --- a/crates/aionui-db/src/repository/sqlite_diagnostics.rs +++ b/crates/aionui-db/src/repository/sqlite_diagnostics.rs @@ -1134,7 +1134,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, @@ -1261,22 +1261,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) @@ -1311,17 +1317,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) From 08284c4c1ba3c65be02b3793c19066fcb5839252 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 15:03:35 +0800 Subject: [PATCH 098/145] db: scope diagnostics detail child queries --- .../src/repository/sqlite_diagnostics.rs | 94 ++++++++++--------- 1 file changed, 52 insertions(+), 42 deletions(-) diff --git a/crates/aionui-db/src/repository/sqlite_diagnostics.rs b/crates/aionui-db/src/repository/sqlite_diagnostics.rs index 0beed89d7..0a3753749 100644 --- a/crates/aionui-db/src/repository/sqlite_diagnostics.rs +++ b/crates/aionui-db/src/repository/sqlite_diagnostics.rs @@ -85,13 +85,13 @@ 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?, + "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?; From 8515f9d478a3d6ac36ee773061667d4a1ed443ba Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 15:16:23 +0800 Subject: [PATCH 099/145] cron: scope delete hook cleanup by user --- crates/aionui-ai-agent/src/task_manager.rs | 2 +- crates/aionui-common/src/constants.rs | 14 +- crates/aionui-common/src/hooks.rs | 6 +- crates/aionui-conversation/src/service.rs | 7 +- .../aionui-conversation/src/service_test.rs | 10 +- crates/aionui-cron/src/executor.rs | 3 +- crates/aionui-cron/src/service.rs | 17 +-- .../aionui-cron/tests/service_integration.rs | 121 +++++++++++++++++- 8 files changed, 146 insertions(+), 34 deletions(-) diff --git a/crates/aionui-ai-agent/src/task_manager.rs b/crates/aionui-ai-agent/src/task_manager.rs index 2cc87ee63..73ffb9449 100644 --- a/crates/aionui-ai-agent/src/task_manager.rs +++ b/crates/aionui-ai-agent/src/task_manager.rs @@ -349,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, 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/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/src/service.rs b/crates/aionui-conversation/src/service.rs index a172e5804..5967cbd28 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); @@ -2139,7 +2140,7 @@ 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(user_id, id).await { diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index 18495fc6b..ee153a49f 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -2178,8 +2178,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}")); } } @@ -2191,7 +2191,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] @@ -2205,8 +2205,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("user_1", 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); } } diff --git a/crates/aionui-cron/src/executor.rs b/crates/aionui-cron/src/executor.rs index e992f9e25..d1e2cacf2 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -253,9 +253,10 @@ impl JobExecutor { 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 diff --git a/crates/aionui-cron/src/service.rs b/crates/aionui-cron/src/service.rs index 270126203..9dd554ede 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -1005,7 +1005,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, @@ -1439,10 +1439,10 @@ impl CronService { } } - 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, @@ -1461,16 +1461,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" @@ -1962,8 +1963,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; } } diff --git a/crates/aionui-cron/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index ca217456a..ee2711c87 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -1081,6 +1081,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, @@ -3090,7 +3136,7 @@ async fn cd1_delete_by_conversation_preserves_jobs() { bc.take_events(); - svc.delete_jobs_by_conversation("conv_cascade").await; + svc.delete_jobs_by_conversation("u1", "conv_cascade").await; assert!(svc.get_job("u1", &job_a.id).await.is_ok()); assert!(svc.get_job("u1", &job_b.id).await.is_ok()); @@ -3117,7 +3163,7 @@ async fn cd2_delete_by_conversation_no_matching_jobs() { .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"); @@ -3139,7 +3185,7 @@ async fn cd3_on_conversation_delete_trait_preserves_jobs() { 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("u1", &job.id).await.is_ok()); @@ -3185,7 +3231,7 @@ async fn cd3b_on_conversation_delete_clears_deleted_workspace_from_jobs() { 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("u1", &job.id).await.is_ok()); let row = cron_repo.get_by_id(&job.id).await.unwrap().unwrap(); @@ -3199,6 +3245,69 @@ 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" + ); + + cron_repo + .insert(&make_cron_row_with_workspace( + "cron_workspace_scope_u1", + "u1", + &conversation_id, + &deleted_workspace, + )) + .await + .unwrap(); + cron_repo + .insert(&make_cron_row_with_workspace( + "cron_workspace_scope_u2", + "u2", + &conversation_id, + &deleted_workspace, + )) + .await + .unwrap(); + + svc.on_conversation_deleted("u1", &conversation_id).await; + + let u1_row = cron_repo.get_by_id("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("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; @@ -3219,7 +3328,7 @@ async fn cd3c_on_conversation_delete_preserves_custom_workspace_on_jobs() { 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 config: CronAgentConfig = serde_json::from_str(row.agent_config.as_deref().unwrap()).unwrap(); @@ -3247,7 +3356,7 @@ async fn cd4_on_conversation_delete_preserves_all_cron_jobs() { bc.take_events(); - svc.on_conversation_deleted("conv_generated_run").await; + svc.on_conversation_deleted("u1", "conv_generated_run").await; assert!( svc.get_job("u1", &new_conversation_job.id).await.is_ok(), From 5c7c711f4d7b885208f7e418f4b1fb15ae804eee Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 15:47:10 +0800 Subject: [PATCH 100/145] assistant: scope rule skill files by user --- crates/aionui-assistant/src/service.rs | 210 +++++++++++++++--- crates/aionui-conversation/src/service.rs | 2 +- .../aionui-conversation/src/service_test.rs | 24 +- crates/aionui-extension/src/classifier.rs | 24 +- crates/aionui-extension/src/skill_routes.rs | 22 +- .../tests/assistant_dispatch_test.rs | 97 +++++--- 6 files changed, 295 insertions(+), 84 deletions(-) diff --git a/crates/aionui-assistant/src/service.rs b/crates/aionui-assistant/src/service.rs index 1a38ca30d..9b803726c 100644 --- a/crates/aionui-assistant/src/service.rs +++ b/crates/aionui-assistant/src/service.rs @@ -909,7 +909,7 @@ impl AssistantService { } 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(id, locale).await?; + let rules_content = self.read_rule_for_user(user_id, id, locale).await?; let projection = self .project_definition(&definition, state.as_ref(), &projections) .await?; @@ -1755,9 +1755,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)) + } } } @@ -1790,8 +1802,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; @@ -1804,18 +1816,39 @@ 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; + } + 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}")))?; @@ -1828,32 +1861,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}")))?; @@ -1865,12 +1926,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)) } } } @@ -1922,13 +1987,29 @@ impl AssistantService { // ----------------------------------------------------------------------- fn user_rules_dir(&self) -> PathBuf { + self.user_rules_dir_for_user(DEFAULT_USER_ID) + } + + fn user_rules_root_dir(&self) -> PathBuf { self.user_data_dir.join("assistant-rules") } + 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_dir(&self) -> PathBuf { + self.user_skills_dir_for_user(DEFAULT_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") } @@ -2165,12 +2246,12 @@ 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_for_user( @@ -2246,38 +2327,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) } @@ -6056,6 +6149,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 @@ -6133,7 +6277,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!( diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index 5967cbd28..c0216473a 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -1439,7 +1439,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 { diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index ee153a49f..d82ebbd7f 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) } } 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/skill_routes.rs b/crates/aionui-extension/src/skill_routes.rs index d762b02bc..a4a900040 100644 --- a/crates/aionui-extension/src/skill_routes.rs +++ b/crates/aionui-extension/src/skill_routes.rs @@ -410,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 = @@ -427,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))); } @@ -449,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?; @@ -468,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 = @@ -485,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))); } @@ -507,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/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())]); } From 95e3a364272911e0d5cb8d59e43bf025fa405d3f Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 15:54:04 +0800 Subject: [PATCH 101/145] cron: reject cross-user conversation targets --- crates/aionui-cron/src/executor.rs | 46 ++++++++++++++++++++++++++++-- 1 file changed, 43 insertions(+), 3 deletions(-) diff --git a/crates/aionui-cron/src/executor.rs b/crates/aionui-cron/src/executor.rs index d1e2cacf2..9a61a6b58 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -537,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, @@ -546,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()) } @@ -1473,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)); @@ -2914,7 +2954,7 @@ mod tests { ) -> 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!({ @@ -2933,7 +2973,7 @@ mod tests { } async fn owner_user_id(&self, _id: &str) -> Result, aionui_db::DbError> { - Ok(Some("cron".into())) + Ok(Some("user1".into())) } async fn create(&self, _row: &aionui_db::models::ConversationRow) -> Result<(), aionui_db::DbError> { From ceac80333ada93461954def0620a86ecd59ee2ba Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 16:18:26 +0800 Subject: [PATCH 102/145] acp: scope session persistence by user --- crates/aionui-ai-agent/src/factory/acp.rs | 4 +- .../src/persistence/acp_session_sync.rs | 97 +++-- .../src/acp_error_recovery.rs | 15 +- crates/aionui-conversation/src/service.rs | 8 +- .../aionui-conversation/src/service_test.rs | 29 ++ .../src/session_context.rs | 92 +++-- .../aionui-db/src/repository/acp_session.rs | 29 ++ .../src/repository/sqlite_acp_session.rs | 369 ++++++++++++++---- 8 files changed, 513 insertions(+), 130 deletions(-) diff --git a/crates/aionui-ai-agent/src/factory/acp.rs b/crates/aionui-ai-agent/src/factory/acp.rs index 69ed262d3..cfc386535 100644 --- a/crates/aionui-ai-agent/src/factory/acp.rs +++ b/crates/aionui-ai-agent/src/factory/acp.rs @@ -165,7 +165,9 @@ 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) } 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 be58f3f16..0cd0216b5 100644 --- a/crates/aionui-ai-agent/src/persistence/acp_session_sync.rs +++ b/crates/aionui-ai-agent/src/persistence/acp_session_sync.rs @@ -43,11 +43,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" @@ -61,12 +66,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" @@ -99,12 +105,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" @@ -118,10 +125,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) { @@ -191,6 +198,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, @@ -205,7 +213,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; } @@ -217,13 +225,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" @@ -236,26 +248,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" @@ -274,6 +298,23 @@ 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 { conversation_id: "conv-1", @@ -299,7 +340,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")); } @@ -310,7 +351,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 @@ -332,7 +373,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() }) @@ -353,7 +394,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::SessionOpened).await.unwrap(); sleep(Duration::from_millis(600)).await; @@ -369,7 +410,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 @@ -396,7 +437,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(), @@ -415,7 +456,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(), @@ -437,7 +478,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::DesiredModeChanged { mode: "plan".into() }) .await @@ -460,7 +501,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(), @@ -484,7 +525,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"), @@ -497,5 +538,7 @@ mod tests { sleep(Duration::from_millis(100)).await; let row = repo.get("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-conversation/src/acp_error_recovery.rs b/crates/aionui-conversation/src/acp_error_recovery.rs index 9e0f3f05f..5b30ca63b 100644 --- a/crates/aionui-conversation/src/acp_error_recovery.rs +++ b/crates/aionui-conversation/src/acp_error_recovery.rs @@ -135,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, ) { @@ -148,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" @@ -167,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, @@ -181,6 +188,7 @@ impl ConversationService { } Ok(false) => { warn!( + user_id, conversation_id, ?previous_model_id, error_code = ?error_code, @@ -190,6 +198,7 @@ impl ConversationService { } Err(err) => { warn!( + user_id, conversation_id, ?previous_model_id, error = %err, @@ -227,7 +236,7 @@ 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(user_id, conversation_id, error_code) .await; diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index c0216473a..1a633b658 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -1318,7 +1318,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}")))?; } @@ -2092,7 +2092,7 @@ impl ConversationService { 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 @@ -2108,7 +2108,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(()) @@ -2153,7 +2153,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" diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index d82ebbd7f..c1c7d78a2 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -994,6 +994,9 @@ impl IAcpSessionRepository for StubAcpSessionRepo { async fn get(&self, conversation_id: &str) -> Result, DbError> { Ok(Some(self.row_for(conversation_id))) } + async fn get_for_user(&self, _user_id: &str, conversation_id: &str) -> Result, DbError> { + self.get(conversation_id).await + } async fn create(&self, params: &CreateAcpSessionParams<'_>) -> Result { self.create_calls.lock().unwrap().push(CreateAcpSessionCall { conversation_id: params.conversation_id.to_owned(), @@ -1017,9 +1020,20 @@ impl IAcpSessionRepository for StubAcpSessionRepo { *self.session_id.lock().unwrap() = Some(session_id.to_owned()); Ok(true) } + async fn update_session_id_for_user( + &self, + _user_id: &str, + conversation_id: &str, + session_id: &str, + ) -> Result { + self.update_session_id(conversation_id, session_id).await + } async fn delete(&self, _conversation_id: &str) -> Result { Ok(false) } + async fn delete_for_user(&self, _user_id: &str, conversation_id: &str) -> Result { + self.delete(conversation_id).await + } async fn load_runtime_state(&self, _conversation_id: &str) -> Result, DbError> { Ok(Some(self.runtime_state.lock().unwrap().clone().unwrap_or( PersistedSessionState { @@ -1028,6 +1042,13 @@ impl IAcpSessionRepository for StubAcpSessionRepo { }, ))) } + async fn load_runtime_state_for_user( + &self, + _user_id: &str, + conversation_id: &str, + ) -> Result, DbError> { + self.load_runtime_state(conversation_id).await + } async fn save_runtime_state( &self, conversation_id: &str, @@ -1041,6 +1062,14 @@ impl IAcpSessionRepository for StubAcpSessionRepo { }); Ok(true) } + async fn save_runtime_state_for_user( + &self, + _user_id: &str, + conversation_id: &str, + params: &SaveRuntimeStateParams<'_>, + ) -> Result { + self.save_runtime_state(conversation_id, params).await + } } fn make_service() -> ( diff --git a/crates/aionui-conversation/src/session_context.rs b/crates/aionui-conversation/src/session_context.rs index b611ddd9d..80a445be6 100644 --- a/crates/aionui-conversation/src/session_context.rs +++ b/crates/aionui-conversation/src/session_context.rs @@ -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()) @@ -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); @@ -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(), } } @@ -802,6 +841,16 @@ 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 { @@ -813,12 +862,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")), @@ -829,15 +879,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); @@ -852,6 +893,16 @@ 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 { @@ -863,7 +914,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")), @@ -874,15 +926,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); @@ -896,6 +939,8 @@ 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 { @@ -905,7 +950,6 @@ mod tests { }) .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); diff --git a/crates/aionui-db/src/repository/acp_session.rs b/crates/aionui-db/src/repository/acp_session.rs index ccf40e24d..ada54f195 100644 --- a/crates/aionui-db/src/repository/acp_session.rs +++ b/crates/aionui-db/src/repository/acp_session.rs @@ -69,6 +69,9 @@ 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 /// surfaces as `DbError::Conflict`. @@ -78,16 +81,34 @@ pub trait IAcpSessionRepository: Send + Sync { /// `session/load` succeeds. Returns `true` when the row existed. async fn update_session_id(&self, conversation_id: &str, session_id: &str) -> Result; + /// User-scoped variant of [`IAcpSessionRepository::update_session_id`]. + async fn update_session_id_for_user( + &self, + user_id: &str, + conversation_id: &str, + session_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; + /// User-scoped variant of [`IAcpSessionRepository::delete`]. + 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>; + /// User-scoped variant of [`IAcpSessionRepository::load_runtime_state`]. + 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. @@ -96,4 +117,12 @@ pub trait IAcpSessionRepository: Send + Sync { conversation_id: &str, params: &SaveRuntimeStateParams<'_>, ) -> Result; + + /// User-scoped variant of [`IAcpSessionRepository::save_runtime_state`]. + 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/sqlite_acp_session.rs b/crates/aionui-db/src/repository/sqlite_acp_session.rs index 7c7980be8..671c8e9fe 100644 --- a/crates/aionui-db/src/repository/sqlite_acp_session.rs +++ b/crates/aionui-db/src/repository/sqlite_acp_session.rs @@ -25,6 +25,83 @@ 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.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"); + } + } + } + + 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(&self, conversation_id: &str) -> Result, DbError> { @@ -35,6 +112,19 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { Ok(row) } + 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(); sqlx::query( @@ -76,6 +166,27 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { Ok(result.rows_affected() > 0) } + 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 delete(&self, conversation_id: &str) -> Result { let result = sqlx::query("DELETE FROM acp_session WHERE conversation_id = ?") .bind(conversation_id) @@ -84,6 +195,19 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { 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(&self, conversation_id: &str) -> Result, DbError> { let raw: Option = sqlx::query_scalar("SELECT session_config FROM acp_session WHERE conversation_id = ?") @@ -95,21 +219,29 @@ 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 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( @@ -134,71 +266,55 @@ 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"); - } - } - } - 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.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"); - } - } + 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) + } + + 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 new_config = - serde_json::to_string(&parsed).map_err(|e| DbError::Init(format!("encode session_config: {e}")))?; + 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,6 +340,35 @@ mod tests { } } + 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 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; @@ -263,6 +408,28 @@ mod tests { assert!(!repo.update_session_id("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("conv-1").await.unwrap().unwrap(); + assert_eq!(fetched.session_id.as_deref(), Some("sess-owner")); + } + #[tokio::test] async fn delete_removes_row() { let (repo, _db) = setup().await; @@ -272,6 +439,20 @@ mod tests { assert!(!repo.delete("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("conv-1").await.unwrap().is_some()); + assert!(repo.delete_for_user("user-1", "conv-1").await.unwrap()); + assert!(repo.get("conv-1").await.unwrap().is_none()); + } + #[tokio::test] async fn load_runtime_state_missing_row() { let (repo, _db) = setup().await; @@ -318,6 +499,52 @@ 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; From 2fbe3c29164147494a7f061f949028d41c537fc6 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 16:24:46 +0800 Subject: [PATCH 103/145] office: scope preview proxy by user --- crates/aionui-app/src/router/routes.rs | 6 ++- crates/aionui-office/src/proxy.rs | 27 +++++++++++- crates/aionui-office/src/routes.rs | 9 ++-- crates/aionui-office/src/watch_manager.rs | 22 ++++++++++ .../aionui-office/tests/proxy_integration.rs | 42 ++++++++++++++++++- 5 files changed, 98 insertions(+), 8 deletions(-) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index 65e887789..679efde02 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -253,8 +253,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 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 fa25be108..d35088a52 100644 --- a/crates/aionui-office/src/routes.rs +++ b/crates/aionui-office/src/routes.rs @@ -270,15 +270,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 { @@ -290,7 +292,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); @@ -307,6 +309,7 @@ async fn office_watch_proxy( async fn proxy_forward( state: OfficeRouterState, + user_id: &str, port: u16, path: &str, doc_type: DocType, @@ -319,7 +322,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/watch_manager.rs b/crates/aionui-office/src/watch_manager.rs index 690b4b3a9..a081e65d5 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, @@ -141,6 +142,7 @@ impl OfficecliWatchManager { self.sessions.insert( key, WatchSession { + user_id: user_id.to_owned(), port, process, file_path: resolved.to_owned(), @@ -232,12 +234,24 @@ impl OfficecliWatchManager { .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() } @@ -744,7 +758,10 @@ mod tests { 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); @@ -836,6 +853,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)); } @@ -860,6 +879,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 // --------------------------------------------------------------------------- From 001c41dd6f487faa993918c6652e2be0b05d54e3 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 16:33:28 +0800 Subject: [PATCH 104/145] runtime: clear user tasks on session revoke --- crates/aionui-ai-agent/src/protocol/error.rs | 1 + crates/aionui-app/src/router/routes.rs | 12 ++++++ crates/aionui-common/src/enums.rs | 3 ++ .../aionui-conversation/src/runtime_state.rs | 41 +++++++++++++++++++ crates/aionui-conversation/src/service.rs | 29 +++++++++++++ .../aionui-conversation/src/service_test.rs | 33 +++++++++++++++ 6 files changed, 119 insertions(+) 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-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index 679efde02..036b12eee 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -168,8 +168,20 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates bootstrap_secret: services.bootstrap_secret.clone(), session_revoked_hook: { let ws_manager = services.ws_manager.clone(); + let conversation_service = states.conversation.service.clone(); Some(Arc::new(move |user_id: &str| { ws_manager.disconnect_user(user_id, "session revoked"); + let user_id = user_id.to_owned(); + let conversation_service = conversation_service.clone(); + tokio::spawn(async move { + 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" + ); + } + }); })) }, local: services.local, 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-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 1a633b658..73fb886e6 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -601,6 +601,35 @@ impl ConversationService { 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, diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index c1c7d78a2..207a9ddbd 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -3702,6 +3702,39 @@ async fn active_count_for_user_counts_only_owned_active_tasks() { 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()); From 39715d245f7a400f45a751f7099eb70c98cc2c67 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 16:44:26 +0800 Subject: [PATCH 105/145] app: clear file watches on session revoke --- crates/aionui-app/src/router/routes.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index 036b12eee..fd1d0f01d 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -169,10 +169,12 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates session_revoked_hook: { let ws_manager = services.ws_manager.clone(); let conversation_service = states.conversation.service.clone(); + let file_watch_service = states.file.watch_service.clone(); Some(Arc::new(move |user_id: &str| { ws_manager.disconnect_user(user_id, "session revoked"); let user_id = user_id.to_owned(); let conversation_service = conversation_service.clone(); + let file_watch_service = file_watch_service.clone(); tokio::spawn(async move { if let Err(err) = conversation_service.terminate_runtime_for_user(&user_id).await { tracing::warn!( @@ -181,6 +183,20 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates "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" + ); + } }); })) }, From a8160c994bf84146ba15df315e6a99f8e38700a9 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 16:49:03 +0800 Subject: [PATCH 106/145] team: clear user sessions on revoke --- crates/aionui-app/src/router/routes.rs | 9 +++++++ crates/aionui-team/src/service.rs | 37 ++++++++++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index fd1d0f01d..ad1b77ba3 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -169,9 +169,18 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates 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 file_watch_service = states.file.watch_service.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 file_watch_service = file_watch_service.clone(); diff --git a/crates/aionui-team/src/service.rs b/crates/aionui-team/src/service.rs index 56baa58f1..9a8ce3526 100644 --- a/crates/aionui-team/src/service.rs +++ b/crates/aionui-team/src/service.rs @@ -1817,6 +1817,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(); @@ -2480,6 +2494,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) = From 6307d81ffa2fef73ea6dc5c97be698d0f7d61cda Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 16:59:18 +0800 Subject: [PATCH 107/145] channel: stop user runtimes on session revoke --- crates/aionui-app/src/router/routes.rs | 12 +++++++ crates/aionui-channel/src/manager.rs | 43 ++++++++++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index ad1b77ba3..261feb7d1 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -170,6 +170,8 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates 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(); Some(Arc::new(move |user_id: &str| { ws_manager.disconnect_user(user_id, "session revoked"); @@ -183,8 +185,18 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates } 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(); tokio::spawn(async move { + channel_manager.shutdown_for_user(&user_id).await; + 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, diff --git a/crates/aionui-channel/src/manager.rs b/crates/aionui-channel/src/manager.rs index 8608c98d9..cea27ea03 100644 --- a/crates/aionui-channel/src/manager.rs +++ b/crates/aionui-channel/src/manager.rs @@ -330,6 +330,29 @@ impl ChannelManager { 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 prefix = format!("{owner_user_id}:"); + let keys: Vec = self + .plugins + .iter() + .filter(|entry| entry.key().starts_with(&prefix)) + .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() @@ -1457,6 +1480,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] From 1e079b809f6f28c263d5a9ea5defb85071ebdec0 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 17:07:18 +0800 Subject: [PATCH 108/145] db: validate user scope migration copies --- .../aionui-db/migrations/026_user_scope.sql | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) diff --git a/crates/aionui-db/migrations/026_user_scope.sql b/crates/aionui-db/migrations/026_user_scope.sql index 710b3bc80..a9a02df6e 100644 --- a/crates/aionui-db/migrations/026_user_scope.sql +++ b/crates/aionui-db/migrations/026_user_scope.sql @@ -133,6 +133,10 @@ 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), @@ -163,6 +167,13 @@ SELECT 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); @@ -190,6 +201,13 @@ SELECT 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; @@ -214,6 +232,13 @@ SELECT 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; @@ -229,6 +254,13 @@ 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; @@ -300,6 +332,13 @@ SELECT 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 @@ -339,6 +378,13 @@ 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 @@ -421,6 +467,13 @@ SELECT 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; @@ -465,6 +518,13 @@ SELECT 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 @@ -501,6 +561,13 @@ SELECT 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; @@ -527,6 +594,13 @@ SELECT 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 @@ -567,6 +641,13 @@ SELECT 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 @@ -593,11 +674,57 @@ SELECT 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), @@ -620,9 +747,18 @@ SELECT 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); +DROP TABLE user_scope_rebuild_checks; + PRAGMA foreign_keys = ON; From dc69be5c115af2f4e034fea55ae95fe144daa779 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 17:15:46 +0800 Subject: [PATCH 109/145] office: scope preview snapshots by user --- crates/aionui-office/src/routes.rs | 15 +++--- crates/aionui-office/src/snapshot.rs | 75 ++++++++++++++++++++++++++-- 2 files changed, 79 insertions(+), 11 deletions(-) diff --git a/crates/aionui-office/src/routes.rs b/crates/aionui-office/src/routes.rs index d35088a52..60ac89154 100644 --- a/crates/aionui-office/src/routes.rs +++ b/crates/aionui-office/src/routes.rs @@ -177,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))) } 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, From e31dbc26581c4b4865ca3e730ab8d2fe8003c61c Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 17:25:30 +0800 Subject: [PATCH 110/145] agent: avoid cross-user registry cache overlay --- crates/aionui-ai-agent/src/registry.rs | 54 ++++++++++++++++--- crates/aionui-db/src/models/agent_metadata.rs | 1 + .../src/repository/sqlite_agent_metadata.rs | 1 + .../aionui-team/src/service/spawn_support.rs | 1 + .../tests/session_service_integration.rs | 1 + 5 files changed, 50 insertions(+), 8 deletions(-) diff --git a/crates/aionui-ai-agent/src/registry.rs b/crates/aionui-ai-agent/src/registry.rs index 5f8ff2b94..393ed668c 100644 --- a/crates/aionui-ai-agent/src/registry.rs +++ b/crates/aionui-ai-agent/src/registry.rs @@ -156,7 +156,8 @@ impl AgentRegistry { 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() @@ -404,10 +405,14 @@ impl AgentRegistry { let cached_reasons = self.unavailable_reasons.read().await.clone(); let mut rows: Vec = rows .into_iter() - .filter_map(|row| decode_row(row, AvailabilityProjection::Cached)) - .map(|(mut meta, mut reason)| { - let cached_meta = cached_rows.get(&meta.id); - let cached_reason = cached_reasons.get(&meta.id); + .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) { @@ -449,11 +454,23 @@ impl AgentRegistry { .get_for_user(user_id, id) .await .map_err(|e| AgentError::internal(format!("load agent_metadata '{id}' for user: {e}")))?; - let Some((mut meta, mut reason)) = row.and_then(|row| decode_row(row, AvailabilityProjection::Cached)) else { + let Some(row) = row else { return Ok(None); }; - let cached_meta = self.by_id.read().await.get(&meta.id).cloned(); - let cached_reason = self.unavailable_reasons.read().await.get(&meta.id).cloned(); + 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) { @@ -1706,6 +1723,24 @@ mod tests { .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. @@ -1862,6 +1897,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, @@ -1908,6 +1944,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, @@ -1963,6 +2000,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-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/repository/sqlite_agent_metadata.rs b/crates/aionui-db/src/repository/sqlite_agent_metadata.rs index 5a20814fa..e09baea02 100644 --- a/crates/aionui-db/src/repository/sqlite_agent_metadata.rs +++ b/crates/aionui-db/src/repository/sqlite_agent_metadata.rs @@ -178,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, diff --git a/crates/aionui-team/src/service/spawn_support.rs b/crates/aionui-team/src/service/spawn_support.rs index e72a629ac..6e8c41c5f 100644 --- a/crates/aionui-team/src/service/spawn_support.rs +++ b/crates/aionui-team/src/service/spawn_support.rs @@ -236,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, diff --git a/crates/aionui-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index 5848aaaa0..ed1dc0a07 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -1961,6 +1961,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, From fa196c1cea6ae46609ff8b78686c8a1773349eba Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 17:31:53 +0800 Subject: [PATCH 111/145] db: test channel session owner migration guard --- .../aionui-db/tests/user_scope_migration.rs | 60 ++++++++++++++++++- 1 file changed, 59 insertions(+), 1 deletion(-) diff --git a/crates/aionui-db/tests/user_scope_migration.rs b/crates/aionui-db/tests/user_scope_migration.rs index 8660781c0..5290e5cae 100644 --- a/crates/aionui-db/tests/user_scope_migration.rs +++ b/crates/aionui-db/tests/user_scope_migration.rs @@ -41,6 +41,10 @@ async fn run_migrations_through(pool: &sqlx::SqlitePool, max_version: i64) { } 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 @@ -54,7 +58,7 @@ async fn run_migration(pool: &sqlx::SqlitePool, version: i64) { locking: true, no_tx: false, }; - migrator.run(pool).await.unwrap(); + migrator.run(pool).await } #[tokio::test] @@ -155,3 +159,57 @@ async fn migration_026_keeps_new_conversation_cron_jobs_unanchored_until_run() { assert_eq!(row.get::("user_id"), "system_default_user"); assert_eq!(row.get::("conversation_id"), ""); } + +#[tokio::test] +async fn migration_026_rejects_channel_session_cross_user_conversation() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 25).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, 26).await.unwrap_err(); + assert!( + err.to_string().contains("CHECK constraint failed"), + "unexpected migration error: {err}" + ); +} From 6a8782beaa720c0fbaa50ee88773c93ceade9f38 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 17:44:25 +0800 Subject: [PATCH 112/145] db: reject cross-user assistant upsert takeover --- .../src/repository/sqlite_assistant.rs | 31 +++++++++++++++++-- 1 file changed, 28 insertions(+), 3 deletions(-) diff --git a/crates/aionui-db/src/repository/sqlite_assistant.rs b/crates/aionui-db/src/repository/sqlite_assistant.rs index 9c7d7379e..4904b0d58 100644 --- a/crates/aionui-db/src/repository/sqlite_assistant.rs +++ b/crates/aionui-db/src/repository/sqlite_assistant.rs @@ -188,14 +188,13 @@ impl IAssistantRepository for SqliteAssistantRepository { ) -> Result { let now = now_ms(); - sqlx::query( + let result = sqlx::query( "INSERT INTO assistants \ (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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) \ ON CONFLICT(id) DO UPDATE SET \ - user_id = excluded.user_id, \ name = excluded.name, \ description = excluded.description, \ avatar = excluded.avatar, \ @@ -207,7 +206,8 @@ 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) @@ -227,6 +227,13 @@ 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_for_user(user_id, params.id) .await? @@ -1272,6 +1279,24 @@ mod tests { 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; From 8d3362fef153cea046952f7b192158f56b01ab29 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 19:26:48 +0800 Subject: [PATCH 113/145] db: renumber user scope migration after main --- .../{026_user_scope.sql => 028_user_scope.sql} | 2 +- crates/aionui-db/tests/provider_repository.rs | 1 + crates/aionui-db/tests/user_scope_migration.rs | 18 +++++++++--------- 3 files changed, 11 insertions(+), 10 deletions(-) rename crates/aionui-db/migrations/{026_user_scope.sql => 028_user_scope.sql} (99%) diff --git a/crates/aionui-db/migrations/026_user_scope.sql b/crates/aionui-db/migrations/028_user_scope.sql similarity index 99% rename from crates/aionui-db/migrations/026_user_scope.sql rename to crates/aionui-db/migrations/028_user_scope.sql index a9a02df6e..afc32bb32 100644 --- a/crates/aionui-db/migrations/026_user_scope.sql +++ b/crates/aionui-db/migrations/028_user_scope.sql @@ -1,4 +1,4 @@ --- Migration 026: add user scope for local, aggregate, and configuration data. +-- Migration 028: add user scope for local, aggregate, and configuration data. PRAGMA foreign_keys = OFF; diff --git a/crates/aionui-db/tests/provider_repository.rs b/crates/aionui-db/tests/provider_repository.rs index ca3dadfc4..37ae86c81 100644 --- a/crates/aionui-db/tests/provider_repository.rs +++ b/crates/aionui-db/tests/provider_repository.rs @@ -184,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"}}"#), diff --git a/crates/aionui-db/tests/user_scope_migration.rs b/crates/aionui-db/tests/user_scope_migration.rs index 5290e5cae..917646929 100644 --- a/crates/aionui-db/tests/user_scope_migration.rs +++ b/crates/aionui-db/tests/user_scope_migration.rs @@ -62,7 +62,7 @@ async fn run_migration_result(pool: &sqlx::SqlitePool, version: i64) -> Result<( } #[tokio::test] -async fn migration_026_adds_core_user_projection_columns() { +async fn migration_028_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"] { @@ -80,7 +80,7 @@ async fn migration_026_adds_core_user_projection_columns() { } #[tokio::test] -async fn migration_026_adds_user_scope_to_independent_roots() { +async fn migration_028_adds_user_scope_to_independent_roots() { let db = init_database_memory().await.unwrap(); for (table, column) in [ ("cron_jobs", "user_id"), @@ -103,13 +103,13 @@ async fn migration_026_adds_user_scope_to_independent_roots() { } #[tokio::test] -async fn migration_026_migrates_cron_skills_as_global_rows() { +async fn migration_028_migrates_cron_skills_as_global_rows() { let pool = SqlitePoolOptions::new() .max_connections(1) .connect("sqlite::memory:") .await .unwrap(); - run_migrations_through(&pool, 25).await; + run_migrations_through(&pool, 26).await; sqlx::query( "INSERT INTO skills (id, name, description, path, source, enabled, created_at, updated_at) @@ -119,7 +119,7 @@ async fn migration_026_migrates_cron_skills_as_global_rows() { .await .unwrap(); - run_migration(&pool, 26).await; + run_migration(&pool, 28).await; let row = sqlx::query("SELECT user_id, source FROM skills WHERE id = 'legacy-cron-skill'") .fetch_one(&pool) @@ -130,7 +130,7 @@ async fn migration_026_migrates_cron_skills_as_global_rows() { } #[tokio::test] -async fn migration_026_keeps_new_conversation_cron_jobs_unanchored_until_run() { +async fn migration_028_keeps_new_conversation_cron_jobs_unanchored_until_run() { let db = init_database_memory().await.unwrap(); let now = aionui_common::now_ms(); @@ -161,13 +161,13 @@ async fn migration_026_keeps_new_conversation_cron_jobs_unanchored_until_run() { } #[tokio::test] -async fn migration_026_rejects_channel_session_cross_user_conversation() { +async fn migration_028_rejects_channel_session_cross_user_conversation() { let pool = SqlitePoolOptions::new() .max_connections(1) .connect("sqlite::memory:") .await .unwrap(); - run_migrations_through(&pool, 25).await; + run_migrations_through(&pool, 26).await; sqlx::query( "INSERT INTO users (id, username, password_hash, created_at, updated_at) @@ -207,7 +207,7 @@ async fn migration_026_rejects_channel_session_cross_user_conversation() { .await .unwrap(); - let err = run_migration_result(&pool, 26).await.unwrap_err(); + let err = run_migration_result(&pool, 28).await.unwrap_err(); assert!( err.to_string().contains("CHECK constraint failed"), "unexpected migration error: {err}" From 5a18dab138e1606a5b3806c8be4c2a6136ab6adb Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 20:28:30 +0800 Subject: [PATCH 114/145] db: reject cross-user scoped upsert takeovers --- .../src/repository/sqlite_assistant.rs | 68 +++++++++++- .../src/repository/sqlite_conversation.rs | 34 +++++- .../tests/conversation_repository.rs | 102 ++++++++++++++++-- 3 files changed, 186 insertions(+), 18 deletions(-) diff --git a/crates/aionui-db/src/repository/sqlite_assistant.rs b/crates/aionui-db/src/repository/sqlite_assistant.rs index 4904b0d58..fd59bf401 100644 --- a/crates/aionui-db/src/repository/sqlite_assistant.rs +++ b/crates/aionui-db/src/repository/sqlite_assistant.rs @@ -296,7 +296,7 @@ impl SqliteAssistantDefinitionRepository { ) -> Result { let now = now_ms(); - sqlx::query( + 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, @@ -310,7 +310,6 @@ impl SqliteAssistantDefinitionRepository { created_at, updated_at, deleted_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL) ON CONFLICT(id) DO UPDATE SET - user_id = excluded.user_id, assistant_id = excluded.assistant_id, source = excluded.source, owner_type = excluded.owner_type, @@ -339,7 +338,8 @@ impl SqliteAssistantDefinitionRepository { default_mcps_mode = excluded.default_mcps_mode, default_mcp_ids = excluded.default_mcp_ids, updated_at = excluded.updated_at, - deleted_at = NULL", + deleted_at = NULL + WHERE assistant_definitions.user_id IS excluded.user_id", ) .bind(params.id) .bind(user_id) @@ -375,6 +375,13 @@ impl SqliteAssistantDefinitionRepository { .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?, @@ -1509,6 +1516,61 @@ mod tests { 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; diff --git a/crates/aionui-db/src/repository/sqlite_conversation.rs b/crates/aionui-db/src/repository/sqlite_conversation.rs index 1906314c1..cc94a81f7 100644 --- a/crates/aionui-db/src/repository/sqlite_conversation.rs +++ b/crates/aionui-db/src/repository/sqlite_conversation.rs @@ -68,7 +68,7 @@ impl SqliteConversationRepository { async fn upsert_message_once(&self, user_id: &str, message: &MessageRow) -> Result<(), DbError> { self.ensure_conversation_for_user(user_id, &message.conversation_id) .await?; - sqlx::query( + let result = sqlx::query( "INSERT INTO messages \ (id, conversation_id, msg_id, type, content, position, \ status, hidden, created_at) \ @@ -96,7 +96,12 @@ impl SqliteConversationRepository { END, \ position = COALESCE(messages.position, excluded.position), \ hidden = excluded.hidden, \ - created_at = MIN(messages.created_at, excluded.created_at)", + 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) @@ -107,9 +112,17 @@ impl SqliteConversationRepository { .bind(&message.status) .bind(message.hidden) .bind(message.created_at) + .bind(user_id) .execute(&self.pool) .await?; + if result.rows_affected() == 0 { + return Err(DbError::Conflict(format!( + "Message with id '{}' already exists outside the requested conversation", + message.id + ))); + } + Ok(()) } @@ -995,7 +1008,7 @@ impl IConversationRepository for SqliteConversationRepository { ) -> Result { self.ensure_conversation_for_user(user_id, &artifact.conversation_id) .await?; - sqlx::query( + let result = sqlx::query( "INSERT INTO conversation_artifacts \ (id, conversation_id, cron_job_id, kind, status, payload, created_at, updated_at) \ VALUES (?, ?, ?, ?, ?, ?, ?, ?) \ @@ -1005,7 +1018,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) @@ -1015,9 +1033,17 @@ impl IConversationRepository for SqliteConversationRepository { .bind(&artifact.payload) .bind(artifact.created_at) .bind(artifact.updated_at) + .bind(user_id) .execute(&self.pool) .await?; + 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))) diff --git a/crates/aionui-db/tests/conversation_repository.rs b/crates/aionui-db/tests/conversation_repository.rs index c9c57f3cd..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 { @@ -1172,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(); @@ -1202,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()); +} From d628eb87c5a8737d704b80cc1ccdb7e774ec3681 Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 20:46:53 +0800 Subject: [PATCH 115/145] extension: isolate user skill import storage --- crates/aionui-extension/src/skill_service.rs | 48 +++++++- .../tests/skill_integration_test.rs | 103 +++++++++++++++++- 2 files changed, 142 insertions(+), 9 deletions(-) diff --git a/crates/aionui-extension/src/skill_service.rs b/crates/aionui-extension/src/skill_service.rs index edcd6713f..89b0d92a0 100644 --- a/crates/aionui-extension/src/skill_service.rs +++ b/crates/aionui-extension/src/skill_service.rs @@ -533,7 +533,7 @@ pub async fn import_skill_with_repo_for_user( user_id: &str, skill_path: &Path, ) -> Result { - let copied = copy_skill_into_user_dir(paths, skill_path).await?; + 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_for_user( @@ -570,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 { @@ -1050,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() )) @@ -1742,10 +1762,17 @@ async fn sync_disk_user_skills_into_repo_for_user( 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, @@ -1782,6 +1809,15 @@ async fn sync_disk_user_skills_into_repo_for_user( 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); diff --git a/crates/aionui-extension/tests/skill_integration_test.rs b/crates/aionui-extension/tests/skill_integration_test.rs index 515d04ce3..2b3ecd105 100644 --- a/crates/aionui-extension/tests/skill_integration_test.rs +++ b/crates/aionui-extension/tests/skill_integration_test.rs @@ -6,12 +6,14 @@ use std::path::Path; +use aionui_db::{ISkillRepository, IUserRepository, SqliteSkillRepository, SqliteUserRepository, 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 +49,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 +192,83 @@ 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)); +} + /// SM-5: Export skill (symlink). #[tokio::test] async fn sm5_export_skill_symlink() { From 401635e84168be689be27d55fda0510d616123fa Mon Sep 17 00:00:00 2001 From: zk <> Date: Tue, 21 Jul 2026 21:12:51 +0800 Subject: [PATCH 116/145] auth: keep external session token cookie-only --- crates/aionui-api-types/src/auth.rs | 1 - crates/aionui-auth/src/routes.rs | 8 ++++---- crates/aionui-auth/src/service.rs | 15 +++++++++++---- crates/aionui-auth/tests/route_tests.rs | 25 +++++++++++++++++++++---- 4 files changed, 36 insertions(+), 13 deletions(-) diff --git a/crates/aionui-api-types/src/auth.rs b/crates/aionui-api-types/src/auth.rs index 889d60acb..4f6b31b06 100644 --- a/crates/aionui-api-types/src/auth.rs +++ b/crates/aionui-api-types/src/auth.rs @@ -126,7 +126,6 @@ pub struct EnsureExternalSessionRequest { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct EnsureExternalSessionResponse { pub user: PublicUser, - pub token: String, pub session_generation: i64, } diff --git a/crates/aionui-auth/src/routes.rs b/crates/aionui-auth/src/routes.rs index 118a967e4..6026dbfb5 100644 --- a/crates/aionui-auth/src/routes.rs +++ b/crates/aionui-auth/src/routes.rs @@ -354,16 +354,16 @@ async fn create_external_session_handler( 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 + let exchange = service .create_external_session(req) .await .map_err(provision_error_to_api_error)?; tracing::info!( - user_id = %response.user.id, + user_id = %exchange.response.user.id, "external core session exchange succeeded" ); - let cookie = state.cookie_config.build_session_cookie(&response.token); - Ok(([(header::SET_COOKIE, cookie)], Json(ApiResponse::ok(response))).into_response()) + let cookie = state.cookie_config.build_session_cookie(&exchange.token); + Ok(([(header::SET_COOKIE, cookie)], Json(ApiResponse::ok(exchange.response))).into_response()) } // --------------------------------------------------------------------------- diff --git a/crates/aionui-auth/src/service.rs b/crates/aionui-auth/src/service.rs index 59b3614e9..e7406e41c 100644 --- a/crates/aionui-auth/src/service.rs +++ b/crates/aionui-auth/src/service.rs @@ -28,6 +28,11 @@ pub struct AuthProvisionService { 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 } @@ -62,7 +67,7 @@ impl AuthProvisionService { pub async fn create_external_session( &self, request: EnsureExternalSessionRequest, - ) -> Result { + ) -> Result { let user_type = map_external_user_type(request.user_type)?; let user = self .user_repo @@ -80,10 +85,12 @@ impl AuthProvisionService { .sign_with_session_generation(&user.id, &username, user.session_generation)?; self.user_repo.update_last_login(&user.id).await?; - Ok(EnsureExternalSessionResponse { - user: PublicUser { id: user.id, username }, + Ok(ExternalSessionExchange { + response: EnsureExternalSessionResponse { + user: PublicUser { id: user.id, username }, + session_generation: user.session_generation, + }, token, - session_generation: user.session_generation, }) } diff --git a/crates/aionui-auth/tests/route_tests.rs b/crates/aionui-auth/tests/route_tests.rs index 1c088d4f7..9d0b9d46a 100644 --- a/crates/aionui-auth/tests/route_tests.rs +++ b/crates/aionui-auth/tests/route_tests.rs @@ -175,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( @@ -885,7 +900,7 @@ async fn external_user_provision_is_idempotent() { } #[tokio::test] -async fn external_session_exchange_returns_cookie_and_token() { +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 @@ -910,11 +925,12 @@ async fn external_session_exchange_returns_cookie_and_token() { .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; - let token = json["data"]["token"].as_str().unwrap(); + 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(); + 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"); @@ -973,9 +989,10 @@ async fn external_session_revoke_invalidates_existing_token() { .await .unwrap(); assert_eq!(session.status(), StatusCode::OK); + let token = extract_session_token(&session).unwrap(); let session_json = body_json(session).await; - let token = session_json["data"]["token"].as_str().unwrap().to_owned(); assert_eq!(session_json["data"]["session_generation"], 0); + assert!(session_json["data"].get("token").is_none()); let revoke = app .clone() From 385fbd638fb6bbbb96cf27998db72eb45fbc6150 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 09:20:48 +0800 Subject: [PATCH 117/145] conversation: align user-scoped snapshot calls after rebase --- crates/aionui-conversation/src/service.rs | 2 +- .../aionui-conversation/src/service_test.rs | 43 +++++++++++-------- 2 files changed, 25 insertions(+), 20 deletions(-) diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index 73fb886e6..974b57361 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -3345,7 +3345,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!( diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index 207a9ddbd..5f8406ad5 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -735,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, @@ -835,6 +836,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, @@ -7777,25 +7779,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 From b6bf9d8557594c37cf42b24ffcf7d45da665b407 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 10:18:29 +0800 Subject: [PATCH 118/145] auth: redact internal user route secrets --- crates/aionui-auth/src/routes.rs | 59 ++++++++++++++++++++----- crates/aionui-auth/tests/route_tests.rs | 4 ++ 2 files changed, 51 insertions(+), 12 deletions(-) diff --git a/crates/aionui-auth/src/routes.rs b/crates/aionui-auth/src/routes.rs index 6026dbfb5..57bc03c21 100644 --- a/crates/aionui-auth/src/routes.rs +++ b/crates/aionui-auth/src/routes.rs @@ -10,7 +10,7 @@ use axum::middleware::from_fn_with_state; use axum::response::{Html, IntoResponse, Response}; use axum::routing::{get, post, put}; use axum::{Extension, Router}; -use serde::Deserialize; +use serde::{Deserialize, Serialize}; use aionui_api_types::{ ApiResponse, AuthStatusResponse, ChangePasswordRequest, EnsureExternalSessionRequest, EnsureExternalUserRequest, @@ -21,7 +21,7 @@ use aionui_api_types::{ }; 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; @@ -105,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(()); @@ -531,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 @@ -578,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( diff --git a/crates/aionui-auth/tests/route_tests.rs b/crates/aionui-auth/tests/route_tests.rs index 9d0b9d46a..9c2446287 100644 --- a/crates/aionui-auth/tests/route_tests.rs +++ b/crates/aionui-auth/tests/route_tests.rs @@ -1074,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() @@ -1082,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 From 04a45515f7248efa6ed7412c3822fedb21af860e Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 10:36:55 +0800 Subject: [PATCH 119/145] db: remove unscoped acp session trait methods --- .../src/persistence/acp_session_sync.rs | 59 ++++- .../aionui-conversation/src/service_test.rs | 41 +--- .../tests/conversation_crud.rs | 7 +- crates/aionui-cron/src/executor.rs | 14 +- .../aionui-db/src/repository/acp_session.rs | 45 ++-- .../src/repository/sqlite_acp_session.rs | 209 +++++++++--------- 6 files changed, 196 insertions(+), 179 deletions(-) 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 0cd0216b5..96e49ca12 100644 --- a/crates/aionui-ai-agent/src/persistence/acp_session_sync.rs +++ b/crates/aionui-ai-agent/src/persistence/acp_session_sync.rs @@ -330,7 +330,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")), @@ -358,11 +359,19 @@ mod tests { .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")); } @@ -383,7 +392,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")); } @@ -399,7 +412,11 @@ mod tests { 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()); } @@ -418,7 +435,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"), @@ -446,7 +467,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")); } @@ -465,7 +490,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*", @@ -485,7 +514,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_mode_id.is_none(), "DesiredModeChanged is reconcile/UI-only; persistence only follows Observed*", @@ -510,7 +543,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); @@ -536,7 +573,7 @@ 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-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index 5f8406ad5..c4da93d1a 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -993,11 +993,8 @@ struct CreateAcpSessionCall { #[async_trait::async_trait] impl IAcpSessionRepository for StubAcpSessionRepo { - async fn get(&self, conversation_id: &str) -> Result, DbError> { - Ok(Some(self.row_for(conversation_id))) - } async fn get_for_user(&self, _user_id: &str, conversation_id: &str) -> Result, DbError> { - self.get(conversation_id).await + Ok(Some(self.row_for(conversation_id))) } async fn create(&self, params: &CreateAcpSessionParams<'_>) -> Result { self.create_calls.lock().unwrap().push(CreateAcpSessionCall { @@ -1018,25 +1015,23 @@ impl IAcpSessionRepository for StubAcpSessionRepo { suspended_at: None, }) } - async fn update_session_id(&self, _conversation_id: &str, session_id: &str) -> Result { - *self.session_id.lock().unwrap() = Some(session_id.to_owned()); - Ok(true) - } async fn update_session_id_for_user( &self, _user_id: &str, - conversation_id: &str, + _conversation_id: &str, session_id: &str, ) -> Result { - self.update_session_id(conversation_id, session_id).await + *self.session_id.lock().unwrap() = Some(session_id.to_owned()); + 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 delete_for_user(&self, _user_id: &str, conversation_id: &str) -> Result { - self.delete(conversation_id).await - } - 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()), @@ -1044,16 +1039,10 @@ impl IAcpSessionRepository for StubAcpSessionRepo { }, ))) } - async fn load_runtime_state_for_user( + async fn save_runtime_state_for_user( &self, _user_id: &str, conversation_id: &str, - ) -> Result, DbError> { - self.load_runtime_state(conversation_id).await - } - async fn save_runtime_state( - &self, - conversation_id: &str, params: &SaveRuntimeStateParams<'_>, ) -> Result { self.runtime_state_saves.lock().unwrap().push(RuntimeStateSaveCall { @@ -1064,14 +1053,6 @@ impl IAcpSessionRepository for StubAcpSessionRepo { }); Ok(true) } - async fn save_runtime_state_for_user( - &self, - _user_id: &str, - conversation_id: &str, - params: &SaveRuntimeStateParams<'_>, - ) -> Result { - self.save_runtime_state(conversation_id, params).await - } } fn make_service() -> ( diff --git a/crates/aionui-conversation/tests/conversation_crud.rs b/crates/aionui-conversation/tests/conversation_crud.rs index 0a26e3ff5..b9e5dd97f 100644 --- a/crates/aionui-conversation/tests/conversation_crud.rs +++ b/crates/aionui-conversation/tests/conversation_crud.rs @@ -881,7 +881,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 +927,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-cron/src/executor.rs b/crates/aionui-cron/src/executor.rs index 9a61a6b58..408e7625a 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -3411,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) @@ -3423,24 +3424,27 @@ 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 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 { diff --git a/crates/aionui-db/src/repository/acp_session.rs b/crates/aionui-db/src/repository/acp_session.rs index ada54f195..563595ec3 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,7 +19,7 @@ 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 conversation_id: &'a str, @@ -43,7 +43,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,9 +66,6 @@ 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>; @@ -78,10 +75,8 @@ 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; - - /// User-scoped variant of [`IAcpSessionRepository::update_session_id`]. + /// `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, @@ -89,36 +84,26 @@ pub trait IAcpSessionRepository: Send + Sync { session_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; - - /// User-scoped variant of [`IAcpSessionRepository::delete`]. + /// 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>; - - /// User-scoped variant of [`IAcpSessionRepository::load_runtime_state`]. + /// 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( - &self, - conversation_id: &str, - params: &SaveRuntimeStateParams<'_>, - ) -> Result; - - /// User-scoped variant of [`IAcpSessionRepository::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, diff --git a/crates/aionui-db/src/repository/sqlite_acp_session.rs b/crates/aionui-db/src/repository/sqlite_acp_session.rs index 671c8e9fe..a165d26b5 100644 --- a/crates/aionui-db/src/repository/sqlite_acp_session.rs +++ b/crates/aionui-db/src/repository/sqlite_acp_session.rs @@ -19,6 +19,90 @@ impl SqliteAcpSessionRepository { pub fn new(pool: SqlitePool) -> Self { Self { pool } } + + 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) + .await?; + Ok(row) + } + + #[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) + .bind(now) + .bind(conversation_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + #[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) + .await?; + Ok(result.rows_affected() > 0) + } + + #[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) + .fetch_optional(&self.pool) + .await?; + + let Some(raw) = raw else { + return Ok(None); + }; + + Ok(Some(decode_runtime_state(&raw)?)) + } + + #[cfg(test)] + async fn save_runtime_state_unscoped( + &self, + conversation_id: &str, + params: &SaveRuntimeStateParams<'_>, + ) -> Result { + if params.is_empty() { + return Ok(true); + } + + // Read-modify-write. The service layer serialises writes per + // conversation_id through a single consumer task, so a naive + // RMW is race-free for our callers. + let raw: Option = + sqlx::query_scalar("SELECT session_config FROM acp_session WHERE conversation_id = ?") + .bind(conversation_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 = ?", + ) + .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 { @@ -104,14 +188,6 @@ fn merge_runtime_state(raw: &str, params: &SaveRuntimeStateParams<'_>) -> Result #[async_trait::async_trait] impl IAcpSessionRepository for SqliteAcpSessionRepository { - async fn get(&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) - .await?; - Ok(row) - } - 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 \ @@ -147,7 +223,7 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { _ => DbError::Query(e), })?; - self.get(params.conversation_id).await?.ok_or_else(|| { + self.get_unscoped(params.conversation_id).await?.ok_or_else(|| { DbError::Init(format!( "create did not produce acp_session row for '{}'", params.conversation_id @@ -155,17 +231,6 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { }) } - async fn update_session_id(&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) - .bind(now) - .bind(conversation_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - async fn update_session_id_for_user( &self, user_id: &str, @@ -187,14 +252,6 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { Ok(result.rows_affected() > 0) } - async fn delete(&self, conversation_id: &str) -> Result { - let result = sqlx::query("DELETE FROM acp_session WHERE conversation_id = ?") - .bind(conversation_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 \ @@ -208,20 +265,6 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { Ok(result.rows_affected() > 0) } - async fn load_runtime_state(&self, conversation_id: &str) -> Result, DbError> { - let raw: Option = - sqlx::query_scalar("SELECT session_config FROM acp_session WHERE conversation_id = ?") - .bind(conversation_id) - .fetch_optional(&self.pool) - .await?; - - let Some(raw) = raw else { - return Ok(None); - }; - - Ok(Some(decode_runtime_state(&raw)?)) - } - async fn load_runtime_state_for_user( &self, user_id: &str, @@ -244,42 +287,6 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { Ok(Some(decode_runtime_state(&raw)?)) } - async fn save_runtime_state( - &self, - conversation_id: &str, - params: &SaveRuntimeStateParams<'_>, - ) -> Result { - if params.is_empty() { - return Ok(true); - } - - // Read-modify-write. The service layer serialises writes per - // conversation_id through a single consumer task, so a naive - // RMW is race-free for our callers. - let raw: Option = - sqlx::query_scalar("SELECT session_config FROM acp_session WHERE conversation_id = ?") - .bind(conversation_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 = ?", - ) - .bind(new_config) - .bind(now) - .bind(conversation_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - async fn save_runtime_state_for_user( &self, user_id: &str, @@ -379,7 +386,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"); } @@ -395,9 +402,9 @@ mod tests { 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()); } @@ -405,7 +412,7 @@ 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] @@ -426,7 +433,7 @@ mod tests { .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-owner")); } @@ -434,9 +441,9 @@ mod tests { 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] @@ -448,22 +455,22 @@ mod tests { 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("conv-1").await.unwrap().is_some()); + assert!(repo.get_unscoped("conv-1").await.unwrap().is_some()); assert!(repo.delete_for_user("user-1", "conv-1").await.unwrap()); - assert!(repo.get("conv-1").await.unwrap().is_none()); + 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()); } @@ -473,7 +480,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")), @@ -486,7 +493,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 @@ -550,7 +557,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")), @@ -562,7 +569,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")), @@ -572,7 +579,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"), @@ -586,7 +593,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")), @@ -595,7 +602,7 @@ mod tests { ) .await .unwrap(); - repo.save_runtime_state( + repo.save_runtime_state_unscoped( "conv-1", &SaveRuntimeStateParams { current_mode_id: Some(None), @@ -605,7 +612,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()); } @@ -614,11 +621,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()); } @@ -626,7 +633,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")), From 45ae67686fd3ddb40979d7cdc20bc88388f62619 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 10:53:31 +0800 Subject: [PATCH 120/145] team: scope mailbox and task repository access --- .../aionui-db/src/repository/sqlite_team.rs | 193 ++++++++--- crates/aionui-db/src/repository/team.rs | 47 ++- crates/aionui-db/tests/team_repository.rs | 318 +++++++++++++----- crates/aionui-team/src/mailbox.rs | 22 +- crates/aionui-team/src/session.rs | 4 +- crates/aionui-team/src/task_board.rs | 34 +- crates/aionui-team/src/test_utils.rs | 72 +++- crates/aionui-team/tests/common/mod.rs | 47 ++- .../tests/session_service_integration.rs | 132 +++++--- 9 files changed, 630 insertions(+), 239 deletions(-) diff --git a/crates/aionui-db/src/repository/sqlite_team.rs b/crates/aionui-db/src/repository/sqlite_team.rs index 4277fb134..1ab168274 100644 --- a/crates/aionui-db/src/repository/sqlite_team.rs +++ b/crates/aionui-db/src/repository/sqlite_team.rs @@ -143,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) @@ -159,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?; @@ -178,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?; } @@ -200,34 +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, team_id: &str, 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 team_id = ? AND 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(()) @@ -235,6 +262,7 @@ impl ITeamRepository for SqliteTeamRepository { async fn get_history( &self, + user_id: &str, team_id: &str, to_agent_id: &str, limit: Option, @@ -245,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? @@ -259,10 +289,12 @@ 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? }; @@ -284,12 +316,13 @@ impl ITeamRepository for SqliteTeamRepository { // ── 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) @@ -302,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, team_id: &str, 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 = ?"); @@ -340,7 +394,9 @@ impl ITeamRepository for SqliteTeamRepository { set_clauses.push("updated_at = ?"); let sql = format!( - "UPDATE team_tasks SET {} WHERE team_id = ? AND id = ?", + "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(", ") ); @@ -363,6 +419,7 @@ impl ITeamRepository for SqliteTeamRepository { 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 { @@ -371,25 +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, team_id: &str, 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 team_id = ? AND id = ?") - .bind(team_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()) { @@ -397,13 +470,18 @@ 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 team_id = ? AND id = ?") - .bind(&new_blocks) - .bind(now_ms()) - .bind(team_id) - .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(()) @@ -411,6 +489,7 @@ impl ITeamRepository for SqliteTeamRepository { async fn remove_from_blocked_by( &self, + user_id: &str, team_id: &str, task_id: &str, unblocked_task_id: &str, @@ -418,24 +497,34 @@ impl ITeamRepository for SqliteTeamRepository { // 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 team_id = ? AND id = ?") - .bind(team_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 team_id = ? AND id = ?") - .bind(&new_blocked_by) - .bind(now_ms()) - .bind(team_id) - .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(()) diff --git a/crates/aionui-db/src/repository/team.rs b/crates/aionui-db/src/repository/team.rs index da0d4f694..f407d025a 100644 --- a/crates/aionui-db/src/repository/team.rs +++ b/crates/aionui-db/src/repository/team.rs @@ -55,22 +55,33 @@ pub trait ITeamRepository: Send + Sync { // ── 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, team_id: &str, 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, @@ -82,26 +93,44 @@ pub trait ITeamRepository: Send + Sync { // ── 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, team_id: &str, 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, team_id: &str, 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, + user_id: &str, team_id: &str, task_id: &str, unblocked_task_id: &str, diff --git a/crates/aionui-db/tests/team_repository.rs b/crates/aionui-db/tests/team_repository.rs index 1cbfc6fa9..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 { @@ -296,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()); } @@ -315,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()); } @@ -326,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")); @@ -340,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"); } @@ -354,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); } @@ -368,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); } @@ -380,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()); } @@ -390,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); } @@ -410,18 +412,18 @@ 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("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); } @@ -434,18 +436,45 @@ async fn scoped_mailbox_delete_and_mark_read_do_not_cross_team_owner() { repo.create_team(&make_team_for_user("t2", "user-b", "Team B")) .await .unwrap(); - repo.write_message(&make_mailbox_msg("m1", "t1", "a1", "a2", "message")) + repo.write_message("user-a", &make_mailbox_msg("m1", "t1", "a1", "a2", "message")) .await .unwrap(); - repo.write_message(&make_mailbox_msg("m2", "t2", "a1", "a2", "message")) + repo.write_message("user-b", &make_mailbox_msg("m2", "t2", "a1", "a2", "message")) .await .unwrap(); - repo.mark_read_batch("t1", &["m2".to_owned()]).await.unwrap(); - assert!(!repo.get_history("t2", "a1", None).await.unwrap()[0].read); + 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("t1", "a1", None).await.unwrap().len(), 1); + 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 ───────────────────────────────────────────────── @@ -456,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"); @@ -469,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()); } @@ -479,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"); } @@ -491,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()); } @@ -501,9 +533,10 @@ 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 { @@ -514,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"); } @@ -524,9 +561,10 @@ 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 { @@ -538,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")); } @@ -548,6 +590,7 @@ async fn update_nonexistent_task_returns_not_found() { let (repo, _db) = repo().await; let result = repo .update_task( + DEFAULT_USER_ID, "t1", "nonexistent", &UpdateTaskParams { @@ -568,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("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!(blocks.contains(&"tkB".to_string())); // Now complete taskA: remove tkA from taskB's blocked_by - repo.remove_from_blocked_by("t1", "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())); } @@ -592,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("t1", "tkA", "tkB").await.unwrap(); - repo.append_to_blocks("t1", "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 } @@ -614,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("t1", "tkA", "tkB").await.unwrap(); - repo.append_to_blocks("t1", "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("t1", "tkB", "tkA").await.unwrap(); - repo.remove_from_blocked_by("t1", "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()); @@ -643,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("t1", "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"]); } @@ -660,10 +745,11 @@ 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 { @@ -674,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()); @@ -686,15 +776,19 @@ 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("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); } @@ -707,13 +801,18 @@ async fn scoped_task_updates_and_deletes_stay_within_team_owner() { repo.create_team(&make_team_for_user("t2", "user-b", "Team B")) .await .unwrap(); - repo.create_task(&make_task("tk1", "t1", "A Task")).await.unwrap(); - repo.create_task(&make_task("tk2", "t2", "B Task")).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", - "tk2", + "tk1", &UpdateTaskParams { status: Some("completed".into()), ..Default::default() @@ -723,8 +822,49 @@ async fn scoped_task_updates_and_deletes_stay_within_team_owner() { assert!(matches!(update, Err(DbError::NotFound(_)))); repo.delete_tasks_by_team("user-b", "t1").await.unwrap(); - assert_eq!(repo.list_tasks("t1").await.unwrap().len(), 1); - assert_eq!(repo.list_tasks("t2").await.unwrap().len(), 1); + 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] @@ -736,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("t1", "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(); @@ -762,9 +904,9 @@ 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("system_default_user", "t1").await.unwrap(); @@ -774,9 +916,9 @@ async fn delete_team_cascades_mailbox_and_tasks() { // Verify all cleaned up 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()); } @@ -789,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("t1", "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-team/src/mailbox.rs b/crates/aionui-team/src/mailbox.rs index dc4d14298..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,7 +91,7 @@ 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) @@ -91,7 +99,7 @@ impl Mailbox { /// Marks the given message IDs as read. Called after successful prompt delivery. pub async fn mark_read_batch(&self, team_id: &str, ids: &[String]) -> Result<(), TeamError> { - self.repo.mark_read_batch(team_id, ids).await?; + self.repo.mark_read_batch(&self.user_id, team_id, ids).await?; Ok(()) } @@ -120,13 +128,13 @@ 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)) } diff --git a/crates/aionui-team/src/session.rs b/crates/aionui-team/src/session.rs index 534702fe6..73bc3068d 100644 --- a/crates/aionui-team/src/session.rs +++ b/crates/aionui-team/src/session.rs @@ -166,8 +166,8 @@ 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(), diff --git a/crates/aionui-team/src/task_board.rs b/crates/aionui-team/src/task_board.rs index 35cc43bca..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(team_id, 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,7 +97,7 @@ impl TaskBoard { metadata: update.metadata.as_ref().map(serde_json::to_string).transpose()?, }; - self.repo.update_task(team_id, 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(team_id, task_id, &existing).await?; @@ -95,7 +105,7 @@ impl TaskBoard { 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,7 +115,7 @@ 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) } @@ -119,7 +129,7 @@ impl TaskBoard { let blocks: Vec = serde_json::from_str(&completed_row.blocks)?; for downstream_id in &blocks { self.repo - .remove_from_blocked_by(team_id, 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, @@ -182,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/test_utils.rs b/crates/aionui-team/src/test_utils.rs index 07508b3c1..b8cf3c7e2 100644 --- a/crates/aionui-team/src/test_utils.rs +++ b/crates/aionui-team/src/test_utils.rs @@ -63,7 +63,7 @@ impl ITeamRepository for MockTeamRepo { // ── 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())); @@ -72,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 { @@ -84,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 @@ -95,7 +105,7 @@ impl ITeamRepository for MockTeamRepo { Ok(result) } - async fn mark_read_batch(&self, team_id: &str, 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 msg.team_id == team_id && ids.contains(&msg.id) { @@ -107,6 +117,7 @@ impl ITeamRepository for MockTeamRepo { async fn get_history( &self, + _user_id: &str, team_id: &str, to_agent_id: &str, limit: Option, @@ -130,12 +141,17 @@ impl ITeamRepository for MockTeamRepo { // ── 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 @@ -145,7 +161,13 @@ impl ITeamRepository for MockTeamRepo { Ok(found) } - async fn update_task(&self, team_id: &str, 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 @@ -171,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())); @@ -180,7 +202,13 @@ impl ITeamRepository for MockTeamRepo { Ok(tasks) } - async fn append_to_blocks(&self, team_id: &str, 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 @@ -195,6 +223,7 @@ impl ITeamRepository for MockTeamRepo { async fn remove_from_blocked_by( &self, + _user_id: &str, team_id: &str, task_id: &str, unblocked_task_id: &str, @@ -512,12 +541,17 @@ pub(crate) mod workspace_harness { 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> { @@ -526,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, _team_id: &str, _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, @@ -549,16 +585,22 @@ pub(crate) mod workspace_harness { 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, + _user_id: &str, _team_id: &str, _task_id: &str, _params: &aionui_db::UpdateTaskParams, @@ -566,12 +608,13 @@ pub(crate) mod workspace_harness { 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, + _user_id: &str, _team_id: &str, _task_id: &str, _blocked_task_id: &str, @@ -581,6 +624,7 @@ pub(crate) mod workspace_harness { async fn remove_from_blocked_by( &self, + _user_id: &str, _team_id: &str, _task_id: &str, _unblocked_task_id: &str, diff --git a/crates/aionui-team/tests/common/mod.rs b/crates/aionui-team/tests/common/mod.rs index 27d64927d..d6d3d821f 100644 --- a/crates/aionui-team/tests/common/mod.rs +++ b/crates/aionui-team/tests/common/mod.rs @@ -45,12 +45,17 @@ impl ITeamRepository for MockTeamRepo { 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 { @@ -62,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 @@ -73,7 +83,7 @@ impl ITeamRepository for MockTeamRepo { Ok(result) } - async fn mark_read_batch(&self, team_id: &str, 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 msg.team_id == team_id && ids.contains(&msg.id) { @@ -85,6 +95,7 @@ impl ITeamRepository for MockTeamRepo { async fn get_history( &self, + _user_id: &str, team_id: &str, to_agent_id: &str, limit: Option, @@ -106,12 +117,17 @@ impl ITeamRepository for MockTeamRepo { 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 @@ -121,7 +137,13 @@ impl ITeamRepository for MockTeamRepo { Ok(found) } - async fn update_task(&self, team_id: &str, 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 @@ -147,13 +169,19 @@ 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, team_id: &str, 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 @@ -168,6 +196,7 @@ impl ITeamRepository for MockTeamRepo { async fn remove_from_blocked_by( &self, + _user_id: &str, team_id: &str, task_id: &str, unblocked_task_id: &str, diff --git a/crates/aionui-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index ed1dc0a07..14c6e0aea 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -809,73 +809,87 @@ impl ITeamRepository for FullMockTeamRepo { 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, team_id: &str, ids: &[String]) -> Result<(), DbError> { - self.inner.mark_read_batch(team_id, 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, 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, + user_id: &str, team_id: &str, task_id: &str, params: &aionui_db::UpdateTaskParams, ) -> Result<(), DbError> { - self.inner.update_task(team_id, task_id, params).await + 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, team_id: &str, task_id: &str, blocked_task_id: &str) -> Result<(), DbError> { - self.inner.append_to_blocks(team_id, 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, + user_id: &str, team_id: &str, task_id: &str, unblocked_task_id: &str, ) -> Result<(), DbError> { self.inner - .remove_from_blocked_by(team_id, task_id, unblocked_task_id) + .remove_from_blocked_by(user_id, team_id, task_id, unblocked_task_id) .await } async fn delete_tasks_by_team(&self, user_id: &str, team_id: &str) -> Result<(), DbError> { @@ -1724,18 +1738,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 +1798,21 @@ async fn teammate_first_wake_uses_canonical_prompt_at_service_boundary() { .expect("clear existing session"); team_repo - .write_message(&aionui_db::models::MailboxMessageRow { - id: "mailbox-worker-1".into(), - team_id: created.id.clone(), - to_agent_id: worker_slot_id.clone(), - from_agent_id: "user".into(), - msg_type: "message".into(), - content: "do X".into(), - summary: None, - files: None, - read: false, - created_at: aionui_common::now_ms(), - }) + .write_message( + "user1", + &aionui_db::models::MailboxMessageRow { + id: "mailbox-worker-1".into(), + team_id: created.id.clone(), + to_agent_id: worker_slot_id.clone(), + from_agent_id: "user".into(), + msg_type: "message".into(), + content: "do X".into(), + summary: None, + files: None, + read: false, + created_at: aionui_common::now_ms(), + }, + ) .await .expect("seed teammate mailbox"); @@ -1851,18 +1871,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"); @@ -3887,7 +3910,10 @@ async fn manual_add_agent_attach_failure_marks_slot_error_and_notifies_leader() ); 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() From afc51740aa7a33f3ab75c7097b0f427afda714ec Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 11:02:47 +0800 Subject: [PATCH 121/145] cron: reject cross-user conversation bindings --- .../aionui-db/src/repository/sqlite_cron.rs | 32 +++++++++++++------ crates/aionui-db/tests/cron_repository.rs | 28 +++++++++++++--- 2 files changed, 46 insertions(+), 14 deletions(-) diff --git a/crates/aionui-db/src/repository/sqlite_cron.rs b/crates/aionui-db/src/repository/sqlite_cron.rs index e9df2747f..cc5d039b9 100644 --- a/crates/aionui-db/src/repository/sqlite_cron.rs +++ b/crates/aionui-db/src/repository/sqlite_cron.rs @@ -22,6 +22,9 @@ impl SqliteCronRepository { #[async_trait::async_trait] impl ICronRepository for SqliteCronRepository { async fn insert(&self, row: &CronJobRow) -> Result<(), DbError> { + self.validate_conversation_owner(&row.user_id, &row.conversation_id) + .await?; + sqlx::query( "INSERT INTO cron_jobs (\ id, user_id, name, enabled, schedule_kind, schedule_value, schedule_tz, \ @@ -386,6 +389,23 @@ 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(); @@ -464,15 +484,7 @@ impl SqliteCronRepository { if let (Some(user_id), Some(conversation_id)) = (user_id, params.conversation_id.as_deref()) && !conversation_id.trim().is_empty() { - let owns_new_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_new_conversation { - return Err(DbError::NotFound(format!("conversation '{conversation_id}'"))); - } + self.validate_conversation_owner(user_id, conversation_id).await?; } let sql = if user_id.is_some() { @@ -557,7 +569,7 @@ mod tests { let now = now_ms(); CronJobRow { id: id.into(), - user_id: "user1".into(), + user_id: "user_1".into(), name: "Test Job".into(), enabled: true, schedule_kind: "every".into(), diff --git a/crates/aionui-db/tests/cron_repository.rs b/crates/aionui-db/tests/cron_repository.rs index 72635ecf6..0061178b1 100644 --- a/crates/aionui-db/tests/cron_repository.rs +++ b/crates/aionui-db/tests/cron_repository.rs @@ -243,18 +243,38 @@ async fn scoped_crud_filters_by_job_user_id() { ); } +#[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 orphan = make_job("cron_orphan"); - orphan.conversation_id = "missing_conversation".into(); - r.insert(&orphan).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().await.unwrap(); assert!(enabled.iter().any(|job| job.id == "cron_owned")); - assert!(enabled.iter().any(|job| job.id == "cron_orphan")); + assert!(enabled.iter().any(|job| job.id == "cron_new_conversation")); } #[tokio::test] From a2b0a5fca0a5f77df1c1f57239919affae95651b Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 11:23:53 +0800 Subject: [PATCH 122/145] assistant: resolve agent bindings within user scope --- crates/aionui-assistant/src/service.rs | 44 +++++++---- crates/aionui-db/src/agent_binding.rs | 12 ++- crates/aionui-db/src/lib.rs | 4 +- .../aionui-db/tests/agent_binding_resolver.rs | 75 ++++++++++++++++++- 4 files changed, 114 insertions(+), 21 deletions(-) diff --git a/crates/aionui-assistant/src/service.rs b/crates/aionui-assistant/src/service.rs index 9b803726c..0e014a9e6 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; @@ -244,7 +244,9 @@ impl AssistantService { .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") @@ -714,7 +716,7 @@ impl AssistantService { 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_for_user( @@ -850,7 +852,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)?); } @@ -884,7 +886,7 @@ impl AssistantService { } 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); } @@ -911,7 +913,7 @@ impl AssistantService { 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, @@ -957,7 +959,7 @@ impl AssistantService { .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, \ @@ -967,12 +969,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 { @@ -981,9 +987,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 { @@ -994,12 +1000,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); @@ -1051,7 +1058,8 @@ impl AssistantService { Some(s) if !s.trim().is_empty() => s.trim().to_string(), _ => 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, @@ -1149,7 +1157,8 @@ 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.resolve_runtime_backend_for_agent_id(user_id, requested_agent_id) + .await?; } self.apply_detail_overrides_for_user(user_id, id, detail_overrides, reset_model_and_permission) .await?; @@ -1259,7 +1268,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() @@ -1689,7 +1698,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, 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 7ee9a5b9e..f9f987184 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, 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, + } +} From b80394aaa023bc4bb130ecf97f476f88b0b429c9 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 11:28:00 +0800 Subject: [PATCH 123/145] db: require user-scoped agent metadata methods --- .../src/repository/agent_metadata.rs | 50 ++++--------------- 1 file changed, 10 insertions(+), 40 deletions(-) diff --git a/crates/aionui-db/src/repository/agent_metadata.rs b/crates/aionui-db/src/repository/agent_metadata.rs index 055aba11b..621c78e47 100644 --- a/crates/aionui-db/src/repository/agent_metadata.rs +++ b/crates/aionui-db/src/repository/agent_metadata.rs @@ -17,18 +17,12 @@ 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> { - let _ = user_id; - self.list_all().await - } + 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> { - let _ = user_id; - self.get(id).await - } + 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( @@ -42,10 +36,7 @@ pub trait IAgentMetadataRepository: Send + Sync { user_id: &str, agent_source: &str, name: &str, - ) -> Result, DbError> { - let _ = user_id; - self.find_by_source_and_name(agent_source, name).await - } + ) -> Result, DbError>; /// Look up the first `builtin` row whose vendor label matches. /// Useful when the caller only has the legacy `backend` string and @@ -56,10 +47,7 @@ pub trait IAgentMetadataRepository: Send + Sync { &self, user_id: &str, backend: &str, - ) -> Result, DbError> { - let _ = user_id; - self.find_builtin_by_backend(backend).await - } + ) -> Result, DbError>; /// Insert or replace a row. Returns the row as stored. async fn upsert(&self, params: &UpsertAgentMetadataParams<'_>) -> Result; @@ -68,10 +56,7 @@ pub trait IAgentMetadataRepository: Send + Sync { &self, user_id: &str, params: &UpsertAgentMetadataParams<'_>, - ) -> Result { - let _ = user_id; - self.upsert(params).await - } + ) -> Result; async fn upsert_global(&self, params: &UpsertAgentMetadataParams<'_>) -> Result { self.upsert(params).await @@ -90,10 +75,7 @@ pub trait IAgentMetadataRepository: Send + Sync { user_id: &str, id: &str, params: &UpdateAgentHandshakeParams<'_>, - ) -> Result, DbError> { - let _ = user_id; - self.apply_handshake(id, params).await - } + ) -> Result, DbError>; /// Persist the latest availability snapshot for an existing row. /// Returns `Ok(None)` if no row matches `id`. @@ -108,10 +90,7 @@ pub trait IAgentMetadataRepository: Send + Sync { user_id: &str, id: &str, params: &UpdateAgentAvailabilitySnapshotParams<'_>, - ) -> Result, DbError> { - let _ = user_id; - self.update_availability_snapshot(id, params).await - } + ) -> Result, DbError>; /// Write only the self-repair override columns for an agent, leaving all /// other columns (seed truth + availability snapshot) untouched. Kept @@ -130,24 +109,15 @@ pub trait IAgentMetadataRepository: Send + Sync { id: &str, command_override: Option<&str>, env_override: Option<&str>, - ) -> Result<(), DbError> { - let _ = user_id; - self.update_agent_overrides(id, command_override, env_override).await - } + ) -> 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 { - let _ = user_id; - self.set_enabled(id, enabled).await - } + 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 { - let _ = user_id; - self.delete(id).await - } + async fn delete_for_user(&self, user_id: &str, id: &str) -> Result; } From cb4bd70a10dd66e68cd3163fb8542ffd2261b305 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 11:30:33 +0800 Subject: [PATCH 124/145] db: require user-scoped skill repository methods --- crates/aionui-db/src/repository/skill.rs | 35 +++++------------------- 1 file changed, 7 insertions(+), 28 deletions(-) diff --git a/crates/aionui-db/src/repository/skill.rs b/crates/aionui-db/src/repository/skill.rs index bcde126f5..3f4d09202 100644 --- a/crates/aionui-db/src/repository/skill.rs +++ b/crates/aionui-db/src/repository/skill.rs @@ -8,37 +8,25 @@ pub trait ISkillRepository: Send + Sync { 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> { - let _ = user_id; - self.list().await - } + 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> { - let _ = user_id; - self.find_by_name(name).await - } + 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> { - let _ = user_id; - self.find_by_name_any(name).await - } + 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 { - let _ = user_id; - self.upsert(params).await - } + 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 { @@ -49,10 +37,7 @@ pub trait ISkillRepository: Send + Sync { 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 { - let _ = user_id; - self.delete_by_name(name).await - } + async fn delete_by_name_for_user(&self, user_id: &str, name: &str) -> Result; /// Appends one import record. async fn create_import_record( @@ -65,10 +50,7 @@ pub trait ISkillRepository: Send + Sync { &self, user_id: &str, params: CreateSkillImportRecordParams<'_>, - ) -> Result { - let _ = user_id; - self.create_import_record(params).await - } + ) -> Result; /// Lists recent import records ordered by creation time descending. async fn list_import_records(&self, limit: i64) -> Result, DbError>; @@ -78,10 +60,7 @@ pub trait ISkillRepository: Send + Sync { &self, user_id: &str, limit: i64, - ) -> Result, DbError> { - let _ = user_id; - self.list_import_records(limit).await - } + ) -> Result, DbError>; } /// Parameters for creating or updating a skill row. From d3b906c36c057a6749d8cff086bde6c8bccad2a9 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 11:40:04 +0800 Subject: [PATCH 125/145] db: require user-scoped assistant repository methods --- crates/aionui-channel/src/channel_settings.rs | 94 +++++++ .../tests/message_service_integration.rs | 2 +- crates/aionui-db/src/repository/assistant.rs | 145 ++-------- crates/aionui-team/src/provisioning.rs | 59 ++++ .../aionui-team/src/service/spawn_support.rs | 257 ++++++++++++++++++ crates/aionui-team/src/test_utils.rs | 163 +++++++++++ .../tests/session_service_integration.rs | 233 ++++++++++++++++ 7 files changed, 836 insertions(+), 117 deletions(-) diff --git a/crates/aionui-channel/src/channel_settings.rs b/crates/aionui-channel/src/channel_settings.rs index b278a9273..40cebfcb6 100644 --- a/crates/aionui-channel/src/channel_settings.rs +++ b/crates/aionui-channel/src/channel_settings.rs @@ -615,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, @@ -635,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<'_>, @@ -642,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 { @@ -661,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 { diff --git a/crates/aionui-channel/tests/message_service_integration.rs b/crates/aionui-channel/tests/message_service_integration.rs index 7ec82fc80..6920735ed 100644 --- a/crates/aionui-channel/tests/message_service_integration.rs +++ b/crates/aionui-channel/tests/message_service_integration.rs @@ -351,7 +351,7 @@ async fn send_to_agent_persists_assistant_snapshot_for_channel_bound_assistant() .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"); diff --git a/crates/aionui-db/src/repository/assistant.rs b/crates/aionui-db/src/repository/assistant.rs index 6b108e273..d53f969f1 100644 --- a/crates/aionui-db/src/repository/assistant.rs +++ b/crates/aionui-db/src/repository/assistant.rs @@ -15,31 +15,19 @@ 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> { - let _ = user_id; - self.list().await - } + 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> { - let _ = user_id; - self.get(id).await - } + 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 { - let _ = user_id; - self.create(params).await - } + 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. @@ -50,32 +38,20 @@ pub trait IAssistantRepository: Send + Sync { user_id: &str, id: &str, params: &UpdateAssistantParams<'_>, - ) -> Result, DbError> { - let _ = user_id; - self.update(id, params).await - } + ) -> 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 { - let _ = user_id; - self.delete(id).await - } + 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 { - let _ = user_id; - self.upsert(params).await - } + async fn upsert_for_user(&self, user_id: &str, params: &CreateAssistantParams<'_>) + -> Result; } /// Per-assistant user state (enabled flag, sort order, last-used timestamp). @@ -84,18 +60,12 @@ 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> { - let _ = user_id; - self.get(assistant_id).await - } + 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> { - let _ = user_id; - self.get_all().await - } + 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; @@ -104,54 +74,36 @@ pub trait IAssistantOverrideRepository: Send + Sync { &self, user_id: &str, params: &UpsertOverrideParams<'_>, - ) -> Result { - let _ = user_id; - self.upsert(params).await - } + ) -> 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 { - let _ = user_id; - self.delete(assistant_id).await - } + 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 { - let _ = user_id; - self.delete_orphans(valid_ids).await - } + 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> { - let _ = user_id; - self.list().await - } + 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> { - let _ = user_id; - self.list_including_deleted().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> { - let _ = user_id; - self.get_by_assistant_id(assistant_id).await - } + ) -> Result, DbError>; async fn get_by_assistant_id_including_deleted( &self, assistant_id: &str, @@ -162,15 +114,9 @@ pub trait IAssistantDefinitionRepository: Send + Sync { &self, user_id: &str, assistant_id: &str, - ) -> Result, DbError> { - let _ = user_id; - self.get_by_assistant_id_including_deleted(assistant_id).await - } + ) -> 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> { - let _ = user_id; - self.get_by_id(id).await - } + async fn get_by_id_for_user(&self, user_id: &str, id: &str) -> Result, DbError>; async fn get_by_source_ref( &self, source: &str, @@ -181,10 +127,7 @@ pub trait IAssistantDefinitionRepository: Send + Sync { user_id: &str, source: &str, source_ref: &str, - ) -> Result, DbError> { - let _ = user_id; - self.get_by_source_ref(source, source_ref).await - } + ) -> Result, DbError>; async fn get_by_source_ref_including_deleted( &self, source: &str, @@ -197,10 +140,7 @@ pub trait IAssistantDefinitionRepository: Send + Sync { user_id: &str, source: &str, source_ref: &str, - ) -> Result, DbError> { - let _ = user_id; - self.get_by_source_ref_including_deleted(source, source_ref).await - } + ) -> Result, DbError>; async fn get_global_by_source_ref_including_deleted( &self, source: &str, @@ -219,10 +159,7 @@ pub trait IAssistantDefinitionRepository: Send + Sync { &self, user_id: &str, params: &UpsertAssistantDefinitionParams<'_>, - ) -> Result { - let _ = user_id; - self.upsert(params).await - } + ) -> Result; async fn upsert_global( &self, params: &UpsertAssistantDefinitionParams<'_>, @@ -241,10 +178,7 @@ 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 { - let _ = user_id; - self.soft_delete(id, deleted_at).await - } + 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. @@ -255,29 +189,17 @@ pub trait IAssistantOverlayRepository: Send + Sync { &self, user_id: &str, assistant_definition_id: &str, - ) -> Result, DbError> { - let _ = user_id; - self.get(assistant_definition_id).await - } + ) -> Result, DbError>; async fn list(&self) -> Result, DbError>; - async fn list_for_user(&self, user_id: &str) -> Result, DbError> { - let _ = user_id; - self.list().await - } + 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 { - let _ = user_id; - self.upsert(params).await - } + ) -> Result; async fn delete(&self, assistant_definition_id: &str) -> Result; - async fn delete_for_user(&self, user_id: &str, assistant_definition_id: &str) -> Result { - let _ = user_id; - self.delete(assistant_definition_id).await - } + async fn delete_for_user(&self, user_id: &str, assistant_definition_id: &str) -> Result; } /// Assistant-scoped "auto remember last" preferences. @@ -288,22 +210,13 @@ pub trait IAssistantPreferenceRepository: Send + Sync { &self, user_id: &str, assistant_definition_id: &str, - ) -> Result, DbError> { - let _ = user_id; - self.get(assistant_definition_id).await - } + ) -> Result, DbError>; async fn upsert(&self, params: &UpsertAssistantPreferenceParams<'_>) -> Result; async fn upsert_for_user( &self, user_id: &str, params: &UpsertAssistantPreferenceParams<'_>, - ) -> Result { - let _ = user_id; - self.upsert(params).await - } + ) -> Result; async fn delete(&self, assistant_definition_id: &str) -> Result; - async fn delete_for_user(&self, user_id: &str, assistant_definition_id: &str) -> Result { - let _ = user_id; - self.delete(assistant_definition_id).await - } + async fn delete_for_user(&self, user_id: &str, assistant_definition_id: &str) -> Result; } diff --git a/crates/aionui-team/src/provisioning.rs b/crates/aionui-team/src/provisioning.rs index d682f99f8..5cf80999f 100644 --- a/crates/aionui-team/src/provisioning.rs +++ b/crates/aionui-team/src/provisioning.rs @@ -881,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, @@ -891,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, @@ -904,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, @@ -911,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, @@ -919,12 +963,27 @@ 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; diff --git a/crates/aionui-team/src/service/spawn_support.rs b/crates/aionui-team/src/service/spawn_support.rs index 6e8c41c5f..6ac91ae47 100644 --- a/crates/aionui-team/src/service/spawn_support.rs +++ b/crates/aionui-team/src/service/spawn_support.rs @@ -287,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, @@ -303,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<'_>, @@ -310,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)] @@ -326,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, @@ -342,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<'_>, @@ -349,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)] @@ -365,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)] @@ -393,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 { @@ -471,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, @@ -487,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 @@ -495,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, @@ -507,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, @@ -515,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, @@ -524,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 { diff --git a/crates/aionui-team/src/test_utils.rs b/crates/aionui-team/src/test_utils.rs index b8cf3c7e2..badaa5dc8 100644 --- a/crates/aionui-team/src/test_utils.rs +++ b/crates/aionui-team/src/test_utils.rs @@ -949,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, @@ -961,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, @@ -977,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, @@ -985,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, @@ -994,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; @@ -1011,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, @@ -1027,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<'_>, @@ -1034,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; @@ -1047,17 +1186,41 @@ 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; diff --git a/crates/aionui-team/tests/session_service_integration.rs b/crates/aionui-team/tests/session_service_integration.rs index 14c6e0aea..343238646 100644 --- a/crates/aionui-team/tests/session_service_integration.rs +++ b/crates/aionui-team/tests/session_service_integration.rs @@ -931,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, @@ -945,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, @@ -958,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, @@ -965,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, @@ -973,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 + } } // --------------------------------------------------------------------------- @@ -1406,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, @@ -1422,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; @@ -1439,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 { @@ -1462,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, @@ -1478,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 { @@ -1497,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) { From 7f00788384a44fff5758fa739a84c740a4a79146 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 11:44:22 +0800 Subject: [PATCH 126/145] assistant: scope legacy bootstrap to default user --- crates/aionui-assistant/src/service.rs | 68 ++++++++++++++------------ 1 file changed, 37 insertions(+), 31 deletions(-) diff --git a/crates/aionui-assistant/src/service.rs b/crates/aionui-assistant/src/service.rs index 0e014a9e6..6eb16e791 100644 --- a/crates/aionui-assistant/src/service.rs +++ b/crates/aionui-assistant/src/service.rs @@ -322,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}")))? { @@ -352,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, @@ -364,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" { @@ -424,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!( @@ -439,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}")))?; @@ -453,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}")))?; } @@ -674,15 +689,6 @@ impl AssistantService { Ok(()) } - async fn upsert_definition_from_legacy_user_row( - &self, - row: &AssistantRow, - requested_agent_id: Option<&str>, - ) -> Result<(), AssistantError> { - self.upsert_definition_from_legacy_user_row_for_user(DEFAULT_USER_ID, row, requested_agent_id) - .await - } - async fn upsert_definition_from_legacy_user_row_for_user( &self, user_id: &str, From 8894d1c2245a3a3b3e668ecade0160fc4cdd92a6 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 11:57:08 +0800 Subject: [PATCH 127/145] channel: use structured runtime plugin keys --- crates/aionui-channel/src/manager.rs | 38 ++++++++++++++++++++++------ 1 file changed, 30 insertions(+), 8 deletions(-) diff --git a/crates/aionui-channel/src/manager.rs b/crates/aionui-channel/src/manager.rs index cea27ea03..3cc32de50 100644 --- a/crates/aionui-channel/src/manager.rs +++ b/crates/aionui-channel/src/manager.rs @@ -29,7 +29,7 @@ pub struct ChannelManager { broadcaster: Arc, encryption_key: [u8; 32], /// Active plugin instances keyed by owner user ID + plugin ID. - plugins: DashMap>, + 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 @@ -322,7 +337,7 @@ 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_by_key(&key).await; @@ -335,11 +350,10 @@ impl ChannelManager { /// 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 prefix = format!("{owner_user_id}:"); - let keys: Vec = self + let keys: Vec = self .plugins .iter() - .filter(|entry| entry.key().starts_with(&prefix)) + .filter(|entry| entry.key().owner_user_id == owner_user_id) .map(|entry| entry.key().clone()) .collect(); let stopped = keys.len(); @@ -451,8 +465,8 @@ impl ChannelManager { } /// Stops and removes an active plugin instance. - fn runtime_key(owner_user_id: &str, plugin_id: &str) -> String { - format!("{owner_user_id}:{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 { @@ -480,7 +494,7 @@ impl ChannelManager { } /// Stops and removes an active plugin instance by runtime key. - async fn stop_plugin_by_key(&self, key: &str) { + 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 { @@ -1554,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(); From c1dd23aaea425d3748fe8b5940e24b47d5cb83d5 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 12:00:53 +0800 Subject: [PATCH 128/145] office: use structured preview session keys --- crates/aionui-office/src/watch_manager.rs | 36 +++++++++++++++++++---- 1 file changed, 31 insertions(+), 5 deletions(-) diff --git a/crates/aionui-office/src/watch_manager.rs b/crates/aionui-office/src/watch_manager.rs index a081e65d5..adfff972b 100644 --- a/crates/aionui-office/src/watch_manager.rs +++ b/crates/aionui-office/src/watch_manager.rs @@ -62,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 { @@ -462,8 +479,8 @@ fn resolve_path(file_path: &str) -> Result { Ok(resolved.to_string_lossy().into_owned()) } -fn session_key(user_id: &str, resolved_path: &str, doc_type: DocType) -> String { - format!("{user_id}:{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 { @@ -642,9 +659,11 @@ mod tests { } #[test] - fn session_key_format() { + fn session_key_preserves_fields() { let key = session_key("user-1", "/path/to/doc.docx", DocType::Word); - assert_eq!(key, "user-1:word:/path/to/doc.docx"); + assert_eq!(key.user_id, "user-1"); + assert_eq!(key.doc_type, DocType::Word); + assert_eq!(key.resolved_path, "/path/to/doc.docx"); } #[test] @@ -661,6 +680,13 @@ mod tests { 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()); From 15cde422188991f3ffbad42fb4ffca6bf82d262a Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 12:04:52 +0800 Subject: [PATCH 129/145] office: stop preview sessions on user revoke --- crates/aionui-app/src/router/routes.rs | 3 ++ crates/aionui-office/src/watch_manager.rs | 43 +++++++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index 261feb7d1..03882c54a 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -173,6 +173,7 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates 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); @@ -188,8 +189,10 @@ pub fn create_router_with_all_state(services: &AppServices, states: ModuleStates 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, diff --git a/crates/aionui-office/src/watch_manager.rs b/crates/aionui-office/src/watch_manager.rs index adfff972b..d6b2ce8a5 100644 --- a/crates/aionui-office/src/watch_manager.rs +++ b/crates/aionui-office/src/watch_manager.rs @@ -245,6 +245,27 @@ 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() @@ -867,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()); From 9936ea3dd40afa61ccab3f7df18fc56a97b19551 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 14:11:12 +0800 Subject: [PATCH 130/145] db: require user scope for acp session create --- .../src/persistence/acp_session_sync.rs | 1 + crates/aionui-conversation/src/service.rs | 1 + .../aionui-conversation/src/service_test.rs | 121 ++++++++++++++++++ .../src/session_context.rs | 3 + .../aionui-db/src/repository/acp_session.rs | 1 + .../src/repository/sqlite_acp_session.rs | 52 ++++++-- 6 files changed, 169 insertions(+), 10 deletions(-) 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 96e49ca12..56a913b68 100644 --- a/crates/aionui-ai-agent/src/persistence/acp_session_sync.rs +++ b/crates/aionui-ai-agent/src/persistence/acp_session_sync.rs @@ -317,6 +317,7 @@ mod tests { .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", diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index 974b57361..784154205 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -1318,6 +1318,7 @@ impl ConversationService { }; let params = CreateAcpSessionParams { + user_id, conversation_id, agent_source, agent_id: &resolved_agent_id, diff --git a/crates/aionui-conversation/src/service_test.rs b/crates/aionui-conversation/src/service_test.rs index c4da93d1a..d7543f66b 100644 --- a/crates/aionui-conversation/src/service_test.rs +++ b/crates/aionui-conversation/src/service_test.rs @@ -781,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, @@ -791,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, @@ -806,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, @@ -813,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, @@ -821,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 @@ -881,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, @@ -891,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, @@ -904,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, @@ -911,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, @@ -919,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)] @@ -986,6 +1104,7 @@ impl StubAcpSessionRepo { #[derive(Debug, Clone, PartialEq, Eq)] struct CreateAcpSessionCall { + user_id: String, conversation_id: String, agent_source: String, agent_id: String, @@ -998,6 +1117,7 @@ impl IAcpSessionRepository for StubAcpSessionRepo { } 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(), @@ -7149,6 +7269,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"); diff --git a/crates/aionui-conversation/src/session_context.rs b/crates/aionui-conversation/src/session_context.rs index 80a445be6..6af20e2eb 100644 --- a/crates/aionui-conversation/src/session_context.rs +++ b/crates/aionui-conversation/src/session_context.rs @@ -854,6 +854,7 @@ mod tests { repos .acp_session_repo .create(&CreateAcpSessionParams { + user_id: "user-1", conversation_id: "conv-1", agent_source: "builtin", agent_id: "builtin-claude-test", @@ -906,6 +907,7 @@ mod tests { repos .acp_session_repo .create(&CreateAcpSessionParams { + user_id: "user-1", conversation_id: "conv-1", agent_source: "builtin", agent_id: "builtin-codex-test", @@ -944,6 +946,7 @@ mod tests { repos .acp_session_repo .create(&CreateAcpSessionParams { + user_id: "user-1", conversation_id: "conv-1", agent_source: "builtin", agent_id: "builtin-codex-test", diff --git a/crates/aionui-db/src/repository/acp_session.rs b/crates/aionui-db/src/repository/acp_session.rs index 563595ec3..c80cc9917 100644 --- a/crates/aionui-db/src/repository/acp_session.rs +++ b/crates/aionui-db/src/repository/acp_session.rs @@ -22,6 +22,7 @@ use crate::models::AcpSessionRow; /// 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, diff --git a/crates/aionui-db/src/repository/sqlite_acp_session.rs b/crates/aionui-db/src/repository/sqlite_acp_session.rs index a165d26b5..cea6a33d8 100644 --- a/crates/aionui-db/src/repository/sqlite_acp_session.rs +++ b/crates/aionui-db/src/repository/sqlite_acp_session.rs @@ -20,6 +20,7 @@ impl SqliteAcpSessionRepository { Self { pool } } + #[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) @@ -203,16 +204,19 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { async fn create(&self, params: &CreateAcpSessionParams<'_>) -> Result { let now = now_ms(); - sqlx::query( + let result = sqlx::query( "INSERT INTO acp_session \ (conversation_id, agent_source, agent_id, \ session_id, session_status, session_config, last_active_at) \ - VALUES (?, ?, ?, NULL, 'idle', '{}', ?)", + SELECT c.id, ?, ?, NULL, 'idle', '{}', ? \ + FROM conversations c \ + WHERE c.id = ? AND c.user_id = ?", ) - .bind(params.conversation_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 { @@ -223,12 +227,21 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { _ => DbError::Query(e), })?; - self.get_unscoped(params.conversation_id).await?.ok_or_else(|| { - DbError::Init(format!( - "create did not produce acp_session row for '{}'", - params.conversation_id - )) - }) + if result.rows_affected() == 0 { + return Err(DbError::NotFound(format!( + "conversation '{}' for user '{}'", + params.conversation_id, params.user_id + ))); + } + + 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( @@ -336,11 +349,13 @@ 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", @@ -362,7 +377,7 @@ mod tests { .await .unwrap(); sqlx::query( - "INSERT INTO conversations \ + "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, ?, ?)", @@ -398,6 +413,23 @@ 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::NotFound(_))); + 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; From 6b63ce1661fcbf41cf41cf9c53c539e447cc31c0 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 14:36:10 +0800 Subject: [PATCH 131/145] test: cover user scope migration integrity --- .../aionui-db/tests/user_scope_migration.rs | 253 ++++++++++++++++++ 1 file changed, 253 insertions(+) diff --git a/crates/aionui-db/tests/user_scope_migration.rs b/crates/aionui-db/tests/user_scope_migration.rs index 917646929..b97b97561 100644 --- a/crates/aionui-db/tests/user_scope_migration.rs +++ b/crates/aionui-db/tests/user_scope_migration.rs @@ -129,6 +129,232 @@ async fn migration_028_migrates_cron_skills_as_global_rows() { assert_eq!(row.get::, _>("user_id"), None); } +#[tokio::test] +async fn migration_028_backfills_existing_independent_roots_to_default_user() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 27).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, 28).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_028_classifies_global_and_user_catalog_rows() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 27).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, 28).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_028_keeps_new_conversation_cron_jobs_unanchored_until_run() { let db = init_database_memory().await.unwrap(); @@ -160,6 +386,33 @@ async fn migration_028_keeps_new_conversation_cron_jobs_unanchored_until_run() { assert_eq!(row.get::("conversation_id"), ""); } +#[tokio::test] +async fn migration_028_rejects_channel_session_without_channel_user() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 27).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, 28).await.unwrap_err(); + assert!( + err.to_string().contains("CHECK constraint failed"), + "unexpected migration error: {err}" + ); +} + #[tokio::test] async fn migration_028_rejects_channel_session_cross_user_conversation() { let pool = SqlitePoolOptions::new() From 427a55a18170679d77b2d7730a70fbc8a32efac0 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 15:11:30 +0800 Subject: [PATCH 132/145] cron: report cross-account conversation links --- crates/aionui-cron/src/error.rs | 3 + crates/aionui-cron/src/executor.rs | 71 +++++++++++++++++++ crates/aionui-cron/src/routes.rs | 10 +++ crates/aionui-cron/src/service.rs | 42 +++++++---- .../aionui-cron/tests/service_integration.rs | 40 +++++++++++ 5 files changed, 153 insertions(+), 13 deletions(-) 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/executor.rs b/crates/aionui-cron/src/executor.rs index 408e7625a..e92351390 100644 --- a/crates/aionui-cron/src/executor.rs +++ b/crates/aionui-cron/src/executor.rs @@ -3459,9 +3459,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, @@ -3469,18 +3482,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, @@ -3488,6 +3523,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, @@ -3495,6 +3538,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, @@ -3503,11 +3554,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 979b4439e..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), @@ -369,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/service.rs b/crates/aionui-cron/src/service.rs index 9dd554ede..924521d9b 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -267,19 +267,9 @@ impl CronService { 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) - && (conversation_id.is_empty() - || self - .executor - .get_conversation_row_for_user(user_id, conversation_id) - .await? - .is_none()) - { - return Err(CronError::Conversation( - aionui_conversation::ConversationError::NotFound { - id: req.conversation_id, - }, - )); + if matches!(execution_mode, ExecutionMode::Existing) { + self.require_existing_conversation_scope(user_id, conversation_id) + .await?; } let agent_config = match req.agent_config { @@ -890,6 +880,32 @@ impl CronService { } } + 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 { let conversation_id = job.conversation_id.trim(); if conversation_id.is_empty() { diff --git a/crates/aionui-cron/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index ee2711c87..c174ee8f6 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -2368,6 +2368,46 @@ async fn oc2_rejects_existing_jobs_with_missing_conversation() { assert!(matches!(err, aionui_cron::error::CronError::Conversation(_))); } +#[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 mut req = make_create_req("Cross User Conversation", every_60s()); + req.conversation_id = "conv_user_b".into(); + + 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_is_rejected() { let (svc, _repo, _bc, _conv_repo) = setup_with_conv_repo().await; From 61c6aba0e7e099c3da3b04037519fc7ef0ad6f2c Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 15:34:30 +0800 Subject: [PATCH 133/145] channel: report cross-account session links --- crates/aionui-channel/src/routes.rs | 14 +++++++++ .../src/repository/sqlite_channel.rs | 31 ++++++++++++++++++- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/crates/aionui-channel/src/routes.rs b/crates/aionui-channel/src/routes.rs index bdb6c6c72..244c2b36c 100644 --- a/crates/aionui-channel/src/routes.rs +++ b/crates/aionui-channel/src/routes.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use axum::Router; use axum::extract::rejection::JsonRejection; use axum::extract::{Extension, Json, Path, State}; +use axum::http::StatusCode; use axum::routing::{get, post, put}; use tracing::warn; @@ -48,6 +49,9 @@ pub struct ChannelRouterState { fn db_error_to_api_error(err: DbError) -> 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}")), @@ -882,6 +886,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-db/src/repository/sqlite_channel.rs b/crates/aionui-db/src/repository/sqlite_channel.rs index 0376f3e00..1e33b4911 100644 --- a/crates/aionui-db/src/repository/sqlite_channel.rs +++ b/crates/aionui-db/src/repository/sqlite_channel.rs @@ -365,6 +365,32 @@ impl IChannelRepository for SqliteChannelRepository { .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(()) @@ -1159,7 +1185,10 @@ mod tests { .update_session_conversation(OWNER_A, "sess-1", "conv-other") .await .unwrap_err(); - assert!(matches!(err, DbError::NotFound(_))); + assert!(matches!( + err, + DbError::Conflict(msg) if msg.starts_with("CROSS_ACCOUNT_REFERENCE:") + )); assert!( repo.get_session(OWNER_A, "sess-1") .await From 37aa08274f17bc56f659391a7d69dc860bc61ca8 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 15:54:59 +0800 Subject: [PATCH 134/145] conversation: report cross-account acp session links --- crates/aionui-conversation/src/routes.rs | 13 ++++++++++++ .../src/repository/sqlite_acp_session.rs | 20 ++++++++++++++++++- 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/crates/aionui-conversation/src/routes.rs b/crates/aionui-conversation/src/routes.rs index dbebe2993..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), @@ -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-db/src/repository/sqlite_acp_session.rs b/crates/aionui-db/src/repository/sqlite_acp_session.rs index cea6a33d8..5fce9340f 100644 --- a/crates/aionui-db/src/repository/sqlite_acp_session.rs +++ b/crates/aionui-db/src/repository/sqlite_acp_session.rs @@ -228,6 +228,21 @@ impl IAcpSessionRepository for SqliteAcpSessionRepository { })?; 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 @@ -426,7 +441,10 @@ mod tests { .await .unwrap_err(); - assert!(matches!(err, DbError::NotFound(_))); + 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()); } From dbb28211d30c187c1a8d5d8098993003f670c999 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 16:09:21 +0800 Subject: [PATCH 135/145] conversation: ignore forged extra user scope --- crates/aionui-conversation/src/service.rs | 8 ++++ .../tests/conversation_crud.rs | 38 +++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/crates/aionui-conversation/src/service.rs b/crates/aionui-conversation/src/service.rs index 784154205..79bffa983 100644 --- a/crates/aionui-conversation/src/service.rs +++ b/crates/aionui-conversation/src/service.rs @@ -762,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 @@ -1971,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() @@ -3640,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) } diff --git a/crates/aionui-conversation/tests/conversation_crud.rs b/crates/aionui-conversation/tests/conversation_crud.rs index b9e5dd97f..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; From bbba3e4ad2d5152407977a8b6cf8ced0ffac814f Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 16:31:00 +0800 Subject: [PATCH 136/145] db: validate user scope migration parent links --- .../aionui-db/migrations/028_user_scope.sql | 85 ++++++++ .../aionui-db/tests/user_scope_migration.rs | 198 ++++++++++++++++++ 2 files changed, 283 insertions(+) diff --git a/crates/aionui-db/migrations/028_user_scope.sql b/crates/aionui-db/migrations/028_user_scope.sql index afc32bb32..346b36ec7 100644 --- a/crates/aionui-db/migrations/028_user_scope.sql +++ b/crates/aionui-db/migrations/028_user_scope.sql @@ -72,6 +72,91 @@ SELECT CASE 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; diff --git a/crates/aionui-db/tests/user_scope_migration.rs b/crates/aionui-db/tests/user_scope_migration.rs index b97b97561..d2fce6973 100644 --- a/crates/aionui-db/tests/user_scope_migration.rs +++ b/crates/aionui-db/tests/user_scope_migration.rs @@ -386,6 +386,204 @@ async fn migration_028_keeps_new_conversation_cron_jobs_unanchored_until_run() { 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_028_preserves_valid_aggregate_child_rows() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + run_migrations_through(&pool, 27).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, 28).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_028_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, 27).await; + sqlx::query(setup_sql).execute(&pool).await.unwrap(); + + let err = run_migration_result(&pool, 28).await.unwrap_err(); + assert!( + err.to_string().contains("CHECK constraint failed"), + "unexpected migration error for {name}: {err}" + ); + } +} + #[tokio::test] async fn migration_028_rejects_channel_session_without_channel_user() { let pool = SqlitePoolOptions::new() From 19fd4e6a9ae05cd4a6168a1559eae289ce3c1e8e Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 16:50:54 +0800 Subject: [PATCH 137/145] system: cover user-scoped settings and providers --- crates/aionui-system/tests/provider_routes.rs | 103 +++++++++++++--- crates/aionui-system/tests/settings_routes.rs | 110 +++++++++++++++--- 2 files changed, 185 insertions(+), 28 deletions(-) diff --git a/crates/aionui-system/tests/provider_routes.rs b/crates/aionui-system/tests/provider_routes.rs index dd0ded32b..4e233346f 100644 --- a/crates/aionui-system/tests/provider_routes.rs +++ b/crates/aionui-system/tests/provider_routes.rs @@ -29,6 +29,7 @@ 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())); @@ -49,15 +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(); - 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(); + 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) } @@ -68,10 +71,14 @@ async fn body_json(resp: axum::response::Response) -> serde_json::Value { } fn get_request(uri: &str) -> Request { + 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: TEST_USER_ID.to_owned(), - username: TEST_USER_ID.to_owned(), + id: user_id.to_owned(), + username: user_id.to_owned(), user_type: UserType::Local, status: UserStatus::Active, }); @@ -79,6 +86,10 @@ fn get_request(uri: &str) -> Request { } fn json_request(method: &str, uri: &str, body: serde_json::Value) -> Request { + 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) @@ -86,8 +97,8 @@ fn json_request(method: &str, uri: &str, body: serde_json::Value) -> Request Request Request { + 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(); req.extensions_mut().insert(CurrentUser { - id: TEST_USER_ID.to_owned(), - username: TEST_USER_ID.to_owned(), + id: user_id.to_owned(), + username: user_id.to_owned(), user_type: UserType::Local, status: UserStatus::Active, }); @@ -215,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; @@ -479,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 51c17540d..49a4b3d0d 100644 --- a/crates/aionui-system/tests/settings_routes.rs +++ b/crates/aionui-system/tests/settings_routes.rs @@ -28,6 +28,7 @@ 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())); @@ -48,15 +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(); - 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(); + 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) } @@ -67,10 +70,14 @@ async fn body_json(resp: axum::response::Response) -> serde_json::Value { } fn get_request(uri: &str) -> Request { + 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: TEST_USER_ID.to_owned(), - username: TEST_USER_ID.to_owned(), + id: user_id.to_owned(), + username: user_id.to_owned(), user_type: UserType::Local, status: UserStatus::Active, }); @@ -78,6 +85,10 @@ fn get_request(uri: &str) -> Request { } fn json_request(method: &str, uri: &str, body: serde_json::Value) -> Request { + 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) @@ -85,8 +96,8 @@ fn json_request(method: &str, uri: &str, body: serde_json::Value) -> Request Date: Wed, 22 Jul 2026 17:06:48 +0800 Subject: [PATCH 138/145] cron: mark unscoped repository paths as system only --- crates/aionui-cron/src/service.rs | 6 +- .../aionui-cron/tests/service_integration.rs | 46 ++++++----- crates/aionui-db/src/repository/cron.rs | 36 +++++---- .../aionui-db/src/repository/sqlite_cron.rs | 80 ++++++++++--------- crates/aionui-db/tests/cron_repository.rs | 68 ++++++++-------- 5 files changed, 128 insertions(+), 108 deletions(-) diff --git a/crates/aionui-cron/src/service.rs b/crates/aionui-cron/src/service.rs index 924521d9b..01a5ab622 100644 --- a/crates/aionui-cron/src/service.rs +++ b/crates/aionui-cron/src/service.rs @@ -496,7 +496,7 @@ impl CronService { warn!(error = %error, "Failed to clean up old cron run records"); } - let rows = match self.repo.list_enabled().await { + let rows = match self.repo.list_enabled_system().await { Ok(rows) => rows, Err(e) => { error!(error = %e, "Failed to load enabled cron jobs"); @@ -538,7 +538,7 @@ impl CronService { } 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"); @@ -652,7 +652,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"); diff --git a/crates/aionui-cron/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index c174ee8f6..b35d29514 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -1262,7 +1262,7 @@ async fn create_job_allows_missing_task_description() { 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); } @@ -1750,7 +1750,7 @@ async fn update_existing_job_to_new_conversation_keeps_owner_anchor() { .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, "conv_mode_switch"); assert!(row.conversation_title.is_none()); @@ -1810,7 +1810,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(), @@ -1855,7 +1855,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())); } @@ -2441,7 +2441,7 @@ async fn run_now_on_running_existing_conversation_returns_active_conversation_wi 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!( @@ -2502,7 +2502,7 @@ 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"); @@ -2534,7 +2534,7 @@ 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")); @@ -2563,7 +2563,7 @@ 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")); } @@ -2620,7 +2620,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")); } @@ -2641,7 +2641,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()); } @@ -2681,7 +2681,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"); @@ -2729,7 +2729,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 * * * *"); @@ -2888,7 +2888,7 @@ 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 * * * *"); @@ -3116,7 +3116,7 @@ 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; @@ -3264,7 +3264,7 @@ async fn cd3b_on_conversation_delete_clears_deleted_workspace_from_jobs() { .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(), @@ -3274,7 +3274,7 @@ async fn cd3b_on_conversation_delete_clears_deleted_workspace_from_jobs() { svc.on_conversation_deleted("u1", &conversation_id).await; assert!(svc.get_job("u1", &job.id).await.is_ok()); - 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!( config.workspace.is_none(), @@ -3332,14 +3332,22 @@ async fn cd3b_on_conversation_delete_only_clears_current_user_jobs() { svc.on_conversation_deleted("u1", &conversation_id).await; - let u1_row = cron_repo.get_by_id("cron_workspace_scope_u1").await.unwrap().unwrap(); + 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("cron_workspace_scope_u2").await.unwrap().unwrap(); + 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(), @@ -3370,7 +3378,7 @@ async fn cd3c_on_conversation_delete_preserves_custom_workspace_on_jobs() { 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())); diff --git a/crates/aionui-db/src/repository/cron.rs b/crates/aionui-db/src/repository/cron.rs index ff17d58a6..b135cd910 100644 --- a/crates/aionui-db/src/repository/cron.rs +++ b/crates/aionui-db/src/repository/cron.rs @@ -70,36 +70,40 @@ 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>; /// 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>; - /// Deletes a cron job by ID. Returns `DbError::NotFound` if absent. - async fn delete(&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>; /// 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 a single cron job by ID, or `None` if not found. - async fn get_by_id(&self, id: &str) -> 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 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>; - /// Returns all cron jobs ordered by creation time ascending. - async fn list_all(&self) -> Result, DbError>; + /// 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>; - /// Returns all enabled cron jobs. - async fn list_enabled(&self) -> Result, DbError>; + /// System-only enabled-job scan used by the scheduler. + async fn list_enabled_system(&self) -> Result, DbError>; - /// Returns all cron jobs for a given conversation. - async fn list_by_conversation(&self, conversation_id: &str) -> 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( @@ -108,9 +112,9 @@ pub trait ICronRepository: Send + Sync { conversation_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 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/sqlite_cron.rs b/crates/aionui-db/src/repository/sqlite_cron.rs index cc5d039b9..91436008d 100644 --- a/crates/aionui-db/src/repository/sqlite_cron.rs +++ b/crates/aionui-db/src/repository/sqlite_cron.rs @@ -67,7 +67,7 @@ impl ICronRepository for SqliteCronRepository { Ok(()) } - async fn update(&self, id: &str, params: &UpdateCronJobParams) -> Result<(), DbError> { + async fn update_system(&self, id: &str, params: &UpdateCronJobParams) -> Result<(), DbError> { self.update_inner(None, id, params).await } @@ -75,7 +75,7 @@ impl ICronRepository for SqliteCronRepository { self.update_inner(Some(user_id), id, params).await } - async fn delete(&self, id: &str) -> Result<(), DbError> { + async fn delete_system(&self, id: &str) -> Result<(), DbError> { let result = sqlx::query("DELETE FROM cron_jobs WHERE id = ?") .bind(id) .execute(&self.pool) @@ -102,7 +102,7 @@ impl ICronRepository for SqliteCronRepository { 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) @@ -119,7 +119,7 @@ impl ICronRepository for SqliteCronRepository { Ok(row) } - async fn list_all(&self) -> Result, DbError> { + 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?; @@ -134,14 +134,14 @@ impl ICronRepository for SqliteCronRepository { Ok(rows) } - async fn list_enabled(&self) -> Result, DbError> { + 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", ) @@ -168,7 +168,7 @@ impl ICronRepository for SqliteCronRepository { Ok(rows) } - async fn delete_by_conversation(&self, conversation_id: &str) -> Result { + 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) @@ -603,7 +603,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); @@ -614,7 +614,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()); } @@ -624,7 +624,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); } @@ -637,7 +637,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"); } @@ -658,11 +658,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"); } @@ -678,9 +678,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); @@ -692,7 +692,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), @@ -702,7 +702,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] @@ -967,9 +973,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()); @@ -980,9 +986,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()); @@ -995,7 +1001,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(_))); } @@ -1004,9 +1010,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); } @@ -1016,15 +1024,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(_))); } @@ -1034,17 +1042,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); } @@ -1061,9 +1069,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")); @@ -1085,7 +1093,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); } @@ -1096,7 +1104,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")); } @@ -1107,7 +1115,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/tests/cron_repository.rs b/crates/aionui-db/tests/cron_repository.rs index 0061178b1..98213b56e 100644 --- a/crates/aionui-db/tests/cron_repository.rs +++ b/crates/aionui-db/tests/cron_repository.rs @@ -102,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); @@ -137,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")); } @@ -152,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()); } @@ -170,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); } @@ -193,10 +193,10 @@ 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); } @@ -271,7 +271,7 @@ async fn scheduler_enabled_scan_uses_job_user_id() { new_conversation.execution_mode = "new_conversation".into(); r.insert(&new_conversation).await.unwrap(); - let enabled = r.list_enabled().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")); @@ -287,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); @@ -307,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")); @@ -323,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(_))); } @@ -331,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(_))); } @@ -405,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"); } @@ -421,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")); } @@ -435,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()); } @@ -444,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()); } @@ -455,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()); } @@ -468,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); } @@ -498,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); @@ -518,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); @@ -535,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")); @@ -550,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"); } From 8d8b7706dd59889c9a3af409d3701f19383c04f7 Mon Sep 17 00:00:00 2001 From: zk <> Date: Wed, 22 Jul 2026 17:11:36 +0800 Subject: [PATCH 139/145] mcp: cover cross-user server isolation --- .../aionui-mcp/tests/service_integration.rs | 53 ++++++++++++++++++- 1 file changed, 51 insertions(+), 2 deletions(-) diff --git a/crates/aionui-mcp/tests/service_integration.rs b/crates/aionui-mcp/tests/service_integration.rs index 33db49ee2..0b93f686d 100644 --- a/crates/aionui-mcp/tests/service_integration.rs +++ b/crates/aionui-mcp/tests/service_integration.rs @@ -6,17 +6,24 @@ 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 { @@ -145,6 +152,48 @@ async fn get_not_found() { 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 // --------------------------------------------------------------------------- From 94b04dfc60873bccef598177c103826501cdb087 Mon Sep 17 00:00:00 2001 From: zk <> Date: Thu, 23 Jul 2026 11:33:27 +0800 Subject: [PATCH 140/145] fix: complete multi-account user scoping --- crates/aionui-app/src/router/routes.rs | 10 +- crates/aionui-app/tests/assistants_e2e.rs | 146 +++++++- crates/aionui-app/tests/extension_e2e.rs | 105 +++++- .../tests/runtime_team_tools_e2e.rs | 140 ++++++++ crates/aionui-assistant/src/service.rs | 20 +- crates/aionui-channel/src/routes.rs | 14 +- crates/aionui-extension/src/registry.rs | 337 +++++++++++++----- crates/aionui-extension/src/routes.rs | 54 ++- crates/aionui-extension/src/skill_service.rs | 20 +- crates/aionui-extension/src/state.rs | 156 +++++++- .../aionui-extension/tests/registry_test.rs | 43 ++- .../tests/skill_integration_test.rs | 42 ++- .../tests/state_persistence_test.rs | 36 ++ 13 files changed, 955 insertions(+), 168 deletions(-) create mode 100644 crates/aionui-app/tests/runtime_team_tools_e2e.rs diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index 03882c54a..f522b8011 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -400,7 +400,7 @@ fn auth_identity_mode(identity_mode: crate::config::IdentityMode) -> AuthIdentit fn is_global_websocket_event(event_name: &str) -> bool { matches!( event_name, - "runtime.statusChanged" | "extensions.state-changed" | "extensions.lifecycle" | "hub.state-changed" + "runtime.statusChanged" | "extensions.lifecycle" | "hub.state-changed" ) } @@ -461,7 +461,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; @@ -490,6 +490,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/tests/assistants_e2e.rs b/crates/aionui-app/tests/assistants_e2e.rs index 67fe09db3..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 // --------------------------------------------------------------------------- @@ -943,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} // =========================================================================== @@ -953,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(); @@ -1322,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 // =========================================================================== @@ -1340,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"); } @@ -1380,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(); @@ -1506,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"); } @@ -1546,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/extension_e2e.rs b/crates/aionui-app/tests/extension_e2e.rs index c620ab23a..c00760ab6 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; @@ -821,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 // --------------------------------------------------------------------------- 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-assistant/src/service.rs b/crates/aionui-assistant/src/service.rs index 6eb16e791..bd989055c 100644 --- a/crates/aionui-assistant/src/service.rs +++ b/crates/aionui-assistant/src/service.rs @@ -1530,7 +1530,7 @@ impl AssistantService { } // Best-effort filesystem cleanup. - self.cleanup_user_assets(id); + self.cleanup_user_assets_for_user(user_id, id); Ok(()) } @@ -2004,10 +2004,6 @@ impl AssistantService { // Internal helpers // ----------------------------------------------------------------------- - fn user_rules_dir(&self) -> PathBuf { - self.user_rules_dir_for_user(DEFAULT_USER_ID) - } - fn user_rules_root_dir(&self) -> PathBuf { self.user_data_dir.join("assistant-rules") } @@ -2016,10 +2012,6 @@ impl AssistantService { self.user_rules_root_dir().join(encode_filename_component(user_id)) } - fn user_skills_dir(&self) -> PathBuf { - self.user_skills_dir_for_user(DEFAULT_USER_ID) - } - fn user_skills_root_dir(&self) -> PathBuf { self.user_data_dir.join("assistant-skills") } @@ -2329,9 +2321,13 @@ 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); + 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); } } diff --git a/crates/aionui-channel/src/routes.rs b/crates/aionui-channel/src/routes.rs index 244c2b36c..5f366df8f 100644 --- a/crates/aionui-channel/src/routes.rs +++ b/crates/aionui-channel/src/routes.rs @@ -140,7 +140,7 @@ async fn get_plugin_status( Extension(user): Extension, ) -> Result>>, ApiError> { let statuses = state.manager.get_plugin_status(&user.id).await?; - let extension_plugins = state.extension_registry.get_channel_plugins().await; + let extension_plugins = state.extension_registry.get_channel_plugins_for_user(&user.id).await; let extension_map: HashMap = extension_plugins .into_iter() @@ -318,7 +318,7 @@ async fn enable_plugin( ) -> 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 @@ -372,7 +372,9 @@ async fn disable_plugin( ) -> 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(&user.id, &req.plugin_id) @@ -407,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, @@ -762,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) 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_service.rs b/crates/aionui-extension/src/skill_service.rs index 89b0d92a0..fe59b2428 100644 --- a/crates/aionui-extension/src/skill_service.rs +++ b/crates/aionui-extension/src/skill_service.rs @@ -1491,14 +1491,6 @@ async fn resolve_skill_source_path_with_repo_for_user( user_id: &str, name: &str, ) -> 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_for_user(user_id, name).await? { let path = PathBuf::from(&row.path); if path.is_dir() { @@ -1510,7 +1502,17 @@ async fn resolve_skill_source_path_with_repo_for_user( 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 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)); } let cron = paths.cron_skills_dir.join(name); if cron.is_dir() { 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/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 2b3ecd105..a9e595e4c 100644 --- a/crates/aionui-extension/tests/skill_integration_test.rs +++ b/crates/aionui-extension/tests/skill_integration_test.rs @@ -6,7 +6,10 @@ use std::path::Path; -use aionui_db::{ISkillRepository, IUserRepository, SqliteSkillRepository, SqliteUserRepository, init_database_memory}; +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, @@ -269,6 +272,43 @@ async fn user_scoped_imports_with_same_name_use_distinct_storage() { 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 // --------------------------------------------------------------------------- From d62c2f5b41f7f24ed04dae36dbcaa52738b7037e Mon Sep 17 00:00:00 2001 From: zk <> Date: Thu, 23 Jul 2026 12:53:49 +0800 Subject: [PATCH 141/145] cron: scope per-job skills to job owner - migrate legacy cron-source skill rows to system_default_user instead of global; only builtin skills stay user_id NULL - stop syncing cron skill dirs into the machine-level catalog and remove the unauthenticated cron-dir fallback from both skill resolvers; cron skills resolve only through a repo row owned by the requesting user - own the skill-row lifecycle in CronService: save_skill upserts an owner-scoped row named cron-{job_id}, delete_skill/remove_job soft-delete it, and init() reconciles rows from disk per job owner while cleaning legacy frontmatter-named duplicates - make cron insert conversation-owner validation mode-aware: enforced for existing-mode rows, skipped for new-conversation rows whose stale anchor is re-validated by the executor before dispatch - fix stale tests never run at full-suite scope: stream_relay skill resolver mock now records the user-scoped hook; extension_e2e skill listing/import tests updated to user-scoped skill storage --- crates/aionui-app/src/router/state.rs | 1 + crates/aionui-app/tests/extension_e2e.rs | 74 ++++++- .../aionui-conversation/src/stream_relay.rs | 8 +- crates/aionui-cron/src/service.rs | 139 +++++++++++- .../aionui-cron/tests/service_integration.rs | 208 ++++++++++++++++-- .../aionui-db/migrations/028_user_scope.sql | 2 +- .../aionui-db/src/repository/sqlite_cron.rs | 11 +- .../aionui-db/tests/user_scope_migration.rs | 9 +- crates/aionui-extension/src/skill_service.rs | 61 ++--- .../tests/cron_skill_resolve_test.rs | 90 +++++++- 10 files changed, 509 insertions(+), 94 deletions(-) diff --git a/crates/aionui-app/src/router/state.rs b/crates/aionui-app/src/router/state.rs index 7c5bdc01a..9ee09e7f8 100644 --- a/crates/aionui-app/src/router/state.rs +++ b/crates/aionui-app/src/router/state.rs @@ -773,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, diff --git a/crates/aionui-app/tests/extension_e2e.rs b/crates/aionui-app/tests/extension_e2e.rs index c00760ab6..8b60bebab 100644 --- a/crates/aionui-app/tests/extension_e2e.rs +++ b/crates/aionui-app/tests/extension_e2e.rs @@ -1416,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)) @@ -1454,13 +1479,29 @@ 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, "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); @@ -1493,13 +1534,36 @@ 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, "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); diff --git a/crates/aionui-conversation/src/stream_relay.rs b/crates/aionui-conversation/src/stream_relay.rs index 25b8a7256..dd6c39c02 100644 --- a/crates/aionui-conversation/src/stream_relay.rs +++ b/crates/aionui-conversation/src/stream_relay.rs @@ -849,6 +849,7 @@ mod tests { #[derive(Default)] struct RecordingSkillResolverForRelay { requested: Mutex>, + requested_user_ids: Mutex>, } #[async_trait::async_trait] @@ -861,7 +862,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() @@ -897,6 +899,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] diff --git a/crates/aionui-cron/src/service.rs b/crates/aionui-cron/src/service.rs index 01a5ab622..a898fdff1 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}; @@ -23,7 +24,10 @@ use crate::events::CronEventEmitter; use crate::error::CronError; use crate::executor::{ExecutionResult, JobExecutor, PreparedRunNow, RETRY_INTERVAL_MS}; use crate::scheduler::{CronScheduler, compute_next_run, compute_next_run_after_occurrence, validate_schedule}; -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, @@ -50,6 +54,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, @@ -62,6 +67,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, @@ -75,6 +81,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, @@ -453,6 +460,7 @@ impl CronService { 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.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"); @@ -496,6 +504,8 @@ impl CronService { warn!(error = %error, "Failed to clean up old cron run records"); } + self.reconcile_skill_rows().await; + let rows = match self.repo.list_enabled_system().await { Ok(rows) => rows, Err(e) => { @@ -537,6 +547,76 @@ 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_system(job_id).await { Ok(Some(r)) => r, @@ -758,6 +838,7 @@ impl CronService { ..Default::default() }; 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?; @@ -766,6 +847,55 @@ impl CronService { Ok(()) } + /// 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 @@ -792,6 +922,7 @@ impl CronService { ..Default::default() }; 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(()) diff --git a/crates/aionui-cron/tests/service_integration.rs b/crates/aionui-cron/tests/service_integration.rs index b35d29514..c63aad457 100644 --- a/crates/aionui-cron/tests/service_integration.rs +++ b/crates/aionui-cron/tests/service_integration.rs @@ -28,8 +28,8 @@ use aionui_db::{ IConversationRepository, ICronRepository, MessagePageParams, MessagePageResult, MessageRowUpdate, MessageSearchRow, SqliteAcpSessionRepository, SqliteAgentMetadataRepository, SqliteAssistantDefinitionRepository, SqliteAssistantOverlayRepository, SqliteAssistantPreferenceRepository, SqliteConversationRepository, - SqliteCronRepository, UpsertAgentMetadataParams, UpsertAssistantDefinitionParams, UpsertAssistantOverlayParams, - UpsertConversationAssistantSnapshotParams, init_database_memory, + SqliteCronRepository, SqliteSkillRepository, UpsertAgentMetadataParams, UpsertAssistantDefinitionParams, + UpsertAssistantOverlayParams, UpsertConversationAssistantSnapshotParams, init_database_memory, models::{ConversationAssistantSnapshotRow, CronJobRow, MessageRow}, }; use aionui_realtime::EventBroadcaster; @@ -840,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) } @@ -851,6 +851,7 @@ 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(); @@ -923,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, @@ -944,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() -> ( @@ -1030,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, @@ -2243,6 +2260,156 @@ async fn sk7_delete_cleans_skill() { 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] @@ -2570,7 +2737,8 @@ async fn create_for_conversation_helper_uses_assistant_metadata_full_auto_mode() #[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 @@ -3311,24 +3479,18 @@ async fn cd3b_on_conversation_delete_only_clears_current_user_jobs() { "test setup should use a workspace ConversationService will delete" ); - cron_repo - .insert(&make_cron_row_with_workspace( - "cron_workspace_scope_u1", - "u1", - &conversation_id, - &deleted_workspace, - )) - .await - .unwrap(); - cron_repo - .insert(&make_cron_row_with_workspace( - "cron_workspace_scope_u2", - "u2", - &conversation_id, - &deleted_workspace, - )) - .await - .unwrap(); + // 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; diff --git a/crates/aionui-db/migrations/028_user_scope.sql b/crates/aionui-db/migrations/028_user_scope.sql index 346b36ec7..eabc25ce2 100644 --- a/crates/aionui-db/migrations/028_user_scope.sql +++ b/crates/aionui-db/migrations/028_user_scope.sql @@ -675,7 +675,7 @@ INSERT INTO skills_new ( ) SELECT id, - CASE WHEN source IN ('builtin', 'cron') THEN NULL ELSE 'system_default_user' END, + CASE WHEN source = 'builtin' THEN NULL ELSE 'system_default_user' END, name, description, path, source, enabled, deleted_at, created_at, updated_at FROM skills; diff --git a/crates/aionui-db/src/repository/sqlite_cron.rs b/crates/aionui-db/src/repository/sqlite_cron.rs index 91436008d..5c6b87d16 100644 --- a/crates/aionui-db/src/repository/sqlite_cron.rs +++ b/crates/aionui-db/src/repository/sqlite_cron.rs @@ -22,8 +22,15 @@ impl SqliteCronRepository { #[async_trait::async_trait] impl ICronRepository for SqliteCronRepository { async fn insert(&self, row: &CronJobRow) -> Result<(), DbError> { - self.validate_conversation_owner(&row.user_id, &row.conversation_id) - .await?; + // 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 (\ diff --git a/crates/aionui-db/tests/user_scope_migration.rs b/crates/aionui-db/tests/user_scope_migration.rs index d2fce6973..da0f09f33 100644 --- a/crates/aionui-db/tests/user_scope_migration.rs +++ b/crates/aionui-db/tests/user_scope_migration.rs @@ -103,7 +103,7 @@ async fn migration_028_adds_user_scope_to_independent_roots() { } #[tokio::test] -async fn migration_028_migrates_cron_skills_as_global_rows() { +async fn migration_028_migrates_cron_skills_to_default_user() { let pool = SqlitePoolOptions::new() .max_connections(1) .connect("sqlite::memory:") @@ -126,7 +126,12 @@ async fn migration_028_migrates_cron_skills_as_global_rows() { .await .unwrap(); assert_eq!(row.get::("source"), "cron"); - assert_eq!(row.get::, _>("user_id"), None); + // 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] diff --git a/crates/aionui-extension/src/skill_service.rs b/crates/aionui-extension/src/skill_service.rs index fe59b2428..31af1b0ad 100644 --- a/crates/aionui-extension/src/skill_service.rs +++ b/crates/aionui-extension/src/skill_service.rs @@ -1282,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, @@ -1478,10 +1480,8 @@ async fn resolve_skill_source_path(paths: &SkillPaths, name: &str) -> Result Result<(), ExtensionError> { - if let Some(existing) = repo.find_by_name_any_for_user(DEFAULT_USER_ID, &skill.name).await? - && existing.source == "user" - && existing.deleted_at.is_none() - && existing.enabled - { - return Ok(()); - } - - repo.upsert_global(UpsertSkillParams { - name: &skill.name, - description: Some(&skill.description), - path: &skill.path, - source, - enabled: true, - }) - .await?; - Ok(()) -} - async fn sync_global_skill_into_repo( repo: &dyn ISkillRepository, skill: &ScannedSkill, @@ -3015,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; @@ -3044,12 +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"); - assert_eq!(scheduled_row.user_id, None); + // 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/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(); } From 33e94f616b50f01a9bb5a0c2ff2bcecaf4c2f569 Mon Sep 17 00:00:00 2001 From: zk <> Date: Thu, 23 Jul 2026 13:17:23 +0800 Subject: [PATCH 142/145] test: authenticate office proxy ssrf e2e coverage Office proxy routes require auth on this branch (preview ports are scoped to the active Core user; unauthenticated access is pinned to 401 by au2_unauthenticated_all_office_endpoints). The four proxy SSRF/root-path tests still asserted 403 without credentials and were tripped by the auth middleware first. They now log in and assert the SSRF port rejection as authenticated users, keeping their original intent. --- crates/aionui-app/tests/office_e2e.rs | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) 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); From c55938a7342aa801838db0df06fcdb960598042d Mon Sep 17 00:00:00 2001 From: zk <> Date: Thu, 23 Jul 2026 13:47:26 +0800 Subject: [PATCH 143/145] test: seed users for websocket e2e and remove sleep races The WS handshake now resolves the token's user against the users table (active row + session generation), so sign_token seeds an active Core user row before signing instead of minting tokens for nonexistent users. Fixed post-connect races by replacing fixed sleeps before client_count assertions and broadcasts with a bounded wait_for_clients poll. --- crates/aionui-app/tests/websocket_e2e.rs | 122 +++++++++++++---------- 1 file changed, 68 insertions(+), 54 deletions(-) diff --git a/crates/aionui-app/tests/websocket_e2e.rs b/crates/aionui-app/tests/websocket_e2e.rs index 01ebf2215..28ec2e0f9 100644 --- a/crates/aionui-app/tests/websocket_e2e.rs +++ b/crates/aionui-app/tests/websocket_e2e.rs @@ -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); @@ -352,14 +378,12 @@ async fn t4_1_broadcast_reaches_all_clients() { #[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"); - let token_b = sign_token(&app, "user-b"); + 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; - - assert_eq!(ws_manager(&app).client_count(), 2); + 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); @@ -379,7 +403,7 @@ async fn t4_1_scoped_event_reaches_only_matching_user() { #[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"); + 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; @@ -396,8 +420,8 @@ async fn t4_1_unscoped_business_event_is_dropped_by_bridge() { #[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"); - let token_b = sign_token(&app, "user-b"); + 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; @@ -429,13 +453,11 @@ 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})); @@ -451,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!({})); @@ -480,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; @@ -495,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; @@ -516,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; @@ -536,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; @@ -560,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; } // =========================================================================== @@ -579,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 { @@ -593,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; } // =========================================================================== @@ -604,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); From dc847432cc3e65f414660382c508dafa19e575c9 Mon Sep 17 00:00:00 2001 From: zk <> Date: Thu, 23 Jul 2026 14:10:49 +0800 Subject: [PATCH 144/145] assistant: migrate legacy locale-less rules on locale reads The layered rule-read fallback migrated legacy-named files only in the user-scoped dir path; the legacy root branch skipped the locale-less retry, so a locale-specific read served the legacy file via the glob last-resort without migrating it to the encoded path. Mirror the scoped fallback in the legacy root branch so the migration side effect applies. Fixes generated_rule_with_requested_locale_falls_back_to_legacy_locale_less_path. --- crates/aionui-assistant/src/service.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/crates/aionui-assistant/src/service.rs b/crates/aionui-assistant/src/service.rs index bd989055c..6fc294a89 100644 --- a/crates/aionui-assistant/src/service.rs +++ b/crates/aionui-assistant/src/service.rs @@ -1844,6 +1844,16 @@ impl AssistantService { 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) } From 564471606e826637cfffbb59f5bd44b7dd673ce8 Mon Sep 17 00:00:00 2001 From: zk <> Date: Thu, 23 Jul 2026 17:54:03 +0800 Subject: [PATCH 145/145] feat(auth): adopt system_default_user data into first external user Upgrade path for AionUi -> AionPro over the same database: legacy data was backfilled to system_default_user by 028_user_scope, so the first AionPro login used to see an empty library. On provision, while the caller is the machine's only external user, re-own every user-scoped row held by system_default_user to it. - Ownership tables are discovered from the live schema (any table with a user_id column) so future migrations cannot be silently missed; a sentinel test pins the convention against the core tables. - Self-idempotent, no marker: the source set empties after adoption, and a second provisioned account closes the window permanently. - UPDATE OR IGNORE skips rows colliding with the new owner's own rows on per-user PK/UNIQUE tables (owner's data stays authoritative). - Global template rows (user_id IS NULL) stay shared. Also mount a credentialed CORS layer for non-local identity mode (mirror origin + allow-credentials + x-csrf-token): the desktop renderer is a cross-origin browser context and preflights were failing with no Access-Control headers. --- crates/aionui-app/src/router/routes.rs | 28 ++- crates/aionui-auth/src/service.rs | 9 + .../aionui-db/src/repository/sqlite_user.rs | 199 ++++++++++++++++++ crates/aionui-db/src/repository/user.rs | 13 ++ 4 files changed, 246 insertions(+), 3 deletions(-) diff --git a/crates/aionui-app/src/router/routes.rs b/crates/aionui-app/src/router/routes.rs index fe4e04373..74bd08832 100644 --- a/crates/aionui-app/src/router/routes.rs +++ b/crates/aionui-app/src/router/routes.rs @@ -6,12 +6,12 @@ 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; @@ -393,7 +393,29 @@ 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) } } diff --git a/crates/aionui-auth/src/service.rs b/crates/aionui-auth/src/service.rs index e7406e41c..a47904f76 100644 --- a/crates/aionui-auth/src/service.rs +++ b/crates/aionui-auth/src/service.rs @@ -61,6 +61,15 @@ impl AuthProvisionService { 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)) } diff --git a/crates/aionui-db/src/repository/sqlite_user.rs b/crates/aionui-db/src/repository/sqlite_user.rs index 4e6026cf9..9bb53829e 100644 --- a/crates/aionui-db/src/repository/sqlite_user.rs +++ b/crates/aionui-db/src/repository/sqlite_user.rs @@ -190,6 +190,60 @@ impl IUserRepository for SqliteUserRepository { 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) @@ -691,4 +745,149 @@ mod tests { 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/user.rs b/crates/aionui-db/src/repository/user.rs index 0ed780bd3..c547b8045 100644 --- a/crates/aionui-db/src/repository/user.rs +++ b/crates/aionui-db/src/repository/user.rs @@ -50,6 +50,19 @@ pub trait IUserRepository: Send + Sync { 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>;