diff --git a/openspec/changes/bound-operation-query-deadlines/.openspec.yaml b/openspec/changes/bound-operation-query-deadlines/.openspec.yaml new file mode 100644 index 0000000..cbd245e --- /dev/null +++ b/openspec/changes/bound-operation-query-deadlines/.openspec.yaml @@ -0,0 +1,2 @@ +schema: spec-driven +created: 2026-09-20 diff --git a/openspec/changes/bound-operation-query-deadlines/proposal.md b/openspec/changes/bound-operation-query-deadlines/proposal.md new file mode 100644 index 0000000..d76c5c4 --- /dev/null +++ b/openspec/changes/bound-operation-query-deadlines/proposal.md @@ -0,0 +1,27 @@ +# Bound operation-ledger database waits + +## Problem + +The durable operation coordinator issues database statements directly on the +SurrealDB client. The client connection has a server-query timeout, but these +callers do not have an application deadline while waiting for the shared +connection. In the deployed service, reconciliation stopped making progress +while both `/health` and `/ready` remained healthy. + +## Change + +- Apply the configured query timeout to every direct operation-ledger database + future. +- Report the failed database stage and deadline when the bound expires. +- Keep embedding and executor work outside this deadline so timeout + cancellation cannot abandon an in-flight executor protocol request. + +## Non-goals + +- Changing the storage retry policy or SDK query timeout. +- Timing out a complete operation or embedding request. + +## Capability + +- `operation-query-deadlines`: direct operation-ledger database waits are + bounded on the production API path. diff --git a/openspec/changes/bound-operation-query-deadlines/specs/operation-query-deadlines/spec.md b/openspec/changes/bound-operation-query-deadlines/specs/operation-query-deadlines/spec.md new file mode 100644 index 0000000..e249c6f --- /dev/null +++ b/openspec/changes/bound-operation-query-deadlines/specs/operation-query-deadlines/spec.md @@ -0,0 +1,25 @@ +## Purpose + +Prevent a direct operation-ledger database wait from freezing durable operation +reconciliation while the service remains healthy. + +## ADDED Requirements + +### Requirement: Direct operation-ledger database waits are bounded + +The operation service SHALL apply the configured query timeout to every direct +operation-ledger database future. The deadline MUST NOT enclose embedding or +executor protocol work. + +#### Scenario: A receipt query does not complete within the configured deadline + +- **WHEN** a receipt request reaches the production API router +- **AND** its direct database future does not complete within the configured query timeout +- **THEN** the request returns an error naming the database stage and elapsed deadline +- **AND** the operation coordinator remains able to process later work + +#### Scenario: The coordinator performs embedding work + +- **WHEN** an operation invokes the supervised embedding executor +- **THEN** the operation database deadline is not active around that executor request +- **AND** executor completion remains governed by the executor watchdog contract diff --git a/openspec/changes/bound-operation-query-deadlines/tasks.md b/openspec/changes/bound-operation-query-deadlines/tasks.md new file mode 100644 index 0000000..da03ae9 --- /dev/null +++ b/openspec/changes/bound-operation-query-deadlines/tasks.md @@ -0,0 +1,5 @@ +## 1. Bound operation-ledger database waits + +- [x] 1.1 Carry the configured query timeout into the production operation service. +- [x] 1.2 Apply the deadline to every direct operation-ledger database future. +- [ ] 1.3 Pass formatting, focused integration coverage, compilation, strict OpenSpec validation, deployed backlog recovery, and review. diff --git a/src/api/mod.rs b/src/api/mod.rs index e12bcab..051bf51 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -18,7 +18,7 @@ use axum::{Json, http::StatusCode}; use axum::{Router, routing::get}; use serde::Serialize; use serde_json::json; -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; use tower_http::cors::CorsLayer; use tower_http::trace::TraceLayer; @@ -90,9 +90,18 @@ pub fn build_router( storage: Arc, embedding_service: Arc, ) -> Router { - let operations = crate::operations::OperationService::start( + build_router_with_query_timeout(storage, embedding_service, Duration::from_secs(10)) +} + +pub fn build_router_with_query_timeout( + storage: Arc, + embedding_service: Arc, + query_timeout: Duration, +) -> Router { + let operations = crate::operations::OperationService::start_with_query_timeout( Arc::clone(&storage), Arc::clone(&embedding_service), + query_timeout, ); let state = AppState { storage: Arc::clone(&storage), diff --git a/src/main.rs b/src/main.rs index ca22a41..aa00589 100644 --- a/src/main.rs +++ b/src/main.rs @@ -103,6 +103,7 @@ async fn main() -> Result<()> { let config = load_config().await?; let embedding_service = init_embedding_service(&config).await?; let retry_config = parse_retry_config_from_env(); + let operation_query_timeout = std::time::Duration::from_millis(retry_config.query_timeout_ms); let mlx_backend = matches!( &config.embedding_provider, @@ -153,8 +154,15 @@ async fn main() -> Result<()> { // ── Axum REST API + HTTP/SSE MCP ───────────────────────────────────────── let api_storage = Arc::clone(&storage); - let api_handle = - tokio::spawn(async move { run_api_server(api_storage, api_port, health_embedding).await }); + let api_handle = tokio::spawn(async move { + run_api_server( + api_storage, + api_port, + health_embedding, + operation_query_timeout, + ) + .await + }); // ── MCP stdio ───────────────────────────────────────────────────────────── let enable_stdio_mcp = std::env::var("MCP_STDIO") @@ -596,10 +604,12 @@ async fn run_api_server( storage: Arc, port: u16, embedding_service: Arc, + operation_query_timeout: std::time::Duration, ) -> Result<()> { let addr = std::net::SocketAddr::from(([0, 0, 0, 0], port)); tracing::info!("🌐 Starting REST API + HTTP MCP server on http://{}", addr); - let router = api::build_router(storage, embedding_service); + let router = + api::build_router_with_query_timeout(storage, embedding_service, operation_query_timeout); let listener = tokio::net::TcpListener::bind(addr) .await .context("Failed to bind REST API port")?; diff --git a/src/operations.rs b/src/operations.rs index 1685885..04e3c10 100644 --- a/src/operations.rs +++ b/src/operations.rs @@ -8,8 +8,10 @@ use std::{ collections::{HashSet, VecDeque}, convert::Infallible, + future::IntoFuture, pin::Pin, sync::Arc, + time::Duration, }; use anyhow::{Context, Result}; @@ -321,6 +323,7 @@ impl From for OperationEvent { pub struct OperationService { storage: Arc, embedding_service: Arc, + query_timeout: Duration, wake_tx: mpsc::Sender, events_tx: broadcast::Sender, } @@ -337,12 +340,21 @@ impl OperationService { storage: Arc, embedding_service: Arc, ) -> Self { - Self::start_with_capacities(storage, embedding_service, 256, 1024) + Self::start_with_query_timeout(storage, embedding_service, Duration::from_secs(10)) + } + + pub fn start_with_query_timeout( + storage: Arc, + embedding_service: Arc, + query_timeout: Duration, + ) -> Self { + Self::start_with_capacities(storage, embedding_service, query_timeout, 256, 1024) } fn start_with_capacities( storage: Arc, embedding_service: Arc, + query_timeout: Duration, wake_capacity: usize, event_capacity: usize, ) -> Self { @@ -351,6 +363,7 @@ impl OperationService { let service = Self { storage, embedding_service, + query_timeout, wake_tx, events_tx, }; @@ -370,6 +383,22 @@ impl OperationService { .context("durable operations require SurrealStorage") } + async fn await_database(&self, stage: &'static str, future: F) -> Result + where + F: IntoFuture>, + anyhow::Error: From, + { + tokio::time::timeout(self.query_timeout, future.into_future()) + .await + .with_context(|| { + format!( + "operation database {stage} timed out after {}ms", + self.query_timeout.as_millis() + ) + })? + .map_err(anyhow::Error::from) + } + pub async fn submit( &self, request: OperationRequest, @@ -429,20 +458,23 @@ impl OperationService { .map_err(SubmitError::Storage)? .db() .map_err(SubmitError::Storage)?; - let response = db - .query( - "BEGIN TRANSACTION;\n\ + let response = self + .await_database( + "submit", + db.query( + "BEGIN TRANSACTION;\n\ CREATE type::record('memory_operation', $key) CONTENT $operation;\n\ CREATE type::record('memory_operation_event', $event_key) CONTENT $event;\n\ COMMIT TRANSACTION;", + ) + .bind(("key", key)) + .bind(("operation", db_record)) + .bind(("event_key", event_key)) + .bind(("event", event)), ) - .bind(("key", key)) - .bind(("operation", db_record)) - .bind(("event_key", event_key)) - .bind(("event", event)) .await; - if let Err(error) = response.and_then(|response| response.check()) { + if let Err(error) = response.and_then(|response| response.check().map_err(Into::into)) { // A concurrent submit may have won the unique-index race. Re-read // the authoritative row and apply the same hash rule. if let Some(existing) = self @@ -480,9 +512,12 @@ impl OperationService { pub async fn get(&self, operation_id: &str) -> Result> { let db = self.surreal()?.db()?; - let mut rows: Vec = db - .query(GET_OPERATION_RECEIPT_QUERY) - .bind(("id", operation_id.to_owned())) + let mut rows: Vec = self + .await_database( + "receipt lookup", + db.query(GET_OPERATION_RECEIPT_QUERY) + .bind(("id", operation_id.to_owned())), + ) .await? .check()? .take(0)?; @@ -491,15 +526,20 @@ impl OperationService { async fn get_db(&self, operation_id: &str) -> Result> { let db = self.surreal()?.db()?; - Ok(db - .select(("memory_operation", record_key(operation_id))) - .await?) + self.await_database( + "operation lookup", + db.select(("memory_operation", record_key(operation_id))), + ) + .await } async fn list_nonterminal_ids(&self) -> Result> { let db = self.surreal()?.db()?; - let rows: Vec = db - .query(LIST_NONTERMINAL_OPERATION_IDS_QUERY) + let rows: Vec = self + .await_database( + "reconciliation discovery", + db.query(LIST_NONTERMINAL_OPERATION_IDS_QUERY), + ) .await? .check()? .take(0)?; @@ -508,12 +548,15 @@ impl OperationService { async fn events_after(&self, operation_id: &str, sequence: u64) -> Result> { let db = self.surreal()?.db()?; - let rows: Vec = db - .query( + let rows: Vec = self + .await_database( + "event history lookup", + db.query( "SELECT * FROM memory_operation_event WHERE operation_id = $id AND sequence > $sequence ORDER BY sequence ASC", ) .bind(("id", operation_id.to_owned())) - .bind(("sequence", sequence)) + .bind(("sequence", sequence)), + ) .await? .check()? .take(0)?; @@ -548,7 +591,9 @@ impl OperationService { }; let event_key = format!("{}-{sequence:016}", record_key(operation_id)); let db = self.surreal()?.db()?; - db.query( + self.await_database( + "state transition", + db.query( "BEGIN TRANSACTION;\n\ UPDATE memory_operation SET state = $state, blocked_by = $blocked_by, result = $result, error = $error, progress_seq = $sequence, updated_at = $now WHERE operation_id = $id;\n\ CREATE type::record('memory_operation_event', $event_key) CONTENT $event;\n\ @@ -562,7 +607,8 @@ impl OperationService { .bind(("now", now)) .bind(("id", operation_id.to_owned())) .bind(("event_key", event_key)) - .bind(("event", event)) + .bind(("event", event)), + ) .await? .check()?; let published = OperationEvent { @@ -595,11 +641,14 @@ impl OperationService { async fn operation_parts(&self, operation_id: &str) -> Result> { let db = self.surreal()?.db()?; - let parts: Vec = db - .query( + let parts: Vec = self + .await_database( + "part lookup", + db.query( "SELECT * FROM memory_operation_part WHERE operation_id = $id ORDER BY part_index ASC", ) - .bind(("id", operation_id.to_owned())) + .bind(("id", operation_id.to_owned())), + ) .await? .check()? .take(0)?; @@ -642,12 +691,15 @@ impl OperationService { } }) .collect::>(); - db.query( - "BEGIN TRANSACTION;\n\ + self.await_database( + "plan persistence", + db.query( + "BEGIN TRANSACTION;\n\ INSERT INTO memory_operation_part $parts;\n\ COMMIT TRANSACTION;", + ) + .bind(("parts", rows)), ) - .bind(("parts", rows)) .await? .check()?; Ok(()) @@ -660,13 +712,16 @@ impl OperationService { embedding: Vec, ) -> Result<()> { let db = self.surreal()?.db()?; - db.query( + self.await_database( + "part persistence", + db.query( "UPDATE memory_operation_part SET state = 'indexed', embedding = $embedding, updated_at = $now WHERE operation_id = $id AND part_index = $index", ) .bind(("embedding", embedding)) .bind(("now", Datetime::default())) .bind(("id", operation_id.to_owned())) - .bind(("index", part_index)) + .bind(("index", part_index)), + ) .await? .check()?; Ok(()) @@ -744,13 +799,15 @@ impl OperationService { event.generation, event.progress_seq ); - self.surreal()? - .db()? - .query("CREATE type::record('memory_executor_event', $key) CONTENT $event") - .bind(("key", key)) - .bind(("event", row)) - .await? - .check()?; + let db = self.surreal()?.db()?; + self.await_database( + "executor event persistence", + db.query("CREATE type::record('memory_executor_event', $key) CONTENT $event") + .bind(("key", key)) + .bind(("event", row)), + ) + .await? + .check()?; Ok(()) } @@ -761,9 +818,10 @@ impl OperationService { else { return Ok(()); }; - self.surreal()? - .db()? - .query( + let db = self.surreal()?.db()?; + self.await_database( + "executor snapshot persistence", + db.query( "UPDATE memory_operation SET executor_generation = $generation, executor_progress_seq = $progress, executor_exit_count = $exit_count, executor_last_exit = $last_exit, executor_error = $executor_error, updated_at = $now WHERE operation_id = $id", ) .bind(("generation", snapshot.generation)) @@ -772,7 +830,8 @@ impl OperationService { .bind(("last_exit", snapshot.last_exit)) .bind(("executor_error", snapshot.error)) .bind(("now", Datetime::default())) - .bind(("id", operation_id.to_owned())) + .bind(("id", operation_id.to_owned())), + ) .await? .check()?; Ok(()) @@ -2017,6 +2076,7 @@ mod tests { let service = OperationService::start_with_capacities( Arc::clone(&storage) as Arc, embedder, + Duration::from_secs(10), 4, 16, ); @@ -2055,6 +2115,7 @@ mod tests { let service = OperationService::start_with_capacities( Arc::clone(&storage) as Arc, embedder, + Duration::from_secs(10), 4, 16, ); @@ -2131,6 +2192,7 @@ mod tests { let service = OperationService::start_with_capacities( storage as Arc, embedder, + Duration::from_secs(10), 16, 4, ); diff --git a/tests/operation_query_deadline.rs b/tests/operation_query_deadline.rs new file mode 100644 index 0000000..e0b347b --- /dev/null +++ b/tests/operation_query_deadline.rs @@ -0,0 +1,233 @@ +use std::{ + net::{Ipv4Addr, SocketAddrV4, TcpListener}, + process::{Child, Command, Stdio}, + sync::Arc, + time::Duration, +}; + +use axum::{ + body::{Body, to_bytes}, + http::{Request, StatusCode, header}, +}; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; +use surreal_memory::{ + EmbeddingService, MemoryStorage, RetryConfig, SurrealConfig, SurrealStorage, + storage::surreal::SurrealMode, +}; +use surreal_memory_server::api; +use tower::ServiceExt; + +struct NoOpEmbedder; + +struct ServerGuard(Child); + +impl Drop for ServerGuard { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} + +async fn start_server() -> (ServerGuard, String) { + let listener = TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0)).unwrap(); + let address = listener.local_addr().unwrap(); + drop(listener); + let child = Command::new("surreal") + .args([ + "start", + "--no-banner", + "--unauthenticated", + "--allow-all", + "--bind", + &address.to_string(), + "memory", + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .expect("installed surreal CLI starts the isolated fixture"); + let guard = ServerGuard(child); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if tokio::net::TcpStream::connect(address).await.is_ok() { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .expect("isolated surreal fixture becomes reachable"); + (guard, format!("ws://{address}")) +} + +#[async_trait::async_trait] +impl EmbeddingService for NoOpEmbedder { + async fn embed(&self, _text: &str) -> anyhow::Result> { + Ok(vec![0.0; 1536]) + } + + async fn embed_batch(&self, texts: Vec) -> anyhow::Result>> { + Ok(texts.into_iter().map(|_| vec![0.0; 1536]).collect()) + } + + fn dimensions(&self) -> usize { + 1536 + } +} + +#[tokio::test] +async fn receipt_timeout_leaves_the_same_coordinator_able_to_commit_later_work() { + let (_server, endpoint) = start_server().await; + let embedder: Arc = Arc::new(NoOpEmbedder); + let storage = Arc::new( + SurrealStorage::new( + &SurrealConfig { + mode: SurrealMode::Server, + endpoint: Some(endpoint), + embedded_path: None, + username: None, + password: None, + namespace: format!("deadline_{}", uuid::Uuid::new_v4().simple()), + database: "operations".to_owned(), + retry: RetryConfig { + max_connect_retries: 0, + query_timeout_ms: 1_000, + ..RetryConfig::default() + }, + }, + Arc::clone(&embedder), + ) + .await + .expect("isolated server-mode SurrealStorage"), + ); + let database = storage.db().unwrap().clone(); + let large_receipt_key = Sha256::digest(b"large-receipt") + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + let empty_payload = json!({}); + let empty_payload_hash = Sha256::digest(serde_json::to_vec(&empty_payload).unwrap()) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + database + .query( + "CREATE type::record('memory_operation', $key) CONTENT { + operation_id: $id, + schema_version: 2, + kind: 'create_task_stream', + dependencies: [], + payload_hash: $payload_hash, + payload: $payload, + state: 'committed', + blocked_by: [], + result: $result, + error: NONE, + executor_generation: 0, + executor_progress_seq: 0, + executor_exit_count: 0, + executor_last_exit: NONE, + executor_error: NONE, + progress_seq: 1, + created_at: time::now(), + updated_at: time::now() + }", + ) + .bind(("key", large_receipt_key)) + .bind(("id", "large-receipt".to_owned())) + .bind(("payload_hash", empty_payload_hash)) + .bind(("payload", empty_payload)) + .bind(("result", json!({"blob":"x".repeat(32 * 1024 * 1024)}))) + .await + .unwrap() + .check() + .unwrap(); + let router = api::build_router_with_query_timeout( + Arc::clone(&storage) as Arc, + embedder, + Duration::from_millis(50), + ); + + let timeout_response = router + .clone() + .oneshot( + Request::builder() + .uri("/api/v2/operations/large-receipt") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(timeout_response.status(), StatusCode::INTERNAL_SERVER_ERROR); + let timeout_body: Value = serde_json::from_slice( + &to_bytes(timeout_response.into_body(), usize::MAX) + .await + .unwrap(), + ) + .unwrap(); + assert_eq!( + timeout_body["error"], + "operation database receipt lookup timed out after 50ms" + ); + tokio::time::sleep(Duration::from_millis(1_100)).await; + + let payload = json!({ + "name": "deadline-probe", + "description": "proves coordinator recovery after query cancellation", + "agent_id": null, + "user_id": "test" + }); + let payload_hash = Sha256::digest(serde_json::to_vec(&payload).unwrap()) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + let request_body = json!({ + "operation_id": "deadline-probe", + "schema_version": 2, + "kind": "create_task_stream", + "dependencies": [], + "payload_hash": payload_hash, + "payload": payload + }); + + let response = router + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/v2/operations") + .header(header::CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_vec(&request_body).unwrap())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::ACCEPTED); + + let receipt = tokio::time::timeout(Duration::from_secs(2), async { + loop { + let response = router + .clone() + .oneshot( + Request::builder() + .uri("/api/v2/operations/deadline-probe") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body: Value = + serde_json::from_slice(&to_bytes(response.into_body(), usize::MAX).await.unwrap()) + .unwrap(); + if body["state"] == "committed" { + break body; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the original coordinator commits later work"); + assert_eq!(receipt["operation_id"], "deadline-probe"); +}