use rmcp::schemars;
use serde::{Deserialize, Serialize};
pub(crate) const NO_TURNS: u32 = 0;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub(crate) enum RunStatus {
Running,
Completed,
Failed,
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum Outcome {
Running,
Completed { value: String, turns: u32 },
Failed { error: String, turns: u32 },
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct RunResult {
run_id: String,
prompt: String,
elapsed_ms: u64,
outcome: Outcome,
}
impl RunResult {
pub(crate) fn completed(
run_id: String,
prompt: &str,
value: String,
turns: u32,
elapsed_ms: u64,
) -> RunResult {
RunResult {
run_id,
prompt: prompt.to_owned(),
elapsed_ms,
outcome: Outcome::Completed { value, turns },
}
}
pub(crate) fn failed(
run_id: String,
prompt: &str,
error: String,
turns: u32,
elapsed_ms: u64,
) -> RunResult {
RunResult {
run_id,
prompt: prompt.to_owned(),
elapsed_ms,
outcome: Outcome::Failed { error, turns },
}
}
pub(crate) fn running(run_id: String, prompt: &str, elapsed_ms: u64) -> RunResult {
RunResult {
run_id,
prompt: prompt.to_owned(),
elapsed_ms,
outcome: Outcome::Running,
}
}
pub(crate) fn run_id(&self) -> &str {
&self.run_id
}
pub(crate) fn prompt(&self) -> &str {
&self.prompt
}
pub(crate) fn elapsed_ms(&self) -> u64 {
self.elapsed_ms
}
pub(crate) fn status(&self) -> RunStatus {
match self.outcome {
Outcome::Running => RunStatus::Running,
Outcome::Completed { .. } => RunStatus::Completed,
Outcome::Failed { .. } => RunStatus::Failed,
}
}
pub(crate) fn turns(&self) -> u32 {
match &self.outcome {
Outcome::Running => NO_TURNS,
Outcome::Completed { turns, .. } | Outcome::Failed { turns, .. } => *turns,
}
}
pub(crate) fn value(&self) -> Option<&str> {
match &self.outcome {
Outcome::Completed { value, .. } => Some(value),
_ => None,
}
}
pub(crate) fn error(&self) -> Option<&str> {
match &self.outcome {
Outcome::Failed { error, .. } => Some(error),
_ => None,
}
}
pub(crate) fn text(&self) -> String {
match &self.outcome {
Outcome::Completed { value, .. } => value.clone(),
Outcome::Failed { error, .. } => error.clone(),
Outcome::Running => format!(
"{} is still running as run {}. Collect it with check_run.",
self.prompt, self.run_id
),
}
}
pub(crate) fn to_wire(&self) -> RunResultWire {
RunResultWire {
run_id: self.run_id.clone(),
prompt: self.prompt.clone(),
status: self.status(),
value: self.value().map(str::to_owned),
turns: self.turns(),
elapsed_ms: self.elapsed_ms,
error: self.error().map(str::to_owned),
}
}
}
#[derive(Debug, Clone, Serialize, schemars::JsonSchema)]
#[non_exhaustive]
pub(crate) struct RunResultWire {
pub(crate) run_id: String,
pub(crate) prompt: String,
pub(crate) status: RunStatus,
pub(crate) value: Option<String>,
pub(crate) turns: u32,
pub(crate) elapsed_ms: u64,
pub(crate) error: Option<String>,
}
#[cfg(test)]
mod tests {
use super::{RunResult, RunStatus};
#[test]
fn a_completed_run_texts_its_value() {
let result = RunResult::completed("r1".into(), "echo", "hello".into(), 2, 4);
assert_eq!(result.status(), RunStatus::Completed);
assert_eq!(result.text(), "hello");
assert_eq!(result.turns(), 2);
assert!(result.error().is_none());
}
#[test]
fn a_failed_run_texts_its_error() {
let result = RunResult::failed("r1".into(), "echo", "lua: boom".into(), 1, 4);
assert_eq!(result.status(), RunStatus::Failed);
assert_eq!(result.text(), "lua: boom");
assert_eq!(result.turns(), 1);
assert!(result.value().is_none());
}
#[test]
fn a_running_run_texts_how_to_collect_it() {
let result = RunResult::running("r1".into(), "echo", 240_000);
assert_eq!(result.status(), RunStatus::Running);
let text = result.text();
assert!(text.contains("r1"), "the id to collect by: {text}");
assert!(
text.contains("check_run"),
"the tool to collect with: {text}"
);
assert!(result.value().is_none());
assert!(result.error().is_none());
}
#[test]
fn equal_runs_compare_equal_and_differ_from_other_states() {
let a = RunResult::completed("r1".into(), "echo", "hello".into(), 2, 4);
let b = RunResult::completed("r1".into(), "echo", "hello".into(), 2, 4);
let failed = RunResult::failed("r1".into(), "echo", "boom".into(), 2, 4);
assert_eq!(a, b);
assert_ne!(a, failed);
}
#[test]
fn status_serializes_in_snake_case() {
let json = serde_json::to_string(&RunStatus::Running).expect("a unit enum serializes");
assert_eq!(json, "\"running\"");
}
#[test]
fn a_completed_wire_object_is_exactly_its_fields() {
let wire = RunResult::completed("r1".into(), "echo", "hello".into(), 2, 4).to_wire();
assert_eq!(
serde_json::to_value(&wire).expect("the wire form serializes"),
serde_json::json!({
"run_id": "r1",
"prompt": "echo",
"status": "completed",
"value": "hello",
"turns": 2,
"elapsed_ms": 4,
"error": null,
})
);
}
#[test]
fn a_failed_wire_object_carries_its_error_and_a_null_value() {
let wire = RunResult::failed("r1".into(), "echo", "lua: boom".into(), 1, 4).to_wire();
assert_eq!(
serde_json::to_value(&wire).expect("the wire form serializes"),
serde_json::json!({
"run_id": "r1",
"prompt": "echo",
"status": "failed",
"value": null,
"turns": 1,
"elapsed_ms": 4,
"error": "lua: boom",
})
);
}
#[test]
fn a_running_wire_object_reports_no_turns_and_no_payload() {
let wire = RunResult::running("r1".into(), "echo", 240_000).to_wire();
assert_eq!(
serde_json::to_value(&wire).expect("the wire form serializes"),
serde_json::json!({
"run_id": "r1",
"prompt": "echo",
"status": "running",
"value": null,
"turns": 0,
"elapsed_ms": 240_000,
"error": null,
})
);
}
}