use std::collections::BTreeSet;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::budget::{RunBudgetLimits, MAX_WIRE_INTEGER};
use crate::checkpoint::{
validate_checkpoint_key, validate_extension_namespace, validate_sha256, AmbiguousModelPolicy,
AmbiguousToolPolicy, ClaimMode, ResumePolicy, DEFAULT_MAX_EXTENSION_STATE_BYTES,
RUN_DEFINITION_SCHEMA,
};
use crate::tools::{ToolPolicy, ToolSideEffect};
use crate::types::AgentTask;
use super::super::RuntimeRecipe;
pub(crate) use super::contract_helpers::now_unix_ms;
use super::contract_helpers::{
normalize_integral_float, require_non_empty, utf16_cmp, validate_current_discriminator_fields,
validate_json_pointer, validate_sorted_unique,
};
pub const DISTRIBUTED_RUN_SCHEMA_VERSION: &str = "vv-agent.distributed-run.v2";
pub const DEFAULT_TOOLSET_ID: &str = "vv-agent.builtin-tools";
pub const DEFAULT_TOOLSET_VERSION: &str = "1";
pub const DEFAULT_TOOLSET_SCHEMA_DIGEST: &str =
"24d8f7bde18b11374820f742cfa244c83666626a315e09d4b6e1b69e899a70aa";
pub const DEFAULT_CYCLE_NAME: &str = "vv_agent.distributed.run_single_cycle";
pub const DEFAULT_LEASE_DURATION_MS: u64 = 5 * 60 * 1000;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct CapabilityRef {
pub id: String,
pub version: String,
}
impl CapabilityRef {
pub fn new(id: impl Into<String>, version: impl Into<String>) -> Result<Self, String> {
let reference = Self {
id: id.into(),
version: version.into(),
};
reference.validate("capability_ref")?;
Ok(reference)
}
pub fn validate(&self, field_name: &str) -> Result<(), String> {
require_non_empty(&self.id, &format!("{field_name}.id"))?;
require_non_empty(&self.version, &format!("{field_name}.version"))
}
pub fn to_dict(&self) -> Value {
serde_json::json!({"id": self.id, "version": self.version})
}
pub fn from_dict(payload: &Value, field_name: &str) -> Result<Self, String> {
let reference: Self = serde_json::from_value(payload.clone())
.map_err(|_| format!("{field_name} must be an object"))?;
reference.validate(field_name)?;
Ok(reference)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolsetRef {
pub id: String,
pub version: String,
pub schema_digest: String,
}
impl Default for ToolsetRef {
fn default() -> Self {
Self {
id: DEFAULT_TOOLSET_ID.to_string(),
version: DEFAULT_TOOLSET_VERSION.to_string(),
schema_digest: DEFAULT_TOOLSET_SCHEMA_DIGEST.to_string(),
}
}
}
impl ToolsetRef {
pub fn validate(&self) -> Result<(), String> {
CapabilityRef {
id: self.id.clone(),
version: self.version.clone(),
}
.validate("toolset_ref")?;
if self.schema_digest.len() != 64
|| self
.schema_digest
.bytes()
.any(|byte| !byte.is_ascii_hexdigit() || byte.is_ascii_uppercase())
{
return Err(
"toolset_ref.schema_digest must be a lowercase SHA-256 hex digest".to_string(),
);
}
Ok(())
}
pub fn capability_ref(&self) -> CapabilityRef {
CapabilityRef {
id: self.id.clone(),
version: self.version.clone(),
}
}
pub fn to_dict(&self) -> Value {
serde_json::json!({
"id": self.id,
"version": self.version,
"schema_digest": self.schema_digest,
})
}
pub fn from_dict(payload: &Value) -> Result<Self, String> {
let reference: Self = serde_json::from_value(payload.clone())
.map_err(|_| "toolset_ref must be an object".to_string())?;
reference.validate()?;
Ok(reference)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DistributedToolPolicy {
pub allowed_tools: Option<Vec<String>>,
#[serde(default)]
pub disallowed_tools: Vec<String>,
#[serde(default = "default_approval_policy")]
pub approval: String,
#[serde(default)]
pub predicate_ref: Option<CapabilityRef>,
#[serde(default)]
pub denied_side_effects: Vec<ToolSideEffect>,
#[serde(default)]
pub denied_capability_tags: Vec<String>,
#[serde(default)]
pub deny_terminal_tools: bool,
#[serde(default)]
pub denied_cost_dimensions: Vec<String>,
}
impl Default for DistributedToolPolicy {
fn default() -> Self {
Self {
allowed_tools: None,
disallowed_tools: Vec::new(),
approval: default_approval_policy(),
predicate_ref: None,
denied_side_effects: Vec::new(),
denied_capability_tags: Vec::new(),
deny_terminal_tools: false,
denied_cost_dimensions: Vec::new(),
}
}
}
impl DistributedToolPolicy {
pub(crate) fn set_metadata_denials(&mut self, policy: &ToolPolicy) {
self.denied_side_effects = policy.denied_side_effects.clone();
self.denied_capability_tags = policy.denied_capability_tags.clone();
self.deny_terminal_tools = policy.deny_terminal_tools;
self.denied_cost_dimensions = policy.denied_cost_dimensions.clone();
}
fn metadata_denials(&self) -> Result<ToolPolicy, String> {
ToolPolicy {
denied_side_effects: self.denied_side_effects.clone(),
denied_capability_tags: self.denied_capability_tags.clone(),
deny_terminal_tools: self.deny_terminal_tools,
denied_cost_dimensions: self.denied_cost_dimensions.clone(),
..ToolPolicy::default()
}
.normalized()
.map_err(|error| error.to_string())
}
pub fn validate(&self) -> Result<(), String> {
if !matches!(
self.approval.as_str(),
"default" | "always" | "never" | "on_request"
) {
return Err("tool_policy.approval is unsupported".to_string());
}
for (field_name, values) in [
("tool_policy.allowed_tools", self.allowed_tools.as_deref()),
(
"tool_policy.disallowed_tools",
Some(self.disallowed_tools.as_slice()),
),
] {
if values.is_some_and(|values| values.iter().any(|value| value.trim().is_empty())) {
return Err(format!("{field_name} must contain non-empty strings"));
}
}
if let Some(reference) = &self.predicate_ref {
reference.validate("tool_policy.predicate_ref")?;
}
self.metadata_denials()?;
Ok(())
}
pub fn to_dict(&self) -> Value {
serde_json::json!({
"allowed_tools": self.allowed_tools,
"disallowed_tools": self.disallowed_tools,
"approval": self.approval,
"predicate_ref": self.predicate_ref,
"denied_side_effects": self.denied_side_effects,
"denied_capability_tags": self.denied_capability_tags,
"deny_terminal_tools": self.deny_terminal_tools,
"denied_cost_dimensions": self.denied_cost_dimensions,
})
}
pub fn from_dict(payload: &Value) -> Result<Self, String> {
let policy: Self = serde_json::from_value(payload.clone())
.map_err(|_| "tool_policy must be an object".to_string())?;
policy.validate()?;
Ok(policy)
}
}
fn default_approval_policy() -> String {
"default".to_string()
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DistributedCheckpointExtensionRef {
pub namespace: String,
pub reference: CapabilityRef,
#[serde(default)]
pub required: bool,
}
impl DistributedCheckpointExtensionRef {
pub fn validate(&self) -> Result<(), String> {
validate_extension_namespace(&self.namespace).map_err(|error| error.to_string())?;
self.reference
.validate("checkpoint_extension_ref.reference")
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DistributedCheckpointConfig {
pub key: String,
#[serde(default)]
pub resume_policy: ResumePolicy,
#[serde(default)]
pub ambiguous_model_policy: AmbiguousModelPolicy,
#[serde(default)]
pub ambiguous_tool_policy: AmbiguousToolPolicy,
#[serde(default)]
pub required_extension_namespaces: Vec<String>,
#[serde(default = "default_max_extension_state_bytes")]
pub max_extension_state_bytes: u64,
#[serde(default)]
pub credential_slots: Vec<String>,
}
impl DistributedCheckpointConfig {
pub fn validate(&self) -> Result<(), String> {
validate_checkpoint_key(&self.key).map_err(|error| error.to_string())?;
if self.max_extension_state_bytes > MAX_WIRE_INTEGER {
return Err(format!(
"max_extension_state_bytes must be between 0 and {MAX_WIRE_INTEGER}"
));
}
validate_sorted_unique(
&self.required_extension_namespaces,
"required_extension_namespaces",
|namespace| validate_extension_namespace(namespace).map_err(|error| error.to_string()),
)?;
for pointer in &self.credential_slots {
validate_json_pointer(pointer)?;
}
if self
.credential_slots
.windows(2)
.any(|window| utf16_cmp(&window[0], &window[1]) != std::cmp::Ordering::Less)
{
return Err("credential_slots must be sorted and unique".to_string());
}
Ok(())
}
pub fn to_dict(&self) -> Value {
serde_json::to_value(self).expect("validated checkpoint config always serializes")
}
}
fn default_max_extension_state_bytes() -> u64 {
DEFAULT_MAX_EXTENSION_STATE_BYTES
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct DistributedCapabilities {
pub toolset_ref: ToolsetRef,
pub tool_policy: DistributedToolPolicy,
pub llm_client_ref: Option<CapabilityRef>,
pub workspace_backend_ref: Option<CapabilityRef>,
pub approval_provider_ref: Option<CapabilityRef>,
pub approval_broker_ref: Option<CapabilityRef>,
pub approval_timeout_seconds: Option<f64>,
pub cancellation_ref: Option<CapabilityRef>,
pub event_sink_ref: Option<CapabilityRef>,
pub host_cost_meter_ref: Option<CapabilityRef>,
pub app_state_ref: Option<CapabilityRef>,
pub sub_task_manager_ref: Option<CapabilityRef>,
#[serde(default)]
pub memory_provider_refs: Vec<CapabilityRef>,
#[serde(default)]
pub hook_refs: Vec<CapabilityRef>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub after_cycle_hook_refs: Vec<CapabilityRef>,
#[serde(default)]
pub observer_refs: Vec<CapabilityRef>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub checkpoint_store_ref: Option<CapabilityRef>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub checkpoint_event_store_ref: Option<CapabilityRef>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub checkpoint_extension_refs: Vec<DistributedCheckpointExtensionRef>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reconciliation_provider_ref: Option<CapabilityRef>,
}
impl DistributedCapabilities {
pub fn validate(&self) -> Result<(), String> {
self.toolset_ref.validate()?;
self.tool_policy.validate()?;
if self.approval_provider_ref.is_some() != self.approval_broker_ref.is_some() {
return Err(
"approval_provider_ref and approval_broker_ref must be declared together"
.to_string(),
);
}
if self
.approval_timeout_seconds
.is_some_and(|value| !value.is_finite() || value <= 0.0)
{
return Err(
"approval_timeout_seconds must be a finite positive number or null".to_string(),
);
}
for (field_name, reference) in [
("llm_client_ref", self.llm_client_ref.as_ref()),
("workspace_backend_ref", self.workspace_backend_ref.as_ref()),
("approval_provider_ref", self.approval_provider_ref.as_ref()),
("approval_broker_ref", self.approval_broker_ref.as_ref()),
("cancellation_ref", self.cancellation_ref.as_ref()),
("event_sink_ref", self.event_sink_ref.as_ref()),
("host_cost_meter_ref", self.host_cost_meter_ref.as_ref()),
("app_state_ref", self.app_state_ref.as_ref()),
("sub_task_manager_ref", self.sub_task_manager_ref.as_ref()),
("checkpoint_store_ref", self.checkpoint_store_ref.as_ref()),
(
"checkpoint_event_store_ref",
self.checkpoint_event_store_ref.as_ref(),
),
(
"reconciliation_provider_ref",
self.reconciliation_provider_ref.as_ref(),
),
] {
if let Some(reference) = reference {
reference.validate(field_name)?;
}
}
for (field_name, references) in [
("memory_provider_refs", self.memory_provider_refs.as_slice()),
("hook_refs", self.hook_refs.as_slice()),
(
"after_cycle_hook_refs",
self.after_cycle_hook_refs.as_slice(),
),
("observer_refs", self.observer_refs.as_slice()),
] {
for (index, reference) in references.iter().enumerate() {
reference.validate(&format!("capabilities.{field_name}[{index}]"))?;
}
}
let mut namespaces = BTreeSet::new();
for reference in &self.checkpoint_extension_refs {
reference.validate()?;
if !namespaces.insert(reference.namespace.as_str()) {
return Err(format!(
"duplicate checkpoint extension namespace {}",
reference.namespace
));
}
}
Ok(())
}
pub fn to_dict(&self) -> Value {
let mut value = serde_json::json!({
"toolset_ref": self.toolset_ref,
"tool_policy": self.tool_policy,
"llm_client_ref": self.llm_client_ref,
"workspace_backend_ref": self.workspace_backend_ref,
"approval_provider_ref": self.approval_provider_ref,
"approval_broker_ref": self.approval_broker_ref,
"approval_timeout_seconds": self.approval_timeout_seconds,
"cancellation_ref": self.cancellation_ref,
"event_sink_ref": self.event_sink_ref,
"host_cost_meter_ref": self.host_cost_meter_ref,
"app_state_ref": self.app_state_ref,
"sub_task_manager_ref": self.sub_task_manager_ref,
"memory_provider_refs": self.memory_provider_refs,
"hook_refs": self.hook_refs,
"observer_refs": self.observer_refs,
});
let object = value
.as_object_mut()
.expect("distributed capabilities are always an object");
for (name, reference) in [
("checkpoint_store_ref", self.checkpoint_store_ref.as_ref()),
(
"checkpoint_event_store_ref",
self.checkpoint_event_store_ref.as_ref(),
),
(
"reconciliation_provider_ref",
self.reconciliation_provider_ref.as_ref(),
),
] {
if let Some(reference) = reference {
object.insert(name.to_string(), reference.to_dict());
}
}
if !self.checkpoint_extension_refs.is_empty() {
object.insert(
"checkpoint_extension_refs".to_string(),
serde_json::to_value(&self.checkpoint_extension_refs)
.expect("validated checkpoint extension references always serialize"),
);
}
if !self.after_cycle_hook_refs.is_empty() {
object.insert(
"after_cycle_hook_refs".to_string(),
serde_json::to_value(&self.after_cycle_hook_refs)
.expect("validated after-cycle hook references always serialize"),
);
}
value
}
pub fn from_dict(payload: &Value) -> Result<Self, String> {
let capabilities: Self = serde_json::from_value(payload.clone())
.map_err(|_| "capabilities must be an object".to_string())?;
capabilities.validate()?;
Ok(capabilities)
}
}
#[derive(Debug, Clone, PartialEq, Deserialize)]
pub struct DistributedRunEnvelope {
pub schema_version: String,
pub job_id: String,
pub run_id: String,
pub task: AgentTask,
#[serde(default)]
pub budget_limits: Option<RunBudgetLimits>,
pub recipe: RuntimeRecipe,
pub cycle_name: String,
pub cycle_index: u32,
pub idempotency_key: String,
pub deadline_unix_ms: Option<u64>,
pub lease_duration_ms: u64,
pub root_run_id: String,
pub trace_id: String,
pub run_definition_schema: String,
pub run_definition_digest: String,
pub claim_mode: ClaimMode,
pub resume_attempt: u64,
pub checkpoint_config: DistributedCheckpointConfig,
}
impl DistributedRunEnvelope {
#[allow(clippy::too_many_arguments)]
pub fn for_cycle(
task: AgentTask,
recipe: RuntimeRecipe,
cycle_index: u32,
cycle_name: impl Into<String>,
run_id: Option<String>,
deadline_unix_ms: Option<u64>,
lease_duration_ms: u64,
budget_limits: Option<RunBudgetLimits>,
root_run_id: impl Into<String>,
trace_id: impl Into<String>,
run_definition_digest: impl Into<String>,
claim_mode: ClaimMode,
resume_attempt: u64,
checkpoint_config: DistributedCheckpointConfig,
) -> Result<Self, String> {
let run_id = run_id
.filter(|value| !value.trim().is_empty())
.or_else(|| {
task.metadata
.get("_vv_agent_run_id")
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
.map(str::to_string)
})
.unwrap_or_else(|| task.task_id.clone());
let idempotency_key = format!("{run_id}:cycle:{cycle_index}");
let envelope = Self {
schema_version: DISTRIBUTED_RUN_SCHEMA_VERSION.to_string(),
job_id: idempotency_key.clone(),
run_id,
task,
budget_limits,
recipe,
cycle_name: cycle_name.into(),
cycle_index,
idempotency_key,
deadline_unix_ms,
lease_duration_ms,
root_run_id: root_run_id.into(),
trace_id: trace_id.into(),
run_definition_schema: RUN_DEFINITION_SCHEMA.to_string(),
run_definition_digest: run_definition_digest.into(),
claim_mode,
resume_attempt,
checkpoint_config,
};
envelope.validate()?;
Ok(envelope)
}
pub fn validate(&self) -> Result<(), String> {
if self.schema_version != DISTRIBUTED_RUN_SCHEMA_VERSION {
return Err(format!(
"unsupported distributed schema_version: {}",
self.schema_version
));
}
self.validate_current_fields()?;
for (field_name, value) in [
("job_id", self.job_id.as_str()),
("run_id", self.run_id.as_str()),
("cycle_name", self.cycle_name.as_str()),
("idempotency_key", self.idempotency_key.as_str()),
] {
require_non_empty(value, &format!("distributed envelope {field_name}"))?;
}
if self.cycle_index == 0 {
return Err(
"distributed envelope cycle_index must be between 1 and 4294967295".to_string(),
);
}
if self.lease_duration_ms == 0 {
return Err(
"distributed envelope lease_duration_ms must be a positive integer".to_string(),
);
}
self.recipe.validate()?;
if let Some(limits) = &self.budget_limits {
limits.validate()?;
}
Ok(())
}
fn validate_current_fields(&self) -> Result<(), String> {
for (field_name, value) in [
("root_run_id", self.root_run_id.as_str()),
("trace_id", self.trace_id.as_str()),
] {
require_non_empty(value, &format!("distributed envelope {field_name}"))?;
}
if self.run_definition_schema != RUN_DEFINITION_SCHEMA {
return Err("checkpoint_definition_schema_unsupported".to_string());
}
validate_sha256(&self.run_definition_digest, "run_definition_digest")
.map_err(|error| error.to_string())?;
if self.resume_attempt == 0 || self.resume_attempt > MAX_WIRE_INTEGER {
return Err("checkpoint_resume_attempt_invalid".to_string());
}
self.checkpoint_config.validate()?;
if self.recipe.capabilities.checkpoint_store_ref.is_none() {
return Err("distributed run requires checkpoint_store_ref".to_string());
}
for namespace in &self.checkpoint_config.required_extension_namespaces {
if !self
.recipe
.capabilities
.checkpoint_extension_refs
.iter()
.any(|reference| reference.namespace == *namespace && reference.required)
{
return Err(format!(
"required checkpoint extension {namespace} is unavailable"
));
}
}
if self
.deadline_unix_ms
.is_some_and(|value| value > MAX_WIRE_INTEGER)
|| self.lease_duration_ms > MAX_WIRE_INTEGER
{
return Err("distributed lease values must be JSON-safe integers".to_string());
}
Ok(())
}
pub fn ensure_not_expired_at(&self, now_ms: u64) -> Result<(), String> {
if self
.deadline_unix_ms
.is_some_and(|deadline| deadline <= now_ms)
{
return Err(format!(
"distributed job {} deadline has expired",
self.job_id
));
}
Ok(())
}
pub fn ensure_not_expired(&self) -> Result<(), String> {
self.ensure_not_expired_at(now_unix_ms()?)
}
pub fn remaining_millis_at(&self, now_ms: u64) -> Option<u64> {
self.deadline_unix_ms
.map(|deadline| deadline.saturating_sub(now_ms))
}
pub fn to_dict(&self) -> Value {
let mut value = serde_json::json!({
"schema_version": self.schema_version,
"job_id": self.job_id,
"run_id": self.run_id,
"task": self.task.to_dict(),
"budget_limits": self.budget_limits,
"recipe": self.recipe.to_dict(),
"cycle_name": self.cycle_name,
"cycle_index": self.cycle_index,
"idempotency_key": self.idempotency_key,
"deadline_unix_ms": self.deadline_unix_ms,
"lease_duration_ms": self.lease_duration_ms,
"root_run_id": self.root_run_id,
"trace_id": self.trace_id,
"run_definition_schema": self.run_definition_schema,
"run_definition_digest": self.run_definition_digest,
"claim_mode": self.claim_mode,
"resume_attempt": self.resume_attempt,
"checkpoint_config": self.checkpoint_config,
});
let object = value
.as_object_mut()
.expect("distributed envelope is always an object");
normalize_integral_float(
object
.get_mut("recipe")
.and_then(Value::as_object_mut)
.expect("runtime recipe is always an object"),
"timeout_seconds",
);
if let Some(capabilities) = object
.get_mut("recipe")
.and_then(Value::as_object_mut)
.and_then(|recipe| recipe.get_mut("capabilities"))
.and_then(Value::as_object_mut)
{
normalize_integral_float(capabilities, "approval_timeout_seconds");
}
value
}
pub fn from_dict(payload: &Value) -> Result<Self, String> {
if !payload.is_object() {
return Err("distributed envelope must be an object".to_string());
}
if let Some(budget_limits) = payload
.get("budget_limits")
.filter(|budget_limits| !budget_limits.is_null())
{
serde_json::from_value::<RunBudgetLimits>(budget_limits.clone()).map_err(|error| {
format!(
"distributed envelope budget limit must be between 0 and {MAX_WIRE_INTEGER}: {error}"
)
})?;
}
let schema_version = payload
.get("schema_version")
.and_then(Value::as_str)
.ok_or_else(|| "unsupported distributed schema_version".to_string())?;
if schema_version != DISTRIBUTED_RUN_SCHEMA_VERSION {
return Err(format!(
"unsupported distributed schema_version: {schema_version}"
));
}
validate_current_discriminator_fields(payload)?;
let envelope: Self =
serde_json::from_value(payload.clone()).map_err(|error| error.to_string())?;
envelope.validate()?;
if envelope.to_dict() != *payload {
return Err(
"distributed envelope must use the complete canonical current wire shape"
.to_string(),
);
}
Ok(envelope)
}
}
impl Serialize for DistributedRunEnvelope {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
self.to_dict().serialize(serializer)
}
}