From c8192fed387372751a35941c3c2d50ee9bce74db Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Thu, 23 Jul 2026 01:46:28 -0400 Subject: [PATCH] fix: enable production PostgreSQL persistence Signed-off-by: Francisco Javier Arceo --- .github/workflows/rust.yml | 44 ++ .../0003_index_conversation_sequence.sql | 2 + crates/agentic-server-core/src/config.rs | 35 ++ .../src/executor/engine.rs | 17 +- .../src/executor/request.rs | 106 ++++- crates/agentic-server-core/src/proxy.rs | 1 + .../src/storage/conversation.rs | 1 + crates/agentic-server-core/src/storage/mod.rs | 4 +- .../src/storage/models/conversation.rs | 35 +- .../src/storage/models/item.rs | 62 ++- .../src/storage/models/response.rs | 4 +- .../agentic-server-core/src/storage/pool.rs | 105 ++++- .../agentic-server-core/src/storage/schema.rs | 322 ++++++++++++- .../tests/postgres_storage_integration.rs | 426 ++++++++++++++++++ .../agentic-server/benches/gateway_bench.rs | 1 + crates/agentic-server/benches/proxy_bench.rs | 1 + crates/agentic-server/src/main.rs | 192 +++++++- crates/agentic-server/tests/cli_test.rs | 16 + crates/agentic-server/tests/common/mod.rs | 1 + crates/agentic-server/tests/responses_test.rs | 41 +- docs/deploying/container.md | 41 ++ 21 files changed, 1391 insertions(+), 66 deletions(-) create mode 100644 crates/agentic-server-core/migrations/0003_index_conversation_sequence.sql create mode 100644 crates/agentic-server-core/tests/postgres_storage_integration.rs diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 9ccb342b..48c68199 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -48,3 +48,47 @@ jobs: - name: Run tests run: cargo test + + postgres: + runs-on: ubuntu-latest + services: + postgres: + image: postgres:17-alpine@sha256:742f40ea20b9ff2ff31db5458d127452988a2164df9e17441e191f3b72252193 + env: + POSTGRES_DB: agentic_api_test + POSTGRES_PASSWORD: postgres + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U postgres -d agentic_api_test" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + TEST_POSTGRES_URL: postgresql://postgres:postgres@localhost:5432/agentic_api_test?sslmode=disable + steps: + - name: Checkout code + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@e97e2d8cc328f1b50210efc529dca0028893a2d9 # v1 + with: + toolchain: stable + + - name: Cache cargo registry and build + uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v4 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + target + key: cargo-postgres-${{ runner.os }}-${{ hashFiles('Cargo.lock') }} + restore-keys: | + cargo-postgres-${{ runner.os }}- + cargo-${{ runner.os }}- + + - name: Verify PostgreSQL upgrade preservation + run: cargo test -p agentic-server-core --lib postgres_integer_widening_upgrade_preserves_existing_state -- --ignored + + - name: Verify PostgreSQL migrations and restart persistence + run: cargo test -p agentic-server-core --test postgres_storage_integration -- --ignored diff --git a/crates/agentic-server-core/migrations/0003_index_conversation_sequence.sql b/crates/agentic-server-core/migrations/0003_index_conversation_sequence.sql new file mode 100644 index 00000000..1dbedb3f --- /dev/null +++ b/crates/agentic-server-core/migrations/0003_index_conversation_sequence.sql @@ -0,0 +1,2 @@ +DROP INDEX IF EXISTS idx_items_conversation_id; +CREATE INDEX idx_items_conversation_id ON items (conversation_id, seq); diff --git a/crates/agentic-server-core/src/config.rs b/crates/agentic-server-core/src/config.rs index 46281c26..cca8bafc 100644 --- a/crates/agentic-server-core/src/config.rs +++ b/crates/agentic-server-core/src/config.rs @@ -1,7 +1,41 @@ +use std::time::Duration; + +pub const DEFAULT_POSTGRES_MAX_CONNECTIONS: u32 = 10; +pub const DEFAULT_POSTGRES_ACQUIRE_TIMEOUT_SECONDS: u64 = 30; +pub const DEFAULT_POSTGRES_IDLE_TIMEOUT_SECONDS: u64 = 600; +pub const DEFAULT_POSTGRES_LOCK_TIMEOUT_SECONDS: u64 = 5; +pub const DEFAULT_POSTGRES_MAX_LIFETIME_SECONDS: u64 = 1_800; +pub const DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS: u64 = 300; +pub const DEFAULT_POSTGRES_STATEMENT_TIMEOUT_SECONDS: u64 = 30; pub const DEFAULT_SQLITE_MAX_CONNECTIONS: u32 = 4; pub const DEFAULT_SQLITE_JOURNAL_SIZE_LIMIT_BYTES: u64 = 6_144_000; pub const DEFAULT_SQLITE_MMAP_SIZE_BYTES: u64 = 268_435_456; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PostgresConfig { + pub max_connections: u32, + pub acquire_timeout: Duration, + pub lock_timeout: Duration, + pub migration_timeout: Duration, + pub statement_timeout: Duration, + pub idle_timeout: Option, + pub max_lifetime: Option, +} + +impl Default for PostgresConfig { + fn default() -> Self { + Self { + max_connections: DEFAULT_POSTGRES_MAX_CONNECTIONS, + acquire_timeout: Duration::from_secs(DEFAULT_POSTGRES_ACQUIRE_TIMEOUT_SECONDS), + lock_timeout: Duration::from_secs(DEFAULT_POSTGRES_LOCK_TIMEOUT_SECONDS), + migration_timeout: Duration::from_secs(DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS), + statement_timeout: Duration::from_secs(DEFAULT_POSTGRES_STATEMENT_TIMEOUT_SECONDS), + idle_timeout: Some(Duration::from_secs(DEFAULT_POSTGRES_IDLE_TIMEOUT_SECONDS)), + max_lifetime: Some(Duration::from_secs(DEFAULT_POSTGRES_MAX_LIFETIME_SECONDS)), + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum SqliteTempStore { Default, @@ -50,6 +84,7 @@ pub struct Config { /// Database URL for conversation and response storage. /// `None` means stateful features are disabled; all requests are proxied. pub db_url: Option, + pub postgres: PostgresConfig, pub sqlite: SqliteConfig, } diff --git a/crates/agentic-server-core/src/executor/engine.rs b/crates/agentic-server-core/src/executor/engine.rs index fe4fa127..cc352428 100644 --- a/crates/agentic-server-core/src/executor/engine.rs +++ b/crates/agentic-server-core/src/executor/engine.rs @@ -10,7 +10,7 @@ use std::sync::Arc; use async_stream::stream; use either::Either; use tokio::sync::mpsc; -use tracing::{debug, warn}; +use tracing::debug; use super::gateway::{ GatewayCallResult, LoopDecision, append_gateway_calls_to_new_input, append_output_items_to_input, @@ -332,9 +332,7 @@ async fn run_blocking( let ch = exec_ctx.conv_handler.clone(); let rh = exec_ctx.resp_handler.clone(); - if let Err(e) = persist_if_needed(payload.clone(), ctx, ch, rh).await { - warn!("persist failed: {e}"); - } + persist_if_needed(payload.clone(), ctx, ch, rh).await?; Ok(payload) } @@ -389,12 +387,11 @@ fn run_stream(ctx: RequestContext, exec_ctx: Arc, auth: Option let terminal_chunk = stream_accumulator.terminal_response_chunk(&payload); let ch = exec_ctx.conv_handler.clone(); let rh = exec_ctx.resp_handler.clone(); - if let Err(e) = persist_if_needed(payload, ctx, ch, rh).await { - warn!("persist failed: {e}"); - } - - match terminal_chunk { - Ok(chunk) => yield chunk, + match persist_if_needed(payload, ctx, ch, rh).await { + Ok(()) => match terminal_chunk { + Ok(chunk) => yield chunk, + Err(e) => yield stream_accumulator.error_chunk(&e.to_string()), + }, Err(e) => yield stream_accumulator.error_chunk(&e.to_string()), } yield DONE_MARKER.to_string(); diff --git a/crates/agentic-server-core/src/executor/request.rs b/crates/agentic-server-core/src/executor/request.rs index 554c021a..590c96ef 100644 --- a/crates/agentic-server-core/src/executor/request.rs +++ b/crates/agentic-server-core/src/executor/request.rs @@ -4,7 +4,7 @@ use std::time::Duration; use crate::config::Config; use crate::error::Error; use crate::executor::modes::{ConversationHandler, ResponseHandler}; -use crate::storage::{ConversationStore, ResponseStore, create_pool_with_schema_and_sqlite_config}; +use crate::storage::{ConversationStore, ResponseStore, create_pool_with_schema_and_configs}; use crate::tool::{GatewayExecutor, GatewayExecutors}; use crate::types::io::InputItem; use crate::types::messages::GatewayToolMap; @@ -113,9 +113,10 @@ impl ExecutionContext { /// migration fails. pub async fn from_config(cfg: &Config) -> Result { let db_url = cfg.db_url.as_deref().unwrap_or("sqlite://./agentic_api.db"); - let pool = create_pool_with_schema_and_sqlite_config(Some(db_url), cfg.sqlite) + let database_backend = database_backend(db_url); + let pool = create_pool_with_schema_and_configs(Some(db_url), cfg.sqlite, cfg.postgres) .await - .map_err(|e| Error::Config(format!("failed to open database '{db_url}': {e}")))?; + .map_err(|error| database_open_error(database_backend, &error))?; let conv_handler = ConversationHandler::new(ConversationStore::new(pool.clone())); let resp_handler = ResponseHandler::new(ResponseStore::new(pool)); @@ -134,6 +135,60 @@ impl ExecutionContext { } } +fn database_open_error(database_backend: &str, error: &sqlx::Error) -> Error { + let category = match error { + sqlx::Error::Configuration(_) | sqlx::Error::InvalidArgument(_) => "configuration error".to_owned(), + sqlx::Error::Database(database_error) => database_error + .code() + .filter(|code| code.len() <= 5 && code.bytes().all(|byte| byte.is_ascii_alphanumeric())) + .map_or_else( + || "database error".to_owned(), + |code| format!("database error (SQLSTATE {code})"), + ), + sqlx::Error::Io(_) => "database I/O error".to_owned(), + sqlx::Error::Tls(_) => "database TLS error".to_owned(), + sqlx::Error::Protocol(_) => "database protocol error".to_owned(), + sqlx::Error::PoolTimedOut => "connection pool timeout".to_owned(), + sqlx::Error::PoolClosed => "connection pool closed".to_owned(), + sqlx::Error::WorkerCrashed => "database worker crashed".to_owned(), + sqlx::Error::Migrate(_) => "database migration error".to_owned(), + _ => "database error".to_owned(), + }; + let detail = redact_database_urls(&error.to_string()); + Error::Config(format!( + "failed to open {database_backend} database: {category}: {detail}" + )) +} + +fn redact_database_urls(message: &str) -> String { + const DATABASE_SCHEMES: [&str; 3] = ["postgresql://", "postgres://", "mysql://"]; + + let mut redacted = message.to_owned(); + for scheme in DATABASE_SCHEMES { + let mut search_from = 0; + while let Some(offset) = redacted[search_from..].find(scheme) { + let start = search_from + offset; + let end = redacted[start..] + .find(char::is_whitespace) + .map_or(redacted.len(), |length| start + length); + let replacement = format!("{scheme}[redacted]"); + redacted.replace_range(start..end, &replacement); + search_from = start + replacement.len(); + } + } + redacted +} + +fn database_backend(database_url: &str) -> &'static str { + if database_url.starts_with("postgres://") || database_url.starts_with("postgresql://") { + "PostgreSQL" + } else if database_url.starts_with("sqlite") { + "SQLite" + } else { + "configured" + } +} + /// Load the `/v1/messages` gateway-tool alias map from the environment. fn messages_gateway_tools_from_env() -> GatewayToolMap { std::env::var(GATEWAY_TOOL_ALIASES_ENV) @@ -141,3 +196,48 @@ fn messages_gateway_tools_from_env() -> GatewayToolMap { .map(|raw| GatewayToolMap::from_env_str(&raw)) .unwrap_or_default() } + +#[cfg(test)] +mod tests { + use super::{database_backend, database_open_error}; + + #[test] + fn database_errors_are_actionable_without_exposing_credentials() { + let database_url = "postgresql://agentic-api:super'secret@postgres.example.com/agentic_api"; + let backend = database_backend(database_url); + let tls_error = sqlx::Error::Tls(Box::new(std::io::Error::other(format!( + "{database_url} certificate verify failed" + )))); + let pool_error = sqlx::Error::PoolTimedOut; + let tls_message = database_open_error(backend, &tls_error).to_string(); + let pool_message = database_open_error(backend, &pool_error).to_string(); + + assert!(tls_message.contains("database TLS error")); + assert!(tls_message.contains("certificate verify failed")); + assert!(tls_message.contains("postgresql://[redacted]")); + assert_eq!( + pool_message, + "failed to open PostgreSQL database: connection pool timeout: \ + pool timed out while waiting for an open connection" + ); + assert!(!tls_message.contains("super-secret")); + assert!(!tls_message.contains("secret")); + assert!(!tls_message.contains("agentic-api")); + assert_ne!(tls_message, pool_message); + + let mysql_error = sqlx::Error::Io(std::io::Error::other( + "mysql://gateway:mysql-secret@mysql.example.com/agentic_api refused", + )); + let mysql_message = database_open_error("configured", &mysql_error).to_string(); + assert!(mysql_message.contains("mysql://[redacted]")); + assert!(!mysql_message.contains("mysql-secret")); + + let short_postgres_url = "postgres://gateway:postgres-secret@postgres.example.com/agentic_api"; + let short_postgres_error = sqlx::Error::Io(std::io::Error::other(format!("{short_postgres_url} refused"))); + let short_postgres_message = + database_open_error(database_backend(short_postgres_url), &short_postgres_error).to_string(); + assert!(short_postgres_message.contains("failed to open PostgreSQL database")); + assert!(short_postgres_message.contains("postgres://[redacted]")); + assert!(!short_postgres_message.contains("postgres-secret")); + } +} diff --git a/crates/agentic-server-core/src/proxy.rs b/crates/agentic-server-core/src/proxy.rs index be8f83d8..cd363a79 100644 --- a/crates/agentic-server-core/src/proxy.rs +++ b/crates/agentic-server-core/src/proxy.rs @@ -317,6 +317,7 @@ mod tests { llm_ready_interval_s: 0.1, skip_llm_ready_check: false, db_url: None, + postgres: crate::config::PostgresConfig::default(), sqlite: crate::config::SqliteConfig::default(), } } diff --git a/crates/agentic-server-core/src/storage/conversation.rs b/crates/agentic-server-core/src/storage/conversation.rs index 192250e1..0dd0d238 100644 --- a/crates/agentic-server-core/src/storage/conversation.rs +++ b/crates/agentic-server-core/src/storage/conversation.rs @@ -113,6 +113,7 @@ impl ConversationStore { let mut tx = pool.begin().await?; + conversation::lock_in_tx(&mut tx, conversation_id).await?; item::create_in_tx(&mut tx, items_, Some(conversation_id)).await?; response::create_in_tx( diff --git a/crates/agentic-server-core/src/storage/mod.rs b/crates/agentic-server-core/src/storage/mod.rs index 09314414..57526aa8 100644 --- a/crates/agentic-server-core/src/storage/mod.rs +++ b/crates/agentic-server-core/src/storage/mod.rs @@ -24,8 +24,8 @@ pub use models::Conversation as DbConversation; pub use models::Item; pub use models::Response as DbResponse; pub use pool::{ - DbPool, DbResult, DbTransaction, create_pool, create_pool_with_schema, create_pool_with_schema_and_sqlite_config, - create_pool_with_sqlite_config, + DbPool, DbResult, DbTransaction, create_pool, create_pool_with_configs, create_pool_with_schema, + create_pool_with_schema_and_configs, create_pool_with_schema_and_sqlite_config, create_pool_with_sqlite_config, }; pub use response::ResponseStore; pub use schema::{PoolWithSchema, SchemaManager}; diff --git a/crates/agentic-server-core/src/storage/models/conversation.rs b/crates/agentic-server-core/src/storage/models/conversation.rs index 4d9515a5..2e170dc9 100644 --- a/crates/agentic-server-core/src/storage/models/conversation.rs +++ b/crates/agentic-server-core/src/storage/models/conversation.rs @@ -1,6 +1,6 @@ //! Conversation context and history. -use super::super::pool::{DbPool, DbResult}; +use super::super::pool::{DbPool, DbResult, DbTransaction}; use crate::utils::common::utcnow_str; /// Conversation context and history. @@ -27,7 +27,7 @@ pub async fn create(pool: &DbPool, id: &str) -> DbResult { let now = utcnow_str(); sqlx::query_as::<_, Conversation>( "INSERT INTO conversations (id, created_at) \ - VALUES (?, ?) RETURNING *", + VALUES ($1, $2) RETURNING *", ) .bind(id) .bind(now) @@ -43,7 +43,7 @@ pub async fn get_or_create(pool: &DbPool, id: &str) -> DbResult { let now = utcnow_str(); sqlx::query_as::<_, Conversation>( "INSERT INTO conversations (id, created_at) \ - VALUES (?, ?) \ + VALUES ($1, $2) \ ON CONFLICT (id) DO UPDATE SET created_at = created_at \ RETURNING *", ) @@ -58,12 +58,39 @@ pub async fn get_or_create(pool: &DbPool, id: &str) -> DbResult { /// # Errors /// Returns `DbResult::Err` if the database query fails. pub async fn get(pool: &DbPool, id: &str) -> DbResult> { - sqlx::query_as::<_, Conversation>("SELECT * FROM conversations WHERE id = ?") + sqlx::query_as::<_, Conversation>("SELECT * FROM conversations WHERE id = $1") .bind(id) .fetch_optional(pool) .await } +/// Locks an existing conversation for the lifetime of the transaction. +/// +/// `PostgreSQL` takes a row lock without writing the row. `SQLite` uses a no-op +/// update to acquire its write lock. Both serialize per-conversation sequence +/// allocation when multiple gateway replicas persist turns concurrently. +/// +/// # Errors +/// Returns `DbResult::Err` if the database query fails or the conversation does not exist. +pub async fn lock_in_tx(tx: &mut DbTransaction<'_>, id: &str) -> DbResult<()> { + if tx.as_mut().backend_name() == "PostgreSQL" { + let locked_id = sqlx::query_scalar::<_, String>("SELECT id FROM conversations WHERE id = $1 FOR UPDATE") + .bind(id) + .fetch_optional(&mut **tx) + .await?; + return locked_id.map(|_| ()).ok_or(sqlx::Error::RowNotFound); + } + + let result = sqlx::query("UPDATE conversations SET created_at = created_at WHERE id = $1") + .bind(id) + .execute(&mut **tx) + .await?; + if result.rows_affected() == 0 { + return Err(sqlx::Error::RowNotFound); + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/agentic-server-core/src/storage/models/item.rs b/crates/agentic-server-core/src/storage/models/item.rs index dad71dbf..4ec4409d 100644 --- a/crates/agentic-server-core/src/storage/models/item.rs +++ b/crates/agentic-server-core/src/storage/models/item.rs @@ -1,6 +1,7 @@ //! Conversation history item stored in the database. use serde_json::Value; +use std::fmt::Write; use tracing::warn; use super::super::pool::{DbPool, DbResult, DbTransaction}; @@ -8,6 +9,9 @@ use super::super::types::item::{InOutItem, ItemKind, STORED_ITEM_KIND_KEY}; use crate::types::io::{InputItem, OutputItem}; use crate::utils::common::{deserialize_from_str_opt, utcnow_str}; +const ITEM_COLUMN_COUNT: usize = 5; +const SEQUENCE_COLUMN_INDEX: usize = 4; + /// Conversation history item stored in the database. /// /// Maps to the `items` table and represents a single message/event @@ -88,6 +92,32 @@ impl Item { } } +fn item_values_clause(row_count: usize, first_bind_index: usize, sequence_from_cte: bool) -> String { + let mut clause = String::new(); + let mut bind_index = first_bind_index; + + for row_index in 0..row_count { + if row_index > 0 { + clause.push_str(", "); + } + clause.push('('); + for column_index in 0..ITEM_COLUMN_COUNT { + if column_index > 0 { + clause.push_str(", "); + } + if sequence_from_cte && column_index == SEQUENCE_COLUMN_INDEX { + write!(clause, "(SELECT start + ${bind_index} FROM next_seq)").expect("writing to String cannot fail"); + } else { + write!(clause, "${bind_index}").expect("writing to String cannot fail"); + } + bind_index += 1; + } + clause.push(')'); + } + + clause +} + /// Create items in a transaction with optional conversation context. /// /// If `conversation_id` is provided, the next sequence range is computed in the insert statement so @@ -109,8 +139,7 @@ pub async fn create_in_tx( } let now = utcnow_str(); - let placeholders: Vec<&str> = vec!["(?, ?, ?, ?, ?)"; items.len()]; - let values_clause = placeholders.join(", "); + let values_clause = item_values_clause(items.len(), 1, false); let sql = format!("INSERT INTO items (id, data, created_at, conversation_id, seq) VALUES {values_clause} RETURNING *"); @@ -128,13 +157,12 @@ async fn create_in_tx_with_next_conversation_seq( conversation_id: &str, ) -> DbResult> { let now = utcnow_str(); - let placeholders: Vec<&str> = vec!["(?, ?, ?, ?, (SELECT start + ? FROM next_seq))"; items.len()]; - let values_clause = placeholders.join(", "); + let values_clause = item_values_clause(items.len(), 2, true); let sql = format!( "WITH next_seq AS ( \ SELECT COALESCE(MAX(seq), -1) + 1 AS start \ FROM items \ - WHERE conversation_id = ? \ + WHERE conversation_id = $1 \ ) \ INSERT INTO items (id, data, created_at, conversation_id, seq) \ VALUES {values_clause} \ @@ -163,7 +191,10 @@ pub async fn get_items(pool: &DbPool, ids: &[String]) -> DbResult> { if ids.is_empty() { return Ok(vec![]); } - let placeholders = ids.iter().map(|_| "?").collect::>().join(", "); + let placeholders = (1..=ids.len()) + .map(|index| format!("${index}")) + .collect::>() + .join(", "); let sql = format!("SELECT * FROM items WHERE id IN ({placeholders})"); let mut q = sqlx::query_as::<_, Item>(&sql); for id in ids { @@ -177,7 +208,7 @@ pub async fn get_items(pool: &DbPool, ids: &[String]) -> DbResult> { /// # Errors /// Returns `DbResult::Err` if the database query fails. pub async fn get_items_by_conversation(pool: &DbPool, conversation_id: &str) -> DbResult> { - sqlx::query_as::<_, Item>("SELECT * FROM items WHERE conversation_id = ? ORDER BY seq ASC") + sqlx::query_as::<_, Item>("SELECT * FROM items WHERE conversation_id = $1 ORDER BY seq ASC") .bind(conversation_id) .fetch_all(pool) .await @@ -189,6 +220,23 @@ mod tests { use crate::types::event::MessageStatus; use crate::types::io::{InputItem, OutputItem, ReasoningOutput, ReasoningTextContent}; + #[test] + fn item_values_clause_numbers_plain_rows() { + assert_eq!( + item_values_clause(2, 1, false), + "($1, $2, $3, $4, $5), ($6, $7, $8, $9, $10)" + ); + } + + #[test] + fn item_values_clause_numbers_conversation_rows_after_cte_bind() { + assert_eq!( + item_values_clause(2, 2, true), + "($2, $3, $4, $5, (SELECT start + $6 FROM next_seq)), \ + ($7, $8, $9, $10, (SELECT start + $11 FROM next_seq))" + ); + } + #[test] fn test_item_basic() { let item = Item { diff --git a/crates/agentic-server-core/src/storage/models/response.rs b/crates/agentic-server-core/src/storage/models/response.rs index 5f61f77b..a613b82d 100644 --- a/crates/agentic-server-core/src/storage/models/response.rs +++ b/crates/agentic-server-core/src/storage/models/response.rs @@ -44,7 +44,7 @@ pub async fn create_in_tx( sqlx::query_as::<_, Response>( "INSERT INTO responses \ (id, conversation_id, previous_response_id, history_item_ids, metadata, created_at) \ - VALUES (?, ?, ?, ?, ?, ?) RETURNING *", + VALUES ($1, $2, $3, $4, $5, $6) RETURNING *", ) .bind(id) .bind(conversation_id) @@ -61,7 +61,7 @@ pub async fn create_in_tx( /// # Errors /// Returns `DbResult::Err` if the database query fails. pub async fn get(pool: &DbPool, id: &str) -> DbResult> { - sqlx::query_as::<_, Response>("SELECT * FROM responses WHERE id = ?") + sqlx::query_as::<_, Response>("SELECT * FROM responses WHERE id = $1") .bind(id) .fetch_optional(pool) .await diff --git a/crates/agentic-server-core/src/storage/pool.rs b/crates/agentic-server-core/src/storage/pool.rs index 7b219cae..b0eb2da3 100644 --- a/crates/agentic-server-core/src/storage/pool.rs +++ b/crates/agentic-server-core/src/storage/pool.rs @@ -5,7 +5,7 @@ use std::{sync::Arc, time::Duration}; use sqlx::AnyConnection; use sqlx::any::AnyPoolOptions; -use crate::config::SqliteConfig; +use crate::config::{PostgresConfig, SqliteConfig}; const SQLITE_MEMORY_MAX_CONNECTIONS: u32 = 1; const DEFAULT_MAX_CONNECTIONS: u32 = 10; @@ -86,6 +86,29 @@ fn sqlite_max_connections(url: &str, config: SqliteConfig) -> u32 { } } +fn is_postgres_url(url: &str) -> bool { + url.starts_with("postgres://") || url.starts_with("postgresql://") +} + +fn pool_options(url: &str, sqlite_config: SqliteConfig, postgres_config: PostgresConfig) -> AnyPoolOptions { + if url.starts_with("sqlite") { + return AnyPoolOptions::new() + .max_connections(sqlite_max_connections(url, sqlite_config)) + .after_connect(move |conn, _meta| Box::pin(configure_sqlite_connection(conn, sqlite_config))); + } + + if is_postgres_url(url) { + return AnyPoolOptions::new() + .max_connections(postgres_config.max_connections) + .acquire_timeout(postgres_config.acquire_timeout) + .idle_timeout(postgres_config.idle_timeout) + .max_lifetime(postgres_config.max_lifetime) + .after_connect(move |conn, _meta| Box::pin(configure_postgres_connection(conn, postgres_config))); + } + + AnyPoolOptions::new().max_connections(DEFAULT_MAX_CONNECTIONS) +} + async fn configure_sqlite_connection(conn: &mut AnyConnection, config: SqliteConfig) -> DbResult<()> { sqlx::query(&format!("PRAGMA busy_timeout = {SQLITE_BUSY_TIMEOUT_MS}")) .execute(&mut *conn) @@ -108,6 +131,20 @@ async fn configure_sqlite_connection(conn: &mut AnyConnection, config: SqliteCon Ok(()) } +async fn configure_postgres_connection(conn: &mut AnyConnection, config: PostgresConfig) -> DbResult<()> { + let lock_timeout_ms = format!("{}ms", config.lock_timeout.as_millis()); + sqlx::query("SELECT set_config('lock_timeout', $1, false)") + .bind(lock_timeout_ms) + .execute(&mut *conn) + .await?; + let statement_timeout_ms = format!("{}ms", config.statement_timeout.as_millis()); + sqlx::query("SELECT set_config('statement_timeout', $1, false)") + .bind(statement_timeout_ms) + .execute(conn) + .await?; + Ok(()) +} + async fn enable_sqlite_wal(pool: &DbPool) -> DbResult<()> { enable_sqlite_wal_with_retry(pool, SQLITE_WAL_MAX_ATTEMPTS, SQLITE_WAL_RETRY_DELAY).await } @@ -173,6 +210,19 @@ pub async fn create_pool(db_url: Option<&str>) -> DbResult> { pub async fn create_pool_with_sqlite_config( db_url: Option<&str>, sqlite_config: SqliteConfig, +) -> DbResult> { + create_pool_with_configs(db_url, sqlite_config, PostgresConfig::default()).await +} + +/// Creates a connection pool with explicit database-specific tuning. +/// +/// # Errors +/// +/// Returns [`sqlx::Error`] if pool creation or connection initialization fails. +pub async fn create_pool_with_configs( + db_url: Option<&str>, + sqlite_config: SqliteConfig, + postgres_config: PostgresConfig, ) -> DbResult> { // Install default drivers for auto-detection sqlx::any::install_default_drivers(); @@ -180,15 +230,7 @@ pub async fn create_pool_with_sqlite_config( // Prepare URL with database-specific parameters let url = prepare_db_url(db_url); - let max_connections = if url.starts_with("sqlite") { - sqlite_max_connections(&url, sqlite_config) - } else { - DEFAULT_MAX_CONNECTIONS - }; - let mut options = AnyPoolOptions::new().max_connections(max_connections); - if url.starts_with("sqlite") { - options = options.after_connect(move |conn, _meta| Box::pin(configure_sqlite_connection(conn, sqlite_config))); - } + let options = pool_options(&url, sqlite_config, postgres_config); let pool = options.connect(&url).await?; if sqlite_should_enable_wal(&url) { enable_sqlite_wal(&pool).await?; @@ -221,11 +263,24 @@ pub async fn create_pool_with_schema(db_url: Option<&str>) -> DbResult, sqlite_config: SqliteConfig, +) -> DbResult> { + create_pool_with_schema_and_configs(db_url, sqlite_config, PostgresConfig::default()).await +} + +/// Creates a connection pool with explicit database-specific tuning and initializes the database schema. +/// +/// # Errors +/// +/// Returns error if pool creation or schema initialization fails. +pub async fn create_pool_with_schema_and_configs( + db_url: Option<&str>, + sqlite_config: SqliteConfig, + postgres_config: PostgresConfig, ) -> DbResult> { use crate::storage::PoolWithSchema; - let pool = create_pool_with_sqlite_config(db_url, sqlite_config).await?; - let pool_with_schema = PoolWithSchema::new(pool); + let pool = create_pool_with_configs(db_url, sqlite_config, postgres_config).await?; + let pool_with_schema = PoolWithSchema::with_postgres_migration_timeout(pool, postgres_config.migration_timeout); pool_with_schema.ensure_schema_ready().await?; Ok(pool_with_schema.pool().clone()) @@ -236,7 +291,7 @@ mod tests { use super::*; use crate::config::{ DEFAULT_SQLITE_JOURNAL_SIZE_LIMIT_BYTES, DEFAULT_SQLITE_MAX_CONNECTIONS, DEFAULT_SQLITE_MMAP_SIZE_BYTES, - SqliteTempStore, + PostgresConfig, SqliteTempStore, }; use sqlx::Connection; @@ -309,6 +364,30 @@ mod tests { assert_eq!(prepared, "postgresql://user:pass@localhost/db"); } + #[test] + fn test_postgres_pool_options_use_explicit_config() { + let postgres_config = PostgresConfig { + max_connections: 7, + acquire_timeout: Duration::from_secs(11), + lock_timeout: Duration::from_secs(13), + migration_timeout: Duration::from_secs(17), + statement_timeout: Duration::from_secs(23), + idle_timeout: None, + max_lifetime: Some(Duration::from_secs(19)), + }; + + let options = pool_options( + "postgresql://user:pass@localhost/db", + SqliteConfig::default(), + postgres_config, + ); + + assert_eq!(options.get_max_connections(), 7); + assert_eq!(options.get_acquire_timeout(), Duration::from_secs(11)); + assert_eq!(options.get_idle_timeout(), None); + assert_eq!(options.get_max_lifetime(), Some(Duration::from_secs(19))); + } + #[test] fn test_prepare_mysql_url() { let url = "mysql://user:pass@localhost/db"; diff --git a/crates/agentic-server-core/src/storage/schema.rs b/crates/agentic-server-core/src/storage/schema.rs index 80105c97..916dfec9 100644 --- a/crates/agentic-server-core/src/storage/schema.rs +++ b/crates/agentic-server-core/src/storage/schema.rs @@ -3,13 +3,145 @@ use std::env; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; +use sqlx::Connection; use tracing::{debug, info}; use super::pool::DbPool; +use crate::config::DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS; type DbResult = Result; +const POSTGRES_SCHEMA_ADVISORY_LOCK: i64 = 7_194_963_546_799_751; +const POSTGRES_INTEGER_WIDENING_SQL: &str = " + ALTER TABLE conversations + ALTER COLUMN created_at TYPE BIGINT USING created_at::BIGINT; + ALTER TABLE items + ALTER COLUMN created_at TYPE BIGINT USING created_at::BIGINT, + ALTER COLUMN seq TYPE BIGINT USING seq::BIGINT; + ALTER TABLE responses + ALTER COLUMN created_at TYPE BIGINT USING created_at::BIGINT; +"; + +async fn configure_postgres_migration_timeout( + connection: &mut sqlx::AnyConnection, + migration_timeout: Duration, +) -> DbResult<()> { + let timeout_ms = format!("{}ms", migration_timeout.as_millis()); + sqlx::query("SELECT set_config('lock_timeout', $1, false)") + .bind(&timeout_ms) + .execute(&mut *connection) + .await?; + sqlx::query("SELECT set_config('statement_timeout', $1, false)") + .bind(timeout_ms) + .execute(connection) + .await?; + Ok(()) +} + +async fn widen_postgres_integer_columns(connection: &mut sqlx::AnyConnection) -> DbResult<()> { + let mut transaction = connection.begin().await?; + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(POSTGRES_SCHEMA_ADVISORY_LOCK) + .execute(&mut *transaction) + .await?; + let narrow_column_count = postgres_narrow_integer_column_count(&mut *transaction).await?; + if narrow_column_count > 0 { + sqlx::raw_sql(POSTGRES_INTEGER_WIDENING_SQL) + .execute(&mut *transaction) + .await?; + } + transaction.commit().await +} + +async fn postgres_narrow_integer_column_count<'e, E>(executor: E) -> DbResult +where + E: sqlx::Executor<'e, Database = sqlx::Any>, +{ + sqlx::query_scalar( + "SELECT COUNT(*) FROM information_schema.columns \ + WHERE table_schema = current_schema() AND data_type <> 'bigint' \ + AND ((table_name = 'conversations' AND column_name = 'created_at') \ + OR (table_name = 'items' AND column_name IN ('created_at', 'seq')) \ + OR (table_name = 'responses' AND column_name = 'created_at'))", + ) + .fetch_one(executor) + .await +} + +fn validate_supervisor_integer_width(narrow_column_count: i64) -> DbResult<()> { + if narrow_column_count == 0 { + return Ok(()); + } + Err(sqlx::Error::Configuration( + "supervisor-managed PostgreSQL schema requires BIGINT compatibility upgrade; \ + apply the documented ALTER TABLE statements before setting AGENTIC_API_SCHEMA_READY" + .into(), + )) +} + +async fn verify_supervisor_managed_postgres_schema( + pool: &DbPool, + postgres_migration_timeout: Duration, +) -> DbResult<()> { + let mut connection = pool.acquire().await?; + if connection.backend_name() != "PostgreSQL" { + return Ok(()); + } + + if let Err(error) = configure_postgres_migration_timeout(&mut connection, postgres_migration_timeout).await { + let _ = connection.close().await; + return Err(error); + } + let compatibility_result = match postgres_narrow_integer_column_count(&mut *connection).await { + Ok(narrow_column_count) => validate_supervisor_integer_width(narrow_column_count), + Err(error) => Err(error), + }; + let close_result = connection.close().await; + compatibility_result?; + close_result +} + +async fn apply_postgres_compatibility( + connection: &mut sqlx::AnyConnection, + postgres_migration_timeout: Duration, +) -> DbResult<()> { + configure_postgres_migration_timeout(connection, postgres_migration_timeout).await?; + widen_postgres_integer_columns(connection).await +} + +async fn run_embedded_migrations(pool: &DbPool, postgres_migration_timeout: Duration) -> DbResult<()> { + let mut connection = pool.acquire().await?; + let is_postgres = connection.backend_name() == "PostgreSQL"; + if is_postgres { + if let Err(error) = configure_postgres_migration_timeout(&mut connection, postgres_migration_timeout).await { + let _ = connection.close().await; + return Err(error); + } + } + + let migration_result = sqlx::migrate!("./migrations") + .run(&mut *connection) + .await + .map_err(|error| sqlx::Error::Configuration(error.to_string().into())); + let postgres_result = if migration_result.is_ok() && is_postgres { + apply_postgres_compatibility(&mut connection, postgres_migration_timeout).await + } else { + Ok(()) + }; + let close_result = if is_postgres { + connection.close().await + } else { + drop(connection); + Ok(()) + }; + + migration_result?; + postgres_result?; + close_result +} + fn is_marked_ready() -> bool { matches!( env::var("AGENTIC_API_SCHEMA_READY").as_deref(), @@ -25,15 +157,23 @@ fn is_marked_ready() -> bool { pub struct PoolWithSchema { pool: Arc, schema_ready: AtomicBool, + postgres_migration_timeout: Duration, } impl PoolWithSchema { /// Creates a new pool with schema tracking. #[must_use] pub fn new(pool: Arc) -> Self { + Self::with_postgres_migration_timeout(pool, Duration::from_secs(DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS)) + } + + /// Creates a new pool with schema tracking and a `PostgreSQL` migration timeout. + #[must_use] + pub fn with_postgres_migration_timeout(pool: Arc, postgres_migration_timeout: Duration) -> Self { Self { pool, schema_ready: AtomicBool::new(false), + postgres_migration_timeout, } } @@ -49,26 +189,29 @@ impl PoolWithSchema { /// 2. `AGENTIC_API_SCHEMA_READY` environment variable /// /// If none of the above, runs all pending migrations from the `migrations/` directory. + /// Supervisor-managed `PostgreSQL` schemas still verify required compatibility upgrades. /// /// # Errors /// /// Returns a [`sqlx::Error`] if migrations fail. pub async fn ensure_schema_ready(&self) -> DbResult<()> { + self.ensure_schema_ready_with_marker(is_marked_ready()).await + } + + async fn ensure_schema_ready_with_marker(&self, supervisor_managed: bool) -> DbResult<()> { if self.schema_ready.load(Ordering::SeqCst) { return Ok(()); } - if is_marked_ready() { - debug!("[schema] DDL skipped — marked ready by supervisor."); + if supervisor_managed { + debug!("[schema] Migrations skipped — marked ready by supervisor."); + verify_supervisor_managed_postgres_schema(self.pool.as_ref(), self.postgres_migration_timeout).await?; self.schema_ready.store(true, Ordering::SeqCst); return Ok(()); } debug!("[schema] Running migrations..."); - sqlx::migrate!("./migrations") - .run(self.pool.as_ref()) - .await - .map_err(|e| sqlx::Error::Configuration(e.to_string().into()))?; + run_embedded_migrations(self.pool.as_ref(), self.postgres_migration_timeout).await?; info!("[schema] DB schema ready."); self.schema_ready.store(true, Ordering::SeqCst); Ok(()) @@ -97,10 +240,11 @@ impl<'a> SchemaManager<'a> { /// Returns a [`sqlx::Error`] if migrations fail. pub async fn run_migrations(&self) -> DbResult<()> { debug!("[schema] Running migrations..."); - sqlx::migrate!("./migrations") - .run(self.pool) - .await - .map_err(|e| sqlx::Error::Configuration(e.to_string().into()))?; + run_embedded_migrations( + self.pool, + Duration::from_secs(DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS), + ) + .await?; info!("[schema] DB schema ready."); Ok(()) } @@ -183,4 +327,162 @@ mod tests { pwc1.ensure_schema_ready().await.expect("pool1 repeat failed"); pwc2.ensure_schema_ready().await.expect("pool2 repeat failed"); } + + #[tokio::test] + #[ignore = "requires TEST_POSTGRES_URL pointing to an isolated PostgreSQL database"] + #[allow( + clippy::too_many_lines, + reason = "keeps the complete migration lifecycle in one integration test" + )] + async fn postgres_integer_widening_upgrade_preserves_existing_state() { + let database_url = std::env::var("TEST_POSTGRES_URL").expect("TEST_POSTGRES_URL must be set"); + let pool = crate::storage::pool::create_pool(Some(&database_url)) + .await + .expect("create PostgreSQL upgrade-test pool"); + let schema = format!("upgrade_{}", uuid::Uuid::now_v7().simple()); + let mut connection = pool.acquire().await.expect("acquire PostgreSQL upgrade connection"); + sqlx::query(&format!("CREATE SCHEMA {schema}")) + .execute(&mut *connection) + .await + .expect("create isolated upgrade schema"); + sqlx::query(&format!("SET search_path TO {schema}")) + .execute(&mut *connection) + .await + .expect("select isolated upgrade schema"); + for migration in [ + include_str!("../../migrations/0001_initial.sql"), + include_str!("../../migrations/0002_add_placeholders.sql"), + include_str!("../../migrations/0003_index_conversation_sequence.sql"), + ] { + sqlx::raw_sql(migration) + .execute(&mut *connection) + .await + .expect("apply portable migration"); + } + sqlx::query("INSERT INTO conversations (id, created_at, metadata) VALUES ($1, $2, $3)") + .bind("conv_upgrade") + .bind(1_704_067_200_i64) + .bind("{\"source\":\"upgrade\"}") + .execute(&mut *connection) + .await + .expect("seed conversation"); + sqlx::query("INSERT INTO items (id, data, created_at, conversation_id, seq) VALUES ($1, $2, $3, $4, $5)") + .bind("item_upgrade") + .bind("{}") + .bind(1_704_067_200_i64) + .bind("conv_upgrade") + .bind(0_i64) + .execute(&mut *connection) + .await + .expect("seed item"); + sqlx::query( + "INSERT INTO responses \ + (id, conversation_id, history_item_ids, metadata, created_at) VALUES ($1, $2, $3, $4, $5)", + ) + .bind("resp_upgrade") + .bind("conv_upgrade") + .bind("[\"item_upgrade\"]") + .bind("{\"source\":\"upgrade\"}") + .bind(1_704_067_200_i64) + .execute(&mut *connection) + .await + .expect("seed response"); + + let narrow_column_count = postgres_narrow_integer_column_count(&mut *connection) + .await + .expect("inspect pre-upgrade PostgreSQL columns"); + assert_eq!(narrow_column_count, 4); + assert!(validate_supervisor_integer_width(narrow_column_count).is_err()); + + let supervisor_schema_name = schema.clone(); + let supervisor_pool = sqlx::any::AnyPoolOptions::new() + .max_connections(1) + .after_connect(move |connection, _metadata| { + let supervisor_schema_name = supervisor_schema_name.clone(); + Box::pin(async move { + sqlx::query("SELECT set_config('search_path', $1, false)") + .bind(supervisor_schema_name) + .execute(connection) + .await?; + Ok(()) + }) + }) + .connect(&database_url) + .await + .expect("create supervisor-managed PostgreSQL pool"); + let supervisor_schema = + PoolWithSchema::with_postgres_migration_timeout(Arc::new(supervisor_pool), Duration::from_secs(5)); + let supervisor_error = supervisor_schema + .ensure_schema_ready_with_marker(true) + .await + .expect_err("narrow supervisor-managed schema should fail compatibility check"); + assert!(supervisor_error.to_string().contains("BIGINT compatibility upgrade")); + assert!(!supervisor_schema.schema_ready.load(Ordering::SeqCst)); + + apply_postgres_compatibility(&mut connection, Duration::from_secs(5)) + .await + .expect("widen PostgreSQL integer columns"); + supervisor_schema + .ensure_schema_ready_with_marker(true) + .await + .expect("widened supervisor-managed schema should pass compatibility check"); + assert!(supervisor_schema.schema_ready.load(Ordering::SeqCst)); + + let future_timestamp = i64::from(i32::MAX) + 1; + sqlx::query("UPDATE conversations SET created_at = $1 WHERE id = $2") + .bind(future_timestamp) + .bind("conv_upgrade") + .execute(&mut *connection) + .await + .expect("write timestamp beyond PostgreSQL INT4 range"); + let linked_state: (i64, String, String) = sqlx::query_as( + "SELECT conversations.created_at, items.id, responses.id \ + FROM conversations \ + JOIN items ON items.conversation_id = conversations.id \ + JOIN responses ON responses.conversation_id = conversations.id \ + WHERE conversations.id = $1", + ) + .bind("conv_upgrade") + .fetch_one(&mut *connection) + .await + .expect("load migrated linked state"); + assert_eq!( + linked_state, + (future_timestamp, "item_upgrade".to_owned(), "resp_upgrade".to_owned()) + ); + let bigint_columns: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM information_schema.columns \ + WHERE table_schema = $1 AND data_type = 'bigint' \ + AND ((table_name = 'conversations' AND column_name = 'created_at') \ + OR (table_name = 'items' AND column_name IN ('created_at', 'seq')) \ + OR (table_name = 'responses' AND column_name = 'created_at'))", + ) + .bind(&schema) + .fetch_one(&mut *connection) + .await + .expect("inspect widened PostgreSQL columns"); + assert_eq!(bigint_columns, 4); + assert!(validate_supervisor_integer_width(4 - bigint_columns).is_ok()); + let foreign_key_error = + sqlx::query("INSERT INTO items (id, data, created_at, conversation_id) VALUES ($1, $2, $3, $4)") + .bind("item_invalid") + .bind("{}") + .bind(future_timestamp) + .bind("conv_missing") + .execute(&mut *connection) + .await; + assert!(foreign_key_error.is_err()); + + sqlx::query("SET search_path TO public") + .execute(&mut *connection) + .await + .expect("restore public schema"); + supervisor_schema.pool.close().await; + sqlx::query(&format!("DROP SCHEMA {schema} CASCADE")) + .execute(&mut *connection) + .await + .expect("drop isolated upgrade schema"); + drop(connection); + pool.close().await; + } } diff --git a/crates/agentic-server-core/tests/postgres_storage_integration.rs b/crates/agentic-server-core/tests/postgres_storage_integration.rs new file mode 100644 index 00000000..34eb6efe --- /dev/null +++ b/crates/agentic-server-core/tests/postgres_storage_integration.rs @@ -0,0 +1,426 @@ +use std::sync::Arc; +use std::time::Duration; + +use agentic_core::config::{PostgresConfig, SqliteConfig}; +use agentic_core::storage::{ + ConversationStore, InOutItem, ResponseMetadata, ResponseStore, create_pool_with_configs, + create_pool_with_schema_and_configs, +}; +use agentic_core::types::io::{InputItem, InputMessage, InputMessageContent}; +use tokio::sync::Barrier; +use tokio::time::{sleep, timeout}; + +fn input_item(text: &str) -> InOutItem { + InOutItem::Input(InputItem::Message(InputMessage { + role: "user".to_owned(), + content: InputMessageContent::Text(text.to_owned()), + })) +} + +#[tokio::test] +#[ignore = "requires TEST_POSTGRES_URL pointing to an isolated PostgreSQL database"] +#[allow( + clippy::too_many_lines, + reason = "keeps the complete index migration lifecycle in one integration test" +)] +async fn postgres_index_upgrade_preserves_existing_state() { + let database_url = std::env::var("TEST_POSTGRES_URL").expect("TEST_POSTGRES_URL must be set"); + let pool = create_pool_with_configs(Some(&database_url), SqliteConfig::default(), PostgresConfig::default()) + .await + .expect("create PostgreSQL upgrade-test pool"); + let schema = format!("upgrade_{}", uuid::Uuid::now_v7().simple()); + let mut connection = pool.acquire().await.expect("acquire PostgreSQL upgrade connection"); + sqlx::query(&format!("CREATE SCHEMA {schema}")) + .execute(&mut *connection) + .await + .expect("create isolated upgrade schema"); + sqlx::query(&format!("SET search_path TO {schema}")) + .execute(&mut *connection) + .await + .expect("select isolated upgrade schema"); + sqlx::raw_sql(include_str!("../migrations/0001_initial.sql")) + .execute(&mut *connection) + .await + .expect("apply initial PostgreSQL schema"); + sqlx::raw_sql(include_str!("../migrations/0002_add_placeholders.sql")) + .execute(&mut *connection) + .await + .expect("apply placeholder PostgreSQL schema"); + sqlx::query("INSERT INTO conversations (id, created_at, metadata) VALUES ($1, $2, $3)") + .bind("conv_upgrade") + .bind(1_704_067_200_i64) + .bind("{\"source\":\"upgrade\"}") + .execute(&mut *connection) + .await + .expect("seed conversation"); + sqlx::query("INSERT INTO items (id, data, created_at, conversation_id, seq) VALUES ($1, $2, $3, $4, $5)") + .bind("item_upgrade") + .bind("{}") + .bind(1_704_067_200_i64) + .bind("conv_upgrade") + .bind(0_i64) + .execute(&mut *connection) + .await + .expect("seed item"); + sqlx::query( + "INSERT INTO responses \ + (id, conversation_id, history_item_ids, metadata, created_at) VALUES ($1, $2, $3, $4, $5)", + ) + .bind("resp_upgrade") + .bind("conv_upgrade") + .bind("[\"item_upgrade\"]") + .bind("{\"source\":\"upgrade\"}") + .bind(1_704_067_200_i64) + .execute(&mut *connection) + .await + .expect("seed response"); + + sqlx::raw_sql(include_str!("../migrations/0003_index_conversation_sequence.sql")) + .execute(&mut *connection) + .await + .expect("apply PostgreSQL composite index migration"); + + let linked_state: (i64, String, String) = sqlx::query_as( + "SELECT conversations.created_at, items.id, responses.id \ + FROM conversations \ + JOIN items ON items.conversation_id = conversations.id \ + JOIN responses ON responses.conversation_id = conversations.id \ + WHERE conversations.id = $1", + ) + .bind("conv_upgrade") + .fetch_one(&mut *connection) + .await + .expect("load migrated linked state"); + assert_eq!( + linked_state, + (1_704_067_200, "item_upgrade".to_owned(), "resp_upgrade".to_owned()) + ); + let recreated_indexes: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM pg_indexes \ + WHERE schemaname = $1 AND indexname IN ( \ + 'idx_conversations_tenant_id', 'idx_items_conversation_id', \ + 'idx_items_created_at', 'idx_items_tenant_id', \ + 'idx_responses_conversation_id', 'idx_responses_previous_response_id', \ + 'idx_responses_created_at', 'idx_responses_tenant_id' \ + )", + ) + .bind(&schema) + .fetch_one(&mut *connection) + .await + .expect("inspect recreated PostgreSQL indexes"); + assert_eq!(recreated_indexes, 8); + let conversation_index: String = sqlx::query_scalar( + "SELECT indexdef FROM pg_indexes \ + WHERE schemaname = $1 AND indexname = 'idx_items_conversation_id'", + ) + .bind(&schema) + .fetch_one(&mut *connection) + .await + .expect("inspect conversation sequence index"); + assert!(conversation_index.contains("(conversation_id, seq)")); + let foreign_key_error = + sqlx::query("INSERT INTO items (id, data, created_at, conversation_id) VALUES ($1, $2, $3, $4)") + .bind("item_invalid") + .bind("{}") + .bind(1_704_067_200_i64) + .bind("conv_missing") + .execute(&mut *connection) + .await; + assert!(foreign_key_error.is_err()); + + sqlx::query("SET search_path TO public") + .execute(&mut *connection) + .await + .expect("restore public schema"); + sqlx::query(&format!("DROP SCHEMA {schema} CASCADE")) + .execute(&mut *connection) + .await + .expect("drop isolated upgrade schema"); + drop(connection); + pool.close().await; +} + +#[tokio::test] +#[ignore = "requires TEST_POSTGRES_URL pointing to an isolated PostgreSQL database"] +#[allow( + clippy::too_many_lines, + reason = "keeps the complete restart persistence lifecycle in one integration test" +)] +async fn postgres_migrations_and_state_survive_pool_restarts() { + let database_url = std::env::var("TEST_POSTGRES_URL").expect("TEST_POSTGRES_URL must be set"); + let postgres_config = PostgresConfig { + max_connections: 2, + acquire_timeout: Duration::from_secs(5), + lock_timeout: Duration::from_secs(1), + migration_timeout: Duration::from_secs(5), + statement_timeout: Duration::from_secs(5), + idle_timeout: Some(Duration::from_secs(30)), + max_lifetime: Some(Duration::from_secs(60)), + }; + let response_id = format!("resp_postgres_{}", uuid::Uuid::now_v7()); + let continuation_id = format!("resp_postgres_{}", uuid::Uuid::now_v7()); + let conversation_response_id = format!("resp_postgres_{}", uuid::Uuid::now_v7()); + let conversation_continuation_id = format!("resp_postgres_{}", uuid::Uuid::now_v7()); + let metadata = ResponseMetadata::default(); + + let first_pool = create_pool_with_schema_and_configs(Some(&database_url), SqliteConfig::default(), postgres_config) + .await + .expect("initialize clean PostgreSQL database"); + assert_eq!(first_pool.options().get_max_connections(), 2); + let statement_timeout: String = sqlx::query_scalar("SHOW statement_timeout") + .fetch_one(first_pool.as_ref()) + .await + .expect("inspect PostgreSQL statement timeout"); + assert_eq!(statement_timeout, "5s"); + let created_at_type: String = sqlx::query_scalar( + "SELECT data_type FROM information_schema.columns \ + WHERE table_schema = current_schema() AND table_name = 'conversations' AND column_name = 'created_at'", + ) + .fetch_one(first_pool.as_ref()) + .await + .expect("inspect PostgreSQL timestamp column"); + assert_eq!(created_at_type, "bigint"); + let future_conversation_id = format!("conv_postgres_{}", uuid::Uuid::now_v7()); + let future_timestamp = i64::from(i32::MAX) + 1; + sqlx::query("INSERT INTO conversations (id, created_at) VALUES ($1, $2)") + .bind(&future_conversation_id) + .bind(future_timestamp) + .execute(first_pool.as_ref()) + .await + .expect("persist timestamp beyond PostgreSQL INT4 range"); + let stored_future_timestamp: i64 = sqlx::query_scalar("SELECT created_at FROM conversations WHERE id = $1") + .bind(&future_conversation_id) + .fetch_one(first_pool.as_ref()) + .await + .expect("load timestamp beyond PostgreSQL INT4 range"); + assert_eq!(stored_future_timestamp, future_timestamp); + + ResponseStore::new(first_pool.clone()) + .persist(&response_id, None, vec![input_item("first turn")], &metadata) + .await + .expect("persist first response"); + let first_conversation_store = ConversationStore::new(first_pool.clone()); + let conversation = first_conversation_store.create().await.expect("create conversation"); + first_conversation_store + .persist( + &conversation.conversation_id, + &conversation_response_id, + None, + vec![ + input_item("conversation turn one"), + input_item("conversation turn one follow-up"), + ], + &metadata, + ) + .await + .expect("persist first conversation turn"); + first_pool.close().await; + + let second_pool = + create_pool_with_schema_and_configs(Some(&database_url), SqliteConfig::default(), postgres_config) + .await + .expect("repeat migrations after first restart"); + let second_store = ResponseStore::new(second_pool.clone()); + assert_eq!( + second_store + .rehydrate(&response_id) + .await + .expect("rehydrate first response after restart") + .len(), + 1 + ); + second_store + .persist( + &continuation_id, + Some(&response_id), + vec![input_item("second turn")], + &metadata, + ) + .await + .expect("persist continuation"); + let second_conversation_store = ConversationStore::new(second_pool.clone()); + assert_eq!( + second_conversation_store + .rehydrate(&conversation.conversation_id) + .await + .expect("rehydrate conversation after restart") + .len(), + 2 + ); + second_conversation_store + .persist( + &conversation.conversation_id, + &conversation_continuation_id, + Some(&conversation_response_id), + vec![input_item("conversation turn two")], + &metadata, + ) + .await + .expect("persist second conversation turn"); + second_pool.close().await; + + let third_pool = create_pool_with_schema_and_configs(Some(&database_url), SqliteConfig::default(), postgres_config) + .await + .expect("repeat migrations after second restart"); + assert_eq!( + ResponseStore::new(third_pool.clone()) + .rehydrate(&continuation_id) + .await + .expect("rehydrate continuation after restart") + .len(), + 2 + ); + assert_eq!( + ConversationStore::new(third_pool) + .rehydrate(&conversation.conversation_id) + .await + .expect("rehydrate full conversation after restart") + .len(), + 3 + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "requires TEST_POSTGRES_URL pointing to an isolated PostgreSQL database"] +async fn postgres_concurrent_conversation_writes_have_contiguous_sequences() { + const WRITE_COUNT: usize = 12; + + let database_url = std::env::var("TEST_POSTGRES_URL").expect("TEST_POSTGRES_URL must be set"); + let postgres_config = PostgresConfig { + max_connections: u32::try_from(WRITE_COUNT).expect("test write count must fit in u32"), + acquire_timeout: Duration::from_secs(5), + lock_timeout: Duration::from_secs(1), + migration_timeout: Duration::from_secs(5), + statement_timeout: Duration::from_secs(5), + idle_timeout: Some(Duration::from_secs(30)), + max_lifetime: Some(Duration::from_secs(60)), + }; + let first_pool = create_pool_with_schema_and_configs(Some(&database_url), SqliteConfig::default(), postgres_config) + .await + .expect("initialize PostgreSQL database"); + let second_pool = create_pool_with_configs(Some(&database_url), SqliteConfig::default(), postgres_config) + .await + .expect("create independent PostgreSQL pool"); + let conversation_store = ConversationStore::new(first_pool.clone()); + let conversation = conversation_store.create().await.expect("create conversation"); + let barrier = Arc::new(Barrier::new(WRITE_COUNT + 1)); + let mut tasks = Vec::with_capacity(WRITE_COUNT); + + for index in 0..WRITE_COUNT { + let barrier = Arc::clone(&barrier); + let pool = if index % 2 == 0 { + first_pool.clone() + } else { + second_pool.clone() + }; + let conversation_id = conversation.conversation_id.clone(); + tasks.push(tokio::spawn(async move { + let store = ConversationStore::new(pool); + let response_id = format!("resp_postgres_{}", uuid::Uuid::now_v7()); + barrier.wait().await; + store + .persist( + &conversation_id, + &response_id, + None, + vec![input_item(&format!("concurrent turn {index}"))], + &ResponseMetadata::default(), + ) + .await + })); + } + + barrier.wait().await; + for task in tasks { + task.await + .expect("join concurrent write") + .expect("persist concurrent turn"); + } + + let rows = + agentic_core::storage::models::item::get_items_by_conversation(&first_pool, &conversation.conversation_id) + .await + .expect("load concurrent conversation items"); + let sequences = rows + .iter() + .map(|row| row.seq.expect("conversation item sequence")) + .collect::>(); + let write_count = i64::try_from(WRITE_COUNT).expect("test write count must fit in i64"); + assert_eq!(sequences, (0..write_count).collect::>()); + + first_pool.close().await; + second_pool.close().await; +} + +#[tokio::test] +#[ignore = "requires TEST_POSTGRES_URL pointing to an isolated PostgreSQL database"] +async fn postgres_lock_wait_is_bounded_without_blocking_other_conversations() { + let database_url = std::env::var("TEST_POSTGRES_URL").expect("TEST_POSTGRES_URL must be set"); + let postgres_config = PostgresConfig { + max_connections: 3, + acquire_timeout: Duration::from_secs(5), + lock_timeout: Duration::from_secs(1), + migration_timeout: Duration::from_secs(5), + statement_timeout: Duration::from_secs(5), + idle_timeout: Some(Duration::from_secs(30)), + max_lifetime: Some(Duration::from_secs(60)), + }; + let first_pool = create_pool_with_schema_and_configs(Some(&database_url), SqliteConfig::default(), postgres_config) + .await + .expect("initialize PostgreSQL database"); + let second_pool = create_pool_with_configs(Some(&database_url), SqliteConfig::default(), postgres_config) + .await + .expect("create independent PostgreSQL pool"); + let first_store = ConversationStore::new(first_pool.clone()); + let locked_conversation = first_store.create().await.expect("create locked conversation"); + let unrelated_conversation = first_store.create().await.expect("create unrelated conversation"); + let mut holding_transaction = first_pool.begin().await.expect("begin lock-holding transaction"); + agentic_core::storage::models::conversation::lock_in_tx( + &mut holding_transaction, + &locked_conversation.conversation_id, + ) + .await + .expect("hold conversation row lock"); + + let blocked_store = ConversationStore::new(second_pool.clone()); + let blocked_conversation_id = locked_conversation.conversation_id.clone(); + let blocked_write = tokio::spawn(async move { + blocked_store + .persist( + &blocked_conversation_id, + &format!("resp_postgres_{}", uuid::Uuid::now_v7()), + None, + vec![input_item("blocked turn")], + &ResponseMetadata::default(), + ) + .await + }); + sleep(Duration::from_millis(100)).await; + + timeout( + Duration::from_millis(500), + ConversationStore::new(second_pool.clone()).persist( + &unrelated_conversation.conversation_id, + &format!("resp_postgres_{}", uuid::Uuid::now_v7()), + None, + vec![input_item("unrelated turn")], + &ResponseMetadata::default(), + ), + ) + .await + .expect("unrelated conversation write should not wait for the held lock") + .expect("persist unrelated conversation turn"); + + let blocked_result = timeout(Duration::from_secs(2), blocked_write) + .await + .expect("locked write should respect PostgreSQL lock_timeout") + .expect("join blocked write"); + assert!(blocked_result.is_err(), "locked write unexpectedly succeeded"); + + holding_transaction + .rollback() + .await + .expect("release conversation row lock"); + first_pool.close().await; + second_pool.close().await; +} diff --git a/crates/agentic-server/benches/gateway_bench.rs b/crates/agentic-server/benches/gateway_bench.rs index 488e51b0..f600abe6 100644 --- a/crates/agentic-server/benches/gateway_bench.rs +++ b/crates/agentic-server/benches/gateway_bench.rs @@ -159,6 +159,7 @@ async fn spawn_gateway(llm_url: &str) -> (Arc, String) { llm_ready_interval_s: 0.1, skip_llm_ready_check: false, db_url: Some(format!("sqlite://{}", db_path.display())), + postgres: agentic_core::config::PostgresConfig::default(), sqlite: agentic_core::config::SqliteConfig::default(), }; diff --git a/crates/agentic-server/benches/proxy_bench.rs b/crates/agentic-server/benches/proxy_bench.rs index 4dce033b..5dc45813 100644 --- a/crates/agentic-server/benches/proxy_bench.rs +++ b/crates/agentic-server/benches/proxy_bench.rs @@ -31,6 +31,7 @@ fn bench_config(llm_url: &str) -> Config { llm_ready_interval_s: 0.1, skip_llm_ready_check: false, db_url: None, + postgres: agentic_core::config::PostgresConfig::default(), sqlite: agentic_core::config::SqliteConfig::default(), } } diff --git a/crates/agentic-server/src/main.rs b/crates/agentic-server/src/main.rs index 890fa4d2..3d874472 100644 --- a/crates/agentic-server/src/main.rs +++ b/crates/agentic-server/src/main.rs @@ -1,8 +1,13 @@ +use std::time::Duration; + use clap::{Args, Parser, Subcommand}; use agentic_core::config::{ - Config, DEFAULT_SQLITE_JOURNAL_SIZE_LIMIT_BYTES, DEFAULT_SQLITE_MAX_CONNECTIONS, DEFAULT_SQLITE_MMAP_SIZE_BYTES, - SqliteConfig, SqliteTempStore, normalize_base_url, + Config, DEFAULT_POSTGRES_ACQUIRE_TIMEOUT_SECONDS, DEFAULT_POSTGRES_IDLE_TIMEOUT_SECONDS, + DEFAULT_POSTGRES_LOCK_TIMEOUT_SECONDS, DEFAULT_POSTGRES_MAX_CONNECTIONS, DEFAULT_POSTGRES_MAX_LIFETIME_SECONDS, + DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS, DEFAULT_POSTGRES_STATEMENT_TIMEOUT_SECONDS, + DEFAULT_SQLITE_JOURNAL_SIZE_LIMIT_BYTES, DEFAULT_SQLITE_MAX_CONNECTIONS, DEFAULT_SQLITE_MMAP_SIZE_BYTES, + PostgresConfig, SqliteConfig, SqliteTempStore, normalize_base_url, }; use agentic_core::error::Error; @@ -34,6 +39,7 @@ struct CommonArgs { #[arg( long, env = "DATABASE_URL", + hide_env_values = true, default_value = "sqlite://./agentic_api.db", global = true )] @@ -104,6 +110,51 @@ fn parse_env_u32_value(name: &str, value: Result, de } } +fn parse_env_duration(name: &str, default_seconds: u64) -> Result { + parse_env_duration_value(name, std::env::var(name), default_seconds) +} + +fn parse_env_nonzero_u32(name: &str, default: u32) -> Result { + parse_env_nonzero_u32_value(name, std::env::var(name), default) +} + +fn parse_env_nonzero_u32_value( + name: &str, + value: Result, + default: u32, +) -> Result { + let value = parse_env_u32_value(name, value, default)?; + if value == 0 { + return Err(Error::Config(format!("{name} must be greater than 0"))); + } + Ok(value) +} + +fn parse_env_duration_value( + name: &str, + value: Result, + default_seconds: u64, +) -> Result { + let seconds = parse_env_u64_value(name, value, default_seconds)?; + if seconds == 0 { + return Err(Error::Config(format!("{name} must be greater than 0"))); + } + Ok(Duration::from_secs(seconds)) +} + +fn parse_env_optional_duration(name: &str, default_seconds: u64) -> Result, Error> { + parse_env_optional_duration_value(name, std::env::var(name), default_seconds) +} + +fn parse_env_optional_duration_value( + name: &str, + value: Result, + default_seconds: u64, +) -> Result, Error> { + let seconds = parse_env_u64_value(name, value, default_seconds)?; + Ok((seconds > 0).then(|| Duration::from_secs(seconds))) +} + fn parse_env_temp_store() -> Result { parse_env_temp_store_value(std::env::var("SQLITE_TEMP_STORE")) } @@ -135,7 +186,45 @@ fn sqlite_config_from_env() -> Result { }) } +fn postgres_config_from_env() -> Result { + Ok(PostgresConfig { + max_connections: parse_env_nonzero_u32("POSTGRES_MAX_CONNECTIONS", DEFAULT_POSTGRES_MAX_CONNECTIONS)?, + acquire_timeout: parse_env_duration( + "POSTGRES_ACQUIRE_TIMEOUT_SECONDS", + DEFAULT_POSTGRES_ACQUIRE_TIMEOUT_SECONDS, + )?, + lock_timeout: parse_env_duration("POSTGRES_LOCK_TIMEOUT_SECONDS", DEFAULT_POSTGRES_LOCK_TIMEOUT_SECONDS)?, + migration_timeout: parse_env_duration( + "POSTGRES_MIGRATION_TIMEOUT_SECONDS", + DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS, + )?, + statement_timeout: parse_env_duration( + "POSTGRES_STATEMENT_TIMEOUT_SECONDS", + DEFAULT_POSTGRES_STATEMENT_TIMEOUT_SECONDS, + )?, + idle_timeout: parse_env_optional_duration( + "POSTGRES_IDLE_TIMEOUT_SECONDS", + DEFAULT_POSTGRES_IDLE_TIMEOUT_SECONDS, + )?, + max_lifetime: parse_env_optional_duration( + "POSTGRES_MAX_LIFETIME_SECONDS", + DEFAULT_POSTGRES_MAX_LIFETIME_SECONDS, + )?, + }) +} + +fn database_configs_from_env(database_url: &str) -> Result<(PostgresConfig, SqliteConfig), Error> { + if database_url.starts_with("postgres://") || database_url.starts_with("postgresql://") { + return Ok((postgres_config_from_env()?, SqliteConfig::default())); + } + if database_url.starts_with("sqlite") { + return Ok((PostgresConfig::default(), sqlite_config_from_env()?)); + } + Ok((PostgresConfig::default(), SqliteConfig::default())) +} + fn build_config(llm_api_base: String, common: &CommonArgs) -> Result { + let (postgres, sqlite) = database_configs_from_env(&common.db_url)?; Ok(Config { llm_api_base, openai_api_key: common.openai_api_key.clone(), @@ -143,7 +232,8 @@ fn build_config(llm_api_base: String, common: &CommonArgs) -> Result Result<(), Error> { #[cfg(test)] mod tests { + use std::time::Duration; + use clap::{CommandFactory, Parser}; - use super::{Cli, Commands, parse_env_temp_store_value, parse_env_u32_value, parse_env_u64_value}; - use agentic_core::config::{DEFAULT_SQLITE_MAX_CONNECTIONS, SqliteTempStore}; + use super::{ + Cli, Commands, parse_env_duration_value, parse_env_nonzero_u32_value, parse_env_optional_duration_value, + parse_env_temp_store_value, parse_env_u32_value, parse_env_u64_value, + }; + use agentic_core::config::{ + DEFAULT_POSTGRES_ACQUIRE_TIMEOUT_SECONDS, DEFAULT_POSTGRES_IDLE_TIMEOUT_SECONDS, + DEFAULT_POSTGRES_LOCK_TIMEOUT_SECONDS, DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS, + DEFAULT_POSTGRES_STATEMENT_TIMEOUT_SECONDS, DEFAULT_SQLITE_MAX_CONNECTIONS, SqliteTempStore, + }; #[test] fn serve_uses_common_args_before_subcommand() { @@ -325,4 +424,87 @@ mod tests { ); assert!(parse_env_temp_store_value(Ok("invalid".to_owned())).is_err()); } + + #[test] + fn postgres_timeout_parser_uses_defaults_and_allows_disabling_recycling() { + assert_eq!( + parse_env_duration_value( + "POSTGRES_ACQUIRE_TIMEOUT_SECONDS", + Err(std::env::VarError::NotPresent), + DEFAULT_POSTGRES_ACQUIRE_TIMEOUT_SECONDS, + ) + .expect("default acquire timeout"), + Duration::from_secs(DEFAULT_POSTGRES_ACQUIRE_TIMEOUT_SECONDS) + ); + assert_eq!( + parse_env_duration_value( + "POSTGRES_ACQUIRE_TIMEOUT_SECONDS", + Ok("9".to_owned()), + DEFAULT_POSTGRES_ACQUIRE_TIMEOUT_SECONDS, + ) + .expect("explicit acquire timeout"), + Duration::from_secs(9) + ); + assert!( + parse_env_duration_value( + "POSTGRES_ACQUIRE_TIMEOUT_SECONDS", + Ok("0".to_owned()), + DEFAULT_POSTGRES_ACQUIRE_TIMEOUT_SECONDS, + ) + .is_err() + ); + assert_eq!( + parse_env_optional_duration_value( + "POSTGRES_IDLE_TIMEOUT_SECONDS", + Ok("0".to_owned()), + DEFAULT_POSTGRES_IDLE_TIMEOUT_SECONDS, + ) + .expect("disabled idle timeout"), + None + ); + assert_eq!( + parse_env_optional_duration_value( + "POSTGRES_IDLE_TIMEOUT_SECONDS", + Ok("45".to_owned()), + DEFAULT_POSTGRES_IDLE_TIMEOUT_SECONDS, + ) + .expect("explicit idle timeout"), + Some(Duration::from_secs(45)) + ); + assert_eq!( + parse_env_duration_value( + "POSTGRES_LOCK_TIMEOUT_SECONDS", + Err(std::env::VarError::NotPresent), + DEFAULT_POSTGRES_LOCK_TIMEOUT_SECONDS, + ) + .expect("default lock timeout"), + Duration::from_secs(DEFAULT_POSTGRES_LOCK_TIMEOUT_SECONDS) + ); + assert_eq!( + parse_env_duration_value( + "POSTGRES_MIGRATION_TIMEOUT_SECONDS", + Err(std::env::VarError::NotPresent), + DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS, + ) + .expect("default migration timeout"), + Duration::from_secs(DEFAULT_POSTGRES_MIGRATION_TIMEOUT_SECONDS) + ); + assert_eq!( + parse_env_duration_value( + "POSTGRES_STATEMENT_TIMEOUT_SECONDS", + Err(std::env::VarError::NotPresent), + DEFAULT_POSTGRES_STATEMENT_TIMEOUT_SECONDS, + ) + .expect("default statement timeout"), + Duration::from_secs(DEFAULT_POSTGRES_STATEMENT_TIMEOUT_SECONDS) + ); + assert!( + parse_env_nonzero_u32_value( + "POSTGRES_MAX_CONNECTIONS", + Ok("0".to_owned()), + DEFAULT_SQLITE_MAX_CONNECTIONS, + ) + .is_err() + ); + } } diff --git a/crates/agentic-server/tests/cli_test.rs b/crates/agentic-server/tests/cli_test.rs index 9c85a1f0..413eb308 100644 --- a/crates/agentic-server/tests/cli_test.rs +++ b/crates/agentic-server/tests/cli_test.rs @@ -14,3 +14,19 @@ fn missing_llm_api_base_error_mentions_environment_and_flag() { "unexpected error message: {stderr}" ); } + +#[test] +fn help_does_not_expose_database_url_credentials() { + let database_url = "postgresql://agentic:super-secret@db.example/agentic"; + let output = Command::new(env!("CARGO_BIN_EXE_agentic-server")) + .env("DATABASE_URL", database_url) + .arg("--help") + .output() + .expect("agentic-server help must run"); + + assert!(output.status.success()); + let stdout = String::from_utf8(output.stdout).expect("stdout must be UTF-8"); + assert!(stdout.contains("DATABASE_URL")); + assert!(!stdout.contains(database_url)); + assert!(!stdout.contains("super-secret")); +} diff --git a/crates/agentic-server/tests/common/mod.rs b/crates/agentic-server/tests/common/mod.rs index 983ae205..17f8ff11 100644 --- a/crates/agentic-server/tests/common/mod.rs +++ b/crates/agentic-server/tests/common/mod.rs @@ -21,6 +21,7 @@ pub fn test_config(llm_url: &str) -> Config { llm_ready_interval_s: 0.1, skip_llm_ready_check: false, db_url: None, + postgres: agentic_core::config::PostgresConfig::default(), sqlite: agentic_core::config::SqliteConfig::default(), } } diff --git a/crates/agentic-server/tests/responses_test.rs b/crates/agentic-server/tests/responses_test.rs index cd78fa10..d7020794 100644 --- a/crates/agentic-server/tests/responses_test.rs +++ b/crates/agentic-server/tests/responses_test.rs @@ -214,8 +214,8 @@ async fn test_store_false_proxies_sse_to_vllm() { } #[tokio::test] -async fn test_store_true_reaches_executor_not_proxy() { - // Arrange — mock vLLM returns 200, but executor path will fail at storage layer +async fn test_store_true_returns_error_when_persistence_fails() { + // Arrange — mock vLLM returns 200, but the executor cannot persist into the disabled test store. let (llm_url, _h1) = spawn_mock_vllm_json().await; let (gw_url, _h2) = spawn_gateway(test_state(&test_config(&llm_url))).await; @@ -227,14 +227,35 @@ async fn test_store_true_reaches_executor_not_proxy() { .await .unwrap(); - // Assert — executor path reached: executor assigns a resp_-prefixed id - assert_eq!(resp.status(), 200); - let body: serde_json::Value = resp.json().await.unwrap(); - let id = body["id"].as_str().unwrap_or(""); - assert!( - id.starts_with("resp_"), - "expected executor-assigned id starting with resp_, got: {id}" - ); + // Assert — a stored request never reports success without durable state. + assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR); + let body = resp.text().await.unwrap(); + assert!(body.contains("storage"), "{body}"); +} + +#[tokio::test] +async fn test_streaming_store_true_emits_error_without_completed_event_when_persistence_fails() { + let (llm_url, _h1) = spawn_mock_vllm_sse().await; + let (gw_url, _h2) = spawn_gateway(test_state(&test_config(&llm_url))).await; + + let resp = reqwest::Client::new() + .post(format!("{gw_url}/v1/responses")) + .json(&serde_json::json!({ + "model": "test", + "input": [{"type": "message", "role": "user", "content": "hi"}], + "store": true, + "stream": true + })) + .send() + .await + .unwrap(); + + assert_eq!(resp.status(), StatusCode::OK); + let body = resp.text().await.unwrap(); + assert!(body.contains("\"type\":\"error\""), "{body}"); + assert!(body.contains("storage not configured or disabled"), "{body}"); + assert!(body.contains("data: [DONE]"), "{body}"); + assert!(!body.contains("\"type\":\"response.completed\""), "{body}"); } #[tokio::test] diff --git a/docs/deploying/container.md b/docs/deploying/container.md index ecb736df..862631c7 100644 --- a/docs/deploying/container.md +++ b/docs/deploying/container.md @@ -29,6 +29,13 @@ The image starts `agentic-server` in standalone mode. At minimum, set `LLM_API_B | `GATEWAY_HOST` | `0.0.0.0` | Listen address | | `GATEWAY_PORT` | `9000` | Listen port | | `DATABASE_URL` | `sqlite://./agentic_api.db` | SQLite or PostgreSQL persistence URL | +| `POSTGRES_MAX_CONNECTIONS` | `10` | Maximum PostgreSQL connections per gateway replica | +| `POSTGRES_ACQUIRE_TIMEOUT_SECONDS` | `30` | Maximum wait for a PostgreSQL pool connection | +| `POSTGRES_LOCK_TIMEOUT_SECONDS` | `5` | Maximum wait for a PostgreSQL row or table lock | +| `POSTGRES_MIGRATION_TIMEOUT_SECONDS` | `300` | Lock and statement timeout while running startup migrations | +| `POSTGRES_STATEMENT_TIMEOUT_SECONDS` | `30` | Maximum runtime for normal PostgreSQL statements | +| `POSTGRES_IDLE_TIMEOUT_SECONDS` | `600` | Recycle idle PostgreSQL connections; `0` disables | +| `POSTGRES_MAX_LIFETIME_SECONDS` | `1800` | Recycle PostgreSQL connections after this lifetime; `0` disables | | `OPENAI_API_KEY` | none | Credential sent to the upstream service when the client does not supply one | | `SKIP_LLM_READY_CHECK` | `false` | Skip the startup probe for hosted providers without `/health` | | `CORS_ALLOWED_ORIGINS` | none | Comma-separated browser origins | @@ -50,6 +57,40 @@ The gateway does not provide inbound client authentication. `OPENAI_API_KEY` is If the upstream is running on the Docker host, use `http://host.docker.internal:` on Docker Desktop. On Linux, add `--add-host host.docker.internal:host-gateway`. +### PostgreSQL production settings + +Use the TLS settings required by the managed database provider. Prefer certificate and hostname verification when the provider supplies a CA certificate: + +```console +DATABASE_URL='postgresql://agentic-api:password@postgres.example.com/agentic_api?sslmode=verify-full&sslrootcert=/run/secrets/postgres-ca.pem' +``` + +`sslmode=require` encrypts the connection but does not verify the server hostname. Mount private CA certificates and client keys from runtime secrets; do not copy them into the image. + +Size the pool across the whole deployment, not one process. Keep `replicas * POSTGRES_MAX_CONNECTIONS` below the managed database connection limit, with capacity reserved for migrations, administration, and failover. The acquire timeout bounds how long a request waits when the pool is exhausted. Lock and statement timeouts prevent a stalled transaction, slow query, or hot conversation from holding a connection indefinitely. Idle and lifetime recycling protect against stale connections and can be disabled with `0` only when the provider recommends it. + +Each gateway replica runs the embedded SQLx migrations during startup. PostgreSQL advisory locking serializes concurrent migration attempts, so replica startup is repeatable and only proceeds after the schema is ready. Migration lock waits and statements use the finite migration timeout rather than the shorter application lock timeout. A migration failure or timeout prevents that replica from serving traffic. + +The first PostgreSQL startup on this release widens timestamp and sequence columns in place to 64-bit integers. PostgreSQL takes exclusive table locks for these changes, which prevents concurrent writes from being lost but can briefly pause older replicas. It also rebuilds the conversation sequence index. For an existing large database, schedule the first upgraded replica during a maintenance window and set the migration timeout for the expected lock, rewrite, and index duration. + +Drain replicas running an older release before enabling writes through this release. Older replicas do not take the per-conversation row lock and can allocate duplicate sequence numbers if they write alongside upgraded replicas. + +Stored requests now fail if their response or conversation state cannot be persisted. For streaming requests, the gateway sends an error event instead of `response.completed`. This prevents clients from receiving a response ID that cannot be continued after a lock timeout or other database failure. + +`AGENTIC_API_SCHEMA_READY` keeps schema changes under supervisor control. Startup performs a read-only compatibility check and fails if the four persistence columns still need widening. Apply this upgrade with a DDL-capable migration role before starting the DML-only gateway role: + +```sql +ALTER TABLE conversations + ALTER COLUMN created_at TYPE BIGINT USING created_at::BIGINT; +ALTER TABLE items + ALTER COLUMN created_at TYPE BIGINT USING created_at::BIGINT, + ALTER COLUMN seq TYPE BIGINT USING seq::BIGINT; +ALTER TABLE responses + ALTER COLUMN created_at TYPE BIGINT USING created_at::BIGINT; +``` + +Enable automated backups and point-in-time recovery in the managed database service according to the deployment's recovery-point and recovery-time requirements. Test restores separately; gateway replicas do not create database backups. + ## Smoke test The liveness probe reports whether the gateway process is serving traffic. Readiness also checks the upstream inference service.