use std::collections::BTreeMap;
use serde::de::{self, MapAccess, SeqAccess, Visitor};
use serde::Deserialize;
use serde_json::{Map, Number, Value};
use crate::budget::BudgetUsageSnapshot;
use crate::checkpoint::{
canonical_json_bytes, CheckpointError, CheckpointResult, CheckpointStatus, CHECKPOINT_SCHEMA,
RUN_DEFINITION_SCHEMA,
};
use crate::runtime::state::{
validate_checkpoint, validate_extension_state_size, Checkpoint, EventOutboxEntry,
ExtensionStateEntry, OperationJournalEntry,
};
use crate::types::{CycleRecord, Message, ModelCallRecord};
const KNOWN_FIELDS: &[&str] = &[
"schema_version",
"run_definition_schema",
"run_definition",
"checkpoint_key",
"task_id",
"root_run_id",
"trace_id",
"run_definition_digest",
"resume_attempt",
"cycle_index",
"status",
"messages",
"cycles",
"model_calls",
"shared_state",
"budget_usage",
"event_cursor",
"event_outbox",
"extension_state",
"model_call_journal",
"tool_journal",
"revision",
"claim_token",
"claimed_cycle",
"lease_expires_at_ms",
"terminal_result",
"terminal_acknowledged",
];
pub fn checkpoint_to_value(
checkpoint: &Checkpoint,
max_extension_state_bytes: u64,
) -> CheckpointResult<Value> {
validate_checkpoint(checkpoint)?;
let extension_bytes =
validate_extension_state_size(&checkpoint.extension_state, max_extension_state_bytes);
extension_bytes?;
let mut object = Map::new();
object.insert(
"schema_version".to_string(),
Value::String(CHECKPOINT_SCHEMA.to_string()),
);
object.insert(
"run_definition_schema".to_string(),
Value::String(RUN_DEFINITION_SCHEMA.to_string()),
);
object.insert(
"run_definition".to_string(),
checkpoint.run_definition.clone(),
);
object.insert(
"checkpoint_key".to_string(),
Value::String(checkpoint.checkpoint_key.clone()),
);
object.insert(
"task_id".to_string(),
Value::String(checkpoint.task_id.clone()),
);
object.insert(
"root_run_id".to_string(),
Value::String(checkpoint.root_run_id.clone()),
);
object.insert(
"trace_id".to_string(),
Value::String(checkpoint.trace_id.clone()),
);
object.insert(
"run_definition_digest".to_string(),
Value::String(checkpoint.run_definition_digest.clone()),
);
object.insert(
"resume_attempt".to_string(),
Value::from(checkpoint.resume_attempt),
);
object.insert(
"cycle_index".to_string(),
Value::from(checkpoint.cycle_index),
);
object.insert(
"status".to_string(),
Value::String(checkpoint.status.as_str().to_string()),
);
object.insert(
"messages".to_string(),
Value::Array(checkpoint.messages.iter().map(Message::to_dict).collect()),
);
object.insert(
"cycles".to_string(),
Value::Array(checkpoint.cycles.iter().map(CycleRecord::to_dict).collect()),
);
object.insert(
"model_calls".to_string(),
serde_json::to_value(&checkpoint.model_calls).map_err(|error| {
CheckpointError::new("checkpoint_model_calls_invalid", error.to_string())
})?,
);
object.insert(
"shared_state".to_string(),
Value::Object(checkpoint.shared_state.clone().into_iter().collect()),
);
object.insert(
"budget_usage".to_string(),
checkpoint
.budget_usage
.as_ref()
.map(serde_json::to_value)
.transpose()
.map_err(|error| CheckpointError::new("checkpoint_json_invalid", error.to_string()))?
.unwrap_or(Value::Null),
);
object.insert(
"event_cursor".to_string(),
checkpoint
.event_cursor
.as_ref()
.map(serde_json::to_value)
.transpose()
.map_err(|error| CheckpointError::new("checkpoint_json_invalid", error.to_string()))?
.unwrap_or(Value::Null),
);
object.insert(
"event_outbox".to_string(),
Value::Array(
checkpoint
.event_outbox
.iter()
.map(EventOutboxEntry::to_value)
.collect(),
),
);
object.insert(
"extension_state".to_string(),
Value::Object(
checkpoint
.extension_state
.iter()
.map(|(namespace, entry)| (namespace.clone(), entry.to_value()))
.collect(),
),
);
object.insert(
"model_call_journal".to_string(),
Value::Array(
checkpoint
.model_call_journal
.iter()
.map(OperationJournalEntry::to_value)
.collect(),
),
);
object.insert(
"tool_journal".to_string(),
Value::Array(
checkpoint
.tool_journal
.iter()
.map(OperationJournalEntry::to_value)
.collect(),
),
);
object.insert("revision".to_string(), Value::from(checkpoint.revision));
object.insert(
"claim_token".to_string(),
checkpoint
.claim_token
.clone()
.map_or(Value::Null, Value::String),
);
object.insert(
"claimed_cycle".to_string(),
checkpoint.claimed_cycle.map_or(Value::Null, Value::from),
);
object.insert(
"lease_expires_at_ms".to_string(),
checkpoint
.lease_expires_at_ms
.map_or(Value::Null, Value::from),
);
object.insert(
"terminal_result".to_string(),
checkpoint.terminal_result.clone().unwrap_or(Value::Null),
);
object.insert(
"terminal_acknowledged".to_string(),
Value::Bool(checkpoint.terminal_acknowledged),
);
Ok(Value::Object(object))
}
pub fn checkpoint_from_value(
payload: &Value,
max_extension_state_bytes: u64,
) -> CheckpointResult<Checkpoint> {
let object = payload.as_object().ok_or_else(|| {
CheckpointError::new(
"checkpoint_payload_invalid",
"checkpoint v3 payload must be an object",
)
})?;
if let Some(field) = object
.keys()
.find(|field| !KNOWN_FIELDS.contains(&field.as_str()))
{
return Err(CheckpointError::new(
"checkpoint_unknown_field",
format!("checkpoint contains unknown field: {field}"),
));
}
if object.get("schema_version").and_then(Value::as_str) != Some(CHECKPOINT_SCHEMA) {
return Err(CheckpointError::new(
"checkpoint_schema_unsupported",
"checkpoint schema_version is not vv-agent.checkpoint.v3",
));
}
let run_definition_schema = required_string(
object,
"run_definition_schema",
"checkpoint_definition_schema_unsupported",
)?;
if run_definition_schema != RUN_DEFINITION_SCHEMA {
return Err(CheckpointError::new(
"checkpoint_definition_schema_unsupported",
"run_definition_schema is unsupported",
));
}
let run_definition = object.get("run_definition").cloned().ok_or_else(|| {
CheckpointError::new(
"checkpoint_definition_invalid",
"run_definition is required",
)
})?;
if let Some(field) = KNOWN_FIELDS
.iter()
.find(|field| !object.contains_key(**field))
{
return Err(CheckpointError::new(
"checkpoint_field_invalid",
format!("required checkpoint field {field} is missing"),
));
}
let checkpoint = Checkpoint {
schema_version: CHECKPOINT_SCHEMA.to_string(),
run_definition_schema: run_definition_schema.to_string(),
run_definition,
checkpoint_key: required_string(object, "checkpoint_key", "checkpoint_key_invalid")?
.to_string(),
task_id: required_string(object, "task_id", "checkpoint_value_invalid")?.to_string(),
root_run_id: required_string(object, "root_run_id", "checkpoint_value_invalid")?
.to_string(),
trace_id: required_string(object, "trace_id", "checkpoint_value_invalid")?.to_string(),
run_definition_digest: required_string(
object,
"run_definition_digest",
"checkpoint_digest_invalid",
)?
.to_string(),
resume_attempt: required_u64(
object,
"resume_attempt",
"checkpoint_resume_attempt_invalid",
)?,
cycle_index: required_u64(object, "cycle_index", "checkpoint_cycle_invalid")?,
status: parse_status(object, "status")?,
messages: parse_messages(object.get("messages"))?,
cycles: parse_cycles(object.get("cycles"))?,
model_calls: parse_model_calls(object.get("model_calls"))?,
shared_state: parse_object_map(object, "shared_state", "checkpoint_shared_state_invalid")?,
budget_usage: parse_optional_budget(object.get("budget_usage"))?,
event_cursor: parse_optional(object.get("event_cursor"), "event_cursor")?,
event_outbox: parse_array(object, "event_outbox")?
.iter()
.map(EventOutboxEntry::from_value)
.collect::<CheckpointResult<Vec<_>>>()?,
extension_state: parse_extensions(object.get("extension_state"))?,
model_call_journal: parse_array(object, "model_call_journal")?
.iter()
.map(OperationJournalEntry::from_value)
.collect::<CheckpointResult<Vec<_>>>()?,
tool_journal: parse_array(object, "tool_journal")?
.iter()
.map(OperationJournalEntry::from_value)
.collect::<CheckpointResult<Vec<_>>>()?,
revision: required_u64(object, "revision", "checkpoint_revision_invalid")?,
claim_token: parse_optional_string(object, "claim_token")?,
claimed_cycle: parse_optional_u64(object, "claimed_cycle")?,
lease_expires_at_ms: parse_optional_u64(object, "lease_expires_at_ms")?,
terminal_result: parse_optional_value(object, "terminal_result")?,
terminal_acknowledged: object
.get("terminal_acknowledged")
.and_then(Value::as_bool)
.ok_or_else(|| {
CheckpointError::new(
"checkpoint_status_invalid",
"terminal_acknowledged must be a boolean",
)
})?,
};
validate_checkpoint(&checkpoint)?;
crate::runtime::state::validate_extension_state_size(
&checkpoint.extension_state,
max_extension_state_bytes,
)?;
Ok(checkpoint)
}
pub fn checkpoint_to_json(
checkpoint: &Checkpoint,
max_extension_state_bytes: u64,
) -> CheckpointResult<String> {
let value = checkpoint_to_value(checkpoint, max_extension_state_bytes)?;
let bytes = canonical_json_bytes(&value, "checkpoint v3")?;
String::from_utf8(bytes).map_err(|error| {
CheckpointError::new(
"checkpoint_canonicalization_invalid",
format!("checkpoint canonical JSON is not UTF-8: {error}"),
)
})
}
pub fn checkpoint_from_json(
payload: &str,
max_extension_state_bytes: u64,
) -> CheckpointResult<Checkpoint> {
let value = strict_json_value(payload)?;
checkpoint_from_value(&value, max_extension_state_bytes)
}
fn parse_status(object: &Map<String, Value>, field: &str) -> CheckpointResult<CheckpointStatus> {
let value = required_string(object, field, "checkpoint_status_invalid")?;
serde_json::from_value(Value::String(value.to_string())).map_err(|_| {
CheckpointError::new(
"checkpoint_status_invalid",
format!("unknown checkpoint status {value}"),
)
})
}
fn parse_messages(value: Option<&Value>) -> CheckpointResult<Vec<Message>> {
let values = value.and_then(Value::as_array).ok_or_else(|| {
CheckpointError::new("checkpoint_messages_invalid", "messages must be an array")
})?;
values
.iter()
.map(|value| {
Message::from_dict(value)
.map_err(|error| CheckpointError::new("checkpoint_messages_invalid", error))
})
.collect()
}
fn parse_cycles(value: Option<&Value>) -> CheckpointResult<Vec<CycleRecord>> {
let values = value.and_then(Value::as_array).ok_or_else(|| {
CheckpointError::new("checkpoint_cycles_invalid", "cycles must be an array")
})?;
values
.iter()
.map(|value| {
CycleRecord::from_dict(value)
.map_err(|error| CheckpointError::new("checkpoint_cycles_invalid", error))
})
.collect()
}
fn parse_model_calls(value: Option<&Value>) -> CheckpointResult<Vec<ModelCallRecord>> {
let values = value.and_then(Value::as_array).ok_or_else(|| {
CheckpointError::new(
"checkpoint_model_calls_invalid",
"model_calls must be an array",
)
})?;
values
.iter()
.cloned()
.map(|value| {
serde_json::from_value(value).map_err(|error| {
CheckpointError::new("checkpoint_model_calls_invalid", error.to_string())
})
})
.collect()
}
fn parse_object_map(
object: &Map<String, Value>,
field: &str,
code: &str,
) -> CheckpointResult<BTreeMap<String, Value>> {
object
.get(field)
.and_then(Value::as_object)
.cloned()
.map(|map| map.into_iter().collect())
.ok_or_else(|| CheckpointError::new(code, format!("{field} must be an object")))
}
fn parse_optional_budget(value: Option<&Value>) -> CheckpointResult<Option<BudgetUsageSnapshot>> {
match value {
None | Some(Value::Null) => Ok(None),
Some(value) => serde_json::from_value(value.clone())
.map(Some)
.map_err(|error| {
CheckpointError::new("checkpoint_budget_usage_invalid", error.to_string())
}),
}
}
fn parse_optional<T>(value: Option<&Value>, field: &str) -> CheckpointResult<Option<T>>
where
T: serde::de::DeserializeOwned,
{
match value {
None | Some(Value::Null) => Ok(None),
Some(value) => serde_json::from_value(value.clone())
.map(Some)
.map_err(|error| {
CheckpointError::new("checkpoint_field_invalid", format!("{field}: {error}"))
}),
}
}
fn parse_extensions(
value: Option<&Value>,
) -> CheckpointResult<BTreeMap<String, ExtensionStateEntry>> {
let object = value.and_then(Value::as_object).ok_or_else(|| {
CheckpointError::new(
"checkpoint_extension_state_invalid",
"extension_state must be an object",
)
})?;
object
.iter()
.map(|(namespace, value)| Ok((namespace.clone(), ExtensionStateEntry::from_value(value)?)))
.collect()
}
fn parse_optional_string(
object: &Map<String, Value>,
field: &str,
) -> CheckpointResult<Option<String>> {
match object.get(field) {
None | Some(Value::Null) => Ok(None),
Some(Value::String(value)) => Ok(Some(value.clone())),
Some(_) => Err(CheckpointError::new(
"checkpoint_claim_invalid",
format!("{field} must be a string or null"),
)),
}
}
fn parse_optional_u64(object: &Map<String, Value>, field: &str) -> CheckpointResult<Option<u64>> {
match object.get(field) {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(number)) => number.as_u64().map(Some).ok_or_else(|| {
CheckpointError::new(
"checkpoint_integer_invalid",
format!("{field} must be a non-negative integer"),
)
}),
Some(_) => Err(CheckpointError::new(
"checkpoint_integer_invalid",
format!("{field} must be an integer or null"),
)),
}
}
fn parse_optional_value(
object: &Map<String, Value>,
field: &str,
) -> CheckpointResult<Option<Value>> {
Ok(object.get(field).filter(|value| !value.is_null()).cloned())
}
fn parse_array<'a>(
object: &'a Map<String, Value>,
field: &str,
) -> CheckpointResult<&'a Vec<Value>> {
object.get(field).and_then(Value::as_array).ok_or_else(|| {
CheckpointError::new(
"checkpoint_field_invalid",
format!("{field} must be an array"),
)
})
}
fn required_string<'a>(
object: &'a Map<String, Value>,
field: &str,
code: &str,
) -> CheckpointResult<&'a str> {
object
.get(field)
.and_then(Value::as_str)
.ok_or_else(|| CheckpointError::new(code, format!("{field} must be a string")))
}
fn required_u64(object: &Map<String, Value>, field: &str, code: &str) -> CheckpointResult<u64> {
object.get(field).and_then(Value::as_u64).ok_or_else(|| {
CheckpointError::new(code, format!("{field} must be a non-negative integer"))
})
}
fn strict_json_value(payload: &str) -> CheckpointResult<Value> {
let mut deserializer = serde_json::Deserializer::from_str(payload);
let StrictValue(value) = StrictValue::deserialize(&mut deserializer)
.map_err(|error| CheckpointError::new("checkpoint_json_invalid", error.to_string()))?;
deserializer
.end()
.map_err(|error| CheckpointError::new("checkpoint_json_invalid", error.to_string()))?;
Ok(value)
}
struct StrictValue(Value);
impl<'de> Deserialize<'de> for StrictValue {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
deserializer.deserialize_any(StrictValueVisitor)
}
}
struct StrictValueVisitor;
impl<'de> Visitor<'de> for StrictValueVisitor {
type Value = StrictValue;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("a JSON value without duplicate object keys")
}
fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
Ok(StrictValue(Value::Bool(value)))
}
fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
Ok(StrictValue(Value::Number(Number::from(value))))
}
fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
Ok(StrictValue(Value::Number(Number::from(value))))
}
fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E>
where
E: de::Error,
{
Number::from_f64(value)
.map(Value::Number)
.map(StrictValue)
.ok_or_else(|| E::custom("non-finite JSON number"))
}
fn visit_str<E>(self, value: &str) -> Result<Self::Value, E> {
Ok(StrictValue(Value::String(value.to_string())))
}
fn visit_string<E>(self, value: String) -> Result<Self::Value, E> {
Ok(StrictValue(Value::String(value)))
}
fn visit_none<E>(self) -> Result<Self::Value, E> {
Ok(StrictValue(Value::Null))
}
fn visit_unit<E>(self) -> Result<Self::Value, E> {
Ok(StrictValue(Value::Null))
}
fn visit_seq<A>(self, mut sequence: A) -> Result<Self::Value, A::Error>
where
A: SeqAccess<'de>,
{
let mut values = Vec::with_capacity(sequence.size_hint().unwrap_or(0));
while let Some(StrictValue(value)) = sequence.next_element()? {
values.push(value);
}
Ok(StrictValue(Value::Array(values)))
}
fn visit_map<A>(self, mut object: A) -> Result<Self::Value, A::Error>
where
A: MapAccess<'de>,
{
let mut values = Map::new();
while let Some(key) = object.next_key::<String>()? {
if values.contains_key(&key) {
return Err(de::Error::custom(format!(
"duplicate JSON object key {key}"
)));
}
let StrictValue(value) = object.next_value()?;
values.insert(key, value);
}
Ok(StrictValue(Value::Object(values)))
}
}