use crate::blueprint::BindingDigest;
use crate::store::replay::{ReplayCursor, ReplayStore};
use crate::types::{RunId, StepId, TaskId};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use thiserror::Error;
pub mod inmemory;
pub mod sqlite;
pub use inmemory::InMemoryRunStore;
pub use sqlite::SqliteRunStore;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum RunStatus {
Pending,
Running,
Done,
Failed,
Interrupted,
Cancelled,
}
#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
pub struct DegradationEntry {
pub tool: String,
pub error: String,
pub fallback: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub note: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub step_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub attempt: Option<u32>,
pub at: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
pub struct StepEntry {
#[schemars(with = "String")]
pub step_id: StepId,
pub step_ref: Option<String>,
pub status: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub binding_digest: Option<BindingDigest>,
pub at: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub attempt: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub started_at_ms: Option<i64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub completed_at_ms: Option<i64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub duration_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub worker_kind: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub usage: Option<crate::store::trace::TokenUsage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub num_turns: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub adapter_data: Option<serde_json::Value>,
}
impl StepEntry {
pub fn basic(
step_id: StepId,
step_ref: Option<String>,
status: Option<String>,
binding_digest: Option<BindingDigest>,
at: u64,
) -> Self {
Self {
step_id,
step_ref,
status,
binding_digest,
at,
attempt: None,
started_at_ms: None,
completed_at_ms: None,
duration_ms: None,
worker_kind: None,
model: None,
usage: None,
num_turns: None,
adapter_data: None,
}
}
pub fn with_worker_stats(mut self, stats: crate::store::trace::WorkerStats) -> Self {
self.worker_kind = stats.worker_kind.or(self.worker_kind);
self.model = stats.model.or(self.model);
self.usage = stats.usage.or(self.usage);
self.num_turns = stats.num_turns.or(self.num_turns);
self.adapter_data = stats
.adapter_data
.map(crate::store::trace::cap_payload)
.or(self.adapter_data);
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
pub struct RunRecord {
#[schemars(with = "String")]
pub id: RunId,
#[schemars(with = "String")]
pub task_id: TaskId,
pub status: RunStatus,
pub step_entries: Vec<StepEntry>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub degradations: Vec<DegradationEntry>,
pub operator_sid: Option<String>,
#[schemars(with = "Option<serde_json::Value>")]
pub result_ref: Option<serde_json::Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input_json: Option<String>,
pub created_at: u64,
pub updated_at: u64,
}
#[derive(Debug, Clone, Default)]
pub struct RunListFilter {
pub task_id: Option<TaskId>,
pub status: Option<RunStatus>,
pub limit: Option<usize>,
pub offset: Option<usize>,
}
#[derive(Debug, Error)]
pub enum RunStoreError {
#[error("run not found: {0}")]
NotFound(RunId),
#[error("run already exists: {0}")]
Duplicate(RunId),
#[error("other: {0}")]
Other(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum SnapshotOrigin {
Launch,
ResumeBackfill,
}
pub const BOUND_AGENTS_ORIGIN_KEY: &str = "bound_agents_origin";
impl SnapshotOrigin {
pub fn from_snapshot(snapshot: &serde_json::Value) -> Self {
snapshot
.get(BOUND_AGENTS_ORIGIN_KEY)
.and_then(|v| serde_json::from_value::<SnapshotOrigin>(v.clone()).ok())
.unwrap_or(SnapshotOrigin::ResumeBackfill)
}
}
#[derive(Debug, Clone)]
pub struct LastFailure {
pub step_id: StepId,
pub step_ref: Option<String>,
pub verdict_value: serde_json::Value,
}
#[derive(Clone)]
pub struct RunContext {
pub run_id: RunId,
pub run_store: Arc<dyn RunStore>,
pub replay_store: Option<Arc<dyn ReplayStore>>,
pub replay_cursor: Option<Arc<Mutex<ReplayCursor>>>,
pub binding_digests: Arc<HashMap<String, BindingDigest>>,
pub resume: bool,
pub last_failure: Arc<Mutex<Option<LastFailure>>>,
pub trace: Option<crate::store::trace::TraceHandle>,
}
impl RunContext {
pub fn new(run_id: RunId, run_store: Arc<dyn RunStore>) -> Self {
Self {
run_id,
run_store,
replay_store: None,
replay_cursor: None,
binding_digests: Arc::new(HashMap::new()),
resume: false,
last_failure: Arc::new(Mutex::new(None)),
trace: None,
}
}
pub fn with_trace(mut self, trace: crate::store::trace::TraceHandle) -> Self {
self.trace = Some(trace);
self
}
pub fn set_last_failure(&self, failure: LastFailure) {
if let Ok(mut slot) = self.last_failure.lock() {
*slot = Some(failure);
}
}
pub async fn snapshot_partial_ctx(&self) -> serde_json::Value {
let record = match self.run_store.get(&self.run_id).await {
Ok(r) => r,
Err(_) => return serde_json::Value::Null,
};
let mut steps = serde_json::Map::new();
for entry in &record.step_entries {
let mut fields = serde_json::Map::new();
if let Some(ref_) = &entry.step_ref {
fields.insert(
"step_ref".to_string(),
serde_json::Value::String(ref_.clone()),
);
}
if let Some(status) = &entry.status {
fields.insert(
"status".to_string(),
serde_json::Value::String(status.clone()),
);
}
if let Some(digest) = &entry.binding_digest {
fields.insert(
"binding_digest".to_string(),
serde_json::Value::String(digest.to_string()),
);
}
fields.insert("at".to_string(), serde_json::Value::Number(entry.at.into()));
steps.insert(entry.step_id.to_string(), serde_json::Value::Object(fields));
}
let mut out = serde_json::Map::new();
out.insert("steps".to_string(), serde_json::Value::Object(steps));
serde_json::Value::Object(out)
}
pub fn with_replay_store(mut self, store: Arc<dyn ReplayStore>) -> Self {
self.replay_store = Some(store);
self
}
pub fn with_replay_cursor(mut self, cursor: Arc<Mutex<ReplayCursor>>) -> Self {
self.replay_cursor = Some(cursor);
self
}
pub fn with_binding_digests(mut self, digests: HashMap<String, BindingDigest>) -> Self {
self.binding_digests = Arc::new(digests);
self
}
pub fn with_resume(mut self) -> Self {
self.resume = true;
self
}
}
impl std::fmt::Debug for RunContext {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RunContext")
.field("run_id", &self.run_id)
.field("run_store", &self.run_store.name())
.field(
"replay_store",
&self.replay_store.as_ref().map(|s| s.name()),
)
.field("replay_cursor", &self.replay_cursor.is_some())
.field("binding_digests", &self.binding_digests.len())
.field("resume", &self.resume)
.field(
"last_failure",
&self.last_failure.lock().ok().and_then(|slot| slot.clone()),
)
.field("trace", &self.trace.is_some())
.finish()
}
}
#[async_trait]
pub trait RunStore: Send + Sync {
fn name(&self) -> &str;
async fn create(&self, record: RunRecord) -> Result<(), RunStoreError>;
async fn get(&self, id: &RunId) -> Result<RunRecord, RunStoreError>;
async fn list_by_task(&self, task_id: &TaskId) -> Result<Vec<RunRecord>, RunStoreError>;
async fn append_step_entry(&self, id: &RunId, entry: StepEntry) -> Result<(), RunStoreError>;
async fn append_degradation(
&self,
id: &RunId,
entry: DegradationEntry,
) -> Result<(), RunStoreError>;
async fn update_status(&self, id: &RunId, status: RunStatus) -> Result<(), RunStoreError>;
async fn try_transition(
&self,
id: &RunId,
from: RunStatus,
to: RunStatus,
) -> Result<bool, RunStoreError>;
async fn set_result(
&self,
id: &RunId,
result_ref: serde_json::Value,
) -> Result<(), RunStoreError>;
async fn set_input_json(&self, id: &RunId, input_json: String) -> Result<(), RunStoreError>;
async fn list_running(&self) -> Result<Vec<RunRecord>, RunStoreError>;
async fn list(&self, filter: &RunListFilter) -> Result<Vec<RunRecord>, RunStoreError>;
async fn delete(&self, id: &RunId) -> Result<(), RunStoreError>;
}
#[derive(Default)]
pub(crate) struct Inner {
pub(crate) order: Vec<RunId>,
pub(crate) records: HashMap<RunId, RunRecord>,
}
pub(crate) type SharedInner = Mutex<Inner>;