diff --git a/Cargo.lock b/Cargo.lock index d30a03cd..7d48b054 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -59,6 +59,7 @@ dependencies = [ "http", "jsonschema", "jsonwebtoken", + "multer", "rand 0.8.8", "reqwest 0.12.28", "rsa", diff --git a/crates/agentic-server-core/migrations/0006_file_expiration.sql b/crates/agentic-server-core/migrations/0006_file_expiration.sql new file mode 100644 index 00000000..d99907cd --- /dev/null +++ b/crates/agentic-server-core/migrations/0006_file_expiration.sql @@ -0,0 +1,5 @@ +-- Nullable expiry preserves existing uploads, including legacy inline rows. +ALTER TABLE file_search_files ADD COLUMN expires_at BIGINT; +ALTER TABLE file_search_files ADD COLUMN purpose TEXT; +CREATE INDEX file_search_files_expiration ON file_search_files(expires_at); +CREATE TABLE file_search_blob_cleanup (file_id TEXT PRIMARY KEY); diff --git a/crates/agentic-server-core/src/storage/file_search.rs b/crates/agentic-server-core/src/storage/file_search.rs index 7609d25b..ccfa23ff 100644 --- a/crates/agentic-server-core/src/storage/file_search.rs +++ b/crates/agentic-server-core/src/storage/file_search.rs @@ -108,12 +108,117 @@ impl FileSearchStorage { Ok(()) } + async fn file_visibility(&self) -> Result { + let mut connection = self.pool.acquire().await?; + let now = database_now(&mut connection).await?; + Ok(format!("(expires_at IS NULL OR expires_at > {now})")) + } + + pub(crate) async fn delete_file(&self, id: &str) -> Result<(), FileSearchError> { + let mut tx = self.pool.begin().await?; + let changed = sqlx::query("UPDATE file_search_files SET id = id WHERE id = $1") + .bind(id) + .execute(&mut *tx) + .await? + .rows_affected(); + if changed == 0 { + return Err(FileSearchError::NotFound("File not found".into())); + } + delete_file_in_transaction(&mut tx, id).await?; + tx.commit().await?; + Ok(()) + } + + pub(crate) async fn expire_files(&self, limit: usize) -> Result<(), FileSearchError> { + let mut tx = self.pool.begin().await?; + let now = database_now(&mut tx).await?; + let ids: Vec = sqlx::query_scalar( + "SELECT id FROM file_search_files WHERE expires_at <= $1 ORDER BY expires_at, id LIMIT $2", + ) + .bind(now) + .bind(i64::try_from(limit).unwrap_or(1000)) + .fetch_all(&mut *tx) + .await?; + for id in ids { + let locked = sqlx::query("UPDATE file_search_files SET id = id WHERE id = $1") + .bind(&id) + .execute(&mut *tx) + .await? + .rows_affected(); + if locked == 0 { + continue; + } + let now = database_now(&mut tx).await?; + let due: Option = + sqlx::query_scalar("SELECT id FROM file_search_files WHERE id = $1 AND expires_at <= $2") + .bind(&id) + .bind(now) + .fetch_optional(&mut *tx) + .await?; + if due.is_some() { + delete_file_in_transaction(&mut tx, &id).await?; + } + } + tx.commit().await?; + Ok(()) + } + + pub(crate) async fn pending_blob_cleanup(&self, limit: usize) -> Result, FileSearchError> { + Ok( + sqlx::query_scalar("SELECT file_id FROM file_search_blob_cleanup ORDER BY file_id LIMIT $1") + .bind(i64::try_from(limit).unwrap_or(1000)) + .fetch_all(self.pool.as_ref()) + .await?, + ) + } + + pub(crate) async fn acknowledge_blob_cleanup(&self, id: &str) -> Result<(), FileSearchError> { + sqlx::query("DELETE FROM file_search_blob_cleanup WHERE file_id = $1") + .bind(id) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } + + pub(crate) async fn visible_result_files( + &self, + stores: &[String], + results: &[crate::types::file_search::SearchResult], + ) -> Result, FileSearchError> { + if results.is_empty() { + return Ok(std::collections::HashSet::new()); + } + let store_placeholders = (1..=stores.len()) + .map(|index| format!("${index}")) + .collect::>() + .join(", "); + let file_placeholders = (stores.len() + 1..=stores.len() + results.len()) + .map(|index| format!("${index}")) + .collect::>() + .join(", "); + let sql = format!( + "SELECT DISTINCT file_id FROM file_search_attachments WHERE store_id IN ({store_placeholders}) AND file_id IN ({file_placeholders}) AND file_id IN (SELECT id FROM file_search_files WHERE {})", + self.file_visibility().await? + ); + let mut query = sqlx::query_scalar::<_, String>(&sql); + for store in stores { + query = query.bind(store); + } + for result in results { + query = query.bind(&result.file_id); + } + Ok(query.fetch_all(self.pool.as_ref()).await?.into_iter().collect()) + } + pub(crate) async fn has_chunks(&self, stores: &[String]) -> Result { let placeholders = (1..=stores.len()) .map(|index| format!("${index}")) .collect::>() .join(", "); - let sql = format!("SELECT chunk_index FROM file_search_chunks WHERE store_id IN ({placeholders}) LIMIT 1"); + let visibility = self.file_visibility().await?; + let sql = format!( + "SELECT chunk_index FROM file_search_chunks WHERE store_id IN ({placeholders}) AND file_id IN (SELECT id FROM file_search_files WHERE {visibility}) LIMIT 1" + ); let mut query = sqlx::query_scalar::<_, i64>(&sql); for store in stores { query = query.bind(store); @@ -151,8 +256,8 @@ impl FileSearchStorage { () = cancelled.cancelled() => return Err(FilePublicationFailure::SafeToRemove(super::local_files::cancelled_error())), result = self.pool.begin() => result.map_err(|error| FilePublicationFailure::SafeToRemove(error.into()))?, }; - let query = sqlx::query("INSERT INTO file_search_files (id, created_at, data, content_type, content_base64) VALUES ($1, $2, $3, $4, '')") - .bind(&file.id).bind(file.created_at).bind(data).bind(content_type); + let query = sqlx::query("INSERT INTO file_search_files (id, created_at, data, content_type, content_base64, expires_at, purpose) VALUES ($1, $2, $3, $4, '', $5, $6)") + .bind(&file.id).bind(file.created_at).bind(data).bind(content_type).bind(file.expires_at).bind(&file.purpose); let inserted = tokio::select! { biased; () = cancelled.cancelled() => Err(super::local_files::cancelled_error()), @@ -177,19 +282,25 @@ impl FileSearchStorage { } pub(crate) async fn file(&self, id: &str) -> Result { - sqlx::query_as("SELECT data, content_type, content_base64 FROM file_search_files WHERE id = $1") - .bind(id) - .fetch_optional(self.pool.as_ref()) - .await? - .ok_or_else(|| FileSearchError::NotFound("File not found".into())) + sqlx::query_as(&format!( + "SELECT data, content_type, content_base64 FROM file_search_files WHERE id = $1 AND {}", + self.file_visibility().await? + )) + .bind(id) + .fetch_optional(self.pool.as_ref()) + .await? + .ok_or_else(|| FileSearchError::NotFound("File not found".into())) } pub(crate) async fn file_object(&self, id: &str) -> Result { - let data: String = sqlx::query_scalar("SELECT data FROM file_search_files WHERE id = $1") - .bind(id) - .fetch_optional(self.pool.as_ref()) - .await? - .ok_or_else(|| FileSearchError::NotFound("File not found".into()))?; + let data: String = sqlx::query_scalar(&format!( + "SELECT data FROM file_search_files WHERE id = $1 AND {}", + self.file_visibility().await? + )) + .bind(id) + .fetch_optional(self.pool.as_ref()) + .await? + .ok_or_else(|| FileSearchError::NotFound("File not found".into()))?; Ok(serde_json::from_str(&data)?) } @@ -204,7 +315,7 @@ impl FileSearchStorage { pub(crate) async fn store_object(&self, id: &str) -> Result { let row = self.store(id).await?; let mut store: VectorStoreObject = serde_json::from_str(&row.data)?; - let (count, bytes): (i64, i64) = sqlx::query_as("SELECT COUNT(*), CAST(COALESCE(SUM(usage_bytes), 0) AS BIGINT) FROM file_search_attachments WHERE store_id = $1") + let (count, bytes): (i64, i64) = sqlx::query_as(&format!("SELECT COUNT(*), CAST(COALESCE(SUM(usage_bytes), 0) AS BIGINT) FROM file_search_attachments WHERE store_id = $1 AND file_id IN (SELECT id FROM file_search_files WHERE {})", self.file_visibility().await?)) .bind(id).fetch_one(self.pool.as_ref()).await?; store.file_counts.completed = count; store.file_counts.total = count; @@ -258,7 +369,7 @@ impl FileSearchStorage { file_id: &str, ) -> Result, FileSearchError> { let data: Option = - sqlx::query_scalar("SELECT data FROM file_search_attachments WHERE store_id = $1 AND file_id = $2") + sqlx::query_scalar(&format!("SELECT data FROM file_search_attachments WHERE store_id = $1 AND file_id = $2 AND file_id IN (SELECT id FROM file_search_files WHERE {})", self.file_visibility().await?)) .bind(store_id) .bind(file_id) .fetch_optional(self.pool.as_ref()) @@ -326,6 +437,21 @@ impl FileSearchStorage { "DESC" }; let filter = if store_id.is_some() { " AND store_id = $4" } else { "" }; + let visibility = self.file_visibility().await?; + let filter = match collection { + Collection::Files => format!( + "{filter} AND {visibility} AND ($4 = '' OR purpose = $4 OR (purpose IS NULL AND {purpose_json} = $4))", + purpose_json = if self.pool.acquire().await?.backend_name() == "PostgreSQL" { + "data::jsonb ->> 'purpose'" + } else { + "json_extract(data, '$.purpose')" + } + ), + Collection::Attachments => { + format!("{filter} AND file_id IN (SELECT id FROM file_search_files WHERE {visibility})") + } + Collection::Stores => filter.to_owned(), + }; let sql = format!( "SELECT data FROM {table} WHERE ($1 = '' OR created_at {operator} $2 OR (created_at = $2 AND {id} {operator} $1)){filter} ORDER BY created_at {order}, {id} {order} LIMIT $3", table = collection.table(), @@ -337,6 +463,9 @@ impl FileSearchStorage { if let Some(store_id) = store_id { query = query.bind(store_id); } + if matches!(collection, Collection::Files) { + query = query.bind(params.purpose.as_deref().unwrap_or("")); + } let rows: Vec = query.fetch_all(self.pool.as_ref()).await?; rows.iter() .map(|row| serde_json::from_str(row).map_err(Into::into)) @@ -350,8 +479,9 @@ impl FileSearchStorage { .map(|index| format!("${index}")) .collect::>() .join(", "); + let visibility = self.file_visibility().await?; let sql = format!( - "SELECT data FROM file_search_chunks WHERE store_id IN ({placeholders}) ORDER BY store_id, file_id, chunk_index" + "SELECT data FROM file_search_chunks WHERE store_id IN ({placeholders}) AND file_id IN (SELECT id FROM file_search_files WHERE {visibility}) ORDER BY store_id, file_id, chunk_index" ); let mut query = sqlx::query_scalar::<_, String>(&sql); for id in store_ids { @@ -375,6 +505,26 @@ impl FileSearchStorage { } } +/// Database wall clock, refreshed after contended writes rather than transaction start. +pub(crate) async fn database_now(connection: &mut sqlx::AnyConnection) -> Result { + let sql = if connection.backend_name() == "PostgreSQL" { + "SELECT CAST(FLOOR(EXTRACT(EPOCH FROM clock_timestamp())) AS BIGINT)" + } else { + "SELECT CAST(strftime('%s', 'now') AS BIGINT)" + }; + Ok(sqlx::query_scalar(sql).fetch_one(connection).await?) +} + +async fn delete_file_in_transaction(tx: &mut DbTransaction<'_>, id: &str) -> Result<(), FileSearchError> { + 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") + .bind(id) + .execute(&mut **tx) + .await?; + Ok(()) +} + fn validate_capacity(bytes: i64, chunks: i64) -> Result<(), FileSearchError> { if bytes > MAX_CORPUS_BYTES || chunks > MAX_CORPUS_CHUNKS { return Err(FileSearchError::InvalidRequest( @@ -409,6 +559,15 @@ async fn publish_attachment( chunks: &[String], storage_bytes: i64, ) -> Result<(), FileSearchError> { + let file_id = &attachment.object.id; + let locked = sqlx::query("UPDATE file_search_files SET id = id WHERE id = $1") + .bind(file_id) + .execute(&mut **tx) + .await? + .rows_affected(); + if locked != 1 { + return Err(FileSearchError::NotFound("File was deleted during ingestion".into())); + } // The conditional write serializes concurrent ingestions and establishes the // model dimension exactly once; no network work takes place in this transaction. let changed = sqlx::query("UPDATE file_search_stores SET embedding_dimensions = $1 WHERE id = $2 AND embedding_identity = $3 AND (embedding_dimensions = 0 OR embedding_dimensions = $1)") @@ -418,6 +577,19 @@ async fn publish_attachment( "Vector store embedding configuration changed or the vector store was deleted".into(), )); } + // The store guard can wait past the source deadline even while we own the + // file row lock. Refresh wall time after both contended parent writes. + let now = database_now(&mut *tx).await?; + let live: Option = sqlx::query_scalar( + "SELECT id FROM file_search_files WHERE id = $1 AND (expires_at IS NULL OR expires_at > $2)", + ) + .bind(file_id) + .bind(now) + .fetch_optional(&mut **tx) + .await?; + if live.is_none() { + return Err(FileSearchError::NotFound("File expired during ingestion".into())); + } let bytes: i64 = sqlx::query_scalar( "SELECT CAST(COALESCE(SUM(storage_bytes), 0) AS BIGINT) FROM file_search_attachments WHERE store_id = $1", ) diff --git a/crates/agentic-server-core/src/storage/local_files.rs b/crates/agentic-server-core/src/storage/local_files.rs index 6b54d109..db8d3cd8 100644 --- a/crates/agentic-server-core/src/storage/local_files.rs +++ b/crates/agentic-server-core/src/storage/local_files.rs @@ -61,12 +61,24 @@ impl LocalFiles { Ok(self.directory.join(id)) } + #[cfg(test)] pub(crate) async fn publish( &self, id: &str, bytes: &[u8], cancelled: &CancellationToken, ) -> Result<(), FileSearchError> { + self.publish_stream(id, &mut std::io::Cursor::new(bytes), cancelled) + .await + .map(|_| ()) + } + + pub(crate) async fn publish_stream( + &self, + id: &str, + reader: &mut (impl tokio::io::AsyncRead + Unpin), + cancelled: &CancellationToken, + ) -> Result { let destination = self.path(id)?; let temporary = self.directory.join(format!(".upload-{id}-{}", uuid::Uuid::now_v7())); self.ensure_directory().await?; @@ -79,27 +91,41 @@ impl LocalFiles { .await .map_err(|source| io_error("staging creation", source))?; let written = async { - for chunk in bytes.chunks(64 * 1024) { - if cancelled.is_cancelled() { - return Err(cancelled_error()); + let mut total = 0usize; + let mut buffer = vec![0; 64 * 1024]; + loop { + let len = tokio::select! { + biased; + () = cancelled.cancelled() => return Err(cancelled_error()), + result = reader.read(&mut buffer) => result.map_err(|source| io_error("upload read", source))?, + }; + if len == 0 { + break; } - file.write_all(chunk) + total = total.saturating_add(len); + if total > crate::tool::file_search::MAX_FILE_BYTES { + return Err(FileSearchError::InvalidRequest("File exceeds 512 MiB".into())); + } + file.write_all(&buffer[..len]) .await .map_err(|source| io_error("write", source))?; } file.flush().await.map_err(|source| io_error("flush", source))?; file.sync_all().await.map_err(|source| io_error("sync", source))?; - Ok(()) + Ok(total) } .await; // Tokio may have a buffered write in flight when cancellation was observed. // Finish it before unlinking so cleanup also works on Windows. let flushed = file.flush().await.map_err(|source| io_error("flush", source)); drop(file); - if let Err(error) = written.and(flushed) { - remove_if_present(&temporary).await?; - return Err(error); - } + let total = match written.and_then(|total| flushed.map(|()| total)) { + Ok(total) => total, + Err(error) => { + remove_if_present(&temporary).await?; + return Err(error); + } + }; if cancelled.is_cancelled() { remove_if_present(&temporary).await?; return Err(cancelled_error()); @@ -118,7 +144,7 @@ impl LocalFiles { remove_if_present(&destination).await?; return Err(error); } - Ok(()) + Ok(total) } pub(crate) async fn read( @@ -127,6 +153,24 @@ impl LocalFiles { expected_bytes: i64, max_bytes: usize, ) -> Result, FileSearchError> { + let file = self.open_read(id, expected_bytes, max_bytes).await?; + let expected = u64::try_from(expected_bytes) + .map_err(|_| io_error("read", io::Error::new(io::ErrorKind::InvalidData, "invalid blob size")))?; + let mut bytes = Vec::new(); + file.take(max_bytes as u64 + 1) + .read_to_end(&mut bytes) + .await + .map_err(|source| io_error("read", source))?; + if bytes.len() > max_bytes || bytes.len() as u64 != expected { + return Err(io_error( + "read", + io::Error::new(io::ErrorKind::InvalidData, "blob size changed while reading"), + )); + } + Ok(bytes) + } + + async fn open_read(&self, id: &str, expected_bytes: i64, max_bytes: usize) -> Result { let path = self.path(id)?; let metadata = fs::symlink_metadata(&path) .await @@ -151,22 +195,80 @@ impl LocalFiles { io::Error::new(io::ErrorKind::InvalidData, "blob size mismatch or limit exceeded"), )); } - let mut bytes = Vec::new(); - file.take(max_bytes as u64 + 1) - .read_to_end(&mut bytes) + Ok(file) + } + + pub(crate) async fn stream( + &self, + id: &str, + expected_bytes: i64, + sender: tokio::sync::mpsc::Sender>, + ready: tokio::sync::oneshot::Sender>, + ) { + let mut file = match self + .open_read(id, expected_bytes, crate::tool::file_search::MAX_FILE_BYTES) .await - .map_err(|source| io_error("read", source))?; - if bytes.len() > max_bytes || bytes.len() as u64 != expected { - return Err(io_error( - "read", - io::Error::new(io::ErrorKind::InvalidData, "blob size changed while reading"), - )); + { + Ok(file) => file, + Err(error) => { + let _ = ready.send(Err(error)); + return; + } + }; + if ready.send(Ok(())).is_err() { + return; + } + let result = async { + let mut remaining = u64::try_from(expected_bytes) + .map_err(|_| io_error("read", io::Error::new(io::ErrorKind::InvalidData, "invalid blob size")))?; + loop { + if sender.is_closed() { + return Ok(()); + } + let mut buffer = vec![0; 64 * 1024]; + let len = file + .read(&mut buffer) + .await + .map_err(|source| io_error("read", source))?; + if len as u64 > remaining || (len == 0 && remaining != 0) { + return Err(io_error( + "read", + io::Error::new(io::ErrorKind::InvalidData, "blob size changed while reading"), + )); + } + if len == 0 { + break; + } + remaining -= len as u64; + // Verify EOF before emitting the final declared bytes: an HTTP + // client may consider Content-Length satisfied immediately. + if remaining == 0 { + let mut extra = [0]; + let read = file.read(&mut extra).await.map_err(|source| io_error("read", source))?; + let metadata = file.metadata().await.map_err(|source| io_error("read", source))?; + if read != 0 || metadata.len() != u64::try_from(expected_bytes).unwrap_or(u64::MAX) { + return Err(io_error( + "read", + io::Error::new(io::ErrorKind::InvalidData, "blob size changed while reading"), + )); + } + } + buffer.truncate(len); + if sender.send(Ok(bytes::Bytes::from(buffer))).await.is_err() { + return Ok(()); + } + } + Ok(()) + } + .await; + if let Err(error) = result { + let _ = sender.send(Err(error)).await; } - Ok(bytes) } pub(crate) async fn remove(&self, id: &str) -> Result<(), FileSearchError> { remove_if_present(&self.path(id)?).await?; + self.ensure_directory().await?; self.sync_directory().await } diff --git a/crates/agentic-server-core/src/storage/pgvector.rs b/crates/agentic-server-core/src/storage/pgvector.rs index ea894eda..36bfa45f 100644 --- a/crates/agentic-server-core/src/storage/pgvector.rs +++ b/crates/agentic-server-core/src/storage/pgvector.rs @@ -194,7 +194,7 @@ impl PgvectorStorage { } sql.push_bind(store); } - sql.push(")"); + sql.push(") AND file_id IN (SELECT id FROM file_search_files WHERE expires_at IS NULL OR expires_at > EXTRACT(EPOCH FROM clock_timestamp()))"); if let Some(filter) = filter { sql.push(" AND "); push_filter(&mut sql, filter)?; diff --git a/crates/agentic-server-core/src/storage/schema.rs b/crates/agentic-server-core/src/storage/schema.rs index 49895f6b..5c5821b0 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 = 40; -const REQUIRED_POSTGRES_CONSTRAINT_COUNT: i64 = 14; +const REQUIRED_POSTGRES_SCHEMA_COLUMN_COUNT: i64 = 43; +const REQUIRED_POSTGRES_CONSTRAINT_COUNT: i64 = 15; const REQUIRED_POSTGRES_INTEGER_COLUMN_COUNT: i64 = 4; const POSTGRES_INTEGER_WIDENING_SQL: &str = " ALTER TABLE conversations @@ -91,6 +91,9 @@ where ('file_search_files', 'data', 'text', 'NO'), \ ('file_search_files', 'content_type', 'text', 'NO'), \ ('file_search_files', 'content_base64', 'text', 'NO'), \ + ('file_search_files', 'expires_at', 'bigint', 'YES'), \ + ('file_search_files', 'purpose', 'text', 'YES'), \ + ('file_search_blob_cleanup', 'file_id', 'text', 'NO'), \ ('file_search_stores', 'id', 'text', 'NO'), \ ('file_search_stores', 'created_at', 'bigint', 'NO'), \ ('file_search_stores', 'data', 'text', 'NO'), \ @@ -146,6 +149,7 @@ where ('conversations', 'f', \ 'FOREIGN KEY (latest_response_id) REFERENCES responses(id) ON DELETE SET NULL'), \ ('file_search_files', 'p', 'PRIMARY KEY (id)'), \ + ('file_search_blob_cleanup', 'p', 'PRIMARY KEY (file_id)'), \ ('file_search_stores', 'p', 'PRIMARY KEY (id)'), \ ('file_search_attachments', 'p', 'PRIMARY KEY (store_id, file_id)'), \ ('file_search_chunks', 'p', 'PRIMARY KEY (store_id, file_id, chunk_index)'), \ @@ -308,7 +312,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_files', 'file_search_stores', 'file_search_attachments', 'file_search_chunks', 'file_search_blob_cleanup') \ ORDER BY table_namespace.nspname::text", ) .fetch_all(&mut *connection) @@ -399,6 +403,10 @@ pub(crate) async fn verify_persistence_ready(pool: &DbPool) -> DbResult<()> { ('responses', 'INSERT'), \ ('file_search_files', 'SELECT'), \ ('file_search_files', 'INSERT'), \ + ('file_search_files', 'UPDATE'), \ + ('file_search_blob_cleanup', 'SELECT'), \ + ('file_search_blob_cleanup', 'INSERT'), \ + ('file_search_blob_cleanup', 'DELETE'), \ ('file_search_files', 'DELETE'), \ ('file_search_stores', 'SELECT'), \ ('file_search_stores', 'INSERT'), \ @@ -411,7 +419,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) = 19 \ + AND COUNT(table_relation.oid) = 23 \ AND COALESCE(BOOL_AND( \ has_table_privilege(current_user, table_relation.oid, required.privilege) \ ), false) \ @@ -440,7 +448,8 @@ pub(crate) async fn verify_persistence_ready(pool: &DbPool) -> DbResult<()> { "SELECT id FROM conversations LIMIT 0", "SELECT id FROM items LIMIT 0", "SELECT id FROM responses LIMIT 0", - "SELECT id, created_at, data, content_type, content_base64 FROM file_search_files 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, created_at, data, embedding_identity, embedding_dimensions FROM file_search_stores LIMIT 0", "SELECT store_id, file_id, created_at, usage_bytes, storage_bytes, data FROM file_search_attachments LIMIT 0", "SELECT store_id, file_id, chunk_index, data FROM file_search_chunks LIMIT 0", @@ -696,6 +705,15 @@ mod tests { .execute(pool.as_ref()) .await .unwrap(); + assert!( + verify_persistence_ready(pool.as_ref()).await.is_err(), + "expiry migration is required" + ); + assert!(wrapper.ensure_schema_ready_with_marker(true).await.is_err()); + sqlx::raw_sql(include_str!("../../migrations/0006_file_expiration.sql")) + .execute(pool.as_ref()) + .await + .unwrap(); verify_persistence_ready(pool.as_ref()).await.unwrap(); wrapper.ensure_schema_ready_with_marker(true).await.unwrap(); } @@ -798,6 +816,14 @@ mod tests { .execute(&mut *connection) .await .unwrap(); + assert!( + supervisor.ensure_schema_ready_with_marker(true).await.is_err(), + "expiry migration is required" + ); + sqlx::raw_sql(include_str!("../../migrations/0006_file_expiration.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") @@ -839,6 +865,7 @@ mod tests { include_str!("../../migrations/0003_index_conversation_sequence.sql"), include_str!("../../migrations/0004_link_conversation_latest_response.sql"), include_str!("../../migrations/0005_file_search.sql"), + include_str!("../../migrations/0006_file_expiration.sql"), ] { sqlx::raw_sql(migration) .execute(&mut *connection) diff --git a/crates/agentic-server-core/src/tool/file_search/files.rs b/crates/agentic-server-core/src/tool/file_search/files.rs new file mode 100644 index 00000000..5d041a50 --- /dev/null +++ b/crates/agentic-server-core/src/tool/file_search/files.rs @@ -0,0 +1,237 @@ +//! Bounded byte transport and owned publication tasks for the Files service. +use super::{ + FileObject, FileSearchError, FileSearchService, LocalFiles, MAX_FILE_BYTES, decode_content, invalid, size_i64, +}; +use crate::storage::file_search::FilePublicationFailure; +use crate::types::file_search::FileExpiresAfter; +use tokio::io::AsyncWriteExt; +use tokio::sync::{mpsc, oneshot}; +use tokio_util::sync::{CancellationToken, DropGuard}; + +/// A bounded upload whose drop cancels publication and cleans staging bytes. +pub struct FileUpload { + writer: tokio::io::DuplexStream, + metadata: Option)>>, + task: tokio::task::JoinHandle>, + bytes: usize, + _cancel: DropGuard, +} + +impl FileUpload { + /// Writes bytes with filesystem backpressure. + /// # Errors + /// Returns size-limit, cancellation, or filesystem errors. + pub async fn write(&mut self, bytes: &[u8]) -> Result<(), FileSearchError> { + self.bytes = self.bytes.saturating_add(bytes.len()); + if self.bytes > MAX_FILE_BYTES { + return invalid("File exceeds 512 MiB"); + } + self.writer + .write_all(bytes) + .await + .map_err(|source| FileSearchError::FileStorage { + operation: "upload", + source, + }) + } + + /// Publishes only after every multipart field has been validated by the caller. + /// # Errors + /// Returns metadata validation, filesystem, or database errors. + pub async fn finish( + mut self, + purpose: &str, + expires: Option, + ) -> Result { + if self.bytes > MAX_FILE_BYTES { + return invalid("File exceeds 512 MiB"); + } + validate_purpose(purpose)?; + if expires + .as_ref() + .is_some_and(|policy| !(3600..=2_592_000).contains(&policy.seconds)) + { + return invalid("expires_after.seconds must be between 3600 and 2592000"); + } + self.writer + .shutdown() + .await + .map_err(|source| FileSearchError::FileStorage { + operation: "upload finish", + source, + })?; + if let Some(sender) = self.metadata.take() { + let _ = sender.send((purpose.to_owned(), expires)); + } + self.task.await? + } +} + +/// A bounded original-byte stream. Dropping it stops the owned reader after active I/O finishes. +pub struct FileDownload { + pub bytes: u64, + receiver: mpsc::Receiver>, +} + +impl futures::Stream for FileDownload { + type Item = Result; + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + self.receiver.poll_recv(cx) + } +} + +pub(super) fn validate_purpose(purpose: &str) -> Result<(), FileSearchError> { + if !matches!( + purpose, + "assistants" | "batch" | "fine-tune" | "vision" | "user_data" | "evals" + ) { + return invalid("unsupported file purpose"); + } + Ok(()) +} + +impl FileSearchService { + /// Begins bounded staging; no file metadata becomes visible until `finish`. + /// # Errors + /// Returns invalid metadata or capacity errors. + pub fn begin_file_upload(&self, filename: &str, content_type: &str) -> Result { + if filename.is_empty() || filename.len() > 255 || filename.chars().any(char::is_control) { + return invalid("filename must contain 1 to 255 bytes without control characters"); + } + if content_type.is_empty() || content_type.len() > 256 || content_type.chars().any(char::is_control) { + return invalid("content_type must contain 1 to 256 bytes without control characters"); + } + let permit = self.permit()?; + let (writer, mut reader) = tokio::io::duplex(64 * 1024); + let (metadata, receiver) = oneshot::channel::<(String, Option)>(); + let cancelled = CancellationToken::new(); + let guard = cancelled.clone().drop_guard(); + let storage = self.storage.clone(); + let files = self.files.clone(); + let filename = filename.to_owned(); + let content_type = content_type.to_owned(); + let task = tokio::spawn(async move { + let _permit = permit; + let id = format!("file-{}", uuid::Uuid::now_v7()); + let bytes = files.publish_stream(&id, &mut reader, &cancelled).await?; + let metadata = tokio::select! { + biased; + () = cancelled.cancelled() => None, + result = receiver => result.ok(), + }; + let Some((purpose, expires)) = metadata else { + files.remove(&id).await?; + return Err(crate::storage::local_files::cancelled_error()); + }; + let created_at = chrono::Utc::now().timestamp(); + let seconds = expires + .map(|policy| policy.seconds) + .or_else(|| (purpose == "batch").then_some(2_592_000)); + let file = FileObject { + id, + object: "file".into(), + bytes: size_i64(bytes)?, + created_at, + filename, + purpose, + expires_at: seconds.map(|seconds| created_at + i64::from(seconds)), + status: "processed".into(), + }; + match storage.upload(&file, &content_type, &cancelled).await { + Ok(()) => Ok(file), + Err(FilePublicationFailure::SafeToRemove(error)) => { + files.remove(&file.id).await?; + Err(error) + } + Err(FilePublicationFailure::Indeterminate(error)) => Err(error.into()), + } + }); + Ok(FileUpload { + writer, + metadata: Some(metadata), + task, + bytes: 0, + _cancel: guard, + }) + } + + /// Opens a verified stream; visibility is linearized before opening the download. + /// # Errors + /// Returns not-found, integrity, capacity, or storage errors. + pub async fn download_file(&self, id: &str) -> Result { + LocalFiles::validate_id(id)?; + let permit = self.permit()?; + let uploaded = self.storage.file(id).await?; + let object: FileObject = serde_json::from_str(&uploaded.data)?; + let expected = u64::try_from(object.bytes).map_err(|_| FileSearchError::ProviderProtocol)?; + let (sender, receiver) = mpsc::channel(2); + let (ready, opened) = oneshot::channel(); + let files = self.files.clone(); + let id = id.to_owned(); + tokio::spawn(async move { + let retained_permit = permit; + if uploaded.content_base64.is_empty() { + files.stream(&id, object.bytes, sender, ready).await; + } else { + let decode_permit = retained_permit.clone(); + let decoded = tokio::task::spawn_blocking(move || { + let _permit = decode_permit; + decode_content(&uploaded.content_base64) + }) + .await + .map_err(FileSearchError::from) + .and_then(std::convert::identity); + match decoded { + Ok(bytes) if bytes.len() as u64 == expected => { + if ready.send(Ok(())).is_ok() { + for chunk in bytes.chunks(64 * 1024) { + if sender.send(Ok(bytes::Bytes::copy_from_slice(chunk))).await.is_err() { + break; + } + } + } + } + Ok(_) => { + let _ = ready.send(Err(FileSearchError::Unavailable("Legacy blob size mismatch".into()))); + } + Err(error) => { + let _ = ready.send(Err(error)); + } + } + } + }); + opened + .await + .map_err(|_| FileSearchError::Unavailable("Download reader stopped".into()))??; + Ok(FileDownload { + bytes: expected, + receiver, + }) + } + + /// Deletes due files transactionally and replays durable blob deletion intents. + /// # Errors + /// Returns invalid batch size or storage errors. Failed intents remain retryable. + pub async fn cleanup_expired_files(&self, limit: usize) -> Result { + if !(1..=1000).contains(&limit) { + return invalid("cleanup limit must be between 1 and 1000"); + } + let permit = self.permit()?; + let storage = self.storage.clone(); + let files = self.files.clone(); + tokio::spawn(async move { + let _permit = permit; + storage.expire_files(limit).await?; + let pending = storage.pending_blob_cleanup(limit).await?; + for id in &pending { + files.remove(id).await?; + storage.acknowledge_blob_cleanup(id).await?; + } + Ok(pending.len()) + }) + .await? + } +} 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 ca2bb7a3..b2a32472 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::{FileSearchService, MAX_FILE_BYTES}; +pub use service::{FileDownload, 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 cf1c0006..40ac7e90 100644 --- a/crates/agentic-server-core/src/tool/file_search/service.rs +++ b/crates/agentic-server-core/src/tool/file_search/service.rs @@ -10,15 +10,15 @@ use std::{ use base64::{Engine, engine::general_purpose::STANDARD}; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; -use tokio_util::sync::CancellationToken; +#[path = "files.rs"] +mod files; +pub use files::{FileDownload, FileUpload}; use super::{embeddings::Embeddings, ingest, models::Models, ranking}; use crate::{ storage::{ DbPool, - file_search::{ - Collection, FilePublicationFailure, FileSearchStorage, PreparedAttachment, StoredChunk, StoredVectorStore, - }, + file_search::{Collection, FileSearchStorage, PreparedAttachment, StoredChunk, StoredVectorStore}, local_files::LocalFiles, }, types::file_search::{ @@ -30,7 +30,8 @@ use crate::{ }; /// Maximum accepted size for one uploaded file. -pub const MAX_FILE_BYTES: usize = 20 * 1024 * 1024; +pub const MAX_FILE_BYTES: usize = 512 * 1024 * 1024; +const MAX_INGESTION_BYTES: usize = 20 * 1024 * 1024; struct CancelIngestionOnDrop(Arc); @@ -150,57 +151,23 @@ impl FileSearchService { purpose: &str, bytes: Vec, ) -> Result { - let permit = self.permit()?; - if filename.is_empty() || filename.len() > 255 || filename.chars().any(char::is_control) { - return invalid("filename must contain 1 to 255 bytes without control characters"); - } - if !matches!(purpose, "assistants" | "user_data") { - return invalid("file search accepts purpose assistants or user_data"); - } - if bytes.is_empty() || bytes.len() > MAX_FILE_BYTES { - return invalid("file must contain 1 byte to 20 MiB"); - } - if content_type.is_empty() || content_type.len() > 256 || content_type.chars().any(char::is_control) { - return invalid("content_type must contain 1 to 256 bytes without control characters"); - } - let file = FileObject { - id: format!("file-{}", uuid::Uuid::now_v7()), - object: "file".into(), - bytes: size_i64(bytes.len())?, - created_at: chrono::Utc::now().timestamp(), - filename: filename.into(), - purpose: purpose.into(), - status: "processed".into(), - }; - let storage = self.storage.clone(); - let files = self.files.clone(); - let content_type = content_type.to_owned(); - let cancelled = CancellationToken::new(); - let _cancel_on_drop = cancelled.clone().drop_guard(); - // The caller joins this operation normally. On caller cancellation it - // retains its permit and finishes rollback/cleanup without abandoning I/O. - tokio::spawn(async move { - let _permit = permit; - files.publish(&file.id, &bytes, &cancelled).await?; - match storage.upload(&file, &content_type, &cancelled).await { - Ok(()) => Ok(file), - Err(FilePublicationFailure::SafeToRemove(error)) => { - files.remove(&file.id).await?; - Err(error) - } - Err(FilePublicationFailure::Indeterminate(error)) => Err(error.into()), - } - }) - .await? + let mut upload = self.begin_file_upload(filename, content_type)?; + upload.write(&bytes).await?; + upload.finish(purpose, None).await } /// # Errors /// Returns pagination validation or storage errors. pub async fn list_files(&self, params: &ListParams) -> Result, FileSearchError> { - validate_list(params)?; + let mut params = params.clone(); + params.limit = Some(params.limit.unwrap_or(10000)); + validate_pagination(¶ms, 10000)?; + if let Some(purpose) = ¶ms.purpose { + files::validate_purpose(purpose)?; + } Ok(page( - self.storage.list(Collection::Files, None, params).await?, - params, + self.storage.list(Collection::Files, None, ¶ms).await?, + ¶ms, |file: &FileObject| &file.id, )) } @@ -236,12 +203,10 @@ impl FileSearchService { // even if the caller disconnects. A failed SQL commit retains the blob. tokio::spawn(async move { let _permit = permit; - let file = storage.file(&id).await?; - storage.delete(Collection::Files, &id, None).await?; - if file.content_base64.is_empty() { - files.remove(&id).await?; - } - Ok(deleted(&id, "file.deleted")) + storage.delete_file(&id).await?; + files.remove(&id).await?; + storage.acknowledge_blob_cleanup(&id).await?; + Ok(deleted(&id, "file")) }) .await? } @@ -258,7 +223,7 @@ impl FileSearchService { let id = id.to_owned(); return tokio::spawn(async move { let _permit = permit; - files.read(&id, expected_bytes, MAX_FILE_BYTES).await + files.read(&id, expected_bytes, MAX_INGESTION_BYTES).await }) .await?; } @@ -739,6 +704,8 @@ impl FileSearchService { if prepare_context { data = self.prepare_context(data, permit).await?; } + let visible = self.storage.visible_result_files(store_ids, &data).await?; + data.retain(|result| visible.contains(&result.file_id)); Ok(SearchResponse { object: "vector_store.search_results.page".into(), search_query: queries, @@ -761,7 +728,7 @@ fn deleted(id: &str, object: &str) -> DeleteObject { } } fn decode_content(encoded: &str) -> Result, FileSearchError> { - if encoded.len() > MAX_FILE_BYTES.div_ceil(3) * 4 { + if encoded.len() > MAX_INGESTION_BYTES.div_ceil(3) * 4 { return Err(FileSearchError::Unavailable( "Stored file exceeds the file size limit".into(), )); @@ -771,8 +738,14 @@ fn decode_content(encoded: &str) -> Result, FileSearchError> { .map_err(|_| FileSearchError::Unavailable("Stored file content could not be decoded".into())) } fn validate_list(params: &ListParams) -> Result<(), FileSearchError> { - if !(1..=100).contains(¶ms.limit.unwrap_or(20)) { - return invalid("limit must be between 1 and 100"); + if params.purpose.is_some() { + return invalid("purpose applies only to Files lists"); + } + validate_pagination(params, 100) +} +fn validate_pagination(params: &ListParams, maximum: usize) -> Result<(), FileSearchError> { + if !(1..=maximum).contains(¶ms.limit.unwrap_or(20)) { + return invalid(&format!("limit must be between 1 and {maximum}")); } if params.after.is_some() && params.before.is_some() { return invalid("provide only one of after and before"); @@ -876,6 +849,15 @@ mod local_file_tests { #[test] fn cancelled_local_read_retains_capacity_until_filesystem_io_finishes() { + cancelled_read_retains_capacity(false); + } + + #[test] + fn cancelled_streaming_download_retains_capacity_until_filesystem_io_finishes() { + cancelled_read_retains_capacity(true); + } + + fn cancelled_read_retains_capacity(streaming: bool) { let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .max_blocking_threads(1) @@ -906,8 +888,17 @@ mod local_file_tests { resume_receiver.recv().unwrap(); }); waiting.await.unwrap(); - let mut read = - Box::pin(service.read_content(&file.id, file.bytes, String::new(), service.permit().unwrap())); + let mut read: std::pin::Pin>>> = + if streaming { + Box::pin(async { service.download_file(&file.id).await.map(|_| ()) }) + } else { + Box::pin(async { + service + .read_content(&file.id, file.bytes, String::new(), service.permit().unwrap()) + .await + .map(|_| ()) + }) + }; assert!( tokio::time::timeout(std::time::Duration::from_millis(20), &mut read) .await diff --git a/crates/agentic-server-core/src/types/file_search.rs b/crates/agentic-server-core/src/types/file_search.rs index 03603f5b..461b2632 100644 --- a/crates/agentic-server-core/src/types/file_search.rs +++ b/crates/agentic-server-core/src/types/file_search.rs @@ -343,6 +343,22 @@ pub struct SearchContent { pub text: String, } +/// Files expiration policy measured from creation. +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +pub struct FileExpiresAfter { + pub anchor: FileExpirationAnchor, + pub seconds: u32, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] +pub enum FileExpirationAnchor { + CreatedAt, +} + #[derive(Clone, Debug, Serialize, Deserialize)] #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] pub struct FileObject { @@ -352,6 +368,8 @@ pub struct FileObject { pub created_at: i64, pub filename: String, pub purpose: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub expires_at: Option, pub status: String, } @@ -535,6 +553,7 @@ pub enum ListOrder { #[serde(deny_unknown_fields)] #[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] pub struct ListParams { + pub purpose: Option, pub limit: Option, pub after: Option, pub before: Option, diff --git a/crates/agentic-server-core/tests/file_search_models.rs b/crates/agentic-server-core/tests/file_search_models.rs index 6661389f..ddf1d7b8 100644 --- a/crates/agentic-server-core/tests/file_search_models.rs +++ b/crates/agentic-server-core/tests/file_search_models.rs @@ -8,6 +8,9 @@ use std::sync::{Arc, Mutex}; struct Fixture { requests: Arc>>, rerank_response: Arc>>, + rerank_pause: Arc>, + rerank_started: Arc, + rerank_resume: Arc, chat_fail: Arc>, chat_fail_after: Arc>>, chat_response: Arc>>, @@ -79,6 +82,11 @@ async fn rerank(State(state): State, headers: axum::http::HeaderMap, Js .to_owned(), ); state.requests.lock().unwrap().push(("rerank".into(), input.clone())); + let pause = *state.rerank_pause.lock().unwrap(); + if pause { + state.rerank_started.notify_one(); + state.rerank_resume.notified().await; + } if let Some(output) = state.rerank_response.lock().unwrap().clone() { return Json(output); } @@ -869,3 +877,27 @@ fn hyphenated_openai_ranker_is_accepted_in_deployment_configuration() { assert_eq!(ranker, Ranker::Default20240821); assert_eq!(serde_json::to_string(&ranker).unwrap(), "\"default-2024-08-21\""); } + +#[tokio::test] +async fn expiration_during_reranking_discards_cached_candidates() { + let setup = setup(false).await; + let store = setup + .service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + let file_id = attach(&setup.service, &store.id, "coral preferred", None).await; + *setup.state.rerank_pause.lock().unwrap() = true; + let service = setup.service.clone(); + let request: SearchRequest = + serde_json::from_value(json!({"query":"coral","ranking_options":{"ranker":"neural"}})).unwrap(); + let search = tokio::spawn(async move { service.search(&[store.id], &request).await }); + setup.state.rerank_started.notified().await; + sqlx::query("UPDATE file_search_files SET expires_at = 1 WHERE id = $1") + .bind(file_id) + .execute(setup.pool.as_ref()) + .await + .unwrap(); + setup.state.rerank_resume.notify_one(); + assert!(search.await.unwrap().unwrap().data.is_empty()); +} diff --git a/crates/agentic-server-core/tests/file_search_service.rs b/crates/agentic-server-core/tests/file_search_service.rs index d233d0ca..9297931a 100644 --- a/crates/agentic-server-core/tests/file_search_service.rs +++ b/crates/agentic-server-core/tests/file_search_service.rs @@ -7,7 +7,7 @@ use std::sync::Arc; struct TestService { service: FileSearchService, - _files: tempfile::TempDir, + files: tempfile::TempDir, } impl std::ops::Deref for TestService { @@ -28,7 +28,7 @@ async fn service() -> TestService { let pool = create_pool_with_schema(Some("sqlite::memory:")).await.unwrap(); let files = tempfile::tempdir().unwrap(); let service = FileSearchService::new(pool, Arc::new(reqwest::Client::new()), file_config(&files)).unwrap(); - TestService { service, _files: files } + TestService { service, files } } fn query(text: &str) -> SearchRequest { @@ -175,6 +175,7 @@ async fn local_files_preserve_legacy_database_content() { let pool = create_pool_with_schema(Some("sqlite::memory:")).await.unwrap(); let service = FileSearchService::new(pool.clone(), Arc::new(reqwest::Client::new()), file_config(&files)).unwrap(); let file = FileObject { + expires_at: None, id: format!("file-{}", uuid::Uuid::now_v7()), object: "file".into(), bytes: 11, @@ -186,6 +187,9 @@ async fn local_files_preserve_legacy_database_content() { sqlx::query("INSERT INTO file_search_files (id, created_at, data, content_type, content_base64) VALUES ($1, 0, $2, 'text/plain', $3)") .bind(&file.id).bind(serde_json::to_string(&file).unwrap()).bind(base64::engine::general_purpose::STANDARD.encode(b"legacy reef")).execute(pool.as_ref()).await.unwrap(); assert_eq!(service.file_content(&file.id).await.unwrap(), b"legacy reef"); + let download = service.download_file(&file.id).await.unwrap(); + let chunks: Vec<_> = futures::TryStreamExt::try_collect(download).await.unwrap(); + assert_eq!(chunks.concat(), b"legacy reef"); let store = service .create_vector_store(CreateVectorStoreRequest::default()) .await @@ -335,6 +339,7 @@ async fn before_pagination_returns_adjacent_files_stores_and_attachments() { before: Some(id.clone()), order: Some(order), after: None, + purpose: None, }; let page = service.list_files(&before(&file_ids[5])).await.unwrap(); assert!(page.has_more); @@ -725,11 +730,13 @@ enum ProviderMode { WrongDimensions, Fail, Wait, + Pause, } #[derive(Clone)] struct ProviderState { mode: Arc>, started: Arc, + resume: Arc, } #[derive(serde::Deserialize)] struct EmbeddingInput { @@ -746,6 +753,10 @@ async fn provider( state.started.notify_one(); std::future::pending::<()>().await; } + if matches!(mode, ProviderMode::Pause) { + state.started.notify_one(); + state.resume.notified().await; + } if matches!(mode, ProviderMode::Fail) { return ( axum::http::StatusCode::BAD_GATEWAY, @@ -772,6 +783,7 @@ async fn embedding_service() -> ( let state = ProviderState { mode: Arc::new(std::sync::Mutex::new(ProviderMode::Good)), started: Arc::new(tokio::sync::Notify::new()), + resume: Arc::new(tokio::sync::Notify::new()), }; let app = axum::Router::new() .route("/v1/embeddings", axum::routing::post(provider)) @@ -791,7 +803,7 @@ async fn embedding_service() -> ( }; let pool = create_pool_with_schema(Some("sqlite::memory:")).await.unwrap(); let service = FileSearchService::new(pool.clone(), Arc::new(reqwest::Client::new()), config.clone()).unwrap(); - (TestService { service, _files: files }, state, task, pool, config) + (TestService { service, files }, state, task, pool, config) } #[tokio::test] @@ -1532,6 +1544,32 @@ async fn postgres_pgvector_indexed_semantics_restart_filters_and_deletion() { .file_id, coral.id ); + sqlx::query("UPDATE file_search_files SET expires_at = 1 WHERE id = $1") + .bind(&coral.id) + .execute(pool.as_ref()) + .await + .unwrap(); + for mode in [SearchMode::Keyword, SearchMode::Semantic, SearchMode::Hybrid] { + let expired_query = SearchRequest { + query: SearchQuery::Text("coral".into()), + search_mode: Some(mode), + ..query.clone() + }; + assert!( + restarted + .search(std::slice::from_ref(&store.id), &expired_query) + .await + .unwrap() + .data + .iter() + .all(|result| result.file_id != coral.id) + ); + } + sqlx::query("UPDATE file_search_files SET expires_at = NULL WHERE id = $1") + .bind(&coral.id) + .execute(pool.as_ref()) + .await + .unwrap(); restarted.detach_file(&store.id, &coral.id).await.unwrap(); assert!( restarted @@ -1934,3 +1972,355 @@ async fn assert_established_ivfflat_search_ignores_initialization_lock( } blocker.rollback().await.unwrap(); } + +#[tokio::test] +async fn expired_uploads_are_hidden_and_cleanup_replays_durable_intents() { + expiration_cleanup_contract("sqlite::memory:").await; +} + +#[tokio::test] +#[ignore = "requires isolated TEST_POSTGRES_URL"] +async fn postgres_expired_uploads_cleanup_replays_durable_intents() { + expiration_cleanup_contract(&std::env::var("TEST_POSTGRES_URL").unwrap()).await; +} + +async fn expiration_cleanup_contract(url: &str) { + let files = tempfile::tempdir().unwrap(); + let pool = create_pool_with_schema(Some(url)).await.unwrap(); + let service = FileSearchService::new(pool.clone(), Arc::new(reqwest::Client::new()), file_config(&files)).unwrap(); + let file = service + .upload_file( + "expire.txt", + "text/plain", + "assistants", + b"expiration visibility".to_vec(), + ) + .await + .unwrap(); + let retained = service + .upload_file("retained.txt", "text/plain", "assistants", b"retained".to_vec()) + .await + .unwrap(); + let store = service + .create_vector_store(CreateVectorStoreRequest { + file_ids: vec![file.id.clone()], + ..Default::default() + }) + .await + .unwrap(); + sqlx::query("UPDATE file_search_files SET expires_at = 1 WHERE id = $1") + .bind(&file.id) + .execute(pool.as_ref()) + .await + .unwrap(); + assert_eq!(service.get_file(&file.id).await.unwrap_err().status_code(), 404); + assert_eq!(service.file_content(&file.id).await.unwrap_err().status_code(), 404); + assert!( + !service + .list_files(&ListParams::default()) + .await + .unwrap() + .data + .iter() + .any(|item| item.id == file.id) + ); + assert!( + service + .list_vector_store_files(&store.id, &ListParams::default()) + .await + .unwrap() + .data + .is_empty() + ); + assert!( + service + .search(std::slice::from_ref(&store.id), &query("expiration")) + .await + .unwrap() + .data + .is_empty() + ); + assert_eq!( + service + .get_vector_store_file(&store.id, &file.id) + .await + .unwrap_err() + .status_code(), + 404 + ); + // A cleanup failure after SQL commit must preserve durable intent for restart. + let blob = files.path().join(&file.id); + std::fs::remove_file(&blob).unwrap(); + std::fs::create_dir(&blob).unwrap(); + assert!(service.cleanup_expired_files(100).await.is_err()); + let pending: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM file_search_blob_cleanup") + .fetch_one(pool.as_ref()) + .await + .unwrap(); + assert_eq!(pending, 1); + std::fs::remove_dir(&blob).unwrap(); + assert_eq!(service.cleanup_expired_files(100).await.unwrap(), 1); + assert_eq!(service.cleanup_expired_files(100).await.unwrap(), 0); + assert_eq!(service.file_content(&retained.id).await.unwrap(), b"retained"); + service.delete_file(&retained.id).await.unwrap(); + service.delete_vector_store(&store.id).await.unwrap(); +} + +#[tokio::test] +async fn dropped_streaming_upload_never_publishes_metadata_or_leaks_staging() { + let files = tempfile::tempdir().unwrap(); + let pool = create_pool_with_schema(Some("sqlite::memory:")).await.unwrap(); + let service = FileSearchService::new(pool, Arc::new(reqwest::Client::new()), file_config(&files)).unwrap(); + let mut upload = service + .begin_file_upload("cancelled.bin", "application/octet-stream") + .unwrap(); + upload.write(&vec![42; 1024 * 1024]).await.unwrap(); + drop(upload); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if std::fs::read_dir(files.path()).unwrap().count() == 0 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!( + service + .list_files(&ListParams::default()) + .await + .unwrap() + .data + .is_empty() + ); +} + +#[tokio::test] +async fn streaming_upload_accepts_exact_512_mib_and_rejects_one_more_byte() { + use futures::TryStreamExt as _; + let service = service().await; + let mut upload = service + .begin_file_upload("boundary.bin", "application/octet-stream") + .unwrap(); + let chunk = vec![0x6b; 64 * 1024]; + for _ in 0..8192 { + upload.write(&chunk).await.unwrap(); + } + let file = upload.finish("user_data", None).await.unwrap(); + assert_eq!(file.bytes, 536_870_912); + assert!( + service.file_content(&file.id).await.is_err(), + "ingestion convenience read remains bounded" + ); + let mut download = service.download_file(&file.id).await.unwrap(); + let mut total = 0usize; + while let Some(bytes) = download.try_next().await.unwrap() { + assert!(bytes.len() <= 64 * 1024); + assert!(bytes.iter().all(|byte| *byte == 0x6b)); + total += bytes.len(); + } + assert_eq!(total, 536_870_912); + service.delete_file(&file.id).await.unwrap(); + let mut upload = service + .begin_file_upload("oversized.bin", "application/octet-stream") + .unwrap(); + for _ in 0..8192 { + upload.write(&chunk).await.unwrap(); + } + assert!(upload.write(&[0]).await.is_err()); + drop(upload); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if std::fs::read_dir(service.files.path()).unwrap().count() == 0 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); +} + +#[tokio::test] +async fn changed_download_fails_before_emitting_the_final_declared_bytes() { + use futures::StreamExt as _; + use std::io::Write as _; + let service = service().await; + let file = service + .upload_file( + "changed.bin", + "application/octet-stream", + "assistants", + vec![1; 1024 * 1024], + ) + .await + .unwrap(); + let mut download = service.download_file(&file.id).await.unwrap(); + let mut total = download.next().await.unwrap().unwrap().len(); + std::fs::OpenOptions::new() + .append(true) + .open(service.files.path().join(&file.id)) + .unwrap() + .write_all(&[2]) + .unwrap(); + let mut failed = false; + while let Some(chunk) = download.next().await { + if let Ok(bytes) = chunk { + total += bytes.len(); + } else { + failed = true; + break; + } + } + assert!(failed); + assert!( + total < 1024 * 1024, + "Content-Length must not be satisfied before integrity verification" + ); +} + +#[tokio::test] +async fn file_expiration_during_embedding_prevents_publication_and_search() { + let (service, state, task, pool, _) = embedding_service().await; + let store = service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + let file = service + .upload_file("expire.txt", "text/plain", "assistants", b"coral reefs".to_vec()) + .await + .unwrap(); + *state.mode.lock().unwrap() = ProviderMode::Pause; + let worker = service.service.clone(); + let store_id = store.id.clone(); + let file_id = file.id.clone(); + let ingestion = tokio::spawn(async move { + worker + .attach_file( + &store_id, + AttachFileRequest { + file_id, + ..Default::default() + }, + ) + .await + }); + state.started.notified().await; + sqlx::query("UPDATE file_search_files SET expires_at = 1 WHERE id = $1") + .bind(&file.id) + .execute(pool.as_ref()) + .await + .unwrap(); + state.resume.notify_one(); + assert_eq!(ingestion.await.unwrap().unwrap_err().status_code(), 404); + *state.mode.lock().unwrap() = ProviderMode::Good; + assert!( + service + .search(&[store.id], &query("coral")) + .await + .unwrap() + .data + .is_empty() + ); + task.abort(); +} + +#[tokio::test] +#[ignore = "requires isolated TEST_POSTGRES_URL"] +async fn postgres_file_expiring_while_publication_waits_for_store_lock_is_not_attached() { + let pool = create_pool_with_schema(Some(&std::env::var("TEST_POSTGRES_URL").unwrap())) + .await + .unwrap(); + let files = tempfile::tempdir().unwrap(); + let service = FileSearchService::new(pool.clone(), Arc::new(reqwest::Client::new()), file_config(&files)).unwrap(); + let file = service + .upload_file( + "lock-expiry.txt", + "text/plain", + "assistants", + b"coral lock expiration".to_vec(), + ) + .await + .unwrap(); + let store = service + .create_vector_store(CreateVectorStoreRequest::default()) + .await + .unwrap(); + let mut blocker = pool.begin().await.unwrap(); + let blocker_pid: i64 = sqlx::query_scalar("SELECT CAST(pg_backend_pid() AS BIGINT)") + .fetch_one(&mut *blocker) + .await + .unwrap(); + sqlx::query("UPDATE file_search_stores SET embedding_dimensions = embedding_dimensions WHERE id = $1") + .bind(&store.id) + .execute(&mut *blocker) + .await + .unwrap(); + sqlx::query("UPDATE file_search_files SET expires_at = CAST(FLOOR(EXTRACT(EPOCH FROM clock_timestamp())) AS BIGINT) + 2 WHERE id = $1").bind(&file.id).execute(pool.as_ref()).await.unwrap(); + let worker = service.clone(); + let file_id = file.id.clone(); + let store_id = store.id.clone(); + let publication = tokio::spawn(async move { + worker + .attach_file( + &store_id, + AttachFileRequest { + file_id, + ..Default::default() + }, + ) + .await + }); + tokio::time::timeout(std::time::Duration::from_secs(3), async { + loop { + let waiting_on_store: bool = sqlx::query_scalar("SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE wait_event_type = 'Lock' AND CAST($1 AS INTEGER) = ANY(pg_blocking_pids(pid)))") + .bind(blocker_pid).fetch_one(pool.as_ref()).await.unwrap(); + if waiting_on_store { break; } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }).await.expect("publication must reach the contended store row before expiry"); + let still_live: bool = sqlx::query_scalar( + "SELECT expires_at > EXTRACT(EPOCH FROM clock_timestamp()) FROM file_search_files WHERE id = $1", + ) + .bind(&file.id) + .fetch_one(pool.as_ref()) + .await + .unwrap(); + assert!( + still_live, + "publication must be observed blocked while its source file is still live" + ); + tokio::time::timeout(std::time::Duration::from_secs(3), async { + loop { + let expired: bool = sqlx::query_scalar( + "SELECT expires_at <= EXTRACT(EPOCH FROM clock_timestamp()) FROM file_search_files WHERE id = $1", + ) + .bind(&file.id) + .fetch_one(pool.as_ref()) + .await + .unwrap(); + if expired { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + blocker.commit().await.unwrap(); + let result = publication.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(); + // Clean up before asserting, including on the red run. + service.delete_vector_store(&store.id).await.unwrap(); + service.delete_file(&file.id).await.unwrap(); + assert!( + matches!(result, Err(FileSearchError::NotFound(_))), + "publication after file expiration must fail: {result:?}" + ); + assert_eq!(chunks, 0); +} diff --git a/crates/agentic-server/Cargo.toml b/crates/agentic-server/Cargo.toml index c0d0d64d..d1ab0a95 100644 --- a/crates/agentic-server/Cargo.toml +++ b/crates/agentic-server/Cargo.toml @@ -23,6 +23,7 @@ either.workspace = true futures.workspace = true http.workspace = true jsonwebtoken.workspace = true +multer = "3.1" reqwest = { workspace = true, default-features = false, features = ["rustls-tls"] } serde.workspace = true serde_json.workspace = true diff --git a/crates/agentic-server/src/handler/http/file_search.rs b/crates/agentic-server/src/handler/http/file_search.rs index e679cdcd..becc88da 100644 --- a/crates/agentic-server/src/handler/http/file_search.rs +++ b/crates/agentic-server/src/handler/http/file_search.rs @@ -4,15 +4,18 @@ 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, FileSearchError, ListParams, SearchRequest, + AttachFileRequest, CreateVectorStoreRequest, FileExpirationAnchor, FileExpiresAfter, FileSearchError, ListParams, + SearchRequest, }; -use axum::extract::multipart::MultipartRejection; +#[path = "multipart_limits.rs"] +mod multipart_limits; use axum::extract::rejection::{JsonRejection, QueryRejection}; -use axum::extract::{DefaultBodyLimit, Multipart, Path, Query, State}; +use axum::extract::{DefaultBodyLimit, Path, Query, State}; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::{Json, Router}; use http::{StatusCode, header}; +use multipart_limits::BoundedMultipart; use serde::Serialize; use crate::app::AppState; @@ -93,69 +96,112 @@ fn query(params: Result, QueryRejection>) -> Result, - multipart: Result, + multipart: Result, ) -> Response { let search = match service(&state) { Ok(service) => service, Err(error) => return *error, }; - let Ok(mut multipart) = multipart else { - return invalid("Expected multipart/form-data with file and purpose fields"); + let mut multipart = match multipart { + Ok(BoundedMultipart(multipart)) => multipart, + Err(rejection) => return rejection, }; let mut purpose = None; + let mut anchor = None; + let mut seconds = None; let mut uploaded = None; loop { - let field = match multipart.next_field().await { + let mut field = match multipart.next_field().await { Ok(Some(field)) => field, Ok(None) => break, Err(failure) => return multipart_error(&failure), }; match field.name() { - Some("purpose") if purpose.is_none() => { - let value = match field.text().await { - Ok(value) => value, - Err(failure) => return multipart_error(&failure), + Some("purpose" | "expires_after[anchor]" | "expires_after[seconds]") => { + let name = field.name().unwrap_or_default().to_owned(); + let slot = match name.as_str() { + "purpose" => &mut purpose, + "expires_after[anchor]" => &mut anchor, + _ => &mut seconds, }; - if value.len() > 32 { - return invalid("Invalid file purpose"); + if slot.is_some() { + return invalid("Duplicate multipart field"); } - purpose = Some(value); + let mut value = Vec::new(); + loop { + match field.chunk().await { + Ok(Some(chunk)) => { + if value.len().saturating_add(chunk.len()) > 32 { + return invalid("Multipart scalar exceeds 32 bytes"); + } + value.extend_from_slice(&chunk); + } + Ok(None) => break, + Err(failure) => return multipart_error(&failure), + } + } + let Ok(value) = String::from_utf8(value) else { + return invalid("Multipart scalar must be UTF-8"); + }; + *slot = Some(value); } Some("file") if uploaded.is_none() => { - let Some(filename) = field.file_name().map(str::to_owned) else { + let Some(filename) = field.file_name() else { return invalid("File field requires a filename"); }; - let content_type = field.content_type().unwrap_or("application/octet-stream").to_owned(); - let mut field = field; - let mut bytes = Vec::new(); + let mut upload = match search + .begin_file_upload(filename, field.content_type().unwrap_or("application/octet-stream")) + { + Ok(upload) => upload, + Err(failure) => return error(failure), + }; + let mut size = 0usize; loop { match field.chunk().await { Ok(Some(chunk)) => { - if bytes.len().saturating_add(chunk.len()) > MAX_FILE_BYTES { + size = size.saturating_add(chunk.len()); + if size > MAX_FILE_BYTES { return executor_error_response(ExecutorError::PayloadTooLarge( - "File exceeds 20 MiB".into(), + "File exceeds 512 MiB".into(), )); } - bytes.extend_from_slice(&chunk); + if let Err(failure) = upload.write(&chunk).await { + return error(failure); + } } Ok(None) => break, Err(failure) => return multipart_error(&failure), } } - uploaded = Some((filename, content_type, bytes)); + uploaded = Some(upload); } - _ => return invalid("Expected exactly one file and one purpose field"), + _ => return invalid("Expected one file, one purpose, and optional expires_after fields"), } } - let (Some(purpose), Some((filename, content_type, bytes))) = (purpose, uploaded) else { + let (Some(purpose), Some(upload)) = (purpose, uploaded) else { return invalid("Both file and purpose are required"); }; - result(search.upload_file(&filename, &content_type, &purpose, bytes).await) + let expires = match (anchor, seconds) { + (None, None) => None, + (Some(anchor), Some(seconds)) if anchor == "created_at" => { + let Ok(seconds) = seconds.parse::() else { + return invalid("expires_after.seconds must be an integer"); + }; + Some(FileExpiresAfter { + anchor: FileExpirationAnchor::CreatedAt, + seconds, + }) + } + _ => return invalid("expires_after requires anchor=created_at and seconds"), + }; + result(upload.finish(&purpose, expires).await) } fn multipart_error(failure: &axum::extract::multipart::MultipartError) -> Response { - if failure.status() == StatusCode::PAYLOAD_TOO_LARGE { - executor_error_response(ExecutorError::PayloadTooLarge("Upload exceeds 20 MiB".into())) + if multipart_limits::is_framing_limit(failure) { + executor_error_response(ExecutorError::PayloadTooLarge("Multipart framing exceeds 8 KiB".into())) + } else if failure.status() == StatusCode::PAYLOAD_TOO_LARGE { + executor_error_response(ExecutorError::PayloadTooLarge("Upload exceeds 512 MiB".into())) } else { invalid("Malformed multipart upload") } @@ -163,7 +209,7 @@ fn multipart_error(failure: &axum::extract::multipart::MultipartError) -> Respon #[cfg_attr(feature = "openapi", utoipa::path( get, path = "/v1/files", - params(("limit" = Option, Query, description = "Page size, 1 to 100"), ("after" = Option, Query), ("before" = Option, Query), ("order" = Option, Query)), + params(("limit" = Option, Query, description = "Page size, 1 to 10000; defaults to 10000"), ("purpose" = Option, Query), ("after" = Option, Query), ("before" = Option, Query), ("order" = 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", ))] @@ -207,13 +253,14 @@ pub(crate) async fn file_content(State(state): State, Path(id): Path service, Err(error) => return *error, }; - match search.file_content(&id).await { - Ok(bytes) => ( + match search.download_file(&id).await { + Ok(download) => ( [ - (header::CONTENT_TYPE, "application/octet-stream"), - (header::CONTENT_DISPOSITION, "attachment"), + (header::CONTENT_TYPE, "application/octet-stream".to_owned()), + (header::CONTENT_DISPOSITION, "attachment".to_owned()), + (header::CONTENT_LENGTH, download.bytes.to_string()), ], - bytes, + axum::body::Body::from_stream(download), ) .into_response(), Err(failure) => error(failure), diff --git a/crates/agentic-server/src/handler/http/multipart_limits.rs b/crates/agentic-server/src/handler/http/multipart_limits.rs new file mode 100644 index 00000000..af02d35f --- /dev/null +++ b/crates/agentic-server/src/handler/http/multipart_limits.rs @@ -0,0 +1,322 @@ +//! Bound framing before Multer buffers it, and prevent its eager stream drain +//! from collecting an always-ready upload. Multer still owns field parsing. +use std::{ + pin::Pin, + task::{Context, Poll}, +}; + +use axum::{ + RequestExt, + body::Body, + extract::{FromRequest, Multipart, Request}, + response::Response, +}; +use bytes::Bytes; +use futures::Stream; + +const FRAMING_LIMIT: usize = 8 * 1024; +const DELIVERY_LIMIT: usize = 64 * 1024; + +pub(crate) struct BoundedMultipart(pub Multipart); + +impl FromRequest for BoundedMultipart { + type Rejection = Response; + + async fn from_request(request: Request, state: &S) -> Result { + let boundary = request + .headers() + .get(http::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .filter(|value| value.len() <= 1024) + .and_then(|value| multer::parse_boundary(value).ok()) + .filter(|boundary| { + !boundary.is_empty() && boundary.len() <= 70 && !boundary.bytes().any(|byte| byte.is_ascii_control()) + }) + .ok_or_else(|| super::invalid("Expected multipart/form-data with a boundary of 1 to 70 bytes"))?; + let (parts, body) = request.with_limited_body().into_parts(); + let body = Body::from_stream(FramedBody::new(body.into_data_stream(), &boundary)); + Multipart::from_request(Request::from_parts(parts, body), state) + .await + .map(Self) + .map_err(|_| super::invalid("Expected multipart/form-data with file and purpose fields")) + } +} + +#[derive(Debug, thiserror::Error)] +enum FramingError { + #[error("multipart framing exceeds 8 KiB")] + Limit, + #[error("malformed multipart boundary suffix")] + Boundary, +} + +pub(super) fn is_framing_limit(error: &(dyn std::error::Error + 'static)) -> bool { + let mut source = Some(error); + while let Some(error) = source { + if matches!(error.downcast_ref::(), Some(FramingError::Limit)) { + return true; + } + source = error.source(); + } + false +} + +struct FramedBody { + stream: S, + pending: Bytes, + framing: Framing, + yield_before_poll: bool, + ended: bool, +} + +impl FramedBody { + fn new(stream: S, boundary: &str) -> Self { + Self { + stream, + pending: Bytes::new(), + framing: Framing::new(boundary), + yield_before_poll: false, + ended: false, + } + } +} + +impl> + Unpin> Stream for FramedBody { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.ended { + return Poll::Ready(None); + } + if self.yield_before_poll { + self.yield_before_poll = false; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + if self.pending.is_empty() { + match futures::ready!(Pin::new(&mut self.stream).poll_next(cx)) { + Some(Ok(bytes)) => self.pending = bytes, + other => { + self.ended = true; + return Poll::Ready(other); + } + } + } + let len = self.pending.len().min(DELIVERY_LIMIT); + let bytes = self.pending.split_to(len); + if let Err(error) = self.framing.check(&bytes) { + self.ended = true; + return Poll::Ready(Some(Err(axum::Error::new(error)))); + } + // Multer drains until Pending before examining its buffer. A deliberate + // yield after each bounded chunk also covers in-memory/always-ready bodies. + self.yield_before_poll = true; + Poll::Ready(Some(Ok(bytes))) + } +} + +#[derive(Clone, Copy)] +enum Stage { + Preamble, + Boundary(Tail), + Headers, + Data, + Done, +} +#[derive(Clone, Copy)] +enum Tail { + Start, + Hyphen, + Padding, + Lf, +} + +struct Framing { + stage: Stage, + count: usize, + initial: Matcher, + delimiter: Matcher, + header_end: u32, +} + +impl Framing { + fn new(boundary: &str) -> Self { + Self { + stage: Stage::Preamble, + count: 0, + initial: Matcher::new(format!("--{boundary}")), + delimiter: Matcher::new(format!("\r\n--{boundary}")), + header_end: 0, + } + } + + fn check(&mut self, bytes: &[u8]) -> Result<(), FramingError> { + for &byte in bytes { + if !matches!(self.stage, Stage::Data | Stage::Done) { + self.count += 1; + if self.count > FRAMING_LIMIT { + return Err(FramingError::Limit); + } + } + match self.stage { + Stage::Preamble if self.initial.accept(byte) => self.boundary(), + Stage::Preamble | Stage::Done => (), + Stage::Data => { + if self.delimiter.accept(byte) { + self.boundary(); + } + } + Stage::Headers => { + self.header_end = (self.header_end << 8) | u32::from(byte); + if self.header_end == 0x0d0a_0d0a { + self.stage = Stage::Data; + self.delimiter.matched = 0; + } + } + Stage::Boundary(tail) => self.tail(tail, byte)?, + } + } + Ok(()) + } + + fn boundary(&mut self) { + self.stage = Stage::Boundary(Tail::Start); + self.count = 0; + } + + fn tail(&mut self, tail: Tail, byte: u8) -> Result<(), FramingError> { + self.stage = match (tail, byte) { + (Tail::Start, b'-') => Stage::Boundary(Tail::Hyphen), + (Tail::Hyphen, b'-') => Stage::Done, + (Tail::Start | Tail::Padding, b' ' | b'\t') => Stage::Boundary(Tail::Padding), + (Tail::Start | Tail::Padding, b'\r') => Stage::Boundary(Tail::Lf), + (Tail::Lf, b'\n') => { + self.count = 0; + self.header_end = 0; + Stage::Headers + } + _ => return Err(FramingError::Boundary), + }; + Ok(()) + } +} + +/// Incremental KMP matching keeps only a bounded prefix length across chunks. +struct Matcher { + pattern: Vec, + failure: Vec, + matched: usize, +} +impl Matcher { + fn new(pattern: String) -> Self { + let pattern = pattern.into_bytes(); + let mut failure = vec![0; pattern.len()]; + let mut matched = 0; + for index in 1..pattern.len() { + while matched > 0 && pattern[index] != pattern[matched] { + matched = failure[matched - 1]; + } + if pattern[index] == pattern[matched] { + matched += 1; + } + failure[index] = matched; + } + Self { + pattern, + failure, + matched: 0, + } + } + + fn accept(&mut self, byte: u8) -> bool { + while self.matched > 0 && byte != self.pattern[self.matched] { + self.matched = self.failure[self.matched - 1]; + } + if byte == self.pattern[self.matched] { + self.matched += 1; + } + if self.matched == self.pattern.len() { + self.matched = self.failure[self.matched - 1]; + true + } else { + false + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::{StreamExt as _, stream}; + + #[tokio::test] + async fn split_delimiters_padding_and_binary_prefixes_preserve_file_bytes() { + let boundary = "abababab"; + let mut payload = (0..=255u8).collect::>(); + payload.extend_from_slice(b"\r\n--abababax\r\r\n--abababa\x00--abababab\r\n--ababab"); + let mut wire = b"preamble\r\n--abababab \t\r\nContent-Disposition: form-data; name=\"file\"; filename=\"bytes.bin\"\r\n\r\n".to_vec(); + wire.extend_from_slice(&payload); + wire.extend_from_slice(b"\r\n--abababab\t \r\nContent-Disposition: form-data; name=\"purpose\"\r\n\r\nvision\r\n--abababab--\r\nepilogue"); + for split in [1, 2, 3, 7, 16, 64 * 1024] { + let chunks: Vec<_> = wire + .chunks(split) + .map(|bytes| Ok(Bytes::copy_from_slice(bytes))) + .collect(); + let mut multipart = multer::Multipart::new(FramedBody::new(stream::iter(chunks), boundary), boundary); + let field = multipart.next_field().await.unwrap().unwrap(); + assert_eq!(field.name(), Some("file")); + assert_eq!(field.bytes().await.unwrap().as_ref(), payload); + let field = multipart.next_field().await.unwrap().unwrap(); + assert_eq!(field.name(), Some("purpose")); + assert_eq!(field.text().await.unwrap(), "vision"); + assert!(multipart.next_field().await.unwrap().is_none()); + } + } + + #[tokio::test] + async fn always_ready_input_cannot_be_drained_ahead_of_field_consumption() { + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + let consumed = Arc::new(AtomicUsize::new(0)); + let counter = consumed.clone(); + let header = Bytes::from_static( + b"--upload\r\nContent-Disposition: form-data; name=\"file\"; filename=\"ready.bin\"\r\n\r\n", + ); + let input = std::iter::once(header) + .chain((0..16).map(|_| Bytes::from(vec![b'z'; DELIVERY_LIMIT]))) + .chain(std::iter::once(Bytes::from_static(b"\r\n--upload--\r\n"))); + let input = stream::iter(input.map(Ok)).inspect(move |_| { + counter.fetch_add(1, Ordering::SeqCst); + }); + let mut multipart = multer::Multipart::new(FramedBody::new(input, "upload"), "upload"); + let mut field = multipart.next_field().await.unwrap().unwrap(); + assert!( + consumed.load(Ordering::SeqCst) <= 2, + "Multer must not eagerly buffer an always-ready upload" + ); + let mut total = 0; + while let Some(bytes) = field.chunk().await.unwrap() { + assert!(bytes.len() <= 2 * DELIVERY_LIMIT); + assert!(bytes.iter().all(|byte| *byte == b'z')); + total += bytes.len(); + } + assert_eq!(total, 16 * DELIVERY_LIMIT); + } + + #[test] + fn large_input_frames_are_split_and_yield_between_deliveries() { + let mut bytes = b"--upload\r\nContent-Disposition: form-data; name=\"file\"\r\n\r\n".to_vec(); + bytes.resize(3 * DELIVERY_LIMIT, b'z'); + let mut body = FramedBody::new(stream::iter([Ok(Bytes::from(bytes))]), "upload"); + let mut cx = Context::from_waker(futures::task::noop_waker_ref()); + for _ in 0..3 { + let Poll::Ready(Some(Ok(chunk))) = Pin::new(&mut body).poll_next(&mut cx) else { + panic!("expected data"); + }; + assert_eq!(chunk.len(), DELIVERY_LIMIT); + assert!(Pin::new(&mut body).poll_next(&mut cx).is_pending()); + } + } +} diff --git a/crates/agentic-server/src/openapi.rs b/crates/agentic-server/src/openapi.rs index 2d73f7ee..65bbfc0c 100644 --- a/crates/agentic-server/src/openapi.rs +++ b/crates/agentic-server/src/openapi.rs @@ -1036,4 +1036,5 @@ pub struct FileUploadRequest { #[schema(format = Binary)] pub file: String, pub purpose: String, + pub expires_after: Option, } diff --git a/crates/agentic-server/tests/file_search_http_test.rs b/crates/agentic-server/tests/file_search_http_test.rs index 446c21e7..0ffd9681 100644 --- a/crates/agentic-server/tests/file_search_http_test.rs +++ b/crates/agentic-server/tests/file_search_http_test.rs @@ -298,3 +298,231 @@ async fn files_api_preserves_binary_uploads_without_requiring_searchable_content .error_for_status() .unwrap(); } + +#[tokio::test] +async fn files_contract_purposes_expiration_pagination_and_delete() { + let server = gateway().await; + let client = reqwest::Client::new(); + for purpose in ["assistants", "batch", "fine-tune", "vision", "user_data", "evals"] { + let body = multipart("contract.bin", "bytes").replace("\r\nassistants\r\n", &format!("\r\n{purpose}\r\n")); + let file = api_json( + client + .post(format!("{}/v1/files", server.url)) + .header("content-type", "multipart/form-data; boundary=upload") + .body(body), + ) + .await; + if purpose == "batch" { + assert_eq!( + file["expires_at"].as_i64().unwrap() - file["created_at"].as_i64().unwrap(), + 2_592_000 + ); + } + let list = api_json(client.get(format!("{}/v1/files?purpose={purpose}&limit=10000", server.url))).await; + assert_eq!(list["data"].as_array().unwrap().len(), 1); + assert_eq!(list["data"][0]["id"], file["id"]); + let deleted = + api_json(client.delete(format!("{}/v1/files/{}", server.url, file["id"].as_str().unwrap()))).await; + assert_eq!(deleted["object"], "file"); + } + let body = multipart("expiry.bin", "bytes").replace("--upload--\r\n", "--upload\r\nContent-Disposition: form-data; name=\"expires_after[anchor]\"\r\n\r\ncreated_at\r\n--upload\r\nContent-Disposition: form-data; name=\"expires_after[seconds]\"\r\n\r\n3600\r\n--upload--\r\n"); + let file = api_json( + client + .post(format!("{}/v1/files", server.url)) + .header("content-type", "multipart/form-data; boundary=upload") + .body(body), + ) + .await; + assert_eq!( + file["expires_at"].as_i64().unwrap() - file["created_at"].as_i64().unwrap(), + 3600 + ); +} + +#[tokio::test] +async fn files_stream_above_old_limit_and_accept_empty_content() { + let server = gateway().await; + let client = reqwest::Client::new(); + for len in [0, 21 * 1024 * 1024] { + let bytes = "z".repeat(len); + let file = api_json( + client + .post(format!("{}/v1/files", server.url)) + .header("content-type", "multipart/form-data; boundary=upload") + .body(multipart("large.bin", &bytes)), + ) + .await; + assert_eq!(file["bytes"], len); + let response = client + .get(format!( + "{}/v1/files/{}/content", + server.url, + file["id"].as_str().unwrap() + )) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.content_length(), Some(len as u64)); + assert_eq!(response.bytes().await.unwrap().as_ref(), bytes.as_bytes()); + } +} + +#[tokio::test] +async fn file_first_upload_validates_trailing_fields_before_publication() { + let server = gateway().await; + let client = reqwest::Client::new(); + let prefix = format!( + "--upload\r\nContent-Disposition: form-data; name=\"file\"; filename=\"first.bin\"\r\n\r\n{}\r\n", + "a".repeat(1024 * 1024) + ); + for tail in [ + "--upload\r\nContent-Disposition: form-data; name=\"purpose\"\r\n\r\nassistants\r\n--upload\r\nContent-Disposition: form-data; name=\"unknown\"\r\n\r\nbad\r\n--upload--\r\n".to_owned(), + "--upload\r\nContent-Disposition: form-data; name=\"purpose\"\r\n\r\nassistants\r\n--upload\r\nContent-Disposition: form-data; name=\"purpose\"\r\n\r\nassistants\r\n--upload--\r\n".to_owned(), + format!("--upload\r\nContent-Disposition: form-data; name=\"purpose\"\r\n\r\n{}\r\n--upload--\r\n", "x".repeat(1024 * 1024)), + "--upload\r\nContent-Disposition: form-data; name=\"expires_after[seconds]\"\r\n\r\n3600\r\n--upload--\r\n".to_owned(), + ] { + let response = client.post(format!("{}/v1/files", server.url)).header("content-type", "multipart/form-data; boundary=upload").body(format!("{prefix}{tail}")).send().await.unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + } + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if std::fs::read_dir(server.files.path()).unwrap().count() == 0 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!( + api_json(client.get(format!("{}/v1/files", server.url))).await["data"] + .as_array() + .unwrap() + .is_empty() + ); + let tail = "--upload\r\nContent-Disposition: form-data; name=\"purpose\"\r\n\r\nvision\r\n--upload--\r\n"; + let file = api_json( + client + .post(format!("{}/v1/files", server.url)) + .header("content-type", "multipart/form-data; boundary=upload") + .body(format!("{prefix}{tail}")), + ) + .await; + assert_eq!(file["bytes"], 1024 * 1024); +} + +#[tokio::test] +async fn disconnected_multipart_upload_cleans_staging() { + use tokio::io::AsyncWriteExt as _; + let server = gateway().await; + let address = server.url.strip_prefix("http://").unwrap(); + let mut socket = tokio::net::TcpStream::connect(address).await.unwrap(); + let headers = format!( + "POST /v1/files HTTP/1.1\r\nHost: {address}\r\nContent-Type: multipart/form-data; boundary=upload\r\nContent-Length: 2097152\r\n\r\n--upload\r\nContent-Disposition: form-data; name=\"file\"; filename=\"disconnect.bin\"\r\n\r\n" + ); + socket.write_all(headers.as_bytes()).await.unwrap(); + socket.write_all(&vec![7; 1024 * 1024]).await.unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if std::fs::read_dir(server.files.path()).unwrap().count() > 0 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + drop(socket); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + if std::fs::read_dir(server.files.path()).unwrap().count() == 0 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let response = api_json(reqwest::Client::new().get(format!("{}/v1/files", server.url))).await; + assert_eq!(response["data"], json!([])); +} + +#[tokio::test] +async fn files_default_page_and_expiration_bounds_match_contract() { + let server = gateway().await; + let client = reqwest::Client::new(); + for index in 0..21 { + api_json( + client + .post(format!("{}/v1/files", server.url)) + .header("content-type", "multipart/form-data; boundary=upload") + .body(multipart(&format!("{index}.txt"), "x")), + ) + .await; + } + let page = api_json(client.get(format!("{}/v1/files", server.url))).await; + assert_eq!(page["data"].as_array().unwrap().len(), 21); + assert_eq!(page["has_more"], false); + for seconds in ["3599", "2592001", "not-a-number"] { + let extra = format!( + "--upload\r\nContent-Disposition: form-data; name=\"expires_after[anchor]\"\r\n\r\ncreated_at\r\n--upload\r\nContent-Disposition: form-data; name=\"expires_after[seconds]\"\r\n\r\n{seconds}\r\n--upload--\r\n" + ); + let body = multipart("invalid.txt", "x").replace("--upload--\r\n", &extra); + let response = client + .post(format!("{}/v1/files", server.url)) + .header("content-type", "multipart/form-data; boundary=upload") + .body(body) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + } + for limit in [0, 10001] { + let response = client + .get(format!("{}/v1/files?limit={limit}", server.url)) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + } +} + +async fn oversized_framing_is_rejected_before_body_finishes(prefix: &[u8]) { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + let server = gateway().await; + let address = server.url.strip_prefix("http://").unwrap(); + let mut socket = tokio::net::TcpStream::connect(address).await.unwrap(); + let headers = format!( + "POST /v1/files HTTP/1.1\r\nHost: {address}\r\nContent-Type: multipart/form-data; boundary=upload\r\nContent-Length: 536870912\r\n\r\n" + ); + socket.write_all(headers.as_bytes()).await.unwrap(); + socket.write_all(prefix).await.unwrap(); + // Deliberately leave the declared 512 MiB body unfinished. Framing must + // fail after this small prefix rather than await/buffer the remaining body. + socket.write_all(&vec![b'x'; 16 * 1024]).await.unwrap(); + let mut response = [0; 1024]; + let length = tokio::time::timeout(std::time::Duration::from_secs(2), socket.read(&mut response)) + .await + .expect("multipart framing must be rejected before the body finishes") + .unwrap(); + assert!( + std::str::from_utf8(&response[..length]) + .unwrap() + .starts_with("HTTP/1.1 413") + ); + assert_eq!(std::fs::read_dir(server.files.path()).unwrap().count(), 0); +} + +#[tokio::test] +async fn oversized_multipart_preamble_is_rejected_early() { + oversized_framing_is_rejected_before_body_finishes(b"unbounded preamble ").await; +} + +#[tokio::test] +async fn oversized_multipart_part_headers_are_rejected_early() { + oversized_framing_is_rejected_before_body_finishes( + b"--upload\r\nContent-Disposition: form-data; name=\"file\"; filename=\"", + ) + .await; +} diff --git a/docs/api/file-search.md b/docs/api/file-search.md index 0b413106..964ab370 100644 --- a/docs/api/file-search.md +++ b/docs/api/file-search.md @@ -291,8 +291,21 @@ failed or cancelled ingestion publishes no partial chunks. Creating a store with the uploaded file ID, filename, attributes, and chunking strategy. The Files API accepts binary uploads independently of search ingestion. Uploads -are limited to 20 MiB and must use purpose `assistants` or `user_data`. Attaching a -file to a vector store validates its format: UTF-8 text, Markdown, CSV, JSON, source +stream to disk with a 512 MiB limit (including zero-byte files). Accepted purposes are +`assistants`, `batch`, `fine-tune`, `vision`, `user_data`, and `evals`. The multipart +fields may occur in any order. Preambles, individual part-header blocks, and boundary +padding are each limited to 8 KiB; boundary values are limited to 70 bytes. A bounded +HTTP adapter enforces these limits before multipart parsing and yields between +64 KiB deliveries, including when the incoming body is immediately available. Optional `expires_after[anchor]=created_at` and +`expires_after[seconds]` (3600 through 2592000) set expiration. Batch-purpose files +expire after 30 days by default; other purposes persist by default. Responses +include `expires_at` when set. OpenAI Python SDK 3.13.0 accepts `evals` in upload +parameters but omits it from its response-purpose literal; strict response validation +in that SDK version rejects this purpose even though the HTTP contract accepts it. + +Attaching a +file to a vector store retains a separate 20 MiB input limit and validates its format: +UTF-8 text, Markdown, CSV, JSON, source files, and other supported text formats work in the default build. The default chunk size is 800 tokens with a 400-token overlap. Override it with: @@ -389,17 +402,32 @@ parts, output items, and the terminal response. | Retrieve/detach a store file | `GET` / `DELETE /v1/vector_stores/{store_id}/files/{file_id}` | | Search a store | `POST /v1/vector_stores/{store_id}/search` | -Lists accept `limit`, `after`, `before`, and `order`. Detaching a file preserves the +Lists accept `limit`, `after`, `before`, and `order`. Files lists additionally accept +`purpose`, with a limit of 1–10000 and default 10000; vector store lists retain their +1–100 limit and default 20. Files deletion returns `object: "file"`. Detaching a file preserves the original upload. Deleting an upload removes its metadata, attachments, and chunks from all stores, then removes its local file bytes. Deleting a vector store preserves uploaded files. Uploads publish complete, synced files before committing metadata. Failures and cancellation before commit clean up the upload. Filesystem and SQL commits are -separate: a process crash, failed unlink, or uncertain database commit can leave -unreferenced files. Automatic orphan-file cleanup is not included. A missing or -damaged file referenced by metadata returns a storage error instead of a partial -download. Uploads stored inline by an earlier draft remain readable. +separate: a process crash or uncertain upload commit can leave unreferenced files. +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. +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 +must share the same files storage directory/mount. + +Downloads stream through a bounded buffer with verified `Content-Length`. Missing or +already damaged files return a storage error before headers; integrity or I/O errors +during streaming terminate the response. Download visibility is checked when opening; +an already-open download can finish if its metadata is subsequently deleted or expires. +Search revalidates visible attachments after rewriting, embedding, and reranking. +Uploads stored inline by an earlier draft remain readable. The routes use the gateway's configured authentication policy. Retrieval uses bounded exact SQL or the configured pgvector backend and returns explicit capacity errors when limits are exceeded.