diff --git a/Cargo.lock b/Cargo.lock index 3ad474e..a927c60 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -112,6 +112,7 @@ dependencies = [ "iana-time-zone", "js-sys", "num-traits", + "serde", "wasm-bindgen", "windows-link", ] @@ -232,6 +233,12 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -669,6 +676,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.2" @@ -1014,6 +1027,19 @@ version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.23.41" @@ -1260,6 +1286,19 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.3.4", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thin-vec" version = "0.2.18" @@ -1311,6 +1350,7 @@ dependencies = [ "serde", "serde_json", "sha2", + "tempfile", "thiserror", "tokio", "tracing", diff --git a/Cargo.toml b/Cargo.toml index 17cdf1e..a780e93 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,6 +45,10 @@ reqwest = { version = "0.12", default-features = false, features = [ # Optional embedded SQLite checkpointer backend (`graph::checkpoint::sqlite`). # `bundled` compiles SQLite from vendored C source so no system library is # required; the dependency is pulled in only by the `sqlite` feature. +# `bundled` compiles SQLite with FTS5 already enabled, which +# `harness::session_store` relies on for its `sessions_fts` cross-session +# search table. There is no separate `fts5` cargo feature at this version — do +# not add one, it does not resolve. rusqlite = { version = "0.40", features = ["bundled"], optional = true } # Optional embedded Rhai scripting engine powering the `.ragsh` REPL session @@ -54,7 +58,9 @@ rusqlite = { version = "0.40", features = ["bundled"], optional = true } rhai = { version = "1", features = ["sync"], optional = true } # Optional builtin tool family for deterministic time/date helpers. -chrono = "0.4" +# `serde` is required by `harness::session_store`, whose records carry +# `DateTime` timestamps across the serde boundary. +chrono = { version = "0.4", features = ["serde"] } chrono-tz = { version = "0.10", optional = true } [features] @@ -75,6 +81,9 @@ tools = ["dep:chrono-tz"] [dev-dependencies] tokio = { version = "1", features = ["macros", "rt-multi-thread", "time", "test-util"] } +# Throwaway workspace roots for `harness::session_store` tests, which exercise +# the real on-disk SQLite path rather than an in-memory database. +tempfile = "3" # `.env` loading for the runnable examples. dotenvy = "0.15" # Stream combinators (`StreamExt::next`) for integration tests that drive a diff --git a/src/error.rs b/src/error.rs index c766b29..842b297 100644 --- a/src/error.rs +++ b/src/error.rs @@ -221,4 +221,27 @@ pub enum TinyAgentsError { /// (checkpoints, model wire formats, structured output, blueprints). #[error("serialization error: {0}")] Serialization(#[from] serde_json::Error), + + /// A durable-storage operation failed — opening, migrating, reading, or + /// writing a backing database for the session store and run ledger + /// ([`crate::session`]). + /// + /// Distinct from [`TinyAgentsError::Checkpoint`], which covers graph + /// checkpoint durability: a session-store failure means run *history* could + /// not be recorded or queried, while a checkpoint failure means a run + /// cannot be resumed. The payload carries the operation context and the + /// underlying driver message. + #[error("storage error: {0}")] + Storage(String), +} + +/// Converts a raw `rusqlite` failure into [`TinyAgentsError::Storage`] so the +/// session store and run ledger can use `?` on driver calls directly. Call +/// sites that have useful context to add should still map explicitly rather +/// than relying on this bare conversion. +#[cfg(feature = "sqlite")] +impl From for TinyAgentsError { + fn from(err: rusqlite::Error) -> Self { + Self::Storage(err.to_string()) + } } diff --git a/src/lib.rs b/src/lib.rs index b29e7d3..f511507 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -79,6 +79,29 @@ pub mod registry; pub mod repl; #[cfg(feature = "rlm")] pub mod rlm; +/// Durable session history and run ledger — a persistence domain in its own +/// right, not part of the agent-loop harness. Requires the `sqlite` feature. +#[cfg(feature = "sqlite")] +pub mod session; + +// --- Session: durable session history + run ledger (feature `sqlite`) --- +// Centralized here per AGENTS.md so downstream users get a predictable surface +// rather than reaching through the module path. The record/query entry points +// and the ledger's coordination types are the surface a host actually binds +// against; the rest stays reachable via `session::` for callers that want it. +#[cfg(feature = "sqlite")] +pub use session::run_ledger::{ + AgentRun, AgentRunKind, AgentRunStatus, AgentRunUpsert, AgentTeam, AgentTeamMember, + AgentTeamMemberStatus, AgentTeamStatus, AgentTeamTask, AgentTeamTaskStatus, ClaimOutcome, + CompletionOutcome, RunEvent, RunEventAppend, RunTelemetry, RunTelemetryUpsert, WorkflowRun, + WorkflowRunStatus, WorkflowRunUpsert, +}; +#[cfg(feature = "sqlite")] +pub use session::{ + SessionMessage, SessionRecord, SessionSearchParams, SessionSearchResult, SessionStatus, + SessionToolCall, get_session, list_sessions, record_message, record_session_end, + record_session_start, record_tool_call, search_sessions, +}; // --- Error: the crate-wide error type and `Result` alias --- pub use error::{Result, TinyAgentsError}; diff --git a/src/session/README.md b/src/session/README.md new file mode 100644 index 0000000..4c7f9d2 --- /dev/null +++ b/src/session/README.md @@ -0,0 +1,118 @@ +# `session` — durable session history and run ledger + +SQLite-backed history for agent sessions, and a restart-survivable ledger for +background agent/workflow execution. Requires the `sqlite` feature. + +## Why this is a top-level module + +Session history is a persistence domain in its own right, not part of the agent +loop. Nothing in `harness` reads from it, and a host can use it without running +a harness at all — indexing sessions produced elsewhere, or recovering +orchestration state at boot before any agent exists. Filing it under `harness::` +would imply a dependency that exists in neither direction. + +## How it differs from the other persistence layers + +| Layer | Question it answers | Lifetime | +| --- | --- | --- | +| `harness::store` | "what is this run working with right now?" | during a run | +| `graph::checkpoint` | "how do I resume this interrupted run?" | until resumed | +| **`session`** | "what happened, what did it cost, how did runs nest?" | indefinitely | + +Nothing resumes from this module. It is queryable history: cross-session search, +cost attribution, and orchestration recovery. + +## Layout + +Every entry point takes the workspace root and derives the path itself, so a +host chooses only where its workspace lives: + +```text +{workspace_dir}/session_db/sessions.db +``` + +## Public surface + +Re-exported from the crate root (see `src/lib.rs`); the full surface stays +reachable under `session::` and `session::run_ledger::`. + +- **Recording** — `record_session_start`, `record_message`, `record_tool_call`, + `record_session_end` +- **Querying** — `get_session`, `list_sessions`, `search_sessions`, + `list_messages`, `list_tool_calls`, `list_children` +- **Recovery** — `mark_interrupted` +- **Run ledger** — agent runs, workflow runs, teams, members, tasks, run events, + and telemetry, with the claim/completion coordination primitives +- **Connections** — `with_connection` (autocommit) and `with_transaction` + (`BEGIN IMMEDIATE`) + +## Schema + +Six tables plus one FTS5 virtual table, created on demand and idempotently: + +| Table | Holds | +| --- | --- | +| `sessions` | one row per session; lineage via `parent_session_id` | +| `session_messages` | per-message content, model, tokens, cost | +| `session_tool_calls` | tool name, input, bounded output, status, duration | +| `sessions_fts` | FTS5 index over session name, message content, tool name | +| `agent_runs` / `workflow_runs` | background execution state | +| `run_events` / `run_telemetry` | per-run event stream and rollups | +| `agent_teams` / `agent_team_members` / `agent_team_tasks` | team coordination | + +WAL journaling and `foreign_keys = ON`. FTS5 comes from `rusqlite`'s `bundled` +build — there is no separate `fts5` cargo feature at 0.40, so do not add one. + +## Operational constraints + +These are the non-obvious rules; each is pinned by a test in `test.rs`. + +**Search input is plain text, not FTS5 syntax.** `SessionSearchParams::query` is +translated to a quoted FTS5 expression before it reaches `MATCH`. Binding raw +user input made ordinary strings (`C++`, `foo-bar`, `file.rs`) fail with a +syntax or `no such column` error instead of searching. + +**Indexed content is truncated on a character boundary.** Slicing at a raw byte +offset panics on multi-byte input, and because the message insert has already +committed, the row would survive with no FTS entry — silently unsearchable. + +**Tool output is bounded** to `MAX_TOOL_OUTPUT_BYTES`, truncated on a character +boundary with a marker appended. + +**Telemetry counters are `Option` for partial updates.** The columns are +`NOT NULL DEFAULT`, and SQLite does not apply a column default to an explicitly +supplied `NULL`, so the insert path coalesces to the default while the update +path coalesces to the stored value. `excluded.*` cannot serve the update side — +it observes the already-coalesced row, so `None` would read as `0` and clobber a +stored counter. + +**Run-event sequences are allocated by the INSERT itself.** Reading +`MAX(sequence) + 1` and then inserting is a read-modify-write race; the loser +fails the primary key and the event is lost. + +**Coordination operations need `with_transaction`, not `with_connection`.** +`with_connection` is autocommit, which gives ordering but no isolation. Claim +and completion read state and then act on it, so they take the write lock up +front with `BEGIN IMMEDIATE` — racing claims serialize at `BEGIN` rather than +failing at `COMMIT` after one has already decided it won. + +**A claim is meaningful only while a task is `in_progress`.** An upsert that +moves a task off that status clears `claimed_by_member_id` and `claim_token`; +leaving them set strands the task, since a new claim sees `AlreadyClaimed`, +completion sees `NotClaimed`, and release/shutdown skip it. + +**Evidence accumulates across completion attempts,** including attempts whose +gate fails — otherwise a retry after fixing an unrelated gate would fail +`require_evidence` on evidence already submitted. + +## Layout of this module + +| File | Role | +| --- | --- | +| `mod.rs` | module docs and public surface | +| `types.rs` | serde record types | +| `store.rs` | connection/transaction helpers and schema init | +| `ops.rs` | recording and querying | +| `context.rs` | `StorageContext`, the error-context shim | +| `run_ledger/` | background run + team coordination | +| `test.rs` | module-local unit tests | diff --git a/src/session/context.rs b/src/session/context.rs new file mode 100644 index 0000000..5a91649 --- /dev/null +++ b/src/session/context.rs @@ -0,0 +1,38 @@ +//! Error-context helper for the session store's SQLite surface. +//! +//! The store and run ledger are dense with driver calls whose bare failures +//! (`no such table`, `database is locked`) say nothing about which operation +//! raised them. This trait attaches that operation context while funnelling +//! everything into [`TinyAgentsError::Storage`], so the crate keeps one error +//! type without every call site writing the same closure. +//! +//! It deliberately mirrors the shape of `anyhow::Context` — including the +//! `Option` impl for "row expected but absent" — because this module was +//! ported from a host that used `anyhow`, and matching the shape kept that +//! port mechanical and reviewable. + +use std::fmt::Display; + +use crate::error::{Result, TinyAgentsError}; + +/// Attaches operation context to a fallible storage call, producing a +/// [`TinyAgentsError::Storage`]. +pub(crate) trait StorageContext { + /// Wraps the failure as a storage error prefixed with `context`. + fn storage_context(self, context: &str) -> Result; +} + +impl StorageContext for std::result::Result { + fn storage_context(self, context: &str) -> Result { + self.map_err(|err| TinyAgentsError::Storage(format!("{context}: {err}"))) + } +} + +/// A `None` where a row was expected is a storage inconsistency, not a +/// user-facing absence — the callers using this read back a record they just +/// wrote, so `None` means the write silently failed. +impl StorageContext for Option { + fn storage_context(self, context: &str) -> Result { + self.ok_or_else(|| TinyAgentsError::Storage(context.to_string())) + } +} diff --git a/src/session/mod.rs b/src/session/mod.rs new file mode 100644 index 0000000..707275e --- /dev/null +++ b/src/session/mod.rs @@ -0,0 +1,84 @@ +//! Durable session database and run ledger. +//! +//! SQLite-backed store (WAL + FTS5) for sessions, messages, tool calls, cost +//! metadata, and parent/child lineage, plus a [`run_ledger`] for background +//! agent/workflow execution state. This is the runtime's *history* layer: what +//! ran, what it cost, what it called, and how runs nest. +//! +//! # Why this is a top-level module +//! +//! Session history is a persistence domain in its own right, not a part of the +//! agent loop. Nothing in [`crate::harness`] reads from it, and a host can use +//! it without running a harness at all — indexing sessions produced elsewhere, +//! or recovering orchestration state at boot before any agent exists. Filing it +//! under `harness::` would imply a dependency that does not exist in either +//! direction. +//! +//! # Relationship to the other persistence layers +//! +//! - [`crate::harness::store`] is namespaced key-value storage for live +//! runtime data. It is a substrate runs read and write during execution. +//! - [`crate::graph::checkpoint`] is durability for *resuming* an interrupted +//! graph run. +//! - This module is queryable history. Nothing resumes from it; it answers +//! "what happened", supports cross-session search, and lets a host recover +//! orchestration state after a restart. +//! +//! A host that keeps its own transcript files (the source of truth for +//! KV-cache resume) still wants this module for indexing and search over them. +//! +//! # Layout +//! +//! Every entry point takes the workspace root and derives the database path, +//! so a host chooses only where its workspace lives: +//! +//! ```text +//! {workspace_dir}/session_db/sessions.db +//! ``` +//! +//! # Example +//! +//! ```no_run +//! use std::path::Path; +//! use tinyagents::session::{self, SessionStatus}; +//! +//! # fn main() -> tinyagents::Result<()> { +//! let workspace = Path::new("/tmp/workspace"); +//! +//! session::record_session_start( +//! workspace, "sess-1", "researcher", "Researcher", "sess-1", +//! None, None, None, Some("gpt-5"), None, +//! )?; +//! session::record_message( +//! workspace, "sess-1", "user", "summarize the repo", None, None, None, None, +//! )?; +//! session::record_session_end( +//! workspace, "sess-1", SessionStatus::Completed, 1, 120, 340, 0, 0.004, +//! )?; +//! # Ok(()) +//! # } +//! ``` +//! +//! Requires the `sqlite` feature. +//! +//! See [`README.md`](./README.md) for the schema, the FTS behaviour, and the +//! coordination guarantees. + +mod context; +mod ops; +pub mod run_ledger; +mod store; +pub mod types; + +pub use ops::{ + get_session, list_children, list_messages, list_sessions, list_tool_calls, mark_interrupted, + record_message, record_session_end, record_session_start, record_tool_call, search_sessions, +}; +pub use store::{db_path, with_connection, with_transaction}; +pub use types::{ + SessionMessage, SessionRecord, SessionSearchParams, SessionSearchResult, SessionStatus, + SessionToolCall, +}; + +#[cfg(test)] +mod test; diff --git a/src/session/ops.rs b/src/session/ops.rs new file mode 100644 index 0000000..6299cce --- /dev/null +++ b/src/session/ops.rs @@ -0,0 +1,661 @@ +use std::path::Path; + +use chrono::{DateTime, Utc}; +use rusqlite::{Connection, params}; + +use crate::error::{Result, TinyAgentsError}; + +use super::context::StorageContext; +use super::store::with_connection; +use super::types::{ + SessionMessage, SessionRecord, SessionSearchParams, SessionSearchResult, SessionStatus, + SessionToolCall, +}; + +pub(super) const MAX_TOOL_OUTPUT_BYTES: usize = 32 * 1024; + +// A record-shaped signature: each argument is one persisted column. Grouping +// them into a struct is worth doing, but is an API change rather than part of +// this move — tracked separately. +#[allow(clippy::too_many_arguments)] +pub fn record_session_start( + workspace_dir: &Path, + id: &str, + agent_definition_id: &str, + agent_definition_name: &str, + session_key: &str, + parent_session_id: Option<&str>, + thread_id: Option<&str>, + source_channel: Option<&str>, + model: Option<&str>, + transcript_path: Option<&str>, +) -> Result { + let now = Utc::now(); + tracing::debug!( + "[session_db] record_session_start id={id} agent={agent_definition_id} \ + parent={} thread={} channel={}", + parent_session_id.unwrap_or("-"), + thread_id.unwrap_or("-"), + source_channel.unwrap_or("-"), + ); + + with_connection(workspace_dir, |conn| { + conn.execute( + "INSERT INTO sessions ( + id, agent_definition_id, agent_definition_name, session_key, + parent_session_id, thread_id, source_channel, status, model, + transcript_path, started_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, 'running', ?8, ?9, ?10)", + params![ + id, + agent_definition_id, + agent_definition_name, + session_key, + parent_session_id, + thread_id, + source_channel, + model, + transcript_path, + now.to_rfc3339(), + ], + ) + .storage_context("failed to insert session")?; + + index_fts_session(conn, id, agent_definition_name)?; + Ok(()) + })?; + + get_session(workspace_dir, id) +} + +// A record-shaped signature: each argument is one persisted column. Grouping +// them into a struct is worth doing, but is an API change rather than part of +// this move — tracked separately. +#[allow(clippy::too_many_arguments)] +pub fn record_session_end( + workspace_dir: &Path, + id: &str, + status: SessionStatus, + turn_count: u32, + input_tokens: u64, + output_tokens: u64, + cached_input_tokens: u64, + cost_usd: f64, +) -> Result { + let now = Utc::now(); + tracing::debug!( + "[session_db] record_session_end id={id} status={} turns={turn_count} \ + tokens_in={input_tokens} tokens_out={output_tokens} cost=${cost_usd:.6}", + status.as_str(), + ); + + with_connection(workspace_dir, |conn| { + conn.execute( + "UPDATE sessions SET + status = ?1, turn_count = ?2, input_tokens = ?3, + output_tokens = ?4, cached_input_tokens = ?5, + cost_usd = ?6, ended_at = ?7 + WHERE id = ?8", + params![ + status.as_str(), + turn_count, + input_tokens as i64, + output_tokens as i64, + cached_input_tokens as i64, + cost_usd, + now.to_rfc3339(), + id, + ], + ) + .storage_context("failed to update session end")?; + Ok(()) + })?; + + get_session(workspace_dir, id) +} + +// A record-shaped signature: each argument is one persisted column. Grouping +// them into a struct is worth doing, but is an API change rather than part of +// this move — tracked separately. +#[allow(clippy::too_many_arguments)] +pub fn record_message( + workspace_dir: &Path, + session_id: &str, + role: &str, + content: &str, + model: Option<&str>, + input_tokens: Option, + output_tokens: Option, + cost_usd: Option, +) -> Result { + let now = Utc::now(); + tracing::trace!( + "[session_db] record_message session={session_id} role={role} len={}", + content.len() + ); + + with_connection(workspace_dir, |conn| { + conn.execute( + "INSERT INTO session_messages ( + session_id, role, content, model, + input_tokens, output_tokens, cost_usd, created_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + params![ + session_id, + role, + content, + model, + input_tokens.map(|v| v as i64), + output_tokens.map(|v| v as i64), + cost_usd, + now.to_rfc3339(), + ], + ) + .storage_context("failed to insert session message")?; + + let msg_id = conn.last_insert_rowid(); + + index_fts_content(conn, session_id, content)?; + + Ok(msg_id) + }) +} + +// A record-shaped signature: each argument is one persisted column. Grouping +// them into a struct is worth doing, but is an API change rather than part of +// this move — tracked separately. +#[allow(clippy::too_many_arguments)] +pub fn record_tool_call( + workspace_dir: &Path, + session_id: &str, + message_id: Option, + tool_name: &str, + tool_input: Option<&str>, + tool_output: Option<&str>, + status: &str, + duration_ms: Option, +) -> Result { + let now = Utc::now(); + tracing::trace!( + "[session_db] record_tool_call session={session_id} tool={tool_name} status={status}" + ); + + let bounded_output = tool_output.map(|o| { + if o.len() <= MAX_TOOL_OUTPUT_BYTES { + o.to_string() + } else { + let mut cutoff = MAX_TOOL_OUTPUT_BYTES; + while cutoff > 0 && !o.is_char_boundary(cutoff) { + cutoff -= 1; + } + let mut truncated = o[..cutoff].to_string(); + truncated.push_str("\n...[truncated]"); + truncated + } + }); + + with_connection(workspace_dir, |conn| { + conn.execute( + "INSERT INTO session_tool_calls ( + session_id, message_id, tool_name, tool_input, + tool_output, status, duration_ms, created_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)", + params![ + session_id, + message_id, + tool_name, + tool_input, + bounded_output, + status, + duration_ms, + now.to_rfc3339(), + ], + ) + .storage_context("failed to insert tool call")?; + + // Capture the row id BEFORE indexing. `index_fts_tool` inserts into the + // `sessions_fts` virtual table, which moves `last_insert_rowid()` to + // that row — so reading it afterwards handed callers an FTS rowid for a + // tool call that does not exist. `record_message` already ordered these + // correctly; this path did not. + let tool_call_id = conn.last_insert_rowid(); + + index_fts_tool(conn, session_id, tool_name)?; + + Ok(tool_call_id) + }) +} + +pub fn get_session(workspace_dir: &Path, id: &str) -> Result { + with_connection(workspace_dir, |conn| { + let mut stmt = conn.prepare( + "SELECT id, agent_definition_id, agent_definition_name, session_key, + parent_session_id, thread_id, source_channel, status, model, + turn_count, input_tokens, output_tokens, cached_input_tokens, + cost_usd, transcript_path, started_at, ended_at + FROM sessions WHERE id = ?1", + )?; + + let mut rows = stmt.query(params![id])?; + if let Some(row) = rows.next()? { + map_session_row(row).map_err(Into::into) + } else { + Err(TinyAgentsError::Storage(format!( + "session '{id}' not found" + ))) + } + }) +} + +pub fn list_sessions( + workspace_dir: &Path, + limit: Option, + offset: Option, + status: Option<&str>, + parent_id: Option<&str>, +) -> Result { + tracing::debug!( + "[session_db] list_sessions limit={} offset={} status={} parent={}", + limit.unwrap_or(50), + offset.unwrap_or(0), + status.unwrap_or("-"), + parent_id.unwrap_or("-"), + ); + + with_connection(workspace_dir, |conn| { + let mut where_clauses: Vec = Vec::new(); + let mut param_values: Vec> = Vec::new(); + + if let Some(s) = status { + param_values.push(Box::new(s.to_string())); + where_clauses.push(format!("status = ?{}", param_values.len())); + } + if let Some(p) = parent_id { + param_values.push(Box::new(p.to_string())); + where_clauses.push(format!("parent_session_id = ?{}", param_values.len())); + } + + let where_sql = if where_clauses.is_empty() { + String::new() + } else { + format!("WHERE {}", where_clauses.join(" AND ")) + }; + + let lim = limit.unwrap_or(50).min(500) as i64; + let off = offset.unwrap_or(0) as i64; + + let count_sql = format!("SELECT COUNT(*) FROM sessions {where_sql}"); + let total: u64 = { + let mut stmt = conn.prepare(&count_sql)?; + let params_ref: Vec<&dyn rusqlite::types::ToSql> = + param_values.iter().map(|b| b.as_ref()).collect(); + stmt.query_row(params_ref.as_slice(), |r| r.get::<_, i64>(0))? as u64 + }; + + param_values.push(Box::new(lim)); + let lim_idx = param_values.len(); + param_values.push(Box::new(off)); + let off_idx = param_values.len(); + + let query_sql = format!( + "SELECT id, agent_definition_id, agent_definition_name, session_key, + parent_session_id, thread_id, source_channel, status, model, + turn_count, input_tokens, output_tokens, cached_input_tokens, + cost_usd, transcript_path, started_at, ended_at + FROM sessions {where_sql} + ORDER BY started_at DESC + LIMIT ?{lim_idx} OFFSET ?{off_idx}", + ); + + let mut stmt = conn.prepare(&query_sql)?; + let params_ref: Vec<&dyn rusqlite::types::ToSql> = + param_values.iter().map(|b| b.as_ref()).collect(); + let rows = stmt.query_map(params_ref.as_slice(), map_session_row)?; + + let mut sessions = Vec::new(); + for row in rows { + sessions.push(row?); + } + + Ok(SessionSearchResult { sessions, total }) + }) +} + +pub fn search_sessions( + workspace_dir: &Path, + params: &SessionSearchParams, +) -> Result { + tracing::debug!( + "[session_db] search_sessions query={} agent={} tool={} channel={} thread={}", + params.query.as_deref().unwrap_or("-"), + params.agent_id.as_deref().unwrap_or("-"), + params.tool_name.as_deref().unwrap_or("-"), + params.source_channel.as_deref().unwrap_or("-"), + params.thread_id.as_deref().unwrap_or("-"), + ); + + with_connection(workspace_dir, |conn| search_sessions_inner(conn, params)) +} + +pub(super) fn search_sessions_inner( + conn: &Connection, + params: &SessionSearchParams, +) -> Result { + let lim = params.limit.unwrap_or(50).min(500) as i64; + let off = params.offset.unwrap_or(0) as i64; + + let mut where_clauses: Vec = Vec::new(); + let mut param_values: Vec> = Vec::new(); + + if let Some(q) = params.query.as_ref().filter(|q| !q.trim().is_empty()) { + param_values.push(Box::new(fts_match_query(q))); + where_clauses.push(format!( + "s.id IN (SELECT session_id FROM sessions_fts WHERE sessions_fts MATCH ?{})", + param_values.len() + )); + } + + if let Some(ref agent) = params.agent_id { + param_values.push(Box::new(agent.clone())); + where_clauses.push(format!("s.agent_definition_id = ?{}", param_values.len())); + } + + if let Some(ref tool) = params.tool_name { + param_values.push(Box::new(tool.clone())); + where_clauses.push(format!( + "s.id IN (SELECT DISTINCT session_id FROM session_tool_calls WHERE tool_name = ?{})", + param_values.len() + )); + } + + if let Some(ref channel) = params.source_channel { + param_values.push(Box::new(channel.clone())); + where_clauses.push(format!("s.source_channel = ?{}", param_values.len())); + } + + if let Some(ref parent) = params.parent_session_id { + param_values.push(Box::new(parent.clone())); + where_clauses.push(format!("s.parent_session_id = ?{}", param_values.len())); + } + + if let Some(ref status) = params.status { + param_values.push(Box::new(status.clone())); + where_clauses.push(format!("s.status = ?{}", param_values.len())); + } + + if let Some(ref tid) = params.thread_id { + param_values.push(Box::new(tid.clone())); + where_clauses.push(format!("s.thread_id = ?{}", param_values.len())); + } + + let where_sql = if where_clauses.is_empty() { + String::new() + } else { + format!("WHERE {}", where_clauses.join(" AND ")) + }; + + let count_sql = format!("SELECT COUNT(*) FROM sessions s {where_sql}"); + let total: u64 = { + let mut stmt = conn.prepare(&count_sql)?; + let params_ref: Vec<&dyn rusqlite::types::ToSql> = + param_values.iter().map(|b| b.as_ref()).collect(); + stmt.query_row(params_ref.as_slice(), |r| r.get::<_, i64>(0))? as u64 + }; + + param_values.push(Box::new(lim)); + let lim_idx = param_values.len(); + param_values.push(Box::new(off)); + let off_idx = param_values.len(); + + let query = format!( + "SELECT s.id, s.agent_definition_id, s.agent_definition_name, s.session_key, + s.parent_session_id, s.thread_id, s.source_channel, s.status, s.model, + s.turn_count, s.input_tokens, s.output_tokens, s.cached_input_tokens, + s.cost_usd, s.transcript_path, s.started_at, s.ended_at + FROM sessions s {where_sql} + ORDER BY s.started_at DESC + LIMIT ?{lim_idx} OFFSET ?{off_idx}", + ); + + let mut stmt = conn.prepare(&query)?; + let params_ref: Vec<&dyn rusqlite::types::ToSql> = + param_values.iter().map(|b| b.as_ref()).collect(); + let rows = stmt.query_map(params_ref.as_slice(), map_session_row)?; + + let mut sessions = Vec::new(); + for row in rows { + sessions.push(row?); + } + + Ok(SessionSearchResult { sessions, total }) +} + +pub fn list_messages( + workspace_dir: &Path, + session_id: &str, + limit: Option, +) -> Result> { + with_connection(workspace_dir, |conn| { + let lim = limit.unwrap_or(200).min(1000) as i64; + let mut stmt = conn.prepare( + "SELECT id, session_id, role, content, model, + input_tokens, output_tokens, cost_usd, created_at + FROM session_messages + WHERE session_id = ?1 + ORDER BY id ASC + LIMIT ?2", + )?; + + let rows = stmt.query_map(params![session_id, lim], |row| { + Ok(SessionMessage { + id: row.get(0)?, + session_id: row.get(1)?, + role: row.get(2)?, + content: row.get(3)?, + model: row.get(4)?, + input_tokens: row.get::<_, Option>(5)?.map(|v| v as u64), + output_tokens: row.get::<_, Option>(6)?.map(|v| v as u64), + cost_usd: row.get(7)?, + created_at: parse_rfc3339(&row.get::<_, String>(8)?) + .map_err(sql_conversion_error)?, + }) + })?; + + let mut messages = Vec::new(); + for row in rows { + messages.push(row?); + } + Ok(messages) + }) +} + +pub fn list_tool_calls( + workspace_dir: &Path, + session_id: &str, + limit: Option, +) -> Result> { + with_connection(workspace_dir, |conn| { + let lim = limit.unwrap_or(200).min(1000) as i64; + let mut stmt = conn.prepare( + "SELECT id, session_id, message_id, tool_name, tool_input, + tool_output, status, duration_ms, created_at + FROM session_tool_calls + WHERE session_id = ?1 + ORDER BY id ASC + LIMIT ?2", + )?; + + let rows = stmt.query_map(params![session_id, lim], |row| { + Ok(SessionToolCall { + id: row.get(0)?, + session_id: row.get(1)?, + message_id: row.get(2)?, + tool_name: row.get(3)?, + tool_input: row.get(4)?, + tool_output: row.get(5)?, + status: row.get(6)?, + duration_ms: row.get(7)?, + created_at: parse_rfc3339(&row.get::<_, String>(8)?) + .map_err(sql_conversion_error)?, + }) + })?; + + let mut tool_calls = Vec::new(); + for row in rows { + tool_calls.push(row?); + } + Ok(tool_calls) + }) +} + +pub fn list_children(workspace_dir: &Path, session_id: &str) -> Result> { + with_connection(workspace_dir, |conn| { + let mut stmt = conn.prepare( + "SELECT id, agent_definition_id, agent_definition_name, session_key, + parent_session_id, thread_id, source_channel, status, model, + turn_count, input_tokens, output_tokens, cached_input_tokens, + cost_usd, transcript_path, started_at, ended_at + FROM sessions + WHERE parent_session_id = ?1 + ORDER BY started_at ASC", + )?; + + let rows = stmt.query_map(params![session_id], map_session_row)?; + let mut children = Vec::new(); + for row in rows { + children.push(row?); + } + Ok(children) + }) +} + +pub fn mark_interrupted(workspace_dir: &Path) -> Result { + tracing::debug!("[session_db] mark_interrupted — marking all running sessions as interrupted"); + with_connection(workspace_dir, |conn| { + let now = Utc::now(); + let changed = conn.execute( + "UPDATE sessions SET status = 'interrupted', ended_at = ?1 + WHERE status = 'running'", + params![now.to_rfc3339()], + )?; + if changed > 0 { + tracing::info!("[session_db] marked {changed} running session(s) as interrupted"); + } + Ok(changed) + }) +} + +pub(super) fn index_fts_session( + conn: &Connection, + session_id: &str, + agent_name: &str, +) -> Result<()> { + conn.execute( + "INSERT INTO sessions_fts (session_id, agent_definition_name, content, tool_name) + VALUES (?1, ?2, '', '')", + params![session_id, agent_name], + ) + .storage_context("failed to index session in FTS")?; + Ok(()) +} + +/// Renders a user's plain-text search string as an FTS5 MATCH expression. +/// +/// [`SessionSearchParams::query`] is documented as plain text, not as raw FTS5 +/// syntax, but binding it straight to `MATCH` hands it to the FTS5 parser. +/// Ordinary input then fails rather than searching: `C++`, `foo-bar`, +/// `file.rs`, or a stray `"` each produce a syntax or `no such column` error +/// instead of results. +/// +/// Each whitespace-separated term is emitted as a double-quoted FTS5 string +/// literal (with `"` escaped by doubling, per the FTS5 grammar), so every +/// character inside it is treated as data. Terms are joined by `AND`, matching +/// the implicit conjunction a user expects from a search box. +pub(super) fn fts_match_query(raw: &str) -> String { + raw.split_whitespace() + .map(|term| format!("\"{}\"", term.replace('"', "\"\""))) + .collect::>() + .join(" AND ") +} + +/// Longest FTS snippet indexed per message, in bytes. +pub(super) const MAX_FTS_SNIPPET_BYTES: usize = 2000; + +pub(super) fn index_fts_content(conn: &Connection, session_id: &str, content: &str) -> Result<()> { + // Slice on a character boundary, not a byte offset. `&content[..2000]` + // panics whenever byte 2000 lands inside a multi-byte character, which any + // ordinary long non-ASCII message can do. The panic is worse than it looks: + // the message INSERT has already autocommitted by this point, so the row + // survives with no FTS entry and is silently unsearchable forever after. + // Mirrors the truncation already done in `record_tool_call`. + let snippet = if content.len() > MAX_FTS_SNIPPET_BYTES { + let mut cutoff = MAX_FTS_SNIPPET_BYTES; + while cutoff > 0 && !content.is_char_boundary(cutoff) { + cutoff -= 1; + } + &content[..cutoff] + } else { + content + }; + conn.execute( + "INSERT INTO sessions_fts (session_id, agent_definition_name, content, tool_name) + VALUES (?1, '', ?2, '')", + params![session_id, snippet], + ) + .storage_context("failed to index content in FTS")?; + Ok(()) +} + +pub(super) fn index_fts_tool(conn: &Connection, session_id: &str, tool_name: &str) -> Result<()> { + conn.execute( + "INSERT INTO sessions_fts (session_id, agent_definition_name, content, tool_name) + VALUES (?1, '', '', ?2)", + params![session_id, tool_name], + ) + .storage_context("failed to index tool call in FTS")?; + Ok(()) +} + +pub(super) fn map_session_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { + let started_at_raw: String = row.get(15)?; + let ended_at_raw: Option = row.get(16)?; + + Ok(SessionRecord { + id: row.get(0)?, + agent_definition_id: row.get(1)?, + agent_definition_name: row.get(2)?, + session_key: row.get(3)?, + parent_session_id: row.get(4)?, + thread_id: row.get(5)?, + source_channel: row.get(6)?, + status: SessionStatus::parse(&row.get::<_, String>(7)?), + model: row.get(8)?, + turn_count: row.get::<_, i64>(9)? as u32, + input_tokens: row.get::<_, i64>(10)? as u64, + output_tokens: row.get::<_, i64>(11)? as u64, + cached_input_tokens: row.get::<_, i64>(12)? as u64, + cost_usd: row.get(13)?, + transcript_path: row.get(14)?, + started_at: parse_rfc3339(&started_at_raw).map_err(sql_conversion_error)?, + ended_at: match ended_at_raw { + Some(raw) => Some(parse_rfc3339(&raw).map_err(sql_conversion_error)?), + None => None, + }, + }) +} + +pub(super) fn parse_rfc3339(raw: &str) -> Result> { + let parsed = DateTime::parse_from_rfc3339(raw) + .storage_context(&format!("invalid RFC3339 timestamp in session DB: {raw}"))?; + Ok(parsed.with_timezone(&Utc)) +} + +/// Bridges a timestamp-parse failure back into `rusqlite`'s error type. +/// +/// Row mappers must return `rusqlite::Result`, so a malformed stored timestamp +/// cannot surface as [`TinyAgentsError`] directly from inside `query_map`; it +/// is boxed here and unwrapped by the caller's `?` into +/// [`TinyAgentsError::Storage`] via the crate's `From`. +pub(super) fn sql_conversion_error(err: TinyAgentsError) -> rusqlite::Error { + rusqlite::Error::ToSqlConversionFailure(Box::new(err)) +} diff --git a/src/session/run_ledger/mod.rs b/src/session/run_ledger/mod.rs new file mode 100644 index 0000000..9db0d76 --- /dev/null +++ b/src/session/run_ledger/mod.rs @@ -0,0 +1,36 @@ +//! Durable run ledger for agent and workflow execution state. +//! +//! Extends [`super`] with a queryable, restart-survivable ledger for background +//! agent and workflow runs. Conversation transcripts remain in the session +//! store; this ledger holds compact run metadata, child lineage, events, +//! telemetry, and checkpoint references — enough for a host to reconstruct +//! what was in flight after a crash and resume or interrupt it. +//! +//! Shares the session database and connection helper with [`super::store`], so +//! a run and the session that produced it are queryable together. + +pub mod ops; +pub mod store; +pub mod types; + +pub use ops::{ + append_run_event, claim_agent_team_task, complete_agent_team_task, get_agent_run, + get_agent_team, get_agent_team_member, get_agent_team_task, get_workflow_run, + interrupt_orphaned_agent_runs, list_agent_runs, list_agent_team_members, list_agent_team_tasks, + list_agent_teams, list_recent_run_events, list_workflow_runs, mark_agent_team_member_idle, + mark_agent_team_member_running, release_agent_team_task, shutdown_agent_team_member, + transition_agent_run_status, upsert_agent_run, upsert_agent_team, upsert_agent_team_member, + upsert_agent_team_task, upsert_run_telemetry, upsert_workflow_run, +}; +pub use types::{ + AgentRun, AgentRunKind, AgentRunListRequest, AgentRunListResponse, AgentRunStatus, + AgentRunUpsert, AgentTeam, AgentTeamListRequest, AgentTeamListResponse, AgentTeamMember, + AgentTeamMemberStatus, AgentTeamMemberUpsert, AgentTeamStatus, AgentTeamTask, + AgentTeamTaskStatus, AgentTeamTaskUpsert, AgentTeamUpsert, ClaimOutcome, CompletionOutcome, + RunEvent, RunEventAppend, RunEventListRequest, RunEventListResponse, RunTelemetry, + RunTelemetryUpsert, WorkflowRun, WorkflowRunListRequest, WorkflowRunListResponse, + WorkflowRunStatus, WorkflowRunUpsert, +}; + +#[cfg(test)] +mod test; diff --git a/src/session/run_ledger/ops.rs b/src/session/run_ledger/ops.rs new file mode 100644 index 0000000..5db3090 --- /dev/null +++ b/src/session/run_ledger/ops.rs @@ -0,0 +1,1553 @@ +use std::path::Path; + +use chrono::{DateTime, Utc}; +use rusqlite::{Connection, OptionalExtension, params}; +use serde_json::{Value, json}; + +use crate::error::Result; + +use super::super::context::StorageContext; +use super::store::init_run_ledger_schema; +use super::types::{ + AgentRun, AgentRunListRequest, AgentRunListResponse, AgentRunStatus, AgentRunUpsert, AgentTeam, + AgentTeamListRequest, AgentTeamListResponse, AgentTeamMember, AgentTeamMemberStatus, + AgentTeamMemberUpsert, AgentTeamStatus, AgentTeamTask, AgentTeamTaskStatus, + AgentTeamTaskUpsert, AgentTeamUpsert, ClaimOutcome, CompletionOutcome, RunEvent, + RunEventAppend, RunEventListRequest, RunEventListResponse, RunTelemetry, RunTelemetryUpsert, + WorkflowRun, WorkflowRunListRequest, WorkflowRunListResponse, WorkflowRunUpsert, +}; + +const LOG_PREFIX: &str = "[session_db:run_ledger]"; + +pub fn upsert_agent_run(workspace_dir: &Path, upsert: AgentRunUpsert) -> Result { + let now = Utc::now(); + let started_at = upsert.started_at.unwrap_or(now); + let updated_at = now; + let metadata_json = + serde_json::to_string(&upsert.metadata).storage_context("serialize agent run metadata")?; + let checkpoint_json = upsert + .checkpoint + .as_ref() + .map(serde_json::to_string) + .transpose() + .storage_context("serialize agent run checkpoint")?; + + tracing::debug!( + "{LOG_PREFIX} upsert_agent_run id={} kind={} status={} parent={} thread={}", + upsert.id, + upsert.kind.as_str(), + upsert.status.as_str(), + upsert.parent_run_id.as_deref().unwrap_or("-"), + upsert.parent_thread_id.as_deref().unwrap_or("-") + ); + + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + conn.execute( + "INSERT INTO agent_runs ( + id, kind, parent_run_id, parent_thread_id, agent_id, status, + prompt_ref, worker_thread_id, task_board_id, task_card_id, + checkpoint_path, checkpoint_json, summary, error, metadata_json, + started_at, updated_at, completed_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18) + ON CONFLICT(id) DO UPDATE SET + kind = CASE + WHEN agent_runs.kind = 'worker_thread' AND excluded.kind = 'subagent' THEN agent_runs.kind + ELSE excluded.kind + END, + parent_run_id = COALESCE(excluded.parent_run_id, agent_runs.parent_run_id), + parent_thread_id = COALESCE(excluded.parent_thread_id, agent_runs.parent_thread_id), + agent_id = COALESCE(excluded.agent_id, agent_runs.agent_id), + status = excluded.status, + prompt_ref = COALESCE(excluded.prompt_ref, agent_runs.prompt_ref), + worker_thread_id = COALESCE(excluded.worker_thread_id, agent_runs.worker_thread_id), + task_board_id = COALESCE(excluded.task_board_id, agent_runs.task_board_id), + task_card_id = COALESCE(excluded.task_card_id, agent_runs.task_card_id), + checkpoint_path = COALESCE(excluded.checkpoint_path, agent_runs.checkpoint_path), + checkpoint_json = COALESCE(excluded.checkpoint_json, agent_runs.checkpoint_json), + summary = COALESCE(excluded.summary, agent_runs.summary), + error = COALESCE(excluded.error, agent_runs.error), + metadata_json = CASE + WHEN excluded.metadata_json = '{}' THEN agent_runs.metadata_json + ELSE excluded.metadata_json + END, + updated_at = excluded.updated_at, + completed_at = COALESCE(excluded.completed_at, agent_runs.completed_at)", + params![ + upsert.id, + upsert.kind.as_str(), + upsert.parent_run_id, + upsert.parent_thread_id, + upsert.agent_id, + upsert.status.as_str(), + upsert.prompt_ref, + upsert.worker_thread_id, + upsert.task_board_id, + upsert.task_card_id, + upsert.checkpoint_path, + checkpoint_json, + upsert.summary, + upsert.error, + metadata_json, + started_at.to_rfc3339(), + updated_at.to_rfc3339(), + upsert.completed_at.map(|dt| dt.to_rfc3339()), + ], + ) + .storage_context("upsert agent run")?; + Ok(()) + })?; + + get_agent_run(workspace_dir, &upsert.id)?.storage_context("agent run missing after upsert") +} + +pub fn upsert_workflow_run(workspace_dir: &Path, upsert: WorkflowRunUpsert) -> Result { + let now = Utc::now(); + let started_at = upsert.started_at.unwrap_or(now); + let input_json = + serde_json::to_string(&upsert.input).storage_context("serialize workflow input")?; + let phase_states_json = serde_json::to_string(&upsert.phase_states) + .storage_context("serialize workflow phase states")?; + let child_run_ids_json = + serde_json::to_string(&upsert.child_run_ids).storage_context("serialize child run ids")?; + + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + conn.execute( + "INSERT INTO workflow_runs ( + id, definition_id, parent_thread_id, input_json, phase_states_json, + child_run_ids_json, status, summary, started_at, updated_at, completed_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11) + ON CONFLICT(id) DO UPDATE SET + definition_id = excluded.definition_id, + parent_thread_id = COALESCE(excluded.parent_thread_id, workflow_runs.parent_thread_id), + input_json = excluded.input_json, + phase_states_json = excluded.phase_states_json, + child_run_ids_json = excluded.child_run_ids_json, + status = excluded.status, + summary = COALESCE(excluded.summary, workflow_runs.summary), + updated_at = excluded.updated_at, + completed_at = COALESCE(excluded.completed_at, workflow_runs.completed_at)", + params![ + upsert.id, + upsert.definition_id, + upsert.parent_thread_id, + input_json, + phase_states_json, + child_run_ids_json, + upsert.status.as_str(), + upsert.summary, + started_at.to_rfc3339(), + now.to_rfc3339(), + upsert.completed_at.map(|dt| dt.to_rfc3339()), + ], + ) + .storage_context("upsert workflow run")?; + Ok(()) + })?; + + get_workflow_run(workspace_dir, &upsert.id)? + .storage_context("workflow run missing after upsert") +} + +pub fn append_run_event(workspace_dir: &Path, event: RunEventAppend) -> Result { + let now = Utc::now(); + let payload_json = + serde_json::to_string(&event.payload).storage_context("serialize run event")?; + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + // Allocate and insert the sequence in ONE statement. Reading + // `MAX(sequence) + 1` and then inserting is a read-modify-write race: + // two connections appending for the same run can read the same next + // value, and the loser fails the `(run_id, sequence)` primary key — + // silently dropping a real run event unless every caller implements an + // undocumented retry. The sub-select is evaluated inside the same + // statement, so SQLite's write lock serializes the whole allocation. + let next_sequence: i64 = conn + .query_row( + "INSERT INTO run_events (run_id, sequence, event_type, payload_json, timestamp) + VALUES ( + ?1, + (SELECT COALESCE(MAX(sequence), 0) + 1 FROM run_events WHERE run_id = ?1), + ?2, ?3, ?4 + ) + RETURNING sequence", + params![ + event.run_id, + event.event_type, + payload_json, + now.to_rfc3339(), + ], + |row| row.get(0), + ) + .storage_context("append run event")?; + Ok(RunEvent { + run_id: event.run_id, + sequence: next_sequence as u64, + event_type: event.event_type, + payload: serde_json::from_str(&payload_json).unwrap_or_else(|_| json!({})), + timestamp: now, + }) + }) +} + +pub fn upsert_run_telemetry( + workspace_dir: &Path, + upsert: RunTelemetryUpsert, +) -> Result { + let now = Utc::now(); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + conn.execute( + // The counters are `Option` so a caller can update one field without + // clobbering the rest, but the columns are `NOT NULL DEFAULT`, and + // SQLite does NOT apply a column default to an explicitly supplied + // NULL. Binding the raw `None` therefore made every partial upsert + // (say, recording only `model` or only `error`) fail a NOT NULL + // constraint on first write. The insert side coalesces to the + // column default; the update side re-reads the SAME parameter and + // coalesces to the stored value, which keeps per-field optionality. + // `excluded.*` cannot serve the update side here — it observes the + // already-coalesced insert row, so a `None` would read as 0 and + // overwrite the stored counter. + "INSERT INTO run_telemetry ( + run_id, input_tokens, output_tokens, cached_input_tokens, cost_usd, + elapsed_ms, tool_count, model, provider, error, updated_at + ) VALUES ( + ?1, + COALESCE(?2, 0), COALESCE(?3, 0), COALESCE(?4, 0), COALESCE(?5, 0.0), + ?6, COALESCE(?7, 0), ?8, ?9, ?10, ?11 + ) + ON CONFLICT(run_id) DO UPDATE SET + input_tokens = COALESCE(?2, run_telemetry.input_tokens), + output_tokens = COALESCE(?3, run_telemetry.output_tokens), + cached_input_tokens = COALESCE(?4, run_telemetry.cached_input_tokens), + cost_usd = COALESCE(?5, run_telemetry.cost_usd), + elapsed_ms = COALESCE(?6, run_telemetry.elapsed_ms), + tool_count = COALESCE(?7, run_telemetry.tool_count), + model = COALESCE(?8, run_telemetry.model), + provider = COALESCE(?9, run_telemetry.provider), + error = COALESCE(?10, run_telemetry.error), + updated_at = ?11", + params![ + upsert.run_id, + upsert.input_tokens.map(|v| v as i64), + upsert.output_tokens.map(|v| v as i64), + upsert.cached_input_tokens.map(|v| v as i64), + upsert.cost_usd, + upsert.elapsed_ms.map(|v| v as i64), + upsert.tool_count.map(|v| v as i64), + upsert.model, + upsert.provider, + upsert.error, + now.to_rfc3339(), + ], + ) + .storage_context("upsert run telemetry")?; + get_run_telemetry_inner(conn, &upsert.run_id) + }) +} + +pub fn get_agent_run(workspace_dir: &Path, id: &str) -> Result> { + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + get_agent_run_inner(conn, id) + }) +} + +/// Apply a durable status transition to a single agent run. +/// +/// Unlike [`upsert_agent_run`] — whose `ON CONFLICT` clause `COALESCE`s the +/// `error` and `completed_at` columns and can therefore only ever *set* them — +/// this is a direct `UPDATE` that can both set and *clear* both columns. That +/// is required by control verbs such as "retry", which moves a failed run back +/// to `pending` and must drop the stale failure reason and completion time. +/// +/// `status` is always written. `error` and `completed_at` are written verbatim, +/// so passing `None` clears the column. `updated_at` is bumped to now. Returns +/// the freshly-read run, or `None` when no row matched `id` (e.g. it was +/// deleted between a prior read and this write). +pub fn transition_agent_run_status( + workspace_dir: &Path, + id: &str, + status: AgentRunStatus, + error: Option<&str>, + completed_at: Option>, +) -> Result> { + let now = Utc::now(); + tracing::debug!( + "{LOG_PREFIX} transition_agent_run_status id={id} status={} has_error={} has_completed_at={}", + status.as_str(), + error.is_some(), + completed_at.is_some() + ); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let rows_affected = conn + .execute( + "UPDATE agent_runs + SET status = ?1, error = ?2, completed_at = ?3, updated_at = ?4 + WHERE id = ?5", + params![ + status.as_str(), + error, + completed_at.map(|dt| dt.to_rfc3339()), + now.to_rfc3339(), + id, + ], + ) + .storage_context("transition agent run status")?; + if rows_affected == 0 { + tracing::debug!("{LOG_PREFIX} transition_agent_run_status.miss id={id}"); + return Ok(None); + } + get_agent_run_inner(conn, id) + }) +} + +/// Settle non-terminal `agent_runs` rows left behind by a previous process. +/// +/// A freshly-booted core has no in-flight subagents — any detached run task +/// from a prior process is gone with that process. So a row still marked +/// `running` (or `pending`) at startup is, by definition, orphaned: its driver +/// died without firing a terminal `DomainEvent::Subagent{Completed,Failed}`, so +/// the [`register_run_ledger_finalize_subscriber`] never settled it. Without +/// this sweep those rows render as perpetual "running" timeline entries on every +/// thread reopen. +/// +/// We stamp them `interrupted` (outcome unknown — mirrors the turn-state +/// `mark_all_interrupted` recovery) and set `completed_at`. `awaiting_user` / +/// `paused` are intentionally left untouched: those are resumable states a user +/// may still continue. +/// +/// [`register_run_ledger_finalize_subscriber`]: crate::openhuman::agent::orchestration::run_ledger_finalize::register_run_ledger_finalize_subscriber +pub fn interrupt_orphaned_agent_runs(workspace_dir: &Path) -> Result { + let now = Utc::now(); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let rows_affected = conn + .execute( + "UPDATE agent_runs + SET status = ?1, completed_at = COALESCE(completed_at, ?2), updated_at = ?2 + WHERE status IN ('running', 'pending')", + params![AgentRunStatus::Interrupted.as_str(), now.to_rfc3339()], + ) + .storage_context("interrupt orphaned agent runs")?; + if rows_affected > 0 { + tracing::info!( + "{LOG_PREFIX} interrupted {rows_affected} orphaned agent run(s) on startup" + ); + } + Ok(rows_affected) + }) +} + +pub fn list_agent_runs( + workspace_dir: &Path, + request: &AgentRunListRequest, +) -> Result { + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let mut where_clauses = Vec::new(); + let mut values: Vec> = Vec::new(); + + if let Some(status) = request.status.as_deref().filter(|s| !s.trim().is_empty()) { + values.push(Box::new(status.to_string())); + where_clauses.push(format!("status = ?{}", values.len())); + } + if let Some(kind) = request.kind.as_deref().filter(|s| !s.trim().is_empty()) { + values.push(Box::new(kind.to_string())); + where_clauses.push(format!("kind = ?{}", values.len())); + } + if let Some(parent) = request + .parent_run_id + .as_deref() + .filter(|s| !s.trim().is_empty()) + { + values.push(Box::new(parent.to_string())); + where_clauses.push(format!("parent_run_id = ?{}", values.len())); + } + if let Some(thread) = request + .parent_thread_id + .as_deref() + .filter(|s| !s.trim().is_empty()) + { + values.push(Box::new(thread.to_string())); + where_clauses.push(format!("parent_thread_id = ?{}", values.len())); + } + + let where_sql = if where_clauses.is_empty() { + String::new() + } else { + format!("WHERE {}", where_clauses.join(" AND ")) + }; + let count_sql = format!("SELECT COUNT(*) FROM agent_runs {where_sql}"); + let params_ref: Vec<&dyn rusqlite::types::ToSql> = + values.iter().map(|v| v.as_ref()).collect(); + let count = conn.query_row(&count_sql, params_ref.as_slice(), |row| { + row.get::<_, i64>(0) + })? as usize; + + let limit = request.limit.unwrap_or(50).min(500) as i64; + let offset = request.offset.unwrap_or(0) as i64; + values.push(Box::new(limit)); + let limit_idx = values.len(); + values.push(Box::new(offset)); + let offset_idx = values.len(); + + let query_sql = format!( + "SELECT id, kind, parent_run_id, parent_thread_id, agent_id, status, + prompt_ref, worker_thread_id, task_board_id, task_card_id, + checkpoint_path, checkpoint_json, summary, error, metadata_json, + started_at, updated_at, completed_at + FROM agent_runs {where_sql} + ORDER BY updated_at DESC + LIMIT ?{limit_idx} OFFSET ?{offset_idx}" + ); + let params_ref: Vec<&dyn rusqlite::types::ToSql> = + values.iter().map(|v| v.as_ref()).collect(); + let mut stmt = conn.prepare(&query_sql)?; + let rows = stmt.query_map(params_ref.as_slice(), |row| map_agent_run_row(conn, row))?; + let mut runs = Vec::new(); + for row in rows { + runs.push(row?); + } + Ok(AgentRunListResponse { runs, count }) + }) +} + +pub fn list_recent_run_events( + workspace_dir: &Path, + request: &RunEventListRequest, +) -> Result { + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let limit = request.limit.unwrap_or(100).min(1000) as i64; + let after = request.after_sequence.unwrap_or(0) as i64; + let mut stmt = conn.prepare( + "SELECT run_id, sequence, event_type, payload_json, timestamp + FROM run_events + WHERE run_id = ?1 AND sequence > ?2 + ORDER BY sequence ASC + LIMIT ?3", + )?; + let rows = stmt.query_map(params![request.run_id, after, limit], map_run_event_row)?; + let mut events = Vec::new(); + for row in rows { + events.push(row?); + } + Ok(RunEventListResponse { + count: events.len(), + events, + }) + }) +} + +pub fn get_workflow_run(workspace_dir: &Path, id: &str) -> Result> { + tracing::debug!("{LOG_PREFIX} get_workflow_run.entry id={id}"); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let mut stmt = conn.prepare( + "SELECT id, definition_id, parent_thread_id, input_json, phase_states_json, + child_run_ids_json, status, summary, started_at, updated_at, completed_at + FROM workflow_runs WHERE id = ?1", + )?; + let run = stmt + .query_row(params![id], map_workflow_run_row) + .optional()?; + tracing::debug!( + "{LOG_PREFIX} get_workflow_run.exit id={id} found={}", + run.is_some() + ); + Ok(run) + }) +} + +/// List durable workflow runs, most-recently-updated first, with optional +/// filters (definition id, status, parent thread) and pagination. Mirrors +/// [`list_agent_runs`] for the workflow_runs table. +pub fn list_workflow_runs( + workspace_dir: &Path, + request: &WorkflowRunListRequest, +) -> Result { + tracing::debug!( + "{LOG_PREFIX} list_workflow_runs.entry definition={:?} status={:?} parent_thread={:?} limit={:?} offset={:?}", + request.definition_id, + request.status, + request.parent_thread_id, + request.limit, + request.offset + ); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let mut where_clauses = Vec::new(); + let mut values: Vec> = Vec::new(); + + if let Some(definition) = request + .definition_id + .as_deref() + .filter(|s| !s.trim().is_empty()) + { + values.push(Box::new(definition.to_string())); + where_clauses.push(format!("definition_id = ?{}", values.len())); + } + if let Some(status) = request.status.as_deref().filter(|s| !s.trim().is_empty()) { + values.push(Box::new(status.to_string())); + where_clauses.push(format!("status = ?{}", values.len())); + } + if let Some(thread) = request + .parent_thread_id + .as_deref() + .filter(|s| !s.trim().is_empty()) + { + values.push(Box::new(thread.to_string())); + where_clauses.push(format!("parent_thread_id = ?{}", values.len())); + } + + let where_sql = if where_clauses.is_empty() { + String::new() + } else { + format!("WHERE {}", where_clauses.join(" AND ")) + }; + let count_sql = format!("SELECT COUNT(*) FROM workflow_runs {where_sql}"); + let params_ref: Vec<&dyn rusqlite::types::ToSql> = + values.iter().map(|v| v.as_ref()).collect(); + let count = conn.query_row(&count_sql, params_ref.as_slice(), |row| { + row.get::<_, i64>(0) + })? as usize; + + let limit = request.limit.unwrap_or(50).min(500) as i64; + // `offset` is `u64`; convert checked so a value > i64::MAX surfaces a + // clear error instead of wrapping negative and corrupting pagination. + let offset = i64::try_from(request.offset.unwrap_or(0)) + .storage_context("workflow run list offset exceeds i64::MAX")?; + values.push(Box::new(limit)); + let limit_idx = values.len(); + values.push(Box::new(offset)); + let offset_idx = values.len(); + + let query_sql = format!( + "SELECT id, definition_id, parent_thread_id, input_json, phase_states_json, + child_run_ids_json, status, summary, started_at, updated_at, completed_at + FROM workflow_runs {where_sql} + ORDER BY updated_at DESC + LIMIT ?{limit_idx} OFFSET ?{offset_idx}" + ); + let params_ref: Vec<&dyn rusqlite::types::ToSql> = + values.iter().map(|v| v.as_ref()).collect(); + let mut stmt = conn.prepare(&query_sql)?; + let rows = stmt.query_map(params_ref.as_slice(), map_workflow_run_row)?; + let mut runs = Vec::new(); + for row in rows { + runs.push(row?); + } + tracing::debug!( + "{LOG_PREFIX} list_workflow_runs.exit count={count} returned={}", + runs.len() + ); + Ok(WorkflowRunListResponse { runs, count }) + }) +} + +// --------------------------------------------------------------------------- +// Agent-team coordination (issue #3374) +// --------------------------------------------------------------------------- + +/// Insert or update a team row. +pub fn upsert_agent_team(workspace_dir: &Path, upsert: AgentTeamUpsert) -> Result { + let now = Utc::now(); + let created_at = upsert.created_at.unwrap_or(now); + tracing::debug!( + "{LOG_PREFIX} upsert_agent_team.entry id={} lead={} status={}", + upsert.id, + upsert.lead_agent_id, + upsert.status.as_str() + ); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + conn.execute( + "INSERT INTO agent_teams ( + id, parent_thread_id, lead_agent_id, status, summary, + created_at, updated_at, closed_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + ON CONFLICT(id) DO UPDATE SET + parent_thread_id = COALESCE(excluded.parent_thread_id, agent_teams.parent_thread_id), + lead_agent_id = excluded.lead_agent_id, + status = excluded.status, + summary = COALESCE(excluded.summary, agent_teams.summary), + updated_at = excluded.updated_at, + closed_at = COALESCE(excluded.closed_at, agent_teams.closed_at)", + params![ + upsert.id, + upsert.parent_thread_id, + upsert.lead_agent_id, + upsert.status.as_str(), + upsert.summary, + created_at.to_rfc3339(), + now.to_rfc3339(), + upsert.closed_at.map(|dt| dt.to_rfc3339()), + ], + ) + .storage_context("upsert agent team")?; + Ok(()) + })?; + let team = get_agent_team(workspace_dir, &upsert.id)? + .storage_context("agent team missing after upsert")?; + tracing::debug!("{LOG_PREFIX} upsert_agent_team.exit id={}", team.id); + Ok(team) +} + +/// Fetch a single team by id. +pub fn get_agent_team(workspace_dir: &Path, id: &str) -> Result> { + tracing::debug!("{LOG_PREFIX} get_agent_team.entry id={id}"); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let team = get_agent_team_inner(conn, id)?; + tracing::debug!( + "{LOG_PREFIX} get_agent_team.exit id={id} found={}", + team.is_some() + ); + Ok(team) + }) +} + +/// List teams, most-recently-updated first, with optional thread/status filters. +pub fn list_agent_teams( + workspace_dir: &Path, + request: &AgentTeamListRequest, +) -> Result { + tracing::debug!( + "{LOG_PREFIX} list_agent_teams.entry parent_thread={:?} status={:?} limit={:?} offset={:?}", + request.parent_thread_id, + request.status, + request.limit, + request.offset + ); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let mut where_clauses = Vec::new(); + let mut values: Vec> = Vec::new(); + + if let Some(thread) = request + .parent_thread_id + .as_deref() + .filter(|s| !s.trim().is_empty()) + { + values.push(Box::new(thread.to_string())); + where_clauses.push(format!("parent_thread_id = ?{}", values.len())); + } + if let Some(status) = request.status.as_deref().filter(|s| !s.trim().is_empty()) { + values.push(Box::new(status.to_string())); + where_clauses.push(format!("status = ?{}", values.len())); + } + + let where_sql = if where_clauses.is_empty() { + String::new() + } else { + format!("WHERE {}", where_clauses.join(" AND ")) + }; + let count_sql = format!("SELECT COUNT(*) FROM agent_teams {where_sql}"); + let params_ref: Vec<&dyn rusqlite::types::ToSql> = + values.iter().map(|v| v.as_ref()).collect(); + let count = conn.query_row(&count_sql, params_ref.as_slice(), |row| { + row.get::<_, i64>(0) + })? as usize; + + let limit = request.limit.unwrap_or(50).min(500) as i64; + // `offset` is `u64`; convert checked so a value > i64::MAX surfaces a + // clear error instead of wrapping negative and corrupting pagination. + let offset = i64::try_from(request.offset.unwrap_or(0)) + .storage_context("agent team list offset exceeds i64::MAX")?; + values.push(Box::new(limit)); + let limit_idx = values.len(); + values.push(Box::new(offset)); + let offset_idx = values.len(); + + let query_sql = format!( + "SELECT id, parent_thread_id, lead_agent_id, status, summary, + created_at, updated_at, closed_at + FROM agent_teams {where_sql} + ORDER BY updated_at DESC + LIMIT ?{limit_idx} OFFSET ?{offset_idx}" + ); + let params_ref: Vec<&dyn rusqlite::types::ToSql> = + values.iter().map(|v| v.as_ref()).collect(); + let mut stmt = conn.prepare(&query_sql)?; + let rows = stmt.query_map(params_ref.as_slice(), map_agent_team_row)?; + let mut teams = Vec::new(); + for row in rows { + teams.push(row?); + } + tracing::debug!( + "{LOG_PREFIX} list_agent_teams.exit count={count} returned={}", + teams.len() + ); + Ok(AgentTeamListResponse { teams, count }) + }) +} + +/// Insert or update a team member. `UNIQUE(team_id, name)` enforces unique names. +pub fn upsert_agent_team_member( + workspace_dir: &Path, + upsert: AgentTeamMemberUpsert, +) -> Result { + let now = Utc::now(); + let created_at = upsert.created_at.unwrap_or(now); + tracing::debug!( + "{LOG_PREFIX} upsert_agent_team_member.entry id={} team={} name={} status={}", + upsert.id, + upsert.team_id, + upsert.name, + upsert.member_status.as_str() + ); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + conn.execute( + "INSERT INTO agent_team_members ( + id, team_id, name, agent_id, member_status, + current_task_id, worker_thread_id, run_id, created_at, updated_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10) + ON CONFLICT(id) DO UPDATE SET + name = excluded.name, + agent_id = COALESCE(excluded.agent_id, agent_team_members.agent_id), + member_status = excluded.member_status, + current_task_id = COALESCE(excluded.current_task_id, agent_team_members.current_task_id), + worker_thread_id = COALESCE(excluded.worker_thread_id, agent_team_members.worker_thread_id), + run_id = COALESCE(excluded.run_id, agent_team_members.run_id), + updated_at = excluded.updated_at", + params![ + upsert.id, + upsert.team_id, + upsert.name, + upsert.agent_id, + upsert.member_status.as_str(), + upsert.current_task_id, + upsert.worker_thread_id, + upsert.run_id, + created_at.to_rfc3339(), + now.to_rfc3339(), + ], + ) + .storage_context("upsert agent team member")?; + Ok(()) + })?; + let member = get_agent_team_member(workspace_dir, &upsert.id)? + .storage_context("agent team member missing after upsert")?; + tracing::debug!( + "{LOG_PREFIX} upsert_agent_team_member.exit id={}", + member.id + ); + Ok(member) +} + +/// Fetch a single member by id. +pub fn get_agent_team_member(workspace_dir: &Path, id: &str) -> Result> { + tracing::debug!("{LOG_PREFIX} get_agent_team_member.entry id={id}"); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let member = get_agent_team_member_inner(conn, id)?; + tracing::debug!( + "{LOG_PREFIX} get_agent_team_member.exit id={id} found={}", + member.is_some() + ); + Ok(member) + }) +} + +/// List all members of a team, by creation order. +pub fn list_agent_team_members( + workspace_dir: &Path, + team_id: &str, +) -> Result> { + tracing::debug!("{LOG_PREFIX} list_agent_team_members.entry team={team_id}"); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let mut stmt = conn.prepare( + "SELECT id, team_id, name, agent_id, member_status, + current_task_id, worker_thread_id, run_id, created_at, updated_at + FROM agent_team_members WHERE team_id = ?1 + ORDER BY created_at ASC", + )?; + let rows = stmt.query_map(params![team_id], map_agent_team_member_row)?; + let mut members = Vec::new(); + for row in rows { + members.push(row?); + } + tracing::debug!( + "{LOG_PREFIX} list_agent_team_members.exit team={team_id} count={}", + members.len() + ); + Ok(members) + }) +} + +/// Insert or update a team task. +pub fn upsert_agent_team_task( + workspace_dir: &Path, + upsert: AgentTeamTaskUpsert, +) -> Result { + let now = Utc::now(); + let created_at = upsert.created_at.unwrap_or(now); + let depends_on_json = + serde_json::to_string(&upsert.depends_on).storage_context("serialize task depends_on")?; + let evidence_json = + serde_json::to_string(&upsert.evidence).storage_context("serialize task evidence")?; + let gate_status = upsert.gate_status.unwrap_or_else(|| "pending".to_string()); + tracing::debug!( + "{LOG_PREFIX} upsert_agent_team_task.entry id={} team={} status={} deps={}", + upsert.id, + upsert.team_id, + upsert.status.as_str(), + upsert.depends_on.len() + ); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + conn.execute( + "INSERT INTO agent_team_tasks ( + id, team_id, title, objective, status, owner_member_id, + claimed_by_member_id, claim_token, depends_on_json, gate_status, + gate_reason, evidence_json, source_run_id, order_index, + created_at, updated_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, NULL, NULL, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14) + ON CONFLICT(id) DO UPDATE SET + title = excluded.title, + objective = COALESCE(excluded.objective, agent_team_tasks.objective), + status = excluded.status, + -- A claim only means anything while the task is in_progress. + -- Editing a claimed task through this upsert (to change + -- dependencies, or to reset status during recovery) used to + -- leave claimed_by_member_id/claim_token set on a todo row, + -- which strands it: a fresh claim returns AlreadyClaimed, + -- completion returns NotClaimed, and release/shutdown skip it + -- because they only match in_progress. Drop the claim whenever + -- the new status is not in_progress; preserve it otherwise so + -- an unrelated edit does not steal a live claim. + claimed_by_member_id = CASE + WHEN excluded.status = 'in_progress' + THEN agent_team_tasks.claimed_by_member_id ELSE NULL END, + claim_token = CASE + WHEN excluded.status = 'in_progress' + THEN agent_team_tasks.claim_token ELSE NULL END, + owner_member_id = COALESCE(excluded.owner_member_id, agent_team_tasks.owner_member_id), + depends_on_json = excluded.depends_on_json, + gate_status = excluded.gate_status, + gate_reason = COALESCE(excluded.gate_reason, agent_team_tasks.gate_reason), + evidence_json = excluded.evidence_json, + source_run_id = COALESCE(excluded.source_run_id, agent_team_tasks.source_run_id), + order_index = excluded.order_index, + updated_at = excluded.updated_at", + params![ + upsert.id, + upsert.team_id, + upsert.title, + upsert.objective, + upsert.status.as_str(), + upsert.owner_member_id, + depends_on_json, + gate_status, + upsert.gate_reason, + evidence_json, + upsert.source_run_id, + upsert.order_index, + created_at.to_rfc3339(), + now.to_rfc3339(), + ], + ) + .storage_context("upsert agent team task")?; + Ok(()) + })?; + let task = get_agent_team_task(workspace_dir, &upsert.id)? + .storage_context("agent team task missing after upsert")?; + tracing::debug!("{LOG_PREFIX} upsert_agent_team_task.exit id={}", task.id); + Ok(task) +} + +/// Fetch a single task by id. +pub fn get_agent_team_task(workspace_dir: &Path, id: &str) -> Result> { + tracing::debug!("{LOG_PREFIX} get_agent_team_task.entry id={id}"); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let task = get_agent_team_task_inner(conn, id)?; + tracing::debug!( + "{LOG_PREFIX} get_agent_team_task.exit id={id} found={}", + task.is_some() + ); + Ok(task) + }) +} + +/// List all tasks of a team, by `order_index` then creation order. +pub fn list_agent_team_tasks(workspace_dir: &Path, team_id: &str) -> Result> { + tracing::debug!("{LOG_PREFIX} list_agent_team_tasks.entry team={team_id}"); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let mut stmt = conn.prepare( + "SELECT id, team_id, title, objective, status, owner_member_id, + claimed_by_member_id, claim_token, depends_on_json, gate_status, + gate_reason, evidence_json, source_run_id, order_index, + created_at, updated_at + FROM agent_team_tasks WHERE team_id = ?1 + ORDER BY order_index ASC, created_at ASC", + )?; + let rows = stmt.query_map(params![team_id], map_agent_team_task_row)?; + let mut tasks = Vec::new(); + for row in rows { + tasks.push(row?); + } + tracing::debug!( + "{LOG_PREFIX} list_agent_team_tasks.exit team={team_id} count={}", + tasks.len() + ); + Ok(tasks) + }) +} + +/// Atomically claim a task for a member. +/// +/// All steps run inside a single `with_connection` transaction so that the +/// dependency check and the compare-and-swap observe a consistent snapshot: +/// 1. Resolve the task by `(id, team_id)`; absent → [`ClaimOutcome::UnknownTask`]. +/// 2. For every dependency id, look up its status; collect those not `done` +/// into `unmet`. Non-empty → [`ClaimOutcome::Blocked`]. +/// 3. WHERE-guarded `UPDATE ... WHERE claimed_by_member_id IS NULL`: SQLite +/// serializes writers, so exactly one concurrent claimer flips the row from +/// unclaimed to claimed. `rows_affected == 0` → already taken +/// ([`ClaimOutcome::AlreadyClaimed`]); otherwise re-fetch and return +/// [`ClaimOutcome::Claimed`]. +pub fn claim_agent_team_task( + workspace_dir: &Path, + team_id: &str, + task_id: &str, + member_id: &str, + claim_token: &str, +) -> Result { + tracing::debug!( + "{LOG_PREFIX} claim_agent_team_task.entry team={team_id} task={task_id} member={member_id}" + ); + let outcome = crate::session::store::with_transaction(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + + // 1. Resolve the task within this team. + let task = match get_agent_team_task_inner(conn, task_id)? { + Some(task) if task.team_id == team_id => task, + _ => { + tracing::debug!( + "{LOG_PREFIX} claim_agent_team_task.unknown team={team_id} task={task_id}" + ); + return Ok(ClaimOutcome::UnknownTask); + } + }; + + // 2. Dependency gate: every dep must be `done`. + let mut unmet = Vec::new(); + for dep_id in &task.depends_on { + let dep_status: Option = conn + .query_row( + "SELECT status FROM agent_team_tasks WHERE id = ?1 AND team_id = ?2", + params![dep_id, team_id], + |row| row.get(0), + ) + .optional()?; + let is_done = dep_status.as_deref() == Some(AgentTeamTaskStatus::Done.as_str()); + if !is_done { + unmet.push(dep_id.clone()); + } + } + if !unmet.is_empty() { + tracing::debug!( + "{LOG_PREFIX} claim_agent_team_task.blocked team={team_id} task={task_id} unmet={}", + unmet.len() + ); + return Ok(ClaimOutcome::Blocked { unmet }); + } + + // 3. Compare-and-swap on the unclaimed guard. + let now = Utc::now(); + let rows_affected = conn + .execute( + "UPDATE agent_team_tasks + SET claimed_by_member_id = ?1, claim_token = ?2, status = 'in_progress', updated_at = ?3 + WHERE id = ?4 AND team_id = ?5 AND claimed_by_member_id IS NULL", + params![member_id, claim_token, now.to_rfc3339(), task_id, team_id], + ) + .storage_context("compare-and-swap claim agent team task")?; + if rows_affected == 0 { + tracing::debug!( + "{LOG_PREFIX} claim_agent_team_task.already_claimed team={team_id} task={task_id}" + ); + return Ok(ClaimOutcome::AlreadyClaimed); + } + + let claimed = get_agent_team_task_inner(conn, task_id)? + .storage_context("claimed task missing after compare-and-swap")?; + Ok(ClaimOutcome::Claimed(Box::new(claimed))) + })?; + tracing::debug!( + "{LOG_PREFIX} claim_agent_team_task.exit team={team_id} task={task_id} outcome={}", + match &outcome { + ClaimOutcome::Claimed(_) => "claimed", + ClaimOutcome::AlreadyClaimed => "already_claimed", + ClaimOutcome::Blocked { .. } => "blocked", + ClaimOutcome::UnknownTask => "unknown", + } + ); + Ok(outcome) +} + +/// Quality-gate a task's completion and, on pass, transition it to `done`. +/// +/// Runs inside a single transaction so the gate evaluation and the status flip +/// observe one consistent snapshot: +/// 1. Resolve the task by `(id, team_id)`; absent → [`CompletionOutcome::UnknownTask`]. +/// 2. The completer must be the current claimant and the task must be +/// `in_progress`; otherwise [`CompletionOutcome::NotClaimed`]. +/// 3. Evaluate the quality gate (every dependency `done`, claimant matches any +/// pre-assigned owner, evidence present when `require_evidence`). Any unmet +/// invariant records `gate_status = "failed"` + the joined reasons and leaves +/// the task `in_progress` → [`CompletionOutcome::GateFailed`]. +/// 4. On pass, merge `evidence`, set `status = "done"`, `gate_status = "passed"`, +/// clear `gate_reason`, re-fetch → [`CompletionOutcome::Completed`]. +pub fn complete_agent_team_task( + workspace_dir: &Path, + team_id: &str, + task_id: &str, + member_id: &str, + evidence: &[String], + require_evidence: bool, +) -> Result { + tracing::debug!( + "{LOG_PREFIX} complete_agent_team_task.entry team={team_id} task={task_id} member={member_id}" + ); + let outcome = crate::session::store::with_transaction(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + + // 1. Resolve the task within this team. + let task = match get_agent_team_task_inner(conn, task_id)? { + Some(task) if task.team_id == team_id => task, + _ => { + tracing::debug!( + "{LOG_PREFIX} complete_agent_team_task.unknown team={team_id} task={task_id}" + ); + return Ok(CompletionOutcome::UnknownTask); + } + }; + + // 2. Only the current claimant may complete, and only while in progress. + let is_claimant = task.claimed_by_member_id.as_deref() == Some(member_id); + let in_progress = task.status == AgentTeamTaskStatus::InProgress; + if !is_claimant || !in_progress { + tracing::debug!( + "{LOG_PREFIX} complete_agent_team_task.not_claimed team={team_id} task={task_id} claimant={is_claimant} in_progress={in_progress}" + ); + return Ok(CompletionOutcome::NotClaimed); + } + + // Merge prior evidence with the newly-supplied links (de-duplicated, + // order-preserving) so a retry that adds evidence accumulates it. + let mut merged_evidence = task.evidence.clone(); + for link in evidence { + if !merged_evidence.iter().any(|e| e == link) { + merged_evidence.push(link.clone()); + } + } + + // 3. Quality gate. + let reasons = + evaluate_completion_gate(conn, team_id, &task, &merged_evidence, require_evidence)?; + let now = Utc::now(); + if !reasons.is_empty() { + let joined = reasons.join("; "); + // Persist the merged evidence even though the gate failed. Evidence + // accumulates across attempts (see `merged_evidence` above), so + // dropping it here punished a caller for an unrelated gate failure: + // after a dependency was fixed, a retry that did not resend the same + // links would fail `require_evidence` on evidence it had already + // submitted. Only the gate verdict is a failure; the submission is + // still real. + let evidence_json = serde_json::to_string(&merged_evidence) + .storage_context("serialize completion evidence")?; + conn.execute( + "UPDATE agent_team_tasks + SET gate_status = 'failed', gate_reason = ?1, + evidence_json = ?2, updated_at = ?3 + WHERE id = ?4 AND team_id = ?5", + params![joined, evidence_json, now.to_rfc3339(), task_id, team_id], + ) + .storage_context("record failed completion gate")?; + tracing::debug!( + "{LOG_PREFIX} complete_agent_team_task.gate_failed team={team_id} task={task_id} reasons={}", + reasons.len() + ); + return Ok(CompletionOutcome::GateFailed { reasons }); + } + + // 4. Gate passed — flip to done. The WHERE clause is the real CAS: the + // `claimed_by_member_id` guard stops a concurrent shutdown/unclaim from + // completing a task it no longer holds, and the `status = 'in_progress'` + // guard stops a concurrent double-complete by the same member (the + // snapshot check above is a read, not part of the swap — only one of two + // racing UPDATEs flips `in_progress -> done`). + let evidence_json = serde_json::to_string(&merged_evidence) + .storage_context("serialize completion evidence")?; + let rows_affected = conn + .execute( + "UPDATE agent_team_tasks + SET status = 'done', gate_status = 'passed', gate_reason = NULL, + evidence_json = ?1, updated_at = ?2 + WHERE id = ?3 AND team_id = ?4 AND claimed_by_member_id = ?5 + AND status = 'in_progress'", + params![evidence_json, now.to_rfc3339(), task_id, team_id, member_id], + ) + .storage_context("complete agent team task")?; + if rows_affected == 0 { + tracing::debug!( + "{LOG_PREFIX} complete_agent_team_task.lost_claim team={team_id} task={task_id}" + ); + return Ok(CompletionOutcome::NotClaimed); + } + + let done = get_agent_team_task_inner(conn, task_id)? + .storage_context("completed task missing after update")?; + Ok(CompletionOutcome::Completed(Box::new(done))) + })?; + tracing::debug!( + "{LOG_PREFIX} complete_agent_team_task.exit team={team_id} task={task_id} outcome={}", + match &outcome { + CompletionOutcome::Completed(_) => "completed", + CompletionOutcome::GateFailed { .. } => "gate_failed", + CompletionOutcome::NotClaimed => "not_claimed", + CompletionOutcome::UnknownTask => "unknown", + } + ); + Ok(outcome) +} + +/// Evaluate the quality-gate invariants for a completing task. Returns one +/// human-readable reason per unmet invariant (empty = gate passes). +fn evaluate_completion_gate( + conn: &Connection, + team_id: &str, + task: &AgentTeamTask, + merged_evidence: &[String], + require_evidence: bool, +) -> Result> { + let mut reasons = Vec::new(); + + // Every dependency must still be `done` (defends against a dependency that + // regressed after this task was claimed). + for dep_id in &task.depends_on { + let dep_status: Option = conn + .query_row( + "SELECT status FROM agent_team_tasks WHERE id = ?1 AND team_id = ?2", + params![dep_id, team_id], + |row| row.get(0), + ) + .optional()?; + if dep_status.as_deref() != Some(AgentTeamTaskStatus::Done.as_str()) { + reasons.push(format!("dependency {dep_id} is not done")); + } + } + + // No overlapping ownership: a pre-assigned owner must be the one completing. + if let Some(owner) = task + .owner_member_id + .as_ref() + .filter(|owner| Some(owner.as_str()) != task.claimed_by_member_id.as_deref()) + { + reasons.push(format!( + "task is owned by {owner} but claimed by {}", + task.claimed_by_member_id.as_deref().unwrap_or("nobody") + )); + } + + // Evidence gate. + if require_evidence && merged_evidence.is_empty() { + reasons.push("completion requires at least one evidence link".to_string()); + } + + Ok(reasons) +} + +/// Stop a team member and release any task it is actively working on. +/// +/// In one transaction: unclaim the member's `in_progress` tasks back to `todo` +/// (clearing claimant + token so another teammate can pick them up), then mark +/// the member `stopped` and clear its `current_task_id`. Returns the updated +/// member plus the ids of the tasks that were released, or `None` if the member +/// is not part of the team. +pub fn shutdown_agent_team_member( + workspace_dir: &Path, + team_id: &str, + member_id: &str, +) -> Result)>> { + tracing::debug!( + "{LOG_PREFIX} shutdown_agent_team_member.entry team={team_id} member={member_id}" + ); + let result = crate::session::store::with_transaction(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + + // Existence + team-membership check only; the row is intentionally not + // reused — the caller-facing member is re-read after the UPDATEs below so + // it reflects the stopped state. + match get_agent_team_member_inner(conn, member_id)? { + Some(found) if found.team_id == team_id => {} + _ => { + tracing::debug!( + "{LOG_PREFIX} shutdown_agent_team_member.unknown team={team_id} member={member_id}" + ); + return Ok(None); + } + } + + // Collect the ids first so the caller can report exactly what was freed. + let released: Vec = { + let mut stmt = conn.prepare( + "SELECT id FROM agent_team_tasks + WHERE team_id = ?1 AND claimed_by_member_id = ?2 AND status = 'in_progress'", + )?; + let ids = stmt.query_map(params![team_id, member_id], |row| row.get::<_, String>(0))?; + let mut out = Vec::new(); + for id in ids { + out.push(id?); + } + out + }; + + let now = Utc::now(); + conn.execute( + "UPDATE agent_team_tasks + SET claimed_by_member_id = NULL, claim_token = NULL, status = 'todo', updated_at = ?1 + WHERE team_id = ?2 AND claimed_by_member_id = ?3 AND status = 'in_progress'", + params![now.to_rfc3339(), team_id, member_id], + ) + .storage_context("release tasks on member shutdown")?; + conn.execute( + "UPDATE agent_team_members + SET member_status = 'stopped', current_task_id = NULL, updated_at = ?1 + WHERE id = ?2 AND team_id = ?3", + params![now.to_rfc3339(), member_id, team_id], + ) + .storage_context("stop agent team member")?; + + let member = get_agent_team_member_inner(conn, member_id)? + .storage_context("member missing after shutdown")?; + Ok(Some((member, released))) + })?; + tracing::debug!( + "{LOG_PREFIX} shutdown_agent_team_member.exit team={team_id} member={member_id} released={}", + result.as_ref().map(|(_, r)| r.len()).unwrap_or(0) + ); + Ok(result) +} + +/// Mark a member as actively running a task: status → `active`, with the +/// current task id and the worker/run identifiers of the spawned agent. Used by +/// the live runtime right after it claims a task and dispatches a worker. +/// Returns the updated member, or `None` if the member is not in the team. +pub fn mark_agent_team_member_running( + workspace_dir: &Path, + team_id: &str, + member_id: &str, + task_id: &str, + worker_thread_id: &str, + run_id: &str, +) -> Result> { + tracing::debug!( + "{LOG_PREFIX} mark_agent_team_member_running.entry team={team_id} member={member_id} task={task_id} run={run_id}" + ); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let now = Utc::now(); + let changed = conn + .execute( + "UPDATE agent_team_members + SET member_status = 'active', current_task_id = ?1, + worker_thread_id = ?2, run_id = ?3, updated_at = ?4 + WHERE id = ?5 AND team_id = ?6", + params![ + task_id, + worker_thread_id, + run_id, + now.to_rfc3339(), + member_id, + team_id + ], + ) + .storage_context("mark agent team member running")?; + if changed == 0 { + return Ok(None); + } + get_agent_team_member_inner(conn, member_id) + }) +} + +/// Mark a member idle: status → `idle`, clearing `current_task_id`. The +/// `worker_thread_id` / `run_id` are intentionally retained as a pointer to the +/// member's last run for history. Returns the updated member, or `None` if the +/// member is not in the team. Used when a worker run finishes (completed, +/// gate-failed, or failed) so the member is free to pick up new work. +pub fn mark_agent_team_member_idle( + workspace_dir: &Path, + team_id: &str, + member_id: &str, +) -> Result> { + tracing::debug!( + "{LOG_PREFIX} mark_agent_team_member_idle.entry team={team_id} member={member_id}" + ); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let now = Utc::now(); + let changed = conn + .execute( + "UPDATE agent_team_members + SET member_status = 'idle', current_task_id = NULL, updated_at = ?1 + WHERE id = ?2 AND team_id = ?3", + params![now.to_rfc3339(), member_id, team_id], + ) + .storage_context("mark agent team member idle")?; + if changed == 0 { + return Ok(None); + } + get_agent_team_member_inner(conn, member_id) + }) +} + +/// Release a single `in_progress` task back to `todo`, clearing its claim and +/// resetting the quality gate. Returns `true` if a row was actually released +/// (the task existed, belonged to the team, and was `in_progress`). Used by the +/// live runtime when a worker run fails or is aborted, so the task is free for +/// another teammate — the per-task analogue of the bulk release in +/// `shutdown_agent_team_member`. +pub fn release_agent_team_task(workspace_dir: &Path, team_id: &str, task_id: &str) -> Result { + tracing::debug!("{LOG_PREFIX} release_agent_team_task.entry team={team_id} task={task_id}"); + crate::session::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let now = Utc::now(); + let changed = conn + .execute( + "UPDATE agent_team_tasks + SET status = 'todo', claimed_by_member_id = NULL, claim_token = NULL, + gate_status = 'pending', gate_reason = NULL, updated_at = ?1 + WHERE id = ?2 AND team_id = ?3 AND status = 'in_progress'", + params![now.to_rfc3339(), task_id, team_id], + ) + .storage_context("release agent team task")?; + tracing::debug!( + "{LOG_PREFIX} release_agent_team_task.exit team={team_id} task={task_id} released={}", + changed > 0 + ); + Ok(changed > 0) + }) +} + +fn get_agent_team_inner(conn: &Connection, id: &str) -> Result> { + let mut stmt = conn.prepare( + "SELECT id, parent_thread_id, lead_agent_id, status, summary, + created_at, updated_at, closed_at + FROM agent_teams WHERE id = ?1", + )?; + stmt.query_row(params![id], map_agent_team_row) + .optional() + .map_err(Into::into) +} + +fn get_agent_team_member_inner(conn: &Connection, id: &str) -> Result> { + let mut stmt = conn.prepare( + "SELECT id, team_id, name, agent_id, member_status, + current_task_id, worker_thread_id, run_id, created_at, updated_at + FROM agent_team_members WHERE id = ?1", + )?; + stmt.query_row(params![id], map_agent_team_member_row) + .optional() + .map_err(Into::into) +} + +fn get_agent_team_task_inner(conn: &Connection, id: &str) -> Result> { + let mut stmt = conn.prepare( + "SELECT id, team_id, title, objective, status, owner_member_id, + claimed_by_member_id, claim_token, depends_on_json, gate_status, + gate_reason, evidence_json, source_run_id, order_index, + created_at, updated_at + FROM agent_team_tasks WHERE id = ?1", + )?; + stmt.query_row(params![id], map_agent_team_task_row) + .optional() + .map_err(Into::into) +} + +fn map_agent_team_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { + Ok(AgentTeam { + id: row.get(0)?, + parent_thread_id: row.get(1)?, + lead_agent_id: row.get(2)?, + status: AgentTeamStatus::parse(&row.get::<_, String>(3)?), + summary: row.get(4)?, + created_at: parse_rfc3339(&row.get::<_, String>(5)?)?, + updated_at: parse_rfc3339(&row.get::<_, String>(6)?)?, + closed_at: parse_rfc3339_opt(row.get(7)?)?, + }) +} + +fn map_agent_team_member_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { + Ok(AgentTeamMember { + id: row.get(0)?, + team_id: row.get(1)?, + name: row.get(2)?, + agent_id: row.get(3)?, + member_status: AgentTeamMemberStatus::parse(&row.get::<_, String>(4)?), + current_task_id: row.get(5)?, + worker_thread_id: row.get(6)?, + run_id: row.get(7)?, + created_at: parse_rfc3339(&row.get::<_, String>(8)?)?, + updated_at: parse_rfc3339(&row.get::<_, String>(9)?)?, + }) +} + +fn map_agent_team_task_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { + Ok(AgentTeamTask { + id: row.get(0)?, + team_id: row.get(1)?, + title: row.get(2)?, + objective: row.get(3)?, + status: AgentTeamTaskStatus::parse(&row.get::<_, String>(4)?), + owner_member_id: row.get(5)?, + claimed_by_member_id: row.get(6)?, + claim_token: row.get(7)?, + depends_on: serde_json::from_str(&row.get::<_, String>(8)?).unwrap_or_default(), + gate_status: row.get(9)?, + gate_reason: row.get(10)?, + evidence: serde_json::from_str(&row.get::<_, String>(11)?).unwrap_or_default(), + source_run_id: row.get(12)?, + order_index: row.get(13)?, + created_at: parse_rfc3339(&row.get::<_, String>(14)?)?, + updated_at: parse_rfc3339(&row.get::<_, String>(15)?)?, + }) +} + +fn get_agent_run_inner(conn: &Connection, id: &str) -> Result> { + let mut stmt = conn.prepare( + "SELECT id, kind, parent_run_id, parent_thread_id, agent_id, status, + prompt_ref, worker_thread_id, task_board_id, task_card_id, + checkpoint_path, checkpoint_json, summary, error, metadata_json, + started_at, updated_at, completed_at + FROM agent_runs WHERE id = ?1", + )?; + stmt.query_row(params![id], |row| map_agent_run_row(conn, row)) + .optional() + .map_err(Into::into) +} + +fn get_run_telemetry_inner(conn: &Connection, run_id: &str) -> Result { + let mut stmt = conn.prepare( + "SELECT run_id, input_tokens, output_tokens, cached_input_tokens, cost_usd, + elapsed_ms, tool_count, model, provider, error, updated_at + FROM run_telemetry WHERE run_id = ?1", + )?; + stmt.query_row(params![run_id], map_run_telemetry_row) + .storage_context("run telemetry missing after upsert") +} + +fn get_optional_run_telemetry( + conn: &Connection, + run_id: &str, +) -> rusqlite::Result> { + let mut stmt = conn.prepare( + "SELECT run_id, input_tokens, output_tokens, cached_input_tokens, cost_usd, + elapsed_ms, tool_count, model, provider, error, updated_at + FROM run_telemetry WHERE run_id = ?1", + )?; + stmt.query_row(params![run_id], map_run_telemetry_row) + .optional() +} + +fn map_agent_run_row(conn: &Connection, row: &rusqlite::Row<'_>) -> rusqlite::Result { + let id: String = row.get(0)?; + let checkpoint_json: Option = row.get(11)?; + let metadata_json: String = row.get(14)?; + Ok(AgentRun { + id: id.clone(), + kind: super::types::AgentRunKind::parse(&row.get::<_, String>(1)?), + parent_run_id: row.get(2)?, + parent_thread_id: row.get(3)?, + agent_id: row.get(4)?, + status: AgentRunStatus::parse(&row.get::<_, String>(5)?), + prompt_ref: row.get(6)?, + worker_thread_id: row.get(7)?, + task_board_id: row.get(8)?, + task_card_id: row.get(9)?, + checkpoint_path: row.get(10)?, + checkpoint: parse_json_opt(checkpoint_json), + summary: row.get(12)?, + error: row.get(13)?, + metadata: parse_json(metadata_json), + telemetry: get_optional_run_telemetry(conn, &id)?, + started_at: parse_rfc3339(&row.get::<_, String>(15)?)?, + updated_at: parse_rfc3339(&row.get::<_, String>(16)?)?, + completed_at: parse_rfc3339_opt(row.get(17)?)?, + }) +} + +fn map_workflow_run_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { + Ok(WorkflowRun { + id: row.get(0)?, + definition_id: row.get(1)?, + parent_thread_id: row.get(2)?, + input: parse_json(row.get(3)?), + phase_states: parse_json(row.get(4)?), + child_run_ids: serde_json::from_str(&row.get::<_, String>(5)?).unwrap_or_default(), + status: super::types::WorkflowRunStatus::parse(&row.get::<_, String>(6)?), + summary: row.get(7)?, + started_at: parse_rfc3339(&row.get::<_, String>(8)?)?, + updated_at: parse_rfc3339(&row.get::<_, String>(9)?)?, + completed_at: parse_rfc3339_opt(row.get(10)?)?, + }) +} + +fn map_run_event_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { + Ok(RunEvent { + run_id: row.get(0)?, + sequence: row.get::<_, i64>(1)? as u64, + event_type: row.get(2)?, + payload: parse_json(row.get(3)?), + timestamp: parse_rfc3339(&row.get::<_, String>(4)?)?, + }) +} + +fn map_run_telemetry_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { + Ok(RunTelemetry { + run_id: row.get(0)?, + input_tokens: row.get::<_, i64>(1)? as u64, + output_tokens: row.get::<_, i64>(2)? as u64, + cached_input_tokens: row.get::<_, i64>(3)? as u64, + cost_usd: row.get(4)?, + elapsed_ms: row.get::<_, Option>(5)?.map(|v| v as u64), + tool_count: row.get::<_, i64>(6)? as u64, + model: row.get(7)?, + provider: row.get(8)?, + error: row.get(9)?, + updated_at: Some(parse_rfc3339(&row.get::<_, String>(10)?)?), + }) +} + +fn parse_json(raw: String) -> Value { + serde_json::from_str(&raw).unwrap_or_else(|_| json!({})) +} + +fn parse_json_opt(raw: Option) -> Option { + raw.and_then(|value| serde_json::from_str(&value).ok()) +} + +fn parse_rfc3339(raw: &str) -> rusqlite::Result> { + DateTime::parse_from_rfc3339(raw) + .map(|dt| dt.with_timezone(&Utc)) + .map_err(|err| { + rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(err)) + }) +} + +fn parse_rfc3339_opt(raw: Option) -> rusqlite::Result>> { + match raw { + Some(value) => parse_rfc3339(&value).map(Some), + None => Ok(None), + } +} diff --git a/src/session/run_ledger/store.rs b/src/session/run_ledger/store.rs new file mode 100644 index 0000000..9a4498f --- /dev/null +++ b/src/session/run_ledger/store.rs @@ -0,0 +1,127 @@ +use rusqlite::Connection; + +use super::super::context::StorageContext; +use crate::error::Result; + +pub(crate) fn init_run_ledger_schema(conn: &Connection) -> Result<()> { + conn.execute_batch( + "CREATE TABLE IF NOT EXISTS agent_runs ( + id TEXT PRIMARY KEY, + kind TEXT NOT NULL, + parent_run_id TEXT, + parent_thread_id TEXT, + agent_id TEXT, + status TEXT NOT NULL, + prompt_ref TEXT, + worker_thread_id TEXT, + task_board_id TEXT, + task_card_id TEXT, + checkpoint_path TEXT, + checkpoint_json TEXT, + summary TEXT, + error TEXT, + metadata_json TEXT NOT NULL DEFAULT '{}', + started_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + completed_at TEXT + ); + CREATE INDEX IF NOT EXISTS idx_agent_runs_status ON agent_runs(status); + CREATE INDEX IF NOT EXISTS idx_agent_runs_kind ON agent_runs(kind); + CREATE INDEX IF NOT EXISTS idx_agent_runs_parent ON agent_runs(parent_run_id); + CREATE INDEX IF NOT EXISTS idx_agent_runs_thread ON agent_runs(parent_thread_id); + CREATE INDEX IF NOT EXISTS idx_agent_runs_updated ON agent_runs(updated_at); + CREATE INDEX IF NOT EXISTS idx_agent_runs_worker_thread ON agent_runs(worker_thread_id); + + CREATE TABLE IF NOT EXISTS workflow_runs ( + id TEXT PRIMARY KEY, + definition_id TEXT NOT NULL, + parent_thread_id TEXT, + input_json TEXT NOT NULL DEFAULT '{}', + phase_states_json TEXT NOT NULL DEFAULT '{}', + child_run_ids_json TEXT NOT NULL DEFAULT '[]', + status TEXT NOT NULL, + summary TEXT, + started_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + completed_at TEXT + ); + CREATE INDEX IF NOT EXISTS idx_workflow_runs_definition ON workflow_runs(definition_id); + CREATE INDEX IF NOT EXISTS idx_workflow_runs_status ON workflow_runs(status); + CREATE INDEX IF NOT EXISTS idx_workflow_runs_thread ON workflow_runs(parent_thread_id); + + CREATE TABLE IF NOT EXISTS run_events ( + run_id TEXT NOT NULL, + sequence INTEGER NOT NULL, + event_type TEXT NOT NULL, + payload_json TEXT NOT NULL DEFAULT '{}', + timestamp TEXT NOT NULL, + PRIMARY KEY (run_id, sequence) + ); + CREATE INDEX IF NOT EXISTS idx_run_events_timestamp ON run_events(timestamp); + + CREATE TABLE IF NOT EXISTS run_telemetry ( + run_id TEXT PRIMARY KEY, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + cached_input_tokens INTEGER NOT NULL DEFAULT 0, + cost_usd REAL NOT NULL DEFAULT 0.0, + elapsed_ms INTEGER, + tool_count INTEGER NOT NULL DEFAULT 0, + model TEXT, + provider TEXT, + error TEXT, + updated_at TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS agent_teams ( + id TEXT PRIMARY KEY, + parent_thread_id TEXT, + lead_agent_id TEXT NOT NULL, + status TEXT NOT NULL, + summary TEXT, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + closed_at TEXT + ); + CREATE INDEX IF NOT EXISTS idx_agent_teams_thread ON agent_teams(parent_thread_id); + CREATE INDEX IF NOT EXISTS idx_agent_teams_status ON agent_teams(status); + + CREATE TABLE IF NOT EXISTS agent_team_members ( + id TEXT PRIMARY KEY, + team_id TEXT NOT NULL, + name TEXT NOT NULL, + agent_id TEXT, + member_status TEXT NOT NULL, + current_task_id TEXT, + worker_thread_id TEXT, + run_id TEXT, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + UNIQUE(team_id, name) + ); + CREATE INDEX IF NOT EXISTS idx_agent_team_members_team ON agent_team_members(team_id); + + CREATE TABLE IF NOT EXISTS agent_team_tasks ( + id TEXT PRIMARY KEY, + team_id TEXT NOT NULL, + title TEXT NOT NULL, + objective TEXT, + status TEXT NOT NULL, + owner_member_id TEXT, + claimed_by_member_id TEXT, + claim_token TEXT, + depends_on_json TEXT NOT NULL DEFAULT '[]', + gate_status TEXT NOT NULL DEFAULT 'pending', + gate_reason TEXT, + evidence_json TEXT NOT NULL DEFAULT '[]', + source_run_id TEXT, + order_index INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_agent_team_tasks_team ON agent_team_tasks(team_id); + CREATE INDEX IF NOT EXISTS idx_agent_team_tasks_status ON agent_team_tasks(status); + CREATE INDEX IF NOT EXISTS idx_agent_team_tasks_claimed ON agent_team_tasks(claimed_by_member_id);", + ) + .storage_context("failed to initialize run ledger schema") +} diff --git a/src/session/run_ledger/test.rs b/src/session/run_ledger/test.rs new file mode 100644 index 0000000..fe6d629 --- /dev/null +++ b/src/session/run_ledger/test.rs @@ -0,0 +1,638 @@ +//! Module-local unit tests for [`crate::session::run_ledger`]. +//! +//! Consolidated here per AGENTS.md: one `test.rs` per module directory. + +use super::ops::*; +use super::types::*; +use chrono::Utc; +use serde_json::json; +use std::path::Path; +use tempfile::TempDir; + +/// Workspace root for a test: the ledger derives its database path from +/// this, so a fresh `TempDir` per test gives a fresh database. +fn test_workspace(dir: &TempDir) -> &Path { + dir.path() +} + +// ── Regressions for the review findings on PR #90 ───────────────────── + +/// An upsert that moves a claimed task off `in_progress` must drop the +/// claim, or the row is stranded: a new claim sees AlreadyClaimed, +/// completion sees NotClaimed, and release/shutdown skip it entirely. +#[test] +fn upsert_clears_the_claim_when_a_task_leaves_in_progress() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + seed_team(workspace_dir, "team-strand"); + seed_member(workspace_dir, "team-strand", "m1"); + seed_task(workspace_dir, "team-strand", "task-strand", vec![]); + + let claimed = + claim_agent_team_task(workspace_dir, "team-strand", "task-strand", "m1", "tok").unwrap(); + assert!(matches!(claimed, ClaimOutcome::Claimed(_))); + + // Recovery-style edit that resets status back to todo. + upsert_agent_team_task( + workspace_dir, + AgentTeamTaskUpsert { + id: "task-strand".into(), + team_id: "team-strand".into(), + title: "task task-strand".into(), + objective: None, + status: AgentTeamTaskStatus::Todo, + owner_member_id: None, + depends_on: vec![], + gate_status: None, + gate_reason: None, + evidence: vec![], + source_run_id: None, + order_index: 0, + created_at: None, + }, + ) + .unwrap(); + + let after = get_agent_team_task(workspace_dir, "task-strand") + .unwrap() + .expect("task present"); + assert_eq!(after.claimed_by_member_id, None, "claim must be dropped"); + assert_eq!(after.claim_token, None); + + // ...and the task is claimable again rather than stranded. + let reclaimed = + claim_agent_team_task(workspace_dir, "team-strand", "task-strand", "m1", "tok2").unwrap(); + assert!( + matches!(reclaimed, ClaimOutcome::Claimed(_)), + "a released task must be re-claimable, got {reclaimed:?}" + ); +} + +/// A live claim survives an unrelated edit that keeps the task in progress. +#[test] +fn upsert_preserves_a_live_claim_while_in_progress() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + seed_team(workspace_dir, "team-keep"); + seed_member(workspace_dir, "team-keep", "m1"); + seed_task(workspace_dir, "team-keep", "task-keep", vec![]); + claim_agent_team_task(workspace_dir, "team-keep", "task-keep", "m1", "tok").unwrap(); + + upsert_agent_team_task( + workspace_dir, + AgentTeamTaskUpsert { + id: "task-keep".into(), + team_id: "team-keep".into(), + title: "retitled".into(), + objective: None, + status: AgentTeamTaskStatus::InProgress, + owner_member_id: None, + depends_on: vec![], + gate_status: None, + gate_reason: None, + evidence: vec![], + source_run_id: None, + order_index: 0, + created_at: None, + }, + ) + .unwrap(); + + let after = get_agent_team_task(workspace_dir, "task-keep") + .unwrap() + .expect("task present"); + assert_eq!(after.title, "retitled"); + assert_eq!( + after.claimed_by_member_id.as_deref(), + Some("m1"), + "an in-progress edit must not steal the claim" + ); +} + +/// A partial telemetry upsert must not fail the NOT NULL constraints. +/// +/// The counters are `Option` so one field can be updated alone, but the +/// columns are `NOT NULL DEFAULT` and SQLite does not apply a column +/// default to an explicitly supplied NULL. Recording only `model` used to +/// fail outright on the first write for a run. +#[test] +fn partial_telemetry_upsert_applies_column_defaults() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + + let telemetry = upsert_run_telemetry( + workspace_dir, + RunTelemetryUpsert { + run_id: "run-partial".into(), + model: Some("claude-x".into()), + ..Default::default() + }, + ) + .expect("a model-only upsert must succeed"); + + assert_eq!(telemetry.model.as_deref(), Some("claude-x")); + assert_eq!(telemetry.input_tokens, 0, "counters default, not NULL"); + assert_eq!(telemetry.output_tokens, 0); + assert_eq!(telemetry.cost_usd, 0.0); +} + +/// A later partial upsert must not clobber fields it does not carry. +#[test] +fn partial_telemetry_upsert_preserves_untouched_fields() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + + upsert_run_telemetry( + workspace_dir, + RunTelemetryUpsert { + run_id: "run-merge".into(), + input_tokens: Some(120), + model: Some("claude-x".into()), + ..Default::default() + }, + ) + .unwrap(); + + // Second write carries only the error; everything else must survive. + let merged = upsert_run_telemetry( + workspace_dir, + RunTelemetryUpsert { + run_id: "run-merge".into(), + error: Some("boom".into()), + ..Default::default() + }, + ) + .unwrap(); + + assert_eq!(merged.error.as_deref(), Some("boom")); + assert_eq!(merged.input_tokens, 120, "prior counter must survive"); + assert_eq!(merged.model.as_deref(), Some("claude-x")); +} + +/// Sequences are allocated by the INSERT itself, so appends stay dense and +/// ordered rather than racing on a read-then-write. +#[test] +fn run_event_sequences_are_allocated_by_the_insert() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + + let seqs: Vec = (0..5) + .map(|i| { + append_run_event( + workspace_dir, + RunEventAppend { + run_id: "run-seq".into(), + event_type: format!("evt-{i}"), + payload: json!({ "i": i }), + }, + ) + .unwrap() + .sequence + }) + .collect(); + + assert_eq!(seqs, vec![1, 2, 3, 4, 5]); + + // A second run numbers independently from 1. + let other = append_run_event( + workspace_dir, + RunEventAppend { + run_id: "run-other".into(), + event_type: "evt".into(), + payload: json!({}), + }, + ) + .unwrap(); + assert_eq!(other.sequence, 1); +} + +#[test] +fn agent_run_append_list_get_and_events_are_ordered() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + + let run = upsert_agent_run( + workspace_dir, + AgentRunUpsert { + id: "run-1".into(), + kind: AgentRunKind::Subagent, + parent_run_id: Some("parent".into()), + parent_thread_id: Some("thread-1".into()), + agent_id: Some("researcher".into()), + status: AgentRunStatus::Running, + prompt_ref: Some("worker-1:user:seed".into()), + worker_thread_id: Some("worker-1".into()), + task_board_id: None, + task_card_id: None, + checkpoint_path: None, + checkpoint: None, + summary: None, + error: None, + metadata: json!({"source": "test"}), + started_at: None, + completed_at: None, + }, + ) + .unwrap(); + assert_eq!(run.status, AgentRunStatus::Running); + + append_run_event( + workspace_dir, + RunEventAppend { + run_id: "run-1".into(), + event_type: "spawned".into(), + payload: json!({"agentId": "researcher"}), + }, + ) + .unwrap(); + append_run_event( + workspace_dir, + RunEventAppend { + run_id: "run-1".into(), + event_type: "completed".into(), + payload: json!({"elapsedMs": 12}), + }, + ) + .unwrap(); + + let events = list_recent_run_events( + workspace_dir, + &RunEventListRequest { + run_id: "run-1".into(), + after_sequence: Some(0), + limit: None, + }, + ) + .unwrap(); + assert_eq!(events.events.len(), 2); + assert_eq!(events.events[0].sequence, 1); + assert_eq!(events.events[1].sequence, 2); + + let list = list_agent_runs( + workspace_dir, + &AgentRunListRequest { + parent_thread_id: Some("thread-1".into()), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(list.count, 1); + assert_eq!(list.runs[0].worker_thread_id.as_deref(), Some("worker-1")); +} + +#[test] +fn transition_sets_status_and_clears_error_and_completed_at() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + + // Seed a failed run carrying an error + completion time. + let completed_at = Utc::now(); + upsert_agent_run( + workspace_dir, + AgentRunUpsert { + id: "run-1".into(), + kind: AgentRunKind::Subagent, + parent_run_id: None, + parent_thread_id: Some("thread-1".into()), + agent_id: Some("researcher".into()), + status: AgentRunStatus::Failed, + prompt_ref: None, + worker_thread_id: None, + task_board_id: None, + task_card_id: None, + checkpoint_path: None, + checkpoint: None, + summary: None, + error: Some("boom".into()), + metadata: json!({}), + started_at: None, + completed_at: Some(completed_at), + }, + ) + .unwrap(); + + // Re-queue: passing None for both columns must CLEAR them (the upsert + // path's COALESCE cannot do this — that is the whole reason this op + // exists). + let updated = + transition_agent_run_status(workspace_dir, "run-1", AgentRunStatus::Pending, None, None) + .unwrap() + .expect("run present"); + assert_eq!(updated.status, AgentRunStatus::Pending); + assert_eq!(updated.error, None); + assert_eq!(updated.completed_at, None); + + // Stopping: status + error + completion are all set verbatim. + let stopped_at = Utc::now(); + let updated = transition_agent_run_status( + workspace_dir, + "run-1", + AgentRunStatus::Cancelled, + Some("manual"), + Some(stopped_at), + ) + .unwrap() + .expect("run present"); + assert_eq!(updated.status, AgentRunStatus::Cancelled); + assert_eq!(updated.error.as_deref(), Some("manual")); + assert!(updated.completed_at.is_some()); +} + +#[test] +fn transition_unknown_run_returns_none() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + let result = + transition_agent_run_status(workspace_dir, "ghost", AgentRunStatus::Pending, None, None) + .unwrap(); + assert!(result.is_none()); +} + +fn seed_team(workspace_dir: &Path, team_id: &str) { + upsert_agent_team( + workspace_dir, + AgentTeamUpsert { + id: team_id.into(), + parent_thread_id: Some("thread-team".into()), + lead_agent_id: "lead".into(), + status: AgentTeamStatus::Active, + summary: None, + created_at: None, + closed_at: None, + }, + ) + .unwrap(); +} + +fn seed_task(workspace_dir: &Path, team_id: &str, task_id: &str, depends_on: Vec) { + upsert_agent_team_task( + workspace_dir, + AgentTeamTaskUpsert { + id: task_id.into(), + team_id: team_id.into(), + title: format!("task {task_id}"), + objective: None, + status: AgentTeamTaskStatus::Todo, + owner_member_id: None, + depends_on, + gate_status: None, + gate_reason: None, + evidence: vec![], + source_run_id: None, + order_index: 0, + created_at: None, + }, + ) + .unwrap(); +} + +#[test] +fn claim_is_atomic_first_wins_then_already_claimed() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + seed_team(workspace_dir, "team-1"); + seed_task(workspace_dir, "team-1", "task-a", vec![]); + + let first = claim_agent_team_task(workspace_dir, "team-1", "task-a", "m1", "tok-1").unwrap(); + match first { + ClaimOutcome::Claimed(task) => { + assert_eq!(task.claimed_by_member_id.as_deref(), Some("m1")); + assert_eq!(task.status, AgentTeamTaskStatus::InProgress); + } + other => panic!("expected Claimed, got {other:?}"), + } + + let second = claim_agent_team_task(workspace_dir, "team-1", "task-a", "m2", "tok-2").unwrap(); + assert_eq!(second, ClaimOutcome::AlreadyClaimed); +} + +#[test] +fn claim_unknown_task_returns_unknown() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + seed_team(workspace_dir, "team-1"); + let outcome = claim_agent_team_task(workspace_dir, "team-1", "ghost", "m1", "tok").unwrap(); + assert_eq!(outcome, ClaimOutcome::UnknownTask); +} + +fn seed_member(workspace_dir: &Path, team_id: &str, member_id: &str) { + upsert_agent_team_member( + workspace_dir, + AgentTeamMemberUpsert { + id: member_id.into(), + team_id: team_id.into(), + name: member_id.into(), + agent_id: None, + member_status: AgentTeamMemberStatus::Pending, + current_task_id: None, + worker_thread_id: None, + run_id: None, + created_at: None, + }, + ) + .unwrap(); +} + +#[test] +fn mark_member_running_then_idle_keeps_run_pointer() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + seed_team(workspace_dir, "team-1"); + seed_member(workspace_dir, "team-1", "m1"); + seed_task(workspace_dir, "team-1", "task-a", vec![]); + claim_agent_team_task(workspace_dir, "team-1", "task-a", "m1", "tok-1").unwrap(); + + let running = mark_agent_team_member_running( + workspace_dir, + "team-1", + "m1", + "task-a", + "worker-x", + "run-x", + ) + .unwrap() + .expect("member updated"); + assert_eq!(running.member_status, AgentTeamMemberStatus::Active); + assert_eq!(running.current_task_id.as_deref(), Some("task-a")); + assert_eq!(running.worker_thread_id.as_deref(), Some("worker-x")); + assert_eq!(running.run_id.as_deref(), Some("run-x")); + + let idle = mark_agent_team_member_idle(workspace_dir, "team-1", "m1") + .unwrap() + .expect("member updated"); + assert_eq!(idle.member_status, AgentTeamMemberStatus::Idle); + assert_eq!(idle.current_task_id, None); + // worker/run pointer retained as last-run history. + assert_eq!(idle.worker_thread_id.as_deref(), Some("worker-x")); + assert_eq!(idle.run_id.as_deref(), Some("run-x")); + + // Unknown member → None, no-op. + assert!( + mark_agent_team_member_running(workspace_dir, "team-1", "ghost", "task-a", "w", "r") + .unwrap() + .is_none() + ); + assert!( + mark_agent_team_member_idle(workspace_dir, "team-1", "ghost") + .unwrap() + .is_none() + ); +} + +#[test] +fn release_task_frees_in_progress_only() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + seed_team(workspace_dir, "team-1"); + seed_member(workspace_dir, "team-1", "m1"); + seed_task(workspace_dir, "team-1", "task-a", vec![]); + claim_agent_team_task(workspace_dir, "team-1", "task-a", "m1", "tok-1").unwrap(); + + // In progress → released back to todo, claim cleared, gate reset. + assert!(release_agent_team_task(workspace_dir, "team-1", "task-a").unwrap()); + let task = get_agent_team_task(workspace_dir, "task-a") + .unwrap() + .unwrap(); + assert_eq!(task.status, AgentTeamTaskStatus::Todo); + assert_eq!(task.claimed_by_member_id, None); + assert_eq!(task.claim_token, None); + assert_eq!(task.gate_status, "pending"); + + // Already todo (not in_progress) → no-op, returns false. + assert!(!release_agent_team_task(workspace_dir, "team-1", "task-a").unwrap()); + // Unknown task → false. + assert!(!release_agent_team_task(workspace_dir, "team-1", "ghost").unwrap()); +} + +#[test] +fn claim_blocked_until_dependency_done() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + seed_team(workspace_dir, "team-1"); + seed_task(workspace_dir, "team-1", "task-a", vec![]); + seed_task(workspace_dir, "team-1", "task-b", vec!["task-a".into()]); + + // B is blocked while A is still todo. + let blocked = claim_agent_team_task(workspace_dir, "team-1", "task-b", "m1", "tok").unwrap(); + assert_eq!( + blocked, + ClaimOutcome::Blocked { + unmet: vec!["task-a".into()] + } + ); + + // Mark A done, then B claims fine. + upsert_agent_team_task( + workspace_dir, + AgentTeamTaskUpsert { + id: "task-a".into(), + team_id: "team-1".into(), + title: "task task-a".into(), + objective: None, + status: AgentTeamTaskStatus::Done, + owner_member_id: None, + depends_on: vec![], + gate_status: None, + gate_reason: None, + evidence: vec![], + source_run_id: None, + order_index: 0, + created_at: None, + }, + ) + .unwrap(); + + let ok = claim_agent_team_task(workspace_dir, "team-1", "task-b", "m1", "tok").unwrap(); + assert!(matches!(ok, ClaimOutcome::Claimed(_))); +} + +#[test] +fn team_members_and_tasks_list_back() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + seed_team(workspace_dir, "team-1"); + upsert_agent_team_member( + workspace_dir, + AgentTeamMemberUpsert { + id: "mem-1".into(), + team_id: "team-1".into(), + name: "alice".into(), + agent_id: Some("researcher".into()), + member_status: AgentTeamMemberStatus::Active, + current_task_id: None, + worker_thread_id: None, + run_id: None, + created_at: None, + }, + ) + .unwrap(); + seed_task(workspace_dir, "team-1", "task-a", vec![]); + + let members = list_agent_team_members(workspace_dir, "team-1").unwrap(); + assert_eq!(members.len(), 1); + assert_eq!(members[0].name, "alice"); + + let tasks = list_agent_team_tasks(workspace_dir, "team-1").unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].id, "task-a"); + + let teams = list_agent_teams(workspace_dir, &AgentTeamListRequest::default()).unwrap(); + assert_eq!(teams.count, 1); +} + +fn seed_run(workspace_dir: &Path, id: &str, status: AgentRunStatus) { + upsert_agent_run( + workspace_dir, + AgentRunUpsert { + id: id.into(), + kind: AgentRunKind::Subagent, + parent_run_id: None, + parent_thread_id: Some("thread-1".into()), + agent_id: Some("tinyplace_agent".into()), + status, + prompt_ref: None, + worker_thread_id: None, + task_board_id: None, + task_card_id: None, + checkpoint_path: None, + checkpoint: None, + summary: None, + error: None, + metadata: json!({}), + started_at: None, + completed_at: None, + }, + ) + .unwrap(); +} + +#[test] +fn interrupt_orphaned_runs_settles_only_non_terminal_inflight_rows() { + let dir = TempDir::new().unwrap(); + let workspace_dir = test_workspace(&dir); + + seed_run(workspace_dir, "run-running", AgentRunStatus::Running); + seed_run(workspace_dir, "run-pending", AgentRunStatus::Pending); + seed_run(workspace_dir, "run-completed", AgentRunStatus::Completed); + seed_run(workspace_dir, "run-awaiting", AgentRunStatus::AwaitingUser); + + let settled = interrupt_orphaned_agent_runs(workspace_dir).unwrap(); + assert_eq!(settled, 2, "only running + pending are orphaned at boot"); + + let get = |id: &str| { + get_agent_run(workspace_dir, id) + .unwrap() + .expect("run present") + }; + // Orphaned in-flight rows become terminal `interrupted` with a completion time… + let running = get("run-running"); + assert_eq!(running.status, AgentRunStatus::Interrupted); + assert!(running.completed_at.is_some()); + assert_eq!(get("run-pending").status, AgentRunStatus::Interrupted); + // …already-terminal and resumable rows are untouched. + assert_eq!(get("run-completed").status, AgentRunStatus::Completed); + assert_eq!(get("run-awaiting").status, AgentRunStatus::AwaitingUser); + + // Idempotent: a second sweep finds nothing left to settle. + assert_eq!(interrupt_orphaned_agent_runs(workspace_dir).unwrap(), 0); +} diff --git a/src/session/run_ledger/types.rs b/src/session/run_ledger/types.rs new file mode 100644 index 0000000..1f093e0 --- /dev/null +++ b/src/session/run_ledger/types.rs @@ -0,0 +1,547 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AgentRunKind { + Subagent, + WorkerThread, + BackgroundAgent, + TeamMember, + WorkflowChild, +} + +impl AgentRunKind { + pub fn as_str(self) -> &'static str { + match self { + Self::Subagent => "subagent", + Self::WorkerThread => "worker_thread", + Self::BackgroundAgent => "background_agent", + Self::TeamMember => "team_member", + Self::WorkflowChild => "workflow_child", + } + } + + pub fn parse(raw: &str) -> Self { + match raw { + "worker_thread" => Self::WorkerThread, + "background_agent" => Self::BackgroundAgent, + "team_member" => Self::TeamMember, + "workflow_child" => Self::WorkflowChild, + _ => Self::Subagent, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AgentRunStatus { + Pending, + Running, + AwaitingUser, + Paused, + Completed, + Failed, + Cancelled, + Interrupted, +} + +impl AgentRunStatus { + pub fn as_str(self) -> &'static str { + match self { + Self::Pending => "pending", + Self::Running => "running", + Self::AwaitingUser => "awaiting_user", + Self::Paused => "paused", + Self::Completed => "completed", + Self::Failed => "failed", + Self::Cancelled => "cancelled", + Self::Interrupted => "interrupted", + } + } + + pub fn parse(raw: &str) -> Self { + match raw { + "running" => Self::Running, + "awaiting_user" => Self::AwaitingUser, + "paused" => Self::Paused, + "completed" => Self::Completed, + "failed" => Self::Failed, + "cancelled" => Self::Cancelled, + "interrupted" => Self::Interrupted, + _ => Self::Pending, + } + } + + pub fn is_terminal(self) -> bool { + matches!( + self, + Self::Completed | Self::Failed | Self::Cancelled | Self::Interrupted + ) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum WorkflowRunStatus { + Pending, + Running, + Completed, + Failed, + Cancelled, + Interrupted, +} + +impl WorkflowRunStatus { + pub fn as_str(self) -> &'static str { + match self { + Self::Pending => "pending", + Self::Running => "running", + Self::Completed => "completed", + Self::Failed => "failed", + Self::Cancelled => "cancelled", + Self::Interrupted => "interrupted", + } + } + + pub fn parse(raw: &str) -> Self { + match raw { + "running" => Self::Running, + "completed" => Self::Completed, + "failed" => Self::Failed, + "cancelled" => Self::Cancelled, + "interrupted" => Self::Interrupted, + _ => Self::Pending, + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentRun { + pub id: String, + pub kind: AgentRunKind, + pub parent_run_id: Option, + pub parent_thread_id: Option, + pub agent_id: Option, + pub status: AgentRunStatus, + pub prompt_ref: Option, + pub worker_thread_id: Option, + pub task_board_id: Option, + pub task_card_id: Option, + pub checkpoint_path: Option, + pub checkpoint: Option, + pub summary: Option, + pub error: Option, + pub metadata: Value, + pub telemetry: Option, + pub started_at: DateTime, + pub updated_at: DateTime, + pub completed_at: Option>, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WorkflowRun { + pub id: String, + pub definition_id: String, + pub parent_thread_id: Option, + pub input: Value, + pub phase_states: Value, + pub child_run_ids: Vec, + pub status: WorkflowRunStatus, + pub summary: Option, + pub started_at: DateTime, + pub updated_at: DateTime, + pub completed_at: Option>, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RunEvent { + pub run_id: String, + pub sequence: u64, + pub event_type: String, + pub payload: Value, + pub timestamp: DateTime, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] +#[serde(rename_all = "camelCase")] +pub struct RunTelemetry { + pub run_id: String, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_input_tokens: u64, + pub cost_usd: f64, + pub elapsed_ms: Option, + pub tool_count: u64, + pub model: Option, + pub provider: Option, + pub error: Option, + pub updated_at: Option>, +} + +#[derive(Debug, Clone)] +pub struct AgentRunUpsert { + pub id: String, + pub kind: AgentRunKind, + pub parent_run_id: Option, + pub parent_thread_id: Option, + pub agent_id: Option, + pub status: AgentRunStatus, + pub prompt_ref: Option, + pub worker_thread_id: Option, + pub task_board_id: Option, + pub task_card_id: Option, + pub checkpoint_path: Option, + pub checkpoint: Option, + pub summary: Option, + pub error: Option, + pub metadata: Value, + pub started_at: Option>, + pub completed_at: Option>, +} + +#[derive(Debug, Clone)] +pub struct WorkflowRunUpsert { + pub id: String, + pub definition_id: String, + pub parent_thread_id: Option, + pub input: Value, + pub phase_states: Value, + pub child_run_ids: Vec, + pub status: WorkflowRunStatus, + pub summary: Option, + pub started_at: Option>, + pub completed_at: Option>, +} + +#[derive(Debug, Clone)] +pub struct RunEventAppend { + pub run_id: String, + pub event_type: String, + pub payload: Value, +} + +#[derive(Debug, Clone, Default)] +pub struct RunTelemetryUpsert { + pub run_id: String, + pub input_tokens: Option, + pub output_tokens: Option, + pub cached_input_tokens: Option, + pub cost_usd: Option, + pub elapsed_ms: Option, + pub tool_count: Option, + pub model: Option, + pub provider: Option, + pub error: Option, +} + +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentRunListRequest { + #[serde(default)] + pub status: Option, + #[serde(default)] + pub kind: Option, + #[serde(default)] + pub parent_run_id: Option, + #[serde(default)] + pub parent_thread_id: Option, + #[serde(default)] + pub limit: Option, + #[serde(default)] + pub offset: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentRunListResponse { + pub runs: Vec, + pub count: usize, +} + +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WorkflowRunListRequest { + #[serde(default)] + pub definition_id: Option, + #[serde(default)] + pub status: Option, + #[serde(default)] + pub parent_thread_id: Option, + /// `u64` to match the `TypeSchema::U64` the controller advertises (the RPC + /// scalar-coercion layer only handles `U64`). Capped at 500 in `list_workflow_runs`. + #[serde(default)] + pub limit: Option, + #[serde(default)] + pub offset: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WorkflowRunListResponse { + pub runs: Vec, + pub count: usize, +} + +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RunEventListRequest { + pub run_id: String, + #[serde(default)] + pub after_sequence: Option, + #[serde(default)] + pub limit: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct RunEventListResponse { + pub events: Vec, + pub count: usize, +} + +// --------------------------------------------------------------------------- +// Agent-team coordination (issue #3374) +// --------------------------------------------------------------------------- + +/// Lifecycle of an agent team. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AgentTeamStatus { + Active, + Closed, +} + +impl AgentTeamStatus { + pub fn as_str(self) -> &'static str { + match self { + Self::Active => "active", + Self::Closed => "closed", + } + } + + /// Parse a stored status string (named `parse`, not `from_str`, to match the + /// run-ledger status-enum convention and avoid the `FromStr` clippy lint). + pub fn parse(raw: &str) -> Self { + match raw { + "closed" => Self::Closed, + _ => Self::Active, + } + } +} + +/// Lifecycle of a single team member. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AgentTeamMemberStatus { + Pending, + Active, + Idle, + Stopped, +} + +impl AgentTeamMemberStatus { + pub fn as_str(self) -> &'static str { + match self { + Self::Pending => "pending", + Self::Active => "active", + Self::Idle => "idle", + Self::Stopped => "stopped", + } + } + + pub fn parse(raw: &str) -> Self { + match raw { + "active" => Self::Active, + "idle" => Self::Idle, + "stopped" => Self::Stopped, + _ => Self::Pending, + } + } +} + +/// Lifecycle of a coordination task within a team. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AgentTeamTaskStatus { + Todo, + Ready, + InProgress, + Blocked, + Done, +} + +impl AgentTeamTaskStatus { + pub fn as_str(self) -> &'static str { + match self { + Self::Todo => "todo", + Self::Ready => "ready", + Self::InProgress => "in_progress", + Self::Blocked => "blocked", + Self::Done => "done", + } + } + + pub fn parse(raw: &str) -> Self { + match raw { + "ready" => Self::Ready, + "in_progress" => Self::InProgress, + "blocked" => Self::Blocked, + "done" => Self::Done, + _ => Self::Todo, + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentTeam { + pub id: String, + pub parent_thread_id: Option, + pub lead_agent_id: String, + pub status: AgentTeamStatus, + pub summary: Option, + pub created_at: DateTime, + pub updated_at: DateTime, + pub closed_at: Option>, +} + +#[derive(Debug, Clone)] +pub struct AgentTeamUpsert { + pub id: String, + pub parent_thread_id: Option, + pub lead_agent_id: String, + pub status: AgentTeamStatus, + pub summary: Option, + pub created_at: Option>, + pub closed_at: Option>, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentTeamMember { + pub id: String, + pub team_id: String, + pub name: String, + pub agent_id: Option, + pub member_status: AgentTeamMemberStatus, + pub current_task_id: Option, + pub worker_thread_id: Option, + pub run_id: Option, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Debug, Clone)] +pub struct AgentTeamMemberUpsert { + pub id: String, + pub team_id: String, + pub name: String, + pub agent_id: Option, + pub member_status: AgentTeamMemberStatus, + pub current_task_id: Option, + pub worker_thread_id: Option, + pub run_id: Option, + pub created_at: Option>, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentTeamTask { + pub id: String, + pub team_id: String, + pub title: String, + pub objective: Option, + pub status: AgentTeamTaskStatus, + pub owner_member_id: Option, + pub claimed_by_member_id: Option, + pub claim_token: Option, + pub depends_on: Vec, + pub gate_status: String, + pub gate_reason: Option, + pub evidence: Vec, + pub source_run_id: Option, + pub order_index: i64, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Debug, Clone)] +pub struct AgentTeamTaskUpsert { + pub id: String, + pub team_id: String, + pub title: String, + pub objective: Option, + pub status: AgentTeamTaskStatus, + pub owner_member_id: Option, + pub depends_on: Vec, + pub gate_status: Option, + pub gate_reason: Option, + pub evidence: Vec, + pub source_run_id: Option, + pub order_index: i64, + pub created_at: Option>, +} + +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentTeamListRequest { + #[serde(default)] + pub parent_thread_id: Option, + #[serde(default)] + pub status: Option, + /// `u64` to match the `TypeSchema::U64` the controller advertises (the RPC + /// scalar-coercion layer only handles `U64`). Capped at 500 in + /// `list_agent_teams`. + #[serde(default)] + pub limit: Option, + #[serde(default)] + pub offset: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentTeamListResponse { + pub teams: Vec, + pub count: usize, +} + +/// Outcome of an atomic claim attempt on a team task. +#[derive(Debug, Clone, PartialEq, Serialize)] +#[serde(rename_all = "camelCase", tag = "kind")] +pub enum ClaimOutcome { + /// The claim succeeded; carries the freshly-claimed task. Boxed to keep the + /// enum small (the task payload dwarfs the other variants). + Claimed(Box), + /// Another member already holds the claim. + AlreadyClaimed, + /// One or more dependency tasks are not yet `done`. + Blocked { unmet: Vec }, + /// No task matched the given team + task id. + UnknownTask, +} + +/// Outcome of a completion attempt on a team task. +/// +/// Completion gates a task's transition to `done` behind quality invariants +/// (dependencies done, claimer owns the task, evidence present when required). +/// A failed gate leaves the task `in_progress` with `gate_status = "failed"` +/// and the reasons recorded, so a teammate can fix and retry. +#[derive(Debug, Clone, PartialEq, Serialize)] +#[serde(rename_all = "camelCase", tag = "kind")] +pub enum CompletionOutcome { + /// The task passed its quality gate and is now `done`. Boxed to keep the + /// enum small (the task payload dwarfs the other variants). + Completed(Box), + /// One or more quality-gate invariants failed; carries human-readable + /// reasons for each unmet invariant. + GateFailed { reasons: Vec }, + /// The task is not claimed by the completing member, or is not in progress. + NotClaimed, + /// No task matched the given team + task id. + UnknownTask, +} diff --git a/src/session/store.rs b/src/session/store.rs new file mode 100644 index 0000000..74a61dc --- /dev/null +++ b/src/session/store.rs @@ -0,0 +1,181 @@ +use std::path::{Path, PathBuf}; + +use rusqlite::Connection; + +use super::context::StorageContext; +use crate::error::Result; + +/// Subdirectory of the workspace holding the session database. +const DB_SUBDIR: &str = "session_db"; +/// Database filename inside [`DB_SUBDIR`]. +const DB_FILE: &str = "sessions.db"; + +/// Resolves the session database path for a workspace root. +/// +/// Kept public so hosts can locate the file for backup, inspection, or +/// migration without reproducing the layout. +pub fn db_path(workspace_dir: &Path) -> PathBuf { + workspace_dir.join(DB_SUBDIR).join(DB_FILE) +} + +/// Opens the workspace's session database, applying schema migrations, and +/// runs `f` against the connection. +/// +/// A connection is opened per call rather than pooled: these operations are +/// short, infrequent relative to a run's model calls, and SQLite in WAL mode +/// handles concurrent readers without a shared handle to synchronize. +pub fn with_connection( + workspace_dir: &Path, + f: impl FnOnce(&Connection) -> Result, +) -> Result { + let db_path = db_path(workspace_dir); + if let Some(parent) = db_path.parent() { + std::fs::create_dir_all(parent).storage_context(&format!( + "failed to create session_db directory: {}", + parent.display() + ))?; + } + + let conn = Connection::open(&db_path) + .storage_context(&format!("failed to open session DB: {}", db_path.display()))?; + + init_schema(&conn)?; + f(&conn) +} + +/// Opens the session database and runs `f` inside a single **immediate** +/// write transaction, committing on `Ok` and rolling back on `Err`. +/// +/// [`with_connection`] hands out an autocommit connection: each statement +/// commits on its own, so a multi-statement read-then-write sequence has no +/// isolation at all. Any operation whose correctness depends on the state it +/// read still holding when it writes — a compare-and-swap claim, a gate that +/// checks dependencies before acting — must use this instead. +/// +/// `BEGIN IMMEDIATE` rather than the default deferred begin: it takes the +/// write lock up front, so two racing claims serialize at `BEGIN` instead of +/// discovering the conflict at COMMIT time and failing with `SQLITE_BUSY` +/// after one of them has already decided it won. +pub fn with_transaction( + workspace_dir: &Path, + f: impl FnOnce(&Connection) -> Result, +) -> Result { + with_connection(workspace_dir, |conn| { + conn.execute_batch("BEGIN IMMEDIATE") + .storage_context("begin session DB transaction")?; + match f(conn) { + Ok(value) => { + conn.execute_batch("COMMIT") + .storage_context("commit session DB transaction")?; + Ok(value) + } + Err(err) => { + // Roll back best-effort: the caller's error is the one worth + // reporting, and a failed rollback (connection already gone) + // must not mask it. + if let Err(rollback_err) = conn.execute_batch("ROLLBACK") { + tracing::warn!( + "[session] rollback after error failed: {rollback_err} (original: {err})" + ); + } + Err(err) + } + } + }) +} + +#[cfg(test)] +pub fn with_memory_connection(f: impl FnOnce(&Connection) -> Result) -> Result { + let conn = + Connection::open_in_memory().storage_context("failed to open in-memory session DB")?; + init_schema(&conn)?; + f(&conn) +} + +pub(super) fn init_schema(conn: &Connection) -> Result<()> { + conn.execute_batch( + "PRAGMA journal_mode = WAL; + PRAGMA foreign_keys = ON; + + CREATE TABLE IF NOT EXISTS sessions ( + id TEXT PRIMARY KEY, + agent_definition_id TEXT NOT NULL, + agent_definition_name TEXT NOT NULL, + session_key TEXT NOT NULL, + parent_session_id TEXT, + thread_id TEXT, + source_channel TEXT, + status TEXT NOT NULL DEFAULT 'running', + model TEXT, + turn_count INTEGER NOT NULL DEFAULT 0, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + cached_input_tokens INTEGER NOT NULL DEFAULT 0, + cost_usd REAL NOT NULL DEFAULT 0.0, + transcript_path TEXT, + started_at TEXT NOT NULL, + ended_at TEXT, + FOREIGN KEY (parent_session_id) REFERENCES sessions(id) ON DELETE SET NULL + ); + CREATE INDEX IF NOT EXISTS idx_sessions_agent ON sessions(agent_definition_id); + CREATE INDEX IF NOT EXISTS idx_sessions_status ON sessions(status); + CREATE INDEX IF NOT EXISTS idx_sessions_started ON sessions(started_at); + CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id); + CREATE INDEX IF NOT EXISTS idx_sessions_thread ON sessions(thread_id); + CREATE INDEX IF NOT EXISTS idx_sessions_channel ON sessions(source_channel); + CREATE INDEX IF NOT EXISTS idx_sessions_key ON sessions(session_key); + + CREATE TABLE IF NOT EXISTS session_messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + role TEXT NOT NULL, + content TEXT NOT NULL, + model TEXT, + input_tokens INTEGER, + output_tokens INTEGER, + cost_usd REAL, + created_at TEXT NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS idx_messages_session ON session_messages(session_id); + + CREATE TABLE IF NOT EXISTS session_tool_calls ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + message_id INTEGER, + tool_name TEXT NOT NULL, + tool_input TEXT, + tool_output TEXT, + status TEXT NOT NULL DEFAULT 'pending', + duration_ms INTEGER, + created_at TEXT NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE, + FOREIGN KEY (message_id) REFERENCES session_messages(id) ON DELETE SET NULL + ); + CREATE INDEX IF NOT EXISTS idx_tool_calls_session ON session_tool_calls(session_id); + CREATE INDEX IF NOT EXISTS idx_tool_calls_name ON session_tool_calls(tool_name);", + ) + .storage_context("failed to initialize session_db schema")?; + + init_fts(conn)?; + Ok(()) +} + +fn init_fts(conn: &Connection) -> Result<()> { + let has_fts: bool = conn + .prepare("SELECT 1 FROM sqlite_master WHERE type='table' AND name='sessions_fts'")? + .exists([])?; + + if !has_fts { + conn.execute_batch( + "CREATE VIRTUAL TABLE sessions_fts USING fts5( + session_id, + agent_definition_name, + content, + tool_name + );", + ) + .storage_context("failed to create sessions_fts virtual table")?; + } + Ok(()) +} diff --git a/src/session/test.rs b/src/session/test.rs new file mode 100644 index 0000000..ab470b7 --- /dev/null +++ b/src/session/test.rs @@ -0,0 +1,618 @@ +//! Module-local unit tests for [`crate::session`]. +//! +//! Consolidated here per AGENTS.md: one `test.rs` per module directory rather +//! than per-file inline `mod tests` blocks. Sections mirror the source files. + +use super::context::StorageContext; +use super::ops::*; +use super::store::{init_schema, with_memory_connection}; +use super::types::*; +use crate::error::TinyAgentsError; +use chrono::Utc; +use rusqlite::{Connection, params}; + +// ── ops.rs ────────────────────────────────────────────────────────────── +fn insert_test_session(conn: &Connection, id: &str, agent_id: &str, key: &str) { + let now = Utc::now(); + conn.execute( + "INSERT INTO sessions ( + id, agent_definition_id, agent_definition_name, session_key, + status, started_at + ) VALUES (?1, ?2, ?3, ?4, 'running', ?5)", + params![id, agent_id, agent_id, key, now.to_rfc3339()], + ) + .unwrap(); + index_fts_session(conn, id, agent_id).unwrap(); +} + +fn insert_test_session_with_parent( + conn: &Connection, + id: &str, + agent_id: &str, + key: &str, + parent_id: &str, +) { + let now = Utc::now(); + conn.execute( + "INSERT INTO sessions ( + id, agent_definition_id, agent_definition_name, session_key, + parent_session_id, status, started_at + ) VALUES (?1, ?2, ?3, ?4, ?5, 'running', ?6)", + params![id, agent_id, agent_id, key, parent_id, now.to_rfc3339()], + ) + .unwrap(); + index_fts_session(conn, id, agent_id).unwrap(); +} + +#[test] +fn map_session_row_roundtrip() { + with_memory_connection(|conn| { + insert_test_session(conn, "sess-1", "orchestrator", "1700000000_orchestrator"); + + let mut stmt = conn.prepare( + "SELECT id, agent_definition_id, agent_definition_name, session_key, + parent_session_id, thread_id, source_channel, status, model, + turn_count, input_tokens, output_tokens, cached_input_tokens, + cost_usd, transcript_path, started_at, ended_at + FROM sessions WHERE id = 'sess-1'", + )?; + let session = stmt.query_row([], map_session_row)?; + + assert_eq!(session.id, "sess-1"); + assert_eq!(session.agent_definition_id, "orchestrator"); + assert_eq!(session.session_key, "1700000000_orchestrator"); + assert_eq!(session.status, SessionStatus::Running); + assert!(session.parent_session_id.is_none()); + assert!(session.ended_at.is_none()); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn search_by_agent_id() { + with_memory_connection(|conn| { + insert_test_session(conn, "a1", "orchestrator", "key1"); + insert_test_session(conn, "a2", "researcher", "key2"); + insert_test_session(conn, "a3", "orchestrator", "key3"); + + let params = SessionSearchParams { + agent_id: Some("orchestrator".to_string()), + ..Default::default() + }; + + let result = search_sessions_inner(conn, ¶ms)?; + assert_eq!(result.total, 2); + assert_eq!(result.sessions.len(), 2); + assert!( + result + .sessions + .iter() + .all(|s| s.agent_definition_id == "orchestrator") + ); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn search_by_fts_query() { + with_memory_connection(|conn| { + insert_test_session(conn, "b1", "orchestrator", "key1"); + insert_test_session(conn, "b2", "researcher", "key2"); + + conn.execute( + "INSERT INTO session_messages (session_id, role, content, created_at) + VALUES ('b1', 'user', 'Fix the login bug in authentication', ?1)", + params![Utc::now().to_rfc3339()], + )?; + index_fts_content(conn, "b1", "Fix the login bug in authentication")?; + + conn.execute( + "INSERT INTO session_messages (session_id, role, content, created_at) + VALUES ('b2', 'user', 'Deploy the new feature to production', ?1)", + params![Utc::now().to_rfc3339()], + )?; + index_fts_content(conn, "b2", "Deploy the new feature to production")?; + + let params = SessionSearchParams { + query: Some("login".to_string()), + ..Default::default() + }; + + let result = search_sessions_inner(conn, ¶ms)?; + assert_eq!(result.total, 1); + assert_eq!(result.sessions[0].id, "b1"); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn search_by_tool_name() { + with_memory_connection(|conn| { + insert_test_session(conn, "c1", "orchestrator", "key1"); + insert_test_session(conn, "c2", "researcher", "key2"); + + conn.execute( + "INSERT INTO session_tool_calls (session_id, tool_name, status, created_at) + VALUES ('c1', 'shell', 'ok', ?1)", + params![Utc::now().to_rfc3339()], + )?; + conn.execute( + "INSERT INTO session_tool_calls (session_id, tool_name, status, created_at) + VALUES ('c2', 'file_read', 'ok', ?1)", + params![Utc::now().to_rfc3339()], + )?; + + let params = SessionSearchParams { + tool_name: Some("shell".to_string()), + ..Default::default() + }; + + let result = search_sessions_inner(conn, ¶ms)?; + assert_eq!(result.total, 1); + assert_eq!(result.sessions[0].id, "c1"); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn search_by_parent_session() { + with_memory_connection(|conn| { + insert_test_session(conn, "parent-1", "orchestrator", "key1"); + insert_test_session_with_parent(conn, "child-1", "researcher", "key2", "parent-1"); + insert_test_session_with_parent(conn, "child-2", "coder", "key3", "parent-1"); + insert_test_session(conn, "unrelated", "other", "key4"); + + let params = SessionSearchParams { + parent_session_id: Some("parent-1".to_string()), + ..Default::default() + }; + + let result = search_sessions_inner(conn, ¶ms)?; + assert_eq!(result.total, 2); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn search_pagination() { + with_memory_connection(|conn| { + for i in 0..10 { + insert_test_session(conn, &format!("p{i}"), "agent", &format!("key{i}")); + } + + let params = SessionSearchParams { + limit: Some(3), + offset: Some(0), + ..Default::default() + }; + let result = search_sessions_inner(conn, ¶ms)?; + assert_eq!(result.total, 10); + assert_eq!(result.sessions.len(), 3); + + let params2 = SessionSearchParams { + limit: Some(3), + offset: Some(3), + ..Default::default() + }; + let result2 = search_sessions_inner(conn, ¶ms2)?; + assert_eq!(result2.total, 10); + assert_eq!(result2.sessions.len(), 3); + assert_ne!(result.sessions[0].id, result2.sessions[0].id); + + Ok(()) + }) + .unwrap(); +} + +#[test] +fn search_empty_results() { + with_memory_connection(|conn| { + let params = SessionSearchParams { + agent_id: Some("nonexistent".to_string()), + ..Default::default() + }; + let result = search_sessions_inner(conn, ¶ms)?; + assert_eq!(result.total, 0); + assert!(result.sessions.is_empty()); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn tool_output_truncation() { + with_memory_connection(|conn| { + let session_id = "trunc-sess"; + insert_test_session(conn, session_id, "agent", "key"); + + let large_output = "x".repeat(MAX_TOOL_OUTPUT_BYTES + 1000); + let bounded = if large_output.len() <= MAX_TOOL_OUTPUT_BYTES { + large_output.clone() + } else { + let mut cutoff = MAX_TOOL_OUTPUT_BYTES; + while cutoff > 0 && !large_output.is_char_boundary(cutoff) { + cutoff -= 1; + } + let mut truncated = large_output[..cutoff].to_string(); + truncated.push_str("\n...[truncated]"); + truncated + }; + + conn.execute( + "INSERT INTO session_tool_calls (session_id, tool_name, tool_output, status, created_at) + VALUES (?1, 'test', ?2, 'ok', ?3)", + params![session_id, bounded, Utc::now().to_rfc3339()], + )?; + + let stored: String = conn.query_row( + "SELECT tool_output FROM session_tool_calls WHERE session_id = ?1", + params![session_id], + |r| r.get(0), + )?; + assert!(stored.len() <= MAX_TOOL_OUTPUT_BYTES + 20); + assert!(stored.ends_with("[truncated]")); + + Ok(()) + }) + .unwrap(); +} + +#[test] +fn mark_interrupted_updates_running() { + with_memory_connection(|conn| { + insert_test_session(conn, "run1", "agent", "key1"); + insert_test_session(conn, "run2", "agent", "key2"); + conn.execute( + "UPDATE sessions SET status = 'completed' WHERE id = 'run2'", + [], + )?; + + let now = Utc::now(); + let changed = conn.execute( + "UPDATE sessions SET status = 'interrupted', ended_at = ?1 + WHERE status = 'running'", + params![now.to_rfc3339()], + )?; + assert_eq!(changed, 1); + + let status: String = + conn.query_row("SELECT status FROM sessions WHERE id = 'run1'", [], |r| { + r.get(0) + })?; + assert_eq!(status, "interrupted"); + + let status2: String = + conn.query_row("SELECT status FROM sessions WHERE id = 'run2'", [], |r| { + r.get(0) + })?; + assert_eq!(status2, "completed"); + + Ok(()) + }) + .unwrap(); +} + +#[test] +fn session_end_updates_cost_fields() { + with_memory_connection(|conn| { + insert_test_session(conn, "cost-sess", "agent", "key"); + + let now = Utc::now(); + conn.execute( + "UPDATE sessions SET + status = 'completed', turn_count = 5, input_tokens = 10000, + output_tokens = 2000, cached_input_tokens = 8000, + cost_usd = 0.0345, ended_at = ?1 + WHERE id = 'cost-sess'", + params![now.to_rfc3339()], + )?; + + let mut stmt = conn.prepare( + "SELECT id, agent_definition_id, agent_definition_name, session_key, + parent_session_id, thread_id, source_channel, status, model, + turn_count, input_tokens, output_tokens, cached_input_tokens, + cost_usd, transcript_path, started_at, ended_at + FROM sessions WHERE id = 'cost-sess'", + )?; + let session = stmt.query_row([], map_session_row)?; + + assert_eq!(session.status, SessionStatus::Completed); + assert_eq!(session.turn_count, 5); + assert_eq!(session.input_tokens, 10000); + assert_eq!(session.output_tokens, 2000); + assert_eq!(session.cached_input_tokens, 8000); + assert!((session.cost_usd - 0.0345).abs() < f64::EPSILON); + assert!(session.ended_at.is_some()); + + Ok(()) + }) + .unwrap(); +} + +#[test] +fn combined_filters() { + with_memory_connection(|conn| { + insert_test_session(conn, "cf1", "orchestrator", "key1"); + insert_test_session(conn, "cf2", "orchestrator", "key2"); + insert_test_session(conn, "cf3", "researcher", "key3"); + + conn.execute( + "UPDATE sessions SET status = 'completed' WHERE id = 'cf1'", + [], + )?; + + let params = SessionSearchParams { + agent_id: Some("orchestrator".to_string()), + status: Some("completed".to_string()), + ..Default::default() + }; + + let result = search_sessions_inner(conn, ¶ms)?; + assert_eq!(result.total, 1); + assert_eq!(result.sessions[0].id, "cf1"); + Ok(()) + }) + .unwrap(); +} + +// ── Regressions for the review findings on PR #90 ───────────────────────── + +/// A long non-ASCII message must not panic while being indexed. +/// +/// `&content[..2000]` panicked whenever byte 2000 landed inside a multi-byte +/// character. The message INSERT autocommits before the FTS write, so the panic +/// left a stored message with no FTS row — permanently unsearchable. +#[test] +fn long_multibyte_message_is_indexed_without_panicking() { + with_memory_connection(|conn| { + // '€' is 3 bytes and 2000 is not a multiple of 3, so byte 2000 lands + // strictly inside a character — the case that panicked. A 2-byte char + // would leave 2000 on a boundary and pass vacuously. + let content = "€".repeat(1500); + assert!(content.len() > 2000); + assert!(!content.is_char_boundary(2000)); + + index_fts_content(conn, "sess-utf8", &content)?; + + let indexed: i64 = conn.query_row( + "SELECT COUNT(*) FROM sessions_fts WHERE session_id = 'sess-utf8'", + [], + |r| r.get(0), + )?; + assert_eq!(indexed, 1, "the message must still get an FTS row"); + Ok(()) + }) + .unwrap(); +} + +/// An ASCII message longer than the cap still truncates to exactly the cap. +#[test] +fn long_ascii_message_truncates_at_the_byte_cap() { + with_memory_connection(|conn| { + let content = "a".repeat(5000); + index_fts_content(conn, "sess-ascii", &content)?; + let stored: String = conn.query_row( + "SELECT content FROM sessions_fts WHERE session_id = 'sess-ascii'", + [], + |r| r.get(0), + )?; + assert_eq!(stored.len(), 2000); + Ok(()) + }) + .unwrap(); +} + +/// Punctuation that is meaningful to FTS5 must be searched as literal text. +/// +/// `SessionSearchParams::query` is plain text, not an FTS5 expression. Binding +/// it raw made ordinary input (`C++`, `foo-bar`, `file.rs`, a stray quote) +/// return a syntax or `no such column` error instead of results. +#[test] +fn fts_query_treats_punctuation_as_literal_text() { + for raw in ["C++", "foo-bar", "file.rs", "a\"b", "NOT", "*"] { + with_memory_connection(|conn| { + conn.execute( + "INSERT INTO sessions_fts (session_id, agent_definition_name, content, tool_name) + VALUES ('s1', '', ?1, '')", + params![raw], + )?; + let params = SessionSearchParams { + query: Some(raw.to_string()), + ..Default::default() + }; + // The assertion is that this does not error; FTS tokenization + // decides whether a given punctuation string is recallable. + search_sessions_inner(conn, ¶ms) + .unwrap_or_else(|e| panic!("query {raw:?} must not error: {e}")); + Ok(()) + }) + .unwrap(); + } +} + +/// Multi-term queries stay conjunctive, as a search box implies. +#[test] +fn fts_query_joins_terms_with_and() { + assert_eq!(fts_match_query("alpha beta"), "\"alpha\" AND \"beta\""); + assert_eq!(fts_match_query("C++"), "\"C++\""); + // Embedded quotes are escaped by doubling, per the FTS5 string grammar. + assert_eq!(fts_match_query("a\"b"), "\"a\"\"b\""); +} + +// ── types.rs ─────────────────────────────────────────────────────────── + +#[test] +fn session_status_roundtrip() { + for status in [ + SessionStatus::Running, + SessionStatus::Completed, + SessionStatus::Failed, + SessionStatus::Interrupted, + ] { + assert_eq!(SessionStatus::parse(status.as_str()), status); + } +} + +#[test] +fn session_status_parse_unknown_defaults_to_running() { + assert_eq!(SessionStatus::parse("bogus"), SessionStatus::Running); + assert_eq!(SessionStatus::parse(""), SessionStatus::Running); +} + +#[test] +fn session_status_serde_roundtrip() { + let status = SessionStatus::Completed; + let json = serde_json::to_string(&status).unwrap(); + assert_eq!(json, "\"completed\""); + let parsed: SessionStatus = serde_json::from_str(&json).unwrap(); + assert_eq!(parsed, status); +} + +#[test] +fn session_search_params_defaults() { + let params = SessionSearchParams::default(); + assert!(params.query.is_none()); + assert!(params.agent_id.is_none()); + assert!(params.limit.is_none()); + assert!(params.offset.is_none()); +} + +// ── store.rs ─────────────────────────────────────────────────────────── + +#[test] +fn schema_initializes_without_error() { + with_memory_connection(|conn| { + let count: i64 = conn.query_row("SELECT COUNT(*) FROM sessions", [], |r| r.get(0))?; + assert_eq!(count, 0); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn schema_is_idempotent() { + let conn = Connection::open_in_memory().unwrap(); + init_schema(&conn).unwrap(); + init_schema(&conn).unwrap(); + let count: i64 = conn + .query_row("SELECT COUNT(*) FROM sessions", [], |r| r.get(0)) + .unwrap(); + assert_eq!(count, 0); +} + +#[test] +fn wal_mode_is_set() { + with_memory_connection(|conn| { + let mode: String = conn.query_row("PRAGMA journal_mode", [], |r| r.get(0))?; + // In-memory DBs may report "memory" instead of "wal" + assert!(mode == "wal" || mode == "memory"); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn fts_table_exists_after_init() { + with_memory_connection(|conn| { + let exists: bool = conn + .prepare("SELECT 1 FROM sqlite_master WHERE type='table' AND name='sessions_fts'")? + .exists([])?; + assert!(exists); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn foreign_keys_are_enabled() { + with_memory_connection(|conn| { + let fk: i64 = conn.query_row("PRAGMA foreign_keys", [], |r| r.get(0))?; + assert_eq!(fk, 1); + Ok(()) + }) + .unwrap(); +} + +// ── context.rs ─────────────────────────────────────────────────────────── + +#[test] +fn result_error_is_prefixed_with_context() { + let failed: std::result::Result<(), _> = Err("disk full"); + let err = failed.storage_context("write session").unwrap_err(); + assert!(matches!(err, TinyAgentsError::Storage(_))); + assert_eq!(err.to_string(), "storage error: write session: disk full"); +} + +#[test] +fn result_ok_passes_through() { + let ok: std::result::Result = Ok(7); + assert_eq!(ok.storage_context("read").unwrap(), 7); +} + +#[test] +fn none_becomes_storage_error_without_a_source_suffix() { + let absent: Option = None; + let err = absent + .storage_context("run missing after upsert") + .unwrap_err(); + assert_eq!(err.to_string(), "storage error: run missing after upsert"); +} + +#[test] +fn some_passes_through() { + assert_eq!(Some(3).storage_context("read").unwrap(), 3); +} + +/// `record_tool_call` must return the `session_tool_calls` row id, not the +/// rowid of the FTS row written immediately afterwards. +/// +/// The FTS insert moves `last_insert_rowid()`, so reading it after indexing +/// handed callers an id for a tool call that does not exist. The session row is +/// already in `sessions_fts` in any real session, which is what makes the two +/// counters diverge. +#[test] +fn record_tool_call_returns_the_tool_call_row_id() { + with_memory_connection(|conn| { + // Seed an FTS row first, as a real session always would, so the FTS + // rowid counter is ahead of the tool-call one. + index_fts_session(conn, "s1", "agent")?; + + conn.execute( + "INSERT INTO sessions ( + id, agent_definition_id, agent_definition_name, session_key, started_at + ) VALUES ('s1', 'a', 'agent', 's1', ?1)", + params![Utc::now().to_rfc3339()], + )?; + conn.execute( + "INSERT INTO session_tool_calls (session_id, tool_name, status, created_at) + VALUES ('s1', 'echo', 'ok', ?1)", + params![Utc::now().to_rfc3339()], + )?; + let expected = conn.last_insert_rowid(); + index_fts_tool(conn, "s1", "echo")?; + + // The FTS insert must have moved the connection's rowid... + assert_ne!( + conn.last_insert_rowid(), + expected, + "test is vacuous unless the FTS insert moves last_insert_rowid()" + ); + // ...and the id we hand back must still address a real tool call. + let exists: i64 = conn.query_row( + "SELECT COUNT(*) FROM session_tool_calls WHERE id = ?1", + params![expected], + |r| r.get(0), + )?; + assert_eq!( + exists, 1, + "returned id must address a session_tool_calls row" + ); + Ok(()) + }) + .unwrap(); +} diff --git a/src/session/types.rs b/src/session/types.rs new file mode 100644 index 0000000..d1fa960 --- /dev/null +++ b/src/session/types.rs @@ -0,0 +1,107 @@ +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SessionStatus { + Running, + Completed, + Failed, + Interrupted, +} + +impl SessionStatus { + pub fn as_str(&self) -> &'static str { + match self { + Self::Running => "running", + Self::Completed => "completed", + Self::Failed => "failed", + Self::Interrupted => "interrupted", + } + } + + pub fn parse(s: &str) -> Self { + match s { + "completed" => Self::Completed, + "failed" => Self::Failed, + "interrupted" => Self::Interrupted, + _ => Self::Running, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SessionRecord { + pub id: String, + pub agent_definition_id: String, + pub agent_definition_name: String, + pub session_key: String, + pub parent_session_id: Option, + pub thread_id: Option, + pub source_channel: Option, + pub status: SessionStatus, + pub model: Option, + pub turn_count: u32, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_input_tokens: u64, + pub cost_usd: f64, + pub transcript_path: Option, + pub started_at: DateTime, + pub ended_at: Option>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SessionMessage { + pub id: i64, + pub session_id: String, + pub role: String, + pub content: String, + pub model: Option, + pub input_tokens: Option, + pub output_tokens: Option, + pub cost_usd: Option, + pub created_at: DateTime, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SessionToolCall { + pub id: i64, + pub session_id: String, + pub message_id: Option, + pub tool_name: String, + pub tool_input: Option, + pub tool_output: Option, + pub status: String, + pub duration_ms: Option, + pub created_at: DateTime, +} + +#[derive(Debug, Clone, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionSearchParams { + #[serde(default)] + pub query: Option, + #[serde(default)] + pub agent_id: Option, + #[serde(default)] + pub tool_name: Option, + #[serde(default)] + pub source_channel: Option, + #[serde(default)] + pub parent_session_id: Option, + #[serde(default)] + pub status: Option, + #[serde(default)] + pub thread_id: Option, + #[serde(default)] + pub limit: Option, + #[serde(default)] + pub offset: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SessionSearchResult { + pub sessions: Vec, + pub total: u64, +}