PicoBot/src/storage/agent_run.rs

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<Self, StorageError> {
        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<Self, StorageError> {
        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<String>,
    pub parent_run_id: Option<String>,
    pub caller_agent_id: String,
    pub caller_scope_id: String,
    pub idempotency_key: Option<String>,
    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<String>,
    pub execution_id: String,
    pub task: String,
    pub context_json: Option<String>,
    pub budget_json: String,
    pub signal_contract_json: Option<String>,
    pub signal_delivery: Option<String>,
    pub status: AgentRunStatus,
    pub result: Option<String>,
    pub error: Option<String>,
    pub prompt_tokens: Option<i64>,
    pub completion_tokens: Option<i64>,
    pub cost: Option<f64>,
    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<i64>,
    pub finished_at: Option<i64>,
    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<String>,
    pub tool_call_id: Option<String>,
    pub tool_name: Option<String>,
    pub tool_calls_json: Option<String>,
    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<String>,
    pub parent_run_id: Option<String>,
    pub caller_agent_id: String,
    pub caller_scope_id: String,
    pub idempotency_key: Option<String>,
    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<String>,
    pub execution_id: String,
    pub task: String,
    pub context_json: Option<String>,
    pub budget_json: String,
    pub signal_contract_json: Option<String>,
    pub signal_delivery: Option<String>,
    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<NewAgentRun>,
    pub now: i64,
}

#[derive(Debug)]
pub enum AcceptedAgentRuns {
    Accepted {
        runs: Vec<AgentRunRecord>,
    },
    /// Idempotent retry: the run already existed for this key.
    Existing {
        runs: Vec<AgentRunRecord>,
    },
}

/// 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<i64>,
        completion_tokens: Option<i64>,
        cost: Option<f64>,
        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<String>,
    },
    Failed {
        error: String,
        prompt_tokens: Option<i64>,
        completion_tokens: Option<i64>,
        cost: Option<f64>,
        signal_ids: Vec<String>,
    },
    TimedOut {
        deadline_at: i64,
        signal_ids: Vec<String>,
    },
    Cancelled {
        reason: String,
        signal_ids: Vec<String>,
    },
    Interrupted {
        reason: String,
        signal_ids: Vec<String>,
    },
}

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<AgentRunRecord, StorageError> {
    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::<i64, _>("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<AgentRunMessageRecord, StorageError> {
    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<AcceptedAgentRuns, StorageError> {
        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<AcceptedAgentRuns, StorageError> {
        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<Option<AgentRunRecord>, 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<Vec<AgentRunMessageRecord>, 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<Vec<AgentRunRecord>, 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<Vec<AgentRunRecord>, 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<bool, StorageError> {
        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<bool, StorageError> {
        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<bool, StorageError> {
        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<bool, StorageError> {
        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<bool, StorageError> {
        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::<i64, _>("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<Option<TerminalCommit>, 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<String> = 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<String> = 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, crate::storage::SCHEMA_VERSION);
        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<String>) =
            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"));
    }
}