diff --git a/Cargo.lock b/Cargo.lock index 4b9884f..e157a08 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4125,6 +4125,7 @@ dependencies = [ "serde", "serde_json", "serial_test", + "sha2 0.10.9", "shaide-common", "shaide-db", "sqlx", diff --git a/README.md b/README.md index d92a9e8..3568b1b 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,10 @@ The following table describes the current environment variables that we use. | PORT | Yes | Server bind port (default: `8080`) | | VECTOR_DB_URL | Yes | Vector DB URL (default: `http://localhost:6334`) | | DATABASE_URL | Yes | Database URL used by SQLx tooling/migrations | +| KNOWLEDGE_CENTER_URL | Yes | Optional fixed KC origin; see [upload orchestration](documentation/upload-lifecycle.md) | +| KNOWLEDGE_CENTER_INTERNAL_TOKEN | With KC | Dedicated internal service bearer secret, at least 32 bytes | +| S3_WEBHOOK_TOKEN | With KC | Dedicated notification bearer secret, distinct from the internal token and JWT secret | +| KNOWLEDGE_CENTER_ALLOW_LOCAL_HTTP | Yes | Explicit local HTTP opt-in for KC (default: HTTPS required) | | MCP_SERVERS | Yes | Static MCP server list as JSON (unset: Kubernetes discovery), see [MCP servers](#mcp-servers) | | MCP_NAMESPACE | Yes | Namespace watched for MCP pods by Kubernetes discovery | | MCP_LABEL_SELECTOR | Yes | Label selector of MCP pods for Kubernetes discovery | @@ -112,6 +116,10 @@ KC target participation, dedicated signing and service credentials, fresh collec access, targeted admin discovery and supported limits. Ordinary MCP behavior above remains the default. +See [collection and upload orchestration](documentation/upload-lifecycle.md) for +Control Panel API compatibility, KC-owned storage, upload limits and durable +operation recovery. Ordinary startup remains available without KC. + # Authentication Users and admins exchange their username and password for a one-hour JWT: diff --git a/crates/shaide-db/migrations/0004_upload_operations.down.sql b/crates/shaide-db/migrations/0004_upload_operations.down.sql new file mode 100644 index 0000000..5a57244 --- /dev/null +++ b/crates/shaide-db/migrations/0004_upload_operations.down.sql @@ -0,0 +1,3 @@ +-- Explicit downgrade drops only operation fences; settle remote work first. +DROP TABLE file_operations; +DROP TABLE collection_operations; diff --git a/crates/shaide-db/migrations/0004_upload_operations.up.sql b/crates/shaide-db/migrations/0004_upload_operations.up.sql new file mode 100644 index 0000000..6195031 --- /dev/null +++ b/crates/shaide-db/migrations/0004_upload_operations.up.sql @@ -0,0 +1,30 @@ +-- Durable fences: no lease expiry releases a possibly accepted remote write. +CREATE TABLE collection_operations ( + organization_id INTEGER PRIMARY KEY REFERENCES organizations(id), + token TEXT NOT NULL UNIQUE, + kind TEXT NOT NULL, + payload TEXT NOT NULL, + state TEXT NOT NULL DEFAULT 'active' CHECK(state IN ('active','recovery_required')), + error_code TEXT, + created_at TIMESTAMP NOT NULL DEFAULT (DATETIME('now')) +); +CREATE TABLE file_operations ( + file_hash TEXT PRIMARY KEY REFERENCES files(hash), + organization_id INTEGER NOT NULL REFERENCES organizations(id), + owner_id INTEGER NOT NULL REFERENCES users(id), + token TEXT NOT NULL UNIQUE, + state TEXT NOT NULL CHECK(state IN ('presigning','reserved','uploading','uploaded','ingesting','deleting','recovery_required')), + expires_at INTEGER NOT NULL, + query_digest TEXT, + event_received BOOLEAN NOT NULL DEFAULT 0, + error_code TEXT, + created_at TIMESTAMP NOT NULL DEFAULT (DATETIME('now')) +); +CREATE INDEX file_operations_organization_id ON file_operations(organization_id); + +-- Legacy non-final metadata cannot establish whether a remote write finished. +-- Preserve the rows and require scoped reconciliation before destructive work. +INSERT INTO collection_operations (organization_id, token, kind, payload, state, error_code) +SELECT DISTINCT organization_id, 'migration-0004-' || organization_id, + 'legacy_recovery', '{}', 'recovery_required', 'legacy_file_outcome_unknown' +FROM files WHERE status != 'processed'; diff --git a/crates/shaide-db/src/lib.rs b/crates/shaide-db/src/lib.rs index 85ac21b..8b0e1a5 100644 --- a/crates/shaide-db/src/lib.rs +++ b/crates/shaide-db/src/lib.rs @@ -7,6 +7,7 @@ pub mod error; pub mod files; pub mod models; pub mod organizations; +pub mod upload_lifecycle; mod users; use std::path::Path; diff --git a/crates/shaide-db/src/organizations.rs b/crates/shaide-db/src/organizations.rs index 28bd23c..9965f49 100644 --- a/crates/shaide-db/src/organizations.rs +++ b/crates/shaide-db/src/organizations.rs @@ -85,7 +85,7 @@ async fn select_organization<'e>( .ok_or(ShaideDBError::NotFound(Resource::Organization)) } -async fn ensure_embedding_model_exists( +pub(crate) async fn ensure_embedding_model_exists( transaction: &mut Transaction<'static, Sqlite>, embedding_model_id: i64, ) -> Result<(), ShaideDBError> { @@ -105,7 +105,7 @@ async fn ensure_embedding_model_exists( /// Adds the users to the organization, skipping existing memberships. Fails /// without writing anything else if a user does not exist. Returns the number /// of new memberships. -async fn insert_memberships( +pub(crate) async fn insert_memberships( transaction: &mut Transaction<'static, Sqlite>, organization_id: i64, user_ids: &[i64], diff --git a/crates/shaide-db/src/upload_lifecycle.rs b/crates/shaide-db/src/upload_lifecycle.rs new file mode 100644 index 0000000..de1b7e9 --- /dev/null +++ b/crates/shaide-db/src/upload_lifecycle.rs @@ -0,0 +1,350 @@ +//! Durable fences for public orchestration. Transactions end before KC calls. +//! Abandoned or uncertain operations are never reclaimed by elapsed time. +use crate::{ + DbConn, + error::{Resource, ShaideDBError}, + files::InsertFileDao, + organizations::{ + InsertOrganizationDao, begin_write, ensure_embedding_model_exists, insert_memberships, + }, +}; +use sqlx::{Row, Sqlite, Transaction}; + +fn conflict() -> ShaideDBError { + ShaideDBError::InsertFailedOnConflict("operation requires reconciliation".into()) +} + +async fn collection_available( + tx: &mut Transaction<'static, Sqlite>, + id: i64, +) -> Result<(), ShaideDBError> { + let exists: bool = sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM organizations WHERE id=?)") + .bind(id) + .fetch_one(&mut **tx) + .await?; + if !exists { + return Err(ShaideDBError::NotFound(Resource::Organization)); + } + let locked: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM collection_operations WHERE organization_id=?)", + ) + .bind(id) + .fetch_one(&mut **tx) + .await?; + if locked { + return Err(conflict()); + } + Ok(()) +} + +async fn upload_owner_allowed( + tx: &mut Transaction<'static, Sqlite>, + id: i64, + hash: &str, +) -> Result<(), ShaideDBError> { + let allowed:bool=sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM file_operations f JOIN users u ON u.id=f.owner_id JOIN organizations o ON o.id=f.organization_id WHERE f.organization_id=? AND f.file_hash=? AND (u.role='admin' OR (unixepoch(u.expiry)>unixepoch() AND o.can_users_upload AND EXISTS(SELECT 1 FROM organization_memberships m WHERE m.organization_id=o.id AND m.user_id=u.id))))") + .bind(id).bind(hash).fetch_one(&mut **tx).await?; + if !allowed { + return Err(conflict()); + } + Ok(()) +} + +impl DbConn { + /// Restrict access under the same write lock as upload/ingestion admission. + /// Existing collection, file and recovery fences are deliberately untouched. + pub async fn disable_organization_uploads(&self, id: i64) -> Result<(), ShaideDBError> { + let mut tx = begin_write(self).await?; + let n = sqlx::query("UPDATE organizations SET can_users_upload=false, updated_at=DATETIME('now') WHERE id=?") + .bind(id).execute(&mut *tx).await?.rows_affected(); + if n != 1 { + return Err(ShaideDBError::NotFound(Resource::Organization)); + } + tx.commit().await?; + Ok(()) + } + + /// Bind the exact KC-issued query without persisting any signature/credential text. + pub async fn bind_reservation_query( + &self, + id: i64, + hash: &str, + digest: &str, + ) -> Result<(), ShaideDBError> { + let n=sqlx::query("UPDATE file_operations SET state='reserved',query_digest=? WHERE organization_id=? AND file_hash=? AND state='presigning'") + .bind(digest).bind(id).bind(hash).execute(&self.pool).await?.rows_affected(); + if n != 1 { + return Err(conflict()); + } + Ok(()) + } + pub async fn reservation_query_matches( + &self, + id: i64, + hash: &str, + digest: &str, + ) -> Result { + Ok(sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM file_operations WHERE organization_id=? AND file_hash=? AND query_digest=?)") + .bind(id).bind(hash).bind(digest).fetch_one(&self.pool).await?) + } + + pub async fn create_reserved_organization( + &self, + org: InsertOrganizationDao, + token: &str, + ) -> Result { + let mut tx = begin_write(self).await?; + ensure_embedding_model_exists(&mut tx, org.embedding_model_id).await?; + let id=sqlx::query("INSERT INTO organizations(organization_name,description,can_users_upload,embedding_model_id) VALUES(?,?,?,?)") + .bind(org.name).bind(org.description).bind(org.can_users_upload).bind(org.embedding_model_id).execute(&mut *tx).await?.last_insert_rowid(); + insert_memberships(&mut tx, id, &org.user_ids).await?; + sqlx::query("INSERT INTO collection_operations(organization_id,token,kind,payload) VALUES(?,?,'create','{}')").bind(id).bind(token).execute(&mut *tx).await?; + tx.commit().await?; + Ok(id) + } + pub async fn claim_collection_operation( + &self, + id: i64, + token: &str, + kind: &str, + payload: &str, + ) -> Result<(), ShaideDBError> { + let mut tx = begin_write(self).await?; + collection_available(&mut tx, id).await?; + let busy:bool=sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM file_operations WHERE organization_id=?) OR EXISTS(SELECT 1 FROM files WHERE organization_id=? AND status IN ('uploading','processing'))") + .bind(id).bind(id).fetch_one(&mut *tx).await?; + if busy { + return Err(conflict()); + } + sqlx::query( + "INSERT INTO collection_operations(organization_id,token,kind,payload) VALUES(?,?,?,?)", + ) + .bind(id) + .bind(token) + .bind(kind) + .bind(payload) + .execute(&mut *tx) + .await?; + tx.commit().await?; + Ok(()) + } + pub async fn finish_collection_operation( + &self, + id: i64, + token: &str, + delete: bool, + ) -> Result<(), ShaideDBError> { + let mut tx = begin_write(self).await?; + let n=sqlx::query("DELETE FROM collection_operations WHERE organization_id=? AND token=? AND state='active'").bind(id).bind(token).execute(&mut *tx).await?.rows_affected(); + if n != 1 { + return Err(conflict()); + } + if delete { + sqlx::query("DELETE FROM organizations WHERE id=?") + .bind(id) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; + Ok(()) + } + pub async fn quarantine_collection_operation( + &self, + id: i64, + token: &str, + code: &str, + ) -> Result<(), ShaideDBError> { + sqlx::query("UPDATE collection_operations SET state='recovery_required',error_code=? WHERE organization_id=? AND token=?").bind(code).bind(id).bind(token).execute(&self.pool).await?; + Ok(()) + } + /// Authorization and metadata reservation are one write transaction. + pub async fn reserve_upload( + &self, + file: InsertFileDao, + owner: i64, + token: &str, + expiry: i64, + ) -> Result<(), ShaideDBError> { + let mut tx = begin_write(self).await?; + collection_available(&mut tx, file.organization_id).await?; + let allowed:bool=sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM users u JOIN organizations o ON o.id=? WHERE u.id=? AND (u.role='admin' OR (unixepoch(u.expiry)>unixepoch() AND o.can_users_upload AND EXISTS(SELECT 1 FROM organization_memberships m WHERE m.organization_id=o.id AND m.user_id=u.id))))") + .bind(file.organization_id).bind(owner).fetch_one(&mut *tx).await?; + if !allowed { + return Err(conflict()); + } + sqlx::query("INSERT INTO files(hash,organization_id,file_name,mime_type) VALUES(?,?,?,?)") + .bind(&file.hash) + .bind(file.organization_id) + .bind(file.file_name) + .bind(file.mime_type) + .execute(&mut *tx) + .await?; + sqlx::query("INSERT INTO file_operations(file_hash,organization_id,owner_id,token,state,expires_at) VALUES(?,?,?,?,'presigning',unixepoch()+?)").bind(file.hash).bind(file.organization_id).bind(owner).bind(token).bind(expiry).execute(&mut *tx).await?; + tx.commit().await?; + Ok(()) + } + pub async fn confirm_reservation(&self, id: i64, hash: &str) -> Result<(), ShaideDBError> { + self.transition_file_operation(id, hash, "presigning", "reserved") + .await + } + pub async fn transition_file_operation( + &self, + id: i64, + hash: &str, + from: &str, + to: &str, + ) -> Result<(), ShaideDBError> { + let n=sqlx::query("UPDATE file_operations SET state=? WHERE organization_id=? AND file_hash=? AND state=?").bind(to).bind(id).bind(hash).bind(from).execute(&self.pool).await?.rows_affected(); + if n != 1 { + return Err(conflict()); + } + Ok(()) + } + pub async fn start_upload(&self, id: i64, hash: &str) -> Result<(), ShaideDBError> { + let mut tx = begin_write(self).await?; + collection_available(&mut tx, id).await?; + let n=sqlx::query("UPDATE file_operations SET state='uploading' WHERE organization_id=? AND file_hash=? AND state='reserved' AND expires_at>unixepoch() AND EXISTS(SELECT 1 FROM users u JOIN organizations o ON o.id=file_operations.organization_id WHERE u.id=file_operations.owner_id AND (u.role='admin' OR (unixepoch(u.expiry)>unixepoch() AND o.can_users_upload AND EXISTS(SELECT 1 FROM organization_memberships m WHERE m.organization_id=o.id AND m.user_id=u.id))))") + .bind(id).bind(hash).execute(&mut *tx).await?.rows_affected(); + if n != 1 { + return Err(conflict()); + } + tx.commit().await?; + Ok(()) + } + pub async fn record_upload_event(&self, id: i64, hash: &str) -> Result { + let mut tx = begin_write(self).await?; + collection_available(&mut tx, id).await?; + let processed:bool=sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM files WHERE organization_id=? AND hash=? AND status='processed')").bind(id).bind(hash).fetch_one(&mut *tx).await?; + if processed { + tx.commit().await?; + return Ok("processed".into()); + } + upload_owner_allowed(&mut tx, id, hash).await?; + let state:Option=sqlx::query_scalar("UPDATE file_operations SET event_received=1 WHERE organization_id=? AND file_hash=? AND state IN ('uploading','uploaded','ingesting') RETURNING state").bind(id).bind(hash).fetch_optional(&mut *tx).await?; + let state = state.ok_or_else(conflict)?; + tx.commit().await?; + Ok(state) + } + pub async fn upload_event_received(&self, id: i64, hash: &str) -> Result { + Ok(sqlx::query_scalar( + "SELECT event_received FROM file_operations WHERE organization_id=? AND file_hash=?", + ) + .bind(id) + .bind(hash) + .fetch_optional(&self.pool) + .await? + .unwrap_or(false)) + } + pub async fn claim_ingestion(&self, id: i64, hash: &str) -> Result { + let mut tx = begin_write(self).await?; + collection_available(&mut tx, id).await?; + let exists: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM file_operations WHERE organization_id=? AND file_hash=?)", + ) + .bind(id) + .bind(hash) + .fetch_one(&mut *tx) + .await?; + if !exists { + tx.commit().await?; + return Ok(false); + } + upload_owner_allowed(&mut tx, id, hash).await?; + let n=sqlx::query("UPDATE file_operations SET state='ingesting' WHERE organization_id=? AND file_hash=? AND state='uploaded' AND event_received=1").bind(id).bind(hash).execute(&mut *tx).await?.rows_affected(); + if n == 1 { + let updated=sqlx::query("UPDATE files SET status='processing',updated_at=DATETIME('now') WHERE organization_id=? AND hash=? AND status='uploading'").bind(id).bind(hash).execute(&mut *tx).await?.rows_affected(); + if updated != 1 { + return Err(conflict()); + } + } + tx.commit().await?; + Ok(n == 1) + } + pub async fn complete_ingestion( + &self, + id: i64, + hash: &str, + size: i64, + pages: i64, + ) -> Result<(), ShaideDBError> { + let mut tx = begin_write(self).await?; + let n=sqlx::query("DELETE FROM file_operations WHERE organization_id=? AND file_hash=? AND state='ingesting'").bind(id).bind(hash).execute(&mut *tx).await?.rows_affected(); + if n != 1 { + return Err(conflict()); + } + let updated=sqlx::query("UPDATE files SET status='processed',size=?,pages=?,error_code=NULL,updated_at=DATETIME('now') WHERE organization_id=? AND hash=? AND status='processing'").bind(size).bind(pages).bind(id).bind(hash).execute(&mut *tx).await?.rows_affected(); + if updated != 1 { + return Err(conflict()); + } + tx.commit().await?; + Ok(()) + } + pub async fn quarantine_file_operation( + &self, + id: i64, + hash: &str, + code: &str, + ) -> Result<(), ShaideDBError> { + let mut tx = begin_write(self).await?; + let updated=sqlx::query("UPDATE file_operations SET state='recovery_required',error_code=? WHERE organization_id=? AND file_hash=?").bind(code).bind(id).bind(hash).execute(&mut *tx).await?.rows_affected(); + if updated != 1 { + return Err(conflict()); + } + sqlx::query("UPDATE files SET status=CASE WHEN status IN ('uploading','uploading_failed') THEN 'uploading_failed' ELSE 'processing_failed' END,error_code=?,updated_at=DATETIME('now') WHERE organization_id=? AND hash=?").bind(code).bind(id).bind(hash).execute(&mut *tx).await?; + tx.commit().await?; + Ok(()) + } + pub async fn claim_file_delete( + &self, + id: i64, + hash: &str, + token: &str, + owner: i64, + ) -> Result<(), ShaideDBError> { + let mut tx = begin_write(self).await?; + collection_available(&mut tx, id).await?; + let allowed:bool=sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM users u JOIN organizations o ON o.id=? WHERE u.id=? AND (u.role='admin' OR (unixepoch(u.expiry)>unixepoch() AND o.can_users_upload AND EXISTS(SELECT 1 FROM organization_memberships m WHERE m.organization_id=o.id AND m.user_id=u.id))))") + .bind(id).bind(owner).fetch_one(&mut *tx).await?; + if !allowed { + return Err(conflict()); + } + let row = sqlx::query("SELECT status FROM files WHERE organization_id=? AND hash=?") + .bind(id) + .bind(hash) + .fetch_optional(&mut *tx) + .await? + .ok_or(ShaideDBError::NotFound(Resource::File))?; + if matches!( + row.get::("status").as_str(), + "uploading" | "processing" + ) { + return Err(conflict()); + } + // A failed file may still have a remote writer: never override its fence. + let n=sqlx::query("INSERT INTO file_operations(file_hash,organization_id,owner_id,token,state,expires_at) VALUES(?,?,?,?,'deleting',0) ON CONFLICT(file_hash) DO NOTHING").bind(hash).bind(id).bind(owner).bind(token).execute(&mut *tx).await?.rows_affected(); + if n != 1 { + return Err(conflict()); + } + tx.commit().await?; + Ok(()) + } + pub async fn complete_file_delete( + &self, + id: i64, + hash: &str, + token: &str, + ) -> Result<(), ShaideDBError> { + let mut tx = begin_write(self).await?; + let n=sqlx::query("DELETE FROM file_operations WHERE organization_id=? AND file_hash=? AND token=? AND state='deleting'").bind(id).bind(hash).bind(token).execute(&mut *tx).await?.rows_affected(); + if n != 1 { + return Err(conflict()); + } + sqlx::query("DELETE FROM files WHERE organization_id=? AND hash=?") + .bind(id) + .bind(hash) + .execute(&mut *tx) + .await?; + tx.commit().await?; + Ok(()) + } +} diff --git a/crates/shaide-db/tests/upload_lifecycle.rs b/crates/shaide-db/tests/upload_lifecycle.rs new file mode 100644 index 0000000..e3dc67b --- /dev/null +++ b/crates/shaide-db/tests/upload_lifecycle.rs @@ -0,0 +1,266 @@ +use shaide_db::{ + DbConn, embedding_models::InsertEmbeddingModelDao, organizations::InsertOrganizationDao, +}; +use temp_testdir::TempDir; + +#[tokio::test] +async fn durable_collection_claim_is_exclusive_and_survives_new_connection() { + let temp = TempDir::default(); + let path = temp.join("lifecycle.sqlite"); + let db = DbConn::new(&path).await.unwrap(); + let model = db + .insert_embedding_model(InsertEmbeddingModelDao { + name: "synthetic".into(), + url: "http://invalid.example".into(), + vector_size: 8, + platform: None, + api_schema: None, + daily_input_token_limit: None, + }) + .await + .unwrap(); + let id = db + .create_organization(InsertOrganizationDao { + name: "test".into(), + description: "".into(), + can_users_upload: true, + embedding_model_id: model, + user_ids: vec![], + }) + .await + .unwrap(); + sqlx::query("INSERT INTO collection_operations (organization_id, token, kind, payload) VALUES (?, 'first', 'create', '{}')").bind(id).execute(&db.pool).await.unwrap(); + let other = DbConn::new(&path).await.unwrap(); + assert!(sqlx::query("INSERT INTO collection_operations (organization_id, token, kind, payload) VALUES (?, 'second', 'delete', '{}')").bind(id).execute(&other.pool).await.is_err()); + let token: String = + sqlx::query_scalar("SELECT token FROM collection_operations WHERE organization_id=?") + .bind(id) + .fetch_one(&other.pool) + .await + .unwrap(); + assert_eq!(token, "first"); +} + +#[tokio::test] +async fn file_reservation_fences_collection_mutation_and_uncertain_upload() { + let temp = TempDir::default(); + let db = DbConn::new(&temp.join("fences.sqlite")).await.unwrap(); + let model = db + .insert_embedding_model(InsertEmbeddingModelDao { + name: "synthetic".into(), + url: "http://invalid.example".into(), + vector_size: 8, + platform: None, + api_schema: None, + daily_input_token_limit: None, + }) + .await + .unwrap(); + let owner = db + .create_user( + "synthetic-owner".into(), + "unused".into(), + chrono::Utc::now() + chrono::Duration::days(1), + ) + .await + .unwrap(); + let id = db + .create_organization(InsertOrganizationDao { + name: "test".into(), + description: "".into(), + can_users_upload: true, + embedding_model_id: model, + user_ids: vec![owner], + }) + .await + .unwrap(); + let hash = "a".repeat(64); + db.reserve_upload( + shaide_db::files::InsertFileDao { + hash: hash.clone(), + organization_id: id, + file_name: "synthetic.txt".into(), + mime_type: "text/plain".into(), + }, + owner, + "reservation", + 300, + ) + .await + .unwrap(); + db.confirm_reservation(id, &hash).await.unwrap(); + assert!( + db.claim_collection_operation(id, "delete", "delete", "{}") + .await + .is_err() + ); + let other = DbConn::new(&temp.join("fences.sqlite")).await.unwrap(); + let (a, b) = tokio::join!(db.start_upload(id, &hash), other.start_upload(id, &hash)); + assert_ne!(a.is_ok(), b.is_ok()); + other + .quarantine_file_operation(id, &hash, "upload_outcome_unknown") + .await + .unwrap(); + sqlx::query("UPDATE file_operations SET expires_at=0") + .execute(&other.pool) + .await + .unwrap(); + assert!(db.start_upload(id, &hash).await.is_err()); + assert!( + db.claim_file_delete(id, &hash, "delete", owner) + .await + .is_err() + ); + assert!( + db.claim_collection_operation(id, "delete", "delete", "{}") + .await + .is_err() + ); +} +#[tokio::test] +async fn reservation_and_ingestion_recheck_current_owner_access() { + let temp = TempDir::default(); + let db = DbConn::new(&temp.join("owner.sqlite")).await.unwrap(); + let model = db + .insert_embedding_model(InsertEmbeddingModelDao { + name: "synthetic".into(), + url: "http://invalid.example".into(), + vector_size: 8, + platform: None, + api_schema: None, + daily_input_token_limit: None, + }) + .await + .unwrap(); + let owner = db + .create_user( + "synthetic-owner".into(), + "unused".into(), + chrono::Utc::now() + chrono::Duration::days(1), + ) + .await + .unwrap(); + let stranger = db + .create_user( + "synthetic-stranger".into(), + "unused".into(), + chrono::Utc::now() + chrono::Duration::days(1), + ) + .await + .unwrap(); + let id = db + .create_organization(InsertOrganizationDao { + name: "synthetic".into(), + description: "".into(), + can_users_upload: true, + embedding_model_id: model, + user_ids: vec![owner], + }) + .await + .unwrap(); + let file = shaide_db::files::InsertFileDao { + hash: "b".repeat(64), + organization_id: id, + file_name: "synthetic.pdf".into(), + mime_type: "application/pdf".into(), + }; + assert!( + db.reserve_upload(file.clone(), stranger, "denied", 300) + .await + .is_err() + ); + assert!(db.list_files(id).await.unwrap().is_empty()); + db.reserve_upload(file.clone(), owner, "owner", 300) + .await + .unwrap(); + db.bind_reservation_query(id, &file.hash, "fingerprint") + .await + .unwrap(); + assert!( + db.reservation_query_matches(id, &file.hash, "fingerprint") + .await + .unwrap() + ); + assert!( + !db.reservation_query_matches(id, &file.hash, "altered") + .await + .unwrap() + ); + // Real SQL change simulates revocation in another process, bypassing the public API's own fence. + db.remove_organization_members(id, &[owner]).await.unwrap(); + assert!(db.start_upload(id, &file.hash).await.is_err()); + db.add_organization_members(id, &[owner]).await.unwrap(); + db.start_upload(id, &file.hash).await.unwrap(); + db.transition_file_operation(id, &file.hash, "uploading", "uploaded") + .await + .unwrap(); + db.remove_organization_members(id, &[owner]).await.unwrap(); + assert!(db.record_upload_event(id, &file.hash).await.is_err()); + assert!(db.claim_ingestion(id, &file.hash).await.is_err()); + db.add_organization_members(id, &[owner]).await.unwrap(); + db.record_upload_event(id, &file.hash).await.unwrap(); + assert!(db.claim_ingestion(id, &file.hash).await.unwrap()); + assert!(!db.claim_ingestion(id, &file.hash).await.unwrap()); + db.complete_ingestion(id, &file.hash, 9, 1).await.unwrap(); + // A concurrent callback may finish before the public upload checks the event flag. + assert!(!db.upload_event_received(id, &file.hash).await.unwrap()); +} +#[tokio::test] +async fn migration_preserves_legacy_rows_and_quarantines_unsettled_work() { + use sqlx::{Connection, Executor}; + let temp = TempDir::default(); + let path = temp.join("legacy.sqlite"); + let mut conn = + sqlx::SqliteConnection::connect(&format!("sqlite://{}?mode=rwc", path.display())) + .await + .unwrap(); + conn.execute(include_str!("../migrations/0001_initial_schema.up.sql")) + .await + .unwrap(); + conn.execute(include_str!( + "../migrations/0002_embedding_usage_limits.up.sql" + )) + .await + .unwrap(); + conn.execute(include_str!( + "../migrations/0003_organizations_and_files.up.sql" + )) + .await + .unwrap(); + sqlx::query("INSERT INTO embedding_models(name,url,vector_size) VALUES('synthetic','http://invalid.example',8)").execute(&mut conn).await.unwrap(); + sqlx::query("INSERT INTO organizations(id,organization_name,embedding_model_id) VALUES(1,'unsettled',1),(2,'settled',1)").execute(&mut conn).await.unwrap(); + sqlx::query("INSERT INTO files(hash,organization_id,file_name,mime_type,status) VALUES('legacy-upload',1,'synthetic.pdf','application/pdf','uploading'),('legacy-ready',2,'synthetic.pdf','application/pdf','processed')").execute(&mut conn).await.unwrap(); + conn.execute(include_str!("../migrations/0004_upload_operations.up.sql")) + .await + .unwrap(); + let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM files") + .fetch_one(&mut conn) + .await + .unwrap(); + assert_eq!(count, 2); + let locked: Vec = sqlx::query_scalar( + "SELECT organization_id FROM collection_operations WHERE state='recovery_required'", + ) + .fetch_all(&mut conn) + .await + .unwrap(); + assert_eq!(locked, vec![1]); + conn.execute(include_str!( + "../migrations/0004_upload_operations.down.sql" + )) + .await + .unwrap(); + let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM files") + .fetch_one(&mut conn) + .await + .unwrap(); + assert_eq!(count, 2); + conn.execute(include_str!("../migrations/0004_upload_operations.up.sql")) + .await + .unwrap(); + let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM collection_operations") + .fetch_one(&mut conn) + .await + .unwrap(); + assert_eq!(count, 1); +} diff --git a/crates/shaide/Cargo.toml b/crates/shaide/Cargo.toml index 4f8bfd0..08b8c18 100644 --- a/crates/shaide/Cargo.toml +++ b/crates/shaide/Cargo.toml @@ -26,7 +26,7 @@ chrono.workspace = true axum-prometheus = "0.10.0" uuid.workspace = true color-eyre = { version = "0.6.3" } -reqwest.workspace = true +reqwest = {workspace = true, features = ["stream"]} async-openai.workspace = true base64.workspace = true imagesize.workspace = true @@ -70,3 +70,6 @@ ignored = ["openssl"] [dependencies.aws-lc-rs] version = "1.18.0" + +[dependencies.sha2] +version = "0.10" diff --git a/crates/shaide/src/config.rs b/crates/shaide/src/config.rs index dc7e414..7bf07e6 100644 --- a/crates/shaide/src/config.rs +++ b/crates/shaide/src/config.rs @@ -133,6 +133,7 @@ pub struct RunTimeConfig { pub mcp_label_selector: String, pub mcp_servers: Option>, pub mcp_collection_context_targets: crate::services::collection_context::ContextTargets, + pub knowledge_center: Option, pub is_trial: bool, pub webapp_url: Option, } @@ -181,6 +182,26 @@ impl RunTimeConfig { &jwt_secret, ) .unwrap_or_else(|error| panic!("MCP collection context configuration: {error}")); + let knowledge_center = crate::services::knowledge_center::KcConfig::parse( + env_map.var("KNOWLEDGE_CENTER_URL").ok().as_deref(), + env_map + .var("KNOWLEDGE_CENTER_INTERNAL_TOKEN") + .ok() + .as_deref(), + env_map.var("S3_WEBHOOK_TOKEN").ok().as_deref(), + env_map + .var("KNOWLEDGE_CENTER_ALLOW_LOCAL_HTTP") + .ok() + .as_deref() + == Some("true"), + ) + .unwrap_or_else(|e| panic!("KC configuration: {e}")); + if knowledge_center + .as_ref() + .is_some_and(|config| config.uses_token(&jwt_secret)) + { + panic!("KC tokens must differ from JWT_SECRET"); + } let is_trial = env_map .var("TRIAL") .map(|s| s.eq_ignore_ascii_case("true")) @@ -198,6 +219,7 @@ impl RunTimeConfig { mcp_label_selector, mcp_servers, mcp_collection_context_targets, + knowledge_center, is_trial, webapp_url, } @@ -232,6 +254,24 @@ impl RunTimeConfig { &jwt_secret, ) .unwrap_or_else(|error| panic!("MCP collection context configuration: {error}")); + let knowledge_center = crate::services::knowledge_center::KcConfig::parse( + std::env::var("KNOWLEDGE_CENTER_URL").ok().as_deref(), + std::env::var("KNOWLEDGE_CENTER_INTERNAL_TOKEN") + .ok() + .as_deref(), + std::env::var("S3_WEBHOOK_TOKEN").ok().as_deref(), + std::env::var("KNOWLEDGE_CENTER_ALLOW_LOCAL_HTTP") + .ok() + .as_deref() + == Some("true"), + ) + .unwrap_or_else(|e| panic!("KC configuration: {e}")); + if knowledge_center + .as_ref() + .is_some_and(|config| config.uses_token(&jwt_secret)) + { + panic!("KC tokens must differ from JWT_SECRET"); + } let is_trial = std::env::var("TRIAL") .map(|s| s.eq_ignore_ascii_case("true")) .unwrap_or_default(); @@ -248,6 +288,7 @@ impl RunTimeConfig { mcp_label_selector, mcp_servers, mcp_collection_context_targets, + knowledge_center, is_trial, webapp_url, } diff --git a/crates/shaide/src/main.rs b/crates/shaide/src/main.rs index c65613b..d15cf20 100644 --- a/crates/shaide/src/main.rs +++ b/crates/shaide/src/main.rs @@ -64,7 +64,17 @@ fn install_rustls_crypto_provider() { } pub async fn api_router(health_state: Arc, db: DbConn) -> Router { + let kc = get_environment_config() + .knowledge_center + .clone() + .map(crate::services::knowledge_center::KcClient::new) + .transpose() + .expect("KC transport configuration"); Router::new() + .merge(crate::routes::knowledge_center::knowledge_center_router( + db.clone(), + kc, + )) .merge(crate::routes::health::health_router(health_state)) .merge(crate::routes::trial::trial_router()) .merge(crate::routes::completions::completion_router(db.clone())) diff --git a/crates/shaide/src/openapi.rs b/crates/shaide/src/openapi.rs index 95429d8..53113ce 100644 --- a/crates/shaide/src/openapi.rs +++ b/crates/shaide/src/openapi.rs @@ -8,6 +8,17 @@ use utoipa_swagger_ui::SwaggerUi; #[derive(OpenApi)] #[openapi( paths( + crate::routes::knowledge_center::list, + crate::routes::knowledge_center::create, + crate::routes::knowledge_center::delete_collection, + crate::routes::knowledge_center::update, + crate::routes::knowledge_center::add_members, + crate::routes::knowledge_center::remove_members, + crate::routes::knowledge_center::presign, + crate::routes::knowledge_center::delete_file, + crate::routes::knowledge_center::stale_files, + crate::routes::knowledge_center::event, + crate::routes::knowledge_center::upload, crate::routes::health::health, crate::routes::trial::is_trial, crate::routes::users::login, @@ -64,6 +75,17 @@ impl Modify for SecurityAddon { return; }; + components.add_security_scheme( + "storage_notification_token", + SecurityScheme::Http( + HttpBuilder::new() + .scheme(HttpAuthScheme::Bearer) + .description(Some( + "Dedicated storage notification secret; not an end-user JWT.", + )) + .build(), + ), + ); components.add_security_scheme( "bearer_token", SecurityScheme::Http( diff --git a/crates/shaide/src/routes/knowledge_center.rs b/crates/shaide/src/routes/knowledge_center.rs new file mode 100644 index 0000000..552a8a8 --- /dev/null +++ b/crates/shaide/src/routes/knowledge_center.rs @@ -0,0 +1,740 @@ +//! Control Panel-compatible public metadata and authenticated KC forwarding. +use crate::{ + error::ShaideError, + middlewares::Principal, + services::knowledge_center::{KcClient, MAX_UPLOAD_BYTES}, +}; +use axum::{ + Extension, Json, Router, + body::Body, + extract::{OriginalUri, Path, State}, + routing, +}; +use http::{HeaderMap, Method, StatusCode}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use shaide_common::api::organizations::*; +use shaide_db::{ + DbConn, UserRole, + files::{FileStatus, InsertFileDao}, + organizations::{InsertOrganizationDao, UpdateOrganizationDao}, +}; + +#[derive(Clone)] +pub(crate) struct Integration(Option); +impl Integration { + fn client(&self) -> Result<&KcClient, ShaideError> { + self.0.as_ref().ok_or_else(|| { + reject( + StatusCode::SERVICE_UNAVAILABLE, + "knowledge_center_unavailable", + ) + }) + } +} +fn reject(status: StatusCode, code: &str) -> ShaideError { + ShaideError::request_rejection(status, code.replace('_', " "), code) +} +fn admin(p: &Principal) -> Result<(), ShaideError> { + if p.role == UserRole::Admin { + Ok(()) + } else { + Err(reject(StatusCode::FORBIDDEN, "admin_required")) + } +} +fn query_digest(query: &str) -> String { + use sha2::Digest; + format!("{:x}", sha2::Sha256::digest(query.as_bytes())) +} +fn valid_hash(hash: &str) -> bool { + hash.len() == 64 + && hash + .bytes() + .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)) +} +fn collection_id(name: &str) -> Result { + let id = name + .strip_prefix("collection-") + .and_then(|s| s.parse::().ok()) + .filter(|id| *id > 0) + .ok_or_else(|| reject(StatusCode::BAD_REQUEST, "invalid_collection"))?; + if collection_name(id) != name { + return Err(reject(StatusCode::BAD_REQUEST, "invalid_collection")); + } + Ok(id) +} +async fn upload_access(db: &DbConn, p: &Principal, id: i64) -> Result<(), ShaideError> { + if id <= 0 { + return Err(reject(StatusCode::BAD_REQUEST, "invalid_collection")); + } + let org = db.get_organization(id).await?; + if p.role != UserRole::Admin + && (!org.can_users_upload || !db.is_organization_member(id, p.user.id).await?) + { + return Err(reject(StatusCode::FORBIDDEN, "collection_upload_forbidden")); + } + Ok(()) +} +#[utoipa::path(get, path="/v1/rag/organization-collection",tag="rag",responses((status=200,description="Acknowledged operation",body=OrganizationsResponse),(status=409,description="Conflict or reconciliation required"),(status=503,description="KC unavailable or uncertain outcome")),security(("bearer_token"=[])))] +pub(crate) async fn list( + p: Principal, + State(db): State, +) -> Result, ShaideError> { + let orgs = if p.role == UserRole::Admin { + db.list_organizations().await? + } else { + db.list_organizations_for_user(p.user.id).await? + }; + let mut organization = Vec::new(); + for o in orgs { + let files = db + .list_files(o.id) + .await? + .into_iter() + .map(|f| OrganizationFileResponse { + hash: f.hash, + name: f.file_name, + size: f.size, + mime_type: f.mime_type, + status: f.status.into(), + uploaded_at: f.uploaded_at, + }) + .collect(); + let users = if p.role == UserRole::Admin { + Some(db.list_organization_member_ids(o.id).await?) + } else { + None + }; + organization.push(OrganizationResponse { + id: o.id, + name: o.name, + description: o.description, + can_users_upload: o.can_users_upload, + users, + files, + embedding_model_id: o.embedding_model_id, + collection_name: collection_name(o.id), + bucket_name: collection_name(o.id), + created_at: o.created_at, + updated_at: o.updated_at, + }); + } + Ok(Json(OrganizationsResponse { organization })) +} +fn validate_name(name: &str) -> Result<(), ShaideError> { + if name.trim().is_empty() || name.len() > 255 || name.chars().any(char::is_control) { + Err(reject(StatusCode::BAD_REQUEST, "invalid_name")) + } else { + Ok(()) + } +} +#[utoipa::path(post, path="/v1/rag/organization-collection",tag="rag",request_body=CreateOrganizationRequest,responses((status=200,description="Acknowledged operation",body=CreateOrganizationResponse),(status=409,description="Conflict or reconciliation required"),(status=503,description="KC unavailable or uncertain outcome")),security(("bearer_token"=[])))] +pub(crate) async fn create( + p: Principal, + State(db): State, + Extension(kc): Extension, + Json(r): Json, +) -> Result, ShaideError> { + admin(&p)?; + validate_name(&r.organization_name)?; + if r.organization_description.len() > 8192 { + return Err(reject(StatusCode::BAD_REQUEST, "description_too_long")); + } + let kc = kc.client()?; + let model = db.get_embedding_model(r.embedding_model_id).await?; + if model.vector_size <= 0 { + return Err(reject( + StatusCode::BAD_REQUEST, + "invalid_embedding_dimensions", + )); + } + let token = uuid::Uuid::new_v4().to_string(); + let id = db + .create_reserved_organization( + InsertOrganizationDao { + name: r.organization_name, + description: r.organization_description, + can_users_upload: r.can_users_upload, + embedding_model_id: r.embedding_model_id, + user_ids: r.users, + }, + &token, + ) + .await?; + let result = kc + .collection_ack( + Method::PUT, + &format!("/internal/v1/collections/{}", collection_name(id)), + Some(json!({"vector_size":model.vector_size})), + ) + .await; + if let Err(e) = result { + db.quarantine_collection_operation(id, &token, e.code) + .await?; + return Err(e.public()); + } + db.finish_collection_operation(id, &token, false).await?; + Ok(Json(CreateOrganizationResponse { id })) +} +#[utoipa::path(delete, path="/v1/rag/organization-collection",tag="rag",request_body=DeleteOrganizationRequest,responses((status=200,description="Acknowledged operation"),(status=409,description="Conflict or reconciliation required"),(status=503,description="KC unavailable or uncertain outcome")),security(("bearer_token"=[])))] +pub(crate) async fn delete_collection( + p: Principal, + State(db): State, + Extension(kc): Extension, + Json(r): Json, +) -> Result, ShaideError> { + admin(&p)?; + let kc = kc.client()?; + let token = uuid::Uuid::new_v4().to_string(); + db.claim_collection_operation(r.id, &token, "delete", "{}") + .await?; + if let Err(e) = kc + .collection_ack( + Method::DELETE, + &format!("/internal/v1/collections/{}", collection_name(r.id)), + None, + ) + .await + { + db.quarantine_collection_operation(r.id, &token, e.code) + .await?; + return Err(e.public()); + } + db.finish_collection_operation(r.id, &token, true).await?; + Ok(Json(json!({}))) +} +#[utoipa::path(patch, path="/v1/rag/organization-collection",tag="rag",request_body=UpdateOrganizationRequest,responses((status=200,description="Acknowledged operation"),(status=409,description="Conflict or reconciliation required"),(status=503,description="KC unavailable or uncertain outcome")),security(("bearer_token"=[])))] +pub(crate) async fn update( + p: Principal, + State(db): State, + Extension(kc): Extension, + Json(r): Json, +) -> Result, ShaideError> { + admin(&p)?; + if let Some(name) = &r.name { + validate_name(name)?; + } + if r.description.as_ref().is_some_and(|v| v.len() > 8192) { + return Err(reject(StatusCode::BAD_REQUEST, "description_too_long")); + } + // Only a pure restriction bypasses operation exclusion; mixed updates and + // grants still require the durable collection fence below. + if r.can_users_upload == Some(false) + && r.name.is_none() + && r.description.is_none() + && r.embedding_model_id.is_none() + { + db.disable_organization_uploads(r.organization_id).await?; + return Ok(Json(json!({}))); + } + let token = uuid::Uuid::new_v4().to_string(); + let payload = + serde_json::to_string(&r).map_err(|_| reject(StatusCode::BAD_REQUEST, "invalid_update"))?; + db.claim_collection_operation(r.organization_id, &token, "update", &payload) + .await?; + // Inspect under the durable collection fence; new reservations cannot race this read. + let old = db.get_organization(r.organization_id).await?; + if let Some(model_id) = r + .embedding_model_id + .filter(|id| *id != old.embedding_model_id) + { + if !db.list_files(r.organization_id).await?.is_empty() { + db.finish_collection_operation(r.organization_id, &token, false) + .await?; + return Err(reject( + StatusCode::CONFLICT, + "populated_collection_model_change_requires_reconciliation", + )); + } + let model = match db.get_embedding_model(model_id).await { + Ok(m) => m, + Err(e) => { + db.finish_collection_operation(r.organization_id, &token, false) + .await?; + return Err(e.into()); + } + }; + if model.vector_size <= 0 { + db.finish_collection_operation(r.organization_id, &token, false) + .await?; + return Err(reject( + StatusCode::BAD_REQUEST, + "invalid_embedding_dimensions", + )); + } + let kc = match kc.client() { + Ok(kc) => kc, + Err(e) => { + db.finish_collection_operation(r.organization_id, &token, false) + .await?; + return Err(e); + } + }; + let path = format!( + "/internal/v1/collections/{}", + collection_name(r.organization_id) + ); + let result = async { + kc.collection_ack(Method::DELETE, &path, None).await?; + kc.collection_ack( + Method::PUT, + &path, + Some(json!({"vector_size":model.vector_size})), + ) + .await + } + .await; + if let Err(e) = result { + db.quarantine_collection_operation(r.organization_id, &token, e.code) + .await?; + return Err(e.public()); + } + } + let result = db + .update_organization( + r.organization_id, + UpdateOrganizationDao { + name: r.name, + description: r.description, + can_users_upload: r.can_users_upload, + embedding_model_id: r.embedding_model_id, + }, + ) + .await; + // A failed local update leaves the fenced payload for explicit settlement. + result?; + db.finish_collection_operation(r.organization_id, &token, false) + .await?; + Ok(Json(json!({}))) +} +#[utoipa::path(patch, path="/v1/rag/organization-collection/membership",tag="rag",request_body=OrganizationMembershipRequest,responses((status=200,description="Acknowledged operation"),(status=409,description="Conflict or reconciliation required"),(status=503,description="KC unavailable or uncertain outcome")),security(("bearer_token"=[])))] +pub(crate) async fn add_members( + p: Principal, + State(db): State, + Json(r): Json, +) -> Result, ShaideError> { + membership(p, db, r, true).await +} +#[utoipa::path(delete, path="/v1/rag/organization-collection/membership",tag="rag",request_body=OrganizationMembershipRequest,responses((status=200,description="Acknowledged operation"),(status=409,description="Conflict or reconciliation required"),(status=503,description="KC unavailable or uncertain outcome")),security(("bearer_token"=[])))] +pub(crate) async fn remove_members( + p: Principal, + State(db): State, + Json(r): Json, +) -> Result, ShaideError> { + membership(p, db, r, false).await +} +async fn membership( + p: Principal, + db: DbConn, + r: OrganizationMembershipRequest, + add: bool, +) -> Result, ShaideError> { + admin(&p)?; + if r.user_ids.len() > 1024 { + return Err(reject(StatusCode::BAD_REQUEST, "too_many_members")); + } + if !add { + // This existing metadata transaction uses BEGIN IMMEDIATE, serializing + // revocation with admission without releasing any remote-work fence. + db.remove_organization_members(r.organization_id, &r.user_ids) + .await?; + return Ok(Json(json!({}))); + } + let token = uuid::Uuid::new_v4().to_string(); + db.claim_collection_operation(r.organization_id, &token, "membership", "{}") + .await?; + let result = db + .add_organization_members(r.organization_id, &r.user_ids) + .await; + db.finish_collection_operation(r.organization_id, &token, false) + .await?; + result?; + Ok(Json(json!({}))) +} +#[derive(Deserialize, Serialize, utoipa::ToSchema)] +pub struct PresignRequest { + organization_id: i64, + file_name: String, + content_type: String, +} +#[derive(Deserialize, Serialize, utoipa::ToSchema)] +pub struct PresignResponse { + presigned_url: String, +} +#[utoipa::path(post, path="/v1/object-storage/presigned-url",tag="rag",request_body=PresignRequest,responses((status=200,description="Acknowledged operation",body=PresignResponse),(status=409,description="Conflict or reconciliation required"),(status=503,description="KC unavailable or uncertain outcome")),security(("bearer_token"=[])))] +pub(crate) async fn presign( + p: Principal, + State(db): State, + Extension(kc): Extension, + Json(r): Json, +) -> Result, ShaideError> { + upload_access(&db, &p, r.organization_id).await?; + validate_name(&r.file_name)?; + if r.file_name.contains(['/', '\\']) + || r.file_name == "." + || r.file_name == ".." + || r.file_name.trim() != r.file_name + { + return Err(reject(StatusCode::BAD_REQUEST, "invalid_file_name")); + } + if !matches!( + r.content_type.as_str(), + "application/pdf" + | "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + | "image/png" + | "image/jpeg" + ) || r.content_type.len() > 255 + || !r.content_type.contains('/') + || http::HeaderValue::from_str(&r.content_type).is_err() + { + return Err(reject(StatusCode::BAD_REQUEST, "invalid_content_type")); + } + let kc = kc.client()?; + // A server-issued immutable identity, intentionally independent of document contents. + let hash = format!( + "{}{}", + uuid::Uuid::new_v4().simple(), + uuid::Uuid::new_v4().simple() + ); + let token = uuid::Uuid::new_v4().to_string(); + db.reserve_upload( + InsertFileDao { + hash: hash.clone(), + organization_id: r.organization_id, + file_name: r.file_name, + mime_type: r.content_type.clone(), + }, + p.user.id, + &token, + 300, + ) + .await?; + let url = match kc + .presign(&collection_name(r.organization_id), &hash, &r.content_type) + .await + { + Ok(url) => url, + Err(e) => { + db.quarantine_file_operation(r.organization_id, &hash, e.code) + .await?; + return Err(e.public()); + } + }; + db.bind_reservation_query( + r.organization_id, + &hash, + &query_digest( + url.split_once('?') + .ok_or_else(|| { + reject( + StatusCode::SERVICE_UNAVAILABLE, + "invalid_presign_acknowledgement", + ) + })? + .1, + ), + ) + .await?; + Ok(Json(PresignResponse { presigned_url: url })) +} +#[utoipa::path(put, path="/s3/{collection}/{file_hash}",tag="rag",params(("collection"=String,Path,description="Canonical collection name"),("file_hash"=String,Path,description="Server-issued file identity")),responses((status=200,description="Acknowledged operation"),(status=409,description="Conflict or reconciliation required"),(status=503,description="KC unavailable or uncertain outcome")))] +pub(crate) async fn upload( + State(db): State, + Extension(kc): Extension, + Path((collection, hash)): Path<(String, String)>, + OriginalUri(uri): OriginalUri, + headers: HeaderMap, + body: Body, +) -> Result { + let id = collection_id(&collection)?; + if !valid_hash(&hash) { + return Err(reject(StatusCode::BAD_REQUEST, "invalid_file_hash")); + } + let query = uri + .query() + .filter(|q| !q.is_empty() && q.len() <= 4096) + .ok_or_else(|| reject(StatusCode::BAD_REQUEST, "signed_query_required"))?; + let file = db.get_file(id, &hash).await?; + if !db + .reservation_query_matches(id, &hash, &query_digest(query)) + .await? + { + return Err(reject(StatusCode::FORBIDDEN, "invalid_upload_reservation")); + } + let mime = headers + .get("content-type") + .and_then(|v| v.to_str().ok()) + .ok_or_else(|| reject(StatusCode::BAD_REQUEST, "content_type_required"))?; + if mime != file.mime_type { + return Err(reject(StatusCode::BAD_REQUEST, "content_type_mismatch")); + } + if headers.get_all("content-type").iter().count() != 1 + || headers.get_all("content-length").iter().count() != 1 + { + return Err(reject( + StatusCode::LENGTH_REQUIRED, + "content_length_required", + )); + } + if let Some(length) = headers.get("content-length") { + let size = length + .to_str() + .ok() + .and_then(|s| s.parse::().ok()) + .ok_or_else(|| reject(StatusCode::BAD_REQUEST, "invalid_content_length"))?; + if size == 0 || size > MAX_UPLOAD_BYTES { + return Err(reject(StatusCode::PAYLOAD_TOO_LARGE, "upload_size_limit")); + } + } + if headers.contains_key("content-encoding") { + return Err(reject( + StatusCode::BAD_REQUEST, + "content_encoding_not_supported", + )); + } + let kc = kc.client()?; + db.start_upload(id, &hash).await?; + let status = match kc.upload(&collection, &hash, query, &headers, body).await { + Ok(status) => status, + Err(e) => { + db.quarantine_file_operation(id, &hash, e.code).await?; + return Err(e.public()); + } + }; + db.transition_file_operation(id, &hash, "uploading", "uploaded") + .await?; + if db.upload_event_received(id, &hash).await? { + process(&db, kc, id, &hash).await?; + } + Ok(status) +} +#[derive(Deserialize, Serialize, utoipa::ToSchema)] +pub struct StorageEvent { + #[serde(rename = "Records")] + records: Vec, +} +#[derive(Deserialize, Serialize, utoipa::ToSchema)] +pub struct StorageRecord { + #[serde(rename = "eventName")] + event_name: String, + s3: StorageEntity, +} +#[derive(Deserialize, Serialize, utoipa::ToSchema)] +pub struct StorageEntity { + bucket: StorageBucket, + object: StorageObject, +} +#[derive(Deserialize, Serialize, utoipa::ToSchema)] +pub struct StorageBucket { + name: String, +} +#[derive(Deserialize, Serialize, utoipa::ToSchema)] +pub struct StorageObject { + key: String, + #[serde(default)] + size: Option, +} +#[derive(Deserialize)] +struct ObjectInfo { + exists: bool, + size: i64, + content_type: String, +} +#[derive(Deserialize)] +struct IngestResult { + status: String, + pages: i64, + chunks: i64, + size: i64, + mime_type: String, +} +#[utoipa::path(post, path="/v1/object-storage/event",tag="rag",security(("storage_notification_token"=[])),request_body=StorageEvent,responses((status=200,description="Acknowledged operation"),(status=409,description="Conflict or reconciliation required"),(status=503,description="KC unavailable or uncertain outcome")))] +pub(crate) async fn event( + State(db): State, + Extension(kc): Extension, + headers: HeaderMap, + Json(r): Json, +) -> Result, ShaideError> { + let kc = kc.client()?; + if !kc.authentic_notification(&headers) { + return Err(reject( + StatusCode::UNAUTHORIZED, + "invalid_storage_notification", + )); + } + if r.records.is_empty() || r.records.len() > 16 { + return Err(reject( + StatusCode::BAD_REQUEST, + "invalid_storage_notification", + )); + } + for record in r.records { + if record.event_name != "s3:ObjectCreated:Put" { + return Err(reject(StatusCode::BAD_REQUEST, "unsupported_storage_event")); + } + let id = collection_id(&record.s3.bucket.name)?; + let hash = record.s3.object.key; + if !valid_hash(&hash) { + return Err(reject(StatusCode::BAD_REQUEST, "invalid_file_hash")); + } + let file = db.get_file(id, &hash).await?; + // Never trust notification size/type. KC alone reads storage metadata. + let object = object_info(&db, kc, id, &hash).await?; + if !object.exists + || object.size <= 0 + || object.size as u64 > MAX_UPLOAD_BYTES as u64 + || object.content_type != file.mime_type + { + return Err(reject(StatusCode::BAD_REQUEST, "storage_object_mismatch")); + } + if file.status == FileStatus::Processed { + continue; + } + let state = db.record_upload_event(id, &hash).await?; + if state == "uploaded" { + process(&db, kc, id, &hash).await?; + } + } + Ok(Json(json!({}))) +} +async fn object_info( + _db: &DbConn, + kc: &KcClient, + id: i64, + hash: &str, +) -> Result { + kc.json( + Method::GET, + &format!( + "/internal/v1/storage/objects/{}/{hash}", + collection_name(id) + ), + None, + ) + .await + .map_err(|e| e.public()) +} +async fn process(db: &DbConn, kc: &KcClient, id: i64, hash: &str) -> Result<(), ShaideError> { + if !db.claim_ingestion(id, hash).await? { + return Ok(()); + } + // Any cancellation after the claim leaves an ingestion fence, even before a call starts. + let file = db.get_file(id, hash).await?; + let org = db.get_organization(id).await?; + let model = db.get_embedding_model(org.embedding_model_id).await?; + let collection = collection_name(id); + let result:Result=kc.json(Method::POST,"/internal/v1/ingest",Some(json!({"collection":collection,"bucket":collection,"key":hash,"file_hash":hash,"file_name":file.file_name,"embedding_model":{"name":model.name,"vector_size":model.vector_size}}))).await; + let result = match result { + Ok(r) => r, + Err(e) => { + db.quarantine_file_operation(id, hash, e.code).await?; + return Err(e.public()); + } + }; + if result.status != "processed" + || result.pages <= 0 + || result.chunks <= 0 + || result.size <= 0 + || result.size as u64 > MAX_UPLOAD_BYTES as u64 + || result.mime_type != file.mime_type + { + db.quarantine_file_operation(id, hash, "invalid_ingest_acknowledgement") + .await?; + return Err(reject( + StatusCode::SERVICE_UNAVAILABLE, + "invalid_ingest_acknowledgement", + )); + } + db.complete_ingestion(id, hash, result.size, result.pages) + .await?; + Ok(()) +} +#[utoipa::path(delete, path="/v1/rag/organization-collection/file",tag="rag",request_body=DeleteOrganizationFileRequest,responses((status=200,description="Acknowledged operation"),(status=409,description="Conflict or reconciliation required"),(status=503,description="KC unavailable or uncertain outcome")),security(("bearer_token"=[])))] +pub(crate) async fn delete_file( + p: Principal, + State(db): State, + Extension(kc): Extension, + Json(r): Json, +) -> Result, ShaideError> { + upload_access(&db, &p, r.organization_id).await?; + if !valid_hash(&r.file_hash) { + return Err(reject(StatusCode::BAD_REQUEST, "invalid_file_hash")); + } + let kc = kc.client()?; + let token = uuid::Uuid::new_v4().to_string(); + db.claim_file_delete(r.organization_id, &r.file_hash, &token, p.user.id) + .await?; + let result: Result = kc + .json( + Method::DELETE, + &format!( + "/internal/v1/collections/{}/files/{}", + collection_name(r.organization_id), + r.file_hash + ), + None, + ) + .await; + match result { + Ok(v) if v["deleted"] == true => {} + Ok(_) => { + db.quarantine_file_operation( + r.organization_id, + &r.file_hash, + "invalid_delete_acknowledgement", + ) + .await?; + return Err(reject( + StatusCode::SERVICE_UNAVAILABLE, + "invalid_delete_acknowledgement", + )); + } + Err(e) => { + db.quarantine_file_operation(r.organization_id, &r.file_hash, e.code) + .await?; + return Err(e.public()); + } + } + db.complete_file_delete(r.organization_id, &r.file_hash, &token) + .await?; + Ok(Json(json!({}))) +} +#[utoipa::path(delete, path="/v1/rag/delete-stale-files",tag="rag",responses((status=200,description="Acknowledged operation"),(status=409,description="Conflict or reconciliation required"),(status=503,description="KC unavailable or uncertain outcome")),security(("bearer_token"=[])))] +pub(crate) async fn stale_files( + p: Principal, + State(db): State, +) -> Result, ShaideError> { + admin(&p)?; + let count:i64=sqlx::query_scalar("SELECT (SELECT COUNT(*) FROM file_operations WHERE state='recovery_required' OR expires_at 0 { + return Err(reject( + StatusCode::CONFLICT, + "stale_uploads_require_reconciliation", + )); + } + Ok(Json(json!({}))) +} +pub fn knowledge_center_router(db: DbConn, kc: Option) -> Router { + Router::new() + .route( + "/v1/rag/organization-collection", + routing::get(list) + .post(create) + .patch(update) + .delete(delete_collection), + ) + .route( + "/v1/rag/organization-collection/membership", + routing::patch(add_members).delete(remove_members), + ) + .route( + "/v1/rag/organization-collection/file", + routing::delete(delete_file), + ) + .route("/v1/rag/delete-stale-files", routing::delete(stale_files)) + .route("/v1/object-storage/presigned-url", routing::post(presign)) + .route("/v1/object-storage/event", routing::post(event)) + .route("/s3/{collection}/{file_hash}", routing::put(upload)) + .layer(Extension(Integration(kc))) + .with_state(db) +} diff --git a/crates/shaide/src/routes/knowledge_center_tests.rs b/crates/shaide/src/routes/knowledge_center_tests.rs new file mode 100644 index 0000000..3084ec6 --- /dev/null +++ b/crates/shaide/src/routes/knowledge_center_tests.rs @@ -0,0 +1,1298 @@ +use crate::{ + routes::knowledge_center::knowledge_center_router, + services::auth::init_test_auth_service, + services::knowledge_center::{KcClient, KcConfig}, +}; +use axum::{ + Json, Router, + body::{Body, to_bytes}, + extract::Request, + response::IntoResponse, +}; +use http::{Method, StatusCode}; +use serde_json::{Value, json}; +use shaide_db::{ + DbConn, embedding_models::InsertEmbeddingModelDao, organizations::InsertOrganizationDao, +}; +use std::sync::{Arc, Mutex}; +use temp_testdir::TempDir; +use tower::ServiceExt; + +type RecordedCall = (Method, String, http::HeaderMap, Vec); +#[derive(Clone, Default)] +struct Calls( + Arc>>, + Arc>, + Arc, + Arc, + Arc, + Arc>>, +); +async fn fake_kc( + axum::extract::State(calls): axum::extract::State, + request: Request, +) -> axum::response::Response { + let (parts, body) = request.into_parts(); + let bytes = to_bytes(body, 1024 * 1024).await.unwrap(); + calls.0.lock().unwrap().push(( + parts.method.clone(), + parts.uri.to_string(), + parts.headers.clone(), + bytes.to_vec(), + )); + if parts.headers.get("authorization").unwrap() + != "Bearer synthetic-internal-token-32-bytes-long" + { + return StatusCode::UNAUTHORIZED.into_response(); + } + if parts.uri.path().starts_with("/internal/v1/collections/") + && let Some((method, body)) = calls.5.lock().unwrap().as_ref() + && method == parts.method + { + return Json(body.clone()).into_response(); + } + let fault = *calls.1.lock().unwrap(); + if (fault == "upload_failure" && parts.uri.path().starts_with("/internal/v1/storage/upload/")) + || (fault == "ingest_failure" && parts.uri.path() == "/internal/v1/ingest") + { + return ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({"error_code":"storage_unavailable"})), + ) + .into_response(); + } + if fault == "embedding_quota" && parts.uri.path() == "/internal/v1/ingest" { + return ( + StatusCode::FORBIDDEN, + Json(json!({"error_code":"embedding_quota"})), + ) + .into_response(); + } + if fault == "collection_failure" && parts.uri.path().starts_with("/internal/v1/collections/") { + return ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({"error":"storage_unavailable"})), + ) + .into_response(); + } + if fault == "empty_ingest" && parts.uri.path() == "/internal/v1/ingest" { + return Json(json!({"status":"processed","pages":0,"chunks":0,"size":9,"mime_type":"application/pdf"})).into_response(); + } + if fault == "upload_delay" && parts.uri.path().starts_with("/internal/v1/storage/upload/") { + let accepted = calls.clone(); + let remote = tokio::spawn(async move { + accepted.2.notify_one(); + accepted.3.notified().await; + accepted.4.notify_one(); + }); + remote.await.unwrap(); + } + match parts.uri.path() { + "/internal/v1/storage/presigned-url"=>{let v:Value=serde_json::from_slice(&bytes).unwrap();Json(json!({"presigned_url":format!("/s3/{}/{}?X-Amz-Expires=300&X-Amz-Signature=synthetic",v["collection"].as_str().unwrap(),v["key"].as_str().unwrap())})).into_response()}, + "/internal/v1/ingest"=>Json(json!({"status":"processed","pages":1,"chunks":1,"size":9,"mime_type":"application/pdf"})).into_response(), + p if p.starts_with("/internal/v1/storage/objects/")=>Json(json!({"exists":true,"size":9,"content_type":"application/pdf"})).into_response(), + p if p.starts_with("/internal/v1/storage/upload/")=>StatusCode::OK.into_response(), + p if p.contains("/files/")=>Json(json!({"deleted":true})).into_response(), + _=>Json(json!({})).into_response() + } +} +struct Fixture { + _temp: TempDir, + db: DbConn, + app: Router, + calls: Calls, + task: tokio::task::JoinHandle<()>, + admin: String, + member: String, + stranger: String, + id: i64, +} +impl Drop for Fixture { + fn drop(&mut self) { + self.task.abort(); + } +} +async fn fixture() -> Fixture { + let calls = Calls::default(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().fallback(fake_kc).with_state(calls.clone()); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let temp = TempDir::default(); + let db = DbConn::new(&temp.join("facade.sqlite")).await.unwrap(); + let admin = db + .create_admin("synthetic-admin".into(), "unused".into()) + .await + .unwrap(); + let member = db + .create_user( + "synthetic-member".into(), + "unused".into(), + chrono::Utc::now() + chrono::Duration::days(1), + ) + .await + .unwrap(); + let stranger = db + .create_user( + "synthetic-stranger".into(), + "unused".into(), + chrono::Utc::now() + chrono::Duration::days(1), + ) + .await + .unwrap(); + let model = db + .insert_embedding_model(InsertEmbeddingModelDao { + name: "synthetic".into(), + url: "http://invalid.example".into(), + vector_size: 8, + platform: None, + api_schema: None, + daily_input_token_limit: None, + }) + .await + .unwrap(); + let id = db + .create_organization(InsertOrganizationDao { + name: "synthetic".into(), + description: "".into(), + can_users_upload: true, + embedding_model_id: model, + user_ids: vec![member], + }) + .await + .unwrap(); + let client = KcClient::new( + KcConfig::parse( + Some(&format!("http://{address}")), + Some("synthetic-internal-token-32-bytes-long"), + Some("synthetic-webhook-token-32-bytes-long"), + true, + ) + .unwrap() + .unwrap(), + ) + .unwrap(); + let app = knowledge_center_router(db.clone(), Some(client)); + let auth = init_test_auth_service(); + let token = |id| auth.issue_access_token(id, None).unwrap().0; + Fixture { + _temp: temp, + db, + app, + calls, + task, + admin: token(admin), + member: token(member), + stranger: token(stranger), + id, + } +} +async fn request( + app: &Router, + method: Method, + path: &str, + token: Option<&str>, + body: Value, +) -> (StatusCode, Value) { + let mut r = http::Request::builder() + .method(method) + .uri(path) + .header("content-type", "application/json"); + if let Some(token) = token { + r = r.header("authorization", format!("Bearer {token}")); + } + let response = app + .clone() + .oneshot(r.body(Body::from(body.to_string())).unwrap()) + .await + .unwrap(); + let status = response.status(); + let bytes = to_bytes(response.into_body(), 1024 * 1024).await.unwrap(); + ( + status, + serde_json::from_slice(&bytes).unwrap_or(Value::Null), + ) +} +#[tokio::test] +async fn presign_is_scoped_and_bytes_only_reach_authenticated_kc() { + let f = fixture().await; + let body = json!({"organization_id":f.id,"file_name":"synthetic.txt","content_type":"application/pdf"}); + assert_eq!( + request( + &f.app, + Method::POST, + "/v1/object-storage/presigned-url", + None, + body.clone() + ) + .await + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + request( + &f.app, + Method::POST, + "/v1/object-storage/presigned-url", + Some(&f.stranger), + body.clone() + ) + .await + .0, + StatusCode::FORBIDDEN + ); + assert!(f.calls.0.lock().unwrap().is_empty()); + let (status, v) = request( + &f.app, + Method::POST, + "/v1/object-storage/presigned-url", + Some(&f.member), + body, + ) + .await; + assert_eq!(status, StatusCode::OK); + let url = v["presigned_url"].as_str().unwrap(); + let response = f + .app + .clone() + .oneshot( + http::Request::builder() + .method("PUT") + .uri(url) + .header("content-type", "application/pdf") + .header("content-length", "9") + .body(Body::from("synthetic")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + f.app + .clone() + .oneshot( + http::Request::builder() + .method("GET") + .uri(url) + .body(Body::empty()) + .unwrap() + ) + .await + .unwrap() + .status(), + StatusCode::METHOD_NOT_ALLOWED + ); + let calls = f.calls.0.lock().unwrap(); + let presign: Value = serde_json::from_slice(&calls[0].3).unwrap(); + assert_eq!(presign["expires_in_seconds"], 300); + assert_eq!(calls[1].0, Method::PUT); + assert_eq!(calls[1].3, b"synthetic"); + assert!( + calls[1] + .1 + .starts_with("/internal/v1/storage/upload/collection-") + ); + assert!( + calls[1] + .1 + .ends_with("?X-Amz-Expires=300&X-Amz-Signature=synthetic") + ); + assert_eq!( + calls[1].2["authorization"], + "Bearer synthetic-internal-token-32-bytes-long" + ); +} +#[tokio::test] +async fn collection_listing_preserves_admin_and_member_visibility_without_kc() { + let f = fixture().await; + let app = knowledge_center_router(f.db.clone(), None); + let (status, v) = request( + &app, + Method::GET, + "/v1/rag/organization-collection", + Some(&f.admin), + Value::Null, + ) + .await; + assert_eq!(status, StatusCode::OK); + assert!(v["organization"][0]["users"].is_array()); + let (_, v) = request( + &app, + Method::GET, + "/v1/rag/organization-collection", + Some(&f.member), + Value::Null, + ) + .await; + assert!(v["organization"][0]["users"].is_null()); + let (_, v) = request( + &app, + Method::GET, + "/v1/rag/organization-collection", + Some(&f.stranger), + Value::Null, + ) + .await; + assert_eq!(v["organization"], json!([])); + let body = json!({"organization_id":f.id,"file_name":"synthetic.txt","content_type":"application/pdf"}); + assert_eq!( + request( + &app, + Method::POST, + "/v1/object-storage/presigned-url", + Some(&f.member), + body + ) + .await + .0, + StatusCode::SERVICE_UNAVAILABLE + ); + assert!(f.db.list_files(f.id).await.unwrap().is_empty()); +} +async fn reserve(f: &Fixture) -> String { + let (status,v)=request(&f.app,Method::POST,"/v1/object-storage/presigned-url",Some(&f.member),json!({"organization_id":f.id,"file_name":"synthetic.pdf","content_type":"application/pdf"})).await; + assert_eq!(status, StatusCode::OK); + v["presigned_url"].as_str().unwrap().into() +} +async fn put(app: &Router, url: &str) -> StatusCode { + app.clone() + .oneshot( + http::Request::builder() + .method("PUT") + .uri(url) + .header("content-type", "application/pdf") + .header("content-length", "9") + .body(Body::from("synthetic")) + .unwrap(), + ) + .await + .unwrap() + .status() +} +fn notification(id: i64, hash: &str) -> Value { + json!({"Records":[{"eventName":"s3:ObjectCreated:Put","s3":{"bucket":{"name":format!("collection-{id}")},"object":{"key":hash,"size":999999999}}}]}) +} +fn file_hash(url: &str) -> &str { + url.split('?').next().unwrap().rsplit('/').next().unwrap() +} +#[tokio::test] +async fn forged_query_and_wrong_collection_do_not_consume_reservation() { + let f = fixture().await; + let url = reserve(&f).await; + assert_eq!( + put( + &f.app, + &url.replace("Signature=synthetic", "Signature=forged") + ) + .await, + StatusCode::FORBIDDEN + ); + assert_eq!( + put(&f.app, &url.replace("collection-1", "collection-999")).await, + StatusCode::NOT_FOUND + ); + assert_eq!(f.calls.0.lock().unwrap().len(), 1); + assert_eq!(put(&f.app, &url).await, StatusCode::OK); + assert_eq!(put(&f.app, &url).await, StatusCode::CONFLICT); + assert_eq!(f.calls.0.lock().unwrap().len(), 2); +} +#[tokio::test] +async fn event_uses_trusted_kc_metadata_and_claims_ingestion_once() { + let f = fixture().await; + let url = reserve(&f).await; + let hash = file_hash(&url); + assert_eq!(put(&f.app, &url).await, StatusCode::OK); + let body = notification(f.id, hash); + assert_eq!( + request( + &f.app, + Method::POST, + "/v1/object-storage/event", + None, + body.clone() + ) + .await + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + request( + &f.app, + Method::POST, + "/v1/object-storage/event", + Some("forged"), + body.clone() + ) + .await + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!(f.calls.0.lock().unwrap().len(), 2); + let (a, b) = tokio::join!( + request( + &f.app, + Method::POST, + "/v1/object-storage/event", + Some("synthetic-webhook-token-32-bytes-long"), + body.clone() + ), + request( + &f.app, + Method::POST, + "/v1/object-storage/event", + Some("synthetic-webhook-token-32-bytes-long"), + body.clone() + ) + ); + assert_eq!(a.0, StatusCode::OK); + assert_eq!(b.0, StatusCode::OK); + let file = f.db.get_file(f.id, hash).await.unwrap(); + assert_eq!(file.status, shaide_db::files::FileStatus::Processed); + assert_eq!(file.size, Some(9)); + let count = || { + f.calls + .0 + .lock() + .unwrap() + .iter() + .filter(|c| c.1 == "/internal/v1/ingest") + .count() + }; + assert_eq!(count(), 1); + assert_eq!( + request( + &f.app, + Method::POST, + "/v1/object-storage/event", + Some("synthetic-webhook-token-32-bytes-long"), + body + ) + .await + .0, + StatusCode::OK + ); + assert_eq!(count(), 1); + let calls = f.calls.0.lock().unwrap(); + let ingest = &calls + .iter() + .find(|c| c.1 == "/internal/v1/ingest") + .unwrap() + .3; + let v: Value = serde_json::from_slice(ingest).unwrap(); + assert_eq!(v["key"], hash); + assert_eq!(v["file_hash"], hash); + assert_eq!( + v["embedding_model"], + json!({"name":"synthetic","vector_size":8}) + ); +} +#[tokio::test] +async fn rejected_or_incomplete_ingestion_is_durable_and_blocks_delete() { + for fault in ["ingest_failure", "empty_ingest", "embedding_quota"] { + let f = fixture().await; + let url = reserve(&f).await; + assert_eq!(put(&f.app, &url).await, StatusCode::OK); + *f.calls.1.lock().unwrap() = fault; + let hash = file_hash(&url); + assert!( + !request( + &f.app, + Method::POST, + "/v1/object-storage/event", + Some("synthetic-webhook-token-32-bytes-long"), + notification(f.id, hash) + ) + .await + .0 + .is_success() + ); + let other = DbConn::new(&f._temp.join("facade.sqlite")).await.unwrap(); + assert_eq!( + other.get_file(f.id, hash).await.unwrap().status, + shaide_db::files::FileStatus::ProcessingFailed + ); + let state: String = + sqlx::query_scalar("SELECT state FROM file_operations WHERE file_hash=?") + .bind(hash) + .fetch_one(&other.pool) + .await + .unwrap(); + assert_eq!(state, "recovery_required"); + assert_eq!( + request( + &f.app, + Method::DELETE, + "/v1/rag/organization-collection/file", + Some(&f.member), + json!({"organization_id":f.id,"file_hash":hash}) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!( + request( + &f.app, + Method::DELETE, + "/v1/rag/organization-collection", + Some(&f.admin), + json!({"id":f.id}) + ) + .await + .0, + StatusCode::CONFLICT + ); + } +} +#[tokio::test] +async fn event_arriving_before_upload_ack_is_processed_after_ack() { + let f = fixture().await; + let url = reserve(&f).await; + *f.calls.1.lock().unwrap() = "upload_delay"; + let app = f.app.clone(); + let u = url.clone(); + let upload = tokio::spawn(async move { put(&app, &u).await }); + tokio::time::timeout(std::time::Duration::from_secs(30), f.calls.2.notified()) + .await + .unwrap(); + assert_eq!( + request( + &f.app, + Method::POST, + "/v1/object-storage/event", + Some("synthetic-webhook-token-32-bytes-long"), + notification(f.id, file_hash(&url)) + ) + .await + .0, + StatusCode::OK + ); + f.calls.3.notify_one(); + assert_eq!(upload.await.unwrap(), StatusCode::OK); + assert_eq!( + f.db.get_file(f.id, file_hash(&url)).await.unwrap().status, + shaide_db::files::FileStatus::Processed + ); +} +#[tokio::test] +async fn cancellation_leaves_a_fence_even_when_remote_handler_finishes() { + let f = fixture().await; + let url = reserve(&f).await; + *f.calls.1.lock().unwrap() = "upload_delay"; + let app = f.app.clone(); + let u = url.clone(); + let upload = tokio::spawn(async move { put(&app, &u).await }); + tokio::time::timeout(std::time::Duration::from_secs(30), f.calls.2.notified()) + .await + .unwrap(); + assert_eq!(f.calls.0.lock().unwrap().len(), 2); + upload.abort(); + assert!(upload.await.unwrap_err().is_cancelled()); + f.calls.3.notify_one(); + tokio::time::timeout(std::time::Duration::from_secs(30), f.calls.4.notified()) + .await + .unwrap(); + let other = DbConn::new(&f._temp.join("facade.sqlite")).await.unwrap(); + assert!(other.start_upload(f.id, file_hash(&url)).await.is_err()); + assert_eq!( + request( + &f.app, + Method::DELETE, + "/v1/rag/organization-collection", + Some(&f.admin), + json!({"id":f.id}) + ) + .await + .0, + StatusCode::CONFLICT + ); +} +#[tokio::test] +async fn collection_crud_membership_and_file_delete_follow_the_kc_contract() { + let f = fixture().await; + let body = json!({"organization_name":"synthetic-created","organization_description":"owned fixture","users":[],"can_users_upload":false,"embedding_model_id":1}); + assert_eq!( + request( + &f.app, + Method::POST, + "/v1/rag/organization-collection", + Some(&f.member), + body.clone() + ) + .await + .0, + StatusCode::FORBIDDEN + ); + let (status, v) = request( + &f.app, + Method::POST, + "/v1/rag/organization-collection", + Some(&f.admin), + body, + ) + .await; + assert_eq!(status, StatusCode::OK); + let id = v["id"].as_i64().unwrap(); + assert_eq!( + request( + &f.app, + Method::PATCH, + "/v1/rag/organization-collection", + Some(&f.admin), + json!({"organization_id":id,"name":"renamed"}) + ) + .await + .0, + StatusCode::OK + ); + assert_eq!( + request( + &f.app, + Method::PATCH, + "/v1/rag/organization-collection/membership", + Some(&f.admin), + json!({"organization_id":id,"user_ids":[2,2]}) + ) + .await + .0, + StatusCode::OK + ); + assert_eq!( + f.db.list_organization_member_ids(id).await.unwrap(), + vec![2] + ); + assert_eq!( + request( + &f.app, + Method::DELETE, + "/v1/rag/organization-collection/membership", + Some(&f.admin), + json!({"organization_id":id,"user_ids":[2]}) + ) + .await + .0, + StatusCode::OK + ); + assert_eq!( + request( + &f.app, + Method::DELETE, + "/v1/rag/organization-collection", + Some(&f.admin), + json!({"id":id}) + ) + .await + .0, + StatusCode::OK + ); + assert!(f.db.get_organization(id).await.is_err()); + let url = reserve(&f).await; + assert_eq!(put(&f.app, &url).await, StatusCode::OK); + let hash = file_hash(&url); + assert_eq!( + request( + &f.app, + Method::POST, + "/v1/object-storage/event", + Some("synthetic-webhook-token-32-bytes-long"), + notification(f.id, hash) + ) + .await + .0, + StatusCode::OK + ); + assert_eq!( + request( + &f.app, + Method::DELETE, + "/v1/rag/organization-collection/file", + Some(&f.stranger), + json!({"organization_id":f.id,"file_hash":hash}) + ) + .await + .0, + StatusCode::FORBIDDEN + ); + assert_eq!( + request( + &f.app, + Method::DELETE, + "/v1/rag/organization-collection/file", + Some(&f.member), + json!({"organization_id":f.id,"file_hash":hash}) + ) + .await + .0, + StatusCode::OK + ); + assert!(f.db.get_file(f.id, hash).await.is_err()); + assert!( + f.calls + .0 + .lock() + .unwrap() + .iter() + .any(|c| c.0 == Method::DELETE + && c.1 == format!("/internal/v1/collections/collection-{}/files/{hash}", f.id)) + ); +} +#[tokio::test] +async fn upload_limits_filename_policy_and_remote_error_never_succeed_silently() { + let f = fixture().await; + for (name, mime) in [ + ("../escape.pdf", "application/pdf"), + ("synthetic.exe", "application/octet-stream"), + ("bad\nname.pdf", "application/pdf"), + ] { + assert_eq!( + request( + &f.app, + Method::POST, + "/v1/object-storage/presigned-url", + Some(&f.member), + json!({"organization_id":f.id,"file_name":name,"content_type":mime}) + ) + .await + .0, + StatusCode::BAD_REQUEST + ); + } + assert!(f.db.list_files(f.id).await.unwrap().is_empty()); + let url = reserve(&f).await; + let response = f + .app + .clone() + .oneshot( + http::Request::builder() + .method("PUT") + .uri(&url) + .header("content-type", "application/pdf") + .header("content-length", (16 * 1024 * 1024 + 1).to_string()) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE); + assert_eq!(f.calls.0.lock().unwrap().len(), 1); + *f.calls.1.lock().unwrap() = "upload_failure"; + assert_eq!(put(&f.app, &url).await, StatusCode::SERVICE_UNAVAILABLE); + let file = f.db.get_file(f.id, file_hash(&url)).await.unwrap(); + assert_eq!(file.status, shaide_db::files::FileStatus::UploadingFailed); + assert_eq!(file.error_code.as_deref(), Some("storage_unavailable")); + assert_eq!(put(&f.app, &url).await, StatusCode::CONFLICT); +} + +#[tokio::test] +async fn model_replacement_preserves_populated_data_and_updates_only_empty_collections() { + let f = fixture().await; + let model = + f.db.insert_embedding_model(InsertEmbeddingModelDao { + name: "synthetic-replacement".into(), + url: "http://invalid.example".into(), + vector_size: 16, + platform: None, + api_schema: None, + daily_input_token_limit: None, + }) + .await + .unwrap(); + assert_eq!( + request( + &f.app, + Method::PATCH, + "/v1/rag/organization-collection", + Some(&f.admin), + json!({"organization_id":f.id,"embedding_model_id":model}) + ) + .await + .0, + StatusCode::OK + ); + assert_eq!( + f.db.get_organization(f.id) + .await + .unwrap() + .embedding_model_id, + model + ); + let url = reserve(&f).await; + assert_eq!(put(&f.app, &url).await, StatusCode::OK); + assert_eq!( + request( + &f.app, + Method::POST, + "/v1/object-storage/event", + Some("synthetic-webhook-token-32-bytes-long"), + notification(f.id, file_hash(&url)) + ) + .await + .0, + StatusCode::OK + ); + let before = f.calls.0.lock().unwrap().len(); + assert_eq!( + request( + &f.app, + Method::PATCH, + "/v1/rag/organization-collection", + Some(&f.admin), + json!({"organization_id":f.id,"embedding_model_id":1}) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!(f.calls.0.lock().unwrap().len(), before); + assert_eq!( + f.db.get_organization(f.id) + .await + .unwrap() + .embedding_model_id, + model + ); + assert_eq!( + f.db.get_file(f.id, file_hash(&url)).await.unwrap().status, + shaide_db::files::FileStatus::Processed + ); +} +#[tokio::test] +async fn uncertain_collection_delete_keeps_metadata_and_a_durable_fence() { + let f = fixture().await; + *f.calls.1.lock().unwrap() = "collection_failure"; + assert_eq!( + request( + &f.app, + Method::DELETE, + "/v1/rag/organization-collection", + Some(&f.admin), + json!({"id":f.id}) + ) + .await + .0, + StatusCode::SERVICE_UNAVAILABLE + ); + assert!(f.db.get_organization(f.id).await.is_ok()); + let other = DbConn::new(&f._temp.join("facade.sqlite")).await.unwrap(); + let state: String = + sqlx::query_scalar("SELECT state FROM collection_operations WHERE organization_id=?") + .bind(f.id) + .fetch_one(&other.pool) + .await + .unwrap(); + assert_eq!(state, "recovery_required"); + assert_eq!( + request( + &f.app, + Method::PATCH, + "/v1/rag/organization-collection", + Some(&f.admin), + json!({"organization_id":f.id,"name":"unsafe retry"}) + ) + .await + .0, + StatusCode::CONFLICT + ); +} + +async fn fix1_collection_ack_case(phase: &str) { + for body in [ + Value::Null, + json!([]), + json!([{}]), + json!(false), + json!(7), + json!("ok"), + json!({"extra":true}), + json!({"status":"failed"}), + json!({}), + ] { + let f = fixture().await; + let valid = body == json!({}); + let method = if phase == "create" || phase == "replace_put" { + Method::PUT + } else { + Method::DELETE + }; + *f.calls.5.lock().unwrap() = Some((method, body.clone())); + let mut id = f.id; + let mut old_files = vec![]; + let (verb, payload) = match phase { + "create" => ( + Method::POST, + json!({"organization_name":"ack-fixture","organization_description":"synthetic","can_users_upload":true,"embedding_model_id":1,"users":[]}), + ), + "delete" => { + // Settled file metadata must survive an invalid collection acknowledgement. + sqlx::query("INSERT INTO files(hash,organization_id,file_name,mime_type,status) VALUES(?,?,'synthetic.pdf','application/pdf','processed')") + .bind("a".repeat(64)).bind(id).execute(&f.db.pool).await.unwrap(); + old_files = f.db.list_files(id).await.unwrap(); + (Method::DELETE, json!({"id":id})) + } + _ => { + let model = + f.db.insert_embedding_model(InsertEmbeddingModelDao { + name: "ack-replacement".into(), + url: "http://invalid.example".into(), + vector_size: 16, + platform: None, + api_schema: None, + daily_input_token_limit: None, + }) + .await + .unwrap(); + ( + Method::PATCH, + json!({"organization_id":id,"embedding_model_id":model}), + ) + } + }; + let (status, response) = request( + &f.app, + verb, + "/v1/rag/organization-collection", + Some(&f.admin), + payload, + ) + .await; + assert_eq!( + status, + if valid { + StatusCode::OK + } else { + StatusCode::SERVICE_UNAVAILABLE + }, + "phase={phase} body={body}" + ); + if phase == "create" { + id = if valid { + response["id"].as_i64().unwrap() + } else { + f.db.list_organizations() + .await + .unwrap() + .into_iter() + .find(|o| o.name == "ack-fixture") + .unwrap() + .id + }; + } + let operation: Option<(String, String, Option)> = sqlx::query_as( + "SELECT kind,state,error_code FROM collection_operations WHERE organization_id=?", + ) + .bind(id) + .fetch_optional(&f.db.pool) + .await + .unwrap(); + if valid { + assert!(operation.is_none()); + if phase == "delete" { + assert!(f.db.get_organization(id).await.is_err()); + let remaining: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM files WHERE organization_id=?") + .bind(id) + .fetch_one(&f.db.pool) + .await + .unwrap(); + assert_eq!(remaining, 0); + } + if phase.starts_with("replace") { + assert_eq!( + f.db.get_organization(id).await.unwrap().embedding_model_id, + 2 + ); + } + } else { + let (kind, state, code) = operation.unwrap(); + assert_eq!( + kind, + if phase.starts_with("replace") { + "update" + } else { + phase + } + ); + assert_eq!(state, "recovery_required"); + assert_eq!(code.as_deref(), Some("kc_outcome_unknown")); + assert_eq!(response["error"]["code"], "kc_outcome_unknown"); + assert_eq!( + f.db.get_organization(id).await.unwrap().embedding_model_id, + 1 + ); + assert_eq!(f.db.list_files(id).await.unwrap(), old_files); + assert!( + f.db.claim_collection_operation(id, "retry", "delete", "{}") + .await + .is_err() + ); + } + let calls = f.calls.0.lock().unwrap(); + assert_eq!( + calls.len(), + if phase == "replace_put" || (phase == "replace_delete" && valid) { + 2 + } else { + 1 + } + ); + if phase.starts_with("replace") { + assert_eq!(calls[0].0, Method::DELETE); + if calls.len() == 2 { + assert_eq!(calls[1].0, Method::PUT); + } + } + } +} +#[tokio::test] +async fn fix1_create_requires_empty_collection_ack() { + fix1_collection_ack_case("create").await; +} +#[tokio::test] +async fn fix1_delete_requires_empty_collection_ack() { + fix1_collection_ack_case("delete").await; +} +#[tokio::test] +async fn fix1_replacement_delete_requires_empty_collection_ack() { + fix1_collection_ack_case("replace_delete").await; +} +#[tokio::test] +async fn fix1_replacement_put_requires_empty_collection_ack() { + fix1_collection_ack_case("replace_put").await; +} + +async fn fix1_file_fences(f: &Fixture) -> Vec { + sqlx::query_scalar("SELECT json_array(file_hash,organization_id,owner_id,token,state,expires_at,query_digest,event_received,error_code,created_at) FROM file_operations WHERE organization_id=? ORDER BY file_hash") + .bind(f.id).fetch_all(&f.db.pool).await.unwrap() +} +async fn fix1_collection_fences(f: &Fixture) -> Vec { + sqlx::query_scalar("SELECT json_array(organization_id,token,kind,payload,state,error_code,created_at) FROM collection_operations WHERE organization_id=? ORDER BY organization_id") + .bind(f.id).fetch_all(&f.db.pool).await.unwrap() +} +async fn fix1_restrict_access_case(remove_member: bool) { + for state in [ + "unused", + "expired", + "uncertain", + "uploaded", + "collection_recovery", + "collection_active", + ] { + let f = fixture().await; + let url = reserve(&f).await; + let hash = file_hash(&url); + let owner: i64 = + sqlx::query_scalar("SELECT owner_id FROM file_operations WHERE file_hash=?") + .bind(hash) + .fetch_one(&f.db.pool) + .await + .unwrap(); + if state == "expired" { + sqlx::query("UPDATE file_operations SET expires_at=0 WHERE file_hash=?") + .bind(hash) + .execute(&f.db.pool) + .await + .unwrap(); + } + if state == "uncertain" { + f.db.quarantine_file_operation(f.id, hash, "upload_outcome_unknown") + .await + .unwrap(); + } + if state == "uploaded" { + assert_eq!(put(&f.app, &url).await, StatusCode::OK); + } + if state.starts_with("collection_") { + sqlx::query("INSERT INTO collection_operations(organization_id,token,kind,payload,state) VALUES(?,'retained','update','{}',?)").bind(f.id).bind(if state=="collection_active" {"active"} else {"recovery_required"}).execute(&f.db.pool).await.unwrap(); + } + let files = f.db.list_files(f.id).await.unwrap(); + let fences = fix1_file_fences(&f).await; + let collections = fix1_collection_fences(&f).await; + let (method, path, payload) = if remove_member { + ( + Method::DELETE, + "/v1/rag/organization-collection/membership", + json!({"organization_id":f.id,"user_ids":[owner]}), + ) + } else { + ( + Method::PATCH, + "/v1/rag/organization-collection", + json!({"organization_id":f.id,"can_users_upload":false}), + ) + }; + assert_eq!( + request( + &f.app, + method.clone(), + path, + Some(&f.member), + payload.clone() + ) + .await + .0, + StatusCode::FORBIDDEN + ); + assert_eq!( + request(&f.app, method, path, Some(&f.admin), payload) + .await + .0, + StatusCode::OK, + "state={state} remove_member={remove_member}" + ); + if remove_member { + assert!(!f.db.is_organization_member(f.id, owner).await.unwrap()); + } else { + assert!(!f.db.get_organization(f.id).await.unwrap().can_users_upload); + } + let before = f.calls.0.lock().unwrap().len(); + assert_ne!(put(&f.app, &url).await, StatusCode::OK); + assert!(f.db.start_upload(f.id, hash).await.is_err()); + assert!(f.db.claim_ingestion(f.id, hash).await.is_err()); + assert_ne!( + request( + &f.app, + Method::POST, + "/v1/object-storage/event", + Some("synthetic-webhook-token-32-bytes-long"), + notification(f.id, hash) + ) + .await + .0, + StatusCode::OK + ); + assert!(!f.calls.0.lock().unwrap()[before..].iter().any( + |c| c.1 == "/internal/v1/ingest" || c.1.starts_with("/internal/v1/storage/upload/") + )); + assert_eq!(request(&f.app,Method::POST,"/v1/object-storage/presigned-url",Some(&f.member),json!({"organization_id":f.id,"file_name":"later.pdf","content_type":"application/pdf"})).await.0,StatusCode::FORBIDDEN); + for (verb, payload) in [ + ( + Method::PATCH, + json!({"organization_id":f.id,"can_users_upload":true}), + ), + ( + Method::PATCH, + json!({"organization_id":f.id,"embedding_model_id":1}), + ), + ( + Method::PATCH, + json!({"organization_id":f.id,"can_users_upload":false,"embedding_model_id":1}), + ), + ( + Method::PATCH, + json!({"organization_id":f.id,"can_users_upload":false,"name":"mixed update"}), + ), + ( + Method::PATCH, + json!({"organization_id":f.id,"name":"not authorized by restriction"}), + ), + (Method::DELETE, json!({"id":f.id})), + ] { + assert_eq!( + request( + &f.app, + verb, + "/v1/rag/organization-collection", + Some(&f.admin), + payload + ) + .await + .0, + StatusCode::CONFLICT + ); + } + assert_eq!( + request( + &f.app, + Method::PATCH, + "/v1/rag/organization-collection/membership", + Some(&f.admin), + json!({"organization_id":f.id,"user_ids":[owner]}) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert!( + f.db.claim_file_delete(f.id, hash, "unsafe", owner) + .await + .is_err() + ); + assert_eq!(f.db.list_files(f.id).await.unwrap(), files); + let after = fix1_file_fences(&f).await; + assert_eq!(after, fences); + let after = fix1_collection_fences(&f).await; + assert_eq!(after, collections); + } +} +#[tokio::test] +async fn fix1_public_member_revocation_retains_fences() { + fix1_restrict_access_case(true).await; +} +#[tokio::test] +async fn fix1_public_upload_disable_retains_fences() { + fix1_restrict_access_case(false).await; +} +#[tokio::test] +async fn fix1_stale_files_conflicts_on_ordinary_collection_recovery() { + let f = fixture().await; + assert_eq!( + request( + &f.app, + Method::DELETE, + "/v1/rag/delete-stale-files", + Some(&f.admin), + json!({}) + ) + .await + .0, + StatusCode::OK + ); + for kind in ["create", "delete", "update"] { + f.db.claim_collection_operation(f.id, "recovery", kind, "{\"retained\":true}") + .await + .unwrap(); + f.db.quarantine_collection_operation(f.id, "recovery", "kc_outcome_unknown") + .await + .unwrap(); + let before = f.db.get_organization(f.id).await.unwrap(); + let (status, body) = request( + &f.app, + Method::DELETE, + "/v1/rag/delete-stale-files", + Some(&f.admin), + json!({}), + ) + .await; + assert_eq!(status, StatusCode::CONFLICT, "kind={kind}"); + assert_eq!( + body["error"]["code"], + "stale_uploads_require_reconciliation" + ); + assert_eq!(f.db.get_organization(f.id).await.unwrap(), before); + let operation: (String, String, String, String) = sqlx::query_as( + "SELECT token,kind,payload,state FROM collection_operations WHERE organization_id=?", + ) + .bind(f.id) + .fetch_one(&f.db.pool) + .await + .unwrap(); + assert_eq!( + operation, + ( + "recovery".into(), + kind.into(), + "{\"retained\":true}".into(), + "recovery_required".into() + ) + ); + // Only fixture cleanup: stale route must leave the operation untouched. + sqlx::query("DELETE FROM collection_operations WHERE organization_id=?") + .bind(f.id) + .execute(&f.db.pool) + .await + .unwrap(); + } + assert!(f.calls.0.lock().unwrap().is_empty()); +} diff --git a/crates/shaide/src/routes/mod.rs b/crates/shaide/src/routes/mod.rs index 768844b..13c1ddb 100644 --- a/crates/shaide/src/routes/mod.rs +++ b/crates/shaide/src/routes/mod.rs @@ -6,6 +6,7 @@ pub mod chat; pub mod completions; pub mod embedding; pub mod health; +pub mod knowledge_center; pub mod mcp; pub mod metrics; pub mod models; @@ -17,5 +18,8 @@ pub mod users; pub mod vector_db; pub async fn fallback_404(uri: Uri) -> ShaideError { - ShaideError::not_found_route(uri.to_string()) + ShaideError::not_found_route(uri.path().to_owned()) } + +#[cfg(test)] +mod knowledge_center_tests; diff --git a/crates/shaide/src/services/knowledge_center.rs b/crates/shaide/src/services/knowledge_center.rs new file mode 100644 index 0000000..44abdd1 --- /dev/null +++ b/crates/shaide/src/services/knowledge_center.rs @@ -0,0 +1,305 @@ +//! Fixed-target, authenticated KC transport. No object-storage provider here. +use crate::error::ShaideError; +use axum::body::Body; +use http::{HeaderMap, Method, StatusCode}; +use serde::de::DeserializeOwned; +use serde_json::{Value, json}; +use std::time::Duration; +use url::Url; + +pub const MAX_UPLOAD_BYTES: usize = 16 * 1024 * 1024; +#[derive(Clone)] +pub struct KcConfig { + url: String, + token: String, + webhook_token: String, +} +impl std::fmt::Debug for KcConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("KcConfig").field("enabled", &true).finish() + } +} +impl KcConfig { + pub fn uses_token(&self, token: &str) -> bool { + self.token == token || self.webhook_token == token + } + pub fn parse( + url: Option<&str>, + token: Option<&str>, + webhook_token: Option<&str>, + allow_http: bool, + ) -> Result, &'static str> { + let Some(raw) = url.filter(|s| !s.is_empty()) else { + if token.is_some() || webhook_token.is_some() { + return Err("KC tokens require KNOWLEDGE_CENTER_URL"); + } + return Ok(None); + }; + let url = Url::parse(raw).map_err(|_| "invalid KC URL")?; + if url.host_str().is_none() + || !(url.scheme() == "https" || (url.scheme() == "http" && allow_http)) + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || url.path() != "/" + { + return Err("KC URL must be an origin; HTTP needs explicit local opt-in"); + } + let token = token + .filter(|v| v.len() >= 32 && !v.chars().any(char::is_whitespace)) + .ok_or("KC internal token must have at least 32 bytes")?; + let webhook_token = webhook_token + .filter(|v| v.len() >= 32 && !v.chars().any(char::is_whitespace)) + .ok_or("KC webhook token must have at least 32 bytes")?; + if token == webhook_token { + return Err("internal and notification tokens must differ"); + } + Ok(Some(Self { + url: url.as_str().trim_end_matches('/').into(), + token: token.into(), + webhook_token: webhook_token.into(), + })) + } +} +#[derive(Clone)] +pub struct KcClient { + config: KcConfig, + http: reqwest::Client, +} +#[derive(Debug)] +pub struct KcFailure { + pub code: &'static str, + pub status: StatusCode, +} +impl KcFailure { + fn unknown() -> Self { + Self { + code: "kc_outcome_unknown", + status: StatusCode::SERVICE_UNAVAILABLE, + } + } + pub fn public(self) -> ShaideError { + ShaideError::request_rejection( + self.status, + "Knowledge Center operation failed; reconciliation may be required".into(), + self.code, + ) + } +} +impl KcClient { + pub fn new(config: KcConfig) -> Result { + let http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .no_proxy() + .connect_timeout(Duration::from_secs(5)) + .timeout(Duration::from_secs(180)) + .build() + .map_err(|_| "could not initialize KC transport")?; + Ok(Self { config, http }) + } + pub fn authentic_notification(&self, headers: &HeaderMap) -> bool { + let Some(raw) = headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.strip_prefix("Bearer ")) + else { + return false; + }; + // Fixed-length constant-work comparison; never format either secret. + let expected = self.config.webhook_token.as_bytes(); + let actual = raw.as_bytes(); + if expected.len() != actual.len() { + return false; + } + expected + .iter() + .zip(actual) + .fold(0u8, |acc, (a, b)| acc | (a ^ b)) + == 0 + } + fn request(&self, method: Method, path: &str) -> reqwest::RequestBuilder { + self.http + .request(method, format!("{}{path}", self.config.url)) + .bearer_auth(&self.config.token) + } + async fn checked( + response: Result, + ) -> Result { + let response = response.map_err(|_| KcFailure::unknown())?; + if matches!(response.status(), StatusCode::OK | StatusCode::NO_CONTENT) { + return Ok(response); + } + let status = response.status(); + let bytes = Self::bounded_response(response).await.unwrap_or_default(); + let v: Value = serde_json::from_slice(&bytes).unwrap_or(Value::Null); + let raw = v + .get("error_code") + .or_else(|| v.get("error")) + .and_then(Value::as_str); + let code = match raw { + Some("invalid_request") => "invalid_request", + Some("invalid_upload_signature") => "invalid_upload_signature", + Some("object_not_found") => "object_not_found", + Some("lifecycle_busy") => "lifecycle_busy", + Some("processing_timeout") => "processing_timeout", + Some("resource_limit") => "resource_limit", + Some("storage_unavailable") => "storage_unavailable", + Some("extraction_failed") => "extraction_failed", + Some("empty_document") => "empty_document", + Some("collection_dimension_mismatch") => "collection_dimension_mismatch", + Some("vector_store_unavailable") => "vector_store_unavailable", + Some("collection_not_found") => "collection_not_found", + Some("quota_exceeded") => "quota_exceeded", + Some("embedding_auth") => "embedding_auth", + Some("embedding_quota") => "embedding_quota", + Some("embedding_rejected") => "embedding_rejected", + Some("embedding_model_not_found") => "embedding_model_not_found", + Some("embedding_unavailable") => "embedding_unavailable", + Some("embedding_invalid_response") => "embedding_invalid_response", + _ => "kc_outcome_unknown", + }; + Err(KcFailure { + code, + status: if status.is_client_error() + && ((status != StatusCode::UNAUTHORIZED && status != StatusCode::FORBIDDEN) + || matches!(code, "invalid_upload_signature" | "embedding_quota")) + { + status + } else { + StatusCode::SERVICE_UNAVAILABLE + }, + }) + } + async fn bounded_response(mut response: reqwest::Response) -> Result, KcFailure> { + let mut data = Vec::new(); + while let Some(chunk) = response.chunk().await.map_err(|_| KcFailure::unknown())? { + if data.len() + chunk.len() > 64 * 1024 { + return Err(KcFailure::unknown()); + } + data.extend_from_slice(&chunk); + } + Ok(data) + } + pub async fn json( + &self, + method: Method, + path: &str, + body: Option, + ) -> Result { + let mut request = self.request(method, path); + if let Some(body) = body { + request = request.json(&body); + } + let response = Self::checked(request.send().await).await?; + serde_json::from_slice(&Self::bounded_response(response).await?) + .map_err(|_| KcFailure::unknown()) + } + /// The frozen collection acknowledgement is exactly an empty JSON object. + pub async fn collection_ack( + &self, + method: Method, + path: &str, + body: Option, + ) -> Result<(), KcFailure> { + let acknowledgement: Value = self.json(method, path, body).await?; + if acknowledgement != json!({}) { + return Err(KcFailure::unknown()); + } + Ok(()) + } + pub async fn presign( + &self, + collection: &str, + hash: &str, + content_type: &str, + ) -> Result { + let v:Value=self.json(Method::POST,"/internal/v1/storage/presigned-url",Some(json!({"collection":collection,"key":hash,"content_type":content_type,"expires_in_seconds":300}))).await?; + let url = v["presigned_url"].as_str().ok_or_else(KcFailure::unknown)?; + let prefix = format!("/s3/{collection}/{hash}?"); + if !url.starts_with(&prefix) || url.contains('#') || url.len() > 4096 { + return Err(KcFailure::unknown()); + } + let parsed = + Url::parse(&format!("http://local.invalid{url}")).map_err(|_| KcFailure::unknown())?; + let expiry: Vec<_> = parsed + .query_pairs() + .filter(|(k, _)| k == "X-Amz-Expires") + .collect(); + if expiry.len() != 1 || expiry[0].1 != "300" { + return Err(KcFailure::unknown()); + } + Ok(url.into()) + } + pub async fn upload( + &self, + collection: &str, + hash: &str, + query: &str, + headers: &HeaderMap, + body: Body, + ) -> Result { + let path = format!("/internal/v1/storage/upload/{collection}/{hash}?{query}"); + let mut request = self.request(Method::PUT, &path); + for (name, value) in headers { + if matches!( + name.as_str(), + "content-type" | "content-length" | "content-md5" + ) || name.as_str().starts_with("x-amz-") + { + request = request.header(name, value); + } + } + // Bound actual bytes as well as any declared length. No buffering of an upload. + use futures::StreamExt; + let mut stream = body.into_data_stream(); + let bounded: std::pin::Pin< + Box> + Send>, + > = Box::pin(async_stream::try_stream! { + let mut size=0usize; + while let Some(chunk)=stream.next().await { + let chunk=chunk.map_err(|_|std::io::Error::other("upload body failed"))?; + size=size.checked_add(chunk.len()).ok_or_else(||std::io::Error::other("upload limit exceeded"))?; + if size>MAX_UPLOAD_BYTES {Err(std::io::Error::other("upload limit exceeded"))?;} + yield chunk; + } + }); + let response = Self::checked( + request + .body(reqwest::Body::wrap_stream(bounded)) + .send() + .await, + ) + .await?; + let status = response.status(); + Self::bounded_response(response).await?; + Ok(status) + } +} +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn configuration_is_optional_and_rejects_credential_urls_and_shared_secrets() { + assert!(KcConfig::parse(None, None, None, false).unwrap().is_none()); + let token = "synthetic-internal-token-32-bytes-long"; + let hook = "synthetic-webhook-token-32-bytes-long"; + for url in [ + "http://localhost:5", + "https://u:p@localhost", + "https://localhost?token=x", + "https://localhost/base", + "file:///tmp/kc", + ] { + assert!(KcConfig::parse(Some(url), Some(token), Some(hook), false).is_err()); + } + assert!( + KcConfig::parse(Some("https://localhost"), Some(token), Some(token), false).is_err() + ); + let c = KcConfig::parse(Some("https://localhost"), Some(token), Some(hook), false) + .unwrap() + .unwrap(); + assert!(!format!("{c:?}").contains(token)); + assert!(!format!("{c:?}").contains(hook)); + } +} diff --git a/crates/shaide/src/services/mod.rs b/crates/shaide/src/services/mod.rs index 5eff2ee..1c1079e 100644 --- a/crates/shaide/src/services/mod.rs +++ b/crates/shaide/src/services/mod.rs @@ -8,3 +8,5 @@ pub mod platform_models; pub mod responses; pub mod collection_context; + +pub mod knowledge_center; diff --git a/documentation/upload-lifecycle.md b/documentation/upload-lifecycle.md new file mode 100644 index 0000000..cf88a0a --- /dev/null +++ b/documentation/upload-lifecycle.md @@ -0,0 +1,105 @@ +# Collection and upload orchestration + +shaide exposes the existing Control Panel APIs for collection CRUD, memberships, +file deletion, five-minute upload URLs and object-created notifications. Collection +mutations require an administrator. Administrators see all collections and member +IDs; users see their own memberships with `users: null`. Upload and file deletion +require an administrator or an active member of a collection with uploads enabled. + +KC is optional. Without it, the server's ordinary APIs and collection metadata +listing remain available. Operations that need KC return 503 before reserving or +changing metadata. Local name, description, policy and membership updates use +durable collection fences. Administrative membership removal and a PATCH containing +only `organization_id` and `can_users_upload: false` use serialized metadata +transactions even while reservations or recovery work exist. They retain every +file and operation fence; current owner access is rechecked at upload and ingestion +admission. Grants, mixed updates and model changes still require the collection +fence. Public knowledge-center search is outside this change; +MCP continues through the gateway, KC-MCP and KC. + +## Configuration + +Set `KNOWLEDGE_CENTER_URL` to KC's fixed origin (no path, query, fragment or +credentials), `KNOWLEDGE_CENTER_INTERNAL_TOKEN` to a distinct secret with at least +32 bytes, and `S3_WEBHOOK_TOKEN` to the matching RustFS notification bearer secret +with at least 32 bytes. `KNOWLEDGE_CENTER_ALLOW_LOCAL_HTTP=true` explicitly permits +HTTP for an isolated local Compose network. HTTPS is required otherwise. The +internal and notification tokens must differ from each other and from `JWT_SECRET`. +Leave all three KC settings unset when KC is disabled. + +All RustFS credentials, SDKs, notification provisioning and storage operations +belong to KC. KC configures RustFS to POST notifications to +`/v1/object-storage/event` with `Authorization: Bearer `. +Real RustFS support for this header must be verified with the pinned integration +image before publication; fake HTTP tests do not establish that capability. + +## Upload transport + +`POST /v1/object-storage/presigned-url` accepts the established +`organization_id`, `file_name`, `content_type` body and returns +`{"presigned_url":"/s3/collection-N/?..."}`. The file identity is +server-issued lowercase 64-hex, independent of content. The reservation is durable +and valid for 300 seconds. KC creates the signature. shaide stores only the SHA-256 +fingerprint of the signed query, never the query or storage credentials. + +The client PUTs bytes to that public relative URL with the original query, +`Content-Type` and exact `Content-Length`. shaide checks the reservation, canonical +path, declared size and current uploader access, then streams bytes to authenticated +`PUT /internal/v1/storage/upload/{collection}/{file_hash}` on its configured KC +origin. It forwards the original signed query and content headers. KC validates +and persists the upload to RustFS. Only PUT is exposed; redirects, arbitrary +upstream URLs and bucket browsing are unsupported. Both services enforce a 16 MiB +bound. Supported types are PDF, DOCX, PNG and JPEG. + +shaide validates notifications against scoped metadata and trusted KC +`GET /internal/v1/storage/objects/{collection}/{file_hash}` results. Notification +size claims do not authorize ingestion. A notification arriving before the upload +response is recorded durably and processed after KC acknowledges the write. +Only one durable ingestion claim is allowed. Metadata becomes `processed` only +following a complete, valid KC ingestion acknowledgement. Duplicate completed +notifications do not index again. Ingest, upload and deletion failures are sanitized +and retained; acknowledgements with missing/invalid fields do not count as success. + +## Durable operations and recovery + +Migration 0004 adds `collection_operations` and `file_operations` without changing +existing file status names or deleting existing rows. Collections with legacy non-final file metadata receive a recovery fence because those rows cannot prove the remote outcome. Every SQL transaction ends +before an HTTP call. Collection fences block reservations and collection changes; +file fences block duplicate uploads, overlapping ingestion/deletion and destructive +collection changes. A disconnect or process restart may leave `active`, `uploading`, +`uploaded` or `ingesting` work. A failed/uncertain call leaves `recovery_required`. +Elapsed time never releases those fences. An expired reservation is not proof that +a remote writer stopped. File/collection deletion and model replacement remain +blocked until the recorded operation is settled. + +Recovery is an explicit operator task, scoped to one operation ID and token. Do not +run blanket stale cleanup, retry remote mutations automatically, delete a file's +metadata first, or treat local request cancellation as a remote write fence. + +1. Inspect the operation row, retained file status and update payload. Identify + the exact collection/file and all possibly accepted KC work. +2. Establish that those public facade handlers, KC handlers and downstream writes have finished or been + stopped and fenced. A timeout, elapsed lease or local restart is insufficient. +3. Reconcile the exact object's metadata and complete indexing/deletion outcome + through KC. Object existence alone does not prove vectors are complete. +4. In a reviewed, scoped SQL transaction, settle the metadata to the verified + outcome and remove only the matching operation token. Keep the fence if the + remote outcome cannot be established. There is no public force-unlock API. + +`DELETE /v1/rag/delete-stale-files` returns 409 when stale or recovery work exists; +it does not discard operation records or files. Operators must settle each one. +Downgrade migration 0004 only after all operations have been settled: dropping the +operation tables removes protection against overlapping remote work. + +Populated collection embedding-model replacement returns 409. The frozen KC +collection DELETE removes stored files as well as vectors, so it cannot be used to +rebuild vectors while preserving documents. Empty collections can change models +through acknowledged KC delete/create under a durable fence. Collection create, +delete and both replacement phases require the frozen empty JSON object `{}`. Any +other successful JSON body is an uncertain outcome: metadata and the operation +fence are retained for reconciliation. A vector-only +replacement protocol requires a separately agreed KC contract; no new KC endpoint is assumed. + +Native HTTP contract tests use explicit synthetic KC fakes and disposable SQLite +files. Real storage signature/header verification, extraction/embedding/Qdrant/MCP +E2E and container acceptance remain separate integration gates.