diff --git a/crates/astra-messaging/src/db_transport.rs b/crates/astra-messaging/src/db_transport.rs index b81049a0c8..b1981aa4ba 100644 --- a/crates/astra-messaging/src/db_transport.rs +++ b/crates/astra-messaging/src/db_transport.rs @@ -67,7 +67,7 @@ use std::time::Duration; use async_trait::async_trait; use sqlx::{MySql, Pool, Row, query}; -use tokio::sync::{RwLock, mpsc, oneshot, watch}; +use tokio::sync::{Notify, RwLock, mpsc, oneshot, watch}; use tracing::Instrument; use super::transport::{MessageStream, MessageTransport}; @@ -348,6 +348,7 @@ pub struct DatabaseTransport { struct PollTaskControl { subscription_id: String, abort_handle: tokio::task::AbortHandle, + wake: Arc, } /// Default visibility timeout (how long before an unclaimed message reappears). @@ -638,6 +639,38 @@ impl DatabaseTransport { Ok(()) } + + fn wake_local_consumer(&self, consumer_id: &str) { + let controls = self + .poll_abort_handles + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(control) = controls.get(consumer_id) { + control.wake.notify_one(); + } + } + + async fn wake_local_broadcast_consumers(&self, delegation_id: &str) { + let consumer_ids = self + .registrations + .read() + .await + .iter() + .filter(|(_, registered_delegation_id)| { + registered_delegation_id.as_deref() == Some(delegation_id) + }) + .map(|(address, _)| format!("{}@{}", address.agent_id, address.run_id)) + .collect::>(); + let controls = self + .poll_abort_handles + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + for consumer_id in consumer_ids { + if let Some(control) = controls.get(&consumer_id) { + control.wake.notify_one(); + } + } + } } #[async_trait] @@ -739,6 +772,7 @@ impl MessageTransport for DatabaseTransport { } let (tx, rx) = mpsc::channel(LOCAL_DELIVERY_BUFFER_CAPACITY); let (start_tx, start_rx) = oneshot::channel(); + let wake = Arc::new(Notify::new()); let subscription_id = uuid::Uuid::new_v4().to_string(); let span_delegation_id = delegation_id.clone().unwrap_or_default(); let poll_task = tokio::spawn( @@ -756,6 +790,7 @@ impl MessageTransport for DatabaseTransport { self.instance_id.clone(), subscription_id.clone(), Arc::clone(&self.poll_abort_handles), + Arc::clone(&wake), start_rx, ) .instrument(tracing::info_span!( @@ -775,6 +810,7 @@ impl MessageTransport for DatabaseTransport { PollTaskControl { subscription_id, abort_handle: poll_task.abort_handle(), + wake, }, ); let _ = start_tx.send(()); @@ -854,6 +890,9 @@ impl MessageTransport for DatabaseTransport { self.metrics .messages_sent .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + if let MessageTarget::Direct { address } = &msg.to { + self.wake_local_consumer(&format!("{}@{}", address.agent_id, address.run_id)); + } Ok(()) } Err(e) => { @@ -887,6 +926,7 @@ impl MessageTransport for DatabaseTransport { self.metrics .messages_sent .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.wake_local_broadcast_consumers(delegation_id).await; Ok(()) } Err(e) => { @@ -977,6 +1017,7 @@ async fn poll_loop( instance_id: String, subscription_id: String, poll_abort_handles: Arc>>, + wake: Arc, start_rx: oneshot::Receiver<()>, ) { if start_rx.await.is_err() { @@ -1058,29 +1099,64 @@ async fn poll_loop( // 1. Claim direct messages atomically (UPDATE then SELECT). // This ensures no two consumers process the same direct message. - let now_ms = chrono::Utc::now().timestamp_millis(); - let claim_token = uuid::Uuid::new_v4().to_string(); - let claim_result = query( - "UPDATE agent_message_queue - SET status = 'claimed', claimed_by = ?, claimed_at_ms = ?, claim_token = ?, - attempt_count = attempt_count + 1 + // + // Most mailbox polls are idle. MatrixOne still has to plan and lock + // the ordered UPDATE when it matches no rows, which turns an idle + // mailbox into sustained write pressure. Probe the covering direct + // index first; a message racing an empty probe is observed by the next + // poll exactly as it was when it raced an empty UPDATE. The UPDATE + // remains the sole claim authority when work exists, so competing + // consumers retain the existing claim-token fencing semantics. + let direct_pending = query( + "SELECT 1 AS pending + FROM agent_message_queue WHERE to_run_id = ? AND to_agent_id = ? AND status = 'pending' - ORDER BY created_at ASC, message_id ASC LIMIT ?", + ORDER BY created_at ASC, message_id ASC LIMIT 1", ) - .bind(&consumer_id) - .bind(now_ms) - .bind(&claim_token) .bind(&addr.run_id) .bind(&addr.agent_id) - .bind(POLL_BATCH_SIZE) - .execute(&pool) + .fetch_optional(&pool) .await; - match claim_result { - Ok(result) if result.rows_affected() > 0 => { - had_activity = true; - // Fetch the messages we just claimed. - let fetch_result = query( + let direct_pending = match direct_pending { + Ok(row) => row.is_some(), + Err(error) => { + had_error = true; + metrics + .poll_errors + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + tracing::warn!(target: "astra_runtime::messaging::db_transport", + " ⚠ messaging: direct pending probe error for {}@{}: {:?}", + addr.agent_id, addr.run_id, error + ); + false + } + }; + + if direct_pending { + let now_ms = chrono::Utc::now().timestamp_millis(); + let claim_token = uuid::Uuid::new_v4().to_string(); + let claim_result = query( + "UPDATE agent_message_queue + SET status = 'claimed', claimed_by = ?, claimed_at_ms = ?, claim_token = ?, + attempt_count = attempt_count + 1 + WHERE to_run_id = ? AND to_agent_id = ? AND status = 'pending' + ORDER BY created_at ASC, message_id ASC LIMIT ?", + ) + .bind(&consumer_id) + .bind(now_ms) + .bind(&claim_token) + .bind(&addr.run_id) + .bind(&addr.agent_id) + .bind(POLL_BATCH_SIZE) + .execute(&pool) + .await; + + match claim_result { + Ok(result) if result.rows_affected() > 0 => { + had_activity = true; + // Fetch the messages we just claimed. + let fetch_result = query( "SELECT message_id, payload_json FROM agent_message_queue WHERE to_run_id = ? AND to_agent_id = ? AND status = 'claimed' AND claimed_by = ? AND claim_token = ? @@ -1094,162 +1170,163 @@ async fn poll_loop( .fetch_all(&pool) .await; - if let Ok(rows) = fetch_result { - for row in rows { - let message_id: Option = row.try_get("message_id").ok(); - if message_id.is_none() { - had_error = true; - metrics - .poll_errors - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - if let Err(e) = release_direct_claimed_batch_for_consumer_in_pool( - &pool, - &consumer_id, - &claim_token, - max_delivery_attempts, - ) - .await - { - tracing::warn!(target: "astra_runtime::messaging::db_transport", - " ⚠ messaging: failed to release direct claimed batch after missing message_id for {}@{}: {:?}", - addr.agent_id, - addr.run_id, - e - ); - } - tracing::warn!(target: "astra_runtime::messaging::db_transport", - " ⚠ messaging: claimed direct row without message_id for {}@{}; released current claim batch", - addr.agent_id, - addr.run_id, - ); - break; - } - let json: String = match row.try_get("payload_json") { - Ok(j) => j, - Err(_) => { + if let Ok(rows) = fetch_result { + for row in rows { + let message_id: Option = row.try_get("message_id").ok(); + if message_id.is_none() { + had_error = true; metrics - .messages_dropped + .poll_errors .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - match mark_direct_failed_by_identity( + if let Err(e) = release_direct_claimed_batch_for_consumer_in_pool( &pool, - message_id.as_deref(), &consumer_id, + &claim_token, + max_delivery_attempts, ) .await { - Ok(()) => {} - Err(e) => { - had_error = true; - metrics - .poll_errors - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - tracing::warn!(target: "astra_runtime::messaging::db_transport", - " ⚠ messaging: failed to dead-letter undecodable direct row (message_id: {}) for {}@{}: {:?}", - message_id.as_deref().unwrap_or(""), - addr.agent_id, - addr.run_id, - e - ); - } + tracing::warn!(target: "astra_runtime::messaging::db_transport", + " ⚠ messaging: failed to release direct claimed batch after missing message_id for {}@{}: {:?}", + addr.agent_id, + addr.run_id, + e + ); } - continue; + tracing::warn!(target: "astra_runtime::messaging::db_transport", + " ⚠ messaging: claimed direct row without message_id for {}@{}; released current claim batch", + addr.agent_id, + addr.run_id, + ); + break; } - }; - - match serde_json::from_str::(&json) { - Ok(msg) if !msg.is_expired() => { - if tx.send(Arc::new(msg)).await.is_err() { + let json: String = match row.try_get("payload_json") { + Ok(j) => j, + Err(_) => { metrics - .poll_errors + .messages_dropped .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - if let Err(e) = release_claimed_for_consumer_in_pool( + match mark_direct_failed_by_identity( &pool, + message_id.as_deref(), &consumer_id, - max_delivery_attempts, ) .await { - tracing::warn!(target: "astra_runtime::messaging::db_transport", - " ⚠ messaging: failed to release direct claims after closed channel for {}@{}: {:?}", - addr.agent_id, addr.run_id, e - ); + Ok(()) => {} + Err(e) => { + had_error = true; + metrics + .poll_errors + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + tracing::warn!(target: "astra_runtime::messaging::db_transport", + " ⚠ messaging: failed to dead-letter undecodable direct row (message_id: {}) for {}@{}: {:?}", + message_id.as_deref().unwrap_or(""), + addr.agent_id, + addr.run_id, + e + ); + } } - break 'poll; + continue; } - metrics - .messages_received - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - } - Ok(_) | Err(_) => { - metrics - .messages_dropped - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - match mark_direct_failed_by_identity( - &pool, - message_id.as_deref(), - &consumer_id, - ) - .await - { - Ok(()) => {} - Err(e) => { - had_error = true; + }; + + match serde_json::from_str::(&json) { + Ok(msg) if !msg.is_expired() => { + if tx.send(Arc::new(msg)).await.is_err() { metrics .poll_errors .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - tracing::warn!(target: "astra_runtime::messaging::db_transport", - " ⚠ messaging: failed to dead-letter direct row (message_id: {}) for {}@{}: {:?}", - message_id.as_deref().unwrap_or(""), - addr.agent_id, - addr.run_id, - e - ); + if let Err(e) = release_claimed_for_consumer_in_pool( + &pool, + &consumer_id, + max_delivery_attempts, + ) + .await + { + tracing::warn!(target: "astra_runtime::messaging::db_transport", + " ⚠ messaging: failed to release direct claims after closed channel for {}@{}: {:?}", + addr.agent_id, addr.run_id, e + ); + } + break 'poll; + } + metrics + .messages_received + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + Ok(_) | Err(_) => { + metrics + .messages_dropped + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + match mark_direct_failed_by_identity( + &pool, + message_id.as_deref(), + &consumer_id, + ) + .await + { + Ok(()) => {} + Err(e) => { + had_error = true; + metrics + .poll_errors + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + tracing::warn!(target: "astra_runtime::messaging::db_transport", + " ⚠ messaging: failed to dead-letter direct row (message_id: {}) for {}@{}: {:?}", + message_id.as_deref().unwrap_or(""), + addr.agent_id, + addr.run_id, + e + ); + } } } } } - } - } else { - had_error = true; - metrics - .poll_errors - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - if let Err(e) = release_direct_claimed_batch_for_consumer_in_pool( - &pool, - &consumer_id, - &claim_token, - max_delivery_attempts, - ) - .await - { + } else { + had_error = true; + metrics + .poll_errors + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + if let Err(e) = release_direct_claimed_batch_for_consumer_in_pool( + &pool, + &consumer_id, + &claim_token, + max_delivery_attempts, + ) + .await + { + tracing::warn!(target: "astra_runtime::messaging::db_transport", + " ⚠ messaging: failed to release direct claimed batch after fetch error for {}@{}: {:?}", + addr.agent_id, + addr.run_id, + e + ); + } tracing::warn!(target: "astra_runtime::messaging::db_transport", - " ⚠ messaging: failed to release direct claimed batch after fetch error for {}@{}: {:?}", + " ⚠ messaging: direct fetch error for {}@{}: {:?}", addr.agent_id, addr.run_id, - e + fetch_result.unwrap_err() ); } + } + Ok(_) => { + // No pending messages — normal idle. + } + Err(e) => { + had_error = true; + metrics + .poll_errors + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); tracing::warn!(target: "astra_runtime::messaging::db_transport", - " ⚠ messaging: direct fetch error for {}@{}: {:?}", - addr.agent_id, - addr.run_id, - fetch_result.unwrap_err() + " ⚠ messaging: direct claim error for {}@{}: {:?}", + addr.agent_id, addr.run_id, e ); } } - Ok(_) => { - // No pending messages — normal idle. - } - Err(e) => { - had_error = true; - metrics - .poll_errors - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - tracing::warn!(target: "astra_runtime::messaging::db_transport", - " ⚠ messaging: direct claim error for {}@{}: {:?}", - addr.agent_id, addr.run_id, e - ); - } } // 2. Poll broadcast messages (if in a delegation group). @@ -1415,6 +1492,7 @@ async fn poll_loop( // Wait with shutdown awareness. tokio::select! { _ = tokio::time::sleep(sleep_duration) => {} + _ = wake.notified() => {} _ = shutdown_rx.changed() => { break; } @@ -2063,6 +2141,34 @@ mod tests { })); } + #[tokio::test] + async fn local_consumer_wake_interrupts_idle_poll_wait() { + let pool = MySqlPoolOptions::new().connect_lazy_with(MySqlConnectOptions::new()); + let transport = DatabaseTransport::new(pool); + let consumer_id = "agent-a@run-a"; + let wake = Arc::new(Notify::new()); + let poll_task = tokio::spawn(std::future::pending::<()>()); + transport + .poll_abort_handles + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .insert( + consumer_id.to_string(), + PollTaskControl { + subscription_id: "subscription-a".to_string(), + abort_handle: poll_task.abort_handle(), + wake: Arc::clone(&wake), + }, + ); + + transport.wake_local_consumer(consumer_id); + + tokio::time::timeout(Duration::from_millis(50), wake.notified()) + .await + .expect("local wake must retain a permit until the poll loop observes it"); + poll_task.abort(); + } + #[test] fn queue_cleanup_sql_uses_ordered_bounded_batches() { for (name, sql) in [ diff --git a/crates/astra-turn-core/src/pipeline/session.rs b/crates/astra-turn-core/src/pipeline/session.rs index ebcdd1174c..8aabc77cd8 100644 --- a/crates/astra-turn-core/src/pipeline/session.rs +++ b/crates/astra-turn-core/src/pipeline/session.rs @@ -31,6 +31,7 @@ use crate::recovery_state::RecoveryState; use crate::session_latches::SessionLatches; use crate::shadow_diff::{ShadowDiffResult, diff_pipeline_outputs}; use crate::working_memory::WorkingMemoryState; +use std::sync::Arc; /// Per-turn input provided by the agentic loop to `PipelineSession::run_turn()`. pub struct TurnInput<'a> { @@ -77,6 +78,7 @@ pub struct ShadowTurnOutput { /// `run_turn()` before each LLM request and `record_feedback()` after. pub struct PipelineSession { pipeline: ContextPipeline, + static_sections: Option>, pub stats: PipelineStats, pub latches: SessionLatches, pub emergent: EmergentContext, @@ -168,6 +170,7 @@ impl PipelineSession { ) -> Self { Self { pipeline: ContextPipeline::new(config), + static_sections: None, stats: PipelineStats::default(), latches: SessionLatches::default(), emergent: EmergentContext::default(), @@ -193,6 +196,7 @@ impl PipelineSession { ) -> Self { Self { pipeline: ContextPipeline::new(config), + static_sections: None, stats, latches: SessionLatches::default(), emergent: EmergentContext::default(), @@ -224,6 +228,7 @@ impl PipelineSession { ) -> Self { Self { pipeline: ContextPipeline::new(config), + static_sections: None, stats, latches, emergent: EmergentContext::default(), @@ -246,6 +251,18 @@ impl PipelineSession { self.turns_completed } + /// Return the immutable prompt sections owned by this pipeline session, + /// building them once on first use. The returned `Arc` lets the caller + /// borrow the sections while mutably advancing the rest of the session. + pub fn static_sections_or_init( + &mut self, + init: impl FnOnce() -> StaticSections, + ) -> Arc { + self.static_sections + .get_or_insert_with(|| Arc::new(init())) + .clone() + } + /// Run the pipeline for one turn. Returns the serialized provider request /// and associated metadata, or an abort if the session is in an /// unrecoverable error state. @@ -847,6 +864,7 @@ impl PipelineSession { Self { pipeline: ContextPipeline::new(config), + static_sections: None, stats, latches, emergent, @@ -1078,6 +1096,32 @@ mod tests { assert_eq!(actual.cache_eligible_tokens, 7_225); } + #[test] + fn static_sections_cache_is_scoped_to_pipeline_session() { + let mut first_session = PipelineSession::new(PipelineConfig::default()); + let first = first_session.static_sections_or_init(|| { + let mut sections = StaticSections::test_default(); + sections.core_rules.text = "first session".into(); + sections + }); + let reused = first_session.static_sections_or_init(|| { + panic!("a pipeline session must build static sections only once") + }); + + assert!(Arc::ptr_eq(&first, &reused)); + assert_eq!(reused.core_rules.text, "first session"); + + let mut second_session = PipelineSession::new(PipelineConfig::default()); + let second = second_session.static_sections_or_init(|| { + let mut sections = StaticSections::test_default(); + sections.core_rules.text = "second session".into(); + sections + }); + + assert!(!Arc::ptr_eq(&first, &second)); + assert_eq!(second.core_rules.text, "second session"); + } + fn test_external() -> ExternalSources { ExternalSources { memory_entries: vec![], diff --git a/crates/runtime/src/data_layer/storage.rs b/crates/runtime/src/data_layer/storage.rs index b875666cb6..ca4d3b35dc 100644 --- a/crates/runtime/src/data_layer/storage.rs +++ b/crates/runtime/src/data_layer/storage.rs @@ -3,12 +3,12 @@ pub use astra_services::storage::*; use std::time::Duration; use serde_json::Value; -use sqlx::{MySql, query}; +use sqlx::{MySql, QueryBuilder, Row, query}; use astra_core::canonical_names::{ metadata_duration_ms, metadata_tool_call_id, metadata_tool_name, }; -use astra_core::matrixone_statement_with_null_shape; +use astra_core::{matrixone_null_shape_comment, matrixone_statement_with_null_shape}; use astra_turn_core::contracts::{ TurnCoreEventRecord, TurnDecisionAuditRecord, TurnSkillSelectionRecord, TurnToolEventRecord, }; @@ -211,68 +211,148 @@ fn trace_event_insert_values(event: &TraceEvent) -> Result, +) -> Vec { + let mut new_ids = std::collections::BTreeSet::new(); + events + .iter() + .enumerate() + .filter_map(|(index, event)| { + (!existing_ids.contains(&event.event_id) && new_ids.insert(&event.event_id)) + .then_some(index) + }) + .collect() +} + +pub(crate) async fn insert_trace_events( tx: &mut sqlx::Transaction<'_, MySql>, - event: &TraceEvent, -) -> Result { - let values = trace_event_insert_values(event)?; - let insert_sql = matrixone_statement_with_null_shape( - "INSERT IGNORE INTO agent_events \ - (event_id, session_id, user_id, agent_id, agent_version, event_type, content, \ - parent_event_id, causal_chain_id, run_id, parent_run_id, turn_id, turn_seq, \ - round_index, tool_call_id, parent_agent_id, trace_kind, token_usage, \ - llm_model_used, reasoning_content, token_input, token_output, token_total, \ - meta_tool_name, meta_duration_ms, metadata, created_at) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - values.nullable_shape(event), + events: &[TraceEvent], +) -> Result<(u64, Option), sqlx::Error> { + let Some(first) = events.first() else { + return Ok((0, None)); + }; + if events + .iter() + .any(|event| event.user_id != first.user_id || event.session_id != first.session_id) + { + return Err(sqlx::Error::Protocol( + "trace event batch must belong to one user session".to_string(), + )); + } + + // `admit_session_event_write` holds the session write fence for this + // transaction before this function is called. Resolve replayed identities + // once, then retain only the first occurrence of each new event id. This + // preserves one batched insert while making the returned tail identify the + // final row actually inserted by this delta. + let mut existing_query = QueryBuilder::::new( + "SELECT event_id, session_id FROM agent_events WHERE user_id = ", ); - let result = query(&insert_sql) - .bind(&event.event_id) - .bind(&event.session_id) - .bind(&event.user_id) - .bind(event.agent_id.as_deref().unwrap_or("astra-server")) - .bind(env!("CARGO_PKG_VERSION")) - .bind(&event.event_type) - .bind(&event.content) - .bind(&event.parent_event_id) - .bind(&event.causal_chain_id) - .bind(&event.run_id) - .bind(&event.parent_run_id) - .bind(&event.turn_id) - .bind(event.turn_seq) - .bind(event.round_index) - .bind(&event.tool_call_id) - .bind(&event.parent_agent_id) - .bind(&event.trace_kind) - .bind(values.token_usage_json) - .bind(&event.llm_model_used) - .bind(&event.reasoning_content) - .bind(values.token_input) - .bind(values.token_output) - .bind(values.token_total) - .bind(&event.meta_tool_name) - .bind(event.meta_duration_ms) - .bind(values.metadata_json) - .bind(values.created_at) - .execute(&mut **tx) - .await?; - let inserted = result.rows_affected() > 0; - if inserted { - insert_agent_event_edges( - &mut **tx, - &event.user_id, - &event.session_id, - &event.event_id, - event.parent_event_id.as_deref(), - event - .parent_event_id - .as_ref() - .map(|id| std::slice::from_ref(id)) - .unwrap_or(&[]), - ) - .await?; + existing_query + .push_bind(&first.user_id) + .push(" AND event_id IN ("); + { + let mut ids = existing_query.separated(", "); + for event in events { + ids.push_bind(&event.event_id); + } } - Ok(inserted) + existing_query.push(")"); + let existing_rows = existing_query.build().fetch_all(&mut **tx).await?; + let mut existing_ids = std::collections::BTreeSet::new(); + for row in existing_rows { + let event_id = row.try_get::("event_id")?; + let existing_session_id = row.try_get::("session_id")?; + if existing_session_id != first.session_id { + return Err(sqlx::Error::Protocol(format!( + "trace event id {event_id} already belongs to another session" + ))); + } + existing_ids.insert(event_id); + } + let prepared = new_trace_event_indices(events, &existing_ids) + .into_iter() + .map(|index| trace_event_insert_values(&events[index]).map(|values| (index, values))) + .collect::, _>>()?; + let (inserted, last_inserted_event_id) = if prepared.is_empty() { + (0, None) + } else { + let mut insert = QueryBuilder::::new( + "INSERT INTO agent_events \ + (event_id, session_id, user_id, agent_id, agent_version, event_type, content, \ + parent_event_id, causal_chain_id, run_id, parent_run_id, turn_id, turn_seq, \ + round_index, tool_call_id, parent_agent_id, trace_kind, token_usage, \ + llm_model_used, reasoning_content, token_input, token_output, token_total, \ + meta_tool_name, meta_duration_ms, metadata, created_at) ", + ); + insert.push_values(prepared.iter(), |mut row, (index, values)| { + let event = &events[*index]; + row.push_bind(&event.event_id) + .push_bind(&event.session_id) + .push_bind(&event.user_id) + .push_bind(event.agent_id.as_deref().unwrap_or("astra-server")) + .push_bind(env!("CARGO_PKG_VERSION")) + .push_bind(&event.event_type) + .push_bind(&event.content) + .push_bind(&event.parent_event_id) + .push_bind(&event.causal_chain_id) + .push_bind(&event.run_id) + .push_bind(&event.parent_run_id) + .push_bind(&event.turn_id) + .push_bind(event.turn_seq) + .push_bind(event.round_index) + .push_bind(&event.tool_call_id) + .push_bind(&event.parent_agent_id) + .push_bind(&event.trace_kind) + .push_bind(&values.token_usage_json) + .push_bind(&event.llm_model_used) + .push_bind(&event.reasoning_content) + .push_bind(values.token_input) + .push_bind(values.token_output) + .push_bind(values.token_total) + .push_bind(&event.meta_tool_name) + .push_bind(event.meta_duration_ms) + .push_bind(&values.metadata_json) + .push_bind(&values.created_at); + }); + insert.push(matrixone_null_shape_comment( + prepared + .iter() + .flat_map(|(index, values)| values.nullable_shape(&events[*index])), + )); + let inserted = insert.build().execute(&mut **tx).await?.rows_affected(); + if inserted != prepared.len() as u64 { + return Err(sqlx::Error::Protocol(format!( + "trace event batch inserted {inserted} rows, expected {} under the session write fence", + prepared.len() + ))); + } + let last_inserted_event_id = prepared + .last() + .map(|(index, _)| events[*index].event_id.clone()); + (inserted, last_inserted_event_id) + }; + + // Both immutable event rows and their edge projection are idempotent + // inserts. Sending the complete deterministic batch lets MatrixOne report + // the number of newly inserted events and removes the preceding existence + // read. Replayed edges are ignored by their unique key, while a partial + // replay still repairs any missing edge rows. + let edge_inputs = events + .iter() + .map(|event| astra_services::storage::AgentEventEdgeInsert { + user_id: &event.user_id, + session_id: &event.session_id, + child_event_id: &event.event_id, + primary_parent_event_id: event.parent_event_id.as_deref(), + parent_event_ids: event.parent_event_id.as_slice(), + }) + .collect::>(); + astra_services::storage::insert_agent_event_edges_batch(&mut **tx, &edge_inputs).await?; + + Ok((inserted, last_inserted_event_id)) } pub(crate) async fn insert_core_turn_event( @@ -429,7 +509,7 @@ pub(crate) async fn insert_turn_decision_audit( mod tests { use super::{ INSERT_CORE_TURN_EVENT_SQL, core_turn_event_insert_values, metadata_string, - metadata_tool_name, trace_event_insert_values, + metadata_tool_name, new_trace_event_indices, trace_event_insert_values, }; use astra_core::matrixone_statement_with_null_shape; use astra_turn_core::contracts::TurnCoreEventRecord; @@ -562,6 +642,31 @@ mod tests { assert_eq!(persisted["total"], 22); } + #[test] + fn trace_batch_identifies_the_last_new_event_across_mixed_replays() { + let event = |id| TraceEvent::new(id, "session-1", "user-1", "trace", "runtime"); + + let existing_first = std::collections::BTreeSet::from(["existing".to_string()]); + let existing_then_new = vec![event("existing"), event("new")]; + assert_eq!( + new_trace_event_indices(&existing_then_new, &existing_first), + vec![1] + ); + + let new_then_existing = vec![event("new"), event("existing")]; + assert_eq!( + new_trace_event_indices(&new_then_existing, &existing_first), + vec![0] + ); + + let repeated_new = vec![event("first"), event("last"), event("first")]; + assert_eq!( + new_trace_event_indices(&repeated_new, &std::collections::BTreeSet::new()), + vec![0, 1], + "the repeated tail id must not replace the final row actually inserted" + ); + } + #[test] fn trace_event_statement_identity_includes_causal_chain_nullness() { let mut event = TraceEvent::new( diff --git a/crates/runtime/src/server/agent_binding_skill_runtime.rs b/crates/runtime/src/server/agent_binding_skill_runtime.rs index ff907dbe6f..7a2319398d 100644 --- a/crates/runtime/src/server/agent_binding_skill_runtime.rs +++ b/crates/runtime/src/server/agent_binding_skill_runtime.rs @@ -690,6 +690,66 @@ pub(crate) async fn prepare_agent_binding_skill_resolver( build_resolver_from_catalogs(server_id, endpoint_url, authorization, catalogs) } +pub(crate) fn prepare_agent_binding_skill_resolver_from_snapshot( + server_id: &str, + endpoint_url: &str, + authorization: &str, + agent_binding_ids: &[String], + snapshot_catalogs: &[astra_services::runs::RuntimeCapabilitySkillCatalogSnapshotRequest], +) -> Result)> { + validate_skill_endpoint(endpoint_url)?; + let mut catalogs_by_binding = HashMap::with_capacity(snapshot_catalogs.len()); + for catalog in snapshot_catalogs { + if catalog.agent_binding_id.trim().is_empty() + || catalog.agent_binding_id != catalog.agent_binding_id.trim() + || catalogs_by_binding + .insert(catalog.agent_binding_id.clone(), &catalog.skills) + .is_some() + { + return Err(skill_error( + StatusCode::BAD_REQUEST, + "runtime capability discovery snapshot has an invalid or duplicate agent_binding_id", + "agent_binding_discovery_snapshot_invalid", + )); + } + } + if catalogs_by_binding.len() != agent_binding_ids.len() + || catalogs_by_binding + .keys() + .any(|id| !agent_binding_ids.contains(id)) + { + return Err(skill_error( + StatusCode::BAD_REQUEST, + "runtime capability discovery snapshot does not match the Agent Binding Set", + "agent_binding_discovery_snapshot_invalid", + )); + } + let mut catalogs = Vec::with_capacity(agent_binding_ids.len()); + for agent_binding_id in agent_binding_ids { + let Some(raw_skills) = catalogs_by_binding.get(agent_binding_id) else { + return Err(skill_error( + StatusCode::BAD_REQUEST, + "runtime capability discovery snapshot is missing an Agent Binding skill catalog", + "agent_binding_discovery_snapshot_invalid", + )); + }; + let mut skills = Vec::with_capacity(raw_skills.len()); + for raw_skill in *raw_skills { + let skill = + serde_json::from_value::(raw_skill.clone()).map_err(|error| { + skill_error( + StatusCode::BAD_REQUEST, + format!("runtime capability discovery Skill is invalid: {error}"), + "agent_binding_discovery_snapshot_invalid", + ) + })?; + skills.push(skill); + } + catalogs.push((Some(agent_binding_id.clone()), skills)); + } + build_resolver_from_catalogs(server_id, endpoint_url, authorization, catalogs) +} + pub(crate) async fn prepare_runtime_skill_resolver( server_id: &str, endpoint_url: &str, @@ -720,6 +780,31 @@ mod tests { } } + #[test] + fn prepares_agent_binding_skills_from_frozen_snapshot() { + let prepared = prepare_agent_binding_skill_resolver_from_snapshot( + "skills", + "https://catalog.example.test/api/v1/skills/http", + "Bearer runtime-grant", + &["binding_1".to_string()], + &[ + astra_services::runs::RuntimeCapabilitySkillCatalogSnapshotRequest { + agent_binding_id: "binding_1".to_string(), + skills: vec![json!({ + "name": "pdf", + "description": "Read PDF files", + "allowed_tools": ["file_download"] + })], + }, + ], + ) + .expect("frozen Skill discovery snapshot should prepare"); + + assert_eq!(prepared.catalogs.len(), 1); + assert_eq!(prepared.catalogs[0].agent_binding_id, "binding_1"); + assert!(prepared.resolver.is_some()); + } + #[test] fn build_resolver_maps_discovered_skill_to_lazy_catalog_entry() { let resolver = build_resolver( diff --git a/crates/runtime/src/server/chat_handlers.rs b/crates/runtime/src/server/chat_handlers.rs index 655da60954..aabe0b41a4 100644 --- a/crates/runtime/src/server/chat_handlers.rs +++ b/crates/runtime/src/server/chat_handlers.rs @@ -58,8 +58,13 @@ pub(super) async fn validate_conversation_authority( "execution_grant_verifier_unavailable", ) })?; - let active_lease = coordinator - .load_active_writer(&authority.key) + // The authority lease and canonical head must come from one database + // snapshot. Besides avoiding two serialized round trips, this prevents + // validating the grant against a lease and cursor observed at different + // instants. The run lifecycle still revalidates and reserves the turn + // immediately before execution. + let admission_snapshot = coordinator + .load_admission_snapshot(&authority.key) .await .map_err(|error| { error_response_coded( @@ -67,14 +72,14 @@ pub(super) async fn validate_conversation_authority( format!("failed to load current session authority: {error}"), "session_authority_unavailable", ) - })? - .ok_or_else(|| { - error_response_coded( - StatusCode::CONFLICT, - "conversation authority no longer owns the active writer lease", - "conversation_authority_fenced", - ) })?; + let active_lease = admission_snapshot.active_writer.ok_or_else(|| { + error_response_coded( + StatusCode::CONFLICT, + "conversation authority no longer owns the active writer lease", + "conversation_authority_fenced", + ) + })?; let now_unix_ms = chrono::Utc::now().timestamp_millis(); let claims = signer .verify( @@ -103,16 +108,7 @@ pub(super) async fn validate_conversation_authority( "conversation_authority_fenced", )); } - let head = coordinator - .load_head(&authority.key) - .await - .map_err(|error| { - error_response_coded( - StatusCode::SERVICE_UNAVAILABLE, - format!("failed to load canonical conversation head: {error}"), - "session_head_unavailable", - ) - })?; + let head = admission_snapshot.head; if head.as_ref().map(|head| &head.cursor) != authority.expected_cursor.as_ref() || authority.prompt_manifest_root.as_deref() != head.as_ref().map(|head| head.latest_manifest_root.as_str()) diff --git a/crates/runtime/src/server/completions.rs b/crates/runtime/src/server/completions.rs index a4c450bbcd..b9b52d252b 100644 --- a/crates/runtime/src/server/completions.rs +++ b/crates/runtime/src/server/completions.rs @@ -910,33 +910,83 @@ mod tests { .is_cancelled() ); - tokio::time::timeout(std::time::Duration::from_secs(5), async { - let mut poll = tokio::time::interval(std::time::Duration::from_millis(20)); - loop { - poll.tick().await; - let row = sqlx::query( - "SELECT i.status AS invocation_status, a.status AS attempt_status + let cancellation_converged = + tokio::time::timeout(std::time::Duration::from_secs(5), async { + let mut poll = tokio::time::interval(std::time::Duration::from_millis(20)); + loop { + poll.tick().await; + let row = sqlx::query( + "SELECT i.status AS invocation_status, a.status AS attempt_status FROM inference_invocations i JOIN inference_provider_attempts a ON a.user_id = i.user_id AND a.invocation_id = i.invocation_id WHERE i.user_id = 'test-user' AND i.session_id = ? AND i.operation_id = 'completion_proxy:memory_extraction' AND i.logical_attempt = 3", - ) - .bind(&session_id) - .fetch_optional(pool) - .await - .expect("poll cancelled durable inference"); - if row.as_ref().is_some_and(|row| { - row.get::("invocation_status") == "delivery_unknown" - && row.get::("attempt_status") == "delivery_unknown" - }) { - break; + ) + .bind(&session_id) + .fetch_optional(pool) + .await + .expect("poll cancelled durable inference"); + if row.as_ref().is_some_and(|row| { + row.get::("invocation_status") == "delivery_unknown" + && row.get::("attempt_status") == "delivery_unknown" + }) { + break; + } } - } - }) - .await - .expect("detached settlement must converge after caller cancellation"); + }) + .await; + if let Err(error) = cancellation_converged { + let invocation_status = sqlx::query_scalar::<_, String>( + "SELECT status FROM inference_invocations + WHERE user_id = 'test-user' AND session_id = ? + AND operation_id = 'completion_proxy:memory_extraction' + AND logical_attempt = 3", + ) + .bind(&session_id) + .fetch_optional(pool) + .await + .expect("load cancelled invocation diagnostic"); + let attempt_status = sqlx::query_scalar::<_, String>( + "SELECT a.status FROM inference_provider_attempts a + JOIN inference_invocations i + ON a.user_id = i.user_id AND a.invocation_id = i.invocation_id + WHERE i.user_id = 'test-user' AND i.session_id = ? + AND i.operation_id = 'completion_proxy:memory_extraction' + AND i.logical_attempt = 3", + ) + .bind(&session_id) + .fetch_optional(pool) + .await + .expect("load cancelled attempt diagnostic"); + let debt = sqlx::query( + "SELECT debt.terminal_status, debt.provider_delivery_state, + debt.reconciliation_status + FROM inference_invocation_settlement_debts debt + JOIN inference_invocations i + ON debt.user_id = i.user_id AND debt.invocation_id = i.invocation_id + WHERE i.user_id = 'test-user' AND i.session_id = ? + AND i.operation_id = 'completion_proxy:memory_extraction' + AND i.logical_attempt = 3", + ) + .bind(&session_id) + .fetch_optional(pool) + .await + .expect("load cancelled settlement debt diagnostic") + .map(|row| { + ( + row.get::("terminal_status"), + row.get::("provider_delivery_state"), + row.get::("reconciliation_status"), + ) + }); + panic!( + "detached settlement must converge after caller cancellation: {error}; \ + invocation_status={invocation_status:?}, attempt_status={attempt_status:?}, \ + debt={debt:?}" + ); + } for statement in [ "DELETE FROM inference_invocation_settlement_debts WHERE user_id = 'test-user' AND session_id = ?", diff --git a/crates/runtime/src/server/run/lifecycle/mod.rs b/crates/runtime/src/server/run/lifecycle/mod.rs index abef370565..bab2042c65 100644 --- a/crates/runtime/src/server/run/lifecycle/mod.rs +++ b/crates/runtime/src/server/run/lifecycle/mod.rs @@ -223,6 +223,18 @@ fn durable_settlement_fence_closed( control_terminal_settlement_committed == Some(true) || settlement_finished_committed } +async fn run_owner_fenced_terminal_projections( + owner_terminal_committed: bool, + projections: F, +) where + F: FnOnce() -> Fut, + Fut: std::future::Future, +{ + if owner_terminal_committed { + projections().await; + } +} + /// A normal loop completion is not a durable completion until every provider /// tool attempt has one canonical terminal class. Apply this before terminal /// events/status are derived so storage, stream clients, and receipts share a @@ -5161,13 +5173,18 @@ impl AgenticRunLifecycleService { session_id, astra_turn_types::DEFAULT_CONVERSATION_BRANCH_ID, ); - let head = coordinator.load_head(&key).await.map_err(|error| { - error_response_coded( - StatusCode::SERVICE_UNAVAILABLE, - format!("failed to load canonical session head: {error}"), - "session_head_unavailable", - ) - })?; + let admission_snapshot = + coordinator + .load_admission_snapshot(&key) + .await + .map_err(|error| { + error_response_coded( + StatusCode::SERVICE_UNAVAILABLE, + format!("failed to load canonical session head: {error}"), + "session_head_unavailable", + ) + })?; + let head = admission_snapshot.head; let prior_canonical_bytes = head.as_ref().map_or(0, |head| head.total_canonical_bytes); let current_bytes = fresh_request_admission_bytes(request).map_err(|error| { error_response_coded( @@ -5196,64 +5213,59 @@ impl AgenticRunLifecycleService { "weighted_session_admission_rejected", ) })?; - let distributed_permit = self - .distributed_weighted_admission - .as_ref() - .ok_or_else(|| { - error_response_coded( - StatusCode::SERVICE_UNAVAILABLE, - "cross-pod weighted session admission is unavailable", - "distributed_session_admission_unavailable", - ) - })? - .try_reserve( - &key, - weighted_work, - Duration::from_secs(15 * 60), - &format!("server-run:{run_id}:weighted-admission"), - ) - .await - .map_err(|error| { - let status = if matches!( - &error, - astra_services::DistributedAdmissionError::Capacity(_) - ) { - StatusCode::TOO_MANY_REQUESTS - } else { - StatusCode::SERVICE_UNAVAILABLE - }; - error_response_coded( - status, - format!("distributed weighted session admission rejected this turn: {error}"), - "distributed_session_admission_rejected", - ) - })?; - let prior_messages = match &head { - Some(head) => coordinator - .materialize(head) - .await - .map(|materialized| materialized.messages) - .map_err(|error| { + let distributed_admission = + self.distributed_weighted_admission + .as_ref() + .ok_or_else(|| { error_response_coded( - StatusCode::CONFLICT, - format!("canonical session materialization requires repair: {error}"), - "session_context_needs_repair", + StatusCode::SERVICE_UNAVAILABLE, + "cross-pod weighted session admission is unavailable", + "distributed_session_admission_unavailable", ) - })?, - None => Vec::new(), - }; - let (lease, release_writer_on_finish) = - if let Some(authority) = request.conversation_authority.as_ref() { - let active = coordinator - .load_active_writer(&key) + })?; + let admission_ttl = Duration::from_secs(15 * 60); + let distributed_reservation_id = format!("server-run:{run_id}:weighted-admission"); + let materialize_prior_messages = async { + match &head { + Some(head) => coordinator + .materialize(head) .await + .map(|materialized| materialized.messages) .map_err(|error| { error_response_coded( - StatusCode::SERVICE_UNAVAILABLE, - format!("failed to load canonical writer: {error}"), - "session_writer_unavailable", + StatusCode::CONFLICT, + format!("canonical session materialization requires repair: {error}"), + "session_context_needs_repair", ) - })? + }), + None => Ok(Vec::new()), + } + }; + let map_distributed_error = |error: astra_services::DistributedAdmissionError| { + let status = if matches!( + &error, + astra_services::DistributedAdmissionError::Capacity(_) + ) { + StatusCode::TOO_MANY_REQUESTS + } else { + StatusCode::SERVICE_UNAVAILABLE + }; + error_response_coded( + status, + format!("distributed weighted session admission rejected this turn: {error}"), + "distributed_session_admission_rejected", + ) + }; + + let (distributed_permit, prior_messages, lease, reservation, release_writer_on_finish) = + if let Some(authority) = request.conversation_authority.as_ref() { + // Request-supplied authority already owns a writer. Avoid + // speculatively installing its turn reservation until the + // independent distributed admission and history reads have + // both succeeded, because this caller must not release that + // externally owned writer on a sibling failure. + let active = admission_snapshot + .active_writer .filter(|lease| { lease.key == key && lease.lease_id == authority.execution_grant.claims.lease_id @@ -5270,19 +5282,61 @@ impl AgenticRunLifecycleService { "conversation_authority_fenced", ) })?; - (active, false) - } else { - let authority_epochs = coordinator - .load_authority_epochs(&key) + let distributed_reservation = distributed_admission.try_reserve( + &key, + weighted_work, + admission_ttl, + &distributed_reservation_id, + ); + let (distributed_permit, prior_messages) = + tokio::join!(distributed_reservation, materialize_prior_messages); + let distributed_permit = distributed_permit.map_err(map_distributed_error)?; + let prior_messages = match prior_messages { + Ok(messages) => messages, + Err(error) => { + let _ = distributed_permit.release().await; + return Err(error); + } + }; + let reservation = match coordinator + .reserve_turn( + &active, + head.as_ref().map(|head| &head.cursor), + admission_ttl, + &format!("server-run:{run_id}:turn"), + ) .await - .map_err(|error| { - error_response_coded( + { + Ok(astra_services::ReserveTurnOutcome::Reserved(reservation)) + | Ok(astra_services::ReserveTurnOutcome::AlreadyReserved(reservation)) => { + reservation + } + Ok(astra_services::ReserveTurnOutcome::Conflict { .. }) => { + let _ = distributed_permit.release().await; + return Err(error_response_coded( + StatusCode::CONFLICT, + "canonical session cursor changed before turn reservation", + "conversation_cursor_conflict", + )); + } + Err(error) => { + let _ = distributed_permit.release().await; + return Err(error_response_coded( StatusCode::SERVICE_UNAVAILABLE, - format!("failed to load canonical authority epochs: {error}"), - "session_authority_unavailable", - ) - })? - .unwrap_or_default(); + format!("failed to reserve canonical turn: {error}"), + "session_turn_reservation_unavailable", + )); + } + }; + ( + distributed_permit, + prior_messages, + active, + reservation, + false, + ) + } else { + let authority_epochs = admission_snapshot.authority_epochs; let actor = astra_turn_types::ActorContextV1::owner_user( user_id, format!("server-run:{run_id}"), @@ -5291,69 +5345,95 @@ impl AgenticRunLifecycleService { None, authority_epochs, ); - let acquired = coordinator - .acquire_writer( - &key, - head.as_ref().map(|head| &head.cursor), - &actor, - Duration::from_secs(15 * 60), - &format!("server-run:{run_id}:writer"), - ) - .await - .map_err(|error| { - error_response_coded( + // These three stores are independent at admission time. The + // database coordinator installs writer+turn facts atomically, + // while distributed capacity and immutable history are read in + // parallel. Any sibling failure releases facts acquired here. + let distributed_reservation = distributed_admission.try_reserve( + &key, + weighted_work, + admission_ttl, + &distributed_reservation_id, + ); + let writer_idempotency_key = format!("server-run:{run_id}:writer"); + let reservation_idempotency_key = format!("server-run:{run_id}:turn"); + let acquire_and_reserve = coordinator.acquire_writer_and_reserve_turn( + &key, + head.as_ref().map(|head| &head.cursor), + &actor, + admission_ttl, + &writer_idempotency_key, + &reservation_idempotency_key, + ); + let (distributed_result, prior_messages_result, canonical_result) = tokio::join!( + distributed_reservation, + materialize_prior_messages, + acquire_and_reserve, + ); + let canonical_outcome = match canonical_result { + Ok(outcome) => outcome, + Err(error) => { + if let Ok(distributed_permit) = distributed_result { + let _ = distributed_permit.release().await; + } + return Err(error_response_coded( StatusCode::SERVICE_UNAVAILABLE, - format!("failed to acquire canonical writer: {error}"), - "session_writer_unavailable", - ) - })?; - let lease = match acquired { - astra_services::AcquireWriterOutcome::Acquired(lease) - | astra_services::AcquireWriterOutcome::AlreadyAcquired(lease) => lease, - astra_services::AcquireWriterOutcome::Conflict { .. } => { + format!("failed to acquire canonical turn authority: {error}"), + "session_turn_reservation_unavailable", + )); + } + }; + let (lease, reservation) = match canonical_outcome { + astra_services::AcquireWriterAndReserveTurnOutcome::Ready { + lease, + reservation, + } => (lease, reservation), + astra_services::AcquireWriterAndReserveTurnOutcome::WriterConflict { + .. + } => { + if let Ok(distributed_permit) = distributed_result { + let _ = distributed_permit.release().await; + } return Err(error_response_coded( StatusCode::CONFLICT, "another controller owns this canonical session branch", "session_writer_conflict", )); } + astra_services::AcquireWriterAndReserveTurnOutcome::ReservationConflict { + lease, + .. + } => { + let _ = coordinator.release_writer(&lease).await; + if let Ok(distributed_permit) = distributed_result { + let _ = distributed_permit.release().await; + } + return Err(error_response_coded( + StatusCode::CONFLICT, + "canonical session cursor changed before turn reservation", + "conversation_cursor_conflict", + )); + } + }; + let distributed_permit = match distributed_result { + Ok(permit) => permit, + Err(error) => { + let _ = coordinator.release_writer(&lease).await; + return Err(map_distributed_error(error)); + } + }; + let prior_messages = match prior_messages_result { + Ok(messages) => messages, + Err(error) => { + let (_, _) = tokio::join!( + coordinator.release_writer(&lease), + distributed_permit.release(), + ); + return Err(error); + } }; - (lease, true) + (distributed_permit, prior_messages, lease, reservation, true) }; - let reservation_outcome = coordinator - .reserve_turn( - &lease, - head.as_ref().map(|head| &head.cursor), - Duration::from_secs(15 * 60), - &format!("server-run:{run_id}:turn"), - ) - .await; - let reservation = match reservation_outcome { - Err(error) => { - if release_writer_on_finish { - let _ = coordinator.release_writer(&lease).await; - } - let _ = distributed_permit.release().await; - return Err(error_response_coded( - StatusCode::SERVICE_UNAVAILABLE, - format!("failed to reserve canonical turn: {error}"), - "session_turn_reservation_unavailable", - )); - } - Ok(astra_services::ReserveTurnOutcome::Reserved(reservation)) - | Ok(astra_services::ReserveTurnOutcome::AlreadyReserved(reservation)) => reservation, - Ok(astra_services::ReserveTurnOutcome::Conflict { .. }) => { - if release_writer_on_finish { - let _ = coordinator.release_writer(&lease).await; - } - let _ = distributed_permit.release().await; - return Err(error_response_coded( - StatusCode::CONFLICT, - "canonical session cursor changed before turn reservation", - "conversation_cursor_conflict", - )); - } - }; let renewal_cancel = CancellationToken::new(); let heartbeat_cancel = renewal_cancel.clone(); let heartbeat_run_cancel = authority_loss_cancel; @@ -5465,7 +5545,7 @@ impl AgenticRunLifecycleService { let cursor = base.ok_or_else(|| { "canonical replacement is missing its admitted base cursor".to_string() })?; - if proof.base_root() != cursor.canonical_root_hash { + if proof.base_manifest_root() != cursor.canonical_root_hash { return Err( "canonical rewrite proof does not match the admitted base root".into(), ); @@ -9536,25 +9616,54 @@ impl AgenticRunLifecycleService { .iter() .map(|binding| binding.binding.id.clone()) .collect::>(); - // Tool and skill discovery are independent reads from the same - // provider runtime. Keep binding validation ahead of both calls, then - // overlap their network latency before constructing the shared prompt. - let (bundle, prepared_skills) = tokio::join!( - runtime_mcp::prepare_agent_binding_mcp_bundle( - &mcp_descriptor.id, - &mcp_endpoint_url, - &runtime_auth.authorization, - mcp_descriptor.semantic_read.as_ref(), - ), - agent_binding_skill_runtime::prepare_agent_binding_skill_resolver( - &skill_descriptor.id, - &skill_endpoint_url, - &runtime_auth.authorization, - &binding_ids, - ), - ); - let bundle = bundle?; - let prepared_skills = prepared_skills?; + let (bundle, prepared_skills) = + if let Some(snapshot) = descriptors.discovery_snapshot.as_ref() { + if snapshot.version + != astra_services::runs::RUNTIME_CAPABILITY_DISCOVERY_SNAPSHOT_VERSION + { + return Err(error_response_coded( + StatusCode::BAD_REQUEST, + "unsupported runtime capability discovery snapshot version", + "agent_binding_discovery_snapshot_invalid", + )); + } + ( + runtime_mcp::prepare_agent_binding_mcp_bundle_from_snapshot( + &mcp_descriptor.id, + &mcp_endpoint_url, + &runtime_auth.authorization, + mcp_descriptor.semantic_read.as_ref(), + &snapshot.tools, + )?, + agent_binding_skill_runtime::prepare_agent_binding_skill_resolver_from_snapshot( + &skill_descriptor.id, + &skill_endpoint_url, + &runtime_auth.authorization, + &binding_ids, + &snapshot.skill_catalogs, + )?, + ) + } else { + // Tool and skill discovery are independent reads from the same + // provider runtime. Keep binding validation ahead of both calls, + // then overlap their network latency before constructing the + // shared prompt. + let (bundle, prepared_skills) = tokio::join!( + runtime_mcp::prepare_agent_binding_mcp_bundle( + &mcp_descriptor.id, + &mcp_endpoint_url, + &runtime_auth.authorization, + mcp_descriptor.semantic_read.as_ref(), + ), + agent_binding_skill_runtime::prepare_agent_binding_skill_resolver( + &skill_descriptor.id, + &skill_endpoint_url, + &runtime_auth.authorization, + &binding_ids, + ), + ); + (bundle?, prepared_skills?) + }; let skill_resolver = apply_normalized_skill_allowlist(prepared_skills.resolver, request_constraints) .map_err(|detail| error_response(StatusCode::BAD_REQUEST, detail))?; @@ -16638,7 +16747,7 @@ impl RunLifecycleService for AgenticRunLifecycleService { // Drop event_tx — signals end-of-stream to the HTTP handler. drop(event_tx); - if owner_terminal_committed { + run_owner_fenced_terminal_projections(owner_terminal_committed, || async { persist_turn_evaluation_journal( &bg_user_id, &bg_session_id, @@ -16725,7 +16834,8 @@ impl RunLifecycleService for AgenticRunLifecycleService { ) .await; } - } + }) + .await; if let Some(guard) = bg_root_runtime_context_guard.as_mut() { guard.settle().await; } diff --git a/crates/runtime/src/server/run/lifecycle/persistence.rs b/crates/runtime/src/server/run/lifecycle/persistence.rs index f4cf2564a1..cb45fb56d7 100644 --- a/crates/runtime/src/server/run/lifecycle/persistence.rs +++ b/crates/runtime/src/server/run/lifecycle/persistence.rs @@ -317,11 +317,15 @@ impl PostLoopPersistContext { self.persist_csl_if_canonical_ready(state, canonical_context_persisted, &mut errors) .await; - // 2. Persist decision audit + skill selection to hook DB. Canonical - // per-call tool lifecycle events were already written atomically with - // core events above; no aggregate shadow events are emitted here. - if let Some(ref writer) = self.hook_db_writer { - if let Err(e) = persist_server_loop_hook_events( + // The remaining consumers all read the immutable completed loop state + // and write independent sinks. Streaming callers complete this phase + // before publishing terminal SSE; the writes are awaited together so + // their independent database latency remains overlapped. + let hook_persist = async { + let Some(writer) = self.hook_db_writer.as_ref() else { + return Ok(()); + }; + persist_server_loop_hook_events( writer.as_ref(), &self.user_id, &self.session_id, @@ -330,14 +334,12 @@ impl PostLoopPersistContext { self.model_name.as_deref(), ) .await - { - errors.push(format!("hook events persist failed: {}", e)); - } - } - - // 3. Fire Memoria observer (cross-session knowledge extraction). - if let Some(worker) = self.observer_worker.clone() { - if let Err(e) = fire_server_loop_observer( + }; + let observer_dispatch = async { + let Some(worker) = self.observer_worker.clone() else { + return Ok(()); + }; + fire_server_loop_observer( worker, &self.user_id, &self.session_id, @@ -345,34 +347,20 @@ impl PostLoopPersistContext { self.metrics_registry.clone(), ) .await - { - errors.push(format!("observer fire failed: {}", e)); - } - } - - // 4. Fire SessionEnd hooks. - crate::skills::hooks::fire_session_end( + }; + let session_end = crate::skills::hooks::fire_session_end( &state.skills.session_event_hooks, state.current_session_id.as_deref().unwrap_or(""), - ) - .await; - - // 5. Persist runtime promotion events. - if let Err(e) = persist_runtime_promotion_events( + ); + let promotion_persist = persist_runtime_promotion_events( &self.matrixone, self.shared_pool.as_ref(), &self.user_id, &self.session_id, &self.run_id, &state.telemetry.promotion_events, - ) - .await - { - errors.push(format!("promotion events persist failed: {}", e)); - } - - // 6. Persist web-agent state projection rows generated by the agentic loop. - if let Err(e) = persist_server_loop_projection_state( + ); + let projection_persist = persist_server_loop_projection_state( self.shared_pool.as_ref(), &self.user_id, &self.session_id, @@ -380,10 +368,25 @@ impl PostLoopPersistContext { self.agent_id.as_deref(), self.model_name.as_deref(), state, - ) - .await - { - errors.push(format!("projection state persist failed: {}", e)); + ); + let (hook_result, observer_result, (), promotion_result, projection_result) = tokio::join!( + hook_persist, + observer_dispatch, + session_end, + promotion_persist, + projection_persist, + ); + if let Err(error) = hook_result { + errors.push(format!("hook events persist failed: {error}")); + } + if let Err(error) = observer_result { + errors.push(format!("observer fire failed: {error}")); + } + if let Err(error) = promotion_result { + errors.push(format!("promotion events persist failed: {error}")); + } + if let Err(error) = projection_result { + errors.push(format!("projection state persist failed: {error}")); } // Use loop_success to conditionally log severity diff --git a/crates/runtime/src/server/run/lifecycle/tests.rs b/crates/runtime/src/server/run/lifecycle/tests.rs index d3df1fe2b0..99083238a6 100644 --- a/crates/runtime/src/server/run/lifecycle/tests.rs +++ b/crates/runtime/src/server/run/lifecycle/tests.rs @@ -601,8 +601,8 @@ fn admitted_proof_allows_successful_turn_to_normalize_prior_execution_scratch() json!({"role": "user", "content": "new request"}), json!({"role": "assistant", "content": "new result"}), ]); - let base_root = astra_turn_types::canonical_conversation_root(&prior); - let proof = CanonicalRewriteProof::new(&prior, &base_root, 0); + let base_manifest_root = "a".repeat(64); + let proof = CanonicalRewriteProof::from_materialized_admission(&prior, &base_manifest_root, 0); let (mode, packs) = canonical_commit_delta(&prior, true, &messages, Some(&proof), false) .unwrap() @@ -658,8 +658,14 @@ fn missing_proof_cannot_replace_prior_execution_scratch() { #[test] fn compaction_commits_the_complete_replacement_projection() { let prior = vec![json!({"role": "user", "content": "old"})]; - let base_root = astra_turn_types::canonical_conversation_root(&prior); - let mut proof = CanonicalRewriteProof::new(&prior, &base_root, 0); + let base_manifest_root = "a".repeat(64); + assert_ne!( + base_manifest_root, + astra_turn_types::canonical_conversation_root(&prior), + "the regression requires distinct manifest and conversation hash domains" + ); + let mut proof = + CanonicalRewriteProof::from_materialized_admission(&prior, &base_manifest_root, 0); let permit = proof.begin(&prior); let compacted = vec![ json!({"role": "user", "content": "summary"}), @@ -672,6 +678,7 @@ fn compaction_commits_the_complete_replacement_projection() { assert_eq!(mode, astra_turn_types::CanonicalDeltaModeV1::Replace); assert_eq!(packs.concat(), compacted); + assert_eq!(proof.base_manifest_root(), base_manifest_root); } #[test] @@ -825,8 +832,14 @@ fn typed_objective_relations_survive_real_tiered_compaction() { // Admit an actual prefix, then prove the real compaction rewrites it. let prior = messages[..7].to_vec(); - let root = astra_turn_types::canonical_conversation_root(&prior); - let mut proof = CanonicalRewriteProof::new(&prior, &root, 0); + let base_manifest_root = "a".repeat(64); + assert_ne!( + base_manifest_root, + astra_turn_types::canonical_conversation_root(&prior), + "the regression requires distinct manifest and conversation hash domains" + ); + let mut proof = + CanonicalRewriteProof::from_materialized_admission(&prior, &base_manifest_root, 0); let permit = proof.begin(&messages); let mut engine = crate::turn::CompactionEngine::new(); engine.add_layer(Box::new(crate::turn::cloud::TieredCompaction::new(2, 0.0))); @@ -905,8 +918,9 @@ fn unexplained_canonical_prefix_shrink_remains_rejected() { fn unrelated_prefix_mutation_after_compaction_is_rejected() { let prior = vec![json!({"role": "user", "content": "committed"})]; let compacted = vec![json!({"role": "system", "content": "summary"})]; - let base_root = astra_turn_types::canonical_conversation_root(&prior); - let mut proof = CanonicalRewriteProof::new(&prior, &base_root, 0); + let base_manifest_root = "a".repeat(64); + let mut proof = + CanonicalRewriteProof::from_materialized_admission(&prior, &base_manifest_root, 0); let permit = proof.begin(&prior); proof.finish(permit, &compacted, None); @@ -921,8 +935,9 @@ fn compaction_cannot_authorize_an_already_mutated_prefix() { let prior = vec![json!({"role": "user", "content": "committed"})]; let mutated_before_compaction = vec![json!({"role": "user", "content": "unrelated mutation"})]; let compacted = vec![json!({"role": "system", "content": "summary"})]; - let base_root = astra_turn_types::canonical_conversation_root(&prior); - let mut proof = CanonicalRewriteProof::new(&prior, &base_root, 0); + let base_manifest_root = "a".repeat(64); + let mut proof = + CanonicalRewriteProof::from_materialized_admission(&prior, &base_manifest_root, 0); let permit = proof.begin(&mutated_before_compaction); proof.finish(permit, &compacted, None); @@ -12785,6 +12800,7 @@ fn authorized_edge_dispatch_request() -> astra_services::runs::ChatRequestData { mcp: None, skills: None, edge_agent: Some(descriptor), + discovery_snapshot: None, }); request } @@ -12991,6 +13007,7 @@ async fn prepare_chat_request_normalizes_provider_descriptor_without_registered_ mcp: None, skills: None, edge_agent: None, + discovery_snapshot: None, }); let prepared = service @@ -13034,6 +13051,7 @@ async fn validate_request_constraints_rejects_descriptor_without_provider_author mcp: None, skills: None, edge_agent: None, + discovery_snapshot: None, }); let err = service @@ -15138,6 +15156,30 @@ fn terminal_events_for_persistence_keeps_only_terminal_lifecycle_events() { assert_eq!(persisted[6]["event_type"], "run_finished"); } +#[tokio::test] +async fn superseded_terminal_skips_all_owner_derived_projections() { + let journal_calls = Arc::new(AtomicUsize::new(0)); + let transcript_calls = Arc::new(AtomicUsize::new(0)); + let hook_calls = Arc::new(AtomicUsize::new(0)); + let observed = ( + Arc::clone(&journal_calls), + Arc::clone(&transcript_calls), + Arc::clone(&hook_calls), + ); + + // A Superseded terminal transition leaves owner_terminal_committed false. + run_owner_fenced_terminal_projections(false, move || async move { + observed.0.fetch_add(1, Ordering::SeqCst); + observed.1.fetch_add(1, Ordering::SeqCst); + observed.2.fetch_add(1, Ordering::SeqCst); + }) + .await; + + assert_eq!(journal_calls.load(Ordering::SeqCst), 0); + assert_eq!(transcript_calls.load(Ordering::SeqCst), 0); + assert_eq!(hook_calls.load(Ordering::SeqCst), 0); +} + #[test] fn terminal_handoff_event_uses_live_persistence_without_terminal_replay_duplication() { let event = json!({ diff --git a/crates/runtime/src/server/runtime_mcp.rs b/crates/runtime/src/server/runtime_mcp.rs index d111c7ecc1..75d03ba268 100644 --- a/crates/runtime/src/server/runtime_mcp.rs +++ b/crates/runtime/src/server/runtime_mcp.rs @@ -1507,6 +1507,59 @@ pub(crate) async fn prepare_agent_binding_mcp_bundle( "agent_binding_discovery_failed", ) })?; + build_agent_binding_mcp_bundle( + server_id, + endpoint_url, + authorization, + semantic_read_capability, + tools, + ) +} + +pub(crate) fn prepare_agent_binding_mcp_bundle_from_snapshot( + server_id: &str, + endpoint_url: &str, + authorization: &str, + semantic_read_capability: Option<&astra_services::runs::RuntimeSemanticReadCapabilityRequest>, + raw_tools: &[Value], +) -> Result)> { + let tool_namespace = sanitize_tool_name(server_id); + if tool_namespace.is_empty() { + return Err(mcp_error( + StatusCode::BAD_REQUEST, + "agent binding MCP server id must not be empty after sanitization", + "agent_binding_capability_ref_invalid", + )); + } + astra_services::validate_registered_endpoint_url( + "agent_binding.capability_server.endpoint_url", + endpoint_url, + "agent_binding_capability_ref_invalid", + )?; + let tools = parse_agent_binding_mcp_tools(json!({ "tools": raw_tools })).map_err(|error| { + agent_binding_mcp_error_response( + StatusCode::BAD_REQUEST, + error, + "agent_binding_discovery_snapshot_invalid", + ) + })?; + build_agent_binding_mcp_bundle( + server_id, + endpoint_url, + authorization, + semantic_read_capability, + tools, + ) +} + +fn build_agent_binding_mcp_bundle( + server_id: &str, + endpoint_url: &str, + authorization: &str, + semantic_read_capability: Option<&astra_services::runs::RuntimeSemanticReadCapabilityRequest>, + tools: Vec, +) -> Result)> { + let tool_namespace = sanitize_tool_name(server_id); let (_adapter_schemas, control_tools, stop_after_success_tools) = agent_binding_tools_to_schemas_checked(&tool_namespace, &tools) .map_err(|(error, code)| mcp_error(StatusCode::BAD_GATEWAY, error, code))?; @@ -1643,6 +1696,26 @@ pub(crate) async fn prepare_agent_binding_mcp_bundle( mod tests { use super::*; + #[test] + fn prepares_agent_binding_mcp_bundle_from_frozen_snapshot() { + let bundle = prepare_agent_binding_mcp_bundle_from_snapshot( + "tools", + "https://catalog.example.test/api/v1/mcp/http", + "Bearer runtime-grant", + None, + &[json!({ + "name": "file_list", + "description": "List files", + "inputSchema": {"type": "object", "properties": {}}, + "side_effect_class": "read" + })], + ) + .expect("frozen MCP discovery snapshot should prepare"); + + assert_eq!(bundle.schemas.len(), 1); + assert!(bundle.agent_binding_mcp.is_some()); + } + #[test] fn agent_binding_timeout_preserves_unknown_outcome_evidence() { let error = agent_binding_mcp_timeout_error("timed out"); diff --git a/crates/runtime/src/server/server_loop_host.rs b/crates/runtime/src/server/server_loop_host.rs index 54f3944a87..fd7b0ccfae 100644 --- a/crates/runtime/src/server/server_loop_host.rs +++ b/crates/runtime/src/server/server_loop_host.rs @@ -21613,11 +21613,14 @@ mod tests { })]; let durable_base = astra_turn_types::ProviderCanonicalWalBaseV2::from_messages(&durable).unwrap(); - let mut proof = crate::turn::canonical_commit::CanonicalRewriteProof::new( - &durable, - &durable_base.canonical.root_hash, - 3, - ); + let base_manifest_root = "a".repeat(64); + assert_ne!(base_manifest_root, durable_base.canonical.root_hash); + let mut proof = + crate::turn::canonical_commit::CanonicalRewriteProof::from_materialized_admission( + &durable, + &base_manifest_root, + 3, + ); let permit = proof.begin(&durable); let mut rewritten = durable.clone(); rewritten.push(json!({ diff --git a/crates/runtime/src/turn/agentic_loop/host.rs b/crates/runtime/src/turn/agentic_loop/host.rs index cb1559fa61..f9d53bb3fa 100644 --- a/crates/runtime/src/turn/agentic_loop/host.rs +++ b/crates/runtime/src/turn/agentic_loop/host.rs @@ -3499,15 +3499,16 @@ impl AgenticLoopState { pub(crate) fn initialize_canonical_rewrite_proof( &mut self, admitted_prefix: &[Value], - base_root: &str, + base_manifest_root: &str, base_compaction_generation: u64, ) { - self.canonical_rewrite_state.proof = - Some(crate::turn::canonical_commit::CanonicalRewriteProof::new( + self.canonical_rewrite_state.proof = Some( + crate::turn::canonical_commit::CanonicalRewriteProof::from_materialized_admission( admitted_prefix, - base_root, + base_manifest_root, base_compaction_generation, - )); + ), + ); } pub(crate) fn initialize_provider_canonical_wal_base(&mut self, durable_prefix: &[Value]) { diff --git a/crates/runtime/src/turn/canonical_commit.rs b/crates/runtime/src/turn/canonical_commit.rs index 43b7aeb3e8..c4358508a4 100644 --- a/crates/runtime/src/turn/canonical_commit.rs +++ b/crates/runtime/src/turn/canonical_commit.rs @@ -2,9 +2,12 @@ use serde_json::Value; #[derive(Debug, Clone)] pub(crate) struct CanonicalRewriteProof { - base_root: String, + // This is the manifest-chain root used to fence the eventual commit. It is + // intentionally a different hash domain from `authorized_prefix_root`. + base_manifest_root: String, base_compaction_generation: u64, base_prefix_len: usize, + base_prefix_root: String, authorized_prefix_len: usize, authorized_prefix_root: String, rewritten: bool, @@ -50,21 +53,27 @@ pub(crate) fn sanitize_provider_canonical_wal_snapshot( } impl CanonicalRewriteProof { - pub(crate) fn new( + /// Binds a rewrite proof to history materialized from the admitted manifest. + /// + /// The caller must pass the messages returned by coordinator materialization + /// for `base_manifest_root`. `begin` proves that compaction starts from those + /// messages, while commit separately fences the manifest root. + pub(crate) fn from_materialized_admission( admitted_prefix: &[Value], - base_root: &str, + base_manifest_root: &str, base_compaction_generation: u64, ) -> Self { let admitted_root = astra_turn_types::canonical_conversation_root(admitted_prefix); Self { - base_root: base_root.to_string(), + base_manifest_root: base_manifest_root.to_string(), base_compaction_generation, base_prefix_len: admitted_prefix.len(), + base_prefix_root: admitted_root.clone(), authorized_prefix_len: admitted_prefix.len(), - authorized_prefix_root: admitted_root.clone(), + authorized_prefix_root: admitted_root, rewritten: false, pending_provider_wal_predecessor: None, - valid: admitted_root == base_root, + valid: true, } } @@ -133,8 +142,8 @@ impl CanonicalRewriteProof { .then(|| self.base_compaction_generation.saturating_add(1)) } - pub(crate) fn base_root(&self) -> &str { - &self.base_root + pub(crate) fn base_manifest_root(&self) -> &str { + &self.base_manifest_root } pub(crate) fn provider_wal_replacement_authorization( @@ -145,7 +154,7 @@ impl CanonicalRewriteProof { let base_count = usize::try_from(durable_base.canonical.message_count).ok()?; let durable_predecessor = self.pending_provider_wal_predecessor.as_ref()?; if base_count != self.base_prefix_len - || durable_base.canonical.root_hash != self.base_root + || durable_base.canonical.root_hash != self.base_prefix_root || !self.authorizes(messages) { return None; @@ -170,7 +179,7 @@ impl CanonicalRewriteProof { != Some(self.base_compaction_generation.saturating_add(1)) || usize::try_from(durable_base.canonical.message_count).ok() != Some(self.base_prefix_len) - || durable_base.canonical.root_hash != self.base_root + || durable_base.canonical.root_hash != self.base_prefix_root { return Err("provider WAL replacement does not match the admitted rewrite base".into()); } @@ -211,7 +220,7 @@ impl CanonicalRewriteProof { || &transition.predecessor != expected_predecessor || usize::try_from(durable_base.canonical.message_count).ok() != Some(self.base_prefix_len) - || durable_base.canonical.root_hash != self.base_root + || durable_base.canonical.root_hash != self.base_prefix_root { return Err("provider WAL replacement does not match the admitted rewrite base".into()); } @@ -383,7 +392,10 @@ mod tests { let wal_base = astra_turn_types::ProviderCanonicalWalBaseV2::from_messages(&durable).unwrap(); - let mut invalid = CanonicalRewriteProof::new(&durable, &base.root_hash, 4); + let base_manifest_root = "a".repeat(64); + assert_ne!(base_manifest_root, base.root_hash); + let mut invalid = + CanonicalRewriteProof::from_materialized_admission(&durable, &base_manifest_root, 4); let rewritten = vec![json!({ "role": "system", "content": "summary with hf_AbCdEfGhIjKlMnOpQrStUvWxYz0123456789" @@ -395,7 +407,8 @@ mod tests { None ); - let mut valid = CanonicalRewriteProof::new(&durable, &base.root_hash, 4); + let mut valid = + CanonicalRewriteProof::from_materialized_admission(&durable, &base_manifest_root, 4); let permit = valid.begin(&durable); valid.finish(permit, &rewritten, Some(&wal_base)); let authorization = valid @@ -453,7 +466,10 @@ mod tests { let base = astra_turn_types::CanonicalPrefixIdentityV1::from_messages(&durable).unwrap(); let wal_base = astra_turn_types::ProviderCanonicalWalBaseV2::from_messages(&durable).unwrap(); - let mut live = CanonicalRewriteProof::new(&durable, &base.root_hash, 7); + let base_manifest_root = "a".repeat(64); + assert_ne!(base_manifest_root, base.root_hash); + let mut live = + CanonicalRewriteProof::from_materialized_admission(&durable, &base_manifest_root, 7); let mut source = durable.clone(); source.push(json!({"role": "user", "content": "current"})); let permit = live.begin(&source); @@ -474,7 +490,8 @@ mod tests { let mut recovered = durable.clone(); transition.apply_to(&mut recovered).unwrap(); - let mut restored_proof = CanonicalRewriteProof::new(&durable, &base.root_hash, 7); + let mut restored_proof = + CanonicalRewriteProof::from_materialized_admission(&durable, &base_manifest_root, 7); restored_proof .recover_provider_wal_replacement(&wal_base, &transition, &recovered) .unwrap(); diff --git a/crates/runtime/src/turn/llm/context.rs b/crates/runtime/src/turn/llm/context.rs index 8ffee7d3ff..edf32a5ae7 100644 --- a/crates/runtime/src/turn/llm/context.rs +++ b/crates/runtime/src/turn/llm/context.rs @@ -1283,7 +1283,14 @@ pub(crate) fn assemble_context_pipeline( if !input.tool_surface.deferred_tools_block.is_empty() { session_ctx.deferred_tools_block = input.tool_surface.deferred_tools_block.to_string(); } - let statics = crate::prompts::build_pipeline_static_sections(); + // Prompt overrides are stable within one pipeline session, but a newly + // created/restored session must observe the current override files. Cache + // the compiled sections on PipelineSession rather than for the process. + let statics = state + .pipeline_session + .as_mut() + .expect("pipeline_session checked before context assembly") + .static_sections_or_init(crate::prompts::build_pipeline_static_sections); let agent = AgentContext { tool_schemas: effective_tools, ..Default::default() diff --git a/crates/runtime/src/turn/llm/durable.rs b/crates/runtime/src/turn/llm/durable.rs index 5e4760f87f..5f5ddf34e1 100644 --- a/crates/runtime/src/turn/llm/durable.rs +++ b/crates/runtime/src/turn/llm/durable.rs @@ -283,6 +283,13 @@ pub(crate) trait InferenceLedgerPersistence: Send + Sync { attempt: &astra_services::InferenceProviderAttemptPlan, terminal: &astra_services::InferenceInvocationTerminal, ) -> astra_services::ServiceResult<()>; + + async fn finish_successful_provider_attempt_and_invocation( + &self, + plan: &astra_services::InferenceInvocationPlan, + attempt: &astra_services::InferenceProviderAttemptPlan, + terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()>; } struct DatabaseInferenceLedgerPersistence { @@ -377,6 +384,21 @@ impl InferenceLedgerPersistence for DatabaseInferenceLedgerPersistence { astra_services::finish_inference_provider_attempt(&self.shared_pool, attempt, terminal) .await } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + plan: &astra_services::InferenceInvocationPlan, + attempt: &astra_services::InferenceProviderAttemptPlan, + terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + astra_services::finish_successful_inference_provider_attempt_and_invocation( + &self.shared_pool, + plan, + attempt, + terminal, + ) + .await + } } #[cfg(not(test))] @@ -2469,6 +2491,75 @@ impl InferenceLedgerPersistence for TestInferenceLedgerPersistence { } Ok(()) } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + plan: &astra_services::InferenceInvocationPlan, + attempt: &astra_services::InferenceProviderAttemptPlan, + terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + if terminal.status != astra_services::InferenceTerminalStatus::Succeeded { + return Err(astra_services::ServiceError::invalid( + "combined test inference settlement requires a successful terminal", + )); + } + let mut state = self.lock(); + let provider_attempt = state + .attempts + .get_mut(attempt.attempt_id()) + .ok_or_else(|| { + astra_services::ServiceError::conflict(format!( + "test provider attempt {} was not admitted", + attempt.attempt_id() + )) + })?; + if provider_attempt.invocation_id != plan.invocation_id() { + return Err(astra_services::ServiceError::conflict(format!( + "test provider attempt {} belongs to a different invocation", + attempt.attempt_id() + ))); + } + match provider_attempt.terminal.as_ref() { + Some(existing) if existing != terminal => { + return Err(astra_services::ServiceError::conflict(format!( + "test provider attempt {} has a conflicting terminal", + attempt.attempt_id() + ))); + } + Some(_) => {} + None => provider_attempt.terminal = Some(terminal.clone()), + } + if state.attempts.values().any(|candidate| { + candidate.invocation_id == plan.invocation_id() && candidate.terminal.is_none() + }) { + return Err(astra_services::ServiceError::conflict(format!( + "test inference invocation {} still has an open provider attempt", + plan.invocation_id() + ))); + } + let invocation = state + .invocations + .get_mut(plan.invocation_id()) + .ok_or_else(|| { + astra_services::ServiceError::conflict(format!( + "test inference invocation {} was not admitted", + plan.invocation_id() + )) + })?; + match invocation.terminal.as_ref() { + Some(existing) if existing != terminal => { + Err(astra_services::ServiceError::conflict(format!( + "test inference invocation {} has a conflicting terminal", + plan.invocation_id() + ))) + } + Some(_) => Ok(()), + None => { + invocation.terminal = Some(terminal.clone()); + Ok(()) + } + } + } } #[derive(Clone)] @@ -3678,6 +3769,20 @@ impl DurableInferenceInvocation { terminal: &astra_services::InferenceInvocationTerminal, ) -> Result<(), astra_core::ClassifiedError> { self.owner_lease.ensure_live("logical terminal")?; + if let Some(committed) = self.observer.committed_logical_terminal().await { + if committed != *terminal { + return Err(contract_error( + "terminal commit", + "logical terminal conflicts with the successful provider settlement", + )); + } + drop( + self.take_settlement_reservation("combined logical terminal") + .await?, + ); + self.owner_lease.stop(); + return Ok(()); + } // Successful provider terminalization atomically creates the matching // debt in the services layer. Other terminal kinds are retryable at // the physical layer, so only the logical owner may declare them final. @@ -3779,6 +3884,7 @@ struct ProviderAttemptState { settlement_handed_off: BTreeSet, pending_terminals: BTreeMap, terminals: BTreeMap, + logical_terminal: Option, } impl ProviderAttemptState { @@ -3966,6 +4072,12 @@ impl DurableProviderAttemptObserver { } } + async fn committed_logical_terminal( + &self, + ) -> Option { + self.state.lock().await.logical_terminal.clone() + } + async fn finish_open_attempts( &self, terminal: &astra_services::InferenceInvocationTerminal, @@ -4246,15 +4358,31 @@ impl ProviderAttemptObserver for DurableProviderAttemptObserver { .insert(attempt_index, terminal.clone()); attempt }; - self.persistence - .finish_provider_attempt(&attempt, terminal) - .await - .map_err(|error| service_error("provider attempt terminal commit", error))?; + if terminal.status == astra_services::InferenceTerminalStatus::Succeeded { + self.persistence + .finish_successful_provider_attempt_and_invocation( + &self.invocation, + &attempt, + terminal, + ) + .await + .map_err(|error| { + service_error("combined provider and invocation terminal commit", error) + })?; + } else { + self.persistence + .finish_provider_attempt(&attempt, terminal) + .await + .map_err(|error| service_error("provider attempt terminal commit", error))?; + } let mut state = self.state.lock().await; state.open_attempts.remove(&attempt_index); state.delivery_authorized.remove(&attempt_index); state.pending_terminals.remove(&attempt_index); state.terminals.insert(attempt_index, terminal.clone()); + if terminal.status == astra_services::InferenceTerminalStatus::Succeeded { + state.logical_terminal = Some(terminal.clone()); + } Ok(()) } @@ -4915,6 +5043,15 @@ mod tests { ) -> astra_services::ServiceResult<()> { Ok(()) } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + _plan: &astra_services::InferenceInvocationPlan, + _attempt: &astra_services::InferenceProviderAttemptPlan, + _terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + Ok(()) + } } #[async_trait] @@ -4997,6 +5134,17 @@ mod tests { ) -> astra_services::ServiceResult<()> { self.inner.finish_provider_attempt(attempt, terminal).await } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + plan: &astra_services::InferenceInvocationPlan, + attempt: &astra_services::InferenceProviderAttemptPlan, + terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + self.inner + .finish_successful_provider_attempt_and_invocation(plan, attempt, terminal) + .await + } } #[async_trait] @@ -5066,6 +5214,17 @@ mod tests { ) -> astra_services::ServiceResult<()> { self.inner.finish_provider_attempt(attempt, terminal).await } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + plan: &astra_services::InferenceInvocationPlan, + attempt: &astra_services::InferenceProviderAttemptPlan, + terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + self.inner + .finish_successful_provider_attempt_and_invocation(plan, attempt, terminal) + .await + } } #[async_trait] @@ -5131,6 +5290,18 @@ mod tests { let _active = ActiveWorkerGuard(&self.active_finish_workers); std::future::pending().await } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + _plan: &astra_services::InferenceInvocationPlan, + _attempt: &astra_services::InferenceProviderAttemptPlan, + _terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + self.finish_entered.fetch_add(1, Ordering::SeqCst); + self.active_finish_workers.fetch_add(1, Ordering::SeqCst); + let _active = ActiveWorkerGuard(&self.active_finish_workers); + std::future::pending().await + } } #[async_trait] @@ -5207,6 +5378,17 @@ mod tests { ) -> astra_services::ServiceResult<()> { self.inner.finish_provider_attempt(attempt, terminal).await } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + plan: &astra_services::InferenceInvocationPlan, + attempt: &astra_services::InferenceProviderAttemptPlan, + terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + self.inner + .finish_successful_provider_attempt_and_invocation(plan, attempt, terminal) + .await + } } #[async_trait] @@ -5280,6 +5462,19 @@ mod tests { self.release_finish.notified().await; self.inner.finish_provider_attempt(attempt, terminal).await } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + plan: &astra_services::InferenceInvocationPlan, + attempt: &astra_services::InferenceProviderAttemptPlan, + terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + self.finish_entered.fetch_add(1, Ordering::SeqCst); + self.release_finish.notified().await; + self.inner + .finish_successful_provider_attempt_and_invocation(plan, attempt, terminal) + .await + } } #[async_trait] @@ -5343,6 +5538,16 @@ mod tests { self.finished.fetch_add(1, Ordering::SeqCst); Ok(()) } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + _plan: &astra_services::InferenceInvocationPlan, + _attempt: &astra_services::InferenceProviderAttemptPlan, + _terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + self.finished.fetch_add(1, Ordering::SeqCst); + Ok(()) + } } #[async_trait] @@ -5418,6 +5623,17 @@ mod tests { ) -> astra_services::ServiceResult<()> { self.inner.finish_provider_attempt(attempt, terminal).await } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + plan: &astra_services::InferenceInvocationPlan, + attempt: &astra_services::InferenceProviderAttemptPlan, + terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + self.inner + .finish_successful_provider_attempt_and_invocation(plan, attempt, terminal) + .await + } } #[async_trait] @@ -5500,6 +5716,15 @@ mod tests { ) -> astra_services::ServiceResult<()> { Ok(()) } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + _plan: &astra_services::InferenceInvocationPlan, + _attempt: &astra_services::InferenceProviderAttemptPlan, + _terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + Ok(()) + } } fn test_invocation_plan() -> astra_services::InferenceInvocationPlan { diff --git a/crates/runtime/src/turn/llm/summary_client.rs b/crates/runtime/src/turn/llm/summary_client.rs index b3c1a53cde..a8490753fb 100644 --- a/crates/runtime/src/turn/llm/summary_client.rs +++ b/crates/runtime/src/turn/llm/summary_client.rs @@ -454,6 +454,17 @@ mod tests { ) -> astra_services::ServiceResult<()> { self.inner.finish_provider_attempt(attempt, terminal).await } + + async fn finish_successful_provider_attempt_and_invocation( + &self, + plan: &astra_services::InferenceInvocationPlan, + attempt: &astra_services::InferenceProviderAttemptPlan, + terminal: &astra_services::InferenceInvocationTerminal, + ) -> astra_services::ServiceResult<()> { + self.inner + .finish_successful_provider_attempt_and_invocation(plan, attempt, terminal) + .await + } } async fn spawn_summary_test_server(app: Router) -> String { diff --git a/crates/runtime/src/turn/services.rs b/crates/runtime/src/turn/services.rs index 909e8df574..54ef4b79ff 100644 --- a/crates/runtime/src/turn/services.rs +++ b/crates/runtime/src/turn/services.rs @@ -1,4 +1,4 @@ -use crate::data_layer::storage::{insert_trace_event, touch_agent_session_activity}; +use crate::data_layer::storage::{insert_trace_events, touch_agent_session_activity}; use crate::server::run::lifecycle::{ TranscriptPersistItem, TranscriptPersistPayload, persist_session_transcript_items_inner_in_tx, }; @@ -502,17 +502,23 @@ impl DatabaseTraceEventWriter { ) .await .map_err(TraceWriteError::Persist)?; + let mut by_session = std::collections::BTreeMap::<(String, String), Vec>::new(); + for event in events { + by_session + .entry((event.user_id.clone(), event.session_id.clone())) + .or_default() + .push(event); + } let mut deltas = SessionEventDeltas::new(); - for event in &events { - if insert_trace_event(tx, event) + for ((user_id, session_id), events) in by_session { + let (inserted, last_event_id) = insert_trace_events(tx, &events) .await - .map_err(|error| TraceWriteError::Persist(error.to_string()))? - { - let entry = deltas - .entry((event.user_id.clone(), event.session_id.clone())) - .or_default(); - entry.0 += 1; - entry.1 = Some(event.event_id.clone()); + .map_err(|error| TraceWriteError::Persist(error.to_string()))?; + if inserted > 0 { + deltas.insert( + (user_id, session_id), + (i64::try_from(inserted).unwrap_or(i64::MAX), last_event_id), + ); } } // Session summary updates are deliberately deferred until the owning @@ -1065,6 +1071,10 @@ mod tests { } } + fn trace_event(event_id: &str, user_id: &str, session_id: &str) -> TraceEvent { + TraceEvent::new(event_id, session_id, user_id, "trace", "runtime") + } + fn tool_event( event_id: &str, user_id: &str, @@ -1321,6 +1331,103 @@ mod tests { .expect("cleanup event count fixture agent_sessions"); } + #[tokio::test] + #[ignore = "requires MatrixOne; run with ASTRA_TEST_DB_IT=1"] + async fn trace_batch_tail_tracks_the_last_new_event_in_both_replay_orders() { + let shared = setup_live_pool_for_test().await; + let pool = shared.get().clone(); + let settings = MatrixOneSettings::from_env(); + let suffix = Uuid::new_v4().to_string(); + let user_id = format!("trace-tail-user-{suffix}"); + let existing_first_session = format!("trace-tail-existing-first-{suffix}"); + let existing_last_session = format!("trace-tail-existing-last-{suffix}"); + let existing_first = format!("trace-existing-first-{suffix}"); + let new_after = format!("trace-new-after-{suffix}"); + let new_before = format!("trace-new-before-{suffix}"); + let existing_last = format!("trace-existing-last-{suffix}"); + + for session_id in [&existing_first_session, &existing_last_session] { + sqlx::query( + "INSERT INTO agent_sessions (session_id, user_id, title, status, event_count) \ + VALUES (?, ?, 'trace-tail-it', 'active', 0)", + ) + .bind(session_id) + .bind(&user_id) + .execute(&pool) + .await + .expect("insert trace-tail session"); + } + + let writer = DatabaseTraceEventWriter::new(settings).with_pool(shared); + writer + .write(trace_event( + &existing_first, + &user_id, + &existing_first_session, + )) + .await + .expect("persist existing-first fixture"); + writer + .write_many(vec![ + trace_event(&existing_first, &user_id, &existing_first_session), + trace_event(&new_after, &user_id, &existing_first_session), + ]) + .await + .expect("persist existing-then-new batch"); + + writer + .write(trace_event( + &existing_last, + &user_id, + &existing_last_session, + )) + .await + .expect("persist existing-last fixture"); + writer + .write_many(vec![ + trace_event(&new_before, &user_id, &existing_last_session), + trace_event(&existing_last, &user_id, &existing_last_session), + ]) + .await + .expect("persist new-then-existing batch"); + + for (session_id, expected_tail) in [ + (&existing_first_session, &new_after), + (&existing_last_session, &new_before), + ] { + let row = sqlx::query( + "SELECT event_count, last_event_id FROM agent_sessions \ + WHERE user_id = ? AND session_id = ?", + ) + .bind(&user_id) + .bind(session_id) + .fetch_one(&pool) + .await + .expect("load trace-tail session summary"); + assert_eq!(row.try_get::("event_count").unwrap(), 2); + assert_eq!( + row.try_get::("last_event_id").unwrap(), + *expected_tail + ); + } + + sqlx::query("DELETE FROM agent_event_edges WHERE user_id = ?") + .bind(&user_id) + .execute(&pool) + .await + .expect("cleanup trace-tail edges"); + sqlx::query("DELETE FROM agent_events WHERE user_id = ?") + .bind(&user_id) + .execute(&pool) + .await + .expect("cleanup trace-tail events"); + sqlx::query("DELETE FROM agent_sessions WHERE user_id = ?") + .bind(&user_id) + .execute(&pool) + .await + .expect("cleanup trace-tail sessions"); + } + /// Verify that all Database*Writer structs fail instantly when no pool is /// configured, rather than blocking on a 2s connect_matrixone() timeout. #[tokio::test] diff --git a/crates/services/src/auth/external.rs b/crates/services/src/auth/external.rs index f09cf64f13..83128cdefa 100644 --- a/crates/services/src/auth/external.rs +++ b/crates/services/src/auth/external.rs @@ -446,6 +446,7 @@ impl ExternalRuntimeCapabilityDescriptors { .edge_agent .as_ref() .map(ExternalRuntimeCapabilityDescriptor::to_request_descriptor), + discovery_snapshot: None, } } } diff --git a/crates/services/src/inference_execution.rs b/crates/services/src/inference_execution.rs index 671ef37b3c..e4edd3f22c 100644 --- a/crates/services/src/inference_execution.rs +++ b/crates/services/src/inference_execution.rs @@ -948,7 +948,12 @@ async fn insert_model_request_context_event_with_expiry( ModelRequestContextScope::HarnessRun(harness_run_id) } }; - compact_model_request_context_scope(connection, &attempt.user_id, scope).await?; + // A normal request appends an accepted/terminal pair. Compact once at + // the terminal boundary; orphaned accepted rows remain covered by the + // existing expiry sweeper. + if stage == ModelRequestEventStage::Terminal { + compact_model_request_context_scope(connection, &attempt.user_id, scope).await?; + } } if let (Some(status), Some(usage)) = (event.terminal_status.as_deref(), usage) { sqlx::query( @@ -4108,6 +4113,221 @@ async fn recover_provider_terminal_after_unknown_write( } } +async fn combined_successful_settlement_is_durable( + db: &sqlx::Pool, + plan: &InferenceInvocationPlan, + attempt: &InferenceProviderAttemptPlan, + provider_wire_bytes: i64, + terminal: &InferenceInvocationTerminal, + fingerprint: &str, +) -> ServiceResult { + let Some(persisted_attempt) = load_provider_attempt_fact(db, attempt).await? else { + return Ok(false); + }; + if classify_persisted_provider_terminal( + &persisted_attempt, + attempt, + provider_wire_bytes, + terminal, + fingerprint, + )? != PersistedProviderTerminalMatch::ExactTerminal + { + return Ok(false); + } + Ok(existing_terminal_fingerprint(db, plan).await?.as_deref() == Some(fingerprint)) +} + +/// Atomically settle a successful physical attempt and its logical invocation. +/// Failure and delivery-unknown outcomes retain the explicit debt protocol +/// because they may still own retryable/open attempts. +pub async fn finish_successful_inference_provider_attempt_and_invocation( + pool: &SharedPool, + plan: &InferenceInvocationPlan, + attempt: &InferenceProviderAttemptPlan, + terminal: &InferenceInvocationTerminal, +) -> ServiceResult<()> { + validate_first_provider_attempt_binding(plan, attempt)?; + if terminal.status != InferenceTerminalStatus::Succeeded { + return Err(ServiceError::invalid( + "combined inference settlement requires a successful terminal", + )); + } + let fingerprint = terminal_fingerprint(terminal)?; + let terminal_state = DurableInferenceTerminal::from_terminal(terminal, fingerprint.clone())?; + let provider_wire_bytes = checked_i64(attempt.wire.provider_wire_bytes, "provider_wire_bytes")?; + let owner_generation = i64::try_from(plan.owner_generation).map_err(|_| { + ServiceError::invalid("inference owner generation exceeds the durable BIGINT range") + })?; + let db = pool.get(); + let mut tx = db.begin().await.map_err(|error| { + ServiceError::with_source( + ServiceErrorKind::Persistence, + "begin combined successful inference settlement", + error, + ) + })?; + lock_admitted_inference_invocation( + &mut tx, + &plan.input.user_id, + &plan.invocation_id, + &plan.owner_token, + plan.owner_generation, + "commit a successful provider and logical terminal", + ) + .await?; + + let attempt_update = sqlx::query( + "UPDATE inference_provider_attempts + SET status = ?, terminal_fingerprint = ?, provider_response_id = ?, + usage_status = ?, input_tokens = ?, output_tokens = ?, cache_read_tokens = ?, + cache_creation_tokens = ?, error_kind = ?, error_message = ?, terminal_at = NOW(6) + WHERE user_id = ? AND attempt_id = ? + AND invocation_id = ? AND attempt_index = ? AND provider = ? + AND admission_token = ? AND provider_protocol = ? + AND provider_wire_hash = ? AND provider_wire_bytes = ? + AND status = 'started'", + ) + .bind(&terminal_state.status) + .bind(&fingerprint) + .bind(&terminal_state.provider_response_id) + .bind(&terminal_state.usage_status) + .bind(terminal_state.input_tokens) + .bind(terminal_state.output_tokens) + .bind(terminal_state.cache_read_tokens) + .bind(terminal_state.cache_creation_tokens) + .bind(&terminal_state.error_kind) + .bind(&terminal_state.error_message) + .bind(&attempt.user_id) + .bind(&attempt.attempt_id) + .bind(&attempt.invocation_id) + .bind(i64::from(attempt.attempt_index)) + .bind(&attempt.provider) + .bind(&attempt.admission_token) + .bind(&attempt.wire.protocol) + .bind(&attempt.wire.provider_wire_hash) + .bind(provider_wire_bytes) + .execute(&mut *tx) + .await + .map_err(|error| { + ServiceError::with_source( + ServiceErrorKind::Persistence, + "finish provider attempt in combined successful inference settlement", + error, + ) + })?; + + if attempt_update.rows_affected() != 1 { + rollback_inference_tx(tx, "classify combined successful inference settlement").await; + return if combined_successful_settlement_is_durable( + db, + plan, + attempt, + provider_wire_bytes, + terminal, + &fingerprint, + ) + .await? + { + Ok(()) + } else { + Err(ServiceError::conflict(format!( + "inference provider attempt {} is unavailable for combined successful settlement", + attempt.attempt_id + ))) + }; + } + + insert_model_request_context_event( + &mut tx, + attempt, + ModelRequestEventStage::Terminal, + Some(terminal), + ) + .await?; + + let invocation_update = sqlx::query( + "UPDATE inference_invocations AS invocation + SET status = ?, terminal_fingerprint = ?, usage_status = ?, + provider_delivery_state = 'delivery_authorized', + input_tokens = ?, output_tokens = ?, cache_read_tokens = ?, + cache_creation_tokens = ?, provider_response_id = ?, + error_kind = ?, error_message = ?, terminal_at = NOW(6) + WHERE invocation.user_id = ? AND invocation.invocation_id = ? + AND invocation.admission_token = ? AND invocation.status = 'admitted' + AND invocation.owner_token = ? AND invocation.owner_generation = ? + AND invocation.owner_lease_expires_at > NOW(6) + AND EXISTS ( + SELECT 1 FROM inference_provider_attempts AS succeeded_attempt + WHERE succeeded_attempt.user_id = invocation.user_id + AND succeeded_attempt.invocation_id = invocation.invocation_id + AND succeeded_attempt.attempt_id = ? + AND succeeded_attempt.status = 'succeeded' + AND succeeded_attempt.terminal_fingerprint = ? + ) + AND NOT EXISTS ( + SELECT 1 FROM inference_provider_attempts AS open_attempt + WHERE open_attempt.user_id = invocation.user_id + AND open_attempt.invocation_id = invocation.invocation_id + AND open_attempt.status = 'started' + )", + ) + .bind(&terminal_state.status) + .bind(&fingerprint) + .bind(&terminal_state.usage_status) + .bind(terminal_state.input_tokens) + .bind(terminal_state.output_tokens) + .bind(terminal_state.cache_read_tokens) + .bind(terminal_state.cache_creation_tokens) + .bind(&terminal_state.provider_response_id) + .bind(&terminal_state.error_kind) + .bind(&terminal_state.error_message) + .bind(&plan.input.user_id) + .bind(&plan.invocation_id) + .bind(&plan.admission_token) + .bind(&plan.owner_token) + .bind(owner_generation) + .bind(&attempt.attempt_id) + .bind(&fingerprint) + .execute(&mut *tx) + .await + .map_err(|error| { + ServiceError::with_source( + ServiceErrorKind::Persistence, + "finish invocation in combined successful inference settlement", + error, + ) + })?; + if invocation_update.rows_affected() != 1 { + rollback_inference_tx(tx, "finish combined successful inference settlement").await; + return Err(ServiceError::conflict(format!( + "inference invocation {} is unavailable for combined successful settlement", + plan.invocation_id + ))); + } + + if let Err(error) = tx.commit().await { + let commit_error = ServiceError::with_source( + ServiceErrorKind::Persistence, + "commit combined successful inference settlement", + error, + ); + if combined_successful_settlement_is_durable( + db, + plan, + attempt, + provider_wire_bytes, + terminal, + &fingerprint, + ) + .await? + { + return Ok(()); + } + return Err(commit_error); + } + Ok(()) +} + pub async fn finish_inference_provider_attempt( pool: &SharedPool, attempt: &InferenceProviderAttemptPlan, diff --git a/crates/services/src/lib.rs b/crates/services/src/lib.rs index 2b9f3f1724..e7128c4aa5 100644 --- a/crates/services/src/lib.rs +++ b/crates/services/src/lib.rs @@ -206,6 +206,7 @@ pub use inference_execution::{ admit_inference_invocation_with_first_provider_attempt, begin_inference_provider_attempt, declare_inference_attempt_settlement, declare_inference_settlement, finish_inference_invocation, finish_inference_provider_attempt, + finish_successful_inference_provider_attempt_and_invocation, load_inference_canonical_transitions_for_session, next_inference_logical_attempt_pair_base, plan_inference_invocation, plan_inference_provider_attempt, plan_inference_provider_attempt_with_context, reconcile_inference_settlement, @@ -324,8 +325,9 @@ pub use session_artifact_store::{ local_owner_scope, local_owner_user_id, local_session_artifact_store, }; pub use session_context_coordinator::{ - AcquireWriterOutcome, DatabaseSessionContextCoordinator, MaterializedConversationV1, - RenewedTurnAuthority, ReserveTurnOutcome, SessionAuthorityEventV1, SessionContextCoordinator, + AcquireWriterAndReserveTurnOutcome, AcquireWriterOutcome, DatabaseSessionContextCoordinator, + MaterializedConversationV1, RenewedTurnAuthority, ReserveTurnOutcome, + SessionAdmissionSnapshotV1, SessionAuthorityEventV1, SessionContextCoordinator, SessionContextCoordinatorError, TransferWriterOutcome, WriterTransferConflictV1, WriterTransferRequestV1, }; diff --git a/crates/services/src/resource_governor.rs b/crates/services/src/resource_governor.rs index 94a3343786..3525fbb5f8 100644 --- a/crates/services/src/resource_governor.rs +++ b/crates/services/src/resource_governor.rs @@ -134,8 +134,11 @@ pub trait ResourceGovernor: Send + Sync + 'static { /// run. Counting every run as a newly-created session would make a normal /// conversation exhaust the daily session cap. async fn check_run_start(&self, user_id: &str) -> LimitCheck { - let limits = self.get_limits(user_id).await; - let usage = self.get_usage(user_id).await; + // Limits and usage are independent snapshots. A quota update or usage + // write may race either read regardless of ordering, so serializing + // them provides no stronger admission guarantee and adds a complete + // database round trip to every turn. + let (limits, usage) = tokio::join!(self.get_limits(user_id), self.get_usage(user_id)); if limits.max_concurrent_sessions > 0 && usage.active_sessions >= limits.max_concurrent_sessions @@ -174,11 +177,13 @@ pub trait ResourceGovernor: Send + Sync + 'static { /// Check whether the user's daily token budget allows further LLM calls. /// Called before each LLM invocation for mid-session enforcement. async fn check_token_budget(&self, user_id: &str) -> LimitCheck { - let limits = self.get_limits(user_id).await; + // Read both independent snapshots together. Even when the effective + // limit is unlimited, avoiding the usage read would make the latency + // depend on which snapshot happens to arrive first. + let (limits, usage) = tokio::join!(self.get_limits(user_id), self.get_usage(user_id)); if limits.max_tokens_per_day == 0 { return LimitCheck::Allowed; } - let usage = self.get_usage(user_id).await; if usage.tokens_consumed >= limits.max_tokens_per_day { LimitCheck::Denied { limit: ResourceLimitKind::DailyTokens, @@ -331,17 +336,18 @@ impl ResourceGovernor for DatabaseResourceGovernor { async fn get_usage(&self, user_id: &str) -> ResourceUsage { let today = Self::today(); - let row: Option<(i32, i64, i64)> = sqlx::query_as( + let usage = sqlx::query_as( "SELECT sessions_created, tool_calls, tokens_consumed \ FROM resource_usage WHERE user_id = ? AND usage_date = ?", ) .bind(user_id) .bind(&today) - .fetch_optional(self.pool.get()) - .await - .unwrap_or(None); - - let active = self.count_active_sessions(user_id).await; + .fetch_optional(self.pool.get()); + // Daily counters and the active-run count live in different tables + // and are observational snapshots. Fetching them concurrently keeps + // the exact fail-open behavior while removing one serialized query. + let (row, active) = tokio::join!(usage, self.count_active_sessions(user_id)); + let row: Option<(i32, i64, i64)> = row.unwrap_or(None); match row { Some((sc, tc, tk)) => ResourceUsage { @@ -357,6 +363,100 @@ impl ResourceGovernor for DatabaseResourceGovernor { } } + async fn check_run_start(&self, user_id: &str) -> LimitCheck { + // Run admission needs three authoritative facts, but not the rest of + // ResourceUsage. Fetch them in one statement so a remote database does + // not turn an observational quota check into three network round trips. + let today = Self::today(); + let row: Result<(i32, i64, i64, i64), sqlx::Error> = sqlx::query_as( + "SELECT \ + CAST(COALESCE((SELECT max_concurrent_sessions FROM resource_limits WHERE user_id = ?), ?) AS SIGNED), \ + CAST(COALESCE((SELECT max_tokens_per_day FROM resource_limits WHERE user_id = ?), ?) AS SIGNED), \ + CAST(COALESCE((SELECT tokens_consumed FROM resource_usage WHERE user_id = ? AND usage_date = ?), 0) AS SIGNED), \ + CAST((SELECT COUNT(DISTINCT session_id) FROM agent_runs WHERE user_id = ? AND status IN ('running', 'paused', 'waiting')) AS SIGNED)", + ) + .bind(user_id) + .bind(ResourceLimits::DEFAULT_MAX_CONCURRENT_SESSIONS as i32) + .bind(user_id) + .bind(ResourceLimits::DEFAULT_MAX_TOKENS_PER_DAY as i64) + .bind(user_id) + .bind(&today) + .bind(user_id) + .fetch_one(self.pool.get()) + .await; + let (max_concurrent_sessions, max_tokens_per_day, tokens_consumed, active_sessions) = + match row { + Ok(row) => row, + Err(error) => { + tracing::warn!( + target: "astra_services::resource_governor", + user_id, + error = %error, + "failed to read run-start quota snapshot; preserving fail-open admission" + ); + return LimitCheck::Allowed; + } + }; + + if max_concurrent_sessions > 0 && active_sessions >= i64::from(max_concurrent_sessions) { + return LimitCheck::Denied { + limit: ResourceLimitKind::ConcurrentSessions, + reason: format!( + "concurrent session limit reached ({active_sessions}/{max_concurrent_sessions})" + ), + }; + } + if max_tokens_per_day > 0 && tokens_consumed >= max_tokens_per_day { + return LimitCheck::Denied { + limit: ResourceLimitKind::DailyTokens, + reason: format!( + "daily token budget exhausted ({tokens_consumed}/{max_tokens_per_day})" + ), + }; + } + LimitCheck::Allowed + } + + async fn check_token_budget(&self, user_id: &str) -> LimitCheck { + // Mid-turn enforcement only depends on the token limit and today's + // token counter. In particular it must not count active sessions. + let today = Self::today(); + let row: Result<(i64, i64), sqlx::Error> = sqlx::query_as( + "SELECT \ + CAST(COALESCE((SELECT max_tokens_per_day FROM resource_limits WHERE user_id = ?), ?) AS SIGNED), \ + CAST(COALESCE((SELECT tokens_consumed FROM resource_usage WHERE user_id = ? AND usage_date = ?), 0) AS SIGNED)", + ) + .bind(user_id) + .bind(ResourceLimits::DEFAULT_MAX_TOKENS_PER_DAY as i64) + .bind(user_id) + .bind(&today) + .fetch_one(self.pool.get()) + .await; + let (max_tokens_per_day, tokens_consumed) = match row { + Ok(row) => row, + Err(error) => { + tracing::warn!( + target: "astra_services::resource_governor", + user_id, + error = %error, + "failed to read token quota snapshot; preserving fail-open admission" + ); + return LimitCheck::Allowed; + } + }; + + if max_tokens_per_day > 0 && tokens_consumed >= max_tokens_per_day { + LimitCheck::Denied { + limit: ResourceLimitKind::DailyTokens, + reason: format!( + "daily token budget exhausted ({tokens_consumed}/{max_tokens_per_day})" + ), + } + } else { + LimitCheck::Allowed + } + } + async fn check_session_create(&self, user_id: &str) -> LimitCheck { let limits = self.get_limits(user_id).await; let usage = self.get_usage(user_id).await; diff --git a/crates/services/src/runs.rs b/crates/services/src/runs.rs index bb992f1e9d..9b2aa81518 100644 --- a/crates/services/src/runs.rs +++ b/crates/services/src/runs.rs @@ -554,6 +554,28 @@ pub struct RuntimeCapabilityDescriptorsRequest { // the edge agent identified by id via the existing edge WebSocket registry. #[serde(default, skip_serializing_if = "Option::is_none")] pub edge_agent: Option, + /// Immutable discovery view frozen by the provider for this turn. Endpoint + /// descriptors remain authoritative for actual calls and lazy Skill reads. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub discovery_snapshot: Option, +} + +pub const RUNTIME_CAPABILITY_DISCOVERY_SNAPSHOT_VERSION: &str = + "moi-runtime-capability-discovery-v1"; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RuntimeCapabilityDiscoverySnapshotRequest { + pub version: String, + pub tools: Vec, + pub skill_catalogs: Vec, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RuntimeCapabilitySkillCatalogSnapshotRequest { + pub agent_binding_id: String, + pub skills: Vec, } /// Request-scoped provider authorization exposed to each bash subprocess on diff --git a/crates/services/src/session_context_coordinator.rs b/crates/services/src/session_context_coordinator.rs index 34ff42ac73..c2e747e237 100644 --- a/crates/services/src/session_context_coordinator.rs +++ b/crates/services/src/session_context_coordinator.rs @@ -90,6 +90,22 @@ pub enum ReserveTurnOutcome { }, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AcquireWriterAndReserveTurnOutcome { + Ready { + lease: ConversationWriterLeaseV1, + reservation: TurnReservationV1, + }, + WriterConflict { + current_head: Option, + active_lease_expires_at_unix_ms: Option, + }, + ReservationConflict { + lease: ConversationWriterLeaseV1, + current_head: Option, + }, +} + /// One atomically renewed writer/reservation pair. /// /// A canonical turn is writable only while both authorities are live. Renewing @@ -142,6 +158,16 @@ pub struct MaterializedConversationV1 { pub canonical_segment_bytes: u64, } +/// Read-only facts used to prepare a canonical turn. Mutating admission still +/// revalidates the writer, cursor, epochs, and reservation while holding the +/// database row lock; this snapshot only removes duplicate preflight reads. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SessionAdmissionSnapshotV1 { + pub head: Option, + pub active_writer: Option, + pub authority_epochs: AuthorityEpochsV1, +} + #[async_trait] pub trait SessionContextCoordinator: Send + Sync { async fn load_head( @@ -149,6 +175,17 @@ pub trait SessionContextCoordinator: Send + Sync { key: &SessionKeyV1, ) -> Result, SessionContextCoordinatorError>; + async fn load_admission_snapshot( + &self, + key: &SessionKeyV1, + ) -> Result { + Ok(SessionAdmissionSnapshotV1 { + head: self.load_head(key).await?, + active_writer: self.load_active_writer(key).await?, + authority_epochs: self.load_authority_epochs(key).await?.unwrap_or_default(), + }) + } + async fn materialize( &self, head: &SessionContextHeadV1, @@ -244,6 +281,61 @@ pub trait SessionContextCoordinator: Send + Sync { idempotency_key: &str, ) -> Result; + /// Acquire the branch writer and reserve its next turn as one logical + /// admission. Stores that can transact both facts together should + /// override this method; other stores preserve the same behavior through + /// the two primitive operations. + async fn acquire_writer_and_reserve_turn( + &self, + key: &SessionKeyV1, + expected_cursor: Option<&SessionCursorV1>, + actor: &ActorContextV1, + ttl: Duration, + writer_idempotency_key: &str, + reservation_idempotency_key: &str, + ) -> Result { + let lease = match self + .acquire_writer(key, expected_cursor, actor, ttl, writer_idempotency_key) + .await? + { + AcquireWriterOutcome::Acquired(lease) + | AcquireWriterOutcome::AlreadyAcquired(lease) => lease, + AcquireWriterOutcome::Conflict { + current_head, + active_lease_expires_at_unix_ms, + } => { + return Ok(AcquireWriterAndReserveTurnOutcome::WriterConflict { + current_head, + active_lease_expires_at_unix_ms, + }); + } + }; + match self + .reserve_turn(&lease, expected_cursor, ttl, reservation_idempotency_key) + .await + { + Ok(ReserveTurnOutcome::Reserved(reservation)) + | Ok(ReserveTurnOutcome::AlreadyReserved(reservation)) => { + Ok(AcquireWriterAndReserveTurnOutcome::Ready { lease, reservation }) + } + Ok(ReserveTurnOutcome::Conflict { current_head }) => { + Ok(AcquireWriterAndReserveTurnOutcome::ReservationConflict { + lease, + current_head, + }) + } + Err(error) => { + if let Err(release_error) = self.release_writer(&lease).await { + tracing::warn!( + %release_error, + "failed to release canonical writer after turn reservation failure" + ); + } + Err(error) + } + } + } + async fn renew_turn_reservation( &self, reservation: &TurnReservationV1, @@ -419,6 +511,77 @@ impl SessionContextCoordinator for DatabaseSessionContextCoordinator { Ok(head) } + async fn load_admission_snapshot( + &self, + key: &SessionKeyV1, + ) -> Result { + key.validate() + .map_err(|error| SessionContextCoordinatorError::Invalid(error.to_string()))?; + let row = sqlx::query( + "SELECT head_json, active_writer_json, authorization_epoch, + device_trust_epoch, permission_epoch, + CAST(UNIX_TIMESTAMP(NOW(6)) * 1000 AS SIGNED) AS database_now_unix_ms + FROM session_context_heads + WHERE isolation_domain = ? AND owner_user_id = ? + AND session_id = ? AND branch_id = ?", + ) + .bind(&key.isolation_domain) + .bind(&key.owner_user_id) + .bind(&key.session_id) + .bind(&key.branch_id) + .fetch_optional(self.pool.get()) + .await + .map_err(|source| database_error("load_admission_snapshot", source))?; + let Some(row) = row else { + return Ok(SessionAdmissionSnapshotV1 { + head: None, + active_writer: None, + authority_epochs: AuthorityEpochsV1::default(), + }); + }; + let head = row + .try_get::, _>("head_json") + .map_err(|source| database_error("decode_admission_head", source))? + .as_deref() + .map(|json| database_json("head", json)) + .transpose()?; + if let Some(head) = &head { + validate_head(head)?; + if head.key != *key { + return Err(SessionContextCoordinatorError::NeedsRepair( + "database admission head key mismatch".into(), + )); + } + } + let now = row + .try_get::("database_now_unix_ms") + .map_err(|source| database_error("decode_admission_database_time", source))?; + let active_writer = row + .try_get::, _>("active_writer_json") + .map_err(|source| database_error("decode_admission_writer", source))? + .as_deref() + .map(|json| database_json::("active_writer", json)) + .transpose()? + .filter(|lease| lease.expires_at_unix_ms > now); + if active_writer + .as_ref() + .is_some_and(|lease| lease.key != *key) + { + return Err(SessionContextCoordinatorError::NeedsRepair( + "admission writer owner-scoped key mismatch".into(), + )); + } + Ok(SessionAdmissionSnapshotV1 { + head, + active_writer, + authority_epochs: AuthorityEpochsV1 { + authorization_epoch: database_u64(&row, "authorization_epoch")?, + device_trust_epoch: database_u64(&row, "device_trust_epoch")?, + permission_epoch: database_u64(&row, "permission_epoch")?, + }, + }) + } + async fn materialize( &self, head: &SessionContextHeadV1, @@ -999,10 +1162,9 @@ impl SessionContextCoordinator for DatabaseSessionContextCoordinator { .begin() .await .map_err(|source| database_error("begin_acquire_writer", source))?; - let now = database_now_ms(&mut tx).await?; - let expires_at = checked_expiry(now, ttl)?; ensure_database_state(&mut tx, key, actor.authority_epochs).await?; - let mut state = lock_database_state(&mut tx, key).await?; + let (mut state, now) = lock_database_state_at_now(&mut tx, key).await?; + let expires_at = checked_expiry(now, ttl)?; let request_hash = lease_request_hash(key, expected_cursor, actor); if let Some(receipt) = load_database_receipt::( &mut tx, @@ -1178,8 +1340,7 @@ impl SessionContextCoordinator for DatabaseSessionContextCoordinator { .begin() .await .map_err(|source| database_error("begin_renew_writer", source))?; - let now = database_now_ms(&mut tx).await?; - let mut state = lock_database_state(&mut tx, &lease.key).await?; + let (mut state, now) = lock_database_state_at_now(&mut tx, &lease.key).await?; if let Err(error) = validate_active_lease(&state, lease, now) { record_database_authority_event( &mut tx, @@ -1597,8 +1758,7 @@ impl SessionContextCoordinator for DatabaseSessionContextCoordinator { .begin() .await .map_err(|source| database_error("begin_reserve_turn", source))?; - let now = database_now_ms(&mut tx).await?; - let mut state = lock_database_state(&mut tx, &lease.key).await?; + let (mut state, now) = lock_database_state_at_now(&mut tx, &lease.key).await?; let request_hash = reservation_request_hash(lease, expected_cursor); if let Some(receipt) = load_database_receipt::( &mut tx, @@ -1790,6 +1950,260 @@ impl SessionContextCoordinator for DatabaseSessionContextCoordinator { Ok(ReserveTurnOutcome::Reserved(reservation)) } + async fn acquire_writer_and_reserve_turn( + &self, + key: &SessionKeyV1, + expected_cursor: Option<&SessionCursorV1>, + actor: &ActorContextV1, + ttl: Duration, + writer_idempotency_key: &str, + reservation_idempotency_key: &str, + ) -> Result { + validate_ttl(ttl, MAX_LEASE_TTL.min(MAX_RESERVATION_TTL))?; + validate_idempotency_key(writer_idempotency_key)?; + validate_idempotency_key(reservation_idempotency_key)?; + actor + .validate_for(key) + .map_err(|_| SessionContextCoordinatorError::Unauthorized)?; + validate_optional_cursor(key, expected_cursor)?; + + let mut tx = self + .pool + .get() + .begin() + .await + .map_err(|source| database_error("begin_acquire_and_reserve_turn", source))?; + ensure_database_state(&mut tx, key, actor.authority_epochs).await?; + let (mut state, now) = lock_database_state_at_now(&mut tx, key).await?; + let expected_cursor_owned = expected_cursor.cloned(); + + let (lease, acquire_outcome) = if let Some(active) = state.active_writer.clone() + && active.idempotency_key == writer_idempotency_key + { + validate_lease_request(&active, key, &expected_cursor_owned, actor)?; + let expires_at = refreshed_live_expiry(now, ttl, active.expires_at_unix_ms, None)?; + let refreshed = state + .active_writer + .as_mut() + .expect("matched active writer lease"); + refreshed.expires_at_unix_ms = expires_at; + (refreshed.clone(), "idempotent_refreshed") + } else { + if state.head.as_ref().map(|head| &head.cursor) != expected_cursor { + let current_head = state.head.clone(); + let active_lease_expires_at_unix_ms = state + .active_writer + .as_ref() + .filter(|lease| lease.expires_at_unix_ms > now) + .map(|lease| lease.expires_at_unix_ms); + record_database_authority_event( + &mut tx, + &state, + AuthorityAuditFact { + operation: "acquire_writer", + outcome: "cursor_conflict", + actor: Some(actor), + lease_id: None, + reservation_id: None, + expected_cursor, + }, + ) + .await?; + tx.commit().await.map_err(|source| { + database_error("commit_acquire_and_reserve_conflict", source) + })?; + return Ok(AcquireWriterAndReserveTurnOutcome::WriterConflict { + current_head, + active_lease_expires_at_unix_ms, + }); + } + if state + .active_writer + .as_ref() + .is_some_and(|lease| lease.expires_at_unix_ms > now) + { + let current_head = state.head.clone(); + let active_lease_expires_at_unix_ms = state + .active_writer + .as_ref() + .map(|lease| lease.expires_at_unix_ms); + record_database_authority_event( + &mut tx, + &state, + AuthorityAuditFact { + operation: "acquire_writer", + outcome: "writer_conflict", + actor: Some(actor), + lease_id: None, + reservation_id: None, + expected_cursor, + }, + ) + .await?; + tx.commit().await.map_err(|source| { + database_error("commit_acquire_and_reserve_conflict", source) + })?; + return Ok(AcquireWriterAndReserveTurnOutcome::WriterConflict { + current_head, + active_lease_expires_at_unix_ms, + }); + } + if actor.authority_epochs != state.authority_epochs { + record_database_authority_event( + &mut tx, + &state, + AuthorityAuditFact { + operation: "acquire_writer", + outcome: "stale_fenced", + actor: Some(actor), + lease_id: None, + reservation_id: None, + expected_cursor, + }, + ) + .await?; + tx.commit().await.map_err(|source| { + database_error("commit_acquire_and_reserve_fenced", source) + })?; + return Err(SessionContextCoordinatorError::Fenced); + } + archive_database_state_receipts(&mut tx, &state).await?; + state.active_reservation = None; + state.writer_epoch = state.writer_epoch.checked_add(1).ok_or_else(|| { + SessionContextCoordinatorError::NeedsRepair("writer epoch overflow".into()) + })?; + let lease = ConversationWriterLeaseV1 { + schema_version: SESSION_COORDINATION_SCHEMA_VERSION, + key: key.clone(), + lease_id: Uuid::new_v4().to_string(), + writer_epoch: state.writer_epoch, + actor: actor.clone(), + expected_cursor: expected_cursor_owned.clone(), + acquired_at_unix_ms: now, + expires_at_unix_ms: checked_expiry(now, ttl)?, + idempotency_key: writer_idempotency_key.to_owned(), + }; + state.active_writer = Some(lease.clone()); + (lease, "acquired") + }; + + let (reservation, reserve_outcome) = if let Some(active) = state.active_reservation.clone() + && active.idempotency_key == reservation_idempotency_key + { + validate_reservation_request(&active, &lease, &expected_cursor_owned)?; + validate_active_lease(&state, &lease, now)?; + if fence_expired_reservation_authority(&mut state, &lease, now) { + update_database_state(&mut tx, &state).await?; + record_database_authority_event( + &mut tx, + &state, + AuthorityAuditFact { + operation: "reserve_turn", + outcome: "expired_authority_fenced", + actor: Some(&lease.actor), + lease_id: Some(&lease.lease_id), + reservation_id: Some(&active.reservation_id), + expected_cursor, + }, + ) + .await?; + tx.commit().await.map_err(|source| { + database_error("commit_acquire_and_reserve_expiry_fence", source) + })?; + return Err(SessionContextCoordinatorError::Expired); + } + let expires_at = refreshed_live_expiry( + now, + ttl, + active.expires_at_unix_ms, + Some(lease.expires_at_unix_ms), + )?; + let refreshed = state + .active_reservation + .as_mut() + .expect("matched active turn reservation"); + refreshed.expires_at_unix_ms = expires_at; + (refreshed.clone(), "idempotent_refreshed") + } else { + validate_active_lease(&state, &lease, now)?; + if state.head.as_ref().map(|head| &head.cursor) != expected_cursor + || state + .active_reservation + .as_ref() + .is_some_and(|reservation| reservation.expires_at_unix_ms > now) + { + let current_head = state.head.clone(); + record_database_authority_event( + &mut tx, + &state, + AuthorityAuditFact { + operation: "reserve_turn", + outcome: "reservation_conflict", + actor: Some(&lease.actor), + lease_id: Some(&lease.lease_id), + reservation_id: None, + expected_cursor, + }, + ) + .await?; + tx.commit().await.map_err(|source| { + database_error("commit_acquire_and_reserve_conflict", source) + })?; + return Ok(AcquireWriterAndReserveTurnOutcome::ReservationConflict { + lease, + current_head, + }); + } + if let Some(previous) = &state.active_reservation { + archive_database_reservation(&mut tx, previous).await?; + } + let reservation = TurnReservationV1 { + schema_version: SESSION_COORDINATION_SCHEMA_VERSION, + reservation_id: Uuid::new_v4().to_string(), + key: lease.key.clone(), + lease_id: lease.lease_id.clone(), + writer_epoch: lease.writer_epoch, + expected_cursor: expected_cursor_owned, + reserved_turn: expected_cursor + .map_or(1, |cursor| cursor.completed_turn.saturating_add(1)), + created_at_unix_ms: now, + expires_at_unix_ms: checked_expiry(now, ttl)?.min(lease.expires_at_unix_ms), + idempotency_key: reservation_idempotency_key.to_owned(), + }; + state.active_reservation = Some(reservation.clone()); + (reservation, "reserved") + }; + + update_database_state(&mut tx, &state).await?; + record_database_authority_events( + &mut tx, + &state, + &[ + AuthorityAuditFact { + operation: "acquire_writer", + outcome: acquire_outcome, + actor: Some(actor), + lease_id: Some(&lease.lease_id), + reservation_id: None, + expected_cursor, + }, + AuthorityAuditFact { + operation: "reserve_turn", + outcome: reserve_outcome, + actor: Some(&lease.actor), + lease_id: Some(&lease.lease_id), + reservation_id: Some(&reservation.reservation_id), + expected_cursor, + }, + ], + ) + .await?; + tx.commit() + .await + .map_err(|source| database_error("commit_acquire_and_reserve_turn", source))?; + Ok(AcquireWriterAndReserveTurnOutcome::Ready { lease, reservation }) + } + async fn commit_turn( &self, reservation: &TurnReservationV1, @@ -1801,57 +2215,6 @@ impl SessionContextCoordinator for DatabaseSessionContextCoordinator { .map_err(|error| SessionContextCoordinatorError::Invalid(error.to_string()))?; validate_idempotency_key(idempotency_key)?; let request_hash = commit_request_hash(reservation, &delta); - if let Some(receipt) = load_database_receipt_pool::( - self.pool.get(), - &reservation.key, - "commit", - idempotency_key, - &request_hash, - ) - .await? - { - record_database_authority_event_pool( - self.pool.get(), - &reservation.key, - AuthorityAuditFact { - operation: "commit_turn", - outcome: "idempotent_replay", - actor: None, - lease_id: Some(&reservation.lease_id), - reservation_id: Some(&reservation.reservation_id), - expected_cursor: reservation.expected_cursor.as_ref(), - }, - ) - .await?; - return Ok(CoordinatorMutationV1::AlreadyApplied { - cursor: receipt.cursor, - }); - } - - let base_head = self.load_head(&reservation.key).await?; - if base_head.as_ref().map(|head| &head.cursor) != reservation.expected_cursor.as_ref() { - record_database_authority_event_pool( - self.pool.get(), - &reservation.key, - AuthorityAuditFact { - operation: "commit_turn", - outcome: "cursor_conflict", - actor: None, - lease_id: Some(&reservation.lease_id), - reservation_id: Some(&reservation.reservation_id), - expected_cursor: reservation.expected_cursor.as_ref(), - }, - ) - .await?; - return Ok(CoordinatorMutationV1::Conflict { - current_cursor: base_head.map(|head| head.cursor), - safe_options: vec![ - CoordinatorConflictOptionV1::Refresh, - CoordinatorConflictOptionV1::Fork, - ], - }); - } - validate_delta_advance(base_head.as_ref(), reservation, &delta)?; let mut segments = Vec::with_capacity(delta.logical_segments.len()); for messages in delta.logical_segments.iter().cloned() { segments.push( @@ -1859,27 +2222,13 @@ impl SessionContextCoordinator for DatabaseSessionContextCoordinator { .map_err(|error| SessionContextCoordinatorError::Invalid(error.to_string()))?, ); } - let node = manifest_node_for_delta(&reservation.key, base_head.as_ref(), &delta, &segments) - .map_err(|error| SessionContextCoordinatorError::Invalid(error.to_string()))?; - let (prepared_total_canonical_bytes, prepared_total_message_count) = - next_head_totals(base_head.as_ref(), &segments, delta.mode)?; - self.persist_database_immutables( - &reservation.key, - &segments, - &node, - prepared_total_canonical_bytes, - prepared_total_message_count, - ) - .await?; - let mut tx = self .pool .get() .begin() .await .map_err(|source| database_error("begin_commit_turn", source))?; - let now = database_now_ms(&mut tx).await?; - let mut state = lock_database_state(&mut tx, &reservation.key).await?; + let (mut state, now) = lock_database_state_at_now(&mut tx, &reservation.key).await?; if let Some(last) = state.last_commit.clone() && last.idempotency_key == idempotency_key { @@ -1904,36 +2253,39 @@ impl SessionContextCoordinator for DatabaseSessionContextCoordinator { cursor: last.cursor.clone(), }); } - if let Some(receipt) = load_database_receipt::( - &mut tx, - &reservation.key, - "commit", - idempotency_key, - &request_hash, - ) - .await? - { - record_database_authority_event( + if let Err(error) = validate_active_reservation(&state, reservation, now) { + // Archived receipts are only relevant when this is not the current + // active reservation. The successful hot path already proves that + // the commit has not been applied, so avoid an extra database read. + if let Some(receipt) = load_database_receipt::( &mut tx, - &state, - AuthorityAuditFact { - operation: "commit_turn", - outcome: "idempotent_replay", - actor: None, - lease_id: Some(&reservation.lease_id), - reservation_id: Some(&reservation.reservation_id), - expected_cursor: reservation.expected_cursor.as_ref(), - }, + &reservation.key, + "commit", + idempotency_key, + &request_hash, ) - .await?; - tx.commit() - .await - .map_err(|source| database_error("commit_turn_replay", source))?; - return Ok(CoordinatorMutationV1::AlreadyApplied { - cursor: receipt.cursor, - }); - } - if let Err(error) = validate_active_reservation(&state, reservation, now) { + .await? + { + record_database_authority_event( + &mut tx, + &state, + AuthorityAuditFact { + operation: "commit_turn", + outcome: "idempotent_replay", + actor: None, + lease_id: Some(&reservation.lease_id), + reservation_id: Some(&reservation.reservation_id), + expected_cursor: reservation.expected_cursor.as_ref(), + }, + ) + .await?; + tx.commit() + .await + .map_err(|source| database_error("commit_turn_replay", source))?; + return Ok(CoordinatorMutationV1::AlreadyApplied { + cursor: receipt.cursor, + }); + } record_database_authority_event( &mut tx, &state, @@ -1979,43 +2331,27 @@ impl SessionContextCoordinator for DatabaseSessionContextCoordinator { }); } validate_delta_advance(state.head.as_ref(), reservation, &delta)?; - if node.parent_manifest_root - != state - .head - .as_ref() - .map(|head| head.latest_manifest_root.clone()) - { - record_database_authority_event( - &mut tx, - &state, - AuthorityAuditFact { - operation: "commit_turn", - outcome: "stale_fenced", - actor: None, - lease_id: Some(&reservation.lease_id), - reservation_id: Some(&reservation.reservation_id), - expected_cursor: reservation.expected_cursor.as_ref(), - }, - ) - .await?; - tx.commit() - .await - .map_err(|source| database_error("commit_manifest_fenced_audit", source))?; - return Err(SessionContextCoordinatorError::Fenced); - } + // Build and persist immutable canonical objects from the head held by + // this transaction. This removes the duplicate pre-BEGIN head read + // while retaining the same row-lock fence. + let node = + manifest_node_for_delta(&reservation.key, state.head.as_ref(), &delta, &segments) + .map_err(|error| SessionContextCoordinatorError::Invalid(error.to_string()))?; + let (total_canonical_bytes, total_message_count) = + next_head_totals(state.head.as_ref(), &segments, delta.mode)?; + self.persist_database_immutables_in_tx( + &mut tx, + &reservation.key, + &segments, + &node, + total_canonical_bytes, + total_message_count, + ) + .await?; if let Some(previous) = &state.last_commit { archive_database_commit(&mut tx, previous).await?; } let cursor = node.cursor(); - let (total_canonical_bytes, total_message_count) = - next_head_totals(state.head.as_ref(), &segments, delta.mode)?; - if (total_canonical_bytes, total_message_count) - != (prepared_total_canonical_bytes, prepared_total_message_count) - { - return Err(SessionContextCoordinatorError::NeedsRepair( - "prepared manifest totals do not match the fenced parent head".into(), - )); - } state.head = Some(SessionContextHeadV1 { schema_version: SESSION_COORDINATION_SCHEMA_VERSION, key: reservation.key.clone(), @@ -2034,39 +2370,6 @@ impl SessionContextCoordinator for DatabaseSessionContextCoordinator { }); state.active_reservation = None; update_database_state(&mut tx, &state).await?; - let reachable = sqlx::query( - "UPDATE conversation_manifest_nodes - SET reachable = 1 - WHERE isolation_domain = ? AND owner_user_id = ? - AND session_id = ? AND branch_id = ? AND manifest_root = ? - AND compaction_generation = ? - AND total_canonical_bytes = ? AND total_message_count = ?", - ) - .bind(&reservation.key.isolation_domain) - .bind(&reservation.key.owner_user_id) - .bind(&reservation.key.session_id) - .bind(&reservation.key.branch_id) - .bind(&cursor.canonical_root_hash) - .bind(i64_from_u64( - "reachable compaction generation", - cursor.compaction_generation, - )?) - .bind(i64_from_u64( - "reachable total canonical bytes", - total_canonical_bytes, - )?) - .bind(i64_from_u64( - "reachable total message count", - total_message_count, - )?) - .execute(&mut *tx) - .await - .map_err(|source| database_error("activate_reachable_manifest", source))?; - if reachable.rows_affected() != 1 { - return Err(SessionContextCoordinatorError::NeedsRepair( - "canonical manifest was not durably activated".into(), - )); - } record_database_authority_event( &mut tx, &state, @@ -2098,8 +2401,7 @@ impl SessionContextCoordinator for DatabaseSessionContextCoordinator { .begin() .await .map_err(|source| database_error("begin_renew_turn_reservation", source))?; - let now = database_now_ms(&mut tx).await?; - let mut state = lock_database_state(&mut tx, &reservation.key).await?; + let (mut state, now) = lock_database_state_at_now(&mut tx, &reservation.key).await?; if let Err(error) = validate_active_reservation(&state, reservation, now) { record_database_authority_event( &mut tx, @@ -2271,58 +2573,17 @@ impl DatabaseSessionContextCoordinator { Ok(segments) } - async fn persist_database_immutables( + async fn persist_database_immutables_in_tx( &self, + tx: &mut Transaction<'_, MySql>, key: &SessionKeyV1, segments: &[ConversationSegmentV1], node: &ContextManifestNodeV1, total_canonical_bytes: u64, total_message_count: u64, ) -> Result<(), SessionContextCoordinatorError> { - self.persist_database_segments(key, segments).await?; - let mut tx = self - .pool - .get() - .begin() - .await - .map_err(|source| database_error("begin_persist_manifest", source))?; - let mut lock_segments = QueryBuilder::::new( - "SELECT segment_hash FROM conversation_segments - WHERE isolation_domain = ", - ); - lock_segments - .push_bind(&key.isolation_domain) - .push(" AND owner_user_id = ") - .push_bind(&key.owner_user_id) - .push(" AND segment_hash IN ("); - { - let mut separated = lock_segments.separated(", "); - for segment in segments { - separated.push_bind(&segment.segment_hash); - } - } - lock_segments.push(") FOR UPDATE"); - let locked_segments = lock_segments - .build() - .fetch_all(&mut *tx) - .await - .map_err(|source| database_error("lock_manifest_segments", source))?; - let locked_hashes = locked_segments - .into_iter() - .map(|row| { - row.try_get::("segment_hash") - .map_err(|source| database_error("decode_locked_manifest_segment", source)) - }) - .collect::, _>>()?; - if locked_hashes.len() != segments.len() - || segments - .iter() - .any(|segment| !locked_hashes.contains(&segment.segment_hash)) - { - return Err(SessionContextCoordinatorError::NeedsRepair( - "manifest references a segment that is not durably staged".into(), - )); - } + self.persist_database_segments_in_tx(tx, key, segments) + .await?; let canonical_segment_bytes = node.appended_segments .iter() @@ -2334,15 +2595,15 @@ impl DatabaseSessionContextCoordinator { }) })?; let manifest_insert_sql = matrixone_statement_with_null_shape( - "INSERT IGNORE INTO conversation_manifest_nodes + "INSERT INTO conversation_manifest_nodes (isolation_domain, owner_user_id, session_id, branch_id, manifest_root, parent_manifest_root, completed_turn, conversation_seq, compaction_generation, canonical_segment_bytes, total_canonical_bytes, - total_message_count, manifest_json) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + total_message_count, manifest_json, reachable) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1)", [node.parent_manifest_root.is_some()], ); - let result = sqlx::query(&manifest_insert_sql) + let manifest_already_exists = match sqlx::query(&manifest_insert_sql) .bind(&key.isolation_domain) .bind(&key.owner_user_id) .bind(&key.session_id) @@ -2371,32 +2632,56 @@ impl DatabaseSessionContextCoordinator { total_message_count, )?) .bind(database_to_json("manifest", node)?) - .execute(&mut *tx) + .execute(&mut **tx) .await - .map_err(|source| database_error("persist_manifest", source))?; - if result.rows_affected() == 0 { + { + Ok(_) => false, + Err(source) if astra_core::is_duplicate_key_error(&source) => true, + Err(source) => return Err(database_error("persist_manifest", source)), + }; + let existing_manifest_reachable = if manifest_already_exists { let stored = sqlx::query( - "SELECT manifest_json FROM conversation_manifest_nodes + "SELECT parent_manifest_root, completed_turn, conversation_seq, + compaction_generation, canonical_segment_bytes, + total_canonical_bytes, total_message_count, manifest_json, reachable + FROM conversation_manifest_nodes WHERE isolation_domain = ? AND owner_user_id = ? - AND session_id = ? AND branch_id = ? AND manifest_root = ?", + AND session_id = ? AND branch_id = ? AND manifest_root = ? + FOR UPDATE", ) .bind(&key.isolation_domain) .bind(&key.owner_user_id) .bind(&key.session_id) .bind(&key.branch_id) .bind(&node.manifest_root) - .fetch_one(&mut *tx) + .fetch_one(&mut **tx) .await - .map_err(|source| database_error("verify_existing_manifest", source))? - .try_get::("manifest_json") - .map_err(|source| database_error("decode_existing_manifest", source))?; - let stored: ContextManifestNodeV1 = database_json("existing_manifest", &stored)?; - if stored != *node { + .map_err(|source| database_error("verify_existing_manifest", source))?; + let stored_manifest = stored + .try_get::("manifest_json") + .map_err(|source| database_error("decode_existing_manifest", source))?; + let stored_manifest: ContextManifestNodeV1 = + database_json("existing_manifest", &stored_manifest)?; + let stored_parent = stored + .try_get::, _>("parent_manifest_root") + .map_err(|source| database_error("decode_existing_manifest_parent", source))?; + if stored_manifest != *node + || stored_parent != node.parent_manifest_root + || database_u64(&stored, "completed_turn")? != u64::from(node.completed_turn) + || database_u64(&stored, "conversation_seq")? != node.conversation_seq + || database_u64(&stored, "compaction_generation")? != node.compaction_generation + || database_u64(&stored, "canonical_segment_bytes")? != canonical_segment_bytes + || database_u64(&stored, "total_canonical_bytes")? != total_canonical_bytes + || database_u64(&stored, "total_message_count")? != total_message_count + { return Err(SessionContextCoordinatorError::NeedsRepair( "existing immutable manifest does not match its content-addressed key".into(), )); } - } + Some(database_u64(&stored, "reachable")?) + } else { + None + }; let mut insert_references = QueryBuilder::::new( "INSERT IGNORE INTO conversation_manifest_segments @@ -2416,47 +2701,149 @@ impl DatabaseSessionContextCoordinator { .push_bind(&segment.segment_hash); }, ); - insert_references + let inserted_references = insert_references .build() - .execute(&mut *tx) + .execute(&mut **tx) .await .map_err(|source| database_error("persist_manifest_segment_references", source))?; - let stored_references = sqlx::query( - "SELECT segment_position, segment_hash - FROM conversation_manifest_segments - WHERE isolation_domain = ? AND owner_user_id = ? - AND session_id = ? AND branch_id = ? AND manifest_root = ? - ORDER BY segment_position ASC", - ) - .bind(&key.isolation_domain) - .bind(&key.owner_user_id) - .bind(&key.session_id) - .bind(&key.branch_id) - .bind(&node.manifest_root) - .fetch_all(&mut *tx) - .await - .map_err(|source| database_error("verify_manifest_segment_references", source))?; - if stored_references.len() != node.appended_segments.len() { - return Err(SessionContextCoordinatorError::NeedsRepair( - "immutable manifest segment reference count is inconsistent".into(), - )); - } - for (position, row) in stored_references.iter().enumerate() { - let stored_position = database_u64(row, "segment_position")?; - let stored_hash = row - .try_get::("segment_hash") - .map_err(|source| database_error("decode_manifest_segment_reference", source))?; - if stored_position != u64::try_from(position).unwrap_or(u64::MAX) - || stored_hash != node.appended_segments[position].segment_hash - { + if manifest_already_exists + || inserted_references.rows_affected() + != u64::try_from(node.appended_segments.len()).unwrap_or(u64::MAX) + { + let stored_references = sqlx::query( + "SELECT segment_position, segment_hash + FROM conversation_manifest_segments + WHERE isolation_domain = ? AND owner_user_id = ? + AND session_id = ? AND branch_id = ? AND manifest_root = ? + ORDER BY segment_position ASC FOR UPDATE", + ) + .bind(&key.isolation_domain) + .bind(&key.owner_user_id) + .bind(&key.session_id) + .bind(&key.branch_id) + .bind(&node.manifest_root) + .fetch_all(&mut **tx) + .await + .map_err(|source| database_error("verify_manifest_segment_references", source))?; + if stored_references.len() != node.appended_segments.len() { return Err(SessionContextCoordinatorError::NeedsRepair( - "immutable manifest segment reference does not match the manifest".into(), + "immutable manifest segment reference count is inconsistent".into(), )); } + for (position, row) in stored_references.iter().enumerate() { + let stored_position = database_u64(row, "segment_position")?; + let stored_hash = row.try_get::("segment_hash").map_err(|source| { + database_error("decode_manifest_segment_reference", source) + })?; + if stored_position != u64::try_from(position).unwrap_or(u64::MAX) + || stored_hash != node.appended_segments[position].segment_hash + { + return Err(SessionContextCoordinatorError::NeedsRepair( + "immutable manifest segment reference does not match the manifest".into(), + )); + } + } } - tx.commit() + if let Some(reachable) = existing_manifest_reachable { + match reachable { + 0 => { + let activated = sqlx::query( + "UPDATE conversation_manifest_nodes SET reachable = 1 + WHERE isolation_domain = ? AND owner_user_id = ? + AND session_id = ? AND branch_id = ? AND manifest_root = ? + AND reachable = 0 AND compaction_generation = ? + AND canonical_segment_bytes = ? AND total_canonical_bytes = ? + AND total_message_count = ?", + ) + .bind(&key.isolation_domain) + .bind(&key.owner_user_id) + .bind(&key.session_id) + .bind(&key.branch_id) + .bind(&node.manifest_root) + .bind(i64_from_u64( + "manifest compaction generation", + node.compaction_generation, + )?) + .bind(i64_from_u64( + "manifest segment bytes", + canonical_segment_bytes, + )?) + .bind(i64_from_u64( + "manifest total canonical bytes", + total_canonical_bytes, + )?) + .bind(i64_from_u64( + "manifest total message count", + total_message_count, + )?) + .execute(&mut **tx) + .await + .map_err(|source| database_error("activate_existing_manifest", source))?; + if activated.rows_affected() != 1 { + return Err(SessionContextCoordinatorError::NeedsRepair( + "verified staged manifest could not be activated".into(), + )); + } + } + 1 => {} + _ => { + return Err(SessionContextCoordinatorError::NeedsRepair( + "existing immutable manifest has an invalid reachability state".into(), + )); + } + } + } + Ok(()) + } + + async fn persist_database_segments_in_tx( + &self, + tx: &mut Transaction<'_, MySql>, + key: &SessionKeyV1, + segments: &[ConversationSegmentV1], + ) -> Result<(), SessionContextCoordinatorError> { + for segment in segments { + let json = database_to_json("segment", segment)?; + let result = sqlx::query( + "INSERT IGNORE INTO conversation_segments + (isolation_domain, owner_user_id, segment_hash, canonical_root_hash, + canonical_bytes, message_count, segment_json) + VALUES (?, ?, ?, ?, ?, ?, ?)", + ) + .bind(&key.isolation_domain) + .bind(&key.owner_user_id) + .bind(&segment.segment_hash) + .bind(&segment.canonical_root_hash) + .bind(i64_from_u64("segment bytes", segment.canonical_bytes)?) + .bind(i64::from(segment.message_count)) + .bind(json) + .execute(&mut **tx) .await - .map_err(|source| database_error("commit_persist_manifest", source)) + .map_err(|source| database_error("persist_segment", source))?; + if result.rows_affected() == 0 { + let stored = sqlx::query( + "SELECT segment_json FROM conversation_segments + WHERE isolation_domain = ? AND owner_user_id = ? AND segment_hash = ? + FOR UPDATE", + ) + .bind(&key.isolation_domain) + .bind(&key.owner_user_id) + .bind(&segment.segment_hash) + .fetch_one(&mut **tx) + .await + .map_err(|source| database_error("verify_existing_segment", source))? + .try_get::("segment_json") + .map_err(|source| database_error("decode_existing_segment", source))?; + let stored: ConversationSegmentV1 = database_json("existing_segment", &stored)?; + if stored != *segment { + return Err(SessionContextCoordinatorError::NeedsRepair( + "existing immutable segment does not match its content-addressed key" + .into(), + )); + } + } + } + Ok(()) } async fn persist_database_segments( @@ -2865,10 +3252,20 @@ async fn lock_database_state( tx: &mut Transaction<'_, MySql>, key: &SessionKeyV1, ) -> Result { + lock_database_state_at_now(tx, key) + .await + .map(|(state, _)| state) +} + +async fn lock_database_state_at_now( + tx: &mut Transaction<'_, MySql>, + key: &SessionKeyV1, +) -> Result<(CoordinatorStateV1, i64), SessionContextCoordinatorError> { let row = sqlx::query( "SELECT head_json, writer_epoch, authorization_epoch, device_trust_epoch, permission_epoch, active_writer_json, active_reservation_json, - last_commit_json, fork_base_json + last_commit_json, fork_base_json, + CAST(UNIX_TIMESTAMP(NOW(6)) * 1000 AS SIGNED) AS database_now_unix_ms FROM session_context_heads WHERE isolation_domain = ? AND owner_user_id = ? AND session_id = ? AND branch_id = ? @@ -2922,7 +3319,10 @@ async fn lock_database_state( .map(|json| database_json("fork_base", json)) .transpose()?; state.validate_for(key)?; - Ok(state) + let now = row + .try_get::("database_now_unix_ms") + .map_err(|source| database_error("decode_locked_database_time", source))?; + Ok((state, now)) } async fn update_database_state( @@ -3064,6 +3464,7 @@ async fn database_now_ms( .map_err(|source| database_error("load_database_time", source)) } +#[derive(Clone, Copy)] struct AuthorityAuditFact<'a> { operation: &'static str, outcome: &'static str, @@ -3089,71 +3490,73 @@ async fn record_database_authority_event( state: &CoordinatorStateV1, fact: AuthorityAuditFact<'_>, ) -> Result<(), SessionContextCoordinatorError> { - let actor = fact - .actor - .or_else(|| state.active_writer.as_ref().map(|lease| &lease.actor)); - sqlx::query( - "INSERT INTO session_context_authority_events - (isolation_domain, owner_user_id, event_id, session_id, branch_id, - operation_kind, outcome, writer_epoch, actor_id, device_id, lease_id, - reservation_id, expected_root, observed_root, authorization_epoch, - device_trust_epoch, permission_epoch) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind(&state.key.isolation_domain) - .bind(&state.key.owner_user_id) - .bind(Uuid::new_v4().to_string()) - .bind(&state.key.session_id) - .bind(&state.key.branch_id) - .bind(fact.operation) - .bind(fact.outcome) - .bind(i64_from_u64("audit writer epoch", state.writer_epoch)?) - .bind(actor.map(|actor| actor.actor_id.as_str())) - .bind(actor.and_then(|actor| actor.device_id.as_deref())) - .bind(fact.lease_id) - .bind(fact.reservation_id) - .bind( - fact.expected_cursor - .map(|cursor| cursor.canonical_root_hash.as_str()), - ) - .bind( - state - .head - .as_ref() - .map(|head| head.cursor.canonical_root_hash.as_str()), - ) - .bind(i64_from_u64( + record_database_authority_events(tx, state, &[fact]).await +} + +async fn record_database_authority_events( + tx: &mut Transaction<'_, MySql>, + state: &CoordinatorStateV1, + facts: &[AuthorityAuditFact<'_>], +) -> Result<(), SessionContextCoordinatorError> { + if facts.is_empty() { + return Ok(()); + } + let writer_epoch = i64_from_u64("audit writer epoch", state.writer_epoch)?; + let authorization_epoch = i64_from_u64( "audit authorization epoch", state.authority_epochs.authorization_epoch, - )?) - .bind(i64_from_u64( + )?; + let device_trust_epoch = i64_from_u64( "audit device trust epoch", state.authority_epochs.device_trust_epoch, - )?) - .bind(i64_from_u64( + )?; + let permission_epoch = i64_from_u64( "audit permission epoch", state.authority_epochs.permission_epoch, - )?) - .execute(&mut **tx) - .await - .map_err(|source| database_error("record_authority_event", source))?; - Ok(()) -} - -async fn record_database_authority_event_pool( - pool: &sqlx::Pool, - key: &SessionKeyV1, - fact: AuthorityAuditFact<'_>, -) -> Result<(), SessionContextCoordinatorError> { - let mut tx = pool - .begin() - .await - .map_err(|source| database_error("begin_authority_event", source))?; - let state = lock_database_state(&mut tx, key).await?; - record_database_authority_event(&mut tx, &state, fact).await?; - tx.commit() + )?; + let observed_root = state + .head + .as_ref() + .map(|head| head.cursor.canonical_root_hash.as_str()); + let mut insert = QueryBuilder::::new( + "INSERT INTO session_context_authority_events + (isolation_domain, owner_user_id, event_id, session_id, branch_id, + operation_kind, outcome, writer_epoch, actor_id, device_id, lease_id, + reservation_id, expected_root, observed_root, authorization_epoch, + device_trust_epoch, permission_epoch) + ", + ); + insert.push_values(facts, |mut row, fact| { + let actor = fact + .actor + .or_else(|| state.active_writer.as_ref().map(|lease| &lease.actor)); + row.push_bind(&state.key.isolation_domain) + .push_bind(&state.key.owner_user_id) + .push_bind(Uuid::new_v4().to_string()) + .push_bind(&state.key.session_id) + .push_bind(&state.key.branch_id) + .push_bind(fact.operation) + .push_bind(fact.outcome) + .push_bind(writer_epoch) + .push_bind(actor.map(|actor| actor.actor_id.as_str())) + .push_bind(actor.and_then(|actor| actor.device_id.as_deref())) + .push_bind(fact.lease_id) + .push_bind(fact.reservation_id) + .push_bind( + fact.expected_cursor + .map(|cursor| cursor.canonical_root_hash.as_str()), + ) + .push_bind(observed_root) + .push_bind(authorization_epoch) + .push_bind(device_trust_epoch) + .push_bind(permission_epoch); + }); + insert + .build() + .execute(&mut **tx) .await - .map_err(|source| database_error("commit_authority_event", source)) + .map_err(|source| database_error("record_authority_event", source))?; + Ok(()) } async fn load_database_receipt( @@ -3181,31 +3584,6 @@ async fn load_database_receipt( decode_database_receipt(row, request_hash) } -async fn load_database_receipt_pool( - pool: &sqlx::Pool, - key: &SessionKeyV1, - operation: &'static str, - idempotency_key: &str, - request_hash: &str, -) -> Result, SessionContextCoordinatorError> { - let row = sqlx::query( - "SELECT request_hash, receipt_json FROM session_context_operation_receipts - WHERE isolation_domain = ? AND owner_user_id = ? - AND session_id = ? AND branch_id = ? - AND operation_kind = ? AND idempotency_hash = ?", - ) - .bind(&key.isolation_domain) - .bind(&key.owner_user_id) - .bind(&key.session_id) - .bind(&key.branch_id) - .bind(operation) - .bind(hash_receipt(operation, idempotency_key)) - .fetch_optional(pool) - .await - .map_err(|source| database_error("load_operation_receipt", source))?; - decode_database_receipt(row, request_hash) -} - fn decode_database_receipt( row: Option, request_hash: &str, diff --git a/crates/services/src/storage.rs b/crates/services/src/storage.rs index ae510c93bb..4d5a3e42eb 100644 --- a/crates/services/src/storage.rs +++ b/crates/services/src/storage.rs @@ -1469,18 +1469,30 @@ where Ok(()) } -pub async fn insert_agent_event_edges<'e, E>( +pub struct AgentEventEdgeInsert<'a> { + pub user_id: &'a str, + pub session_id: &'a str, + pub child_event_id: &'a str, + pub primary_parent_event_id: Option<&'a str>, + pub parent_event_ids: &'a [String], +} + +pub async fn insert_agent_event_edges_batch<'e, E>( executor: E, - user_id: &str, - session_id: &str, - child_event_id: &str, - primary_parent_event_id: Option<&str>, - parent_event_ids: &[String], + edges: &[AgentEventEdgeInsert<'_>], ) -> Result<(), sqlx::Error> where E: sqlx::Executor<'e, Database = MySql>, { - let normalized = normalized_parent_event_ids(primary_parent_event_id, Some(parent_event_ids)); + let normalized = edges + .iter() + .flat_map(|edge| { + normalized_parent_event_ids(edge.primary_parent_event_id, Some(edge.parent_event_ids)) + .into_iter() + .enumerate() + .map(move |(index, parent_event_id)| (edge, index, parent_event_id)) + }) + .collect::>(); if normalized.is_empty() { return Ok(()); } @@ -1489,17 +1501,14 @@ where "INSERT INTO agent_event_edges \ (user_id, session_id, child_event_id, parent_event_id, relation_kind, parent_order) ", ); - builder.push_values( - normalized.iter().enumerate(), - |mut row, (idx, parent_event_id)| { - row.push_bind(user_id) - .push_bind(session_id) - .push_bind(child_event_id) - .push_bind(parent_event_id) - .push_bind(CAUSAL_EDGE_KIND) - .push_bind(i32::try_from(idx).unwrap_or(i32::MAX)); - }, - ); + builder.push_values(normalized, |mut row, (edge, index, parent_event_id)| { + row.push_bind(edge.user_id) + .push_bind(edge.session_id) + .push_bind(edge.child_event_id) + .push_bind(parent_event_id) + .push_bind(CAUSAL_EDGE_KIND) + .push_bind(i32::try_from(index).unwrap_or(i32::MAX)); + }); builder.push( " ON DUPLICATE KEY UPDATE \ session_id = VALUES(session_id), \ @@ -1509,6 +1518,30 @@ where Ok(()) } +pub async fn insert_agent_event_edges<'e, E>( + executor: E, + user_id: &str, + session_id: &str, + child_event_id: &str, + primary_parent_event_id: Option<&str>, + parent_event_ids: &[String], +) -> Result<(), sqlx::Error> +where + E: sqlx::Executor<'e, Database = MySql>, +{ + insert_agent_event_edges_batch( + executor, + &[AgentEventEdgeInsert { + user_id, + session_id, + child_event_id, + primary_parent_event_id, + parent_event_ids, + }], + ) + .await +} + pub async fn load_agent_event_parent_ids<'e, E>( executor: E, user_id: &str, @@ -4456,6 +4489,10 @@ async fn ensure_core_schema_while_leased( ) .execute(&pool) .await?; + sqlx::query("INSERT IGNORE INTO session_weighted_admission_gates (scope_name) VALUES (?)") + .bind(crate::weighted_admission::DISTRIBUTED_ADMISSION_SCOPE) + .execute(&pool) + .await?; core_schema_create!( pool, diff --git a/crates/services/src/weighted_admission.rs b/crates/services/src/weighted_admission.rs index d9c3c625f8..a9d7e4af7e 100644 --- a/crates/services/src/weighted_admission.rs +++ b/crates/services/src/weighted_admission.rs @@ -22,7 +22,7 @@ use thiserror::Error; use tokio::sync::Notify; use uuid::Uuid; -const DISTRIBUTED_ADMISSION_SCOPE: &str = "canonical_turn_v1"; +pub(crate) const DISTRIBUTED_ADMISSION_SCOPE: &str = "canonical_turn_v1"; const DISTRIBUTED_ADMISSION_MAX_TTL: Duration = Duration::from_secs(15 * 60); const DISTRIBUTED_ADMISSION_IDEMPOTENCY_DOMAIN: &[u8] = b"astra.distributed-weighted-admission-idempotency.v1\0"; @@ -153,8 +153,7 @@ impl DatabaseWeightedAdmissionController { .begin() .await .map_err(|source| distributed_database_error("begin_reservation", source))?; - lock_distributed_admission_gate(&mut tx).await?; - let now = distributed_database_now(&mut tx).await?; + let now = lock_distributed_admission_gate(&mut tx).await?; sqlx::query( "DELETE FROM session_weighted_admission_reservations WHERE scope_name = ? AND expires_at <= ?", @@ -218,8 +217,7 @@ impl DatabaseWeightedAdmissionController { .begin() .await .map_err(|source| distributed_database_error("begin_renewal", source))?; - lock_distributed_admission_gate(&mut tx).await?; - let now = distributed_database_now(&mut tx).await?; + let now = lock_distributed_admission_gate(&mut tx).await?; let expires_at = now .checked_add_signed(chrono::Duration::from_std(ttl).map_err(|_| { DistributedAdmissionError::Invalid("admission TTL is outside clock range".into()) @@ -407,29 +405,20 @@ fn validate_available_work( async fn lock_distributed_admission_gate( tx: &mut Transaction<'_, MySql>, -) -> Result<(), DistributedAdmissionError> { - sqlx::query("INSERT IGNORE INTO session_weighted_admission_gates (scope_name) VALUES (?)") - .bind(DISTRIBUTED_ADMISSION_SCOPE) - .execute(&mut **tx) - .await - .map_err(|source| distributed_database_error("ensure_gate", source))?; - sqlx::query( - "SELECT scope_name FROM session_weighted_admission_gates +) -> Result { + let row = sqlx::query( + "SELECT scope_name, + CAST(UNIX_TIMESTAMP(NOW(6)) * 1000 AS SIGNED) AS database_now_unix_ms + FROM session_weighted_admission_gates WHERE scope_name = ? FOR UPDATE", ) .bind(DISTRIBUTED_ADMISSION_SCOPE) .fetch_one(&mut **tx) .await .map_err(|source| distributed_database_error("lock_gate", source))?; - Ok(()) -} - -async fn distributed_database_now( - tx: &mut Transaction<'_, MySql>, -) -> Result { - let unix_ms = crate::db_row::database_now_unix_ms(tx) - .await - .map_err(|source| distributed_database_error("load_database_time", source))?; + let unix_ms = row + .try_get::("database_now_unix_ms") + .map_err(|source| distributed_database_error("decode_gate_database_time", source))?; chrono::DateTime::from_timestamp_millis(unix_ms) .map(|timestamp| timestamp.naive_utc()) .ok_or_else(|| { diff --git a/crates/services/tests/session_context_authority_db_it.rs b/crates/services/tests/session_context_authority_db_it.rs index 01a65e1e8e..20bfeed0ef 100644 --- a/crates/services/tests/session_context_authority_db_it.rs +++ b/crates/services/tests/session_context_authority_db_it.rs @@ -7,7 +7,9 @@ use astra_services::{ SessionContextCoordinator, }; use astra_turn_types::{ - ActorContextV1, ActorKindV1, AuthorityEpochsV1, SessionKeyV1, SessionSurfaceV1, + ActorContextV1, ActorKindV1, AuthorityEpochsV1, CANONICAL_TURN_DELTA_SCHEMA_VERSION, + CanonicalDeltaModeV1, CanonicalTurnDeltaV1, ContextManifestNodeV1, ConversationSegmentV1, + CoordinatorMutationV1, SessionKeyV1, SessionSurfaceV1, }; use uuid::Uuid; @@ -104,3 +106,173 @@ async fn complete_turn_authority_renews_atomically_in_database() { .expect("clean authority fixture"); } } + +#[tokio::test] +#[ignore = "ASTRA_TEST_DB_IT=1 and live MatrixOne"] +async fn commit_reactivates_matching_legacy_staged_manifest() { + let pool = common::setup_pool().await; + let owner_id = format!("staged-manifest-owner-{}", Uuid::new_v4()); + let session_id = format!("staged-manifest-session-{}", Uuid::new_v4()); + let key = SessionKeyV1::owner_session("server", &owner_id, &session_id, "main"); + let actor = ActorContextV1::owner_user( + &owner_id, + "staged-manifest-db-it", + ActorKindV1::Server, + SessionSurfaceV1::Server, + None, + AuthorityEpochsV1::default(), + ); + let coordinator = DatabaseSessionContextCoordinator::new(pool.clone()); + let lease = match coordinator + .acquire_writer(&key, None, &actor, Duration::from_secs(30), "acquire") + .await + .expect("acquire writer") + { + AcquireWriterOutcome::Acquired(lease) => lease, + other => panic!("unexpected writer outcome: {other:?}"), + }; + let reservation = match coordinator + .reserve_turn(&lease, None, Duration::from_secs(30), "reserve") + .await + .expect("reserve turn") + { + ReserveTurnOutcome::Reserved(reservation) => reservation, + other => panic!("unexpected reservation outcome: {other:?}"), + }; + let messages = vec![serde_json::json!({ + "role": "assistant", + "content": "legacy staged manifest" + })]; + let segment = ConversationSegmentV1::new(&key, messages.clone()).expect("segment"); + let node = ContextManifestNodeV1::new( + key.clone(), + None, + 1, + 1, + 1, + 0, + None, + vec![segment.reference()], + ) + .expect("manifest"); + let segment_json = serde_json::to_string(&segment).expect("serialize segment"); + sqlx::query( + "INSERT INTO conversation_segments + (isolation_domain, owner_user_id, segment_hash, canonical_root_hash, + canonical_bytes, message_count, segment_json) + VALUES (?, ?, ?, ?, ?, ?, ?)", + ) + .bind(&key.isolation_domain) + .bind(&key.owner_user_id) + .bind(&segment.segment_hash) + .bind(&segment.canonical_root_hash) + .bind(i64::try_from(segment.canonical_bytes).expect("segment bytes fit BIGINT")) + .bind(i64::from(segment.message_count)) + .bind(segment_json) + .execute(pool.get()) + .await + .expect("stage legacy segment"); + sqlx::query( + "INSERT INTO conversation_manifest_nodes + (isolation_domain, owner_user_id, session_id, branch_id, manifest_root, + parent_manifest_root, completed_turn, conversation_seq, + compaction_generation, canonical_segment_bytes, total_canonical_bytes, + total_message_count, manifest_json, reachable) + VALUES (?, ?, ?, ?, ?, NULL, ?, ?, ?, ?, ?, ?, ?, 0)", + ) + .bind(&key.isolation_domain) + .bind(&key.owner_user_id) + .bind(&key.session_id) + .bind(&key.branch_id) + .bind(&node.manifest_root) + .bind(i64::from(node.completed_turn)) + .bind(i64::try_from(node.conversation_seq).expect("conversation sequence fits BIGINT")) + .bind(i64::try_from(node.compaction_generation).expect("generation fits BIGINT")) + .bind(i64::try_from(segment.canonical_bytes).expect("segment bytes fit BIGINT")) + .bind(i64::try_from(segment.canonical_bytes).expect("total bytes fit BIGINT")) + .bind(i64::from(segment.message_count)) + .bind(serde_json::to_string(&node).expect("serialize manifest")) + .execute(pool.get()) + .await + .expect("stage legacy unreachable manifest"); + sqlx::query( + "INSERT INTO conversation_manifest_segments + (isolation_domain, owner_user_id, session_id, branch_id, + manifest_root, segment_position, segment_hash) + VALUES (?, ?, ?, ?, ?, 0, ?)", + ) + .bind(&key.isolation_domain) + .bind(&key.owner_user_id) + .bind(&key.session_id) + .bind(&key.branch_id) + .bind(&node.manifest_root) + .bind(&segment.segment_hash) + .execute(pool.get()) + .await + .expect("stage legacy manifest reference"); + + let outcome = coordinator + .commit_turn( + &reservation, + CanonicalTurnDeltaV1 { + schema_version: CANONICAL_TURN_DELTA_SCHEMA_VERSION, + completed_turn: 1, + journal_event_seq: 1, + conversation_seq: 1, + compaction_generation: 0, + config_version_id: None, + mode: CanonicalDeltaModeV1::Append, + logical_segments: vec![messages.clone()], + }, + "commit-staged-manifest", + ) + .await + .expect("commit over legacy staged manifest"); + let cursor = match outcome { + CoordinatorMutationV1::Applied { cursor } => cursor, + other => panic!("unexpected commit outcome: {other:?}"), + }; + assert_eq!(cursor.canonical_root_hash, node.manifest_root); + let reachable: i64 = sqlx::query_scalar( + "SELECT reachable FROM conversation_manifest_nodes + WHERE isolation_domain = ? AND owner_user_id = ? + AND session_id = ? AND branch_id = ? AND manifest_root = ?", + ) + .bind(&key.isolation_domain) + .bind(&key.owner_user_id) + .bind(&key.session_id) + .bind(&key.branch_id) + .bind(&node.manifest_root) + .fetch_one(pool.get()) + .await + .expect("load manifest reachability"); + assert_eq!(reachable, 1); + let head = coordinator + .load_head(&key) + .await + .expect("load committed head") + .expect("committed head"); + let materialized = coordinator + .materialize(&head) + .await + .expect("materialize reactivated manifest"); + assert_eq!(materialized.messages, messages); + + for table in [ + "conversation_manifest_segments", + "conversation_manifest_nodes", + "conversation_segments", + "session_context_operation_receipts", + "session_context_authority_events", + "session_context_heads", + ] { + sqlx::query(&format!( + "DELETE FROM {table} WHERE isolation_domain = ? AND owner_user_id = ?" + )) + .bind(&key.isolation_domain) + .bind(&key.owner_user_id) + .execute(pool.get()) + .await + .expect("clean staged manifest fixture"); + } +}