diff --git a/crates/agentic-server-core/migrations/0008_vector_store_batches.sql b/crates/agentic-server-core/migrations/0008_vector_store_batches.sql new file mode 100644 index 00000000..3588846a --- /dev/null +++ b/crates/agentic-server-core/migrations/0008_vector_store_batches.sql @@ -0,0 +1,16 @@ +-- Membership snapshots survive detach and source deletion. Store deletion removes its batch history. +ALTER TABLE file_search_attachments ADD COLUMN generation TEXT NOT NULL DEFAULT ''; +CREATE TABLE file_search_batches ( + id TEXT PRIMARY KEY, store_id TEXT NOT NULL REFERENCES file_search_stores(id) ON DELETE CASCADE, + created_at BIGINT NOT NULL, cancelled BIGINT NOT NULL DEFAULT 0 +); +CREATE TABLE file_search_jobs ( + id TEXT PRIMARY KEY, batch_id TEXT NOT NULL REFERENCES file_search_batches(id) ON DELETE CASCADE, + store_id TEXT NOT NULL, file_id TEXT NOT NULL, generation TEXT NOT NULL, + created_at BIGINT NOT NULL, updated_at BIGINT NOT NULL, state TEXT NOT NULL, + options TEXT NOT NULL, identity TEXT NOT NULL, snapshot TEXT NOT NULL, + claim_token TEXT, lease_until BIGINT, attempts BIGINT NOT NULL DEFAULT 0, + UNIQUE(batch_id, file_id) +); +CREATE INDEX file_search_job_claims ON file_search_jobs(state, lease_until, created_at, id); +CREATE INDEX file_search_job_members ON file_search_jobs(batch_id, created_at, file_id); diff --git a/crates/agentic-server-core/src/storage/file_search.rs b/crates/agentic-server-core/src/storage/file_search.rs index 85ba5457..29ca5ee6 100644 --- a/crates/agentic-server-core/src/storage/file_search.rs +++ b/crates/agentic-server-core/src/storage/file_search.rs @@ -7,6 +7,8 @@ use sqlx::FromRow; use tokio_util::sync::CancellationToken; use super::{DbPool, DbTransaction}; +#[path = "vector_store_batches.rs"] +pub(crate) mod batches; #[path = "vector_store_lifecycle.rs"] mod lifecycle; use crate::types::file_search::{ @@ -20,6 +22,8 @@ const MAX_CORPUS_CHUNKS: i64 = 10_000; pub(crate) struct FileSearchStorage { pool: Arc, pgvector: Option, + #[cfg(test)] + pub(crate) batch_test_hooks: Option>, } #[derive(FromRow)] @@ -88,7 +92,11 @@ impl Collection { impl FileSearchStorage { #[cfg(test)] pub(crate) fn new(pool: Arc) -> Self { - Self { pool, pgvector: None } + Self { + pool, + pgvector: None, + batch_test_hooks: None, + } } pub(crate) fn with_backend( @@ -96,7 +104,12 @@ impl FileSearchStorage { backend: &crate::types::file_search::FileSearchBackend, ) -> Result { let pgvector = super::pgvector::PgvectorStorage::from_config(&pool, backend)?; - Ok(Self { pool, pgvector }) + Ok(Self { + pool, + pgvector, + #[cfg(test)] + batch_test_hooks: None, + }) } pub(crate) fn vector_dimensions(&self) -> Option { @@ -402,6 +415,12 @@ impl FileSearchStorage { return Err(FileSearchError::NotFound("File not found or expired".into())); } } + if let Some(store_id) = store_id { + batches::invalidate(&mut tx, Some(store_id), Some(id)).await?; + } else if matches!(collection, Collection::Stores) { + lifecycle::lock_store(&mut tx, id).await?; + batches::invalidate(&mut tx, Some(id), None).await?; + } let filter = if store_id.is_some() { " AND store_id = $2" } else { "" }; let sql = format!( "DELETE FROM {} WHERE {} = $1{filter}", @@ -540,6 +559,7 @@ pub(crate) async fn database_now(connection: &mut sqlx::AnyConnection) -> Result } async fn delete_file_in_transaction(tx: &mut DbTransaction<'_>, id: &str) -> Result<(), FileSearchError> { + batches::invalidate(tx, None, Some(id)).await?; sqlx::query("INSERT INTO file_search_blob_cleanup (file_id) SELECT id FROM file_search_files WHERE id = $1 AND content_base64 = '' ON CONFLICT (file_id) DO NOTHING") .bind(id).execute(&mut **tx).await?; sqlx::query("DELETE FROM file_search_files WHERE id = $1") diff --git a/crates/agentic-server-core/src/storage/schema.rs b/crates/agentic-server-core/src/storage/schema.rs index bc943360..5e43dbc9 100644 --- a/crates/agentic-server-core/src/storage/schema.rs +++ b/crates/agentic-server-core/src/storage/schema.rs @@ -15,8 +15,8 @@ use crate::config::DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS; type DbResult = Result; const POSTGRES_SCHEMA_ADVISORY_LOCK: i64 = 7_194_963_546_799_751; -const REQUIRED_POSTGRES_SCHEMA_COLUMN_COUNT: i64 = 49; -const REQUIRED_POSTGRES_CONSTRAINT_COUNT: i64 = 15; +const REQUIRED_POSTGRES_SCHEMA_COLUMN_COUNT: i64 = 68; +const REQUIRED_POSTGRES_CONSTRAINT_COUNT: i64 = 20; const REQUIRED_POSTGRES_INTEGER_COLUMN_COUNT: i64 = 4; const POSTGRES_INTEGER_WIDENING_SQL: &str = " ALTER TABLE conversations @@ -111,6 +111,25 @@ where ('file_search_attachments', 'data', 'text', 'NO'), \ ('file_search_attachments', 'status', 'text', 'NO'), \ ('file_search_attachments', 'parsed_content', 'text', 'YES'), \ + ('file_search_attachments', 'generation', 'text', 'NO'), \ + ('file_search_batches', 'id', 'text', 'NO'), \ + ('file_search_batches', 'store_id', 'text', 'NO'), \ + ('file_search_batches', 'created_at', 'bigint', 'NO'), \ + ('file_search_batches', 'cancelled', 'bigint', 'NO'), \ + ('file_search_jobs', 'id', 'text', 'NO'), \ + ('file_search_jobs', 'batch_id', 'text', 'NO'), \ + ('file_search_jobs', 'store_id', 'text', 'NO'), \ + ('file_search_jobs', 'file_id', 'text', 'NO'), \ + ('file_search_jobs', 'generation', 'text', 'NO'), \ + ('file_search_jobs', 'created_at', 'bigint', 'NO'), \ + ('file_search_jobs', 'updated_at', 'bigint', 'NO'), \ + ('file_search_jobs', 'state', 'text', 'NO'), \ + ('file_search_jobs', 'options', 'text', 'NO'), \ + ('file_search_jobs', 'identity', 'text', 'NO'), \ + ('file_search_jobs', 'snapshot', 'text', 'NO'), \ + ('file_search_jobs', 'claim_token', 'text', 'YES'), \ + ('file_search_jobs', 'lease_until', 'bigint', 'YES'), \ + ('file_search_jobs', 'attempts', 'bigint', 'NO'), \ ('file_search_chunks', 'store_id', 'text', 'NO'), \ ('file_search_chunks', 'file_id', 'text', 'NO'), \ ('file_search_chunks', 'chunk_index', 'bigint', 'NO'), \ @@ -154,6 +173,11 @@ where 'FOREIGN KEY (previous_response_id) REFERENCES responses(id) ON DELETE SET NULL'), \ ('conversations', 'f', \ 'FOREIGN KEY (latest_response_id) REFERENCES responses(id) ON DELETE SET NULL'), \ + ('file_search_batches', 'p', 'PRIMARY KEY (id)'), \ + ('file_search_batches', 'f', 'FOREIGN KEY (store_id) REFERENCES file_search_stores(id) ON DELETE CASCADE'), \ + ('file_search_jobs', 'p', 'PRIMARY KEY (id)'), \ + ('file_search_jobs', 'f', 'FOREIGN KEY (batch_id) REFERENCES file_search_batches(id) ON DELETE CASCADE'), \ + ('file_search_jobs', 'u', 'UNIQUE (batch_id, file_id)'), \ ('file_search_files', 'p', 'PRIMARY KEY (id)'), \ ('file_search_blob_cleanup', 'p', 'PRIMARY KEY (file_id)'), \ ('file_search_stores', 'p', 'PRIMARY KEY (id)'), \ @@ -318,7 +342,7 @@ pub(crate) async fn pin_postgres_persistence_schema(connection: &mut sqlx::AnyCo WHERE table_namespace.nspname = ANY(current_schemas(false)) \ AND table_relation.relkind IN ('r', 'p', 'v', 'm', 'f') \ AND table_relation.relname IN ('_sqlx_migrations', 'conversations', 'items', 'responses', \ - 'file_search_files', 'file_search_stores', 'file_search_attachments', 'file_search_chunks', 'file_search_blob_cleanup') \ + 'file_search_files', 'file_search_stores', 'file_search_attachments', 'file_search_chunks', 'file_search_blob_cleanup', 'file_search_batches', 'file_search_jobs') \ ORDER BY table_namespace.nspname::text", ) .fetch_all(&mut *connection) @@ -407,6 +431,14 @@ pub(crate) async fn verify_persistence_ready(pool: &DbPool) -> DbResult<()> { ('items', 'INSERT'), \ ('responses', 'SELECT'), \ ('responses', 'INSERT'), \ + ('file_search_batches', 'SELECT'), \ + ('file_search_batches', 'INSERT'), \ + ('file_search_batches', 'UPDATE'), \ + ('file_search_batches', 'DELETE'), \ + ('file_search_jobs', 'SELECT'), \ + ('file_search_jobs', 'INSERT'), \ + ('file_search_jobs', 'UPDATE'), \ + ('file_search_jobs', 'DELETE'), \ ('file_search_files', 'SELECT'), \ ('file_search_files', 'INSERT'), \ ('file_search_files', 'UPDATE'), \ @@ -427,7 +459,7 @@ pub(crate) async fn verify_persistence_ready(pool: &DbPool) -> DbResult<()> { ('file_search_chunks', 'INSERT') \ ) \ SELECT current_setting('transaction_read_only') = 'off' \ - AND COUNT(table_relation.oid) = 25 \ + AND COUNT(table_relation.oid) = 33 \ AND COALESCE(BOOL_AND( \ has_table_privilege(current_user, table_relation.oid, required.privilege) \ ), false) \ @@ -458,8 +490,10 @@ pub(crate) async fn verify_persistence_ready(pool: &DbPool) -> DbResult<()> { "SELECT id FROM responses LIMIT 0", "SELECT id, created_at, data, content_type, content_base64, expires_at, purpose FROM file_search_files LIMIT 0", "SELECT file_id FROM file_search_blob_cleanup LIMIT 0", + "SELECT id, store_id, created_at, cancelled FROM file_search_batches LIMIT 0", + "SELECT id,batch_id,store_id,file_id,generation,created_at,updated_at,state,options,identity,snapshot,claim_token,lease_until,attempts FROM file_search_jobs LIMIT 0", "SELECT id, created_at, data, embedding_identity, embedding_dimensions, last_active_at, expires_after_days, expires_at, lifecycle_status FROM file_search_stores LIMIT 0", - "SELECT store_id, file_id, created_at, usage_bytes, storage_bytes, data, status, parsed_content FROM file_search_attachments LIMIT 0", + "SELECT store_id, file_id, created_at, usage_bytes, storage_bytes, data, status, parsed_content, generation FROM file_search_attachments LIMIT 0", "SELECT store_id, file_id, chunk_index, data FROM file_search_chunks LIMIT 0", ] { sqlx::query(statement).execute(&mut *connection).await?; @@ -731,6 +765,14 @@ mod tests { .execute(pool.as_ref()) .await .unwrap(); + assert!( + verify_persistence_ready(pool.as_ref()).await.is_err(), + "batch migration is required" + ); + sqlx::raw_sql(include_str!("../../migrations/0008_vector_store_batches.sql")) + .execute(pool.as_ref()) + .await + .unwrap(); verify_persistence_ready(pool.as_ref()).await.unwrap(); wrapper.ensure_schema_ready_with_marker(true).await.unwrap(); } @@ -849,6 +891,14 @@ mod tests { .execute(&mut *connection) .await .unwrap(); + assert!( + supervisor.ensure_schema_ready_with_marker(true).await.is_err(), + "batch migration is required" + ); + sqlx::raw_sql(include_str!("../../migrations/0008_vector_store_batches.sql")) + .execute(&mut *connection) + .await + .unwrap(); supervisor.ensure_schema_ready_with_marker(true).await.unwrap(); supervisor.pool.close().await; sqlx::query("SET search_path TO public") @@ -892,6 +942,7 @@ mod tests { include_str!("../../migrations/0005_file_search.sql"), include_str!("../../migrations/0006_file_expiration.sql"), include_str!("../../migrations/0007_vector_store_lifecycle.sql"), + include_str!("../../migrations/0008_vector_store_batches.sql"), ] { sqlx::raw_sql(migration) .execute(&mut *connection) diff --git a/crates/agentic-server-core/src/storage/vector_store_batches.rs b/crates/agentic-server-core/src/storage/vector_store_batches.rs new file mode 100644 index 00000000..72543420 --- /dev/null +++ b/crates/agentic-server-core/src/storage/vector_store_batches.rs @@ -0,0 +1,947 @@ +//! Durable batch membership and portable claims. Parent writes serialize every job transition. +use super::{FileSearchStorage, PreparedAttachment, database_now, lifecycle, publish_attachment, serialize_chunks}; +use crate::{ + storage::DbTransaction, + types::file_search::{ + AttachFileRequest, AttachmentStatus, BatchStatus, FileBatchObject, FileCounts, FileSearchError, ListOrder, + ListParams, VectorStoreFileError, VectorStoreFileObject, invalid, + }, +}; + +// Per-storage timing controls keep concurrency regressions deterministic without global failpoints. +#[cfg(test)] +#[derive(Default)] +pub(crate) struct BatchTestHooks { + pub after_commit: Option, + pub renewal_delay: std::time::Duration, + pub renewal_started: tokio::sync::Notify, +} + +#[cfg(test)] +#[derive(Default)] +pub(crate) struct CommitBarrier { + pub reached: tokio::sync::Notify, + pub resume: tokio::sync::Notify, + used: std::sync::atomic::AtomicBool, +} + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub(crate) struct BatchFileOptions { + pub request: AttachFileRequest, + pub contextual_identity: Option, +} + +#[derive(Clone, Copy, Debug)] +enum JobState { + Queued, + Running, + Completed, + Failed, + Cancelled, +} + +impl TryFrom<&str> for JobState { + type Error = FileSearchError; + fn try_from(value: &str) -> Result { + match value { + "queued" => Ok(Self::Queued), + "running" => Ok(Self::Running), + "completed" => Ok(Self::Completed), + "failed" => Ok(Self::Failed), + "cancelled" => Ok(Self::Cancelled), + _ => Err(FileSearchError::Unavailable("Invalid durable job state".into())), + } + } +} + +#[derive(Clone, Debug)] +pub(crate) struct JobId(String); +#[derive(Clone, Debug)] +pub(crate) struct ClaimToken(String); +#[derive(Clone, Debug)] +pub(crate) struct AttachmentGeneration(String); +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum ClaimOutcome { + Applied, + LostClaim, +} +#[derive(Clone, Debug)] +pub(crate) struct ClaimedFileJob { + id: JobId, + token: ClaimToken, + generation: AttachmentGeneration, + pub store_id: String, + pub options: AttachFileRequest, + pub identity: String, + pub contextual_identity: Option, +} +#[derive(sqlx::FromRow)] +struct Job { + id: String, + store_id: String, + generation: String, + options: String, + identity: String, +} + +impl FileSearchStorage { + pub(crate) async fn create_batch( + &self, + store_id: &str, + identity: &str, + members: &[(BatchFileOptions, VectorStoreFileObject)], + ) -> Result { + let mut tx = self.pool.begin().await?; + let mut ids = members + .iter() + .map(|(r, _)| r.request.file_id.as_str()) + .collect::>(); + ids.sort_unstable(); + for id in ids { + sqlx::query("UPDATE file_search_files SET id = id WHERE id = $1") + .bind(id) + .execute(&mut *tx) + .await?; + } + lifecycle::lock_store(&mut tx, store_id).await?; + let now = database_now(&mut tx).await?; + lifecycle::require_live_store(&mut tx, store_id, now).await?; + let id = format!("vsfb_{}", uuid::Uuid::now_v7().simple()); + sqlx::query("INSERT INTO file_search_batches (id,store_id,created_at) VALUES ($1,$2,$3)") + .bind(&id) + .bind(store_id) + .bind(now) + .execute(&mut *tx) + .await?; + let mut counts = FileCounts::default(); + for (saved_options, queued) in members { + let options = &saved_options.request; + let exists: Option = sqlx::query_scalar( + "SELECT id FROM file_search_files WHERE id=$1 AND (expires_at IS NULL OR expires_at>$2)", + ) + .bind(&options.file_id) + .bind(now) + .fetch_optional(&mut *tx) + .await?; + if exists.is_none() { + return Err(FileSearchError::NotFound( + "Batch source file not found or expired".into(), + )); + } + let existing: Option = + sqlx::query_scalar("SELECT data FROM file_search_attachments WHERE store_id=$1 AND file_id=$2") + .bind(store_id) + .bind(&options.file_id) + .fetch_optional(&mut *tx) + .await?; + let mut object = queued.clone(); + object.created_at = now; + let generation = uuid::Uuid::now_v7().to_string(); + let state = if let Some(existing) = existing { + object = serde_json::from_str(&existing)?; + if object.status != AttachmentStatus::Completed { + return Err(FileSearchError::Conflict( + "Batch file already has an unfinished or unsuccessful attachment; detach it before retrying" + .into(), + )); + } + counts.completed += 1; + "completed" + } else { + sqlx::query("INSERT INTO file_search_attachments (store_id,file_id,created_at,usage_bytes,storage_bytes,data,status,generation) VALUES ($1,$2,$3,0,0,$4,'in_progress',$5)") + .bind(store_id).bind(&options.file_id).bind(now).bind(serde_json::to_string(&object)?).bind(&generation).execute(&mut *tx).await?; + counts.in_progress += 1; + "queued" + }; + counts.total += 1; + sqlx::query("INSERT INTO file_search_jobs (id,batch_id,store_id,file_id,generation,created_at,updated_at,state,options,identity,snapshot) VALUES ($1,$2,$3,$4,$5,$6,$6,$7,$8,$9,$10)") + .bind(uuid::Uuid::now_v7().to_string()).bind(&id).bind(store_id).bind(&options.file_id).bind(generation).bind(now).bind(state).bind(serde_json::to_string(saved_options)?).bind(identity).bind(serde_json::to_string(&object)?).execute(&mut *tx).await?; + } + // Build the response from this transaction's membership. After COMMIT succeeds, + // no fallible response read may turn success into a retry of batch creation. + let object = FileBatchObject { + id, + object: "vector_store.files_batch".into(), + created_at: now, + vector_store_id: store_id.into(), + status: if counts.in_progress > 0 { + BatchStatus::InProgress + } else { + BatchStatus::Completed + }, + file_counts: counts, + }; + tx.commit().await?; + #[cfg(test)] + if let Some(barrier) = self + .batch_test_hooks + .as_ref() + .and_then(|hooks| hooks.after_commit.as_ref()) + { + if !barrier.used.swap(true, std::sync::atomic::Ordering::SeqCst) { + barrier.reached.notify_one(); + barrier.resume.notified().await; + } + } + Ok(object) + } + + pub(crate) async fn batch(&self, store_id: &str, id: &str) -> Result { + let mut tx = self.pool.begin().await?; + lifecycle::lock_store(&mut tx, store_id).await?; + let now = database_now(&mut tx).await?; + lifecycle::require_live_store(&mut tx, store_id, now).await?; + let row: Option<(i64, i64)> = + sqlx::query_as("SELECT created_at,cancelled FROM file_search_batches WHERE id=$1 AND store_id=$2") + .bind(id) + .bind(store_id) + .fetch_optional(&mut *tx) + .await?; + let (created_at, cancelled) = row.ok_or_else(|| FileSearchError::NotFound("File batch not found".into()))?; + let totals: Vec<(String, i64)> = + sqlx::query_as("SELECT state,COUNT(*) FROM file_search_jobs WHERE batch_id=$1 GROUP BY state") + .bind(id) + .fetch_all(&mut *tx) + .await?; + let mut counts = FileCounts::default(); + for (state, count) in totals { + match JobState::try_from(state.as_str())? { + JobState::Queued | JobState::Running => counts.in_progress += count, + JobState::Completed => counts.completed += count, + JobState::Failed => counts.failed += count, + JobState::Cancelled => counts.cancelled += count, + } + counts.total += count; + } + let status = if cancelled != 0 { + BatchStatus::Cancelled + } else if counts.in_progress > 0 { + BatchStatus::InProgress + } else { + BatchStatus::Completed + }; + tx.commit().await?; + Ok(FileBatchObject { + id: id.into(), + object: "vector_store.files_batch".into(), + created_at, + vector_store_id: store_id.into(), + status, + file_counts: counts, + }) + } + + pub(crate) async fn batch_files( + &self, + store_id: &str, + id: &str, + params: &ListParams, + ) -> Result, FileSearchError> { + self.batch(store_id, id).await?; + let cursor = params.after.as_deref().or(params.before.as_deref()).unwrap_or(""); + if !cursor.is_empty() { + let found: Option = + sqlx::query_scalar("SELECT file_id FROM file_search_jobs WHERE batch_id=$1 AND file_id=$2") + .bind(id) + .bind(cursor) + .fetch_optional(self.pool.as_ref()) + .await?; + if found.is_none() { + return invalid("Pagination cursor does not exist in this batch"); + } + } + // All members share the atomic batch creation timestamp, so file_id breaks every tie. + let ascending = matches!(params.order.unwrap_or_default(), ListOrder::Asc) == params.before.is_none(); + let operator = if ascending { ">" } else { "<" }; + let order = if ascending { "ASC" } else { "DESC" }; + let filter = params.filter.map_or("", AttachmentStatus::as_str); + let sql = format!( + "SELECT snapshot FROM file_search_jobs WHERE batch_id=$1 AND ($2='' OR file_id {operator} $2) AND ($3='' OR state=$3 OR ($3='in_progress' AND state IN ('queued','running'))) ORDER BY file_id {order} LIMIT $4" + ); + let rows: Vec = sqlx::query_scalar(&sql) + .bind(id) + .bind(cursor) + .bind(filter) + .bind(i64::try_from(params.limit.unwrap_or(20) + 1).unwrap_or(101)) + .fetch_all(self.pool.as_ref()) + .await?; + rows.iter() + .map(|s| serde_json::from_str(s).map_err(Into::into)) + .collect() + } + + pub(crate) async fn cancel_batch(&self, store_id: &str, id: &str) -> Result { + let mut tx = self.pool.begin().await?; + lifecycle::lock_store(&mut tx, store_id).await?; + let now = database_now(&mut tx).await?; + lifecycle::require_live_store(&mut tx, store_id, now).await?; + let found=sqlx::query("UPDATE file_search_batches SET cancelled=1 WHERE id=$1 AND store_id=$2 AND EXISTS (SELECT 1 FROM file_search_jobs WHERE batch_id=$1 AND state IN ('queued','running'))").bind(id).bind(store_id).execute(&mut *tx).await?.rows_affected(); + if found > 0 { + let rows:Vec<(String,String,String)>=sqlx::query_as("SELECT id,generation,snapshot FROM file_search_jobs WHERE batch_id=$1 AND state IN ('queued','running') ORDER BY id").bind(id).fetch_all(&mut *tx).await?; + for (job, generation, data) in rows { + let mut object: VectorStoreFileObject = serde_json::from_str(&data)?; + object.status = AttachmentStatus::Cancelled; + terminal(&mut tx, &job, &generation, &object, "cancelled", now).await?; + } + } + tx.commit().await?; + self.batch(store_id, id).await + } + + pub(crate) async fn claim_next( + &self, + lease: i64, + per_batch: usize, + scan: usize, + ) -> Result, FileSearchError> { + let mut connection = self.pool.acquire().await?; + let now = database_now(&mut connection).await?; + drop(connection); + let candidates:Vec<(String,String)>=sqlx::query_as("SELECT id,store_id FROM file_search_jobs WHERE state='queued' OR (state='running' AND lease_until<=$1) ORDER BY created_at,id LIMIT $2").bind(now).bind(i64::try_from(scan).unwrap_or(1000)).fetch_all(self.pool.as_ref()).await?; + for (id, store) in candidates { + let mut tx = self.pool.begin().await?; + if let Err(error) = lifecycle::lock_store(&mut tx, &store).await { + if matches!(error, FileSearchError::NotFound(_)) { + continue; + } + return Err(error); + } + sqlx::query("UPDATE file_search_jobs SET id=id WHERE id=$1") + .bind(&id) + .execute(&mut *tx) + .await?; + let now = database_now(&mut tx).await?; + let token = uuid::Uuid::now_v7().to_string(); + let updated=sqlx::query("UPDATE file_search_jobs SET state='running',claim_token=$2,lease_until=$3,updated_at=$4,attempts=attempts+1 WHERE id=$1 AND (state='queued' OR (state='running' AND lease_until<=$4)) AND (SELECT COUNT(*) FROM file_search_jobs j WHERE j.batch_id=file_search_jobs.batch_id AND j.state='running' AND j.lease_until>$4)<$5") + .bind(&id).bind(&token).bind(now+lease).bind(now).bind(i64::try_from(per_batch).unwrap_or(32)).execute(&mut *tx).await?.rows_affected(); + if updated == 0 { + continue; + } + let row: Job = + sqlx::query_as("SELECT id,store_id,generation,options,identity FROM file_search_jobs WHERE id=$1") + .bind(&id) + .fetch_one(&mut *tx) + .await?; + let options: BatchFileOptions = serde_json::from_str(&row.options)?; + let job = ClaimedFileJob { + id: JobId(row.id), + token: ClaimToken(token), + generation: AttachmentGeneration(row.generation), + store_id: row.store_id, + options: options.request, + contextual_identity: options.contextual_identity, + identity: row.identity, + }; + // Missing generations/expired parents become durable cancellation, never publication. + if !live(&mut tx, &job, now).await? { + let data: String = sqlx::query_scalar("SELECT snapshot FROM file_search_jobs WHERE id=$1") + .bind(&id) + .fetch_one(&mut *tx) + .await?; + let mut object: VectorStoreFileObject = serde_json::from_str(&data)?; + object.status = AttachmentStatus::Cancelled; + terminal(&mut tx, &id, &job.generation.0, &object, "cancelled", now).await?; + tx.commit().await?; + continue; + } + tx.commit().await?; + return Ok(Some(job)); + } + Ok(None) + } + + pub(crate) async fn renew_claim(&self, job: &ClaimedFileJob, lease: i64) -> Result { + let mut tx = self.pool.begin().await?; + if let Err(error) = lifecycle::lock_store(&mut tx, &job.store_id).await { + if matches!(error, FileSearchError::NotFound(_)) { + return Ok(ClaimOutcome::LostClaim); + } + return Err(error); + } + sqlx::query("UPDATE file_search_jobs SET id=id WHERE id=$1") + .bind(&job.id.0) + .execute(&mut *tx) + .await?; + let now = database_now(&mut tx).await?; + if !live(&mut tx, job, now).await? { + return Ok(ClaimOutcome::LostClaim); + } + let n=sqlx::query("UPDATE file_search_jobs SET lease_until=$3,updated_at=$4 WHERE id=$1 AND claim_token=$2 AND state='running' AND lease_until>$4").bind(&job.id.0).bind(&job.token.0).bind(now+lease).bind(now).execute(&mut *tx).await?.rows_affected(); + tx.commit().await?; + #[cfg(test)] + if let Some(hooks) = &self.batch_test_hooks { + hooks.renewal_started.notify_one(); + tokio::time::sleep(hooks.renewal_delay).await; + } + Ok(if n == 1 { + ClaimOutcome::Applied + } else { + ClaimOutcome::LostClaim + }) + } + + pub(crate) async fn release_claim(&self, job: &ClaimedFileJob) -> Result<(), FileSearchError> { + // A conditional single statement is sufficient: it does not write attachments or terminal state. + sqlx::query("UPDATE file_search_jobs SET state='queued',claim_token=NULL,lease_until=NULL WHERE id=$1 AND claim_token=$2 AND state='running'").bind(&job.id.0).bind(&job.token.0).execute(self.pool.as_ref()).await?; + Ok(()) + } + + #[allow( + clippy::too_many_lines, + reason = "keeps claim validation and shared publication in one transaction" + )] + pub(crate) async fn finish_claim( + &self, + job: &ClaimedFileJob, + result: Result, + ) -> Result { + let mut encoded = if let Ok(prepared) = &result { + Some(serialize_chunks(&prepared.chunks).await?) + } else { + None + }; + let mut tx = self.pool.begin().await?; + sqlx::query("UPDATE file_search_files SET id=id WHERE id=$1") + .bind(&job.options.file_id) + .execute(&mut *tx) + .await?; + if let Err(error) = lifecycle::lock_store(&mut tx, &job.store_id).await { + if matches!(error, FileSearchError::NotFound(_)) { + return Ok(ClaimOutcome::LostClaim); + } + return Err(error); + } + sqlx::query("UPDATE file_search_jobs SET id=id WHERE id=$1") + .bind(&job.id.0) + .execute(&mut *tx) + .await?; + let now = database_now(&mut tx).await?; + if !live(&mut tx, job, now).await? { + return Ok(ClaimOutcome::LostClaim); + } + let data:Option=sqlx::query_scalar("SELECT snapshot FROM file_search_jobs WHERE id=$1 AND claim_token=$2 AND state='running' AND lease_until>$3").bind(&job.id.0).bind(&job.token.0).bind(now).fetch_optional(&mut *tx).await?; + let Some(data) = data else { + return Ok(ClaimOutcome::LostClaim); + }; + let mut object: VectorStoreFileObject = serde_json::from_str(&data)?; + match result { + Ok(mut prepared) => { + prepared.object.created_at = object.created_at; + let current: String = + sqlx::query_scalar("SELECT data FROM file_search_attachments WHERE store_id=$1 AND file_id=$2") + .bind(&job.store_id) + .bind(&job.options.file_id) + .fetch_one(&mut *tx) + .await?; + let current: VectorStoreFileObject = serde_json::from_str(¤t)?; + if prepared.object.attributes != current.attributes { + prepared.object.attributes = current.attributes; + for chunk in &mut prepared.chunks { + chunk.attributes.clone_from(&prepared.object.attributes); + } + encoded = Some(serialize_chunks(&prepared.chunks).await?); + } + // Remove only our pending generation, then reuse the one atomic chunk writer. + sqlx::query("DELETE FROM file_search_attachments WHERE store_id=$1 AND file_id=$2 AND generation=$3") + .bind(&job.store_id) + .bind(&job.options.file_id) + .bind(&job.generation.0) + .execute(&mut *tx) + .await?; + let (chunks, bytes) = + encoded.ok_or_else(|| FileSearchError::Unavailable("Missing prepared chunks".into()))?; + publish_attachment(&mut tx, &job.store_id, &job.identity, &prepared, &chunks, bytes).await?; + sqlx::query("UPDATE file_search_attachments SET generation=$3 WHERE store_id=$1 AND file_id=$2") + .bind(&job.store_id) + .bind(&job.options.file_id) + .bind(&job.generation.0) + .execute(&mut *tx) + .await?; + object = prepared.object; + } + Err(error) => { + object.status = AttachmentStatus::Failed; + object.last_error = Some(error); + } + } + let now = database_now(&mut tx).await?; + let valid: Option = sqlx::query_scalar( + "SELECT id FROM file_search_jobs WHERE id=$1 AND claim_token=$2 AND state='running' AND lease_until>$3", + ) + .bind(&job.id.0) + .bind(&job.token.0) + .bind(now) + .fetch_optional(&mut *tx) + .await?; + if valid.is_none() { + return Ok(ClaimOutcome::LostClaim); + } + lifecycle::require_live_store(&mut tx, &job.store_id, now).await?; + let source: Option = sqlx::query_scalar( + "SELECT id FROM file_search_files WHERE id=$1 AND (expires_at IS NULL OR expires_at>$2)", + ) + .bind(&job.options.file_id) + .bind(now) + .fetch_optional(&mut *tx) + .await?; + if source.is_none() { + return Ok(ClaimOutcome::LostClaim); + } + terminal( + &mut tx, + &job.id.0, + &job.generation.0, + &object, + object.status.as_str(), + now, + ) + .await?; + tx.commit().await?; + Ok(ClaimOutcome::Applied) + } +} + +async fn live(tx: &mut DbTransaction<'_>, job: &ClaimedFileJob, now: i64) -> Result { + let found:Option=sqlx::query_scalar("SELECT a.file_id FROM file_search_attachments a JOIN file_search_files f ON f.id=a.file_id JOIN file_search_stores s ON s.id=a.store_id JOIN file_search_jobs j ON j.id=$4 JOIN file_search_batches b ON b.id=j.batch_id WHERE a.store_id=$1 AND a.file_id=$2 AND a.generation=$3 AND a.status='in_progress' AND b.cancelled=0 AND s.lifecycle_status!='expired' AND (s.expires_at IS NULL OR s.expires_at>$5) AND (f.expires_at IS NULL OR f.expires_at>$5)").bind(&job.store_id).bind(&job.options.file_id).bind(&job.generation.0).bind(&job.id.0).bind(now).fetch_optional(&mut **tx).await?; + Ok(found.is_some()) +} +async fn terminal( + tx: &mut DbTransaction<'_>, + id: &str, + generation: &str, + object: &VectorStoreFileObject, + state: &str, + now: i64, +) -> Result<(), FileSearchError> { + let data = serde_json::to_string(object)?; + sqlx::query( + "UPDATE file_search_jobs SET state=$2,snapshot=$3,updated_at=$4,claim_token=NULL,lease_until=NULL WHERE id=$1", + ) + .bind(id) + .bind(state) + .bind(&data) + .bind(now) + .execute(&mut **tx) + .await?; + sqlx::query( + "UPDATE file_search_attachments SET data=$3,status=$4 WHERE store_id=$1 AND file_id=$2 AND generation=$5", + ) + .bind(&object.vector_store_id) + .bind(&object.id) + .bind(data) + .bind(object.status.as_str()) + .bind(generation) + .execute(&mut **tx) + .await?; + Ok(()) +} + +/// Parent lifecycle writes already own the source or store lock before invalidating unfinished members. +pub(super) async fn invalidate( + tx: &mut DbTransaction<'_>, + store: Option<&str>, + file: Option<&str>, +) -> Result<(), FileSearchError> { + let rows:Vec<(String,String,String)>=sqlx::query_as("SELECT id,generation,snapshot FROM file_search_jobs WHERE ($1='' OR store_id=$1) AND ($2='' OR file_id=$2) AND state IN ('queued','running') ORDER BY id").bind(store.unwrap_or("")).bind(file.unwrap_or("")).fetch_all(&mut **tx).await?; + for (id, generation, data) in rows { + let mut object: VectorStoreFileObject = serde_json::from_str(&data)?; + object.status = AttachmentStatus::Cancelled; + let now = database_now(&mut *tx).await?; + terminal(tx, &id, &generation, &object, "cancelled", now).await?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::file_search::{FileSearchConfig, VectorStoreFileErrorCode}; + use crate::{storage::create_pool_with_schema, tool::file_search::FileSearchService}; + use serde_json::json; + use std::sync::Arc; + + #[allow( + clippy::too_many_lines, + reason = "keeps competing claims and old-generation publication assertions together" + )] + async fn fencing(database: &str) { + let pool = create_pool_with_schema(Some(database)).await.unwrap(); + let dir = tempfile::tempdir().unwrap(); + let service = FileSearchService::new( + pool.clone(), + Arc::new(reqwest::Client::new()), + FileSearchConfig { + files_storage_dir: Some(dir.path().into()), + ..Default::default() + }, + ) + .unwrap(); + let storage = FileSearchStorage::new(pool.clone()); + let store = service + .create_vector_store(serde_json::from_value(json!({})).unwrap()) + .await + .unwrap(); + let file = service + .upload_file("fence.txt", "text/plain", "assistants", b"fenced text".to_vec()) + .await + .unwrap(); + let request = || serde_json::from_value(json!({"file_ids":[file.id]})).unwrap(); + let batch = service.create_file_batch(&store.id, request()).await.unwrap(); + let (one, two) = tokio::join!(storage.claim_next(30, 3, 1000), storage.claim_next(30, 3, 1000)); + let mut claims = [one.unwrap(), two.unwrap()].into_iter().flatten(); + let stale = claims.next().unwrap(); + assert!(claims.next().is_none(), "only one concurrent claimant wins"); + sqlx::query("UPDATE file_search_jobs SET lease_until=1 WHERE id=$1") + .bind(&stale.id.0) + .execute(pool.as_ref()) + .await + .unwrap(); + let current = storage.claim_next(30, 3, 1000).await.unwrap().unwrap(); + assert_ne!(stale.token.0, current.token.0); + assert_eq!(storage.renew_claim(&stale, 30).await.unwrap(), ClaimOutcome::LostClaim); + let mut object = service.get_vector_store_file(&store.id, &file.id).await.unwrap(); + object.status = AttachmentStatus::Completed; + let prepared = || PreparedAttachment { + object: object.clone(), + chunks: vec![super::super::StoredChunk { + file_id: file.id.clone(), + filename: "fence.txt".into(), + chunk_index: 0, + text: "stale".into(), + embedding_text: None, + embedding: None, + attributes: object.attributes.clone(), + }], + dimensions: 0, + parsed_content: "stale".into(), + }; + assert_eq!( + storage.finish_claim(&stale, Ok(prepared())).await.unwrap(), + ClaimOutcome::LostClaim + ); + assert_eq!( + storage.finish_claim(&stale, Err(failure())).await.unwrap(), + ClaimOutcome::LostClaim + ); + storage.release_claim(&stale).await.unwrap(); + assert_eq!(storage.renew_claim(¤t, 30).await.unwrap(), ClaimOutcome::Applied); + service.cancel_file_batch(&store.id, &batch.id).await.unwrap(); + assert_eq!( + storage.finish_claim(¤t, Err(failure())).await.unwrap(), + ClaimOutcome::LostClaim + ); + assert_eq!( + service + .get_file_batch(&store.id, &batch.id) + .await + .unwrap() + .file_counts + .cancelled, + 1 + ); + service.detach_file(&store.id, &file.id).await.unwrap(); + let second = service.create_file_batch(&store.id, request()).await.unwrap(); + let old = storage.claim_next(30, 3, 1000).await.unwrap().unwrap(); + service.detach_file(&store.id, &file.id).await.unwrap(); + service + .attach_file( + &store.id, + AttachFileRequest { + file_id: file.id.clone(), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!( + storage.finish_claim(&old, Ok(prepared())).await.unwrap(), + ClaimOutcome::LostClaim + ); + assert_eq!( + storage.finish_claim(&old, Err(failure())).await.unwrap(), + ClaimOutcome::LostClaim + ); + storage.release_claim(&old).await.unwrap(); + assert!(storage.claim_next(30, 3, 1000).await.unwrap().is_none()); + assert_eq!( + service.get_vector_store_file(&store.id, &file.id).await.unwrap().status, + AttachmentStatus::Completed + ); + assert_eq!( + service + .get_file_batch(&store.id, &second.id) + .await + .unwrap() + .file_counts + .cancelled, + 1 + ); + let reused = service.create_file_batch(&store.id, request()).await.unwrap(); + assert_eq!(reused.file_counts.completed, 1); + service.delete_vector_store(&store.id).await.unwrap(); + service.delete_file(&file.id).await.unwrap(); + } + fn failure() -> VectorStoreFileError { + VectorStoreFileError { + code: VectorStoreFileErrorCode::ServerError, + message: "test failure".into(), + } + } + #[tokio::test] + async fn sqlite_claim_fencing() { + fencing("sqlite::memory:").await; + } + #[tokio::test] + #[ignore = "requires TEST_POSTGRES_URL"] + async fn postgres_claim_fencing() { + fencing(&std::env::var("TEST_POSTGRES_URL").unwrap()).await; + } + #[allow( + clippy::too_many_lines, + reason = "verifies identical parent invalidation invariants on both SQL backends" + )] + async fn parents(database: &str) { + let pool = create_pool_with_schema(Some(database)).await.unwrap(); + let dir = tempfile::tempdir().unwrap(); + let service = FileSearchService::new( + pool.clone(), + Arc::new(reqwest::Client::new()), + FileSearchConfig { + files_storage_dir: Some(dir.path().into()), + ..Default::default() + }, + ) + .unwrap(); + let storage = FileSearchStorage::new(pool.clone()); + for scenario in ["file_expired", "store_expired", "file_deleted", "store_deleted"] { + let store = service + .create_vector_store(serde_json::from_value(json!({})).unwrap()) + .await + .unwrap(); + let file = service + .upload_file("parent.txt", "text/plain", "assistants", b"parent fencing".to_vec()) + .await + .unwrap(); + let batch = service + .create_file_batch( + &store.id, + serde_json::from_value(json!({"file_ids":[file.id]})).unwrap(), + ) + .await + .unwrap(); + let claim = storage.claim_next(30, 3, 1000).await.unwrap().unwrap(); + let mut object = service.get_vector_store_file(&store.id, &file.id).await.unwrap(); + object.status = AttachmentStatus::Completed; + let prepared = PreparedAttachment { + object, + chunks: Vec::new(), + dimensions: 0, + parsed_content: "must not publish".into(), + }; + match scenario { + "file_expired" => { + sqlx::query("UPDATE file_search_files SET expires_at=1 WHERE id=$1") + .bind(&file.id) + .execute(pool.as_ref()) + .await + .unwrap(); + } + "store_expired" => { + sqlx::query("UPDATE file_search_stores SET expires_at=1 WHERE id=$1") + .bind(&store.id) + .execute(pool.as_ref()) + .await + .unwrap(); + } + "file_deleted" => { + service.delete_file(&file.id).await.unwrap(); + } + "store_deleted" => { + service.delete_vector_store(&store.id).await.unwrap(); + } + _ => unreachable!(), + } + assert_eq!( + storage.finish_claim(&claim, Ok(prepared)).await.unwrap(), + ClaimOutcome::LostClaim, + "{scenario}" + ); + assert_eq!( + storage.renew_claim(&claim, 30).await.unwrap(), + ClaimOutcome::LostClaim, + "{scenario}" + ); + service.cleanup_expired_files(1000).await.unwrap(); + service.cleanup_expired_vector_stores(1000).await.unwrap(); + let chunks: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM file_search_chunks WHERE store_id=$1") + .bind(&store.id) + .fetch_one(pool.as_ref()) + .await + .unwrap(); + assert_eq!(chunks, 0); + if scenario.starts_with("file_") { + assert_eq!( + service + .get_file_batch(&store.id, &batch.id) + .await + .unwrap() + .file_counts + .cancelled, + 1 + ); + assert_eq!( + service + .list_file_batch_files(&store.id, &batch.id, &ListParams::default()) + .await + .unwrap() + .data + .len(), + 1 + ); + service.delete_vector_store(&store.id).await.unwrap(); + } else { + assert!( + service.get_file(&file.id).await.is_ok(), + "store deletion preserves uploads" + ); + service.delete_file(&file.id).await.unwrap(); + if scenario == "store_expired" { + service.delete_vector_store(&store.id).await.unwrap(); + } + } + } + } + #[tokio::test] + async fn sqlite_parent_fencing() { + parents("sqlite::memory:").await; + } + #[tokio::test] + #[ignore = "requires TEST_POSTGRES_URL"] + async fn postgres_parent_fencing() { + parents(&std::env::var("TEST_POSTGRES_URL").unwrap()).await; + } + #[tokio::test] + #[ignore = "requires TEST_POSTGRES_URL"] + async fn postgres_job_lock_wait_rechecks_lease_before_publication() { + let pool = create_pool_with_schema(Some(&std::env::var("TEST_POSTGRES_URL").unwrap())) + .await + .unwrap(); + let dir = tempfile::tempdir().unwrap(); + let service = FileSearchService::new( + pool.clone(), + Arc::new(reqwest::Client::new()), + FileSearchConfig { + files_storage_dir: Some(dir.path().into()), + ..Default::default() + }, + ) + .unwrap(); + let storage = FileSearchStorage::new(pool.clone()); + let store = service + .create_vector_store(serde_json::from_value(json!({})).unwrap()) + .await + .unwrap(); + let file = service + .upload_file("contended.txt", "text/plain", "assistants", b"no stale lease".to_vec()) + .await + .unwrap(); + service + .create_file_batch( + &store.id, + serde_json::from_value(json!({"file_ids":[file.id]})).unwrap(), + ) + .await + .unwrap(); + let job = storage.claim_next(30, 3, 1000).await.unwrap().unwrap(); + let mut object = service.get_vector_store_file(&store.id, &file.id).await.unwrap(); + object.status = AttachmentStatus::Completed; + let prepared = PreparedAttachment { + object, + chunks: Vec::new(), + dimensions: 0, + parsed_content: "stale lease".into(), + }; + let mut held = pool.begin().await.unwrap(); + let now = database_now(&mut held).await.unwrap(); + sqlx::query("UPDATE file_search_jobs SET lease_until=$2 WHERE id=$1") + .bind(&job.id.0) + .bind(now + 1) + .execute(&mut *held) + .await + .unwrap(); + let waiting = tokio::spawn(async move { storage.finish_claim(&job, Ok(prepared)).await }); + tokio::time::sleep(std::time::Duration::from_millis(1200)).await; + assert!(!waiting.is_finished(), "publication must wait on the job row"); + held.commit().await.unwrap(); + assert_eq!(waiting.await.unwrap().unwrap(), ClaimOutcome::LostClaim); + assert_eq!( + service.get_vector_store_file(&store.id, &file.id).await.unwrap().status, + AttachmentStatus::InProgress + ); + service.delete_vector_store(&store.id).await.unwrap(); + service.delete_file(&file.id).await.unwrap(); + } + #[tokio::test] + #[ignore = "requires TEST_POSTGRES_URL"] + async fn postgres_store_delete_waits_for_jobs_before_attachments() { + let pool = create_pool_with_schema(Some(&std::env::var("TEST_POSTGRES_URL").unwrap())) + .await + .unwrap(); + let dir = tempfile::tempdir().unwrap(); + let service = FileSearchService::new( + pool.clone(), + Arc::new(reqwest::Client::new()), + FileSearchConfig { + files_storage_dir: Some(dir.path().into()), + ..Default::default() + }, + ) + .unwrap(); + let store = service + .create_vector_store(serde_json::from_value(json!({})).unwrap()) + .await + .unwrap(); + let file = service + .upload_file("order.txt", "text/plain", "assistants", b"consistent deletion".to_vec()) + .await + .unwrap(); + let batch = service + .create_file_batch( + &store.id, + serde_json::from_value(json!({"file_ids":[file.id]})).unwrap(), + ) + .await + .unwrap(); + // Hold the source-deletion prefix: source -> job, before its attachment invalidation. + let mut source = pool.begin().await.unwrap(); + sqlx::query("UPDATE file_search_files SET id=id WHERE id=$1") + .bind(&file.id) + .execute(&mut *source) + .await + .unwrap(); + sqlx::query("UPDATE file_search_jobs SET id=id WHERE batch_id=$1") + .bind(&batch.id) + .execute(&mut *source) + .await + .unwrap(); + let deleting = service.clone(); + let id = store.id.clone(); + let deletion = tokio::spawn(async move { deleting.delete_vector_store(&id).await }); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + assert!(!deletion.is_finished()); + tokio::time::timeout( + std::time::Duration::from_secs(1), + sqlx::query("UPDATE file_search_attachments SET status=status WHERE store_id=$1 AND file_id=$2") + .bind(&store.id) + .bind(&file.id) + .execute(&mut *source), + ) + .await + .expect("store deletion must not own the attachment while waiting for the job") + .unwrap(); + source.commit().await.unwrap(); + deletion.await.unwrap().unwrap(); + assert!(service.get_file(&file.id).await.is_ok()); + service.delete_file(&file.id).await.unwrap(); + } +} diff --git a/crates/agentic-server-core/src/storage/vector_store_lifecycle.rs b/crates/agentic-server-core/src/storage/vector_store_lifecycle.rs index c48b9828..d45bdfe6 100644 --- a/crates/agentic-server-core/src/storage/vector_store_lifecycle.rs +++ b/crates/agentic-server-core/src/storage/vector_store_lifecycle.rs @@ -182,6 +182,7 @@ impl FileSearchStorage { let changed = sqlx::query("UPDATE file_search_stores SET lifecycle_status = 'expired' WHERE id = $1 AND lifecycle_status != 'expired' AND expires_at <= $2") .bind(&id).bind(now).execute(&mut *tx).await?.rows_affected(); if changed == 1 { + super::batches::invalidate(&mut tx, Some(&id), None).await?; sqlx::query("DELETE FROM file_search_attachments WHERE store_id = $1") .bind(&id) .execute(&mut *tx) @@ -217,6 +218,10 @@ impl FileSearchStorage { .bind(serde_json::to_string(&object)?) .execute(&mut *tx) .await?; + if object.status == AttachmentStatus::InProgress { + sqlx::query("UPDATE file_search_jobs SET snapshot=$3 WHERE store_id=$1 AND file_id=$2 AND state IN ('queued','running') AND generation=(SELECT generation FROM file_search_attachments WHERE store_id=$1 AND file_id=$2)") + .bind(store_id).bind(file_id).bind(serde_json::to_string(&object)?).execute(&mut *tx).await?; + } // Bounded by the store's existing serialized corpus budget; preserve every other chunk field. let rows: Vec<(i64, String)> = sqlx::query_as("SELECT chunk_index, data FROM file_search_chunks WHERE store_id = $1 AND file_id = $2") diff --git a/crates/agentic-server-core/src/tool/file_search/batches.rs b/crates/agentic-server-core/src/tool/file_search/batches.rs new file mode 100644 index 00000000..d80dcd96 --- /dev/null +++ b/crates/agentic-server-core/src/tool/file_search/batches.rs @@ -0,0 +1,693 @@ +//! Explicitly owned workers; service clones never create or cancel tasks. +use super::{FileSearchService, ingest, page, validate_list}; +use crate::{ + storage::file_search::batches::{BatchFileOptions, ClaimOutcome, ClaimedFileJob}, + types::file_search::{ + AttachFileRequest, AttachmentStatus, ChunkingStrategy, CreateFileBatchRequest, FileBatchObject, + FileSearchError, ListParams, ListResponse, StaticChunking, VectorStoreFileChunkingStrategy, + VectorStoreFileError, VectorStoreFileErrorCode, VectorStoreFileObject, invalid, validate_attributes, + }, +}; +use std::time::Duration; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; + +const LEASE_SECONDS: i64 = 30; +const HEARTBEAT: Duration = Duration::from_secs(1); +const POLL: Duration = Duration::from_millis(100); + +impl FileSearchService { + /// Atomically records batch membership and immediately visible pending attachments. + /// # Errors + /// Returns validation, missing parent, conflicting attachment, or storage errors. + pub async fn create_file_batch( + &self, + store_id: &str, + request: CreateFileBatchRequest, + ) -> Result { + let files = match (request.file_ids, request.files) { + (Some(ids), None) => ids + .into_iter() + .map(|file_id| AttachFileRequest { + file_id, + attributes: request.attributes.clone(), + chunking_strategy: request.chunking_strategy.clone(), + }) + .collect(), + (None, Some(files)) => files, + _ => return invalid("Supply exactly one of file_ids or files"), + }; + if !(1..=2000).contains(&files.len()) { + return invalid("File batches accept 1 to 2000 members"); + } + let mut ids = std::collections::HashSet::new(); + let mut members = Vec::with_capacity(files.len()); + for mut file in files { + if !ids.insert(file.file_id.clone()) { + return invalid("Duplicate file IDs are not allowed in a batch"); + } + validate_attributes(&file.attributes)?; + let strategy = file.chunking_strategy.as_ref().unwrap_or(&ChunkingStrategy::Auto); + let config = if matches!(strategy, ChunkingStrategy::Auto) { + StaticChunking { + max_chunk_size_tokens: self.config.file_ingestion_params.default_chunk_size_tokens, + chunk_overlap_tokens: self.config.file_ingestion_params.default_chunk_overlap_tokens, + } + } else { + ingest::chunking_config(strategy)? + }; + if matches!(strategy, ChunkingStrategy::Auto) { + file.chunking_strategy = Some(ChunkingStrategy::Static { config: config.clone() }); + } + let contextual_identity = self.contextual_identity(&file)?; + let object = VectorStoreFileObject { + id: file.file_id.clone(), + object: "vector_store.file".into(), + created_at: 0, + vector_store_id: store_id.into(), + status: AttachmentStatus::InProgress, + usage_bytes: 0, + attributes: file.attributes.clone(), + chunking_strategy: VectorStoreFileChunkingStrategy::Static { config }, + last_error: None, + }; + members.push(( + BatchFileOptions { + request: file, + contextual_identity, + }, + object, + )); + } + self.compatible(&self.storage.store(store_id).await?)?; + let identity = self.identity(); + retry(|| self.storage.create_batch(store_id, &identity, &members)).await + } + fn contextual_identity(&self, request: &AttachFileRequest) -> Result, FileSearchError> { + let Some(ChunkingStrategy::Contextual { contextual }) = &request.chunking_strategy else { + return Ok(None); + }; + let (provider, model) = self.config.resolve( + contextual.model_id.as_deref(), + self.config.contextual_retrieval_params.model.as_ref(), + )?; + Ok(Some(format!("{}\n{model}", provider.endpoint("chat/completions")?))) + } + /// Retrieves durable batch counts. + /// # Errors + /// Returns not-found or storage errors. + pub async fn get_file_batch(&self, store_id: &str, id: &str) -> Result { + self.storage.batch(store_id, id).await + } + /// Cancels only unfinished members; completed and failed history remains unchanged. + /// # Errors + /// Returns not-found or storage errors. + pub async fn cancel_file_batch(&self, store_id: &str, id: &str) -> Result { + retry(|| self.storage.cancel_batch(store_id, id)).await + } + /// Lists membership snapshots, including results whose attachment has since been removed. + /// # Errors + /// Returns validation, not-found, or storage errors. + pub async fn list_file_batch_files( + &self, + store_id: &str, + id: &str, + params: &ListParams, + ) -> Result, FileSearchError> { + validate_list(params)?; + Ok(page( + self.storage.batch_files(store_id, id, params).await?, + params, + |file| &file.id, + )) + } +} + +/// Non-clone runtime handle. Call consuming `shutdown` before releasing the server owner. +#[must_use = "the runtime must be explicitly shut down and joined"] +pub struct FileSearchRuntime { + stop: CancellationToken, + tasks: JoinSet<()>, +} +impl FileSearchRuntime { + /// Starts bounded workers after service and server initialization. Polls SQL for restart recovery. + pub fn start(service: FileSearchService) -> Self { + Self::start_with_shutdown(service, CancellationToken::new()) + } + /// Starts workers whose admission and model work stop when the server token is cancelled. + /// The runtime owner must still call `shutdown` to join them. + pub fn start_with_shutdown(service: FileSearchService, stop: CancellationToken) -> Self { + let mut tasks = JoinSet::new(); + // The shared four-operation semaphore also bounds synchronous ingestion and parsing. + for _ in 0..4 { + let service = service.clone(); + let stop = stop.clone(); + tasks.spawn(async move { + worker(service, stop).await; + }); + } + let cleanup_stop = stop.clone(); + tasks.spawn(async move { + let mut tick = tokio::time::interval(Duration::from_secs( + service.config.file_batch_params.cleanup_interval_seconds, + )); + loop { + tokio::select! {biased;()=cleanup_stop.cancelled()=>break,_=tick.tick()=>{}} + let limit = service.config.file_batch_params.file_batch_chunk_size; + if let Err(error) = service.cleanup_expired_files(limit).await { + tracing::warn!(%error,"expired file cleanup failed"); + tick.reset_after(HEARTBEAT); + } + if let Err(error) = service.cleanup_expired_vector_stores(limit).await { + tracing::warn!(%error,"expired store cleanup failed"); + tick.reset_after(HEARTBEAT); + } + } + }); + Self { stop, tasks } + } + /// Stops admission, cooperatively cancels preparation, releases owned claims, and joins all tasks. + /// Shutdown does not cancel the API batch. Committing transactions are allowed to finish. + /// # Errors + /// Reports worker panics after all other owned tasks have joined. + pub async fn shutdown(mut self) -> Result<(), FileSearchError> { + self.stop.cancel(); + let mut failure = None; + while let Some(result) = self.tasks.join_next().await { + if let Err(error) = result { + failure = Some(error); + } + } + tracing::info!("file search workers stopped"); + failure.map_or(Ok(()), |error| Err(error.into())) + } +} +impl Drop for FileSearchRuntime { + fn drop(&mut self) { + self.stop.cancel(); + } +} + +async fn worker(service: FileSearchService, stop: CancellationToken) { + loop { + if stop.is_cancelled() { + break; + } + let Ok(permit) = service.permit() else { + tokio::select! {()=stop.cancelled()=>break,()=tokio::time::sleep(POLL)=>{}} + continue; + }; + let params = &service.config.file_batch_params; + let claim = service + .storage + .claim_next( + LEASE_SECONDS, + params.max_concurrent_files_per_batch, + params.file_batch_chunk_size, + ) + .await; + match claim { + Ok(Some(job)) => { + run_job(&service, &job, permit, &stop).await; + } + Ok(None) => { + drop(permit); + tokio::select! {()=stop.cancelled()=>break,()=tokio::time::sleep(POLL)=>{}} + } + Err(error) => { + drop(permit); + tracing::warn!(%error,"batch claim failed; retrying durable queue"); + tokio::select! {()=stop.cancelled()=>break,()=tokio::time::sleep(POLL)=>{}} + } + } + } +} +async fn run_job( + service: &FileSearchService, + job: &ClaimedFileJob, + permit: std::sync::Arc, + stop: &CancellationToken, +) { + let cancellation = CancellationToken::new(); + let prepare_service = service.clone(); + let prepare_job = job.clone(); + let prepare_cancellation = cancellation.clone(); + let mut prepare = tokio::spawn(async move { + let service = &prepare_service; + let job = &prepare_job; + let cancellation = &prepare_cancellation; + if job.identity != service.identity() || job.contextual_identity != service.contextual_identity(&job.options)? { + return Err(FileSearchError::Conflict( + "Restore the batch ingestion model configuration before ingestion".into(), + )); + } + let store = service.storage.store(&job.store_id).await?; + let dimensions = usize::try_from(store.embedding_dimensions).ok().filter(|n| *n > 0); + service + .prepare_cancellable(&job.store_id, job.options.clone(), dimensions, permit, cancellation) + .await + }); + let mut heartbeat = tokio::time::interval(HEARTBEAT); + let result = loop { + tokio::select! {biased; + ()=stop.cancelled()=>break None, + // A renewal can take longer than the interval: observe ready work before + // another overdue tick, while retaining shutdown as the first priority. + result=&mut prepare=>break Some(result.unwrap_or_else(|error|Err(error.into()))), + _=heartbeat.tick()=>{ + match service.storage.renew_claim(job,LEASE_SECONDS).await { + Ok(ClaimOutcome::Applied)=>{}, + Ok(ClaimOutcome::LostClaim)=>break None, + Err(error)=>{tracing::warn!(%error,"claim renewal failed");break None;} + } + } + } + }; + if let Some(result) = result { + let result = result.map_err(|error| file_error(&error)); + if let Err(error) = service.storage.finish_claim(job, result).await { + tracing::warn!(%error,"batch publication failed"); + if transient(&error) { + if let Err(failure) = service.storage.release_claim(job).await { + tracing::warn!(%failure,"transient publication claim release failed"); + } + return; + } + // A failed/indeterminate COMMIT may have succeeded. Token fencing makes this safe. + if let Err(failure) = service.storage.finish_claim(job, Err(file_error(&error))).await { + tracing::warn!(%failure,"batch failure recording failed"); + } + } + } else { + cancellation.cancel(); + if let Err(error) = prepare.await { + tracing::warn!(%error,"cancelled preparation task failed"); + } + if let Err(error) = service.storage.release_claim(job).await { + tracing::warn!(%error,"claim release failed; lease will expire"); + } + } +} +fn file_error(error: &FileSearchError) -> VectorStoreFileError { + VectorStoreFileError { + code: match error { + FileSearchError::InvalidRequest(_) => VectorStoreFileErrorCode::InvalidFile, + FileSearchError::UnsupportedFile(_) => VectorStoreFileErrorCode::UnsupportedFile, + #[cfg(feature = "file-search-pdf")] + FileSearchError::PdfParse(_) => VectorStoreFileErrorCode::InvalidFile, + _ => VectorStoreFileErrorCode::ServerError, + }, + message: error.to_string(), + } +} + +fn transient(error: &FileSearchError) -> bool { + let FileSearchError::Storage(sqlx::Error::Database(error)) = error else { + return false; + }; + matches!( + error.code().as_deref(), + Some("5" | "6" | "261" | "517" | "40001" | "40P01" | "55P03") + ) +} +async fn retry(mut attempt: F) -> Result +where + F: FnMut() -> Fut, + Fut: std::future::Future>, +{ + for tries in 0..3 { + match attempt().await { + Err(error) if tries < 2 && transient(&error) => tokio::time::sleep(POLL).await, + result => return result, + } + } + unreachable!("the final attempt always returns") +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + #[tokio::test] + async fn cancelled_preparation_joins_parser_and_releases_all_capacity() { + let pool = crate::storage::create_pool_with_schema(Some("sqlite::memory:")) + .await + .unwrap(); + let directory = tempfile::tempdir().unwrap(); + let service = FileSearchService::new( + pool, + Arc::new(reqwest::Client::new()), + crate::types::file_search::FileSearchConfig { + files_storage_dir: Some(directory.path().into()), + ..Default::default() + }, + ) + .unwrap(); + let file = service + .upload_file( + "cancelled.txt", + "text/plain", + "assistants", + b"bounded parser ".repeat(100_000), + ) + .await + .unwrap(); + let store = service + .create_vector_store(serde_json::from_str("{}").unwrap()) + .await + .unwrap(); + let cancellation = CancellationToken::new(); + cancellation.cancel(); + let permit = service.permit().unwrap(); + assert_eq!(service.workers.available_permits(), 3); + let result = service + .prepare_cancellable( + &store.id, + AttachFileRequest { + file_id: file.id.clone(), + ..Default::default() + }, + None, + permit, + &cancellation, + ) + .await; + assert!(result.is_err()); + assert_eq!( + service.workers.available_permits(), + 4, + "parser handle and its capacity must finish before returning" + ); + service.delete_vector_store(&store.id).await.unwrap(); + service.delete_file(&file.id).await.unwrap(); + } + #[tokio::test] + async fn runtime_sweeps_expiration_and_replays_committed_blob_cleanup() { + let pool = crate::storage::create_pool_with_schema(Some("sqlite::memory:")) + .await + .unwrap(); + let directory = tempfile::tempdir().unwrap(); + let service = FileSearchService::new( + pool.clone(), + Arc::new(reqwest::Client::new()), + crate::types::file_search::FileSearchConfig { + files_storage_dir: Some(directory.path().into()), + ..Default::default() + }, + ) + .unwrap(); + let expired = service + .upload_file("expired.txt", "text/plain", "assistants", b"expired".to_vec()) + .await + .unwrap(); + let retained = service + .upload_file("retained.txt", "text/plain", "assistants", b"retained".to_vec()) + .await + .unwrap(); + let replay = service + .upload_file("replay.txt", "text/plain", "assistants", b"replay".to_vec()) + .await + .unwrap(); + let store = service + .create_vector_store(serde_json::from_value(serde_json::json!({"file_ids":[retained.id]})).unwrap()) + .await + .unwrap(); + sqlx::query("UPDATE file_search_files SET expires_at=1 WHERE id=$1") + .bind(&expired.id) + .execute(pool.as_ref()) + .await + .unwrap(); + sqlx::query("UPDATE file_search_stores SET expires_at=1 WHERE id=$1") + .bind(&store.id) + .execute(pool.as_ref()) + .await + .unwrap(); + // Simulate a prior process stopping after the durable delete transaction committed. + let mut tx = pool.begin().await.unwrap(); + sqlx::query("INSERT INTO file_search_blob_cleanup (file_id) VALUES ($1)") + .bind(&replay.id) + .execute(&mut *tx) + .await + .unwrap(); + sqlx::query("DELETE FROM file_search_files WHERE id=$1") + .bind(&replay.id) + .execute(&mut *tx) + .await + .unwrap(); + tx.commit().await.unwrap(); + assert!(directory.path().join(&replay.id).exists()); + let runtime = FileSearchRuntime::start(service.clone()); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM file_search_attachments WHERE store_id=$1") + .bind(&store.id) + .fetch_one(pool.as_ref()) + .await + .unwrap(); + if count == 0 + && !directory.path().join(&expired.id).exists() + && !directory.path().join(&replay.id).exists() + { + break; + } + tokio::time::sleep(POLL).await; + } + }) + .await + .unwrap(); + runtime.shutdown().await.unwrap(); + let pending: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM file_search_blob_cleanup") + .fetch_one(pool.as_ref()) + .await + .unwrap(); + assert_eq!(pending, 0); + assert!(service.get_file(&retained.id).await.is_ok()); + assert_eq!(service.workers.available_permits(), 4); + service.delete_vector_store(&store.id).await.unwrap(); + service.delete_file(&retained.id).await.unwrap(); + } + #[tokio::test] + #[ignore = "requires TEST_POSTGRES_URL"] + #[allow( + clippy::too_many_lines, + reason = "exercises both duplicate-creation outcomes across a real post-commit lock timeout" + )] + async fn postgres_committed_creation_survives_response_read_timeout() { + use crate::storage::file_search::batches::{BatchTestHooks, CommitBarrier}; + let pool = crate::storage::create_pool_with_schema_and_configs( + Some(&std::env::var("TEST_POSTGRES_URL").unwrap()), + crate::config::SqliteConfig::default(), + crate::config::PostgresConfig { + lock_timeout: Duration::from_millis(100), + ..Default::default() + }, + ) + .await + .unwrap(); + let directory = tempfile::tempdir().unwrap(); + for completed in [true, false] { + let mut service = FileSearchService::new( + pool.clone(), + Arc::new(reqwest::Client::new()), + crate::types::file_search::FileSearchConfig { + files_storage_dir: Some(directory.path().into()), + ..Default::default() + }, + ) + .unwrap(); + let store = service + .create_vector_store(serde_json::from_str("{}").unwrap()) + .await + .unwrap(); + let file = service + .upload_file("committed.txt", "text/plain", "assistants", b"durable member".to_vec()) + .await + .unwrap(); + if completed { + service + .attach_file( + &store.id, + AttachFileRequest { + file_id: file.id.clone(), + ..Default::default() + }, + ) + .await + .unwrap(); + } + let hooks = Arc::new(BatchTestHooks { + after_commit: Some(CommitBarrier::default()), + ..Default::default() + }); + service.storage.batch_test_hooks = Some(hooks.clone()); + let creating = service.clone(); + let store_id = store.id.clone(); + let file_id = file.id.clone(); + let creation = tokio::spawn(async move { + creating + .create_file_batch( + &store_id, + serde_json::from_value(serde_json::json!({"file_ids":[file_id]})).unwrap(), + ) + .await + }); + let barrier = hooks.after_commit.as_ref().unwrap(); + tokio::time::timeout(Duration::from_secs(5), barrier.reached.notified()) + .await + .unwrap(); + let committed_id: String = sqlx::query_scalar("SELECT id FROM file_search_batches WHERE store_id=$1") + .bind(&store.id) + .fetch_one(pool.as_ref()) + .await + .unwrap(); + let mut held = pool.begin().await.unwrap(); + sqlx::query("UPDATE file_search_stores SET id=id WHERE id=$1") + .bind(&store.id) + .execute(&mut *held) + .await + .unwrap(); + let read_error = service.get_file_batch(&store.id, &committed_id).await.unwrap_err(); + assert!( + matches!(&read_error,FileSearchError::Storage(sqlx::Error::Database(error)) if error.code().as_deref()==Some("55P03")) + ); + barrier.resume.notify_one(); + // The old follow-up read times out; its retry starts after we release the lock. + tokio::time::sleep(Duration::from_millis(150)).await; + held.commit().await.unwrap(); + let response = tokio::time::timeout(Duration::from_secs(5), creation) + .await + .unwrap() + .unwrap(); + let ids: Vec = sqlx::query_scalar("SELECT id FROM file_search_batches WHERE store_id=$1") + .bind(&store.id) + .fetch_all(pool.as_ref()) + .await + .unwrap(); + let retrieved = service.get_file_batch(&store.id, &committed_id).await.unwrap(); + service.delete_vector_store(&store.id).await.unwrap(); + service.delete_file(&file.id).await.unwrap(); + assert_eq!( + ids, + vec![committed_id.clone()], + "post-commit response failure must not repeat creation" + ); + let response = response.expect("creation must recover its committed ID without another database read"); + assert_eq!(response.id, committed_id); + assert_eq!(response.id, retrieved.id); + assert_eq!(response.file_counts.total, 1); + assert_eq!(response.file_counts.completed, i64::from(completed)); + assert_eq!(response.file_counts.in_progress, i64::from(!completed)); + } + } + + #[tokio::test] + #[allow( + clippy::too_many_lines, + reason = "keeps the owned model barrier and slow-renewal completion assertions in one fixture" + )] + async fn ready_preparation_is_not_starved_by_slow_renewals() { + use crate::storage::file_search::batches::BatchTestHooks; + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}/v1", listener.local_addr().unwrap()); + let entered = started.clone(); + let resume = release.clone(); + let app = axum::Router::new().route( + "/v1/embeddings", + axum::routing::post(move || { + let entered = entered.clone(); + let resume = resume.clone(); + async move { + entered.notify_one(); + resume.notified().await; + axum::Json(serde_json::json!({"model":"fixture","data":[{"index":0,"embedding":[1.0,0.0]}]})) + } + }), + ); + let http = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let pool = crate::storage::create_pool_with_schema(Some("sqlite::memory:")) + .await + .unwrap(); + let directory = tempfile::tempdir().unwrap(); + let mut service = FileSearchService::new( + pool, + Arc::new(reqwest::Client::new()), + crate::types::file_search::FileSearchConfig { + files_storage_dir: Some(directory.path().into()), + embedding_base_url: Some(endpoint), + embedding_model: Some("fixture".into()), + ..Default::default() + }, + ) + .unwrap(); + let store = service + .create_vector_store(serde_json::from_str("{}").unwrap()) + .await + .unwrap(); + let file = service + .upload_file( + "slow-renewal.txt", + "text/plain", + "assistants", + b"successful preparation".to_vec(), + ) + .await + .unwrap(); + let batch = service + .create_file_batch( + &store.id, + serde_json::from_value(serde_json::json!({"file_ids":[file.id]})).unwrap(), + ) + .await + .unwrap(); + let claim = service.storage.claim_next(30, 3, 10).await.unwrap().unwrap(); + let hooks = Arc::new(BatchTestHooks { + renewal_delay: Duration::from_millis(1100), + ..Default::default() + }); + service.storage.batch_test_hooks = Some(hooks.clone()); + let stop = CancellationToken::new(); + let owned_stop = stop.clone(); + let worker = service.clone(); + let permit = service.permit().unwrap(); + let mut task = tokio::spawn(async move { + run_job(&worker, &claim, permit, &owned_stop).await; + }); + tokio::time::timeout(Duration::from_secs(2), hooks.renewal_started.notified()) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(2), started.notified()) + .await + .unwrap(); + release.notify_one(); + tokio::time::timeout(Duration::from_millis(500), async { + while service.workers.available_permits() != 4 { + tokio::task::yield_now().await; + } + }) + .await + .expect("preparation must complete during the first slow renewal"); + let completed = tokio::time::timeout(Duration::from_millis(2500), &mut task).await; + stop.cancel(); + let completed = if let Ok(result) = completed { + result.unwrap(); + true + } else { + task.await.unwrap(); + false + }; + let object = service.get_file_batch(&store.id, &batch.id).await.unwrap(); + service.delete_vector_store(&store.id).await.unwrap(); + service.delete_file(&file.id).await.unwrap(); + http.abort(); + let _ = http.await; + assert!( + completed, + "ready preparation was starved behind overdue successful renewals" + ); + assert_eq!(object.file_counts.completed, 1); + } +} diff --git a/crates/agentic-server-core/src/tool/file_search/ingest.rs b/crates/agentic-server-core/src/tool/file_search/ingest.rs index 158418f7..cdffbb57 100644 --- a/crates/agentic-server-core/src/tool/file_search/ingest.rs +++ b/crates/agentic-server-core/src/tool/file_search/ingest.rs @@ -81,7 +81,9 @@ pub(super) fn validate_content_type(filename: &str, content_type: &str) -> Resul { return Ok(()); } - invalid("Unsupported file type; upload UTF-8 text or a PDF containing extractable text") + Err(FileSearchError::UnsupportedFile( + "Unsupported file type; upload UTF-8 text or a PDF containing extractable text".into(), + )) } pub(super) struct ExtractedDocument { diff --git a/crates/agentic-server-core/src/tool/file_search/mod.rs b/crates/agentic-server-core/src/tool/file_search/mod.rs index b2a32472..6b2ccb04 100644 --- a/crates/agentic-server-core/src/tool/file_search/mod.rs +++ b/crates/agentic-server-core/src/tool/file_search/mod.rs @@ -9,4 +9,4 @@ mod service; pub use crate::types::file_search::FileSearchError; pub use handler::{FileSearchExecutionParams, FileSearchExecutor, FileSearchHandler}; -pub use service::{FileDownload, FileSearchService, FileUpload, MAX_FILE_BYTES}; +pub use service::{FileDownload, FileSearchRuntime, FileSearchService, FileUpload, MAX_FILE_BYTES}; diff --git a/crates/agentic-server-core/src/tool/file_search/service.rs b/crates/agentic-server-core/src/tool/file_search/service.rs index 9148e0cd..8ce260b8 100644 --- a/crates/agentic-server-core/src/tool/file_search/service.rs +++ b/crates/agentic-server-core/src/tool/file_search/service.rs @@ -13,8 +13,11 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore}; #[path = "files.rs"] mod files; pub use files::{FileDownload, FileUpload}; +#[path = "batches.rs"] +mod batches; #[path = "stores.rs"] mod stores; +pub use batches::FileSearchRuntime; use stores::validate_store_fields; use super::{embeddings::Embeddings, ingest, models::Models, ranking}; @@ -383,6 +386,28 @@ impl FileSearchService { request: AttachFileRequest, dimensions: Option, permit: Arc, + ) -> Result { + self.prepare_cancellable( + store_id, + request, + dimensions, + permit, + &tokio_util::sync::CancellationToken::new(), + ) + .await + } + + #[allow( + clippy::too_many_lines, + reason = "keeps bounded extraction and model preparation in the single shared ingestion pipeline" + )] + async fn prepare_cancellable( + &self, + store_id: &str, + request: AttachFileRequest, + dimensions: Option, + permit: Arc, + cancellation: &tokio_util::sync::CancellationToken, ) -> Result { validate_attributes(&request.attributes)?; LocalFiles::validate_id(&request.file_id)?; @@ -416,25 +441,33 @@ impl FileSearchService { let worker_permit = permit.clone(); let cancelled = Arc::new(AtomicBool::new(false)); let _cancel_on_drop = CancelIngestionOnDrop(cancelled.clone()); - let document = tokio::task::spawn_blocking(move || { + let parser_cancelled = cancelled.clone(); + let mut parser = tokio::task::spawn_blocking(move || { let _permit = worker_permit; ingest::extract_and_chunk(bytes, &filename, &uploaded.content_type, &chunking, &cancelled) - }) - .await??; + }); + let document = tokio::select! { + result = &mut parser => result??, + () = cancellation.cancelled() => { + parser_cancelled.store(true, Ordering::Relaxed); + let _ = parser.await?; + return Err(FileSearchError::Unavailable("Ingestion stopped".into())); + } + }; let contextual = if let ChunkingStrategy::Contextual { contextual } = &strategy { - Some( - self.models - .contextualize(&document.text, &document.chunks, contextual) - .await?, - ) + Some(tokio::select! { + result = self.models.contextualize(&document.text, &document.chunks, contextual) => result?, + () = cancellation.cancelled() => return Err(FileSearchError::Unavailable("Ingestion stopped".into())), + }) } else { None }; let texts = document.chunks; let vectors = if let Some(embeddings) = &self.embeddings { - embeddings - .embed(contextual.as_deref().unwrap_or(&texts), dimensions) - .await? + tokio::select! { + result = embeddings.embed(contextual.as_deref().unwrap_or(&texts), dimensions) => result?, + () = cancellation.cancelled() => return Err(FileSearchError::Unavailable("Ingestion stopped".into())), + } } else { Vec::new() }; diff --git a/crates/agentic-server-core/src/types/file_search.rs b/crates/agentic-server-core/src/types/file_search.rs index 9d6b38a7..f7c5fba0 100644 --- a/crates/agentic-server-core/src/types/file_search.rs +++ b/crates/agentic-server-core/src/types/file_search.rs @@ -108,6 +108,8 @@ pub enum FileSearchError { #[error("{0}")] InvalidRequest(String), #[error("{0}")] + UnsupportedFile(String), + #[error("{0}")] NotFound(String), #[error("{0}")] Conflict(String), @@ -142,7 +144,7 @@ impl FileSearchError { #[must_use] pub const fn status_code(&self) -> u16 { match self { - Self::InvalidRequest(_) => 400, + Self::InvalidRequest(_) | Self::UnsupportedFile(_) => 400, #[cfg(feature = "file-search-pdf")] Self::PdfParse(_) => 400, Self::NotFound(_) => 404, @@ -157,6 +159,7 @@ impl FileSearchError { pub fn public_message(&self) -> String { match self { Self::InvalidRequest(message) + | Self::UnsupportedFile(message) | Self::NotFound(message) | Self::Conflict(message) | Self::Unavailable(message) => message.clone(), @@ -944,3 +947,37 @@ pub enum VectorStoreFileChunkingStrategy { #[serde(alias = "auto", alias = "contextual")] Other, } + +/// A batch supplies either shared options with IDs or independent per-file options. +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +pub struct CreateFileBatchRequest { + pub file_ids: Option>, + pub files: Option>, + #[serde(default, deserialize_with = "null_default")] + #[cfg_attr(feature = "openapi", schema(nullable = true))] + pub attributes: FileAttributes, + pub chunking_strategy: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +pub enum BatchStatus { + InProgress, + Completed, + Cancelled, + Failed, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +pub struct FileBatchObject { + pub id: String, + pub object: String, + pub created_at: i64, + pub vector_store_id: String, + pub status: BatchStatus, + pub file_counts: FileCounts, +} diff --git a/crates/agentic-server-core/tests/vector_store_batches.rs b/crates/agentic-server-core/tests/vector_store_batches.rs new file mode 100644 index 00000000..d5b319e5 --- /dev/null +++ b/crates/agentic-server-core/tests/vector_store_batches.rs @@ -0,0 +1,454 @@ +use agentic_core::{ + storage::create_pool_with_schema, + tool::file_search::{FileSearchRuntime, FileSearchService}, + types::file_search::*, +}; +use serde_json::json; +use std::{sync::Arc, time::Duration}; + +#[allow( + clippy::too_many_lines, + reason = "keeps atomic membership, partial failure, and pagination assertions together" +)] +async fn batches(database: &str) { + let pool = create_pool_with_schema(Some(database)).await.unwrap(); + let files = tempfile::tempdir().unwrap(); + let service = FileSearchService::new( + pool.clone(), + Arc::new(reqwest::Client::new()), + FileSearchConfig { + files_storage_dir: Some(files.path().into()), + ..Default::default() + }, + ) + .unwrap(); + let store = service + .create_vector_store(serde_json::from_value(json!({})).unwrap()) + .await + .unwrap(); + let good = service + .upload_file("good.txt", "text/plain", "assistants", b"lunar policy".to_vec()) + .await + .unwrap(); + let bad = service + .upload_file("bad.bin", "application/octet-stream", "assistants", vec![0, 255]) + .await + .unwrap(); + for invalid in [ + json!({}), + json!({"file_ids":[]}), + json!({"file_ids":[good.id],"files":[]}), + json!({"file_ids":[good.id,good.id]}), + json!({"file_ids":[good.id,"missing"]}), + ] { + assert!( + service + .create_file_batch(&store.id, serde_json::from_value(invalid).unwrap()) + .await + .is_err() + ); + assert_eq!(service.get_vector_store(&store.id).await.unwrap().file_counts.total, 0); + } + let batch = service + .create_file_batch( + &store.id, + serde_json::from_value( + json!({"files":[{"file_id":good.id,"attributes":{"team":"moon"}},{"file_id":bad.id}]}), + ) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(batch.object, "vector_store.files_batch"); + assert_eq!(batch.file_counts.in_progress, 2); + assert_eq!( + service.get_vector_store_file(&store.id, &good.id).await.unwrap().status, + AttachmentStatus::InProgress + ); + assert_eq!( + service + .create_file_batch( + &store.id, + serde_json::from_value(json!({"file_ids":[good.id]})).unwrap() + ) + .await + .unwrap_err() + .status_code(), + 409 + ); + let first = service + .list_file_batch_files( + &store.id, + &batch.id, + &ListParams { + limit: Some(1), + order: Some(ListOrder::Asc), + ..Default::default() + }, + ) + .await + .unwrap(); + assert!(first.has_more); + let next = service + .list_file_batch_files( + &store.id, + &batch.id, + &ListParams { + limit: Some(1), + order: Some(ListOrder::Asc), + after: Some(first.data[0].id.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert!(!next.has_more); + assert_ne!(first.data[0].id, next.data[0].id); + let previous = service + .list_file_batch_files( + &store.id, + &batch.id, + &ListParams { + limit: Some(1), + order: Some(ListOrder::Asc), + before: Some(next.data[0].id.clone()), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(previous.data[0].id, first.data[0].id); + let runtime = FileSearchRuntime::start(service.clone()); + let other = FileSearchRuntime::start(service.clone()); + let finished = tokio::time::timeout(Duration::from_secs(10), async { + loop { + let batch = service.get_file_batch(&store.id, &batch.id).await.unwrap(); + if batch.status != BatchStatus::InProgress { + break batch; + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + }) + .await + .unwrap(); + runtime.shutdown().await.unwrap(); + other.shutdown().await.unwrap(); + assert_eq!(finished.status, BatchStatus::Completed); + assert_eq!(finished.file_counts.completed, 1); + assert_eq!(finished.file_counts.failed, 1); + let failed = service.get_vector_store_file(&store.id, &bad.id).await.unwrap(); + assert_eq!( + failed.last_error.unwrap().code, + VectorStoreFileErrorCode::UnsupportedFile + ); + + let page = service + .list_file_batch_files( + &store.id, + &batch.id, + &ListParams { + filter: Some(AttachmentStatus::Completed), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(page.data.len(), 1); + assert_eq!(page.data[0].id, good.id); + let reused = service + .create_file_batch( + &store.id, + serde_json::from_value(json!({"files":[{"file_id":good.id,"attributes":{"team":"changed"}}]})).unwrap(), + ) + .await + .unwrap(); + assert_eq!(reused.status, BatchStatus::Completed); + assert_eq!( + service + .list_file_batch_files(&store.id, &reused.id, &ListParams::default()) + .await + .unwrap() + .data[0] + .attributes, + page.data[0].attributes + ); + assert_eq!( + service + .create_file_batch(&store.id, serde_json::from_value(json!({"file_ids":[bad.id]})).unwrap()) + .await + .unwrap_err() + .status_code(), + 409 + ); + service.detach_file(&store.id, &good.id).await.unwrap(); + assert_eq!( + service + .get_file_batch(&store.id, &batch.id) + .await + .unwrap() + .file_counts + .completed, + 1 + ); + service.delete_vector_store(&store.id).await.unwrap(); + service.delete_file(&good.id).await.unwrap(); + service.delete_file(&bad.id).await.unwrap(); +} +#[tokio::test] +async fn sqlite_batches() { + batches("sqlite::memory:").await; +} +#[tokio::test] +#[ignore = "requires TEST_POSTGRES_URL"] +async fn postgres_batches() { + batches(&std::env::var("TEST_POSTGRES_URL").unwrap()).await; +} + +#[derive(Clone, Default)] +struct Barrier { + blocked: Arc, + started: Arc, + resume: Arc, +} +async fn embedding( + axum::extract::State(barrier): axum::extract::State, + axum::Json(input): axum::Json, +) -> axum::Json { + if barrier.blocked.load(std::sync::atomic::Ordering::SeqCst) { + barrier.started.notify_one(); + barrier.resume.notified().await; + } + axum::Json( + json!({"model":input["model"],"data":input["input"].as_array().unwrap().iter().enumerate().map(|(index,_)|json!({"index":index,"embedding":[1.0,0.0]})).collect::>()}), + ) +} +#[allow( + clippy::too_many_lines, + reason = "keeps the blocked model cancellation and restart sequence in one fixture" +)] +async fn blocked_lifecycle(database: &str) { + let barrier = Barrier::default(); + barrier.blocked.store(true, std::sync::atomic::Ordering::SeqCst); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/v1", listener.local_addr().unwrap()); + let app = axum::Router::new() + .route("/v1/embeddings", axum::routing::post(embedding)) + .with_state(barrier.clone()); + let http = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let pool = create_pool_with_schema(Some(database)).await.unwrap(); + let directory = tempfile::tempdir().unwrap(); + let service = FileSearchService::new( + pool.clone(), + Arc::new(reqwest::Client::new()), + FileSearchConfig { + files_storage_dir: Some(directory.path().into()), + embedding_base_url: Some(url), + embedding_model: Some("fixture".into()), + ..Default::default() + }, + ) + .unwrap(); + let store = service + .create_vector_store(serde_json::from_value(json!({})).unwrap()) + .await + .unwrap(); + let file = service + .upload_file("blocked.txt", "text/plain", "assistants", b"blocked embedding".to_vec()) + .await + .unwrap(); + let request = || serde_json::from_value(json!({"file_ids":[file.id]})).unwrap(); + let batch = service.create_file_batch(&store.id, request()).await.unwrap(); + let runtime = FileSearchRuntime::start(service.clone()); + tokio::time::timeout(Duration::from_secs(5), barrier.started.notified()) + .await + .unwrap(); + // A second service instance issues cancellation while the first owns the model request. + let cancelled = service.clone().cancel_file_batch(&store.id, &batch.id).await.unwrap(); + assert_eq!(cancelled.file_counts.cancelled, 1); + tokio::time::timeout(Duration::from_secs(5), runtime.shutdown()) + .await + .unwrap() + .unwrap(); + let chunks: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM file_search_chunks WHERE store_id=$1") + .bind(&store.id) + .fetch_one(pool.as_ref()) + .await + .unwrap(); + assert_eq!(chunks, 0); + service.detach_file(&store.id, &file.id).await.unwrap(); + let resumable = service.create_file_batch(&store.id, request()).await.unwrap(); + let runtime = FileSearchRuntime::start(service.clone()); + tokio::time::timeout(Duration::from_secs(5), barrier.started.notified()) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), runtime.shutdown()) + .await + .unwrap() + .unwrap(); + let states: Vec = sqlx::query_scalar("SELECT state FROM file_search_jobs WHERE batch_id=$1") + .bind(&resumable.id) + .fetch_all(pool.as_ref()) + .await + .unwrap(); + assert_eq!(states, ["queued"]); + assert_eq!( + service.get_file_batch(&store.id, &resumable.id).await.unwrap().status, + BatchStatus::InProgress + ); + service + .update_vector_store_file( + &store.id, + &file.id, + serde_json::from_value(json!({"attributes":{"team":"updated"}})).unwrap(), + ) + .await + .unwrap(); + barrier.blocked.store(false, std::sync::atomic::Ordering::SeqCst); + barrier.resume.notify_waiters(); + let restarted = FileSearchRuntime::start(service.clone()); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if service + .get_file_batch(&store.id, &resumable.id) + .await + .unwrap() + .file_counts + .completed + == 1 + { + break; + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + }) + .await + .unwrap(); + restarted.shutdown().await.unwrap(); + assert_eq!( + service + .get_vector_store_file(&store.id, &file.id) + .await + .unwrap() + .attributes, + serde_json::from_value(json!({"team":"updated"})).unwrap() + ); + let result = service + .search( + std::slice::from_ref(&store.id), + &serde_json::from_value(json!({"query":"blocked","filters":{"type":"eq","key":"team","value":"updated"}})) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(result.data.len(), 1); + + let chunks: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM file_search_chunks WHERE store_id=$1") + .bind(&store.id) + .fetch_one(pool.as_ref()) + .await + .unwrap(); + assert_eq!(chunks, 1); + // Fresh starts preserve completed work and never reset another runtime's valid claim. + let again = FileSearchRuntime::start(service.clone()); + again.shutdown().await.unwrap(); + assert_eq!( + service + .get_file_batch(&store.id, &resumable.id) + .await + .unwrap() + .file_counts + .completed, + 1 + ); + service.delete_vector_store(&store.id).await.unwrap(); + service.delete_file(&file.id).await.unwrap(); + http.abort(); + let _ = http.await; +} +#[tokio::test] +async fn sqlite_blocked_cancel_shutdown_restart() { + blocked_lifecycle("sqlite::memory:").await; +} +#[tokio::test] +#[ignore = "requires TEST_POSTGRES_URL"] +async fn postgres_blocked_cancel_shutdown_restart() { + blocked_lifecycle(&std::env::var("TEST_POSTGRES_URL").unwrap()).await; +} + +async fn changed_contextual_default(database: &str) { + let pool = create_pool_with_schema(Some(database)).await.unwrap(); + let directory = tempfile::tempdir().unwrap(); + let mut config: FileSearchConfig = serde_json::from_value(json!({"vector_stores":{ + "providers":{"local":{"base_url":"http://127.0.0.1:9/v1","models":["embed","old","new"]}}, + "default_embedding_model":{"provider_id":"local","model_id":"embed","embedding_dimensions":2}, + "contextual_retrieval_params":{"model":{"provider_id":"local","model_id":"old"}} + }})) + .unwrap(); + config.files_storage_dir = Some(directory.path().into()); + let service = FileSearchService::new(pool.clone(), Arc::new(reqwest::Client::new()), config.clone()).unwrap(); + let store = service + .create_vector_store(serde_json::from_value(json!({})).unwrap()) + .await + .unwrap(); + let file = service + .upload_file( + "context.txt", + "text/plain", + "assistants", + b"contextual recovery".to_vec(), + ) + .await + .unwrap(); + let batch = service + .create_file_batch( + &store.id, + serde_json::from_value( + json!({"files":[{"file_id":file.id,"chunking_strategy":{"type":"contextual","contextual":{}}}]}), + ) + .unwrap(), + ) + .await + .unwrap(); + config + .vector_stores + .contextual_retrieval_params + .model + .as_mut() + .unwrap() + .model_id = "new".into(); + let changed = FileSearchService::new(pool.clone(), Arc::new(reqwest::Client::new()), config).unwrap(); + let runtime = FileSearchRuntime::start(changed); + tokio::time::timeout(Duration::from_secs(5), async { + while service + .get_file_batch(&store.id, &batch.id) + .await + .unwrap() + .file_counts + .failed + == 0 + { + tokio::time::sleep(Duration::from_millis(25)).await; + } + }) + .await + .unwrap(); + runtime.shutdown().await.unwrap(); + let failed = service.get_vector_store_file(&store.id, &file.id).await.unwrap(); + assert_eq!( + failed.last_error.unwrap().message, + "Restore the batch ingestion model configuration before ingestion" + ); + service.delete_vector_store(&store.id).await.unwrap(); + service.delete_file(&file.id).await.unwrap(); +} +#[tokio::test] +async fn sqlite_changed_contextual_default_is_rejected() { + changed_contextual_default("sqlite::memory:").await; +} +#[tokio::test] +#[ignore = "requires TEST_POSTGRES_URL"] +async fn postgres_changed_contextual_default_is_rejected() { + changed_contextual_default(&std::env::var("TEST_POSTGRES_URL").unwrap()).await; +} diff --git a/crates/agentic-server/src/handler/http/file_search.rs b/crates/agentic-server/src/handler/http/file_search.rs index 3229b7b0..c76829df 100644 --- a/crates/agentic-server/src/handler/http/file_search.rs +++ b/crates/agentic-server/src/handler/http/file_search.rs @@ -4,8 +4,8 @@ use agentic_core::executor::ExecutorError; use agentic_core::tool::ToolError; use agentic_core::tool::file_search::{FileSearchService, MAX_FILE_BYTES}; use agentic_core::types::file_search::{ - AttachFileRequest, CreateVectorStoreRequest, FileExpirationAnchor, FileExpiresAfter, FileSearchError, ListParams, - SearchRequest, UpdateVectorStoreFileRequest, UpdateVectorStoreRequest, + AttachFileRequest, CreateFileBatchRequest, CreateVectorStoreRequest, FileExpirationAnchor, FileExpiresAfter, + FileSearchError, ListParams, SearchRequest, UpdateVectorStoreFileRequest, UpdateVectorStoreRequest, }; #[path = "multipart_limits.rs"] mod multipart_limits; @@ -52,6 +52,19 @@ pub(crate) fn router() -> Router { "/v1/vector_stores/{store_id}/files/{file_id}/content", get(vector_store_file_content), ) + .route("/v1/vector_stores/{store_id}/file_batches", post(create_file_batch)) + .route( + "/v1/vector_stores/{store_id}/file_batches/{batch_id}", + get(get_file_batch), + ) + .route( + "/v1/vector_stores/{store_id}/file_batches/{batch_id}/cancel", + post(cancel_file_batch), + ) + .route( + "/v1/vector_stores/{store_id}/file_batches/{batch_id}/files", + get(list_file_batch_files), + ) .route("/v1/vector_stores/{store_id}/search", post(search)) } @@ -523,3 +536,82 @@ pub(crate) async fn vector_store_file_content( }; result(search.vector_store_file_content(&store_id, &file_id).await) } + +#[cfg_attr(feature = "openapi", utoipa::path( + post, path = "/v1/vector_stores/{store_id}/file_batches", + params(("store_id" = String, Path)), + request_body = agentic_core::types::file_search::CreateFileBatchRequest, + responses((status = 200, description = "Success", body = agentic_core::types::file_search::FileBatchObject), (status = 400, description = "Invalid request", body = crate::openapi::ApiErrorResponse), (status = 404, description = "Object not found", body = crate::openapi::ApiErrorResponse)), + security(("bearer_auth" = [])), tag = "file_search", +))] +pub(crate) async fn create_file_batch( + State(state): State, + Path(store_id): Path, + request: Result, JsonRejection>, +) -> Response { + let search = match service(&state) { + Ok(service) => service, + Err(error) => return *error, + }; + let request = match body(request) { + Ok(request) => request, + Err(error) => return *error, + }; + result(search.create_file_batch(&store_id, request).await) +} + +#[cfg_attr(feature = "openapi", utoipa::path( + get, path = "/v1/vector_stores/{store_id}/file_batches/{batch_id}", + params(("store_id" = String, Path), ("batch_id" = String, Path)), + responses((status = 200, description = "Success", body = agentic_core::types::file_search::FileBatchObject), (status = 400, description = "Invalid request", body = crate::openapi::ApiErrorResponse), (status = 404, description = "Object not found", body = crate::openapi::ApiErrorResponse)), + security(("bearer_auth" = [])), tag = "file_search", +))] +pub(crate) async fn get_file_batch( + State(state): State, + Path((store_id, batch_id)): Path<(String, String)>, +) -> Response { + let search = match service(&state) { + Ok(service) => service, + Err(error) => return *error, + }; + result(search.get_file_batch(&store_id, &batch_id).await) +} + +#[cfg_attr(feature = "openapi", utoipa::path( + post, path = "/v1/vector_stores/{store_id}/file_batches/{batch_id}/cancel", + params(("store_id" = String, Path), ("batch_id" = String, Path)), + responses((status = 200, description = "Success", body = agentic_core::types::file_search::FileBatchObject), (status = 400, description = "Invalid request", body = crate::openapi::ApiErrorResponse), (status = 404, description = "Object not found", body = crate::openapi::ApiErrorResponse)), + security(("bearer_auth" = [])), tag = "file_search", +))] +pub(crate) async fn cancel_file_batch( + State(state): State, + Path((store_id, batch_id)): Path<(String, String)>, +) -> Response { + let search = match service(&state) { + Ok(service) => service, + Err(error) => return *error, + }; + result(search.cancel_file_batch(&store_id, &batch_id).await) +} + +#[cfg_attr(feature = "openapi", utoipa::path( + get, path = "/v1/vector_stores/{store_id}/file_batches/{batch_id}/files", + params(("store_id" = String, Path), ("batch_id" = String, Path), ("filter" = Option, Query), ("limit" = Option, Query, description="Page size, 1 to 100"), ("order" = Option, Query), ("before" = Option, Query), ("after" = Option, Query)), + responses((status = 200, description = "Success", body = agentic_core::types::file_search::ListResponse), (status = 400, description = "Invalid request", body = crate::openapi::ApiErrorResponse), (status = 404, description = "Object not found", body = crate::openapi::ApiErrorResponse)), + security(("bearer_auth" = [])), tag = "file_search", +))] +pub(crate) async fn list_file_batch_files( + State(state): State, + Path((store_id, batch_id)): Path<(String, String)>, + params: Result, QueryRejection>, +) -> Response { + let search = match service(&state) { + Ok(service) => service, + Err(error) => return *error, + }; + let params = match query(params) { + Ok(params) => params, + Err(error) => return *error, + }; + result(search.list_file_batch_files(&store_id, &batch_id, ¶ms).await) +} diff --git a/crates/agentic-server/src/openapi.rs b/crates/agentic-server/src/openapi.rs index 6e06f289..ae083e79 100644 --- a/crates/agentic-server/src/openapi.rs +++ b/crates/agentic-server/src/openapi.rs @@ -28,6 +28,11 @@ use utoipa::OpenApi; crate::handler::http::file_search::get_vector_store_file, crate::handler::http::file_search::detach_file, crate::handler::http::file_search::search, + crate::handler::http::file_search::create_file_batch, + crate::handler::http::file_search::get_file_batch, + crate::handler::http::file_search::cancel_file_batch, + crate::handler::http::file_search::list_file_batch_files, + crate::handler::http::models::health, crate::handler::http::models::ready, crate::handler::http::models::models, diff --git a/crates/agentic-server/src/server.rs b/crates/agentic-server/src/server.rs index 32d6b07b..3d363c91 100644 --- a/crates/agentic-server/src/server.rs +++ b/crates/agentic-server/src/server.rs @@ -8,6 +8,7 @@ use agentic_core::error::Error as CoreError; use agentic_core::executor::ExecutionContext; use agentic_core::proxy::ProxyState; use agentic_core::readiness::{llm_readiness_client, wait_llm_ready}; +use agentic_core::tool::file_search::{FileSearchError, FileSearchRuntime}; use agentic_server::app::{AppState, ReadinessTracker, ServerConfig, WebSocketTracker, build_router_with_auth}; use agentic_server::auth::{OidcAuthError, OidcAuthenticator, OidcConfig}; use tokio::net::TcpListener; @@ -21,6 +22,8 @@ pub enum ServerError { #[error(transparent)] Core(#[from] CoreError), #[error(transparent)] + FileSearch(#[from] FileSearchError), + #[error(transparent)] Io(#[from] std::io::Error), #[error("failed to initialize OIDC authentication: {0}")] Oidc(#[source] OidcAuthError), @@ -151,7 +154,26 @@ pub async fn run(config: Config, host: &str, port: u16, oidc_config: Option, + result: Result<(), ServerError>, +) -> Result<(), ServerError> { + let shutdown = match runtime { + Some(runtime) => runtime.shutdown().await, + None => Ok(()), + }; + result?; + shutdown?; + Ok(()) } /// Spawn vLLM as a subprocess and run the gateway in the foreground. @@ -201,22 +223,31 @@ pub async fn run_with_llm( } => state?, }; + let runtime = state + .exec_ctx + .file_search + .clone() + .map(|service| FileSearchRuntime::start_with_shutdown(service, state.shutdown_token.child_token())); let gateway = serve_gateway(state, host, port, authenticator); tokio::pin!(gateway); - tokio::select! { - gateway = &mut gateway => gateway, - status = child.wait() => { - shutdown_token.cancel(); - let status = status?; - Err(ServerError::from(CoreError::LlmProcessExited { status: status.to_string() })) - }, - () = &mut shutdown => { - info!("shutdown signal received"); - shutdown_token.cancel(); - drain_gateway(gateway.as_mut()).await + let serving = async { + tokio::select! { + gateway = &mut gateway => gateway, + status = child.wait() => { + shutdown_token.cancel(); + let status = status?; + Err(ServerError::from(CoreError::LlmProcessExited { status: status.to_string() })) + }, + () = &mut shutdown => { + info!("shutdown signal received"); + shutdown_token.cancel(); + drain_gateway(gateway.as_mut()).await + } } } + .await; + shutdown_runtime(runtime, serving).await } .await; diff --git a/crates/agentic-server/tests/openapi_test.rs b/crates/agentic-server/tests/openapi_test.rs index 55f101c0..a9ec7b13 100644 --- a/crates/agentic-server/tests/openapi_test.rs +++ b/crates/agentic-server/tests/openapi_test.rs @@ -237,3 +237,28 @@ async fn vector_store_schema_preserves_required_nullable_contracts() { .collect(); assert_eq!(kinds, ["static", "other"]); } + +#[tokio::test] +async fn file_batch_schema_and_filtered_routes_are_published() { + let spec = fetch_spec().await; + for (path, method) in [ + ("/v1/vector_stores/{store_id}/file_batches", "post"), + ("/v1/vector_stores/{store_id}/file_batches/{batch_id}", "get"), + ("/v1/vector_stores/{store_id}/file_batches/{batch_id}/cancel", "post"), + ("/v1/vector_stores/{store_id}/file_batches/{batch_id}/files", "get"), + ] { + assert!(spec["paths"][path][method].is_object(), "{method} {path}"); + } + let required = spec["components"]["schemas"]["FileBatchObject"]["required"] + .as_array() + .unwrap(); + for field in ["id", "object", "created_at", "vector_store_id", "status", "file_counts"] { + assert!(required.iter().any(|entry| entry == field)); + } + let params = spec["paths"]["/v1/vector_stores/{store_id}/file_batches/{batch_id}/files"]["get"]["parameters"] + .as_array() + .unwrap(); + for field in ["filter", "limit", "order", "after", "before"] { + assert!(params.iter().any(|entry| entry["name"] == field)); + } +} diff --git a/crates/agentic-server/tests/server_lifecycle_test.rs b/crates/agentic-server/tests/server_lifecycle_test.rs index 1ef29def..ffd2e8d5 100644 --- a/crates/agentic-server/tests/server_lifecycle_test.rs +++ b/crates/agentic-server/tests/server_lifecycle_test.rs @@ -20,6 +20,9 @@ struct ServerProcess { impl ServerProcess { fn start(upstream_port: u16, extra_args: &[&str]) -> Self { + Self::start_on(upstream_port, 0, extra_args) + } + fn start_on(upstream_port: u16, gateway_port: u16, extra_args: &[&str]) -> Self { let directory = tempfile::tempdir().expect("temporary server home"); let python = directory.path().join("python"); // exec preserves the child PID without spawning any grandchildren. @@ -38,7 +41,7 @@ impl ServerProcess { "--gateway-host", "127.0.0.1", "--gateway-port", - "0", + &gateway_port.to_string(), "--llm-ready-interval-s", "60", ]) @@ -235,3 +238,45 @@ async fn sigterm_after_gateway_startup_still_reaps_model() { .unwrap_or_else(|_| panic!("gateway did not start: {}", server.log())); server.stop("-TERM").await; } + +async fn wait_serving(server: &ServerProcess) { + timeout(TEST_TIMEOUT, async { + while !server.log().contains("gateway listening") { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("server must finish initialization"); +} + +#[tokio::test] +async fn worker_runtime_joins_on_bind_failure() { + let occupied = listener().await; + let mut server = ServerProcess::start_on( + 0, + occupied.local_addr().unwrap().port(), + &["--skip-llm-ready-check", "--db-url", "sqlite::memory:"], + ); + assert!(!server.wait().await.success()); + assert!(server.log().contains("file search workers stopped"), "{}", server.log()); + assert!(!process_exists(server.model_pid())); +} + +#[tokio::test] +async fn worker_runtime_joins_on_subprocess_exit() { + let mut server = ServerProcess::start(0, &["--skip-llm-ready-check", "--db-url", "sqlite::memory:"]); + wait_serving(&server).await; + send_signal(server.model_pid(), "-KILL"); + assert!(!server.wait().await.success()); + assert!(server.log().contains("file search workers stopped"), "{}", server.log()); +} + +#[tokio::test] +async fn worker_runtime_joins_on_serving_signals() { + for signal in ["-TERM", "-INT"] { + let mut server = ServerProcess::start(0, &["--skip-llm-ready-check", "--db-url", "sqlite::memory:"]); + wait_serving(&server).await; + server.stop(signal).await; + assert!(server.log().contains("file search workers stopped"), "{}", server.log()); + } +} diff --git a/docs/api/file-search.md b/docs/api/file-search.md index c1278a07..5bbc563c 100644 --- a/docs/api/file-search.md +++ b/docs/api/file-search.md @@ -443,8 +443,8 @@ time after both locks. Policy updates and cleanup share the same store guard. stores per call and returns the number expired. It rechecks each deadline after acquiring its store lock, commits expired status and removal of attachments/chunks atomically, and preserves uploaded files, including those used by other stores. -Repeated cleanup is safe after a restart. This layer exposes explicit cleanup; -it does not start a worker from service construction or cloning. +Repeated cleanup is safe after a restart. The server runtime invokes explicit cleanup; +service construction and cloning never start workers. ### Attachment updates and parsed content @@ -486,8 +486,8 @@ separate: a process crash or uncertain upload commit can leave unreferenced file Deletion and expiration atomically persist a blob-cleanup intent with SQL deletion; `FileSearchService::cleanup_expired_files(limit)` retries pending filesystem deletion and acknowledges it only after directory synchronization. It never sweeps arbitrary -unreferenced files that another upload might be publishing. A lifecycle worker must -invoke this method to reclaim expired bytes; visibility does not depend on that worker. +unreferenced files that another upload might be publishing. The server runtime +invokes this method to reclaim expired bytes; visibility does not depend on that worker. Expired files disappear from reads, lists, attachment reads, and search immediately, and publication rechecks expiry after model work. Cleanup removes attachments and chunks from every store while preserving independent uploads. Servers sharing SQL @@ -506,3 +506,58 @@ Each store is limited to 10,000 chunks and 64 MiB of serialized chunk data, including embeddings; ingestion enforces these limits before publication. Searching multiple stores shares the same aggregate retrieval budget. It does not expose OGX's provider catalog or asynchronous file batches. + +### Durable file batches and workers + +`POST /v1/vector_stores/{store_id}/file_batches` accepts exactly one of `file_ids` +or `files`, with 1–2000 members. Shared `attributes` and `chunking_strategy` apply +only to `file_ids`; `files` contains independent attachment requests and ignores +shared options. Auto chunk boundaries and contextual model identity are resolved +at creation. Responses use `object: "vector_store.files_batch"`. Retrieve a batch +at `.../file_batches/{batch_id}`, cancel unfinished work with POST to its `/cancel` +path, and list membership snapshots at `/files` with the usual `filter`, `limit`, +`order`, `after`, and `before` parameters. + +Creation atomically saves every member and its immediately visible `in_progress` +attachment. Duplicate IDs reject the entire request. A completed existing +attachment contributes a completed member with its existing options and no new +model calls. Any existing in-progress, failed, or cancelled attachment rejects the +entire batch with conflict. Detach unsuccessful attachments before retrying them. +Missing or expired source files reject creation atomically. An unsupported or +malformed uploaded document fails independently during processing. A batch becomes +`completed` when every member finishes, including when individual members failed; +its five `file_counts` fields describe those results. + +Cancellation preserves completed and failed members, marks unfinished members +cancelled, and invalidates their claims transactionally. Membership snapshots +survive detach and source deletion; deleting the store removes its batch history. +Source deletion, store expiration, and detach invalidate pending work. Expiration +visibility remains immediate even before cleanup. Pending attribute updates are +applied to both the eventual attachment and its chunks at publication. + +Server startup explicitly owns a `FileSearchRuntime`, separate from the cloneable +request service. SQL is the durable queue. Four worker slots share the service's +four-operation admission limit with synchronous ingestion; capacity is acquired +before claiming. `file_batch_params.max_concurrent_files_per_batch` additionally +limits active claims for each batch across instances, `file_batch_chunk_size` +bounds candidate and cleanup scans, and `cleanup_interval_seconds` schedules +expired-file, expired-store, and durable blob-intent cleanup. Failed cleanup is +retried after one second. + +Claims have a 30-second database-clock lease and renew once per second. Every +attempt receives a fresh token; each pending attachment has a separate generation. +Publication locks source, store, and job, checks fresh time after contention, then +uses the same atomic chunk writer as synchronous ingestion. Reclaimed work may +repeat provider calls, but stale attempts cannot publish. Recovery rejects changed +embedding or contextual model identities instead of mixing configurations. All +instances sharing SQL must share the Files storage mount. + +Signals stop admission and cancel model preparation. Shutdown joins worker and +parser work and releases only still-owned claims for later recovery; it does not +cancel the API batch. SQL transactions and COMMIT are allowed to finish. Worker +drain can outlast the eight-second HTTP drain while bounded extraction/filesystem +operations and configured database waits finish. PostgreSQL defaults are 30 seconds +for pool acquisition/statement execution and five seconds for lock waits; SQLite +uses a five-second busy timeout. Embedding HTTP calls have a 45-second timeout; +contextual calls use the bounded configured timeout, and both respond to runtime +cancellation. Library callers must explicitly consume the runtime with `shutdown`.