use std::collections::{BTreeMap, BTreeSet};
use std::io::Read;
use std::sync::Arc;
use rustev_contract::ids::{ArtifactId, ContentDigest, SnapshotId};
use rustev_contract::live::AuthorizationOutcome;
use rustev_contract::plan::{PlanExecution, PlanStepDetail};
use rustev_contract::remote::Disclosure;
use rustev_contract::run::{AttemptEnd, CostModel, RunRecord};
use rustev_contract::scope::{Handle, Scope};
use rustev_contract::snapshot::Snapshot;
use rustev_contract::time::Timestamp;
use rustev_contract::{Document, Identified};
use rustev_core::Compiled;
use rustev_core::seams::CancelSignal;
use rustev_runtime::{
Clock, LiveAuthorizer, LiveRefusal, LiveRequest, LiveResult, LiveSources, LiveSpend,
PreparedPlan, Runtime, SourceOutcome, SpendBinding, SpendUnit,
};
use serde::{Deserialize, Serialize};
use crate::JournalGate;
use crate::documents::{BindingIdentity, CostUnit, LiveConfig, LiveError, Probe, ProbeProvenance};
#[derive(Debug, Clone)]
pub struct ProbeSnapshot {
snapshot: Snapshot,
id: SnapshotId,
provenance: ProbeProvenance,
}
impl ProbeSnapshot {
pub fn load(probe: &Probe, case_snapshots: &[SnapshotId]) -> Result<Self, LiveError> {
let mut bytes = Vec::new();
std::fs::File::open(&probe.snapshot)
.and_then(|f| {
f.take(Snapshot::LIMITS.max_bytes as u64 + 1)
.read_to_end(&mut bytes)
})
.map_err(|e| LiveError {
status: "io_error",
subject: "probe-snapshot",
detail: e.to_string(),
})?;
Self::from_bytes(probe, &bytes, case_snapshots)
}
pub fn from_bytes(
probe: &Probe,
bytes: &[u8],
case_snapshots: &[SnapshotId],
) -> Result<Self, LiveError> {
let refused = |detail: &str| LiveError::invalid("probe-snapshot", detail);
let snapshot = Snapshot::parse(bytes).map_err(|e| refused(&e.to_string()))?;
let id = snapshot.id().map_err(|e| refused(&e.to_string()))?;
if id != probe.digest {
return Err(refused("the recomputed digest differs from probe.digest"));
}
if case_snapshots.contains(&id) {
return Err(refused("the probe snapshot is a case snapshot"));
}
Ok(Self {
snapshot,
id,
provenance: probe.provenance,
})
}
pub fn id(&self) -> &SnapshotId {
&self.id
}
}
pub struct ProbeInputs<'a> {
pub acquisition_pages: &'a [rustev_contract::acquisition::AcquisitionPage],
pub runtime: &'a Runtime,
pub config: &'a LiveConfig,
pub candidates: &'a BTreeMap<String, Compiled>,
pub snapshot: &'a ProbeSnapshot,
pub evaluation_time: Timestamp,
pub principal_handle: Vec<u8>,
pub subject: Handle,
pub scope: Scope,
pub clock: &'a dyn Clock,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum ProbeStatus {
Probed,
ProbeFailed,
SpendCapped,
Interrupted,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
pub enum IdentityVerdict {
Matched,
Unverified,
Mismatched,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ProbeFailure {
pub call: Option<u64>,
pub code: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ProbeCall {
pub case_digest: Option<ContentDigest>,
pub decision_id: String,
pub attempts: u64,
pub served: Vec<Option<String>>,
pub code: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct BindingProbe {
pub identity: BindingIdentity,
pub candidate: String,
pub calls: Vec<ProbeCall>,
pub reachable: bool,
pub served_identity: Option<IdentityVerdict>,
pub failure: Option<ProbeFailure>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ProbeReport {
pub snapshot_digest: SnapshotId,
pub provenance: ProbeProvenance,
pub status: ProbeStatus,
pub bindings: Vec<BindingProbe>,
pub spend_refusal: Option<String>,
}
const PURPOSE: &str = ".probe";
const MISMATCH: &str = "served-identity-mismatch";
fn decision_prefix(gate: &JournalGate, resume: bool) -> Result<String, LiveError> {
let prefix = format!("live/{}/{PURPOSE}/", gate.run_id());
if !resume {
return Ok(prefix);
}
let largest = gate
.journal()
.attempts_with_prefix(&prefix)
.iter()
.filter_map(|attempt| {
attempt
.strip_prefix(&prefix)?
.strip_prefix("resume-")?
.split_once('/')?
.0
.parse::<u64>()
.ok()
})
.max()
.unwrap_or(0);
let next = largest
.checked_add(1)
.ok_or_else(|| LiveError::invalid("probe-sequence", "probe resume sequence exhausted"))?;
Ok(format!("{prefix}resume-{next}/"))
}
pub fn reached(compiled: &Compiled) -> BTreeSet<BindingIdentity> {
let mut out: BTreeSet<_> = compiled
.plan
.steps
.iter()
.filter_map(|s| match &s.detail {
PlanStepDetail::Semantic(b) => Some((&b.backend_id, &b.artifact)),
PlanStepDetail::Generate(b) => Some((&b.backend_id, &b.artifact)),
PlanStepDetail::Extract(b) => Some((&b.backend_id, &b.artifact)),
_ => None,
})
.map(|(b, a)| identity(b, a))
.collect();
if let PlanExecution::Declared(d) = &compiled.plan.execution {
out.extend(
d.fallbacks
.iter()
.map(|f| identity(&f.backend_id, &f.artifact)),
);
}
out
}
fn identity(backend_id: &str, artifact: &ArtifactId) -> BindingIdentity {
BindingIdentity {
backend_id: backend_id.to_string(),
artifact: artifact.clone(),
}
}
pub async fn probe(
inputs: ProbeInputs<'_>,
authorizer: &dyn LiveAuthorizer,
gate: &JournalGate,
cancel: &CancelSignal,
) -> Result<ProbeReport, LiveError> {
probe_recorded(inputs, authorizer, gate, cancel, false, |_, _, _, _, _| {
Ok(None)
})
.await
}
pub(crate) async fn probe_recorded(
inputs: ProbeInputs<'_>,
authorizer: &dyn LiveAuthorizer,
gate: &JournalGate,
cancel: &CancelSignal,
resume: bool,
mut record: impl FnMut(
usize,
u64,
&str,
&LiveResult,
&SourceOutcome,
) -> Result<Option<ContentDigest>, LiveError>,
) -> Result<ProbeReport, LiveError> {
let ProbeInputs {
runtime: rt,
config,
candidates,
snapshot,
..
} = inputs;
let reach: BTreeMap<&String, BTreeSet<BindingIdentity>> =
candidates.iter().map(|(n, c)| (n, reached(c))).collect();
let configured: BTreeSet<_> = config.bindings.iter().map(|b| &b.identity).collect();
if reach.values().flatten().collect::<BTreeSet<_>>() != configured {
return Err(LiveError::invalid(
"live-config",
"configured bindings differ from reachable bindings",
));
}
for b in &config.bindings {
if rt
.backend_terms(&b.identity.backend_id)
.is_some_and(|(_, model)| model == CostModel::Unknown)
{
return Err(LiveError::invalid(
"spend-unit",
format!("{} discloses an unknown cost model", b.identity.backend_id),
));
}
}
let mut prepared: BTreeMap<&String, Result<Arc<PreparedPlan>, ()>> = BTreeMap::new();
let mut probes = Vec::new();
for b in &config.bindings {
let name = reach
.iter()
.find(|(_, set)| set.contains(&b.identity))
.map(|(n, _)| *n)
.expect("every configured binding is reached");
let plan = prepared
.entry(name)
.or_insert_with(|| {
rt.prepare(candidates[name].clone())
.map(Arc::new)
.map_err(drop)
})
.clone();
let terms = rt.backend_terms(&b.identity.backend_id);
let failure = if terms.is_some_and(|(d, _)| {
d.backend_id != b.identity.backend_id || d.artifact != b.identity.artifact
}) {
Some("descriptor-identity")
} else if plan.is_err() {
Some("check-plan")
} else {
None
};
probes.push(BindingProbe {
identity: b.identity.clone(),
candidate: name.clone(),
calls: vec![],
reachable: false,
served_identity: None,
failure: failure.map(|code| ProbeFailure {
call: None,
code: code.into(),
}),
});
}
let mut report = ProbeReport {
snapshot_digest: snapshot.id.clone(),
provenance: snapshot.provenance,
status: ProbeStatus::ProbeFailed,
bindings: vec![],
spend_refusal: None,
};
if probes.iter().any(|p| p.failure.is_some()) {
report.bindings = probes;
return Ok(report);
}
check_acquisitions(config, candidates)?;
let acquisitions = acquisition_postures(config)?;
let spend_bindings = spend_bindings(config);
let decision_prefix = decision_prefix(gate, resume)?;
let mut stopped = None;
'bindings: for (index, (p, b)) in probes.iter_mut().zip(&config.bindings).enumerate() {
let Ok(plan) = prepared[&p.candidate].clone() else {
unreachable!("static checks passed")
};
let mut last_end = None;
for call in 1..=config.probe.calls {
if let Some(end) = last_end {
let mut sleep = inputs
.clock
.sleep_until(u64::saturating_add(end, config.probe.interval_ms));
std::future::poll_fn(|cx| {
if cancel.poll_raised(cx).is_ready() {
return std::task::Poll::Ready(());
}
sleep.as_mut().poll(cx)
})
.await;
}
if cancel.is_raised() {
stopped = Some(ProbeStatus::Interrupted);
break 'bindings;
}
let decision_id = format!("{decision_prefix}{}/{call}", index + 1);
let result = rt
.reexecute_live_with_sources(
&plan,
LiveRequest {
decision_id: decision_id.clone(),
snapshot: snapshot.snapshot.clone(),
expected_snapshot_id: Some(snapshot.id.clone()),
evaluation_time: inputs.evaluation_time,
deadline_ms: Some(b.case_deadline_ms),
principal_handle: inputs.principal_handle.clone(),
subject: inputs.subject.clone(),
scope: inputs.scope.clone(),
capture: None,
historical: None,
},
authorizer,
cancel,
LiveSpend {
gate,
bindings: &spend_bindings,
},
LiveSources {
acquisitions: &acquisitions,
pages: inputs.acquisition_pages,
},
)
.await;
last_end = Some(inputs.clock.now());
let case_digest = result
.as_ref()
.ok()
.map(|(r, sources)| record(index + 1, call, &p.candidate, r, sources))
.transpose()?
.flatten();
let mut observed = observe(&p.identity, &decision_id, result);
observed.call.case_digest = case_digest;
if let Some(code) = &observed.spend_refusal {
report.spend_refusal = Some(code.clone());
stopped = Some(ProbeStatus::SpendCapped);
}
let failed = observed.call.code.clone();
p.calls.push(observed.call);
if let Some(code) = failed {
p.failure = Some(ProbeFailure {
call: Some(call),
code,
});
}
if stopped.is_none() && cancel.is_raised() {
stopped = Some(ProbeStatus::Interrupted);
}
if stopped.is_some() {
break 'bindings;
}
if p.failure.is_some() {
break;
}
}
}
for p in &mut probes {
let mapped = |c: &ProbeCall| c.code.as_deref().is_none_or(|c| c == MISMATCH);
p.reachable = p.calls.iter().all(mapped)
&& (p.calls.len() as u64 == config.probe.calls
|| p.failure.as_ref().is_some_and(|f| f.code == MISMATCH));
p.served_identity = verdict(p);
}
report.status = match stopped {
Some(s) => s,
None if probes.iter().all(|p| p.failure.is_none()) => ProbeStatus::Probed,
None => ProbeStatus::ProbeFailed,
};
report.bindings = probes;
Ok(report)
}
struct Observed {
call: ProbeCall,
spend_refusal: Option<String>,
}
fn observe(
binding: &BindingIdentity,
decision_id: &str,
result: Result<(LiveResult, SourceOutcome), LiveRefusal>,
) -> Observed {
let mut call = ProbeCall {
case_digest: None,
decision_id: decision_id.to_string(),
attempts: 0,
served: vec![],
code: None,
};
let (r, sources) = match result {
Ok(r) => r,
Err(e) => {
call.code = Some(refusal_code(&e).into());
return Observed {
call,
spend_refusal: None,
};
}
};
let ours = |backend_id: &str, artifact: &ArtifactId| {
backend_id == binding.backend_id && *artifact == binding.artifact
};
let refused_authorization = r.document.effects.iter().any(|e| {
ours(&e.backend_id, &e.artifact)
&& matches!(e.authorization, AuthorizationOutcome::Refused { .. })
});
let mut ends = vec![];
for a in r.record.requests.iter().flat_map(|q| &q.attempts) {
if ours(&a.backend_id, &a.artifact) {
ends.push(a.end.clone());
call.served.push(served(&r.record, &a.attempt_id));
}
}
call.attempts = ends.len() as u64;
let spend_refusal = r.spend_refusal.map(|s| s.code);
call.code = if let Some(code) = &spend_refusal {
Some(code.clone())
} else if let Some(reason) = sources.incomparable {
Some(reason)
} else if !sources.refusals.is_empty() {
Some(sources.refusals[0].1.clone())
} else if refused_authorization {
Some("authorization_refused".into())
} else if let Some(end) = ends.iter().find(|e| **e != AttemptEnd::Output) {
Some(match end {
AttemptEnd::Failed { class, .. } => class.name().into(),
AttemptEnd::Deadline => "deadline".into(),
AttemptEnd::Cancelled => "cancelled".into(),
AttemptEnd::Output => unreachable!(),
})
} else if ends.is_empty() {
Some("probe-unreached".into())
} else if call
.served
.iter()
.flatten()
.any(|s| s != binding.artifact.as_str())
{
Some(MISMATCH.into())
} else {
None
};
Observed {
call,
spend_refusal,
}
}
fn served(record: &RunRecord, attempt_id: &str) -> Option<String> {
let e = record
.remote_evidence
.iter()
.flatten()
.find(|e| e.attempt_id == attempt_id)?;
match e.served.disclosure {
Disclosure::Unknown => None,
Disclosure::Pinned | Disclosure::Reported => e.served.identity.clone(),
}
}
fn verdict(p: &BindingProbe) -> Option<IdentityVerdict> {
let served: Vec<_> = p.calls.iter().flat_map(|c| &c.served).collect();
if served.is_empty() {
None
} else if served.iter().any(|s| {
s.as_deref()
.is_some_and(|s| s != p.identity.artifact.as_str())
}) {
Some(IdentityVerdict::Mismatched)
} else if served.iter().any(|s| s.is_none()) {
Some(IdentityVerdict::Unverified)
} else {
Some(IdentityVerdict::Matched)
}
}
pub(crate) fn refusal_code(e: &LiveRefusal) -> &'static str {
match e {
LiveRefusal::SpendAcquisitionUnsupported => "acquisition-undeclared",
LiveRefusal::Rejected(_) => "rejected",
LiveRefusal::GenerationFailed(_) => "generation-failed",
LiveRefusal::SnapshotIdentity { .. } | LiveRefusal::InvalidSnapshot(_) => {
"identity-mismatch"
}
_ => "live-refused",
}
}
pub(crate) fn spend_bindings(config: &LiveConfig) -> Vec<SpendBinding> {
config
.bindings
.iter()
.map(|b| SpendBinding {
backend_id: b.identity.backend_id.clone(),
artifact: b.identity.artifact.clone(),
cost_unit: match b.cost_unit {
CostUnit::NanoUsd => SpendUnit::NanoUsd,
CostUnit::None => SpendUnit::None,
},
})
.collect()
}
pub(crate) fn check_acquisitions(
config: &LiveConfig,
candidates: &BTreeMap<String, Compiled>,
) -> Result<(), LiveError> {
let declared: BTreeSet<_> = candidates
.values()
.flat_map(|c| c.plan.acquisitions.iter().flatten())
.flat_map(|a| &a.sources)
.map(|s| s.source.id.to_string())
.collect();
if declared != config.acquisitions.keys().cloned().collect() {
return Err(LiveError::invalid(
"acquisition-undeclared",
"configured sources differ from candidate sources",
));
}
Ok(())
}
pub(crate) fn acquisition_postures(
config: &LiveConfig,
) -> Result<Vec<rustev_runtime::LiveAcquisition>, LiveError> {
config
.acquisitions
.iter()
.map(|(source, posture)| {
Ok(rustev_runtime::LiveAcquisition {
source: rustev_contract::ids::SourceId::parse(source)
.map_err(|e| LiveError::invalid("acquisition-undeclared", e))?,
posture: match posture {
crate::documents::AcquisitionMode::Live => {
rustev_runtime::AcquisitionPosture::Live
}
crate::documents::AcquisitionMode::Captured => {
rustev_runtime::AcquisitionPosture::Captured
}
},
})
})
.collect()
}