use std::collections::BTreeMap;
use std::sync::Arc;
use salvor_core::{
Budget, Event, EventEnvelope, PendingCall, RunId, RunStatus, UnresolvedWrite, derive_state,
};
use salvor_store::EventStore;
use serde_json::Value;
use crate::agent::Agent;
use crate::budgets::validate_extension_input;
use crate::ctx::{ClockFn, RandomFn, RunCtx};
use crate::driver::{self, LoopOutcome};
use crate::error::RuntimeError;
use crate::validate::validate_against_schema;
#[derive(Debug, Clone)]
pub enum ParkReason {
Suspended {
reason: String,
input_schema: Value,
},
BudgetExceeded {
budget: Budget,
observed: f64,
},
}
#[derive(Debug, Clone)]
pub enum RunOutcome {
Completed {
run_id: RunId,
output: Value,
},
Parked {
run_id: RunId,
reason: ParkReason,
},
}
pub struct Runtime {
store: Arc<dyn EventStore>,
clock: ClockFn,
random: RandomFn,
record_prompts: bool,
labels: Option<BTreeMap<String, String>>,
}
impl Runtime {
#[must_use]
pub fn new(store: Arc<dyn EventStore>) -> Self {
Self::with_hooks(
store,
Arc::new(time::OffsetDateTime::now_utc),
Arc::new(crate::ctx::os_random),
)
}
#[must_use]
pub fn with_hooks(store: Arc<dyn EventStore>, clock: ClockFn, random: RandomFn) -> Self {
Self {
store,
clock,
random,
record_prompts: false,
labels: None,
}
}
#[must_use]
pub fn with_record_prompts(mut self, record_prompts: bool) -> Self {
self.record_prompts = record_prompts;
self
}
#[must_use]
pub fn with_labels(mut self, labels: BTreeMap<String, String>) -> Self {
self.labels = Some(labels);
self
}
pub async fn start(&self, agent: &Agent, input: Value) -> Result<RunOutcome, RuntimeError> {
self.start_with_id(agent, RunId::new(), input).await
}
pub async fn start_with_id(
&self,
agent: &Agent,
run_id: RunId,
input: Value,
) -> Result<RunOutcome, RuntimeError> {
let log = self.store.read_log(run_id).await?;
if !log.is_empty() {
return Err(RuntimeError::RunAlreadyStarted { run_id });
}
let mut ctx = self.ctx(run_id, log)?;
finish(run_id, driver::drive(&mut ctx, agent, &input).await?)
}
pub async fn recover(&self, agent: &Agent, run_id: RunId) -> Result<RunOutcome, RuntimeError> {
let log = self.read_existing(run_id).await?;
let mut ctx = self.ctx(run_id, log)?;
finish(run_id, driver::drive(&mut ctx, agent, &Value::Null).await?)
}
pub async fn resume(
&self,
agent: &Agent,
run_id: RunId,
input: Value,
) -> Result<RunOutcome, RuntimeError> {
let log = self.read_existing(run_id).await?;
let state = derive_state(&log);
match &state.status {
RunStatus::Suspended { input_schema, .. } => {
validate_against_schema(&input, input_schema)
.map_err(RuntimeError::ResumeInputRejected)?;
}
RunStatus::BudgetExceeded { .. } => {
validate_extension_input(&input).map_err(RuntimeError::ResumeInputRejected)?;
}
other => {
return Err(RuntimeError::NotParked {
run_id,
status: status_name(other).to_owned(),
});
}
}
let mut ctx = self.ctx(run_id, log)?;
ctx.set_resume_input(input);
finish(run_id, driver::drive(&mut ctx, agent, &Value::Null).await?)
}
pub async fn resolve(&self, run_id: RunId, output: Value) -> Result<RunId, RuntimeError> {
let log = self.read_existing(run_id).await?;
let state = derive_state(&log);
let intent_seq = match (&state.status, &state.pending_call) {
(RunStatus::NeedsReconciliation, Some(PendingCall::Tool { seq, .. })) => *seq,
(other, _) => {
return Err(RuntimeError::NotReconcilable {
run_id,
status: status_name(other).to_owned(),
});
}
};
let completion = Event::ToolCallCompleted {
seq: intent_seq,
output,
};
let envelope = EventEnvelope::new(run_id, state.next_seq, (self.clock)(), completion);
self.store.append(&envelope).await?;
crate::progress::emit_step(run_id, envelope.seq, &envelope.event);
Ok(run_id)
}
pub async fn abandon(
&self,
run_id: RunId,
reason: Option<String>,
) -> Result<RunId, RuntimeError> {
let log = self.read_existing(run_id).await?;
let state = derive_state(&log);
if matches!(
state.status,
RunStatus::Completed { .. } | RunStatus::Failed { .. } | RunStatus::Abandoned { .. }
) {
return Err(RuntimeError::AlreadyTerminal {
run_id,
status: status_name(&state.status).to_owned(),
});
}
let unresolved_write = match (&state.status, &state.pending_call) {
(RunStatus::NeedsReconciliation, Some(PendingCall::Tool { seq, tool, .. })) => {
Some(UnresolvedWrite {
seq: *seq,
tool: tool.clone(),
})
}
_ => None,
};
let event = Event::RunAbandoned {
reason,
unresolved_write,
};
let envelope = EventEnvelope::new(run_id, state.next_seq, (self.clock)(), event);
self.store.append(&envelope).await?;
crate::progress::emit_step(run_id, envelope.seq, &envelope.event);
Ok(run_id)
}
async fn read_existing(&self, run_id: RunId) -> Result<Vec<EventEnvelope>, RuntimeError> {
let log = self.store.read_log(run_id).await?;
if log.is_empty() {
return Err(RuntimeError::UnknownRun { run_id });
}
Ok(log)
}
fn ctx(&self, run_id: RunId, log: Vec<EventEnvelope>) -> Result<RunCtx, RuntimeError> {
let mut ctx = RunCtx::with_hooks(
self.store.clone(),
run_id,
log,
self.clock.clone(),
self.random.clone(),
)?
.with_record_prompts(self.record_prompts);
if let Some(labels) = &self.labels {
ctx = ctx.with_labels(labels.clone());
}
Ok(ctx)
}
}
#[allow(clippy::unnecessary_wraps)]
fn finish(run_id: RunId, outcome: LoopOutcome) -> Result<RunOutcome, RuntimeError> {
Ok(match outcome {
LoopOutcome::Completed(output) => RunOutcome::Completed { run_id, output },
LoopOutcome::Parked(reason) => RunOutcome::Parked { run_id, reason },
})
}
fn status_name(status: &RunStatus) -> &'static str {
match status {
RunStatus::NotStarted => "not started",
RunStatus::Running => "running",
RunStatus::AwaitingModel => "awaiting model (interrupted; use recover)",
RunStatus::AwaitingTool => "awaiting tool (interrupted; use recover)",
RunStatus::Suspended { .. } => "suspended",
RunStatus::BudgetExceeded { .. } => "budget exceeded",
RunStatus::NeedsReconciliation => "needs reconciliation",
RunStatus::Completed { .. } => "completed",
RunStatus::Failed { .. } => "failed",
RunStatus::Abandoned { .. } => "abandoned",
}
}