1527 lines
53 KiB
Rust
1527 lines
53 KiB
Rust
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, 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<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"));
|
|
}
|
|
}
|