use sqlx::{Row, SqliteConnection}; use super::StorageError; /// Frozen DDL for the Agent orchestration tables. Executed inside the single /// migration transaction so table creation, column additions and `user_version` /// advance atomically. The inbox table belongs to Phase 3 behavior but its /// shape is frozen together with the run tables. pub const AGENT_SCHEMA_STATEMENTS: &[&str] = &[ r#" CREATE TABLE IF NOT EXISTS agent_runs ( id TEXT PRIMARY KEY, root_session_id TEXT NOT NULL, root_turn_id TEXT, parent_run_id TEXT, caller_agent_id TEXT NOT NULL, caller_scope_id TEXT NOT NULL, idempotency_key TEXT, agent_id TEXT NOT NULL, definition_hash TEXT NOT NULL, provider_profile TEXT NOT NULL, provider_name TEXT NOT NULL, model_id TEXT NOT NULL, mode TEXT NOT NULL, depth INTEGER NOT NULL, plan_item_id TEXT, execution_id TEXT NOT NULL, task TEXT NOT NULL, context_json TEXT, budget_json TEXT NOT NULL, signal_contract_json TEXT, signal_delivery TEXT, status TEXT NOT NULL, result TEXT, error TEXT, prompt_tokens INTEGER, completion_tokens INTEGER, cost REAL, tool_calls_count INTEGER NOT NULL DEFAULT 0, iterations INTEGER NOT NULL DEFAULT 0, runtime_generation INTEGER NOT NULL, attempt INTEGER NOT NULL DEFAULT 1, completion_slot_reserved INTEGER NOT NULL DEFAULT 0, deadline_at INTEGER NOT NULL, revision INTEGER NOT NULL, started_at INTEGER, finished_at INTEGER, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL, CHECK (mode IN ('foreground', 'background')), CHECK (status IN ('queued', 'running', 'waiting_children', 'completed', 'failed', 'timed_out', 'cancelled', 'interrupted')), CHECK (depth >= 1), CHECK (completion_slot_reserved IN (0, 1)), FOREIGN KEY (parent_run_id) REFERENCES agent_runs(id) ON DELETE RESTRICT ) "#, "CREATE UNIQUE INDEX IF NOT EXISTS idx_agent_runs_execution ON agent_runs(execution_id)", "CREATE UNIQUE INDEX IF NOT EXISTS idx_agent_runs_idempotency ON agent_runs(root_session_id, caller_scope_id, idempotency_key) WHERE idempotency_key IS NOT NULL", "CREATE INDEX IF NOT EXISTS idx_agent_runs_session_created ON agent_runs(root_session_id, created_at DESC)", "CREATE INDEX IF NOT EXISTS idx_agent_runs_parent ON agent_runs(parent_run_id, created_at)", "CREATE INDEX IF NOT EXISTS idx_agent_runs_recovery ON agent_runs(runtime_generation, status, deadline_at)", r#" CREATE TABLE IF NOT EXISTS agent_run_messages ( id TEXT PRIMARY KEY, run_id TEXT NOT NULL, seq INTEGER NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL, reasoning_content TEXT, tool_call_id TEXT, tool_name TEXT, tool_calls_json TEXT, created_at INTEGER NOT NULL, FOREIGN KEY (run_id) REFERENCES agent_runs(id) ON DELETE CASCADE ) "#, "CREATE INDEX IF NOT EXISTS idx_agent_run_messages_run_seq ON agent_run_messages(run_id, seq)", r#" CREATE TABLE IF NOT EXISTS agent_session_state ( root_session_id TEXT PRIMARY KEY, revision INTEGER NOT NULL DEFAULT 0, pending_event_count INTEGER NOT NULL DEFAULT 0, reserved_completion_slots INTEGER NOT NULL DEFAULT 0, updated_at INTEGER NOT NULL, CHECK (revision >= 0), CHECK (pending_event_count >= 0), CHECK (reserved_completion_slots >= 0) ) "#, r#" CREATE TABLE IF NOT EXISTS agent_inbox_events ( id TEXT PRIMARY KEY, root_session_id TEXT NOT NULL, run_id TEXT NOT NULL, event_type TEXT NOT NULL, event_key TEXT NOT NULL, delivery TEXT NOT NULL, requires_continuation INTEGER NOT NULL DEFAULT 1, severity TEXT, payload_json TEXT NOT NULL, status TEXT NOT NULL, attempt_count INTEGER NOT NULL DEFAULT 0, lease_token TEXT, lease_until INTEGER, next_attempt_at INTEGER, admitted_turn_id TEXT, last_error TEXT, revision INTEGER NOT NULL, created_at INTEGER NOT NULL, consumed_at INTEGER, superseded_at INTEGER, dead_lettered_at INTEGER, fallback_notified_at INTEGER, fallback_suppressed_reason TEXT, updated_at INTEGER NOT NULL, CHECK (event_type IN ('signal', 'completion')), CHECK (delivery IN ('queue', 'steer')), CHECK (requires_continuation IN (0, 1)), CHECK (status IN ('pending', 'leased', 'admitted', 'consumed', 'superseded', 'dead_letter')), UNIQUE(run_id, event_type, event_key), FOREIGN KEY (run_id) REFERENCES agent_runs(id) ON DELETE RESTRICT ) "#, "CREATE INDEX IF NOT EXISTS idx_agent_inbox_claim ON agent_inbox_events(root_session_id, status, next_attempt_at, created_at)", "CREATE INDEX IF NOT EXISTS idx_agent_inbox_lease ON agent_inbox_events(status, lease_until)", "CREATE INDEX IF NOT EXISTS idx_agent_inbox_revision ON agent_inbox_events(root_session_id, revision)", ]; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum AgentRunMode { Foreground, Background, } impl AgentRunMode { pub fn as_str(&self) -> &'static str { match self { Self::Foreground => "foreground", Self::Background => "background", } } pub fn parse(value: &str) -> Result { match value { "foreground" => Ok(Self::Foreground), "background" => Ok(Self::Background), other => Err(StorageError::Migration(format!( "corrupt agent run mode '{other}'" ))), } } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum AgentRunStatus { Queued, Running, WaitingChildren, Completed, Failed, TimedOut, Cancelled, Interrupted, } impl AgentRunStatus { pub fn as_str(&self) -> &'static str { match self { Self::Queued => "queued", Self::Running => "running", Self::WaitingChildren => "waiting_children", Self::Completed => "completed", Self::Failed => "failed", Self::TimedOut => "timed_out", Self::Cancelled => "cancelled", Self::Interrupted => "interrupted", } } pub fn parse(value: &str) -> Result { match value { "queued" => Ok(Self::Queued), "running" => Ok(Self::Running), "waiting_children" => Ok(Self::WaitingChildren), "completed" => Ok(Self::Completed), "failed" => Ok(Self::Failed), "timed_out" => Ok(Self::TimedOut), "cancelled" => Ok(Self::Cancelled), "interrupted" => Ok(Self::Interrupted), other => Err(StorageError::Migration(format!( "corrupt agent run status '{other}'" ))), } } pub fn is_terminal(self) -> bool { !matches!(self, Self::Queued | Self::Running | Self::WaitingChildren) } } #[derive(Debug, Clone)] pub struct AgentRunRecord { pub id: String, pub root_session_id: String, pub root_turn_id: Option, pub parent_run_id: Option, pub caller_agent_id: String, pub caller_scope_id: String, pub idempotency_key: Option, pub agent_id: String, pub definition_hash: String, pub provider_profile: String, pub provider_name: String, pub model_id: String, pub mode: AgentRunMode, pub depth: i64, pub plan_item_id: Option, pub execution_id: String, pub task: String, pub context_json: Option, pub budget_json: String, pub signal_contract_json: Option, pub signal_delivery: Option, pub status: AgentRunStatus, pub result: Option, pub error: Option, pub prompt_tokens: Option, pub completion_tokens: Option, pub cost: Option, pub tool_calls_count: i64, pub iterations: i64, pub runtime_generation: i64, pub attempt: i64, pub completion_slot_reserved: bool, pub deadline_at: i64, pub revision: i64, pub started_at: Option, pub finished_at: Option, pub created_at: i64, pub updated_at: i64, } /// Raw persisted transcript row for an Agent run. Incrementally appended by /// the run's transcript writer; `tool_calls_json` is stored verbatim and only /// parsed into `providers::ToolCall` at the protocol boundary. #[derive(Debug, Clone)] pub struct AgentRunMessageRecord { pub id: String, pub run_id: String, pub seq: i64, pub role: String, pub content: String, pub reasoning_content: Option, pub tool_call_id: Option, pub tool_name: Option, pub tool_calls_json: Option, pub created_at: i64, } /// One run to admit inside `accept_agent_runs`. #[derive(Debug, Clone)] pub struct NewAgentRun { pub id: String, pub root_session_id: String, pub root_turn_id: Option, pub parent_run_id: Option, pub caller_agent_id: String, pub caller_scope_id: String, pub idempotency_key: Option, pub agent_id: String, pub definition_hash: String, pub provider_profile: String, pub provider_name: String, pub model_id: String, pub mode: AgentRunMode, pub depth: i64, pub plan_item_id: Option, pub execution_id: String, pub task: String, pub context_json: Option, pub budget_json: String, pub signal_contract_json: Option, pub signal_delivery: Option, pub deadline_at: i64, pub runtime_generation: i64, /// Background runs reserve a completion slot at admission so their /// completion can never be lost to inbox capacity exhaustion. pub completion_slot_reserved: bool, } /// Batch admission request. Each run carries its own idempotency key; a /// single-task request is just a one-element batch. #[derive(Debug, Clone)] pub struct AcceptAgentRequest { pub runs: Vec, pub now: i64, } #[derive(Debug)] pub enum AcceptedAgentRuns { Accepted { runs: Vec, }, /// Idempotent retry: the run already existed for this key. Existing { runs: Vec, }, } /// Terminal outcome produced by a runner. The Coordinator persists it; the /// runner itself never writes channels or plan state. #[derive(Debug, Clone)] pub enum AgentTerminalOutcome { Completed { result: String, prompt_tokens: Option, completion_tokens: Option, cost: Option, tool_calls: i64, iterations: i64, /// Signal IDs emitted by this run; included in the completion /// payload so the main Agent can recognise duplicates (design ยง12.3). signal_ids: Vec, }, Failed { error: String, prompt_tokens: Option, completion_tokens: Option, cost: Option, signal_ids: Vec, }, TimedOut { deadline_at: i64, signal_ids: Vec, }, Cancelled { reason: String, signal_ids: Vec, }, Interrupted { reason: String, signal_ids: Vec, }, } impl AgentTerminalOutcome { pub fn status(&self) -> AgentRunStatus { match self { Self::Completed { .. } => AgentRunStatus::Completed, Self::Failed { .. } => AgentRunStatus::Failed, Self::TimedOut { .. } => AgentRunStatus::TimedOut, Self::Cancelled { .. } => AgentRunStatus::Cancelled, Self::Interrupted { .. } => AgentRunStatus::Interrupted, } } } #[derive(Debug, Clone)] pub struct TerminalCommit { pub run: AgentRunRecord, } const RUN_COLUMNS: &str = "id, root_session_id, root_turn_id, parent_run_id, \ caller_agent_id, caller_scope_id, idempotency_key, agent_id, definition_hash, \ provider_profile, provider_name, model_id, mode, depth, plan_item_id, execution_id, \ task, context_json, budget_json, signal_contract_json, signal_delivery, \ status, result, error, prompt_tokens, completion_tokens, cost, tool_calls_count, \ iterations, runtime_generation, attempt, completion_slot_reserved, deadline_at, \ revision, started_at, finished_at, created_at, updated_at"; fn run_record_from_row(row: &sqlx::sqlite::SqliteRow) -> Result { Ok(AgentRunRecord { id: row.get("id"), root_session_id: row.get("root_session_id"), root_turn_id: row.get("root_turn_id"), parent_run_id: row.get("parent_run_id"), caller_agent_id: row.get("caller_agent_id"), caller_scope_id: row.get("caller_scope_id"), idempotency_key: row.get("idempotency_key"), agent_id: row.get("agent_id"), definition_hash: row.get("definition_hash"), provider_profile: row.get("provider_profile"), provider_name: row.get("provider_name"), model_id: row.get("model_id"), mode: AgentRunMode::parse(row.get::<&str, _>("mode"))?, depth: row.get("depth"), plan_item_id: row.get("plan_item_id"), execution_id: row.get("execution_id"), task: row.get("task"), context_json: row.get("context_json"), budget_json: row.get("budget_json"), signal_contract_json: row.get("signal_contract_json"), signal_delivery: row.get("signal_delivery"), status: AgentRunStatus::parse(row.get::<&str, _>("status"))?, result: row.get("result"), error: row.get("error"), prompt_tokens: row.get("prompt_tokens"), completion_tokens: row.get("completion_tokens"), cost: row.get("cost"), tool_calls_count: row.get("tool_calls_count"), iterations: row.get("iterations"), runtime_generation: row.get("runtime_generation"), attempt: row.get("attempt"), completion_slot_reserved: row.get::("completion_slot_reserved") != 0, deadline_at: row.get("deadline_at"), revision: row.get("revision"), started_at: row.get("started_at"), finished_at: row.get("finished_at"), created_at: row.get("created_at"), updated_at: row.get("updated_at"), }) } fn agent_run_message_record_from_row( row: &sqlx::sqlite::SqliteRow, ) -> Result { Ok(AgentRunMessageRecord { id: row.get("id"), run_id: row.get("run_id"), seq: row.get("seq"), role: row.get("role"), content: row.get("content"), reasoning_content: row.get("reasoning_content"), tool_call_id: row.get("tool_call_id"), tool_name: row.get("tool_name"), tool_calls_json: row.get("tool_calls_json"), created_at: row.get("created_at"), }) } impl super::Storage { /// Admit a batch of runs in one transaction, claiming any referenced /// plan items atomically. If any plan item was already taken the whole /// admission rolls back so a run can never diverge from the plan it /// claims to execute. pub async fn accept_agent_runs( &self, request: AcceptAgentRequest, ) -> Result { if request.runs.is_empty() { return Err(StorageError::Conflict( "agent admission requires at least one run".to_string(), )); } let mut tx = self.pool.begin().await?; for run in &request.runs { let inserted = sqlx::query( "INSERT INTO agent_runs (id, root_session_id, root_turn_id, \ parent_run_id, caller_agent_id, caller_scope_id, idempotency_key, \ agent_id, definition_hash, provider_profile, provider_name, model_id, \ mode, depth, plan_item_id, execution_id, task, context_json, budget_json, \ signal_contract_json, signal_delivery, \ status, runtime_generation, attempt, completion_slot_reserved, deadline_at, \ revision, created_at, updated_at) \ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, \ ?, 'queued', ?, 1, ?, ?, 0, ?, ?)", ) .bind(&run.id) .bind(&run.root_session_id) .bind(&run.root_turn_id) .bind(&run.parent_run_id) .bind(&run.caller_agent_id) .bind(&run.caller_scope_id) .bind(&run.idempotency_key) .bind(&run.agent_id) .bind(&run.definition_hash) .bind(&run.provider_profile) .bind(&run.provider_name) .bind(&run.model_id) .bind(run.mode.as_str()) .bind(run.depth) .bind(&run.plan_item_id) .bind(&run.execution_id) .bind(&run.task) .bind(&run.context_json) .bind(&run.budget_json) .bind(&run.signal_contract_json) .bind(&run.signal_delivery) .bind(run.runtime_generation) .bind(i64::from(run.completion_slot_reserved)) .bind(run.deadline_at) .bind(request.now) .bind(request.now) .execute(&mut *tx) .await? .rows_affected() == 1; if !inserted { drop(tx); return self.existing_agent_admission(request).await; } if let Some(item_id) = run.plan_item_id.as_deref() { claim_plan_item( &mut tx, &run.root_session_id, item_id, &run.execution_id, request.now, ) .await?; } } tx.commit().await?; let mut runs = Vec::with_capacity(request.runs.len()); for run in &request.runs { runs.push(self.get_agent_run(&run.id).await?.ok_or_else(|| { StorageError::NotFound(format!("agent run {} vanished after admission", run.id)) })?); } Ok(AcceptedAgentRuns::Accepted { runs }) } async fn existing_agent_admission( &self, request: AcceptAgentRequest, ) -> Result { let mut runs = Vec::new(); for run in &request.runs { if let Some(record) = self.get_agent_run(&run.id).await? { runs.push(record); } } if runs.is_empty() { return Err(StorageError::Conflict( "agent admission conflicted but no existing rows were found".to_string(), )); } Ok(AcceptedAgentRuns::Existing { runs }) } pub async fn get_agent_run( &self, run_id: &str, ) -> Result, StorageError> { let row = sqlx::query(sqlx::AssertSqlSafe(format!( "SELECT {RUN_COLUMNS} FROM agent_runs WHERE id = ?" ))) .bind(run_id) .fetch_optional(&self.pool) .await?; match row { Some(row) => Ok(Some(run_record_from_row(&row)?)), None => Ok(None), } } /// Append one transcript message for an Agent run. The writer owns the /// monotonically increasing `seq`; `provider_state` is expected to have /// been stripped by the caller before this is called. pub async fn append_agent_run_message( &self, run_id: &str, seq: i64, message: &crate::bus::ChatMessage, now: i64, ) -> Result<(), StorageError> { let tool_calls_json = message .tool_calls .as_ref() .map(serde_json::to_string) .transpose() .map_err(|error| StorageError::Migration(format!("serialize tool_calls: {error}")))?; sqlx::query( "INSERT INTO agent_run_messages (id, run_id, seq, role, content, \ reasoning_content, tool_call_id, tool_name, tool_calls_json, created_at) \ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(uuid::Uuid::new_v4().to_string()) .bind(run_id) .bind(seq) .bind(&message.role) .bind(&message.content) .bind(&message.reasoning_content) .bind(&message.tool_call_id) .bind(&message.tool_name) .bind(tool_calls_json) .bind(now) .execute(&self.pool) .await?; Ok(()) } /// List the persisted transcript for a run ordered by `seq`. The /// transcript is naturally bounded by the run's iteration budget; the /// default `limit` is a generous ceiling, not a pagination contract. pub async fn list_agent_run_messages( &self, run_id: &str, limit: i64, ) -> Result, StorageError> { let rows = sqlx::query( "SELECT id, run_id, seq, role, content, reasoning_content, tool_call_id, \ tool_name, tool_calls_json, created_at \ FROM agent_run_messages WHERE run_id = ? ORDER BY seq ASC LIMIT ?", ) .bind(run_id) .bind(limit) .fetch_all(&self.pool) .await?; rows.iter().map(agent_run_message_record_from_row).collect() } /// List runs for a session ordered by `(created_at DESC, id DESC)`. /// The cursor is the pair of the last row the client has seen. pub async fn list_agent_runs( &self, root_session_id: &str, cursor: Option<(i64, String)>, limit: i64, ) -> Result, StorageError> { let limit = limit.clamp(1, 200); let rows = match cursor { Some((created_at, id)) => { sqlx::query(sqlx::AssertSqlSafe(format!( "SELECT {RUN_COLUMNS} FROM agent_runs \ WHERE root_session_id = ? AND (created_at < ? OR (created_at = ? AND id < ?)) \ ORDER BY created_at DESC, id DESC LIMIT ?" ))) .bind(root_session_id) .bind(created_at) .bind(created_at) .bind(id) .bind(limit) .fetch_all(&self.pool) .await? } None => { sqlx::query(sqlx::AssertSqlSafe(format!( "SELECT {RUN_COLUMNS} FROM agent_runs \ WHERE root_session_id = ? ORDER BY created_at DESC, id DESC LIMIT ?" ))) .bind(root_session_id) .bind(limit) .fetch_all(&self.pool) .await? } }; rows.iter().map(run_record_from_row).collect() } /// All durable runs across sessions, newest first (management union). pub async fn list_all_agent_runs( &self, cursor: Option<(i64, String)>, limit: i64, ) -> Result, StorageError> { let limit = limit.clamp(1, 200); let rows = match cursor { Some((created_at, id)) => { sqlx::query(sqlx::AssertSqlSafe(format!( "SELECT {RUN_COLUMNS} FROM agent_runs \ WHERE (created_at < ? OR (created_at = ? AND id < ?)) \ ORDER BY created_at DESC, id DESC LIMIT ?" ))) .bind(created_at) .bind(created_at) .bind(id) .bind(limit) .fetch_all(&self.pool) .await? } None => { sqlx::query(sqlx::AssertSqlSafe(format!( "SELECT {RUN_COLUMNS} FROM agent_runs \ ORDER BY created_at DESC, id DESC LIMIT ?" ))) .bind(limit) .fetch_all(&self.pool) .await? } }; rows.iter().map(run_record_from_row).collect() } /// Conditional `queued -> running` transition owned by this execution. pub async fn mark_agent_run_running( &self, run_id: &str, execution_id: &str, now: i64, ) -> Result { let rows = sqlx::query( "UPDATE agent_runs SET status = 'running', started_at = ?, updated_at = ? \ WHERE id = ? AND execution_id = ? AND status = 'queued'", ) .bind(now) .bind(now) .bind(run_id) .bind(execution_id) .execute(&self.pool) .await? .rows_affected(); Ok(rows == 1) } /// Conditional transition into `waiting_children` from the expected /// nonterminal status while this execution still owns the run. pub async fn mark_agent_run_waiting_children( &self, run_id: &str, execution_id: &str, expected: AgentRunStatus, now: i64, ) -> Result { if expected.is_terminal() { return Err(StorageError::Conflict(format!( "cannot wait on children from terminal status {}", expected.as_str() ))); } let rows = sqlx::query( "UPDATE agent_runs SET status = 'waiting_children', updated_at = ? \ WHERE id = ? AND execution_id = ? AND status = ?", ) .bind(now) .bind(run_id) .bind(execution_id) .bind(expected.as_str()) .execute(&self.pool) .await? .rows_affected(); Ok(rows == 1) } /// Restore a waiting parent to `running` once its children settled. A /// run that was cancelled/timed out in the meantime keeps its terminal /// state. pub async fn restore_agent_run_running( &self, run_id: &str, execution_id: &str, now: i64, ) -> Result { let rows = sqlx::query( "UPDATE agent_runs SET status = 'running', updated_at = ? \ WHERE id = ? AND execution_id = ? AND status = 'waiting_children'", ) .bind(now) .bind(run_id) .bind(execution_id) .execute(&self.pool) .await? .rows_affected(); Ok(rows == 1) } /// Cancel a nonterminal run. Returns true when this call owned the /// transition. pub async fn cancel_agent_run( &self, run_id: &str, reason: &str, now: i64, ) -> Result { let rows = sqlx::query( "UPDATE agent_runs SET status = 'cancelled', error = ?, finished_at = ?, updated_at = ? \ WHERE id = ? AND status IN ('queued', 'running', 'waiting_children')", ) .bind(reason) .bind(now) .bind(now) .bind(run_id) .execute(&self.pool) .await? .rows_affected(); Ok(rows == 1) } /// Cancel a nonterminal run and resolve its completion reservation in the /// same transaction. With `suppress_continuation` (explicit `/stop` or /// lifecycle cancellation) the completion event is written directly as /// `consumed` so no continuation Turn restarts after cancellation; the /// audit fact is preserved either way. pub async fn cancel_agent_run_with_completion( &self, run_id: &str, reason: &str, suppress_continuation: bool, now: i64, ) -> Result { let mut tx = self.pool.begin().await?; let row = sqlx::query( "SELECT completion_slot_reserved FROM agent_runs \ WHERE id = ? AND status IN ('queued', 'running', 'waiting_children')", ) .bind(run_id) .fetch_optional(&mut *tx) .await?; let Some(row) = row else { return Ok(false); }; let reserved: bool = row.get::("completion_slot_reserved") != 0; sqlx::query( "UPDATE agent_runs SET status = 'cancelled', error = ?, finished_at = ?, updated_at = ? \ WHERE id = ? AND status IN ('queued', 'running', 'waiting_children')", ) .bind(reason) .bind(now) .bind(now) .bind(run_id) .execute(&mut *tx) .await?; if reserved { let (session, agent_id, task): (String, String, String) = sqlx::query_as( "SELECT root_session_id, agent_id, task FROM agent_runs WHERE id = ?", ) .bind(run_id) .fetch_one(&mut *tx) .await?; let revision: i64 = sqlx::query_scalar( "UPDATE agent_session_state \ SET reserved_completion_slots = MAX(reserved_completion_slots - 1, 0), \ revision = revision + 1, updated_at = ? \ WHERE root_session_id = ? RETURNING revision", ) .bind(now) .bind(&session) .fetch_one(&mut *tx) .await?; let event = super::agent_inbox::NewInboxEvent { id: uuid::Uuid::new_v4().to_string(), root_session_id: session, run_id: Some(run_id.to_string()), event_type: super::agent_inbox::AgentEventType::Completion, event_key: format!("completion:{run_id}"), delivery: super::agent_inbox::AgentEventDelivery::Queue, requires_continuation: !suppress_continuation, severity: Some("warning".to_string()), payload_json: super::agent_inbox::completion_payload( run_id, &agent_id, &task, None, "cancelled", Some(reason), &[], ), }; if suppress_continuation { // Directly consumed: pending count never grows. sqlx::query( "INSERT INTO agent_inbox_events (id, root_session_id, run_id, event_type, \ event_key, delivery, requires_continuation, severity, payload_json, \ status, attempt_count, revision, next_attempt_at, created_at, \ updated_at, consumed_at) \ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'consumed', 0, ?, NULL, ?, ?, ?)", ) .bind(&event.id) .bind(&event.root_session_id) .bind(&event.run_id) .bind(event.event_type.as_str()) .bind(&event.event_key) .bind(event.delivery.as_str()) .bind(i64::from(event.requires_continuation)) .bind(&event.severity) .bind(&event.payload_json) .bind(revision) .bind(now) .bind(now) .bind(now) .execute(&mut *tx) .await?; } else { super::agent_inbox::insert_event_tx(&mut tx, &event, revision, now).await?; } } tx.commit().await?; Ok(true) } /// Atomically commit a terminal outcome. The conditional update makes /// exactly one writer the owner; late results from stale executions /// update zero rows and return `None`. pub async fn commit_agent_terminal( &self, run_id: &str, execution_id: &str, runtime_generation: i64, outcome: &AgentTerminalOutcome, plan_summary: Option<&str>, now: i64, ) -> Result, StorageError> { let mut tx = self.pool.begin().await?; let (result, error, prompt_tokens, completion_tokens, cost, tool_calls, iterations) = match outcome { AgentTerminalOutcome::Completed { result, prompt_tokens, completion_tokens, cost, tool_calls, iterations, .. } => ( Some(result.as_str()), None, *prompt_tokens, *completion_tokens, *cost, *tool_calls, *iterations, ), AgentTerminalOutcome::Failed { error, prompt_tokens, completion_tokens, cost, .. } => ( None, Some(error.as_str()), *prompt_tokens, *completion_tokens, *cost, 0, 0, ), AgentTerminalOutcome::TimedOut { .. } => { (None, Some("deadline exceeded"), None, None, None, 0, 0) } AgentTerminalOutcome::Cancelled { reason, .. } => { (None, Some(reason.as_str()), None, None, None, 0, 0) } AgentTerminalOutcome::Interrupted { reason, .. } => { (None, Some(reason.as_str()), None, None, None, 0, 0) } }; let updated = sqlx::query( "UPDATE agent_runs SET status = ?, result = ?, error = ?, prompt_tokens = ?, \ completion_tokens = ?, cost = ?, tool_calls_count = ?, iterations = ?, \ finished_at = ?, updated_at = ? \ WHERE id = ? AND execution_id = ? AND runtime_generation = ? \ AND status IN ('queued', 'running', 'waiting_children')", ) .bind(outcome.status().as_str()) .bind(result) .bind(error) .bind(prompt_tokens) .bind(completion_tokens) .bind(cost) .bind(tool_calls) .bind(iterations) .bind(now) .bind(now) .bind(run_id) .bind(execution_id) .bind(runtime_generation) .execute(&mut *tx) .await? .rows_affected(); if updated != 1 { return Ok(None); } let run_row = sqlx::query(sqlx::AssertSqlSafe(format!( "SELECT {RUN_COLUMNS} FROM agent_runs WHERE id = ?" ))) .bind(run_id) .fetch_one(&mut *tx) .await?; let run = run_record_from_row(&run_row)?; if let Some(item_id) = run.plan_item_id.as_deref() { finish_plan_item( &mut tx, &run.root_session_id, item_id, &run.execution_id, matches!(outcome, AgentTerminalOutcome::Completed { .. }), plan_summary, now, ) .await?; } // Background runs that reserved a completion slot convert the // reservation into a durable completion event in the same commit. // The event survives restarts, queue-full conditions and lost wakes. if run.completion_slot_reserved { let (status, error, signal_ids, result) = match outcome { AgentTerminalOutcome::Completed { result, signal_ids, .. } => ( "completed", None, signal_ids.as_slice(), Some(result.as_str()), ), AgentTerminalOutcome::Failed { error, signal_ids, .. } => ("failed", Some(error.as_str()), signal_ids.as_slice(), None), AgentTerminalOutcome::TimedOut { signal_ids, .. } => ( "timed_out", Some("deadline exceeded"), signal_ids.as_slice(), None, ), AgentTerminalOutcome::Cancelled { reason, signal_ids } => ( "cancelled", Some(reason.as_str()), signal_ids.as_slice(), None, ), AgentTerminalOutcome::Interrupted { reason, signal_ids } => ( "interrupted", Some(reason.as_str()), signal_ids.as_slice(), None, ), }; super::agent_inbox::insert_completion_event_tx( &mut tx, &run.id, &run.root_session_id, &run.agent_id, &run.task, result, status, error, signal_ids, now, ) .await?; } tx.commit().await?; Ok(Some(TerminalCommit { run })) } } async fn claim_plan_item( tx: &mut SqliteConnection, session_id: &str, item_id: &str, execution_id: &str, now: i64, ) -> Result<(), StorageError> { let plan_id: Option = sqlx::query_scalar( "SELECT id FROM task_plans WHERE session_id = ? AND status = 'active' LIMIT 1", ) .bind(session_id) .fetch_optional(&mut *tx) .await?; let Some(plan_id) = plan_id else { return Err(StorageError::Conflict(format!( "plan item {item_id} cannot be claimed without an active plan" ))); }; let rows = sqlx::query( "UPDATE task_items SET status = 'in_progress', executor_kind = 'sub_agent', \ execution_id = ?, error = NULL, version = version + 1, updated_at = ? \ WHERE plan_id = ? AND id = ? AND status = 'pending'", ) .bind(execution_id) .bind(now) .bind(&plan_id) .bind(item_id) .execute(&mut *tx) .await? .rows_affected(); if rows != 1 { return Err(StorageError::Conflict(format!( "plan item {item_id} was already claimed by another execution" ))); } bump_plan_version(tx, &plan_id, now).await } async fn finish_plan_item( tx: &mut SqliteConnection, session_id: &str, item_id: &str, execution_id: &str, completed: bool, summary: Option<&str>, now: i64, ) -> Result<(), StorageError> { let plan_id: Option = sqlx::query_scalar( "SELECT id FROM task_plans WHERE session_id = ? AND status = 'active' LIMIT 1", ) .bind(session_id) .fetch_optional(&mut *tx) .await?; let Some(plan_id) = plan_id else { return Ok(()); }; let status = if completed { "completed" } else { "blocked" }; let rows = sqlx::query( "UPDATE task_items SET status = ?, result_summary = ?, error = ?, \ version = version + 1, updated_at = ? \ WHERE plan_id = ? AND id = ? AND execution_id = ? AND status = 'in_progress'", ) .bind(status) .bind(completed.then_some(summary).flatten()) .bind((!completed).then_some(summary).flatten()) .bind(now) .bind(&plan_id) .bind(item_id) .bind(execution_id) .execute(&mut *tx) .await? .rows_affected(); if rows != 1 { return Ok(()); } bump_plan_version(tx, &plan_id, now).await } async fn bump_plan_version( tx: &mut SqliteConnection, plan_id: &str, now: i64, ) -> Result<(), StorageError> { sqlx::query("UPDATE task_plans SET version = version + 1, updated_at = ? WHERE id = ?") .bind(now) .bind(plan_id) .execute(&mut *tx) .await?; Ok(()) } #[cfg(test)] mod tests { use super::*; use tempfile::TempDir; async fn create_test_storage() -> (super::super::Storage, TempDir) { let dir = tempfile::tempdir().unwrap(); let db_path = dir.path().join("agent.db"); let storage = super::super::Storage::new(&db_path).await.unwrap(); (storage, dir) } fn new_run(id: &str, execution_id: &str, session: &str) -> NewAgentRun { NewAgentRun { id: id.to_string(), root_session_id: session.to_string(), root_turn_id: None, parent_run_id: None, caller_agent_id: "ROOT".to_string(), caller_scope_id: "turn-1".to_string(), idempotency_key: None, agent_id: "researcher".to_string(), definition_hash: "hash".to_string(), provider_profile: "research".to_string(), provider_name: "test".to_string(), model_id: "test-model".to_string(), mode: AgentRunMode::Foreground, depth: 1, plan_item_id: None, execution_id: execution_id.to_string(), task: "do the work".to_string(), context_json: None, budget_json: "{\"remaining_runs\":15}".to_string(), signal_contract_json: None, signal_delivery: None, deadline_at: 1_000, runtime_generation: 1, completion_slot_reserved: false, } } #[tokio::test] async fn fresh_database_creates_current_agent_tables() { let (storage, _dir) = create_test_storage().await; let version: i64 = sqlx::query_scalar("PRAGMA user_version") .fetch_one(storage.pool()) .await .unwrap(); assert_eq!(version, 10); for table in [ "agent_runs", "agent_session_state", "agent_inbox_events", "agent_run_messages", ] { let exists: i64 = sqlx::query_scalar( "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name = ?", ) .bind(table) .fetch_one(storage.pool()) .await .unwrap(); assert_eq!(exists, 1, "missing table {table}"); } } #[tokio::test] async fn single_run_admission_persists_queued_without_group() { let (storage, _dir) = create_test_storage().await; let accepted = storage .accept_agent_runs(AcceptAgentRequest { runs: vec![new_run("run-1", "exec-1", "cli:test:d1")], now: 100, }) .await .unwrap(); assert!(matches!(accepted, AcceptedAgentRuns::Accepted { .. })); let run = storage.get_agent_run("run-1").await.unwrap().unwrap(); assert_eq!(run.status, AgentRunStatus::Queued); assert_eq!(run.execution_id, "exec-1"); } #[tokio::test] async fn stale_execution_cannot_commit_terminal() { let (storage, _dir) = create_test_storage().await; storage .accept_agent_runs(AcceptAgentRequest { runs: vec![new_run("run-1", "exec-1", "cli:test:d1")], now: 100, }) .await .unwrap(); let stale = storage .commit_agent_terminal( "run-1", "exec-other", 1, &AgentTerminalOutcome::Completed { result: "late".to_string(), prompt_tokens: None, completion_tokens: None, cost: None, tool_calls: 0, iterations: 0, signal_ids: Vec::new(), }, None, 120, ) .await .unwrap(); assert!(stale.is_none()); let run = storage.get_agent_run("run-1").await.unwrap().unwrap(); assert_eq!(run.status, AgentRunStatus::Queued); assert!(run.result.is_none()); } #[tokio::test] async fn waiting_children_transitions_are_conditional() { let (storage, _dir) = create_test_storage().await; storage .accept_agent_runs(AcceptAgentRequest { runs: vec![new_run("run-1", "exec-1", "cli:test:d1")], now: 100, }) .await .unwrap(); storage .mark_agent_run_running("run-1", "exec-1", 110) .await .unwrap(); assert!( storage .mark_agent_run_waiting_children("run-1", "exec-1", AgentRunStatus::Running, 120) .await .unwrap() ); let run = storage.get_agent_run("run-1").await.unwrap().unwrap(); assert_eq!(run.status, AgentRunStatus::WaitingChildren); assert!( storage .restore_agent_run_running("run-1", "exec-1", 130) .await .unwrap() ); let run = storage.get_agent_run("run-1").await.unwrap().unwrap(); assert_eq!(run.status, AgentRunStatus::Running); } #[tokio::test] async fn plan_item_claim_is_atomic_with_run_admission() { let (storage, _dir) = create_test_storage().await; sqlx::query( "INSERT INTO sessions (id, channel, chat_id, dialog_id, created_at, last_active_at) VALUES ('cli:test:d1', 'cli', 'test', 'd1', 1, 1)", ) .execute(storage.pool()) .await .unwrap(); sqlx::query( "INSERT INTO task_plans (id, session_id, objective, status, version, created_at, updated_at) VALUES ('plan-1', 'cli:test:d1', 'obj', 'active', 1, 1, 1)", ) .execute(storage.pool()) .await .unwrap(); sqlx::query( "INSERT INTO task_items (id, plan_id, ordinal, title, status, version, created_at, updated_at) VALUES ('T1', 'plan-1', 1, 'work', 'pending', 1, 1, 1)", ) .execute(storage.pool()) .await .unwrap(); let mut run = new_run("run-1", "exec-1", "cli:test:d1"); run.plan_item_id = Some("T1".to_string()); storage .accept_agent_runs(AcceptAgentRequest { runs: vec![run.clone()], now: 100, }) .await .unwrap(); let status: String = sqlx::query_scalar("SELECT status FROM task_items WHERE id = 'T1'") .fetch_one(storage.pool()) .await .unwrap(); assert_eq!(status, "in_progress"); // A second admission for the same item must roll back entirely. run.id = "run-2".to_string(); run.execution_id = "exec-2".to_string(); let error = storage .accept_agent_runs(AcceptAgentRequest { runs: vec![run], now: 110, }) .await .unwrap_err(); assert!(matches!(error, StorageError::Conflict(_))); assert!(storage.get_agent_run("run-2").await.unwrap().is_none()); // Terminal commit releases the item as completed with the summary. storage .commit_agent_terminal( "run-1", "exec-1", 1, &AgentTerminalOutcome::Completed { result: "done".to_string(), prompt_tokens: None, completion_tokens: None, cost: None, tool_calls: 0, iterations: 0, signal_ids: Vec::new(), }, Some("finished the work"), 120, ) .await .unwrap(); let (status, summary): (String, Option) = sqlx::query_as("SELECT status, result_summary FROM task_items WHERE id = 'T1'") .fetch_one(storage.pool()) .await .unwrap(); assert_eq!(status, "completed"); assert_eq!(summary.as_deref(), Some("finished the work")); } #[tokio::test] async fn list_agent_runs_paginates_with_created_at_cursor() { let (storage, _dir) = create_test_storage().await; let mut runs = Vec::new(); for index in 0..5 { runs.push(new_run( &format!("run-{index}"), &format!("exec-{index}"), "cli:test:d1", )); } storage .accept_agent_runs(AcceptAgentRequest { runs, now: 100 }) .await .unwrap(); let first_page = storage .list_agent_runs("cli:test:d1", None, 2) .await .unwrap(); assert_eq!(first_page.len(), 2); let last = first_page.last().unwrap(); let second_page = storage .list_agent_runs("cli:test:d1", Some((last.created_at, last.id.clone())), 10) .await .unwrap(); assert_eq!(second_page.len(), 3); let seen: std::collections::HashSet<_> = first_page .iter() .chain(second_page.iter()) .map(|run| run.id.clone()) .collect(); assert_eq!(seen.len(), 5); } #[tokio::test] async fn cancel_agent_run_only_transitions_nonterminal_rows() { let (storage, _dir) = create_test_storage().await; storage .accept_agent_runs(AcceptAgentRequest { runs: vec![new_run("run-1", "exec-1", "cli:test:d1")], now: 100, }) .await .unwrap(); assert!( storage .cancel_agent_run("run-1", "stopped", 110) .await .unwrap() ); assert!( !storage .cancel_agent_run("run-1", "stopped", 120) .await .unwrap() ); let run = storage.get_agent_run("run-1").await.unwrap().unwrap(); assert_eq!(run.status, AgentRunStatus::Cancelled); assert_eq!(run.error.as_deref(), Some("stopped")); } #[tokio::test] async fn suppress_cancel_converts_reservation_to_consumed_completion() { let (storage, _dir) = create_test_storage().await; let mut run = new_run("run-1", "exec-1", "cli:test:d1"); run.mode = AgentRunMode::Background; run.completion_slot_reserved = true; storage .accept_agent_runs(AcceptAgentRequest { runs: vec![run], now: 100, }) .await .unwrap(); storage .reserve_completion_slots("cli:test:d1", 1, 16, 100) .await .unwrap(); assert!( storage .cancel_agent_run_with_completion("run-1", "stopped", true, 110) .await .unwrap() ); let run = storage.get_agent_run("run-1").await.unwrap().unwrap(); assert_eq!(run.status, AgentRunStatus::Cancelled); let events = storage .list_agent_inbox_events("cli:test:d1", 10) .await .unwrap(); assert_eq!(events.len(), 1); assert_eq!( events[0].status, crate::storage::agent_inbox::AgentEventStatus::Consumed ); assert!(!events[0].requires_continuation); let state: (i64, i64) = sqlx::query_as( "SELECT pending_event_count, reserved_completion_slots FROM agent_session_state \ WHERE root_session_id = 'cli:test:d1'", ) .fetch_one(storage.pool()) .await .unwrap(); assert_eq!(state, (0, 0)); // A second cancel is a no-op. assert!( !storage .cancel_agent_run_with_completion("run-1", "again", true, 120) .await .unwrap() ); } #[tokio::test] async fn transcript_messages_round_trip_in_seq_order() { let (storage, _dir) = create_test_storage().await; storage .accept_agent_runs(AcceptAgentRequest { runs: vec![new_run("run-1", "exec-1", "cli:test:d1")], now: 100, }) .await .unwrap(); let mut assistant = crate::bus::ChatMessage::assistant_with_tool_calls( "calling".to_string(), vec![crate::providers::ToolCall { id: "call-1".to_string(), name: "bash".to_string(), arguments: serde_json::json!({}), }], ); assistant.reasoning_content = Some("thinking".to_string()); let tool = crate::bus::ChatMessage::tool("call-1", "bash", "output"); storage .append_agent_run_message("run-1", 0, &assistant, 200) .await .unwrap(); storage .append_agent_run_message("run-1", 1, &tool, 201) .await .unwrap(); let messages = storage .list_agent_run_messages("run-1", 10_000) .await .unwrap(); assert_eq!(messages.len(), 2); assert_eq!(messages[0].seq, 0); assert_eq!(messages[0].role, "assistant"); assert_eq!(messages[0].reasoning_content.as_deref(), Some("thinking")); assert!(messages[0].tool_calls_json.is_some()); assert_eq!(messages[1].seq, 1); assert_eq!(messages[1].role, "tool"); assert_eq!(messages[1].tool_call_id.as_deref(), Some("call-1")); assert_eq!(messages[1].tool_name.as_deref(), Some("bash")); } }