Skip to content

feat: add durable vector store batches and owned workers #283

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
franciscojavierarceo wants to merge 3 commits into
base: codex/vector-store-lifecycle
Choose a base branch
from
Draft

feat: add durable vector store batches and owned workers #283

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
fbe08ca
feat: add durable vector store batches and owned workers
franciscojavierarceo Sep 11, 2026
1109735
fix: preserve batch creation and worker completion progress
franciscojavierarceo Sep 11, 2026
7affaad
chore: sync file search stack layer 283 with main
franciscojavierarceo Sep 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions crates/agentic-server-core/migrations/0008_vector_store_batches.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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);
24 changes: 22 additions & 2 deletions crates/agentic-server-core/src/storage/file_search.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand All @@ -20,6 +22,8 @@ const MAX_CORPUS_CHUNKS: i64 = 10_000;
pub(crate) struct FileSearchStorage {
pool: Arc<DbPool>,
pgvector: Option<super::pgvector::PgvectorStorage>,
#[cfg(test)]
pub(crate) batch_test_hooks: Option<Arc<batches::BatchTestHooks>>,
}

#[derive(FromRow)]
Expand Down Expand Up @@ -88,15 +92,24 @@ impl Collection {
impl FileSearchStorage {
#[cfg(test)]
pub(crate) fn new(pool: Arc<DbPool>) -> Self {
Self { pool, pgvector: None }
Self {
pool,
pgvector: None,
batch_test_hooks: None,
}
}

pub(crate) fn with_backend(
pool: Arc<DbPool>,
backend: &crate::types::file_search::FileSearchBackend,
) -> Result<Self, FileSearchError> {
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<usize> {
Expand Down Expand Up @@ -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}",
Expand Down Expand Up @@ -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")
Expand Down
61 changes: 56 additions & 5 deletions crates/agentic-server-core/src/storage/schema.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ use crate::config::DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS;
type DbResult<T> = Result<T, sqlx::Error>;

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
Expand Down Expand Up @@ -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'), \
Expand Down Expand Up @@ -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)'), \
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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'), \
Expand All @@ -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) \
Expand Down Expand Up @@ -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?;
Expand Down Expand Up @@ -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();
}
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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)
Expand Down
Loading
Loading