use crate::backends::{ChatRequest, Message};
use crate::flow_dispatcher::{DispatchCtx, DispatchError, NodeOutcome};
use crate::flow_execution_event::{now_ms, FlowExecutionEvent};
use crate::ir_nodes::{
IRProbe, IRReasonStep, IRRefineStep, IRStep, IRValidateStep, IRWeaveStep,
};
use crate::stream_effect::BackpressurePolicy;
use futures::StreamExt;
use sha2::{Digest, Sha256};
pub struct PureShapeStep {
pub name: String,
pub user_prompt: String,
pub framing_addendum: Option<String>,
pub kind_slug: &'static str,
pub tools: Vec<crate::backends::ToolSpec>,
pub requires_context: Option<u32>,
pub temperature: Option<f64>,
pub now_tz: Option<String>,
}
pub async fn run_step(
step: &IRStep,
ctx: &mut DispatchCtx,
) -> Result<NodeOutcome, DispatchError> {
let prompt =
crate::exec_context::interpolate_vars(&step.ask, &ctx.let_bindings);
if !step.apply_ref.is_empty() {
if let Some(registry) = ctx.tool_registry.clone() {
if let Some(entry) = registry.get(&step.apply_ref) {
if entry.is_streaming {
return run_step_streaming_tool(step, entry.clone(), &prompt, ctx).await;
}
}
}
}
let tools = synthesize_tools_from_step(step);
let shape = PureShapeStep {
name: if step.name.is_empty() {
"Step".to_string()
} else {
step.name.clone()
},
user_prompt: prompt,
framing_addendum: None,
kind_slug: "step",
tools,
requires_context: step.requires_context,
temperature: None,
now_tz: step.now_tz.clone(),
};
run_pure_shape(shape, ctx).await
}
async fn run_step_streaming_tool(
step: &IRStep,
entry: crate::tool_registry::ToolEntry,
prompt: &str,
ctx: &mut DispatchCtx,
) -> Result<NodeOutcome, DispatchError> {
let step_index = ctx.step_counter;
ctx.step_counter += 1;
if ctx.cancel.is_cancelled() {
return Err(DispatchError::UpstreamCancelled);
}
let policy =
crate::tool_dispatch_bridge::extract_stream_policy(&entry.effect_row);
let step_name = if step.name.is_empty() {
"Step".to_string()
} else {
step.name.clone()
};
ctx.tx
.send(FlowExecutionEvent::StepStart {
step_name: step_name.clone(),
step_index,
step_type: "step".to_string(),
branch_path: ctx.branch_path_string(),
timestamp_ms: now_ms(),
})
.map_err(|_| DispatchError::ChannelClosed)?;
let tool_ctx = crate::tool_dispatch_bridge::build_tool_context(
ctx.cancel.clone(),
0, );
let tool = crate::tool_dispatch_bridge::resolve_streaming_tool(&entry);
if ctx.cancel.is_cancelled() {
return Err(DispatchError::UpstreamCancelled);
}
if let Some(budget) = &ctx.budget {
let now = chrono::Utc::now();
let decision = budget.lock().unwrap().gate(&entry.name, now);
if let crate::runtime::budget_kernel::GateDecision::Deny {
retry_at,
on_exhausted,
} = decision
{
let retry_at_ms = retry_at.timestamp_millis();
match on_exhausted.as_str() {
"shed" => {
ctx.tx
.send(FlowExecutionEvent::StepComplete {
step_name: step_name.clone(),
step_index,
success: true,
full_output: String::new(),
tokens_input: 0,
tokens_output: 0,
branch_path: ctx.branch_path_string(),
timestamp_ms: now_ms(),
})
.map_err(|_| DispatchError::ChannelClosed)?;
{
let mut guard = ctx.step_audit_records.lock().await;
guard.push(crate::axonendpoint_replay::StepAuditRecord {
step_name: step_name.clone(),
step_index,
success: true,
tokens_emitted: 0,
output_hash_hex: String::new(),
effect_policy_applied: Some("budget:shed".to_string()),
chunks_dropped: 0,
chunks_degraded: 0,
timestamp_ms: now_ms(),
tool_name: Some(entry.name.clone()),
tool_chunks_emitted: Some(0),
tool_output_hash_hex: Some(String::new()),
tool_terminator_kind: Some("shed".to_string()),
anchor_breaches: Vec::new(),
});
}
ctx.let_bindings.insert(step_name.clone(), String::new());
return Ok(NodeOutcome::Completed {
output: String::new(),
tokens_emitted: 0,
step_index,
});
}
"defer" => {
return Err(DispatchError::EffectDeferred {
effect: entry.name.clone(),
retry_at_ms,
});
}
_ => {
return Err(DispatchError::EffectQuotaExhausted {
effect: entry.name.clone(),
retry_at_ms,
});
}
}
}
}
let source = tool.stream(prompt.to_string(), tool_ctx).await;
let summary = crate::flow_dispatcher::unified_stream::unified_stream_handler(
source,
policy,
&ctx.cancel,
&ctx.tx,
&step_name,
&ctx.branch_path_string(),
)
.await?;
if let Some(p) = policy {
let wire = crate::axon_server::EnforcementSummaryWire {
policy_slug: p.slug().to_string(),
chunks_pushed: summary.chunks_pushed,
chunks_delivered: summary.chunks_delivered,
drop_oldest_hits: summary.chunks_dropped,
degrade_quality_hits: summary.chunks_degraded,
pause_upstream_blocks: summary.pause_upstream_blocks,
fail_overflows: summary.fail_overflows,
failed: !summary.success,
};
ctx.enforcement_summaries
.lock()
.await
.insert(step_name.clone(), wire);
}
if summary.cancelled && ctx.cancel.is_cancelled() {
return Err(DispatchError::UpstreamCancelled);
}
ctx.tx
.send(FlowExecutionEvent::StepComplete {
step_name: step_name.clone(),
step_index,
success: summary.success,
full_output: summary.accumulated.clone(),
tokens_input: 0,
tokens_output: summary.tokens_emitted,
branch_path: ctx.branch_path_string(),
timestamp_ms: now_ms(),
})
.map_err(|_| DispatchError::ChannelClosed)?;
{
let terminator_kind = if summary.cancelled {
"cancelled"
} else if summary.terminator_message.is_some() {
"error"
} else {
"stop"
};
let record = crate::axonendpoint_replay::StepAuditRecord {
step_name: step_name.clone(),
step_index,
success: summary.success,
tokens_emitted: summary.tokens_emitted,
output_hash_hex: summary.output_hash_hex.clone(),
effect_policy_applied: policy.map(|p| p.slug().to_string()),
chunks_dropped: summary.chunks_dropped,
chunks_degraded: summary.chunks_degraded,
timestamp_ms: now_ms(),
tool_name: Some(entry.name.clone()),
tool_chunks_emitted: Some(summary.chunks_pushed),
tool_output_hash_hex: Some(summary.output_hash_hex.clone()),
tool_terminator_kind: Some(terminator_kind.to_string()),
anchor_breaches: Vec::new(),
};
let mut guard = ctx.step_audit_records.lock().await;
guard.push(record);
}
if let Some(message) = summary.terminator_message {
return Err(DispatchError::BackendError {
name: format!("tool:{}", entry.name),
message,
});
}
ctx.let_bindings
.insert(step_name.clone(), summary.accumulated.clone());
Ok(NodeOutcome::Completed {
output: summary.accumulated,
tokens_emitted: summary.tokens_emitted,
step_index,
})
}
fn synthesize_tools_from_step(step: &IRStep) -> Vec<crate::backends::ToolSpec> {
if step.apply_ref.is_empty() {
return Vec::new();
}
vec![crate::backends::ToolSpec {
name: step.apply_ref.clone(),
description: format!("Tool reference: {}", step.apply_ref),
parameters_json: "{}".to_string(),
}]
}
pub async fn run_probe(
probe: &IRProbe,
ctx: &mut DispatchCtx,
) -> Result<NodeOutcome, DispatchError> {
let shape = PureShapeStep {
name: if probe.target.is_empty() {
"Probe".to_string()
} else {
probe.target.clone()
},
user_prompt: format!("Investigate: {}", probe.target),
framing_addendum: Some(
"You are probing the target. Investigate deeply, surface what's hidden, return concisely.".into(),
),
kind_slug: "probe",
tools: Vec::new(),
requires_context: None,
temperature: None,
now_tz: None,
};
run_pure_shape(shape, ctx).await
}
pub async fn run_reason(
reason: &IRReasonStep,
ctx: &mut DispatchCtx,
) -> Result<NodeOutcome, DispatchError> {
let strategy_clause = if reason.strategy.is_empty() {
String::new()
} else {
format!(" using strategy `{}`", reason.strategy)
};
let shape = PureShapeStep {
name: if reason.target.is_empty() {
"Reason".to_string()
} else {
reason.target.clone()
},
user_prompt: format!("Reason about: {}{}", reason.target, strategy_clause),
framing_addendum: Some(
"You are reasoning deliberately. Show the steps of your reasoning where they bear on the answer.".into(),
),
kind_slug: "reason",
tools: Vec::new(),
requires_context: None,
temperature: None,
now_tz: None,
};
run_pure_shape(shape, ctx).await
}
pub async fn run_validate(
validate: &IRValidateStep,
ctx: &mut DispatchCtx,
) -> Result<NodeOutcome, DispatchError> {
let rule_clause = if validate.rule.is_empty() {
String::new()
} else {
format!(" against rule `{}`", validate.rule)
};
let shape = PureShapeStep {
name: if validate.target.is_empty() {
"Validate".to_string()
} else {
validate.target.clone()
},
user_prompt: format!("Validate: {}{}", validate.target, rule_clause),
framing_addendum: Some(
"You are validating. Return a structured verdict (pass/fail) with the reasoning that supports it.".into(),
),
kind_slug: "validate",
tools: Vec::new(),
requires_context: None,
temperature: None,
now_tz: None,
};
run_pure_shape(shape, ctx).await
}
pub async fn run_refine(
refine: &IRRefineStep,
ctx: &mut DispatchCtx,
) -> Result<NodeOutcome, DispatchError> {
let strategy_clause = if refine.strategy.is_empty() {
String::new()
} else {
format!(" using strategy `{}`", refine.strategy)
};
let shape = PureShapeStep {
name: if refine.target.is_empty() {
"Refine".to_string()
} else {
refine.target.clone()
},
user_prompt: format!("Refine: {}{}", refine.target, strategy_clause),
framing_addendum: Some(
"You are refining. Treat the target as draft input; improve it along the declared strategy without losing fidelity to its intent.".into(),
),
kind_slug: "refine",
tools: Vec::new(),
requires_context: None,
temperature: None,
now_tz: None,
};
run_pure_shape(shape, ctx).await
}
pub async fn run_weave(
weave: &IRWeaveStep,
ctx: &mut DispatchCtx,
) -> Result<NodeOutcome, DispatchError> {
let sources_clause = if weave.sources.is_empty() {
String::new()
} else {
format!(" from sources [{}]", weave.sources.join(", "))
};
let format_clause = if weave.format_type.is_empty() {
String::new()
} else {
format!(" as {}", weave.format_type)
};
let style_clause = if weave.style.is_empty() {
String::new()
} else {
format!(" in {} style", weave.style)
};
let priority_clause = if weave.priority.is_empty() {
String::new()
} else {
format!(" with priority [{}]", weave.priority.join(", "))
};
let shape = PureShapeStep {
name: if weave.target.is_empty() {
"Weave".to_string()
} else {
weave.target.clone()
},
user_prompt: format!(
"Weave: {}{}{}{}{}",
weave.target, sources_clause, format_clause, style_clause, priority_clause
),
framing_addendum: Some(
"You are weaving. Stitch the sources into the target output. Honor the declared priority + format + style.".into(),
),
kind_slug: "weave",
tools: Vec::new(),
requires_context: None,
temperature: None,
now_tz: None,
};
run_pure_shape(shape, ctx).await
}
pub async fn run_pure_shape(
shape: PureShapeStep,
ctx: &mut DispatchCtx,
) -> Result<NodeOutcome, DispatchError> {
let step_index = ctx.step_counter;
ctx.step_counter += 1;
let effect_policy = ctx.take_pending_effect_policy();
if ctx.cancel.is_cancelled() {
return Err(DispatchError::UpstreamCancelled);
}
ctx.tx
.send(FlowExecutionEvent::StepStart {
step_name: shape.name.clone(),
step_index,
step_type: shape.kind_slug.to_string(),
branch_path: ctx.branch_path_string(),
timestamp_ms: now_ms(),
})
.map_err(|_| DispatchError::ChannelClosed)?;
let backend = crate::backends::resolve_streaming_backend_with_key_and_endpoint(
&ctx.backend_name,
ctx.api_key.as_deref(),
ctx.llm_base_url.as_deref(),
ctx.llm_chat_path.as_deref(),
)
.ok_or_else(|| DispatchError::BackendError {
name: ctx.backend_name.clone(),
message: format!(
"not in streaming registry; supported: {}",
crate::backends::STREAMING_BACKEND_NAMES.join(", ")
),
})?;
let system = match &shape.framing_addendum {
Some(addendum) if ctx.system_prompt.is_empty() => addendum.clone(),
Some(addendum) => format!("{}\n\n{}", ctx.system_prompt, addendum),
None => ctx.system_prompt.clone(),
};
let system = {
let mut temporal = ctx.temporal.lock().unwrap();
crate::temporal_context::compose_effective_system(
&system,
shape.now_tz.as_deref(),
ctx.default_now_tz.as_deref(),
&mut temporal,
)
.map_err(|e| DispatchError::BackendError {
name: "temporal_context".to_string(),
message: format!("step '{}': {e}", shape.name),
})?
};
let history_msgs: Vec<Message> = {
let mut conv = ctx.conversation.lock().unwrap();
conv.truncate_to_budget(ctx.context_budget);
conv.messages()
.iter()
.map(|m| {
if m.role == "assistant" {
Message::assistant(m.content.clone())
} else {
Message::user(m.content.clone())
}
})
.collect()
};
let resolved_model = crate::model_resolution::resolve_model(
crate::backends::model_catalog::models_for(&ctx.backend_name),
shape.requires_context,
)
.map_err(|e| DispatchError::BackendError {
name: ctx.backend_name.clone(),
message: format!("model capability unsatisfied: {e}"),
})?;
let make_request = |user_prompt: &str, cancel: &crate::cancel_token::CancellationFlag| {
let mut messages = history_msgs.clone();
messages.push(Message::user(user_prompt.to_string()));
ChatRequest {
model: resolved_model.model.clone(),
messages,
system: if system.is_empty() { None } else { Some(system.clone()) },
max_tokens: None,
temperature: shape.temperature,
top_p: None,
tools: shape.tools.clone(),
stream: true,
trace_id: None,
cancel: cancel.clone(),
}
};
if ctx.cancel.is_cancelled() {
return Err(DispatchError::UpstreamCancelled);
}
let (accumulated, tokens_emitted, drop_count, degrade_count, anchor_breaches): (
String,
u64,
u64,
u64,
Vec<String>,
) = if ctx.anchors.is_empty() {
let cancel = ctx.cancel.clone();
let chunk_stream = backend
.stream(make_request(&shape.user_prompt, &cancel))
.await
.map_err(|e| DispatchError::BackendError {
name: ctx.backend_name.clone(),
message: format!("{e}"),
})?;
let (acc, toks, dc, dg) = match effect_policy {
Some(policy) => {
drain_through_enforcer(chunk_stream, &shape, ctx, policy, step_index).await?
}
None => drain_direct(chunk_stream, &shape, ctx, step_index).await?,
};
(acc, toks, dc, dg, Vec::new())
} else {
let mut user_prompt = shape.user_prompt.clone();
let mut attempt: u32 = 0;
loop {
let cancel = ctx.cancel.clone();
let chunk_stream = backend
.stream(make_request(&user_prompt, &cancel))
.await
.map_err(|e| DispatchError::BackendError {
name: ctx.backend_name.clone(),
message: format!("{e}"),
})?;
let (acc, buffered) = drain_to_buffer(chunk_stream, ctx).await?;
let results = crate::anchor_checker::check_all(&ctx.anchors, &acc);
let error_breaches: Vec<&crate::anchor_checker::AnchorResult> = results
.iter()
.filter(|r| !r.passed && r.severity == "error")
.collect();
if error_breaches.is_empty() || attempt >= MAX_ANCHOR_RETRIES {
let breaches: Vec<String> = results
.iter()
.filter(|r| !r.passed)
.map(|r| {
let first = r.violations.first().cloned().unwrap_or_default();
format!("{} [{}]: {}", r.anchor_name, r.severity, first)
})
.collect();
let toks = emit_buffered(buffered, &shape, ctx).await?;
break (acc, toks, 0, 0, breaches);
}
attempt += 1;
let feedback = error_breaches
.iter()
.enumerate()
.map(|(i, r)| {
let v = r.violations.first().cloned().unwrap_or_else(|| r.anchor_name.clone());
format!("{}. {}", i + 1, v)
})
.collect::<Vec<_>>()
.join("\n");
user_prompt = format!(
"{}\n\nIMPORTANT: Your previous response violated the following \
constraints:\n{}\n\nPlease regenerate your response, strictly \
avoiding the violations listed above.",
shape.user_prompt, feedback
);
}
};
{
let mut conv = ctx.conversation.lock().unwrap();
conv.add_user(&shape.user_prompt);
conv.add_assistant(&accumulated);
}
let output_hash_hex = sha256_hex(&accumulated);
ctx.tx
.send(FlowExecutionEvent::StepComplete {
step_name: shape.name.clone(),
step_index,
success: true,
full_output: accumulated.clone(),
tokens_input: 0,
tokens_output: tokens_emitted,
branch_path: ctx.branch_path_string(),
timestamp_ms: now_ms(),
})
.map_err(|_| DispatchError::ChannelClosed)?;
{
let record = crate::axonendpoint_replay::StepAuditRecord {
step_name: shape.name.clone(),
step_index,
success: true,
tokens_emitted,
output_hash_hex,
effect_policy_applied: effect_policy.map(|p| p.slug().to_string()),
chunks_dropped: drop_count,
chunks_degraded: degrade_count,
timestamp_ms: now_ms(),
tool_name: None,
tool_chunks_emitted: None,
tool_output_hash_hex: None,
tool_terminator_kind: None,
anchor_breaches,
};
let mut guard = ctx.step_audit_records.lock().await;
guard.push(record);
}
ctx.let_bindings
.insert(shape.name.clone(), accumulated.clone());
Ok(NodeOutcome::Completed {
output: accumulated,
tokens_emitted,
step_index,
})
}
async fn drain_direct(
chunk_stream: crate::backends::ChatStream,
shape: &PureShapeStep,
ctx: &mut DispatchCtx,
_step_index: usize,
) -> Result<(String, u64, u64, u64), DispatchError> {
use crate::backends::FinishReason;
let mut accumulated = String::new();
let mut tokens_emitted: u64 = 0;
let mut stream = chunk_stream;
while let Some(chunk_result) = stream.next().await {
if ctx.cancel.is_cancelled() {
return Err(DispatchError::UpstreamCancelled);
}
match chunk_result {
Ok(chunk) => {
if let Some(FinishReason::ToolUse) = &chunk.finish_reason {
let tool_name = shape
.tools
.first()
.map(|t| t.name.clone())
.unwrap_or_else(|| "<unknown>".to_string());
ctx.tx
.send(FlowExecutionEvent::ToolCall {
step_name: shape.name.clone(),
tool_name,
content: chunk.delta.clone(),
timestamp_ms: now_ms(),
})
.map_err(|_| DispatchError::ChannelClosed)?;
}
if !chunk.delta.is_empty() {
tokens_emitted += 1;
accumulated.push_str(&chunk.delta);
ctx.tx
.send(FlowExecutionEvent::StepToken {
step_name: shape.name.clone(),
content: chunk.delta,
token_index: tokens_emitted,
branch_path: ctx.branch_path_string(),
timestamp_ms: now_ms(),
})
.map_err(|_| DispatchError::ChannelClosed)?;
}
}
Err(e) => {
return Err(DispatchError::BackendError {
name: ctx.backend_name.clone(),
message: format!("chunk error: {e}"),
});
}
}
}
Ok((accumulated, tokens_emitted, 0, 0))
}
const MAX_ANCHOR_RETRIES: u32 = 2;
async fn drain_to_buffer(
chunk_stream: crate::backends::ChatStream,
ctx: &DispatchCtx,
) -> Result<(String, Vec<(String, bool)>), DispatchError> {
use crate::backends::FinishReason;
let mut accumulated = String::new();
let mut buffered: Vec<(String, bool)> = Vec::new();
let mut stream = chunk_stream;
while let Some(chunk_result) = stream.next().await {
if ctx.cancel.is_cancelled() {
return Err(DispatchError::UpstreamCancelled);
}
match chunk_result {
Ok(chunk) => {
let is_tool = matches!(chunk.finish_reason, Some(FinishReason::ToolUse));
if is_tool || !chunk.delta.is_empty() {
if !chunk.delta.is_empty() {
accumulated.push_str(&chunk.delta);
}
buffered.push((chunk.delta, is_tool));
}
}
Err(e) => {
return Err(DispatchError::BackendError {
name: ctx.backend_name.clone(),
message: format!("chunk error: {e}"),
});
}
}
}
Ok((accumulated, buffered))
}
async fn emit_buffered(
buffered: Vec<(String, bool)>,
shape: &PureShapeStep,
ctx: &mut DispatchCtx,
) -> Result<u64, DispatchError> {
let mut tokens_emitted: u64 = 0;
for (delta, is_tool) in buffered {
if is_tool {
let tool_name = shape
.tools
.first()
.map(|t| t.name.clone())
.unwrap_or_else(|| "<unknown>".to_string());
ctx.tx
.send(FlowExecutionEvent::ToolCall {
step_name: shape.name.clone(),
tool_name,
content: delta.clone(),
timestamp_ms: now_ms(),
})
.map_err(|_| DispatchError::ChannelClosed)?;
}
if !delta.is_empty() {
tokens_emitted += 1;
ctx.tx
.send(FlowExecutionEvent::StepToken {
step_name: shape.name.clone(),
content: delta,
token_index: tokens_emitted,
branch_path: ctx.branch_path_string(),
timestamp_ms: now_ms(),
})
.map_err(|_| DispatchError::ChannelClosed)?;
}
}
Ok(tokens_emitted)
}
async fn drain_through_enforcer(
chunk_stream: crate::backends::ChatStream,
shape: &PureShapeStep,
ctx: &mut DispatchCtx,
policy: BackpressurePolicy,
_step_index: usize,
) -> Result<(String, u64, u64, u64), DispatchError> {
use crate::stream_effect_dispatcher::{StreamPolicyEnforcer, DEFAULT_STREAM_BUFFER_CAPACITY};
use std::sync::Arc;
let enforcer = Arc::new(match policy {
BackpressurePolicy::DegradeQuality => StreamPolicyEnforcer::with_degrader(
policy,
DEFAULT_STREAM_BUFFER_CAPACITY,
Arc::new(|chunk| chunk),
),
BackpressurePolicy::DropOldest
| BackpressurePolicy::PauseUpstream
| BackpressurePolicy::Fail => StreamPolicyEnforcer::new(policy),
});
let producer_enforcer = enforcer.clone();
let producer = tokio::spawn(async move {
let summary = producer_enforcer
.drain(chunk_stream, |_e| {
})
.await;
producer_enforcer.close().await;
summary
});
let mut accumulated = String::new();
let mut tokens_emitted: u64 = 0;
while let Some(chunk) = enforcer.pop_chunk().await {
if ctx.cancel.is_cancelled() {
return Err(DispatchError::UpstreamCancelled);
}
if let Some(crate::backends::FinishReason::ToolUse) = &chunk.finish_reason {
let tool_name = shape
.tools
.first()
.map(|t| t.name.clone())
.unwrap_or_else(|| "<unknown>".to_string());
ctx.tx
.send(FlowExecutionEvent::ToolCall {
step_name: shape.name.clone(),
tool_name,
content: chunk.delta.clone(),
timestamp_ms: now_ms(),
})
.map_err(|_| DispatchError::ChannelClosed)?;
}
if !chunk.delta.is_empty() {
tokens_emitted += 1;
accumulated.push_str(&chunk.delta);
ctx.tx
.send(FlowExecutionEvent::StepToken {
step_name: shape.name.clone(),
content: chunk.delta,
token_index: tokens_emitted,
branch_path: ctx.branch_path_string(),
timestamp_ms: now_ms(),
})
.map_err(|_| DispatchError::ChannelClosed)?;
}
}
let drain_summary = producer.await.map_err(|e| DispatchError::BackendError {
name: ctx.backend_name.clone(),
message: format!("enforcer producer task join: {e}"),
})?;
let snap = enforcer.metrics_snapshot();
let wire = crate::axon_server::EnforcementSummaryWire {
policy_slug: policy.slug().to_string(),
chunks_pushed: snap.items_pushed,
chunks_delivered: snap.items_delivered,
drop_oldest_hits: snap.drop_oldest_hits,
degrade_quality_hits: snap.degrade_quality_hits,
pause_upstream_blocks: snap.pause_upstream_blocks,
fail_overflows: snap.fail_overflows,
failed: drain_summary.failed,
};
{
let mut guard = ctx.enforcement_summaries.lock().await;
guard.insert(shape.name.clone(), wire);
}
let drop_count = snap.drop_oldest_hits;
let degrade_count = snap.degrade_quality_hits;
Ok((accumulated, tokens_emitted, drop_count, degrade_count))
}
fn sha256_hex(content: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(content.as_bytes());
let digest = hasher.finalize();
let mut hex = String::with_capacity(digest.len() * 2);
for byte in digest.as_slice() {
use std::fmt::Write as _;
let _ = write!(hex, "{byte:02x}");
}
hex
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cancel_token::CancellationFlag;
use tokio::sync::mpsc;
fn fresh_ctx() -> (
DispatchCtx,
mpsc::UnboundedReceiver<FlowExecutionEvent>,
) {
let (tx, rx) = mpsc::unbounded_channel();
let ctx = DispatchCtx::new(
"TestFlow",
"stub",
"system prompt",
CancellationFlag::new(),
tx,
);
(ctx, rx)
}
#[test]
fn sha256_hex_empty_string_is_canonical() {
assert_eq!(
sha256_hex(""),
"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
);
}
#[test]
fn sha256_hex_stub_marker() {
let h = sha256_hex("(stub)");
assert_eq!(h.len(), 64);
assert!(h.chars().all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase()));
}
#[tokio::test]
async fn run_step_with_stub_backend_emits_one_token() {
use crate::ir_nodes::IRStep;
let step = IRStep {
node_type: "step",
source_line: 0,
source_column: 0,
name: "Generate".into(),
persona_ref: String::new(),
given: String::new(),
ask: "hi".into(),
use_tool: None,
probe: None,
reason: None,
weave: None,
output_type: String::new(),
confidence_floor: None,
navigate_ref: String::new(),
apply_ref: String::new(),
requires_context: None, now_tz: None, body: Vec::new(),
};
let (mut ctx, mut rx) = fresh_ctx();
let outcome = run_step(&step, &mut ctx).await.expect("run_step ok");
match outcome {
NodeOutcome::Completed { output, tokens_emitted, step_index } => {
assert_eq!(output, "(stub)");
assert_eq!(tokens_emitted, 1);
assert_eq!(step_index, 0);
}
other => panic!("expected Completed, got {other:?}"),
}
let mut events = Vec::new();
while let Ok(ev) = rx.try_recv() {
events.push(ev);
}
assert_eq!(events.len(), 3, "events: {events:?}");
assert!(matches!(events[0], FlowExecutionEvent::StepStart { .. }));
assert!(matches!(events[1], FlowExecutionEvent::StepToken { .. }));
assert!(matches!(events[2], FlowExecutionEvent::StepComplete { .. }));
}
#[tokio::test]
async fn step_without_requires_context_runs_unchanged() {
use crate::ir_nodes::IRStep;
let step = IRStep {
node_type: "step",
source_line: 0,
source_column: 0,
name: "Generate".into(),
persona_ref: String::new(),
given: String::new(),
ask: "hi".into(),
use_tool: None,
probe: None,
reason: None,
weave: None,
output_type: String::new(),
confidence_floor: None,
navigate_ref: String::new(),
apply_ref: String::new(),
requires_context: None,
now_tz: None,
body: Vec::new(),
};
let (mut ctx, _rx) = fresh_ctx();
let outcome = run_step(&step, &mut ctx).await.expect("run_step ok");
assert!(matches!(outcome, NodeOutcome::Completed { .. }));
}
#[tokio::test]
async fn step_with_unsatisfiable_requires_context_fails_closed() {
use crate::ir_nodes::IRStep;
let step = IRStep {
node_type: "step",
source_line: 0,
source_column: 0,
name: "BigSummary".into(),
persona_ref: String::new(),
given: String::new(),
ask: "summarize a long conversation".into(),
use_tool: None,
probe: None,
reason: None,
weave: None,
output_type: String::new(),
confidence_floor: None,
navigate_ref: String::new(),
apply_ref: String::new(),
requires_context: Some(16_000),
now_tz: None,
body: Vec::new(),
};
let (mut ctx, _rx) = fresh_ctx();
match run_step(&step, &mut ctx).await {
Err(DispatchError::BackendError { message, .. }) => assert!(
message.contains("model capability unsatisfied"),
"fail-closed must name the capability gap, got: {message}"
),
other => panic!("expected fail-closed BackendError, got {other:?}"),
}
}
#[tokio::test]
async fn conversation_history_accumulates_across_steps() {
use crate::ir_nodes::IRStep;
let mk = |ask: &str| IRStep {
node_type: "step",
source_line: 0,
source_column: 0,
name: "Generate".into(),
persona_ref: String::new(),
given: String::new(),
ask: ask.into(),
use_tool: None,
probe: None,
reason: None,
weave: None,
output_type: String::new(),
confidence_floor: None,
navigate_ref: String::new(),
apply_ref: String::new(),
requires_context: None, now_tz: None, body: Vec::new(),
};
let (mut ctx, _rx) = fresh_ctx();
run_step(&mk("first question"), &mut ctx).await.expect("step 1");
run_step(&mk("second question"), &mut ctx).await.expect("step 2");
let conv = ctx.conversation.lock().unwrap();
let msgs = conv.messages();
assert_eq!(msgs.len(), 4, "two turns recorded (user+assistant ×2)");
assert_eq!(msgs[0].role, "user");
assert_eq!(msgs[1].role, "assistant");
assert_eq!(msgs[1].content, "(stub)");
assert_eq!(msgs[2].role, "user");
assert_eq!(msgs[3].role, "assistant");
assert!(msgs[0].content.contains("first question"));
assert!(msgs[2].content.contains("second question"));
}
#[tokio::test]
async fn conversation_history_respects_char_budget() {
use crate::ir_nodes::IRStep;
let mk = |ask: &str| IRStep {
node_type: "step",
source_line: 0,
source_column: 0,
name: "G".into(),
persona_ref: String::new(),
given: String::new(),
ask: ask.into(),
use_tool: None,
probe: None,
reason: None,
weave: None,
output_type: String::new(),
confidence_floor: None,
navigate_ref: String::new(),
apply_ref: String::new(),
requires_context: None, now_tz: None, body: Vec::new(),
};
let (mut ctx, _rx) = fresh_ctx();
ctx.context_budget = 1;
run_step(&mk("turn one is quite long"), &mut ctx).await.expect("s1");
run_step(&mk("turn two"), &mut ctx).await.expect("s2");
run_step(&mk("turn three"), &mut ctx).await.expect("s3");
let conv = ctx.conversation.lock().unwrap();
assert!(conv.messages().len() <= 4, "budget bounds the history: {}", conv.messages().len());
}
fn step_named(name: &str, ask: &str) -> crate::ir_nodes::IRStep {
crate::ir_nodes::IRStep {
node_type: "step",
source_line: 0,
source_column: 0,
name: name.into(),
persona_ref: String::new(),
given: String::new(),
ask: ask.into(),
use_tool: None,
probe: None,
reason: None,
weave: None,
output_type: String::new(),
confidence_floor: None,
navigate_ref: String::new(),
apply_ref: String::new(),
requires_context: None, now_tz: None, body: Vec::new(),
}
}
fn anchor_named(name: &str) -> crate::ir_nodes::IRAnchor {
crate::ir_nodes::IRAnchor {
node_type: "anchor",
source_line: 0,
source_column: 0,
name: name.into(),
description: String::new(),
require: String::new(),
reject: Vec::new(),
enforce: String::new(),
confidence_floor: None,
unknown_response: String::new(),
on_violation: String::new(),
on_violation_target: String::new(),
}
}
#[tokio::test]
async fn anchor_breach_is_surfaced_in_step_audit() {
let (mut ctx, _rx) = fresh_ctx();
ctx.anchors = std::sync::Arc::new(vec![anchor_named("RequiresCitation")]);
run_step(&step_named("Generate", "say hi"), &mut ctx).await.expect("ok");
let audit = ctx.step_audit_records.lock().await;
let rec = audit.last().expect("one audit record");
assert_eq!(rec.anchor_breaches.len(), 1, "the breach is recorded: {:?}", rec.anchor_breaches);
assert!(rec.anchor_breaches[0].contains("RequiresCitation"));
assert!(rec.anchor_breaches[0].contains("[error]"));
}
#[tokio::test]
async fn no_anchors_means_no_breaches_recorded() {
let (mut ctx, _rx) = fresh_ctx();
run_step(&step_named("Generate", "say hi"), &mut ctx).await.expect("ok");
let audit = ctx.step_audit_records.lock().await;
assert!(audit.last().unwrap().anchor_breaches.is_empty());
}
#[tokio::test]
async fn anchored_step_buffers_retries_then_replays_tokens_to_wire() {
let (mut ctx, mut rx) = fresh_ctx();
ctx.anchors = std::sync::Arc::new(vec![anchor_named("RequiresCitation")]);
let outcome = run_step(&step_named("Generate", "say hi"), &mut ctx)
.await
.expect("ok");
match outcome {
NodeOutcome::Completed { output, tokens_emitted, .. } => {
assert_eq!(output, "(stub)", "accepted output after retries");
assert_eq!(tokens_emitted, 1, "the buffered chunk is replayed to the wire");
}
other => panic!("expected Completed, got {other:?}"),
}
let mut events = Vec::new();
while let Ok(ev) = rx.try_recv() {
events.push(ev);
}
assert!(matches!(events[0], FlowExecutionEvent::StepStart { .. }));
assert!(
events.iter().any(|e| matches!(
e,
FlowExecutionEvent::StepToken { content, .. } if content == "(stub)"
)),
"the accepted response's token is replayed to the wire: {events:?}"
);
assert!(events.iter().any(|e| matches!(e, FlowExecutionEvent::StepComplete { .. })));
let audit = ctx.step_audit_records.lock().await;
let breaches = &audit.last().unwrap().anchor_breaches;
assert_eq!(breaches.len(), 1, "the unresolved breach is recorded: {breaches:?}");
}
#[tokio::test]
async fn run_step_cancel_pre_dispatch_short_circuits() {
use crate::ir_nodes::IRStep;
let step = IRStep {
node_type: "step",
source_line: 0,
source_column: 0,
name: "S".into(),
persona_ref: String::new(),
given: String::new(),
ask: "hi".into(),
use_tool: None,
probe: None,
reason: None,
weave: None,
output_type: String::new(),
confidence_floor: None,
navigate_ref: String::new(),
apply_ref: String::new(),
requires_context: None, now_tz: None, body: Vec::new(),
};
let cancel = CancellationFlag::new();
cancel.cancel();
let (tx, _rx) = mpsc::unbounded_channel();
let mut ctx = DispatchCtx::new("F", "stub", "", cancel, tx);
let outcome = run_step(&step, &mut ctx).await;
assert!(matches!(outcome, Err(DispatchError::UpstreamCancelled)));
}
#[tokio::test]
async fn run_step_unknown_backend_returns_backend_error() {
use crate::ir_nodes::IRStep;
let step = IRStep {
node_type: "step",
source_line: 0,
source_column: 0,
name: "S".into(),
persona_ref: String::new(),
given: String::new(),
ask: "hi".into(),
use_tool: None,
probe: None,
reason: None,
weave: None,
output_type: String::new(),
confidence_floor: None,
navigate_ref: String::new(),
apply_ref: String::new(),
requires_context: None, now_tz: None, body: Vec::new(),
};
let (tx, _rx) = mpsc::unbounded_channel();
let mut ctx = DispatchCtx::new(
"F",
"does_not_exist",
"",
CancellationFlag::new(),
tx,
);
let outcome = run_step(&step, &mut ctx).await;
match outcome {
Err(DispatchError::BackendError { name, message }) => {
assert_eq!(name, "does_not_exist");
assert!(message.contains("not in streaming registry"));
}
other => panic!("expected BackendError, got {other:?}"),
}
}
#[tokio::test]
async fn run_step_pending_policy_consumed_on_entry() {
use crate::ir_nodes::IRStep;
let step = IRStep {
node_type: "step",
source_line: 0,
source_column: 0,
name: "S".into(),
persona_ref: String::new(),
given: String::new(),
ask: "hi".into(),
use_tool: None,
probe: None,
reason: None,
weave: None,
output_type: String::new(),
confidence_floor: None,
navigate_ref: String::new(),
apply_ref: String::new(),
requires_context: None, now_tz: None, body: Vec::new(),
};
let (mut ctx, _rx) = fresh_ctx();
ctx.pending_effect_policy = Some(BackpressurePolicy::DropOldest);
let _ = run_step(&step, &mut ctx).await.expect("ok");
assert!(
ctx.pending_effect_policy.is_none(),
"33.y.c contract: handler MUST consume pending_effect_policy on entry"
);
let summaries = ctx.enforcement_summaries.lock().await;
assert!(summaries.contains_key("S"));
assert_eq!(summaries["S"].policy_slug, "drop_oldest");
}
#[tokio::test]
async fn run_step_records_step_audit_row() {
use crate::ir_nodes::IRStep;
let step = IRStep {
node_type: "step",
source_line: 0,
source_column: 0,
name: "Generate".into(),
persona_ref: String::new(),
given: String::new(),
ask: "hi".into(),
use_tool: None,
probe: None,
reason: None,
weave: None,
output_type: String::new(),
confidence_floor: None,
navigate_ref: String::new(),
apply_ref: String::new(),
requires_context: None, now_tz: None, body: Vec::new(),
};
let (mut ctx, _rx) = fresh_ctx();
let _ = run_step(&step, &mut ctx).await.expect("ok");
let audit = ctx.step_audit_records.lock().await;
assert_eq!(audit.len(), 1);
assert_eq!(audit[0].step_name, "Generate");
assert_eq!(audit[0].tokens_emitted, 1);
assert!(audit[0].success);
assert_eq!(audit[0].output_hash_hex.len(), 64);
assert!(audit[0].effect_policy_applied.is_none());
}
#[tokio::test]
async fn run_probe_kind_slug_is_probe() {
use crate::ir_nodes::IRProbe;
let probe = IRProbe {
node_type: "probe",
source_line: 0,
source_column: 0,
target: "market_data".into(),
};
let (mut ctx, mut rx) = fresh_ctx();
let _ = run_probe(&probe, &mut ctx).await.expect("ok");
let ev = rx.try_recv().expect("event");
match ev {
FlowExecutionEvent::StepStart { step_type, step_name, .. } => {
assert_eq!(step_type, "probe");
assert_eq!(step_name, "market_data");
}
other => panic!("expected StepStart, got {other:?}"),
}
}
#[tokio::test]
async fn run_reason_kind_slug_is_reason() {
use crate::ir_nodes::IRReasonStep;
let reason = IRReasonStep {
node_type: "reason",
source_line: 0,
source_column: 0,
strategy: "chain_of_thought".into(),
target: "claim".into(),
};
let (mut ctx, mut rx) = fresh_ctx();
let _ = run_reason(&reason, &mut ctx).await.expect("ok");
let ev = rx.try_recv().expect("event");
match ev {
FlowExecutionEvent::StepStart { step_type, .. } => {
assert_eq!(step_type, "reason");
}
other => panic!("expected StepStart, got {other:?}"),
}
}
#[tokio::test]
async fn run_validate_kind_slug_is_validate() {
use crate::ir_nodes::IRValidateStep;
let validate = IRValidateStep {
node_type: "validate",
source_line: 0,
source_column: 0,
target: "draft".into(),
rule: "no_pii".into(),
};
let (mut ctx, mut rx) = fresh_ctx();
let _ = run_validate(&validate, &mut ctx).await.expect("ok");
let ev = rx.try_recv().expect("event");
match ev {
FlowExecutionEvent::StepStart { step_type, .. } => {
assert_eq!(step_type, "validate");
}
other => panic!("expected StepStart, got {other:?}"),
}
}
#[tokio::test]
async fn run_refine_kind_slug_is_refine() {
use crate::ir_nodes::IRRefineStep;
let refine = IRRefineStep {
node_type: "refine",
source_line: 0,
source_column: 0,
target: "draft".into(),
strategy: "tighten".into(),
};
let (mut ctx, mut rx) = fresh_ctx();
let _ = run_refine(&refine, &mut ctx).await.expect("ok");
let ev = rx.try_recv().expect("event");
match ev {
FlowExecutionEvent::StepStart { step_type, .. } => {
assert_eq!(step_type, "refine");
}
other => panic!("expected StepStart, got {other:?}"),
}
}
#[tokio::test]
async fn run_weave_kind_slug_is_weave() {
use crate::ir_nodes::IRWeaveStep;
let weave = IRWeaveStep {
node_type: "weave",
source_line: 0,
source_column: 0,
sources: vec!["A".into(), "B".into()],
target: "report".into(),
format_type: "markdown".into(),
priority: vec!["A".into()],
style: "formal".into(),
};
let (mut ctx, mut rx) = fresh_ctx();
let _ = run_weave(&weave, &mut ctx).await.expect("ok");
let ev = rx.try_recv().expect("event");
match ev {
FlowExecutionEvent::StepStart { step_type, .. } => {
assert_eq!(step_type, "weave");
}
other => panic!("expected StepStart, got {other:?}"),
}
}
}