simple_agents_workflow/

yaml_runner.rs

use std::collections::{BTreeMap, HashMap, HashSet};
use std::path::Path;
use std::time::Instant;

use async_trait::async_trait;
use futures::StreamExt;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use simple_agent_type::message::{Message, Role};
use simple_agent_type::request::CompletionRequest;
use simple_agent_type::response::FinishReason;
use simple_agent_type::tool::{
    ToolCall, ToolChoice, ToolChoiceFunction, ToolChoiceMode, ToolChoiceTool, ToolDefinition,
    ToolFunction, ToolType,
};
use simple_agents_core::{
    CompletionMode, CompletionOptions, CompletionOutcome, SimpleAgentsClient,
};
use simple_agents_healing::JsonishParser;
use thiserror::Error;

use crate::ir::{Node, NodeKind, RouterRoute, WorkflowDefinition, WORKFLOW_IR_V0};
use crate::observability::tracing::{
    flush_workflow_tracer, workflow_tracer, SpanKind, TraceContext, WorkflowSpan,
};
use crate::runtime::{
    LlmExecutionError, LlmExecutionInput, LlmExecutionOutput, LlmExecutor, ToolExecutionError,
    ToolExecutionInput, ToolExecutor, WorkflowRuntime, WorkflowRuntimeError,
    WorkflowRuntimeOptions,
};
use crate::visualize::workflow_to_mermaid;

const YAML_START_NODE_ID: &str = "__yaml_start";
const YAML_LLM_TOOL_ID: &str = "__yaml_llm_call";

static TRACE_ID_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);

#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct YamlStepTiming {
    pub node_id: String,
    pub node_kind: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub model_name: Option<String>,
    pub elapsed_ms: u128,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prompt_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub completion_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub total_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tokens_per_second: Option<f64>,
}

#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct YamlLlmNodeMetrics {
    pub elapsed_ms: u128,
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
    pub total_tokens: u32,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_tokens: Option<u32>,
    pub tokens_per_second: f64,
}

#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct YamlWorkflowRunOutput {
    pub workflow_id: String,
    pub entry_node: String,
    pub email_text: String,
    pub trace: Vec<String>,
    pub outputs: BTreeMap<String, Value>,
    pub terminal_node: String,
    pub terminal_output: Option<Value>,
    pub step_timings: Vec<YamlStepTiming>,
    pub llm_node_metrics: BTreeMap<String, YamlLlmNodeMetrics>,
    pub llm_node_models: BTreeMap<String, String>,
    pub total_elapsed_ms: u128,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ttft_ms: Option<u128>,
    pub total_input_tokens: u64,
    pub total_output_tokens: u64,
    pub total_tokens: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub total_reasoning_tokens: Option<u64>,
    pub tokens_per_second: f64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub trace_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<Value>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum YamlWorkflowPayloadMode {
    #[default]
    FullPayload,
    RedactedPayload,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum YamlToolTraceMode {
    #[default]
    Full,
    Redacted,
    Off,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct YamlWorkflowTraceContextInput {
    #[serde(default)]
    pub trace_id: Option<String>,
    #[serde(default)]
    pub span_id: Option<String>,
    #[serde(default)]
    pub parent_span_id: Option<String>,
    #[serde(default)]
    pub traceparent: Option<String>,
    #[serde(default)]
    pub tracestate: Option<String>,
    #[serde(default)]
    pub baggage: BTreeMap<String, String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct YamlWorkflowTraceTenantContext {
    #[serde(default)]
    pub workspace_id: Option<String>,
    #[serde(default)]
    pub user_id: Option<String>,
    #[serde(default)]
    pub conversation_id: Option<String>,
    #[serde(default)]
    pub request_id: Option<String>,
    #[serde(default)]
    pub run_id: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct YamlWorkflowTelemetryConfig {
    #[serde(default = "default_true")]
    pub enabled: bool,
    #[serde(default = "default_true")]
    pub nerdstats: bool,
    #[serde(default = "default_sample_rate")]
    pub sample_rate: f32,
    #[serde(default)]
    pub payload_mode: YamlWorkflowPayloadMode,
    #[serde(default = "default_retention_days")]
    pub retention_days: u32,
    #[serde(default = "default_true")]
    pub multi_tenant: bool,
    #[serde(default)]
    pub tool_trace_mode: YamlToolTraceMode,
}

impl Default for YamlWorkflowTelemetryConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            nerdstats: true,
            sample_rate: 1.0,
            payload_mode: YamlWorkflowPayloadMode::FullPayload,
            retention_days: 30,
            multi_tenant: true,
            tool_trace_mode: YamlToolTraceMode::Full,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct YamlWorkflowTraceOptions {
    #[serde(default)]
    pub context: Option<YamlWorkflowTraceContextInput>,
    #[serde(default)]
    pub tenant: YamlWorkflowTraceTenantContext,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct YamlWorkflowRunOptions {
    #[serde(default)]
    pub telemetry: YamlWorkflowTelemetryConfig,
    #[serde(default)]
    pub trace: YamlWorkflowTraceOptions,
    #[serde(default)]
    pub model: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct YamlLlmTokenUsage {
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
    pub total_tokens: u32,
    pub reasoning_tokens: Option<u32>,
}

#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct YamlLlmExecutionResult {
    pub payload: Value,
    pub usage: Option<YamlLlmTokenUsage>,
    pub ttft_ms: Option<u128>,
    pub tool_calls: Vec<YamlToolCallTrace>,
}

#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct YamlToolCallTrace {
    pub id: String,
    pub name: String,
    pub arguments: Value,
    pub output: Option<Value>,
    pub status: String,
    pub elapsed_ms: u128,
    pub error: Option<String>,
}

#[derive(Debug, Clone, Default)]
struct YamlTokenTotals {
    input_tokens: u64,
    output_tokens: u64,
    total_tokens: u64,
    reasoning_tokens: Option<u64>,
}

impl YamlTokenTotals {
    fn add_usage(&mut self, usage: &YamlLlmTokenUsage) {
        self.input_tokens += u64::from(usage.prompt_tokens);
        self.output_tokens += u64::from(usage.completion_tokens);
        self.total_tokens += u64::from(usage.total_tokens);

        if let Some(reasoning_tokens) = usage.reasoning_tokens {
            let next = self.reasoning_tokens.unwrap_or(0) + u64::from(reasoning_tokens);
            self.reasoning_tokens = Some(next);
        }
    }

    fn tokens_per_second(&self, elapsed_ms: u128) -> f64 {
        if elapsed_ms == 0 {
            return 0.0;
        }
        round_two_decimals((self.output_tokens as f64) * 1000.0 / (elapsed_ms as f64))
    }
}

fn round_two_decimals(value: f64) -> f64 {
    (value * 100.0).round() / 100.0
}

fn completion_tokens_per_second(completion_tokens: u32, elapsed_ms: u128) -> f64 {
    if elapsed_ms == 0 {
        return 0.0;
    }
    round_two_decimals((completion_tokens as f64) * 1000.0 / (elapsed_ms as f64))
}

fn resolve_requested_model(run_model_override: Option<&str>, node_model: &str) -> String {
    run_model_override
        .and_then(|model| {
            let trimmed = model.trim();
            if trimmed.is_empty() {
                None
            } else {
                Some(trimmed.to_string())
            }
        })
        .unwrap_or_else(|| node_model.to_string())
}

fn default_true() -> bool {
    true
}

fn default_sample_rate() -> f32 {
    1.0
}

fn default_retention_days() -> u32 {
    30
}

fn validate_sample_rate(sample_rate: f32) -> Result<(), YamlWorkflowRunError> {
    if sample_rate.is_finite() && (0.0..=1.0).contains(&sample_rate) {
        return Ok(());
    }

    Err(YamlWorkflowRunError::InvalidInput {
        message: format!(
            "telemetry.sample_rate must be between 0.0 and 1.0 inclusive; received {sample_rate}"
        ),
    })
}

fn should_sample_trace(trace_id: &str, sample_rate: f32) -> bool {
    if sample_rate >= 1.0 {
        return true;
    }
    if sample_rate <= 0.0 {
        return false;
    }

    let mut hash: u64 = 0xcbf29ce484222325;
    for byte in trace_id.as_bytes() {
        hash ^= u64::from(*byte);
        hash = hash.wrapping_mul(0x100000001b3);
    }

    let ratio = (hash as f64) / (u64::MAX as f64);
    ratio < (sample_rate as f64)
}

fn trace_id_from_traceparent(traceparent: &str) -> Option<String> {
    let mut parts = traceparent.split('-');
    let version = parts.next()?;
    let trace_id = parts.next()?;
    let _span_id = parts.next()?;
    let _flags = parts.next()?;
    if parts.next().is_some() {
        return None;
    }

    if version.len() != 2 || trace_id.len() != 32 {
        return None;
    }

    if !version.chars().all(|ch| ch.is_ascii_hexdigit()) {
        return None;
    }
    if !trace_id.chars().all(|ch| ch.is_ascii_hexdigit()) {
        return None;
    }
    if trace_id.chars().all(|ch| ch == '0') {
        return None;
    }

    Some(trace_id.to_ascii_lowercase())
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TraceIdSource {
    Disabled,
    ExplicitTraceId,
    ParentTraceId,
    Traceparent,
    ParentTraceparent,
    Generated,
}

#[derive(Debug, Clone)]
struct ResolvedTelemetryContext {
    trace_id: Option<String>,
    sampled: bool,
    trace_id_source: TraceIdSource,
}

fn resolve_run_trace_id_with_source(
    options: &YamlWorkflowRunOptions,
    parent_trace_context: Option<&TraceContext>,
) -> (Option<String>, TraceIdSource) {
    if !options.telemetry.enabled {
        return (None, TraceIdSource::Disabled);
    }

    if let Some(trace_id) = options
        .trace
        .context
        .as_ref()
        .and_then(|context| context.trace_id.clone())
    {
        return (Some(trace_id), TraceIdSource::ExplicitTraceId);
    }

    if let Some(trace_id) = parent_trace_context.and_then(|context| context.trace_id.clone()) {
        return (Some(trace_id), TraceIdSource::ParentTraceId);
    }

    if let Some(trace_id) = options
        .trace
        .context
        .as_ref()
        .and_then(|context| context.traceparent.as_deref())
        .and_then(trace_id_from_traceparent)
    {
        return (Some(trace_id), TraceIdSource::Traceparent);
    }

    if let Some(trace_id) = parent_trace_context
        .and_then(|context| context.traceparent.as_deref())
        .and_then(trace_id_from_traceparent)
    {
        return (Some(trace_id), TraceIdSource::ParentTraceparent);
    }

    (Some(generate_trace_id()), TraceIdSource::Generated)
}

fn resolve_telemetry_context(
    options: &YamlWorkflowRunOptions,
    parent_trace_context: Option<&TraceContext>,
) -> ResolvedTelemetryContext {
    let (trace_id, trace_id_source) =
        resolve_run_trace_id_with_source(options, parent_trace_context);
    let sampled = trace_id
        .as_deref()
        .map(|value| should_sample_trace(value, options.telemetry.sample_rate))
        .unwrap_or(false);

    ResolvedTelemetryContext {
        trace_id,
        sampled,
        trace_id_source,
    }
}

fn generate_trace_id() -> String {
    use std::time::{SystemTime, UNIX_EPOCH};

    let now_nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|duration| duration.as_nanos())
        .unwrap_or(0);
    let sequence = u128::from(TRACE_ID_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed));
    format!("{:032x}", now_nanos ^ sequence)
}

fn workflow_metadata_with_trace(
    options: &YamlWorkflowRunOptions,
    trace_id: &str,
    sampled: bool,
    trace_id_source: TraceIdSource,
) -> Value {
    json!({
        "telemetry": {
            "trace_id": trace_id,
            "trace_id_source": match trace_id_source {
                TraceIdSource::Disabled => "disabled",
                TraceIdSource::ExplicitTraceId => "explicit_trace_id",
                TraceIdSource::ParentTraceId => "parent_trace_id",
                TraceIdSource::Traceparent => "traceparent",
                TraceIdSource::ParentTraceparent => "parent_traceparent",
                TraceIdSource::Generated => "generated",
            },
            "enabled": options.telemetry.enabled,
            "sampled": sampled,
            "nerdstats": options.telemetry.nerdstats,
            "sample_rate": options.telemetry.sample_rate,
            "payload_mode": match options.telemetry.payload_mode {
                YamlWorkflowPayloadMode::FullPayload => "full_payload",
                YamlWorkflowPayloadMode::RedactedPayload => "redacted_payload",
            },
            "retention_days": options.telemetry.retention_days,
            "multi_tenant": options.telemetry.multi_tenant,
            "tool_trace_mode": match options.telemetry.tool_trace_mode {
                YamlToolTraceMode::Full => "full",
                YamlToolTraceMode::Redacted => "redacted",
                YamlToolTraceMode::Off => "off",
            },
        },
        "trace": {
            "tenant": {
                "workspace_id": options.trace.tenant.workspace_id,
                "user_id": options.trace.tenant.user_id,
                "conversation_id": options.trace.tenant.conversation_id,
                "request_id": options.trace.tenant.request_id,
                "run_id": options.trace.tenant.run_id,
            }
        },
    })
}

fn apply_trace_tenant_attributes(span: &mut dyn WorkflowSpan, options: &YamlWorkflowRunOptions) {
    if let Some(workspace_id) = options.trace.tenant.workspace_id.as_deref() {
        span.set_attribute("tenant.workspace_id", workspace_id);
    }
    if let Some(user_id) = options.trace.tenant.user_id.as_deref() {
        span.set_attribute("tenant.user_id", user_id);
    }
    if let Some(conversation_id) = options.trace.tenant.conversation_id.as_deref() {
        span.set_attribute("tenant.conversation_id", conversation_id);
    }
    if let Some(request_id) = options.trace.tenant.request_id.as_deref() {
        span.set_attribute("tenant.request_id", request_id);
    }
    if let Some(run_id) = options.trace.tenant.run_id.as_deref() {
        span.set_attribute("tenant.run_id", run_id);
    }
}

fn workflow_nerdstats(output: &YamlWorkflowRunOutput) -> Value {
    let llm_nodes_without_usage: Vec<String> = output
        .step_timings
        .iter()
        .filter(|step| step.node_kind == "llm_call" && step.total_tokens.is_none())
        .map(|step| step.node_id.clone())
        .collect();
    let token_metrics_available = llm_nodes_without_usage.is_empty();
    let token_metrics_source = if token_metrics_available {
        "provider_usage"
    } else {
        "provider_stream_usage_unavailable"
    };
    let total_input_tokens = if token_metrics_available {
        json!(output.total_input_tokens)
    } else {
        Value::Null
    };
    let total_output_tokens = if token_metrics_available {
        json!(output.total_output_tokens)
    } else {
        Value::Null
    };
    let total_tokens = if token_metrics_available {
        json!(output.total_tokens)
    } else {
        Value::Null
    };
    let total_reasoning_tokens = if token_metrics_available {
        json!(output.total_reasoning_tokens)
    } else {
        Value::Null
    };
    let tokens_per_second = if token_metrics_available {
        json!(output.tokens_per_second)
    } else {
        Value::Null
    };

    json!({
        "workflow_id": output.workflow_id,
        "terminal_node": output.terminal_node,
        "total_elapsed_ms": output.total_elapsed_ms,
        "ttft_ms": output.ttft_ms,
        "step_details": output.step_timings,
        "total_input_tokens": total_input_tokens,
        "total_output_tokens": total_output_tokens,
        "total_tokens": total_tokens,
        "total_reasoning_tokens": total_reasoning_tokens,
        "tokens_per_second": tokens_per_second,
        "trace_id": output.trace_id,
        "token_metrics_available": token_metrics_available,
        "token_metrics_source": token_metrics_source,
        "llm_nodes_without_usage": llm_nodes_without_usage,
    })
}

fn payload_for_span(mode: YamlWorkflowPayloadMode, payload: &Value) -> String {
    match mode {
        YamlWorkflowPayloadMode::FullPayload => payload.to_string(),
        YamlWorkflowPayloadMode::RedactedPayload => json!({
            "redacted": true,
            "value_type": match payload {
                Value::Null => "null",
                Value::Bool(_) => "bool",
                Value::Number(_) => "number",
                Value::String(_) => "string",
                Value::Array(_) => "array",
                Value::Object(_) => "object",
            }
        })
        .to_string(),
    }
}

fn payload_for_tool_trace(mode: YamlToolTraceMode, payload: &Value) -> Value {
    match mode {
        YamlToolTraceMode::Full => payload.clone(),
        YamlToolTraceMode::Redacted => json!({
            "redacted": true,
            "value_type": json_type_name(payload),
        }),
        YamlToolTraceMode::Off => Value::Null,
    }
}

fn validate_json_schema(schema: &Value) -> Result<(), String> {
    jsonschema::JSONSchema::compile(schema)
        .map(|_| ())
        .map_err(|error| format!("invalid JSON schema: {error}"))
}

fn validate_schema_instance(schema: &Value, instance: &Value) -> Result<(), String> {
    let validator = jsonschema::JSONSchema::compile(schema)
        .map_err(|error| format!("invalid JSON schema: {error}"))?;
    if let Err(errors) = validator.validate(instance) {
        let message = errors
            .into_iter()
            .next()
            .map(|error| error.to_string())
            .unwrap_or_else(|| "unknown schema validation error".to_string());
        return Err(format!("schema validation failed: {message}"));
    }
    Ok(())
}

fn schema_type(schema: &Value) -> Option<&str> {
    schema.get("type").and_then(Value::as_str)
}

fn schema_expects_object(schema: &Value) -> bool {
    schema_type(schema) == Some("object")
}

fn trace_context_from_options(options: &YamlWorkflowRunOptions) -> Option<TraceContext> {
    options.trace.context.as_ref().map(|input| TraceContext {
        trace_id: input.trace_id.clone(),
        span_id: input.span_id.clone(),
        parent_span_id: input.parent_span_id.clone(),
        traceparent: input.traceparent.clone(),
        tracestate: input.tracestate.clone(),
        baggage: input.baggage.clone(),
    })
}

fn merged_trace_context_for_worker(
    span_context: Option<&TraceContext>,
    resolved_trace_id: Option<&str>,
    options: &YamlWorkflowRunOptions,
) -> TraceContext {
    let input_context = options.trace.context.as_ref();
    let baggage = if let Some(context) = span_context {
        if !context.baggage.is_empty() {
            context.baggage.clone()
        } else {
            input_context
                .map(|value| value.baggage.clone())
                .unwrap_or_default()
        }
    } else {
        input_context
            .map(|value| value.baggage.clone())
            .unwrap_or_default()
    };

    TraceContext {
        trace_id: span_context
            .and_then(|context| context.trace_id.clone())
            .or_else(|| resolved_trace_id.map(|value| value.to_string()))
            .or_else(|| input_context.and_then(|context| context.trace_id.clone())),
        span_id: span_context
            .and_then(|context| context.span_id.clone())
            .or_else(|| input_context.and_then(|context| context.span_id.clone())),
        parent_span_id: span_context
            .and_then(|context| context.parent_span_id.clone())
            .or_else(|| input_context.and_then(|context| context.parent_span_id.clone())),
        traceparent: span_context
            .and_then(|context| context.traceparent.clone())
            .or_else(|| input_context.and_then(|context| context.traceparent.clone())),
        tracestate: span_context
            .and_then(|context| context.tracestate.clone())
            .or_else(|| input_context.and_then(|context| context.tracestate.clone())),
        baggage,
    }
}

fn custom_worker_context_with_trace(
    context: &Value,
    trace_context: &TraceContext,
    tenant_context: &YamlWorkflowTraceTenantContext,
) -> Value {
    let mut context_with_trace = context.clone();
    let Some(root) = context_with_trace.as_object_mut() else {
        return context_with_trace;
    };

    root.insert(
        "trace".to_string(),
        json!({
            "context": {
                "trace_id": trace_context.trace_id,
                "span_id": trace_context.span_id,
                "parent_span_id": trace_context.parent_span_id,
                "traceparent": trace_context.traceparent,
                "tracestate": trace_context.tracestate,
                "baggage": trace_context.baggage,
            },
            "tenant": {
                "workspace_id": tenant_context.workspace_id,
                "user_id": tenant_context.user_id,
                "conversation_id": tenant_context.conversation_id,
                "request_id": tenant_context.request_id,
                "run_id": tenant_context.run_id,
            }
        }),
    );

    context_with_trace
}

fn include_raw_stream_debug_events() -> bool {
    match std::env::var("SIMPLE_AGENTS_WORKFLOW_STREAM_INCLUDE_RAW") {
        Ok(value) => {
            let normalized = value.trim().to_ascii_lowercase();
            normalized == "1" || normalized == "true" || normalized == "yes" || normalized == "on"
        }
        Err(_) => false,
    }
}

#[derive(Debug)]
struct StreamedPayloadResolution {
    payload: Value,
    heal_confidence: Option<f32>,
}

#[derive(Debug, Default)]
struct StreamJsonAsTextFormatter {
    raw_json: String,
    emitted: bool,
}

impl StreamJsonAsTextFormatter {
    fn push(&mut self, chunk: &str) {
        self.raw_json.push_str(chunk);
    }

    fn emit_if_ready(&mut self, complete: bool) -> Option<String> {
        if self.emitted || !complete {
            return None;
        }
        self.emitted = true;
        Some(render_json_object_as_text(self.raw_json.as_str()))
    }
}

fn render_json_object_as_text(raw_json: &str) -> String {
    let value = match serde_json::from_str::<Value>(raw_json) {
        Ok(value) => value,
        Err(_) => return raw_json.to_string(),
    };
    let Some(object) = value.as_object() else {
        return raw_json.to_string();
    };

    let mut lines = Vec::with_capacity(object.len());
    for (key, value) in object {
        let rendered = match value {
            Value::String(text) => text.clone(),
            _ => value.to_string(),
        };
        lines.push(format!("{key}: {rendered}"));
    }
    lines.join("\n")
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum YamlWorkflowTokenKind {
    Output,
    Thinking,
}

#[derive(Debug, Default)]
struct StructuredJsonDeltaFilter {
    started: bool,
    completed: bool,
    depth: u32,
    in_string: bool,
    escape: bool,
}

impl StructuredJsonDeltaFilter {
    fn split(&mut self, delta: &str) -> (Option<String>, Option<String>) {
        if delta.is_empty() {
            return (None, None);
        }

        let mut output = String::new();
        let mut thinking = String::new();

        for ch in delta.chars() {
            if self.completed {
                thinking.push(ch);
                continue;
            }

            if !self.started {
                if ch != '{' {
                    thinking.push(ch);
                    continue;
                }
                self.started = true;
                self.depth = 1;
                output.push(ch);
                continue;
            }

            output.push(ch);
            if self.in_string {
                if self.escape {
                    self.escape = false;
                    continue;
                }
                if ch == '\\' {
                    self.escape = true;
                    continue;
                }
                if ch == '"' {
                    self.in_string = false;
                }
                continue;
            }

            match ch {
                '"' => self.in_string = true,
                '{' => self.depth = self.depth.saturating_add(1),
                '}' => {
                    if self.depth > 0 {
                        self.depth -= 1;
                    }
                    if self.depth == 0 {
                        self.completed = true;
                    }
                }
                _ => {}
            }
        }

        let output = if output.is_empty() {
            None
        } else {
            Some(output)
        };
        let thinking = if thinking.is_empty() {
            None
        } else {
            Some(thinking)
        };

        (output, thinking)
    }

    fn completed(&self) -> bool {
        self.completed
    }
}

fn extract_last_fenced_json_block(raw: &str) -> Option<&str> {
    let start = raw.rfind("```json")?;
    let remainder = &raw[start + "```json".len()..];
    let end = remainder.find("```")?;
    let candidate = remainder[..end].trim();
    if candidate.is_empty() {
        return None;
    }
    Some(candidate)
}

fn extract_balanced_object_from(raw: &str, start_index: usize) -> Option<&str> {
    let mut depth = 0u32;
    let mut in_string = false;
    let mut escape = false;

    for (relative_index, ch) in raw[start_index..].char_indices() {
        if in_string {
            if escape {
                escape = false;
                continue;
            }
            if ch == '\\' {
                escape = true;
                continue;
            }
            if ch == '"' {
                in_string = false;
            }
            continue;
        }

        match ch {
            '"' => in_string = true,
            '{' => depth = depth.saturating_add(1),
            '}' => {
                if depth == 0 {
                    return None;
                }
                depth -= 1;
                if depth == 0 {
                    let end_index = start_index + relative_index + ch.len_utf8();
                    return Some(raw[start_index..end_index].trim());
                }
            }
            _ => {}
        }
    }

    None
}

fn extract_last_parsable_object(raw: &str) -> Option<&str> {
    let starts: Vec<usize> = raw
        .char_indices()
        .filter_map(|(index, ch)| if ch == '{' { Some(index) } else { None })
        .collect();

    for start in starts.into_iter().rev() {
        let Some(candidate) = extract_balanced_object_from(raw, start) else {
            continue;
        };
        if serde_json::from_str::<Value>(candidate).is_ok() {
            return Some(candidate);
        }
    }

    None
}

fn resolve_structured_json_candidate(raw: &str) -> Option<&str> {
    extract_last_fenced_json_block(raw).or_else(|| extract_last_parsable_object(raw))
}

fn parse_streamed_structured_payload(
    raw: &str,
    heal: bool,
) -> Result<StreamedPayloadResolution, String> {
    if !heal {
        if let Ok(payload) = serde_json::from_str::<Value>(raw) {
            return Ok(StreamedPayloadResolution {
                payload,
                heal_confidence: None,
            });
        }

        let candidate = resolve_structured_json_candidate(raw).ok_or_else(|| {
            "failed to parse streamed structured completion JSON: no JSON object candidate found"
                .to_string()
        })?;
        let payload = serde_json::from_str::<Value>(candidate).map_err(|error| {
            format!(
                "failed to parse streamed structured completion JSON: {error}; candidate={candidate}"
            )
        })?;
        return Ok(StreamedPayloadResolution {
            payload,
            heal_confidence: None,
        });
    }

    let candidate = resolve_structured_json_candidate(raw).unwrap_or(raw);
    let parser = JsonishParser::new();
    let healed = parser
        .parse(candidate)
        .map_err(|error| format!("failed to heal streamed structured completion JSON: {error}"))?;

    Ok(StreamedPayloadResolution {
        payload: healed.value,
        heal_confidence: Some(healed.confidence),
    })
}

#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct YamlWorkflowEvent {
    pub event_type: String,
    pub node_id: Option<String>,
    pub step_id: Option<String>,
    pub node_kind: Option<String>,
    pub streamable: Option<bool>,
    pub message: Option<String>,
    pub delta: Option<String>,
    pub token_kind: Option<YamlWorkflowTokenKind>,
    pub is_terminal_node_token: Option<bool>,
    pub elapsed_ms: Option<u128>,
    pub metadata: Option<Value>,
}

pub type WorkflowMessageRole = Role;

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct WorkflowMessage {
    pub role: WorkflowMessageRole,
    pub content: String,
    #[serde(default)]
    pub name: Option<String>,
    #[serde(default, alias = "toolCallId")]
    pub tool_call_id: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct YamlTemplateBinding {
    pub index: usize,
    pub expression: String,
    pub source_path: String,
    pub resolved: Value,
    pub resolved_type: String,
    pub missing: bool,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub enum YamlWorkflowDiagnosticSeverity {
    Error,
    Warning,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct YamlWorkflowDiagnostic {
    pub node_id: Option<String>,
    pub code: String,
    pub severity: YamlWorkflowDiagnosticSeverity,
    pub message: String,
}

#[derive(Debug, Error)]
pub enum YamlWorkflowRunError {
    #[error("failed to read workflow yaml '{path}': {source}")]
    Read {
        path: String,
        source: std::io::Error,
    },
    #[error("failed to parse workflow yaml '{path}': {source}")]
    Parse {
        path: String,
        source: serde_yaml::Error,
    },
    #[error("workflow '{workflow_id}' has no nodes")]
    EmptyNodes { workflow_id: String },
    #[error("entry node '{entry_node}' does not exist")]
    MissingEntry { entry_node: String },
    #[error("unknown node id '{node_id}'")]
    MissingNode { node_id: String },
    #[error("unsupported node type in '{node_id}'")]
    UnsupportedNodeType { node_id: String },
    #[error("unsupported switch condition format: {condition}")]
    UnsupportedCondition { condition: String },
    #[error("switch node '{node_id}' has no valid next target")]
    InvalidSwitchTarget { node_id: String },
    #[error("llm returned non-object payload for node '{node_id}'")]
    LlmPayloadNotObject { node_id: String },
    #[error("custom worker handler '{handler}' is not supported")]
    UnsupportedCustomHandler { handler: String },
    #[error("llm execution failed for node '{node_id}': {message}")]
    Llm { node_id: String, message: String },
    #[error("custom worker execution failed for node '{node_id}': {message}")]
    CustomWorker { node_id: String, message: String },
    #[error("workflow validation failed with {diagnostics_count} error(s)")]
    Validation {
        diagnostics_count: usize,
        diagnostics: Vec<YamlWorkflowDiagnostic>,
    },
    #[error("invalid workflow input: {message}")]
    InvalidInput { message: String },
    #[error("ir runtime execution failed: {message}")]
    IrRuntime { message: String },
    #[error("workflow event stream cancelled: {message}")]
    EventSinkCancelled { message: String },
}

pub trait YamlWorkflowEventSink: Send + Sync {
    fn emit(&self, event: &YamlWorkflowEvent);

    fn is_cancelled(&self) -> bool {
        false
    }
}

pub struct NoopYamlWorkflowEventSink;

impl YamlWorkflowEventSink for NoopYamlWorkflowEventSink {
    fn emit(&self, _event: &YamlWorkflowEvent) {}
}

fn workflow_event_sink_cancelled_message() -> &'static str {
    "workflow event callback cancelled"
}

fn event_sink_is_cancelled(event_sink: Option<&dyn YamlWorkflowEventSink>) -> bool {
    event_sink.map(|sink| sink.is_cancelled()).unwrap_or(false)
}

#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum YamlToIrError {
    #[error("entry node '{entry_node}' does not exist")]
    MissingEntry { entry_node: String },
    #[error("node '{node_id}' has multiple outgoing edges in YAML; IR llm/tool nodes require one")]
    MultipleOutgoingEdge { node_id: String },
    #[error("node '{node_id}' is unsupported for IR conversion: {reason}")]
    UnsupportedNode { node_id: String, reason: String },
}

/// Render a YAML workflow graph as Mermaid flowchart.
pub fn yaml_workflow_to_mermaid(workflow: &YamlWorkflow) -> String {
    if let Ok(ir) = yaml_workflow_to_ir(workflow) {
        return workflow_to_mermaid(&ir);
    }

    yaml_workflow_to_mermaid_fallback(workflow)
}

fn yaml_workflow_to_mermaid_fallback(workflow: &YamlWorkflow) -> String {
    let mut lines = Vec::new();
    lines.push("flowchart TD".to_string());
    let mut tool_node_ids: Vec<String> = Vec::new();

    for node in &workflow.nodes {
        let label = format!("{}\\n({})", node.id, node.kind_name());
        lines.push(format!(
            "  {}[\"{}\"]",
            sanitize_mermaid_id(&node.id),
            escape_mermaid_label(label.as_str()),
        ));

        if let Some(llm) = node.node_type.llm_call.as_ref() {
            for (idx, tool_name) in llm_tool_names(llm).into_iter().enumerate() {
                let tool_id = sanitize_mermaid_id(format!("{}__tool_{}", node.id, idx).as_str());
                lines.push(format!(
                    "  {}([\"{}\"])",
                    tool_id,
                    escape_mermaid_label(format!("tool: {tool_name}").as_str())
                ));
                lines.push(format!(
                    "  {} -.-> {}",
                    sanitize_mermaid_id(&node.id),
                    tool_id
                ));
                tool_node_ids.push(tool_id);
            }
        }
    }

    let mut emitted: HashSet<(String, String, String)> = HashSet::new();

    for edge in &workflow.edges {
        emitted.insert((edge.from.clone(), String::new(), edge.to.clone()));
    }

    for node in &workflow.nodes {
        if let Some(switch) = node.node_type.switch.as_ref() {
            for branch in &switch.branches {
                emitted.insert((
                    node.id.clone(),
                    branch.condition.clone(),
                    branch.target.clone(),
                ));
            }
            emitted.insert((
                node.id.clone(),
                "default".to_string(),
                switch.default.clone(),
            ));
        }
    }

    let mut edges = emitted.into_iter().collect::<Vec<_>>();
    edges.sort();

    for (from, label, to) in edges {
        if label.is_empty() {
            lines.push(format!(
                "  {} --> {}",
                sanitize_mermaid_id(&from),
                sanitize_mermaid_id(&to)
            ));
        } else {
            lines.push(format!(
                "  {} -- \"{}\" --> {}",
                sanitize_mermaid_id(&from),
                escape_mermaid_label(&label),
                sanitize_mermaid_id(&to)
            ));
        }
    }

    if !tool_node_ids.is_empty() {
        lines.push("  classDef toolNode fill:#FFF4D6,stroke:#D97706,color:#7C2D12;".to_string());
        lines.push(format!("  class {} toolNode;", tool_node_ids.join(",")));
    }

    lines.join("\n")
}

fn llm_tool_names(llm: &YamlLlmCall) -> Vec<String> {
    llm.tools
        .iter()
        .map(|tool| match tool {
            YamlToolDeclaration::OpenAi(openai) => openai.function.name.clone(),
            YamlToolDeclaration::Simplified(simple) => simple.name.clone(),
        })
        .collect()
}

/// Load a YAML workflow file and render it as Mermaid flowchart.
pub fn yaml_workflow_file_to_mermaid(workflow_path: &Path) -> Result<String, YamlWorkflowRunError> {
    let contents =
        std::fs::read_to_string(workflow_path).map_err(|source| YamlWorkflowRunError::Read {
            path: workflow_path.display().to_string(),
            source,
        })?;

    let workflow: YamlWorkflow =
        serde_yaml::from_str(&contents).map_err(|source| YamlWorkflowRunError::Parse {
            path: workflow_path.display().to_string(),
            source,
        })?;

    let referenced_subgraphs = discover_referenced_subgraphs(workflow_path, &workflow)?;
    if referenced_subgraphs.is_empty() {
        return Ok(yaml_workflow_to_mermaid(&workflow));
    }

    Ok(yaml_workflow_to_mermaid_with_subgraphs(
        &workflow,
        &referenced_subgraphs,
    ))
}

#[derive(Debug, Clone)]
struct MermaidSubgraphWorkflow {
    alias: String,
    workflow: YamlWorkflow,
}

#[derive(Debug, Default)]
struct MermaidBlockRender {
    lines: Vec<String>,
    tool_node_ids: Vec<String>,
    run_workflow_tool_node_ids: Vec<String>,
    entry_node_id: String,
}

fn yaml_workflow_to_mermaid_with_subgraphs(
    workflow: &YamlWorkflow,
    subgraphs: &[MermaidSubgraphWorkflow],
) -> String {
    let mut lines = vec!["flowchart TD".to_string()];

    let main_block = render_mermaid_block(workflow, "main");
    lines.push(format!(
        "  subgraph main_graph[\"Main: {}\"]",
        escape_mermaid_label(&workflow.id)
    ));
    for line in &main_block.lines {
        lines.push(format!("    {line}"));
    }
    lines.push("  end".to_string());

    let mut all_tool_nodes = main_block.tool_node_ids.clone();

    for (index, subgraph) in subgraphs.iter().enumerate() {
        let block_id = format!("subgraph_{}", index + 1);
        let block = render_mermaid_block(&subgraph.workflow, &block_id);
        lines.push(format!(
            "  subgraph {}[\"Subgraph: {}\"]",
            sanitize_mermaid_id(&format!("{}_cluster", block_id)),
            escape_mermaid_label(&subgraph.alias)
        ));
        for line in &block.lines {
            lines.push(format!("    {line}"));
        }
        lines.push("  end".to_string());

        for tool_node in &main_block.run_workflow_tool_node_ids {
            lines.push(format!(
                "  {} -. \"{}\" .-> {}",
                tool_node,
                escape_mermaid_label(&format!("calls {}", subgraph.alias)),
                block.entry_node_id
            ));
        }

        all_tool_nodes.extend(block.tool_node_ids);
    }

    if !all_tool_nodes.is_empty() {
        lines.push("  classDef toolNode fill:#FFF4D6,stroke:#D97706,color:#7C2D12;".to_string());
        lines.push(format!("  class {} toolNode;", all_tool_nodes.join(",")));
    }

    lines.join("\n")
}

fn render_mermaid_block(workflow: &YamlWorkflow, prefix: &str) -> MermaidBlockRender {
    let mut block = MermaidBlockRender {
        entry_node_id: prefixed_mermaid_id(prefix, &workflow.entry_node),
        ..Default::default()
    };

    for node in &workflow.nodes {
        let node_id = prefixed_mermaid_id(prefix, &node.id);
        let label = format!("{}\\n({})", node.id, node.kind_name());
        block
            .lines
            .push(format!("{}[\"{}\"]", node_id, escape_mermaid_label(&label)));

        if let Some(llm) = node.node_type.llm_call.as_ref() {
            for (idx, tool) in llm.tools.iter().enumerate() {
                let tool_name = tool_declaration_name(tool);
                let tool_id = prefixed_mermaid_id(prefix, &format!("{}__tool_{}", node.id, idx));
                block.lines.push(format!(
                    "{}([\"{}\"])",
                    tool_id,
                    escape_mermaid_label(format!("tool: {tool_name}").as_str())
                ));
                block.lines.push(format!("{} -.-> {}", node_id, tool_id));
                block.tool_node_ids.push(tool_id.clone());

                if tool_name == "run_workflow_graph" {
                    block.run_workflow_tool_node_ids.push(tool_id);
                }
            }
        }
    }

    let mut emitted: HashSet<(String, String, String)> = HashSet::new();

    for edge in &workflow.edges {
        emitted.insert((edge.from.clone(), String::new(), edge.to.clone()));
    }

    for node in &workflow.nodes {
        if let Some(switch) = node.node_type.switch.as_ref() {
            for branch in &switch.branches {
                emitted.insert((
                    node.id.clone(),
                    branch.condition.clone(),
                    branch.target.clone(),
                ));
            }
            emitted.insert((
                node.id.clone(),
                "default".to_string(),
                switch.default.clone(),
            ));
        }
    }

    let mut edges = emitted.into_iter().collect::<Vec<_>>();
    edges.sort();

    for (from, label, to) in edges {
        let from_id = prefixed_mermaid_id(prefix, &from);
        let to_id = prefixed_mermaid_id(prefix, &to);
        if label.is_empty() {
            block.lines.push(format!("{} --> {}", from_id, to_id));
        } else {
            block.lines.push(format!(
                "{} -- \"{}\" --> {}",
                from_id,
                escape_mermaid_label(&label),
                to_id
            ));
        }
    }

    block
}

fn discover_referenced_subgraphs(
    workflow_path: &Path,
    workflow: &YamlWorkflow,
) -> Result<Vec<MermaidSubgraphWorkflow>, YamlWorkflowRunError> {
    let workflow_ids = referenced_workflow_ids(workflow);
    if workflow_ids.is_empty() {
        return Ok(Vec::new());
    }

    let parent_dir = workflow_path.parent().unwrap_or(Path::new("."));
    let sibling_workflows = load_yaml_sibling_workflows(parent_dir, workflow_path)?;

    let mut discovered = Vec::new();
    let mut seen = HashSet::new();

    for workflow_id in workflow_ids {
        let normalized = normalize_workflow_lookup_key(&workflow_id);
        if seen.contains(&normalized) {
            continue;
        }

        if let Some((_, subworkflow)) = sibling_workflows.iter().find(|(key, _)| key == &normalized)
        {
            discovered.push(MermaidSubgraphWorkflow {
                alias: workflow_id.clone(),
                workflow: subworkflow.clone(),
            });
            seen.insert(normalized);
        }
    }

    Ok(discovered)
}

fn load_yaml_sibling_workflows(
    parent_dir: &Path,
    workflow_path: &Path,
) -> Result<Vec<(String, YamlWorkflow)>, YamlWorkflowRunError> {
    let mut results = Vec::new();
    let entries = std::fs::read_dir(parent_dir).map_err(|source| YamlWorkflowRunError::Read {
        path: parent_dir.display().to_string(),
        source,
    })?;

    for entry in entries {
        let entry = entry.map_err(|source| YamlWorkflowRunError::Read {
            path: parent_dir.display().to_string(),
            source,
        })?;
        let path = entry.path();
        if !is_yaml_file(&path) {
            continue;
        }

        if path == workflow_path {
            continue;
        }

        let contents =
            std::fs::read_to_string(&path).map_err(|source| YamlWorkflowRunError::Read {
                path: path.display().to_string(),
                source,
            })?;
        let subworkflow: YamlWorkflow =
            serde_yaml::from_str(&contents).map_err(|source| YamlWorkflowRunError::Parse {
                path: path.display().to_string(),
                source,
            })?;

        if let Some(stem) = path.file_stem().and_then(|value| value.to_str()) {
            results.push((normalize_workflow_lookup_key(stem), subworkflow.clone()));
        }
        results.push((normalize_workflow_lookup_key(&subworkflow.id), subworkflow));
    }

    Ok(results)
}

fn referenced_workflow_ids(workflow: &YamlWorkflow) -> Vec<String> {
    let mut ids = Vec::new();
    let mut seen = HashSet::new();

    for node in &workflow.nodes {
        let prompt = node
            .config
            .as_ref()
            .and_then(|config| config.prompt.as_deref());

        if let Some(llm) = node.node_type.llm_call.as_ref() {
            for tool in &llm.tools {
                if tool_declaration_name(tool) != "run_workflow_graph" {
                    continue;
                }

                for workflow_id in referenced_workflow_ids_from_tool(tool) {
                    let normalized = normalize_workflow_lookup_key(&workflow_id);
                    if seen.insert(normalized) {
                        ids.push(workflow_id);
                    }
                }

                if let Some(prompt_text) = prompt {
                    for workflow_id in referenced_workflow_ids_from_prompt(prompt_text) {
                        let normalized = normalize_workflow_lookup_key(&workflow_id);
                        if seen.insert(normalized) {
                            ids.push(workflow_id);
                        }
                    }
                }
            }
        }
    }

    ids
}

fn referenced_workflow_ids_from_tool(tool: &YamlToolDeclaration) -> Vec<String> {
    let schema = match tool {
        YamlToolDeclaration::OpenAi(openai) => openai.function.parameters.as_ref(),
        YamlToolDeclaration::Simplified(simple) => Some(&simple.input_schema),
    };

    let mut ids = Vec::new();
    if let Some(schema) = schema {
        if let Some(workflow_prop) = schema
            .get("properties")
            .and_then(Value::as_object)
            .and_then(|properties| properties.get("workflow_id"))
        {
            if let Some(value) = workflow_prop.get("const").and_then(Value::as_str) {
                ids.push(value.to_string());
            }

            if let Some(enum_values) = workflow_prop.get("enum").and_then(Value::as_array) {
                for value in enum_values {
                    if let Some(value) = value.as_str() {
                        ids.push(value.to_string());
                    }
                }
            }
        }
    }

    ids
}

fn referenced_workflow_ids_from_prompt(prompt: &str) -> Vec<String> {
    let mut ids = Vec::new();
    let mut search = prompt;

    while let Some(index) = search.find("\"workflow_id\"") {
        let remainder = &search[index + "\"workflow_id\"".len()..];
        let Some(colon_index) = remainder.find(':') else {
            break;
        };

        let candidate = remainder[colon_index + 1..].trim_start();
        if let Some(rest) = candidate.strip_prefix('"') {
            if let Some(end_quote_index) = rest.find('"') {
                let workflow_id = rest[..end_quote_index].trim();
                if !workflow_id.is_empty() {
                    ids.push(workflow_id.to_string());
                }
                search = &rest[end_quote_index + 1..];
                continue;
            }
        }

        search = &remainder[colon_index + 1..];
    }

    ids
}

fn normalize_workflow_lookup_key(value: &str) -> String {
    value
        .chars()
        .map(|ch| match ch {
            '-' => '_',
            _ => ch.to_ascii_lowercase(),
        })
        .collect()
}

fn is_yaml_file(path: &Path) -> bool {
    matches!(
        path.extension().and_then(|ext| ext.to_str()),
        Some("yaml") | Some("yml")
    )
}

fn prefixed_mermaid_id(prefix: &str, id: &str) -> String {
    sanitize_mermaid_id(format!("{}__{}", prefix, id).as_str())
}

fn tool_declaration_name(tool: &YamlToolDeclaration) -> &str {
    match tool {
        YamlToolDeclaration::OpenAi(openai) => &openai.function.name,
        YamlToolDeclaration::Simplified(simple) => &simple.name,
    }
}

pub fn yaml_workflow_to_ir(workflow: &YamlWorkflow) -> Result<WorkflowDefinition, YamlToIrError> {
    let known_ids: HashSet<&str> = workflow.nodes.iter().map(|n| n.id.as_str()).collect();
    if !known_ids.contains(workflow.entry_node.as_str()) {
        return Err(YamlToIrError::MissingEntry {
            entry_node: workflow.entry_node.clone(),
        });
    }

    let mut outgoing: HashMap<&str, Vec<&str>> = HashMap::new();
    for edge in &workflow.edges {
        outgoing
            .entry(edge.from.as_str())
            .or_default()
            .push(edge.to.as_str());
    }

    let mut nodes = Vec::with_capacity(workflow.nodes.len() + 1);
    nodes.push(Node {
        id: YAML_START_NODE_ID.to_string(),
        kind: NodeKind::Start {
            next: workflow.entry_node.clone(),
        },
    });

    for node in &workflow.nodes {
        if let Some(llm) = node.node_type.llm_call.as_ref() {
            if node
                .config
                .as_ref()
                .and_then(|c| c.set_globals.as_ref())
                .is_some()
                || node
                    .config
                    .as_ref()
                    .and_then(|c| c.update_globals.as_ref())
                    .is_some()
            {
                return Err(YamlToIrError::UnsupportedNode {
                    node_id: node.id.clone(),
                    reason: "set_globals/update_globals are not represented in canonical IR llm nodes yet"
                        .to_string(),
                });
            }

            if !llm.tools.is_empty() {
                return Err(YamlToIrError::UnsupportedNode {
                    node_id: node.id.clone(),
                    reason: "llm_call.tools are not represented in canonical IR llm nodes yet"
                        .to_string(),
                });
            }

            let next = single_next_for_node(&outgoing, &node.id)?;
            nodes.push(Node {
                id: node.id.clone(),
                kind: NodeKind::Tool {
                    tool: YAML_LLM_TOOL_ID.to_string(),
                    input: json!({
                        "node_id": node.id,
                        "model": llm.model,
                        "prompt_template": node
                            .config
                            .as_ref()
                            .and_then(|c| c.prompt.clone())
                            .unwrap_or_default(),
                        "stream": llm.stream.unwrap_or(false),
                        "stream_json_as_text": llm.stream_json_as_text.unwrap_or(false),
                        "heal": llm.heal.unwrap_or(false),
                        "messages_path": llm.messages_path,
                        "append_prompt_as_user": llm.append_prompt_as_user.unwrap_or(true),
                        "output_schema": node
                            .config
                            .as_ref()
                            .and_then(|c| c.output_schema.clone())
                            .unwrap_or_else(default_llm_output_schema),
                    }),
                    next,
                },
            });
            continue;
        }

        if let Some(worker) = node.node_type.custom_worker.as_ref() {
            if node
                .config
                .as_ref()
                .and_then(|c| c.set_globals.as_ref())
                .is_some()
                || node
                    .config
                    .as_ref()
                    .and_then(|c| c.update_globals.as_ref())
                    .is_some()
            {
                return Err(YamlToIrError::UnsupportedNode {
                    node_id: node.id.clone(),
                    reason: "set_globals/update_globals are not represented in canonical IR tool nodes yet"
                        .to_string(),
                });
            }

            let next = single_next_for_node(&outgoing, &node.id)?;
            nodes.push(Node {
                id: node.id.clone(),
                kind: NodeKind::Tool {
                    tool: worker.handler.clone(),
                    input: node
                        .config
                        .as_ref()
                        .and_then(|c| c.payload.clone())
                        .unwrap_or_else(|| json!({})),
                    next,
                },
            });
            continue;
        }

        if let Some(switch) = node.node_type.switch.as_ref() {
            nodes.push(Node {
                id: node.id.clone(),
                kind: NodeKind::Router {
                    routes: switch
                        .branches
                        .iter()
                        .map(|b| RouterRoute {
                            when: rewrite_yaml_condition_to_ir(&b.condition),
                            next: b.target.clone(),
                        })
                        .collect(),
                    default: switch.default.clone(),
                },
            });
            continue;
        }

        return Err(YamlToIrError::UnsupportedNode {
            node_id: node.id.clone(),
            reason: "node_type must be llm_call, switch, or custom_worker".to_string(),
        });
    }

    Ok(WorkflowDefinition {
        version: WORKFLOW_IR_V0.to_string(),
        name: workflow.id.clone(),
        nodes,
    })
}

fn single_next_for_node(
    outgoing: &HashMap<&str, Vec<&str>>,
    node_id: &str,
) -> Result<Option<String>, YamlToIrError> {
    match outgoing.get(node_id) {
        None => Ok(None),
        Some(targets) if targets.len() == 1 => Ok(Some(targets[0].to_string())),
        Some(_) => Err(YamlToIrError::MultipleOutgoingEdge {
            node_id: node_id.to_string(),
        }),
    }
}

fn rewrite_yaml_condition_to_ir(expr: &str) -> String {
    let rewritten = expr
        .replace("$.nodes.", "$.node_outputs.")
        .replace(".output.", ".");
    if let Some(prefix) = rewritten.strip_suffix(".output") {
        prefix.to_string()
    } else {
        rewritten
    }
}

fn sanitize_mermaid_id(id: &str) -> String {
    let mut out = String::with_capacity(id.len() + 1);
    if id
        .chars()
        .next()
        .is_some_and(|ch| ch.is_ascii_alphabetic() || ch == '_')
    {
        out.push_str(id);
    } else {
        out.push('n');
        out.push('_');
        out.push_str(id);
    }
    out.chars()
        .map(|ch| {
            if ch.is_ascii_alphanumeric() || ch == '_' {
                ch
            } else {
                '_'
            }
        })
        .collect()
}

fn escape_mermaid_label(label: &str) -> String {
    label.replace('"', "\\\"")
}

#[derive(Debug, Clone)]
pub struct YamlLlmExecutionRequest {
    pub node_id: String,
    pub is_terminal_node: bool,
    pub stream_json_as_text: bool,
    pub model: String,
    pub messages: Option<Vec<Message>>,
    pub append_prompt_as_user: bool,
    pub prompt: String,
    pub prompt_template: String,
    pub prompt_bindings: Vec<YamlTemplateBinding>,
    pub schema: Value,
    pub stream: bool,
    pub heal: bool,
    pub tools: Vec<YamlResolvedTool>,
    pub tool_choice: Option<ToolChoice>,
    pub max_tool_roundtrips: u8,
    pub tool_calls_global_key: Option<String>,
    pub tool_trace_mode: YamlToolTraceMode,
    pub execution_context: Value,
    pub email_text: String,
    pub trace_id: Option<String>,
    pub trace_context: Option<TraceContext>,
    pub trace_sampled: bool,
}

#[derive(Debug, Clone)]
pub struct YamlResolvedTool {
    pub definition: ToolDefinition,
    pub output_schema: Option<Value>,
}

#[async_trait]
pub trait YamlWorkflowLlmExecutor: Send + Sync {
    async fn complete_structured(
        &self,
        request: YamlLlmExecutionRequest,
        event_sink: Option<&dyn YamlWorkflowEventSink>,
    ) -> Result<YamlLlmExecutionResult, String>;
}

#[async_trait]
pub trait YamlWorkflowCustomWorkerExecutor: Send + Sync {
    async fn execute(
        &self,
        handler: &str,
        payload: &Value,
        email_text: &str,
        context: &Value,
    ) -> Result<Value, String>;
}

pub async fn run_workflow_yaml_file(
    workflow_path: &Path,
    workflow_input: &Value,
    executor: &dyn YamlWorkflowLlmExecutor,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let contents =
        std::fs::read_to_string(workflow_path).map_err(|source| YamlWorkflowRunError::Read {
            path: workflow_path.display().to_string(),
            source,
        })?;

    let workflow: YamlWorkflow =
        serde_yaml::from_str(&contents).map_err(|source| YamlWorkflowRunError::Parse {
            path: workflow_path.display().to_string(),
            source,
        })?;

    run_workflow_yaml(&workflow, workflow_input, executor).await
}

pub async fn run_email_workflow_yaml_file(
    workflow_path: &Path,
    email_text: &str,
    executor: &dyn YamlWorkflowLlmExecutor,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let workflow_input = json!({ "email_text": email_text });
    run_workflow_yaml_file(workflow_path, &workflow_input, executor).await
}

pub async fn run_workflow_yaml_file_with_client(
    workflow_path: &Path,
    workflow_input: &Value,
    client: &SimpleAgentsClient,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let contents =
        std::fs::read_to_string(workflow_path).map_err(|source| YamlWorkflowRunError::Read {
            path: workflow_path.display().to_string(),
            source,
        })?;

    let workflow: YamlWorkflow =
        serde_yaml::from_str(&contents).map_err(|source| YamlWorkflowRunError::Parse {
            path: workflow_path.display().to_string(),
            source,
        })?;

    run_workflow_yaml_with_client(&workflow, workflow_input, client).await
}

pub async fn run_email_workflow_yaml_file_with_client(
    workflow_path: &Path,
    email_text: &str,
    client: &SimpleAgentsClient,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let workflow_input = json!({ "email_text": email_text });
    run_workflow_yaml_file_with_client(workflow_path, &workflow_input, client).await
}

pub async fn run_workflow_yaml_with_client(
    workflow: &YamlWorkflow,
    workflow_input: &Value,
    client: &SimpleAgentsClient,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    run_workflow_yaml_with_client_and_custom_worker(workflow, workflow_input, client, None).await
}

pub async fn run_email_workflow_yaml_with_client(
    workflow: &YamlWorkflow,
    email_text: &str,
    client: &SimpleAgentsClient,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let workflow_input = json!({ "email_text": email_text });
    run_workflow_yaml_with_client(workflow, &workflow_input, client).await
}

pub async fn run_workflow_yaml_file_with_client_and_custom_worker(
    workflow_path: &Path,
    workflow_input: &Value,
    client: &SimpleAgentsClient,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let contents =
        std::fs::read_to_string(workflow_path).map_err(|source| YamlWorkflowRunError::Read {
            path: workflow_path.display().to_string(),
            source,
        })?;

    let workflow: YamlWorkflow =
        serde_yaml::from_str(&contents).map_err(|source| YamlWorkflowRunError::Parse {
            path: workflow_path.display().to_string(),
            source,
        })?;

    run_workflow_yaml_with_client_and_custom_worker(
        &workflow,
        workflow_input,
        client,
        custom_worker,
    )
    .await
}

pub async fn run_email_workflow_yaml_file_with_client_and_custom_worker(
    workflow_path: &Path,
    email_text: &str,
    client: &SimpleAgentsClient,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let workflow_input = json!({ "email_text": email_text });
    run_workflow_yaml_file_with_client_and_custom_worker(
        workflow_path,
        &workflow_input,
        client,
        custom_worker,
    )
    .await
}

pub async fn run_workflow_yaml_file_with_client_and_custom_worker_and_events(
    workflow_path: &Path,
    workflow_input: &Value,
    client: &SimpleAgentsClient,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
    event_sink: Option<&dyn YamlWorkflowEventSink>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    run_workflow_yaml_file_with_client_and_custom_worker_and_events_and_options(
        workflow_path,
        workflow_input,
        client,
        custom_worker,
        event_sink,
        &YamlWorkflowRunOptions::default(),
    )
    .await
}

pub async fn run_workflow_yaml_file_with_client_and_custom_worker_and_events_and_options(
    workflow_path: &Path,
    workflow_input: &Value,
    client: &SimpleAgentsClient,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
    event_sink: Option<&dyn YamlWorkflowEventSink>,
    options: &YamlWorkflowRunOptions,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let contents =
        std::fs::read_to_string(workflow_path).map_err(|source| YamlWorkflowRunError::Read {
            path: workflow_path.display().to_string(),
            source,
        })?;

    let workflow: YamlWorkflow =
        serde_yaml::from_str(&contents).map_err(|source| YamlWorkflowRunError::Parse {
            path: workflow_path.display().to_string(),
            source,
        })?;

    run_workflow_yaml_with_client_and_custom_worker_and_events_and_options(
        &workflow,
        workflow_input,
        client,
        custom_worker,
        event_sink,
        options,
    )
    .await
}

pub async fn run_email_workflow_yaml_file_with_client_and_custom_worker_and_events(
    workflow_path: &Path,
    email_text: &str,
    client: &SimpleAgentsClient,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
    event_sink: Option<&dyn YamlWorkflowEventSink>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let workflow_input = json!({ "email_text": email_text });
    run_workflow_yaml_file_with_client_and_custom_worker_and_events(
        workflow_path,
        &workflow_input,
        client,
        custom_worker,
        event_sink,
    )
    .await
}

pub async fn run_workflow_yaml_with_client_and_custom_worker(
    workflow: &YamlWorkflow,
    workflow_input: &Value,
    client: &SimpleAgentsClient,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    run_workflow_yaml_with_client_and_custom_worker_and_events_and_options(
        workflow,
        workflow_input,
        client,
        custom_worker,
        None,
        &YamlWorkflowRunOptions::default(),
    )
    .await
}

pub async fn run_email_workflow_yaml_with_client_and_custom_worker(
    workflow: &YamlWorkflow,
    email_text: &str,
    client: &SimpleAgentsClient,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let workflow_input = json!({ "email_text": email_text });
    run_workflow_yaml_with_client_and_custom_worker(
        workflow,
        &workflow_input,
        client,
        custom_worker,
    )
    .await
}

pub async fn run_workflow_yaml_with_client_and_custom_worker_and_events(
    workflow: &YamlWorkflow,
    workflow_input: &Value,
    client: &SimpleAgentsClient,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
    event_sink: Option<&dyn YamlWorkflowEventSink>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    run_workflow_yaml_with_client_and_custom_worker_and_events_and_options(
        workflow,
        workflow_input,
        client,
        custom_worker,
        event_sink,
        &YamlWorkflowRunOptions::default(),
    )
    .await
}

pub async fn run_workflow_yaml_with_client_and_custom_worker_and_events_and_options(
    workflow: &YamlWorkflow,
    workflow_input: &Value,
    client: &SimpleAgentsClient,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
    event_sink: Option<&dyn YamlWorkflowEventSink>,
    options: &YamlWorkflowRunOptions,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    struct BorrowedClientExecutor<'a> {
        client: &'a SimpleAgentsClient,
        custom_worker: Option<&'a dyn YamlWorkflowCustomWorkerExecutor>,
        run_options: YamlWorkflowRunOptions,
    }

    #[async_trait]
    impl<'a> YamlWorkflowLlmExecutor for BorrowedClientExecutor<'a> {
        async fn complete_structured(
            &self,
            request: YamlLlmExecutionRequest,
            event_sink: Option<&dyn YamlWorkflowEventSink>,
        ) -> Result<YamlLlmExecutionResult, String> {
            let expects_object = schema_expects_object(&request.schema);
            let messages = if let Some(mut history) = request.messages.clone() {
                if request.append_prompt_as_user && !request.prompt.trim().is_empty() {
                    history.push(Message::user(&request.prompt));
                }
                history
            } else {
                vec![
                    Message::system("You execute workflow classification steps."),
                    Message::user(&request.prompt),
                ]
            };

            if !request.tools.is_empty() {
                let mut tool_traces: Vec<YamlToolCallTrace> = Vec::new();
                let mut conversation = messages;
                let mut usage_total: Option<YamlLlmTokenUsage> = None;

                for roundtrip in 0..=request.max_tool_roundtrips {
                    let mut builder = CompletionRequest::builder()
                        .model(&request.model)
                        .messages(conversation.clone())
                        .tools(request.tools.iter().map(|t| t.definition.clone()).collect());

                    if request.heal && expects_object {
                        builder = builder.json_schema("workflow_step", request.schema.clone());
                    }

                    if let Some(choice) = request.tool_choice.clone() {
                        builder = builder.tool_choice(choice);
                    }

                    if request.stream {
                        builder = builder.stream(true);
                    }

                    let completion_request = builder
                        .build()
                        .map_err(|error| format!("failed to build completion request: {error}"))?;

                    let outcome = self
                        .client
                        .complete(&completion_request, CompletionOptions::default())
                        .await
                        .map_err(|error| error.to_string())?;

                    let mut streamed_tool_calls: Option<Vec<ToolCall>> = None;
                    let mut streamed_content = String::new();
                    let mut finish_reason = FinishReason::Stop;

                    match outcome {
                        CompletionOutcome::Response(response) => {
                            if let Some(usage) = usage_total.as_mut() {
                                usage.prompt_tokens += response.usage.prompt_tokens;
                                usage.completion_tokens += response.usage.completion_tokens;
                                usage.total_tokens += response.usage.total_tokens;
                                if let Some(reasoning_tokens) = response.usage.reasoning_tokens {
                                    usage.reasoning_tokens = Some(
                                        usage.reasoning_tokens.unwrap_or(0) + reasoning_tokens,
                                    );
                                }
                            } else {
                                usage_total = Some(YamlLlmTokenUsage {
                                    prompt_tokens: response.usage.prompt_tokens,
                                    completion_tokens: response.usage.completion_tokens,
                                    total_tokens: response.usage.total_tokens,
                                    reasoning_tokens: response.usage.reasoning_tokens,
                                });
                            }

                            let choice = response
                                .choices
                                .first()
                                .ok_or_else(|| "completion returned no choices".to_string())?;
                            streamed_content = choice.message.content.clone();
                            streamed_tool_calls = choice.message.tool_calls.clone();
                            finish_reason = choice.finish_reason;
                        }
                        CompletionOutcome::HealedJson(healed) => {
                            let response = healed.response;
                            if let Some(usage) = usage_total.as_mut() {
                                usage.prompt_tokens += response.usage.prompt_tokens;
                                usage.completion_tokens += response.usage.completion_tokens;
                                usage.total_tokens += response.usage.total_tokens;
                                if let Some(reasoning_tokens) = response.usage.reasoning_tokens {
                                    usage.reasoning_tokens = Some(
                                        usage.reasoning_tokens.unwrap_or(0) + reasoning_tokens,
                                    );
                                }
                            } else {
                                usage_total = Some(YamlLlmTokenUsage {
                                    prompt_tokens: response.usage.prompt_tokens,
                                    completion_tokens: response.usage.completion_tokens,
                                    total_tokens: response.usage.total_tokens,
                                    reasoning_tokens: response.usage.reasoning_tokens,
                                });
                            }

                            let choice = response
                                .choices
                                .first()
                                .ok_or_else(|| "completion returned no choices".to_string())?;
                            streamed_content = choice.message.content.clone();
                            streamed_tool_calls = choice.message.tool_calls.clone();
                            finish_reason = choice.finish_reason;
                        }
                        CompletionOutcome::CoercedSchema(coerced) => {
                            let response = coerced.response;
                            if let Some(usage) = usage_total.as_mut() {
                                usage.prompt_tokens += response.usage.prompt_tokens;
                                usage.completion_tokens += response.usage.completion_tokens;
                                usage.total_tokens += response.usage.total_tokens;
                                if let Some(reasoning_tokens) = response.usage.reasoning_tokens {
                                    usage.reasoning_tokens = Some(
                                        usage.reasoning_tokens.unwrap_or(0) + reasoning_tokens,
                                    );
                                }
                            } else {
                                usage_total = Some(YamlLlmTokenUsage {
                                    prompt_tokens: response.usage.prompt_tokens,
                                    completion_tokens: response.usage.completion_tokens,
                                    total_tokens: response.usage.total_tokens,
                                    reasoning_tokens: response.usage.reasoning_tokens,
                                });
                            }

                            let choice = response
                                .choices
                                .first()
                                .ok_or_else(|| "completion returned no choices".to_string())?;
                            streamed_content = choice.message.content.clone();
                            streamed_tool_calls = choice.message.tool_calls.clone();
                            finish_reason = choice.finish_reason;
                        }
                        CompletionOutcome::Stream(mut stream) => {
                            let mut final_stream_usage: Option<simple_agent_type::response::Usage> =
                                None;
                            let mut delta_filter = StructuredJsonDeltaFilter::default();
                            let include_raw_debug = include_raw_stream_debug_events();
                            let mut json_text_formatter = if request.stream_json_as_text {
                                Some(StreamJsonAsTextFormatter::default())
                            } else {
                                None
                            };
                            let mut tool_calls_by_index: HashMap<u32, ToolCall> = HashMap::new();

                            while let Some(chunk_result) = stream.next().await {
                                if event_sink_is_cancelled(event_sink) {
                                    return Err(workflow_event_sink_cancelled_message().to_string());
                                }

                                let chunk = chunk_result.map_err(|error| error.to_string())?;
                                if let Some(usage) = chunk.usage {
                                    final_stream_usage = Some(usage);
                                }

                                if let Some(choice) = chunk.choices.first() {
                                    if let Some(chunk_finish_reason) = choice.finish_reason {
                                        finish_reason = chunk_finish_reason;
                                    }

                                    if include_raw_debug {
                                        if let Some(reasoning_delta) =
                                            choice.delta.reasoning_content.as_ref()
                                        {
                                            if let Some(sink) = event_sink {
                                                sink.emit(&YamlWorkflowEvent {
                                                    event_type: "node_stream_thinking_delta"
                                                        .to_string(),
                                                    node_id: Some(request.node_id.clone()),
                                                    step_id: Some(request.node_id.clone()),
                                                    node_kind: Some("llm_call".to_string()),
                                                    streamable: Some(true),
                                                    message: None,
                                                    delta: Some(reasoning_delta.clone()),
                                                    token_kind: Some(
                                                        YamlWorkflowTokenKind::Thinking,
                                                    ),
                                                    is_terminal_node_token: Some(
                                                        request.is_terminal_node,
                                                    ),
                                                    elapsed_ms: None,
                                                    metadata: None,
                                                });
                                            }
                                        }
                                    }

                                    if let Some(delta) = choice.delta.content.clone() {
                                        streamed_content.push_str(delta.as_str());
                                        let (output_delta, thinking_delta) = if expects_object {
                                            delta_filter.split(delta.as_str())
                                        } else {
                                            (Some(delta.clone()), None)
                                        };
                                        let rendered_output_delta = if let Some(output_chunk) =
                                            output_delta
                                        {
                                            if let Some(formatter) = json_text_formatter.as_mut() {
                                                formatter.push(output_chunk.as_str());
                                                formatter.emit_if_ready(delta_filter.completed())
                                            } else {
                                                Some(output_chunk)
                                            }
                                        } else {
                                            None
                                        };

                                        if include_raw_debug {
                                            if let Some(sink) = event_sink {
                                                if let Some(raw_thinking_delta) =
                                                    thinking_delta.as_ref()
                                                {
                                                    sink.emit(&YamlWorkflowEvent {
                                                        event_type: "node_stream_thinking_delta"
                                                            .to_string(),
                                                        node_id: Some(request.node_id.clone()),
                                                        step_id: Some(request.node_id.clone()),
                                                        node_kind: Some("llm_call".to_string()),
                                                        streamable: Some(true),
                                                        message: None,
                                                        delta: Some(raw_thinking_delta.clone()),
                                                        token_kind: Some(
                                                            YamlWorkflowTokenKind::Thinking,
                                                        ),
                                                        is_terminal_node_token: Some(
                                                            request.is_terminal_node,
                                                        ),
                                                        elapsed_ms: None,
                                                        metadata: None,
                                                    });
                                                }
                                                if let Some(raw_output_delta) =
                                                    rendered_output_delta.as_ref()
                                                {
                                                    sink.emit(&YamlWorkflowEvent {
                                                        event_type: "node_stream_output_delta"
                                                            .to_string(),
                                                        node_id: Some(request.node_id.clone()),
                                                        step_id: Some(request.node_id.clone()),
                                                        node_kind: Some("llm_call".to_string()),
                                                        streamable: Some(true),
                                                        message: None,
                                                        delta: Some(raw_output_delta.clone()),
                                                        token_kind: Some(
                                                            YamlWorkflowTokenKind::Output,
                                                        ),
                                                        is_terminal_node_token: Some(
                                                            request.is_terminal_node,
                                                        ),
                                                        elapsed_ms: None,
                                                        metadata: None,
                                                    });
                                                }
                                            }
                                        }

                                        if let Some(filtered_delta) = rendered_output_delta {
                                            if let Some(sink) = event_sink {
                                                sink.emit(&YamlWorkflowEvent {
                                                    event_type: "node_stream_delta".to_string(),
                                                    node_id: Some(request.node_id.clone()),
                                                    step_id: Some(request.node_id.clone()),
                                                    node_kind: Some("llm_call".to_string()),
                                                    streamable: Some(true),
                                                    message: None,
                                                    delta: Some(filtered_delta),
                                                    token_kind: Some(YamlWorkflowTokenKind::Output),
                                                    is_terminal_node_token: Some(
                                                        request.is_terminal_node,
                                                    ),
                                                    elapsed_ms: None,
                                                    metadata: None,
                                                });
                                            }
                                        }
                                    }

                                    if let Some(tool_call_deltas) = choice.delta.tool_calls.as_ref()
                                    {
                                        for tool_call_delta in tool_call_deltas {
                                            let entry = tool_calls_by_index
                                                .entry(tool_call_delta.index)
                                                .or_insert_with(|| ToolCall {
                                                    id: tool_call_delta.id.clone().unwrap_or_else(
                                                        || {
                                                            format!(
                                                                "tool_call_{}",
                                                                tool_call_delta.index
                                                            )
                                                        },
                                                    ),
                                                    tool_type: ToolType::Function,
                                                    function:
                                                        simple_agent_type::tool::ToolCallFunction {
                                                            name: String::new(),
                                                            arguments: String::new(),
                                                        },
                                                });

                                            if let Some(id) = tool_call_delta.id.as_ref() {
                                                entry.id = id.clone();
                                            }
                                            if let Some(tool_type) = tool_call_delta.tool_type {
                                                entry.tool_type = tool_type;
                                            }
                                            if let Some(function_delta) =
                                                tool_call_delta.function.as_ref()
                                            {
                                                if let Some(name) = function_delta.name.as_ref() {
                                                    entry.function.name = name.clone();
                                                }
                                                if let Some(arguments) =
                                                    function_delta.arguments.as_ref()
                                                {
                                                    entry.function.arguments.push_str(arguments);
                                                }
                                            }
                                        }
                                    }
                                }

                                if event_sink_is_cancelled(event_sink) {
                                    return Err(workflow_event_sink_cancelled_message().to_string());
                                }
                            }

                            if let Some(usage) = final_stream_usage {
                                if let Some(total) = usage_total.as_mut() {
                                    total.prompt_tokens += usage.prompt_tokens;
                                    total.completion_tokens += usage.completion_tokens;
                                    total.total_tokens += usage.total_tokens;
                                    if let Some(reasoning_tokens) = usage.reasoning_tokens {
                                        total.reasoning_tokens = Some(
                                            total.reasoning_tokens.unwrap_or(0) + reasoning_tokens,
                                        );
                                    }
                                } else {
                                    usage_total = Some(YamlLlmTokenUsage {
                                        prompt_tokens: usage.prompt_tokens,
                                        completion_tokens: usage.completion_tokens,
                                        total_tokens: usage.total_tokens,
                                        reasoning_tokens: usage.reasoning_tokens,
                                    });
                                }
                            }

                            let mut ordered_tool_calls =
                                tool_calls_by_index.into_iter().collect::<Vec<_>>();
                            ordered_tool_calls.sort_by_key(|(index, _)| *index);
                            if !ordered_tool_calls.is_empty() {
                                streamed_tool_calls = Some(
                                    ordered_tool_calls
                                        .into_iter()
                                        .map(|(_, tool_call)| tool_call)
                                        .collect::<Vec<_>>(),
                                );
                            }
                        }
                    }

                    let has_tool_calls = streamed_tool_calls
                        .as_ref()
                        .is_some_and(|calls| !calls.is_empty());
                    if finish_reason != FinishReason::ToolCalls && !has_tool_calls {
                        let payload = if expects_object {
                            parse_streamed_structured_payload(
                                streamed_content.as_str(),
                                request.heal,
                            )
                            .map_err(|error| {
                                format!("failed to parse structured completion JSON: {error}")
                            })?
                            .payload
                        } else {
                            Value::String(streamed_content.clone())
                        };
                        return Ok(YamlLlmExecutionResult {
                            payload,
                            usage: usage_total,
                            ttft_ms: None,
                            tool_calls: tool_traces,
                        });
                    }

                    if roundtrip >= request.max_tool_roundtrips {
                        return Err(format!(
                            "tool call roundtrip limit reached for node '{}' (max={})",
                            request.node_id, request.max_tool_roundtrips
                        ));
                    }

                    let tool_calls: Vec<ToolCall> = streamed_tool_calls.ok_or_else(|| {
                        "finish_reason=tool_calls but no tool calls found".to_string()
                    })?;
                    if tool_calls
                        .iter()
                        .any(|tool_call| tool_call.function.name.trim().is_empty())
                    {
                        return Err("streamed tool call missing function name".to_string());
                    }

                    let assistant_tool_message =
                        Message::assistant(&streamed_content).with_tool_calls(tool_calls.clone());
                    conversation.push(assistant_tool_message);

                    for tool_call in tool_calls {
                        let tool_call_id = tool_call.id.clone();
                        let tool_name = tool_call.function.name.clone();
                        let tool_started = Instant::now();
                        let arguments: Value = serde_json::from_str(&tool_call.function.arguments)
                            .map_err(|error| {
                                format!(
                                    "tool '{}' arguments must be valid JSON: {}",
                                    tool_name, error
                                )
                            })?;
                        let mut tool_span_context: Option<TraceContext> = None;
                        let mut tool_span = if request.trace_sampled {
                            let (span_context, mut span) = workflow_tracer().start_span(
                                "workflow.tool.execute",
                                SpanKind::Node,
                                request.trace_context.as_ref(),
                            );
                            tool_span_context = Some(span_context);
                            span.set_attribute(
                                "trace_id",
                                request.trace_id.as_deref().unwrap_or_default(),
                            );
                            span.set_attribute("node_id", request.node_id.as_str());
                            span.set_attribute("node_kind", "llm_call");
                            span.set_attribute("tool_name", tool_name.as_str());
                            span.set_attribute("tool_call_id", tool_call_id.as_str());
                            let args_for_span =
                                payload_for_tool_trace(request.tool_trace_mode, &arguments)
                                    .to_string();
                            span.set_attribute("tool_arguments", args_for_span.as_str());
                            Some(span)
                        } else {
                            None
                        };

                        if request.tool_trace_mode != YamlToolTraceMode::Off {
                            if let Some(sink) = event_sink {
                                sink.emit(&YamlWorkflowEvent {
                                    event_type: "node_tool_call_requested".to_string(),
                                    node_id: Some(request.node_id.clone()),
                                    step_id: Some(request.node_id.clone()),
                                    node_kind: Some("llm_call".to_string()),
                                    streamable: Some(false),
                                    message: Some(format!(
                                        "tool call requested: {}",
                                        tool_name
                                    )),
                                    delta: None,
                                    token_kind: None,
                                    is_terminal_node_token: None,
                                    elapsed_ms: None,
                                    metadata: Some(json!({
                                        "tool_call_id": tool_call_id.clone(),
                                        "tool_name": tool_name.clone(),
                                        "arguments": payload_for_tool_trace(request.tool_trace_mode, &arguments),
                                    })),
                                });
                            }
                        }

                        let Some(tool_config) = request
                            .tools
                            .iter()
                            .find(|tool| tool.definition.function.name == tool_name)
                        else {
                            return Err(format!("model requested unknown tool '{}'", tool_name));
                        };

                        let tool_output_result = if tool_name == "run_workflow_graph" {
                            execute_subworkflow_tool_call(
                                &arguments,
                                &request.execution_context,
                                self.client,
                                self.custom_worker,
                                &self.run_options,
                                tool_span_context.as_ref(),
                                request.trace_id.as_deref(),
                            )
                            .await
                        } else if let Some(custom_worker) = self.custom_worker {
                            custom_worker
                                .execute(
                                    tool_name.as_str(),
                                    &arguments,
                                    request.email_text.as_str(),
                                    &request.execution_context,
                                )
                                .await
                        } else {
                            Err(format!(
                                "tool '{}' requested but no custom worker executor is configured",
                                tool_name
                            ))
                        };

                        let tool_output = match tool_output_result {
                            Ok(output) => output,
                            Err(message) => {
                                let elapsed_ms = tool_started.elapsed().as_millis();
                                if let Some(span) = tool_span.as_mut() {
                                    span.add_event("workflow.tool.execute.error");
                                    span.set_attribute("tool_status", "error");
                                    span.set_attribute("tool_error", message.as_str());
                                    span.set_attribute(
                                        "elapsed_ms",
                                        elapsed_ms.to_string().as_str(),
                                    );
                                }
                                if request.tool_trace_mode != YamlToolTraceMode::Off {
                                    if let Some(sink) = event_sink {
                                        sink.emit(&YamlWorkflowEvent {
                                            event_type: "node_tool_call_failed".to_string(),
                                            node_id: Some(request.node_id.clone()),
                                            step_id: Some(request.node_id.clone()),
                                            node_kind: Some("llm_call".to_string()),
                                            streamable: Some(false),
                                            message: Some(message.clone()),
                                            delta: None,
                                            token_kind: None,
                                            is_terminal_node_token: None,
                                            elapsed_ms: Some(elapsed_ms),
                                            metadata: Some(json!({
                                                "tool_call_id": tool_call_id.clone(),
                                                "tool_name": tool_name.clone(),
                                            })),
                                        });
                                    }
                                }
                                tool_traces.push(YamlToolCallTrace {
                                    id: tool_call_id.clone(),
                                    name: tool_name.clone(),
                                    arguments,
                                    output: None,
                                    status: "error".to_string(),
                                    elapsed_ms,
                                    error: Some(message.clone()),
                                });
                                if let Some(span) = tool_span.take() {
                                    span.end();
                                }
                                return Err(format!("tool '{}' failed: {}", tool_name, message));
                            }
                        };

                        if let Some(output_schema) = tool_config.output_schema.as_ref() {
                            validate_schema_instance(output_schema, &tool_output).map_err(
                                |message| {
                                    format!(
                                        "tool '{}' output failed schema validation: {}",
                                        tool_name, message
                                    )
                                },
                            )?;
                        }

                        let elapsed_ms = tool_started.elapsed().as_millis();
                        if let Some(span) = tool_span.as_mut() {
                            span.add_event("workflow.tool.execute.completed");
                            span.set_attribute("tool_status", "ok");
                            span.set_attribute("elapsed_ms", elapsed_ms.to_string().as_str());
                            let output_for_span =
                                payload_for_tool_trace(request.tool_trace_mode, &tool_output)
                                    .to_string();
                            span.set_attribute("tool_output", output_for_span.as_str());
                        }
                        if request.tool_trace_mode != YamlToolTraceMode::Off {
                            if let Some(sink) = event_sink {
                                sink.emit(&YamlWorkflowEvent {
                                    event_type: "node_tool_call_completed".to_string(),
                                    node_id: Some(request.node_id.clone()),
                                    step_id: Some(request.node_id.clone()),
                                    node_kind: Some("llm_call".to_string()),
                                    streamable: Some(false),
                                    message: Some(format!(
                                        "tool call completed: {}",
                                        tool_name
                                    )),
                                    delta: None,
                                    token_kind: None,
                                    is_terminal_node_token: None,
                                    elapsed_ms: Some(elapsed_ms),
                                    metadata: Some(json!({
                                        "tool_call_id": tool_call_id.clone(),
                                        "tool_name": tool_name.clone(),
                                        "arguments": payload_for_tool_trace(request.tool_trace_mode, &arguments),
                                        "output": payload_for_tool_trace(request.tool_trace_mode, &tool_output),
                                    })),
                                });
                            }
                        }

                        tool_traces.push(YamlToolCallTrace {
                            id: tool_call_id.clone(),
                            name: tool_name.clone(),
                            arguments: arguments.clone(),
                            output: Some(tool_output.clone()),
                            status: "ok".to_string(),
                            elapsed_ms,
                            error: None,
                        });

                        conversation.push(Message::tool(
                            serde_json::to_string(&tool_output).map_err(|error| {
                                format!("failed to serialize tool output: {error}")
                            })?,
                            tool_call_id,
                        ));
                        if let Some(span) = tool_span.take() {
                            span.end();
                        }
                    }

                    if request.tool_trace_mode != YamlToolTraceMode::Off {
                        if let Some(sink) = event_sink {
                            sink.emit(&YamlWorkflowEvent {
                                event_type: "node_tool_roundtrip_completed".to_string(),
                                node_id: Some(request.node_id.clone()),
                                step_id: Some(request.node_id.clone()),
                                node_kind: Some("llm_call".to_string()),
                                streamable: Some(false),
                                message: Some(format!(
                                    "tool roundtrip {} completed",
                                    roundtrip + 1
                                )),
                                delta: None,
                                token_kind: None,
                                is_terminal_node_token: None,
                                elapsed_ms: None,
                                metadata: Some(json!({
                                    "roundtrip": roundtrip + 1,
                                    "max_tool_roundtrips": request.max_tool_roundtrips,
                                })),
                            });
                        }
                    }
                }

                return Err(format!(
                    "tool-enabled llm_call '{}' exhausted loop without final payload",
                    request.node_id
                ));
            }

            let mut builder = CompletionRequest::builder()
                .model(&request.model)
                .messages(messages);

            if request.heal && !request.stream && expects_object {
                builder = builder.json_schema("workflow_step", request.schema.clone());
            }

            if request.stream {
                builder = builder.stream(true);
            }

            let completion_request = builder
                .build()
                .map_err(|error| format!("failed to build completion request: {error}"))?;

            let completion_options = if request.heal && !request.stream && expects_object {
                CompletionOptions {
                    mode: CompletionMode::HealedJson,
                }
            } else {
                CompletionOptions::default()
            };

            let outcome = self
                .client
                .complete(&completion_request, completion_options)
                .await
                .map_err(|error| error.to_string())?;

            match outcome {
                CompletionOutcome::Stream(mut stream) => {
                    let mut aggregated = String::new();
                    let mut final_stream_usage: Option<simple_agent_type::response::Usage> = None;
                    let stream_started = Instant::now();
                    let mut ttft_ms: Option<u128> = None;
                    let mut delta_filter = StructuredJsonDeltaFilter::default();
                    let include_raw_debug = include_raw_stream_debug_events();
                    let mut json_text_formatter = if request.stream_json_as_text {
                        Some(StreamJsonAsTextFormatter::default())
                    } else {
                        None
                    };
                    while let Some(chunk_result) = stream.next().await {
                        if event_sink_is_cancelled(event_sink) {
                            return Err(workflow_event_sink_cancelled_message().to_string());
                        }
                        let chunk = chunk_result.map_err(|error| error.to_string())?;
                        if let Some(usage) = chunk.usage {
                            final_stream_usage = Some(usage);
                        }
                        if let Some(choice) = chunk.choices.first() {
                            if ttft_ms.is_none()
                                && (choice
                                    .delta
                                    .content
                                    .as_ref()
                                    .is_some_and(|delta| !delta.is_empty())
                                    || choice
                                        .delta
                                        .reasoning_content
                                        .as_ref()
                                        .is_some_and(|delta| !delta.is_empty()))
                            {
                                ttft_ms = Some(stream_started.elapsed().as_millis());
                            }
                            if include_raw_debug {
                                if let Some(reasoning_delta) =
                                    choice.delta.reasoning_content.as_ref()
                                {
                                    if let Some(sink) = event_sink {
                                        sink.emit(&YamlWorkflowEvent {
                                            event_type: "node_stream_thinking_delta".to_string(),
                                            node_id: Some(request.node_id.clone()),
                                            step_id: Some(request.node_id.clone()),
                                            node_kind: Some("llm_call".to_string()),
                                            streamable: Some(true),
                                            message: None,
                                            delta: Some(reasoning_delta.clone()),
                                            token_kind: Some(YamlWorkflowTokenKind::Thinking),
                                            is_terminal_node_token: Some(request.is_terminal_node),
                                            elapsed_ms: None,
                                            metadata: None,
                                        });
                                    }
                                }
                            }
                            if let Some(delta) = choice.delta.content.clone() {
                                aggregated.push_str(delta.as_str());
                                let (output_delta, thinking_delta) = if expects_object {
                                    delta_filter.split(delta.as_str())
                                } else {
                                    (Some(delta.clone()), None)
                                };
                                let rendered_output_delta = if let Some(output_chunk) = output_delta
                                {
                                    if let Some(formatter) = json_text_formatter.as_mut() {
                                        formatter.push(output_chunk.as_str());
                                        formatter.emit_if_ready(delta_filter.completed())
                                    } else {
                                        Some(output_chunk)
                                    }
                                } else {
                                    None
                                };
                                if include_raw_debug {
                                    if let Some(sink) = event_sink {
                                        if let Some(raw_thinking_delta) = thinking_delta.as_ref() {
                                            sink.emit(&YamlWorkflowEvent {
                                                event_type: "node_stream_thinking_delta"
                                                    .to_string(),
                                                node_id: Some(request.node_id.clone()),
                                                step_id: Some(request.node_id.clone()),
                                                node_kind: Some("llm_call".to_string()),
                                                streamable: Some(true),
                                                message: None,
                                                delta: Some(raw_thinking_delta.clone()),
                                                token_kind: Some(YamlWorkflowTokenKind::Thinking),
                                                is_terminal_node_token: Some(
                                                    request.is_terminal_node,
                                                ),
                                                elapsed_ms: None,
                                                metadata: None,
                                            });
                                        }
                                        if let Some(raw_output_delta) =
                                            rendered_output_delta.as_ref()
                                        {
                                            sink.emit(&YamlWorkflowEvent {
                                                event_type: "node_stream_output_delta".to_string(),
                                                node_id: Some(request.node_id.clone()),
                                                step_id: Some(request.node_id.clone()),
                                                node_kind: Some("llm_call".to_string()),
                                                streamable: Some(true),
                                                message: None,
                                                delta: Some(raw_output_delta.clone()),
                                                token_kind: Some(YamlWorkflowTokenKind::Output),
                                                is_terminal_node_token: Some(
                                                    request.is_terminal_node,
                                                ),
                                                elapsed_ms: None,
                                                metadata: None,
                                            });
                                        }
                                    }
                                }
                                if let Some(filtered_delta) = rendered_output_delta {
                                    if let Some(sink) = event_sink {
                                        sink.emit(&YamlWorkflowEvent {
                                            event_type: "node_stream_delta".to_string(),
                                            node_id: Some(request.node_id.clone()),
                                            step_id: Some(request.node_id.clone()),
                                            node_kind: Some("llm_call".to_string()),
                                            streamable: Some(true),
                                            message: None,
                                            delta: Some(filtered_delta),
                                            token_kind: Some(YamlWorkflowTokenKind::Output),
                                            is_terminal_node_token: Some(request.is_terminal_node),
                                            elapsed_ms: None,
                                            metadata: None,
                                        });
                                    }
                                }
                            }
                        }

                        if event_sink_is_cancelled(event_sink) {
                            return Err(workflow_event_sink_cancelled_message().to_string());
                        }
                    }

                    let payload = if expects_object {
                        let resolved =
                            parse_streamed_structured_payload(aggregated.as_str(), request.heal)?;
                        if let Some(confidence) = resolved.heal_confidence {
                            if let Some(sink) = event_sink {
                                sink.emit(&YamlWorkflowEvent {
                                    event_type: "node_healed".to_string(),
                                    node_id: Some(request.node_id.clone()),
                                    step_id: Some(request.node_id.clone()),
                                    node_kind: Some("llm_call".to_string()),
                                    streamable: Some(true),
                                    message: Some(format!(
                                        "healed streamed structured response confidence={confidence}"
                                    )),
                                    delta: None,
                                    token_kind: None,
                                    is_terminal_node_token: None,
                                    elapsed_ms: None,
                                    metadata: None,
                                });
                            }
                        }
                        resolved.payload
                    } else {
                        Value::String(aggregated)
                    };

                    Ok(YamlLlmExecutionResult {
                        payload,
                        usage: final_stream_usage.map(|usage| YamlLlmTokenUsage {
                            prompt_tokens: usage.prompt_tokens,
                            completion_tokens: usage.completion_tokens,
                            total_tokens: usage.total_tokens,
                            reasoning_tokens: usage.reasoning_tokens,
                        }),
                        ttft_ms,
                        tool_calls: Vec::new(),
                    })
                }
                CompletionOutcome::Response(response) => {
                    let payload = if expects_object {
                        let content = response
                            .content()
                            .ok_or_else(|| "completion returned empty content".to_string())?;
                        serde_json::from_str(content).map_err(|error| {
                            format!("failed to parse structured completion JSON: {error}")
                        })?
                    } else {
                        Value::String(response.content().unwrap_or_default().to_string())
                    };

                    Ok(YamlLlmExecutionResult {
                        payload,
                        usage: Some(YamlLlmTokenUsage {
                            prompt_tokens: response.usage.prompt_tokens,
                            completion_tokens: response.usage.completion_tokens,
                            total_tokens: response.usage.total_tokens,
                            reasoning_tokens: response.usage.reasoning_tokens,
                        }),
                        ttft_ms: None,
                        tool_calls: Vec::new(),
                    })
                }
                CompletionOutcome::HealedJson(healed) => {
                    if !expects_object {
                        return Err(
                            "healed json outcome is unsupported for non-object schema".to_string()
                        );
                    }
                    if let Some(sink) = event_sink {
                        sink.emit(&YamlWorkflowEvent {
                            event_type: "node_healed".to_string(),
                            node_id: Some(request.node_id.clone()),
                            step_id: Some(request.node_id.clone()),
                            node_kind: Some("llm_call".to_string()),
                            streamable: Some(request.stream),
                            message: Some(format!(
                                "healed structured response confidence={}",
                                healed.parsed.confidence
                            )),
                            delta: None,
                            token_kind: None,
                            is_terminal_node_token: None,
                            elapsed_ms: None,
                            metadata: None,
                        });
                    }
                    Ok(YamlLlmExecutionResult {
                        payload: healed.parsed.value,
                        usage: Some(YamlLlmTokenUsage {
                            prompt_tokens: healed.response.usage.prompt_tokens,
                            completion_tokens: healed.response.usage.completion_tokens,
                            total_tokens: healed.response.usage.total_tokens,
                            reasoning_tokens: healed.response.usage.reasoning_tokens,
                        }),
                        ttft_ms: None,
                        tool_calls: Vec::new(),
                    })
                }
                CompletionOutcome::CoercedSchema(coerced) => {
                    if !expects_object {
                        return Err(
                            "coerced schema outcome is unsupported for non-object schema"
                                .to_string(),
                        );
                    }
                    Ok(YamlLlmExecutionResult {
                        payload: coerced.coerced.value,
                        usage: Some(YamlLlmTokenUsage {
                            prompt_tokens: coerced.response.usage.prompt_tokens,
                            completion_tokens: coerced.response.usage.completion_tokens,
                            total_tokens: coerced.response.usage.total_tokens,
                            reasoning_tokens: coerced.response.usage.reasoning_tokens,
                        }),
                        ttft_ms: None,
                        tool_calls: Vec::new(),
                    })
                }
            }
        }
    }

    let executor = BorrowedClientExecutor {
        client,
        custom_worker,
        run_options: options.clone(),
    };
    run_workflow_yaml_with_custom_worker_and_events_and_options(
        workflow,
        workflow_input,
        &executor,
        custom_worker,
        event_sink,
        options,
    )
    .await
}

pub async fn run_email_workflow_yaml_with_client_and_custom_worker_and_events(
    workflow: &YamlWorkflow,
    email_text: &str,
    client: &SimpleAgentsClient,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
    event_sink: Option<&dyn YamlWorkflowEventSink>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let workflow_input = json!({ "email_text": email_text });
    run_workflow_yaml_with_client_and_custom_worker_and_events(
        workflow,
        &workflow_input,
        client,
        custom_worker,
        event_sink,
    )
    .await
}

pub async fn run_workflow_yaml(
    workflow: &YamlWorkflow,
    workflow_input: &Value,
    executor: &dyn YamlWorkflowLlmExecutor,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    run_workflow_yaml_with_custom_worker_and_events(workflow, workflow_input, executor, None, None)
        .await
}

pub async fn run_email_workflow_yaml(
    workflow: &YamlWorkflow,
    email_text: &str,
    executor: &dyn YamlWorkflowLlmExecutor,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let workflow_input = json!({ "email_text": email_text });
    run_workflow_yaml(workflow, &workflow_input, executor).await
}

pub async fn run_workflow_yaml_with_custom_worker(
    workflow: &YamlWorkflow,
    workflow_input: &Value,
    executor: &dyn YamlWorkflowLlmExecutor,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    run_workflow_yaml_with_custom_worker_and_events(
        workflow,
        workflow_input,
        executor,
        custom_worker,
        None,
    )
    .await
}

pub async fn run_email_workflow_yaml_with_custom_worker(
    workflow: &YamlWorkflow,
    email_text: &str,
    executor: &dyn YamlWorkflowLlmExecutor,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let workflow_input = json!({ "email_text": email_text });
    run_workflow_yaml_with_custom_worker(workflow, &workflow_input, executor, custom_worker).await
}

pub async fn run_workflow_yaml_with_custom_worker_and_events(
    workflow: &YamlWorkflow,
    workflow_input: &Value,
    executor: &dyn YamlWorkflowLlmExecutor,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
    event_sink: Option<&dyn YamlWorkflowEventSink>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    run_workflow_yaml_with_custom_worker_and_events_and_options(
        workflow,
        workflow_input,
        executor,
        custom_worker,
        event_sink,
        &YamlWorkflowRunOptions::default(),
    )
    .await
}

pub async fn run_workflow_yaml_with_custom_worker_and_events_and_options(
    workflow: &YamlWorkflow,
    workflow_input: &Value,
    executor: &dyn YamlWorkflowLlmExecutor,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
    event_sink: Option<&dyn YamlWorkflowEventSink>,
    options: &YamlWorkflowRunOptions,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    if !workflow_input.is_object() {
        return Err(YamlWorkflowRunError::InvalidInput {
            message: "workflow input must be a JSON object".to_string(),
        });
    }

    validate_sample_rate(options.telemetry.sample_rate)?;

    let email_text = workflow_input
        .get("email_text")
        .and_then(Value::as_str)
        .unwrap_or_default();

    let diagnostics = verify_yaml_workflow(workflow);
    let errors: Vec<YamlWorkflowDiagnostic> = diagnostics
        .iter()
        .filter(|d| d.severity == YamlWorkflowDiagnosticSeverity::Error)
        .cloned()
        .collect();
    if !errors.is_empty() {
        return Err(YamlWorkflowRunError::Validation {
            diagnostics_count: errors.len(),
            diagnostics: errors,
        });
    }

    if let Some(output) =
        try_run_yaml_via_ir_runtime(workflow, workflow_input, executor, custom_worker, options)
            .await?
    {
        return Ok(output);
    }

    let parent_trace_context = trace_context_from_options(options);
    let telemetry_context = resolve_telemetry_context(options, parent_trace_context.as_ref());

    let tracer = workflow_tracer();
    let mut workflow_span_context: Option<TraceContext> = None;
    let mut workflow_span = if telemetry_context.sampled {
        let (span_context, mut span) = tracer.start_span(
            "workflow.run",
            SpanKind::Workflow,
            parent_trace_context.as_ref(),
        );
        if let Some(trace_id_value) = telemetry_context.trace_id.as_deref() {
            span.set_attribute("trace_id", trace_id_value);
        }
        apply_trace_tenant_attributes(span.as_mut(), options);
        workflow_span_context = Some(span_context);
        Some(span)
    } else {
        None
    };

    if workflow.nodes.is_empty() {
        return Err(YamlWorkflowRunError::EmptyNodes {
            workflow_id: workflow.id.clone(),
        });
    }

    let node_map: HashMap<&str, &YamlNode> = workflow
        .nodes
        .iter()
        .map(|node| (node.id.as_str(), node))
        .collect();
    if !node_map.contains_key(workflow.entry_node.as_str()) {
        return Err(YamlWorkflowRunError::MissingEntry {
            entry_node: workflow.entry_node.clone(),
        });
    }

    let edge_map: HashMap<&str, &str> = workflow
        .edges
        .iter()
        .map(|edge| (edge.from.as_str(), edge.to.as_str()))
        .collect();

    let mut current = workflow.entry_node.clone();
    let mut trace = Vec::new();
    let mut outputs: BTreeMap<String, Value> = BTreeMap::new();
    let mut globals = serde_json::Map::new();
    let mut step_timings = Vec::new();
    let mut llm_node_metrics: BTreeMap<String, YamlLlmNodeMetrics> = BTreeMap::new();
    let mut llm_node_models: BTreeMap<String, String> = BTreeMap::new();
    let mut token_totals = YamlTokenTotals::default();
    let mut workflow_ttft_ms: Option<u128> = None;
    let started = Instant::now();

    if let Some(sink) = event_sink {
        sink.emit(&YamlWorkflowEvent {
            event_type: "workflow_started".to_string(),
            node_id: None,
            step_id: None,
            node_kind: None,
            streamable: None,
            message: Some(format!("workflow_id={}", workflow.id)),
            delta: None,
            token_kind: None,
            is_terminal_node_token: None,
            elapsed_ms: Some(0),
            metadata: None,
        });
    }

    if event_sink_is_cancelled(event_sink) {
        return Err(YamlWorkflowRunError::EventSinkCancelled {
            message: workflow_event_sink_cancelled_message().to_string(),
        });
    }

    loop {
        if event_sink_is_cancelled(event_sink) {
            return Err(YamlWorkflowRunError::EventSinkCancelled {
                message: workflow_event_sink_cancelled_message().to_string(),
            });
        }

        let node =
            *node_map
                .get(current.as_str())
                .ok_or_else(|| YamlWorkflowRunError::MissingNode {
                    node_id: current.clone(),
                })?;

        trace.push(node.id.clone());
        let step_started = Instant::now();

        let mut node_span_context: Option<TraceContext> = None;
        let mut node_span = if telemetry_context.sampled {
            let (span_context, mut span) = tracer.start_span(
                "workflow.node.execute",
                SpanKind::Node,
                workflow_span_context.as_ref(),
            );
            node_span_context = Some(span_context);
            span.set_attribute(
                "trace_id",
                telemetry_context.trace_id.as_deref().unwrap_or_default(),
            );
            span.set_attribute("node_id", node.id.as_str());
            span.set_attribute("node_kind", node.kind_name());
            Some(span)
        } else {
            None
        };

        let node_streamable = node
            .node_type
            .llm_call
            .as_ref()
            .map(|llm| llm.stream.unwrap_or(false) && !llm.heal.unwrap_or(false));
        let workflow_elapsed_before_node_ms = started.elapsed().as_millis();

        if let Some(sink) = event_sink {
            sink.emit(&YamlWorkflowEvent {
                event_type: "node_started".to_string(),
                node_id: Some(node.id.clone()),
                step_id: Some(node.id.clone()),
                node_kind: Some(node.kind_name().to_string()),
                streamable: node_streamable,
                message: if node_streamable == Some(false) {
                    Some("Node is not streamable; status events only".to_string())
                } else {
                    None
                },
                delta: None,
                token_kind: None,
                is_terminal_node_token: None,
                elapsed_ms: Some(workflow_elapsed_before_node_ms),
                metadata: None,
            });
        }

        if event_sink_is_cancelled(event_sink) {
            return Err(YamlWorkflowRunError::EventSinkCancelled {
                message: workflow_event_sink_cancelled_message().to_string(),
            });
        }

        let mut node_usage: Option<YamlLlmTokenUsage> = None;
        let mut node_model_name: Option<String> = None;
        let is_terminal_node = !edge_map.contains_key(node.id.as_str());
        let next = if let Some(llm) = &node.node_type.llm_call {
            let prompt_template = node
                .config
                .as_ref()
                .and_then(|cfg| cfg.prompt.as_deref())
                .unwrap_or_default();
            let context = json!({
                "input": workflow_input,
                "nodes": outputs,
                "globals": Value::Object(globals.clone())
            });
            let messages = if let Some(path) = llm.messages_path.as_deref() {
                Some(
                    parse_messages_from_context(path, &context).map_err(|message| {
                        YamlWorkflowRunError::Llm {
                            node_id: node.id.clone(),
                            message,
                        }
                    })?,
                )
            } else {
                None
            };
            let prompt_bindings = collect_template_bindings(prompt_template, &context);
            let prompt = interpolate_template(prompt_template, &context);
            let schema = llm_output_schema_for_node(node);

            let request = YamlLlmExecutionRequest {
                node_id: node.id.clone(),
                is_terminal_node,
                stream_json_as_text: llm.stream_json_as_text.unwrap_or(false),
                model: resolve_requested_model(options.model.as_deref(), &llm.model),
                messages,
                append_prompt_as_user: llm.append_prompt_as_user.unwrap_or(true),
                prompt,
                prompt_template: prompt_template.to_string(),
                prompt_bindings,
                schema,
                stream: llm.stream.unwrap_or(false),
                heal: llm.heal.unwrap_or(false),
                tools: normalize_llm_tools(llm).map_err(|message| YamlWorkflowRunError::Llm {
                    node_id: node.id.clone(),
                    message,
                })?,
                tool_choice: normalize_tool_choice(llm.tool_choice.clone()).map_err(|message| {
                    YamlWorkflowRunError::Llm {
                        node_id: node.id.clone(),
                        message,
                    }
                })?,
                max_tool_roundtrips: llm.max_tool_roundtrips.unwrap_or(1),
                tool_calls_global_key: llm.tool_calls_global_key.clone(),
                tool_trace_mode: options.telemetry.tool_trace_mode,
                execution_context: context.clone(),
                email_text: email_text.to_string(),
                trace_id: telemetry_context.trace_id.clone(),
                trace_context: node_span_context.clone(),
                trace_sampled: telemetry_context.sampled,
            };

            if let Some(span) = node_span.as_mut() {
                span.set_attribute(
                    "node_input",
                    payload_for_span(options.telemetry.payload_mode, &context).as_str(),
                );
            }

            if let Some(sink) = event_sink {
                sink.emit(&YamlWorkflowEvent {
                    event_type: "node_llm_input_resolved".to_string(),
                    node_id: Some(node.id.clone()),
                    step_id: Some(node.id.clone()),
                    node_kind: Some("llm_call".to_string()),
                    streamable: Some(request.stream),
                    message: Some("resolved llm input for telemetry".to_string()),
                    delta: None,
                    token_kind: None,
                    is_terminal_node_token: None,
                    elapsed_ms: Some(started.elapsed().as_millis()),
                    metadata: Some(json!({
                        "model": request.model.clone(),
                        "stream_requested": request.stream,
                        "stream_json_as_text": request.stream_json_as_text,
                        "heal_requested": request.heal,
                        "effective_stream": request.stream,
                        "prompt_template": request.prompt_template.clone(),
                        "prompt": request.prompt.clone(),
                        "schema": request.schema.clone(),
                        "bindings": request.prompt_bindings.clone(),
                        "tools_count": request.tools.len(),
                        "max_tool_roundtrips": request.max_tool_roundtrips,
                    })),
                });
            }

            node_model_name = Some(request.model.clone());
            llm_node_models.insert(node.id.clone(), request.model.clone());

            if event_sink_is_cancelled(event_sink) {
                return Err(YamlWorkflowRunError::EventSinkCancelled {
                    message: workflow_event_sink_cancelled_message().to_string(),
                });
            }

            let llm_result = executor
                .complete_structured(request, event_sink)
                .await
                .map_err(|message| YamlWorkflowRunError::Llm {
                    node_id: node.id.clone(),
                    message,
                })?;

            if let Some(usage) = llm_result.usage.as_ref() {
                token_totals.add_usage(usage);
            }
            if workflow_ttft_ms.is_none() {
                workflow_ttft_ms = llm_result
                    .ttft_ms
                    .map(|node_ttft_ms| workflow_elapsed_before_node_ms + node_ttft_ms);
            }
            node_usage = llm_result.usage;

            let payload = llm_result.payload;
            let tool_calls = llm_result.tool_calls;

            let mut node_output = json!({ "output": payload });
            if !tool_calls.is_empty() {
                if let Some(output_obj) = node_output.as_object_mut() {
                    output_obj.insert("tool_calls".to_string(), json!(tool_calls));
                }
            }
            outputs.insert(node.id.clone(), node_output);
            if let Some(span) = node_span.as_mut() {
                if let Some(output_payload) = outputs.get(node.id.as_str()) {
                    span.set_attribute(
                        "node_output",
                        payload_for_span(options.telemetry.payload_mode, output_payload).as_str(),
                    );
                }
            }
            apply_set_globals(node, &outputs, workflow_input, &mut globals);
            apply_update_globals(node, &outputs, workflow_input, &mut globals);
            if let Some(global_key) = llm.tool_calls_global_key.as_ref() {
                if let Some(node_tool_calls) = outputs
                    .get(node.id.as_str())
                    .and_then(|value| value.get("tool_calls"))
                    .cloned()
                {
                    globals.insert(global_key.clone(), node_tool_calls);
                }
            }
            edge_map
                .get(node.id.as_str())
                .map(|value| value.to_string())
        } else if let Some(switch) = &node.node_type.switch {
            let context = json!({
                "input": workflow_input,
                "nodes": outputs,
                "globals": Value::Object(globals.clone())
            });
            let mut chosen = Some(switch.default.clone());
            for branch in &switch.branches {
                if evaluate_switch_condition(branch.condition.as_str(), &context)? {
                    chosen = Some(branch.target.clone());
                    break;
                }
            }
            let chosen = chosen.ok_or_else(|| YamlWorkflowRunError::InvalidSwitchTarget {
                node_id: node.id.clone(),
            })?;
            Some(chosen)
        } else if let Some(custom) = &node.node_type.custom_worker {
            let payload = node
                .config
                .as_ref()
                .and_then(|cfg| cfg.payload.as_ref())
                .cloned()
                .unwrap_or_else(|| json!({}));
            let context = json!({
                "input": workflow_input,
                "nodes": outputs,
                "globals": Value::Object(globals.clone())
            });

            if let Some(span) = node_span.as_mut() {
                span.set_attribute("handler_name", custom.handler.as_str());
                span.set_attribute(
                    "node_input",
                    payload_for_span(options.telemetry.payload_mode, &payload).as_str(),
                );
            }

            let mut handler_span_context: Option<TraceContext> = None;
            let mut handler_span = if telemetry_context.sampled {
                let (span_context, mut span) = tracer.start_span(
                    "handler.invoke",
                    SpanKind::Node,
                    workflow_span_context.as_ref(),
                );
                handler_span_context = Some(span_context);
                span.set_attribute(
                    "trace_id",
                    telemetry_context.trace_id.as_deref().unwrap_or_default(),
                );
                span.set_attribute("handler_name", custom.handler.as_str());
                apply_trace_tenant_attributes(span.as_mut(), options);
                Some(span)
            } else {
                None
            };

            let worker_trace_context = merged_trace_context_for_worker(
                handler_span_context.as_ref(),
                telemetry_context.trace_id.as_deref(),
                options,
            );
            let worker_context = custom_worker_context_with_trace(
                &context,
                &worker_trace_context,
                &options.trace.tenant,
            );

            let worker_output_result = if let Some(custom_worker_executor) = custom_worker {
                custom_worker_executor
                    .execute(
                        custom.handler.as_str(),
                        &payload,
                        email_text,
                        &worker_context,
                    )
                    .await
                    .map_err(|message| YamlWorkflowRunError::CustomWorker {
                        node_id: node.id.clone(),
                        message,
                    })
            } else {
                mock_custom_worker_output(custom.handler.as_str(), &payload)
            };

            if let Some(span) = handler_span.take() {
                span.end();
            }

            let worker_output = worker_output_result?;

            outputs.insert(node.id.clone(), json!({ "output": worker_output }));
            if let Some(span) = node_span.as_mut() {
                if let Some(output_payload) = outputs.get(node.id.as_str()) {
                    span.set_attribute(
                        "node_output",
                        payload_for_span(options.telemetry.payload_mode, output_payload).as_str(),
                    );
                }
            }
            apply_set_globals(node, &outputs, workflow_input, &mut globals);
            apply_update_globals(node, &outputs, workflow_input, &mut globals);
            edge_map
                .get(node.id.as_str())
                .map(|value| value.to_string())
        } else {
            return Err(YamlWorkflowRunError::UnsupportedNodeType {
                node_id: node.id.clone(),
            });
        };

        let node_kind = node.kind_name().to_string();
        let elapsed_ms = step_started.elapsed().as_millis();
        step_timings.push(YamlStepTiming {
            node_id: node.id.clone(),
            node_kind,
            model_name: node_model_name,
            elapsed_ms,
            prompt_tokens: node_usage.as_ref().map(|usage| usage.prompt_tokens),
            completion_tokens: node_usage.as_ref().map(|usage| usage.completion_tokens),
            total_tokens: node_usage.as_ref().map(|usage| usage.total_tokens),
            reasoning_tokens: node_usage.as_ref().and_then(|usage| usage.reasoning_tokens),
            tokens_per_second: node_usage
                .as_ref()
                .map(|usage| completion_tokens_per_second(usage.completion_tokens, elapsed_ms)),
        });

        if let Some(usage) = node_usage.as_ref() {
            llm_node_metrics.insert(
                node.id.clone(),
                YamlLlmNodeMetrics {
                    elapsed_ms,
                    prompt_tokens: usage.prompt_tokens,
                    completion_tokens: usage.completion_tokens,
                    total_tokens: usage.total_tokens,
                    reasoning_tokens: usage.reasoning_tokens,
                    tokens_per_second: completion_tokens_per_second(
                        usage.completion_tokens,
                        elapsed_ms,
                    ),
                },
            );
        }

        if let Some(mut span) = node_span.take() {
            span.set_attribute("elapsed_ms", elapsed_ms.to_string().as_str());
            span.add_event("node_completed");
            span.end();
        }

        if let Some(sink) = event_sink {
            sink.emit(&YamlWorkflowEvent {
                event_type: "node_completed".to_string(),
                node_id: Some(node.id.clone()),
                step_id: Some(node.id.clone()),
                node_kind: Some(node.kind_name().to_string()),
                streamable: node_streamable,
                message: None,
                delta: None,
                token_kind: None,
                is_terminal_node_token: None,
                elapsed_ms: Some(elapsed_ms),
                metadata: None,
            });
        }

        if event_sink_is_cancelled(event_sink) {
            return Err(YamlWorkflowRunError::EventSinkCancelled {
                message: workflow_event_sink_cancelled_message().to_string(),
            });
        }

        if let Some(next) = next {
            current = next;
            continue;
        }
        break;
    }

    let terminal_node = trace
        .last()
        .cloned()
        .ok_or_else(|| YamlWorkflowRunError::EmptyNodes {
            workflow_id: workflow.id.clone(),
        })?;

    let terminal_output = outputs
        .get(terminal_node.as_str())
        .and_then(|value| value.get("output"))
        .cloned();

    let total_elapsed_ms = started.elapsed().as_millis();
    let output = YamlWorkflowRunOutput {
        workflow_id: workflow.id.clone(),
        entry_node: workflow.entry_node.clone(),
        email_text: email_text.to_string(),
        trace,
        outputs,
        terminal_node,
        terminal_output,
        step_timings,
        llm_node_metrics,
        llm_node_models,
        total_elapsed_ms,
        ttft_ms: workflow_ttft_ms,
        total_input_tokens: token_totals.input_tokens,
        total_output_tokens: token_totals.output_tokens,
        total_tokens: token_totals.total_tokens,
        total_reasoning_tokens: token_totals.reasoning_tokens,
        tokens_per_second: token_totals.tokens_per_second(total_elapsed_ms),
        trace_id: telemetry_context.trace_id.clone(),
        metadata: telemetry_context.trace_id.as_ref().map(|value| {
            workflow_metadata_with_trace(
                options,
                value,
                telemetry_context.sampled,
                telemetry_context.trace_id_source,
            )
        }),
    };

    if let Some(sink) = event_sink {
        let event_metadata = if options.telemetry.nerdstats {
            Some(json!({
                "nerdstats": workflow_nerdstats(&output),
            }))
        } else {
            None
        };
        sink.emit(&YamlWorkflowEvent {
            event_type: "workflow_completed".to_string(),
            node_id: None,
            step_id: None,
            node_kind: None,
            streamable: None,
            message: Some(format!("terminal_node={}", output.terminal_node)),
            delta: None,
            token_kind: None,
            is_terminal_node_token: None,
            elapsed_ms: Some(output.total_elapsed_ms),
            metadata: event_metadata,
        });
    }

    if event_sink_is_cancelled(event_sink) {
        return Err(YamlWorkflowRunError::EventSinkCancelled {
            message: workflow_event_sink_cancelled_message().to_string(),
        });
    }

    if let Some(mut span) = workflow_span.take() {
        span.set_attribute("workflow_id", workflow.id.as_str());
        if let Some(trace_id_value) = telemetry_context.trace_id.as_ref() {
            span.set_attribute("trace_id", trace_id_value.as_str());
        }
        span.end();
        flush_workflow_tracer();
    }

    Ok(output)
}

async fn try_run_yaml_via_ir_runtime(
    workflow: &YamlWorkflow,
    workflow_input: &Value,
    executor: &dyn YamlWorkflowLlmExecutor,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
    options: &YamlWorkflowRunOptions,
) -> Result<Option<YamlWorkflowRunOutput>, YamlWorkflowRunError> {
    let ir = match yaml_workflow_to_ir(workflow) {
        Ok(def) => def,
        Err(YamlToIrError::UnsupportedNode { .. })
        | Err(YamlToIrError::MultipleOutgoingEdge { .. }) => return Ok(None),
        Err(err) => {
            return Err(YamlWorkflowRunError::InvalidInput {
                message: err.to_string(),
            });
        }
    };

    let parent_trace_context = trace_context_from_options(options);
    let telemetry_context = resolve_telemetry_context(options, parent_trace_context.as_ref());

    let tracer = workflow_tracer();
    let mut workflow_span_context: Option<TraceContext> = None;
    let mut workflow_span = if telemetry_context.sampled {
        let (span_context, mut span) = tracer.start_span(
            "workflow.run",
            SpanKind::Workflow,
            parent_trace_context.as_ref(),
        );
        if let Some(trace_id_value) = telemetry_context.trace_id.as_deref() {
            span.set_attribute("trace_id", trace_id_value);
        }
        apply_trace_tenant_attributes(span.as_mut(), options);
        workflow_span_context = Some(span_context);
        Some(span)
    } else {
        None
    };

    struct NoopLlm;
    #[async_trait]
    impl LlmExecutor for NoopLlm {
        async fn execute(
            &self,
            _input: LlmExecutionInput,
        ) -> Result<LlmExecutionOutput, LlmExecutionError> {
            Err(LlmExecutionError::UnexpectedOutcome(
                "yaml_ir_uses_tool_path",
            ))
        }
    }

    struct YamlIrToolExecutor<'a> {
        llm_executor: &'a dyn YamlWorkflowLlmExecutor,
        custom_worker: Option<&'a dyn YamlWorkflowCustomWorkerExecutor>,
        token_totals: std::sync::Mutex<YamlTokenTotals>,
        node_usage: std::sync::Mutex<BTreeMap<String, YamlLlmTokenUsage>>,
        node_models: std::sync::Mutex<BTreeMap<String, String>>,
        model_override: Option<String>,
        trace_id: Option<String>,
        trace_context: Option<TraceContext>,
        trace_input_context: Option<YamlWorkflowTraceContextInput>,
        tenant_context: YamlWorkflowTraceTenantContext,
        payload_mode: YamlWorkflowPayloadMode,
        trace_sampled: bool,
    }

    #[async_trait]
    impl ToolExecutor for YamlIrToolExecutor<'_> {
        async fn execute_tool(
            &self,
            input: ToolExecutionInput,
        ) -> Result<Value, ToolExecutionError> {
            let context = build_yaml_context_from_ir_scope(&input.scoped_input);

            if input.tool == YAML_LLM_TOOL_ID {
                let node_id = input
                    .input
                    .get("node_id")
                    .and_then(Value::as_str)
                    .ok_or_else(|| {
                        ToolExecutionError::Failed("yaml llm call missing node_id".to_string())
                    })?
                    .to_string();
                let node_id_for_metrics = node_id.clone();
                let model = input
                    .input
                    .get("model")
                    .and_then(Value::as_str)
                    .ok_or_else(|| {
                        ToolExecutionError::Failed("yaml llm call missing model".to_string())
                    })?
                    .to_string();
                let resolved_model =
                    resolve_requested_model(self.model_override.as_deref(), &model);
                let prompt_template = input
                    .input
                    .get("prompt_template")
                    .and_then(Value::as_str)
                    .unwrap_or_default()
                    .to_string();
                let stream = input
                    .input
                    .get("stream")
                    .and_then(Value::as_bool)
                    .unwrap_or(false);
                let heal = input
                    .input
                    .get("heal")
                    .and_then(Value::as_bool)
                    .unwrap_or(false);
                let append_prompt_as_user = input
                    .input
                    .get("append_prompt_as_user")
                    .and_then(Value::as_bool)
                    .unwrap_or(true);
                let messages_path = input
                    .input
                    .get("messages_path")
                    .and_then(Value::as_str)
                    .map(str::to_string);

                let messages = if let Some(path) = messages_path.as_deref() {
                    Some(
                        parse_messages_from_context(path, &context)
                            .map_err(ToolExecutionError::Failed)?,
                    )
                } else {
                    None
                };

                let prompt_bindings = collect_template_bindings(&prompt_template, &context);
                let prompt = interpolate_template(&prompt_template, &context);
                let email_text = context
                    .get("input")
                    .and_then(|v| v.get("email_text"))
                    .and_then(Value::as_str)
                    .unwrap_or_default();
                let schema = input
                    .input
                    .get("output_schema")
                    .cloned()
                    .unwrap_or_else(default_llm_output_schema);

                let request = YamlLlmExecutionRequest {
                    node_id,
                    is_terminal_node: false,
                    stream_json_as_text: input
                        .input
                        .get("stream_json_as_text")
                        .and_then(Value::as_bool)
                        .unwrap_or(false),
                    model: resolved_model.clone(),
                    messages,
                    append_prompt_as_user,
                    prompt,
                    prompt_template,
                    prompt_bindings,
                    schema,
                    stream,
                    heal,
                    tools: Vec::new(),
                    tool_choice: None,
                    max_tool_roundtrips: 1,
                    tool_calls_global_key: None,
                    tool_trace_mode: YamlToolTraceMode::Off,
                    execution_context: context.clone(),
                    email_text: email_text.to_string(),
                    trace_id: self.trace_id.clone(),
                    trace_context: self.trace_context.clone(),
                    trace_sampled: self.trace_sampled,
                };

                let llm_result = self
                    .llm_executor
                    .complete_structured(request, None)
                    .await
                    .map_err(ToolExecutionError::Failed);

                if let Ok(ref result) = llm_result {
                    if let Some(usage) = result.usage.as_ref() {
                        if let Ok(mut totals) = self.token_totals.lock() {
                            totals.add_usage(usage);
                        }
                        if let Ok(mut usage_map) = self.node_usage.lock() {
                            usage_map.insert(node_id_for_metrics.clone(), usage.clone());
                        }
                    }
                    if let Ok(mut model_map) = self.node_models.lock() {
                        model_map.insert(node_id_for_metrics, resolved_model);
                    }
                }

                return llm_result.map(|result| result.payload);
            }

            let worker = self
                .custom_worker
                .ok_or_else(|| ToolExecutionError::NotFound {
                    tool: input.tool.clone(),
                })?;

            let payload = input.input.clone();
            let email_text = context
                .get("input")
                .and_then(|v| v.get("email_text"))
                .and_then(Value::as_str)
                .unwrap_or_default();

            let tracer = workflow_tracer();
            let mut handler_span_context: Option<TraceContext> = None;
            let mut handler_span = if self.trace_sampled {
                let (span_context, mut span) = tracer.start_span(
                    "handler.invoke",
                    SpanKind::Node,
                    self.trace_context.as_ref(),
                );
                handler_span_context = Some(span_context);
                if let Some(trace_id) = self.trace_id.as_ref() {
                    span.set_attribute("trace_id", trace_id.as_str());
                }
                span.set_attribute("handler_name", input.tool.as_str());
                if let Some(workspace_id) = self.tenant_context.workspace_id.as_deref() {
                    span.set_attribute("tenant.workspace_id", workspace_id);
                }
                if let Some(user_id) = self.tenant_context.user_id.as_deref() {
                    span.set_attribute("tenant.user_id", user_id);
                }
                if let Some(conversation_id) = self.tenant_context.conversation_id.as_deref() {
                    span.set_attribute("tenant.conversation_id", conversation_id);
                }
                if let Some(request_id) = self.tenant_context.request_id.as_deref() {
                    span.set_attribute("tenant.request_id", request_id);
                }
                if let Some(run_id) = self.tenant_context.run_id.as_deref() {
                    span.set_attribute("tenant.run_id", run_id);
                }
                span.set_attribute(
                    "node_input",
                    payload_for_span(self.payload_mode, &payload).as_str(),
                );
                Some(span)
            } else {
                None
            };

            let trace_options = YamlWorkflowRunOptions {
                telemetry: YamlWorkflowTelemetryConfig::default(),
                trace: YamlWorkflowTraceOptions {
                    context: self.trace_input_context.clone(),
                    tenant: self.tenant_context.clone(),
                },
                model: None,
            };
            let worker_trace_context = merged_trace_context_for_worker(
                handler_span_context.as_ref(),
                self.trace_id.as_deref(),
                &trace_options,
            );
            let worker_context = custom_worker_context_with_trace(
                &context,
                &worker_trace_context,
                &self.tenant_context,
            );

            let output_result = worker
                .execute(&input.tool, &payload, email_text, &worker_context)
                .await
                .map_err(ToolExecutionError::Failed);

            if let Some(span) = handler_span.as_mut() {
                if output_result.is_ok() {
                    span.add_event("handler.success");
                } else {
                    span.add_event("handler.error");
                }
            }

            if let Some(span) = handler_span.take() {
                span.end();
            }

            output_result
        }
    }

    let tool_executor = YamlIrToolExecutor {
        llm_executor: executor,
        custom_worker,
        token_totals: std::sync::Mutex::new(YamlTokenTotals::default()),
        node_usage: std::sync::Mutex::new(BTreeMap::new()),
        node_models: std::sync::Mutex::new(BTreeMap::new()),
        model_override: options.model.clone(),
        trace_id: telemetry_context.trace_id.clone(),
        trace_context: workflow_span_context.clone(),
        trace_input_context: options.trace.context.clone(),
        tenant_context: options.trace.tenant.clone(),
        payload_mode: options.telemetry.payload_mode,
        trace_sampled: telemetry_context.sampled,
    };

    let runtime = WorkflowRuntime::new(
        ir,
        &NoopLlm,
        Some(&tool_executor),
        WorkflowRuntimeOptions::default(),
    );

    let started = Instant::now();
    let result = match runtime.execute(workflow_input.clone(), None).await {
        Ok(result) => result,
        Err(WorkflowRuntimeError::Validation(_)) => return Ok(None),
        Err(error) => {
            return Err(YamlWorkflowRunError::IrRuntime {
                message: error.to_string(),
            });
        }
    };
    let total_elapsed_ms = started.elapsed().as_millis();

    let mut outputs: BTreeMap<String, Value> = BTreeMap::new();
    for (node_id, output) in result.node_outputs {
        if node_id == YAML_START_NODE_ID {
            continue;
        }
        outputs.insert(node_id, json!({"output": output}));
    }

    let mut trace = Vec::new();
    let mut step_timings = Vec::new();
    let node_usage_map = tool_executor
        .node_usage
        .lock()
        .map(|usage| usage.clone())
        .unwrap_or_default();
    let llm_node_models = tool_executor
        .node_models
        .lock()
        .map(|models| models.clone())
        .unwrap_or_default();
    let mut llm_node_metrics: BTreeMap<String, YamlLlmNodeMetrics> = BTreeMap::new();
    for execution in result.node_executions {
        if execution.node_id == YAML_START_NODE_ID {
            continue;
        }
        let node_id = execution.node_id;
        trace.push(node_id.clone());
        let usage = node_usage_map.get(&node_id);
        if let Some(usage) = usage {
            llm_node_metrics.insert(
                node_id.clone(),
                YamlLlmNodeMetrics {
                    elapsed_ms: 0,
                    prompt_tokens: usage.prompt_tokens,
                    completion_tokens: usage.completion_tokens,
                    total_tokens: usage.total_tokens,
                    reasoning_tokens: usage.reasoning_tokens,
                    tokens_per_second: completion_tokens_per_second(usage.completion_tokens, 0),
                },
            );
        }
        step_timings.push(YamlStepTiming {
            node_id: node_id.clone(),
            node_kind: "ir_runtime".to_string(),
            model_name: llm_node_models.get(&node_id).cloned(),
            elapsed_ms: 0,
            prompt_tokens: usage.map(|value| value.prompt_tokens),
            completion_tokens: usage.map(|value| value.completion_tokens),
            total_tokens: usage.map(|value| value.total_tokens),
            reasoning_tokens: usage.and_then(|value| value.reasoning_tokens),
            tokens_per_second: usage
                .map(|value| completion_tokens_per_second(value.completion_tokens, 0)),
        });
    }

    let terminal_node = result.terminal_node_id;
    let terminal_output = outputs
        .get(&terminal_node)
        .and_then(|v| v.get("output"))
        .cloned();

    let email_text = workflow_input
        .get("email_text")
        .and_then(Value::as_str)
        .unwrap_or_default()
        .to_string();

    let token_totals = tool_executor
        .token_totals
        .lock()
        .map(|totals| totals.clone())
        .unwrap_or_default();

    if let Some(mut span) = workflow_span.take() {
        span.set_attribute("workflow_id", workflow.id.as_str());
        if let Some(trace_id_value) = telemetry_context.trace_id.as_ref() {
            span.set_attribute("trace_id", trace_id_value.as_str());
        }
        span.end();
        flush_workflow_tracer();
    }

    Ok(Some(YamlWorkflowRunOutput {
        workflow_id: workflow.id.clone(),
        entry_node: workflow.entry_node.clone(),
        email_text,
        trace,
        outputs,
        terminal_node,
        terminal_output,
        step_timings,
        llm_node_metrics,
        llm_node_models,
        total_elapsed_ms,
        ttft_ms: None,
        total_input_tokens: token_totals.input_tokens,
        total_output_tokens: token_totals.output_tokens,
        total_tokens: token_totals.total_tokens,
        total_reasoning_tokens: token_totals.reasoning_tokens,
        tokens_per_second: token_totals.tokens_per_second(total_elapsed_ms),
        trace_id: telemetry_context.trace_id.clone(),
        metadata: telemetry_context.trace_id.as_ref().map(|value| {
            workflow_metadata_with_trace(
                options,
                value,
                telemetry_context.sampled,
                telemetry_context.trace_id_source,
            )
        }),
    }))
}

fn build_yaml_context_from_ir_scope(scoped_input: &Value) -> Value {
    let input = scoped_input.get("input").cloned().unwrap_or(Value::Null);

    let mut nodes = serde_json::Map::new();
    if let Some(node_outputs) = scoped_input.get("node_outputs").and_then(Value::as_object) {
        for (node_id, output) in node_outputs {
            nodes.insert(node_id.clone(), json!({"output": output.clone()}));
        }
    }

    json!({
        "input": input,
        "nodes": Value::Object(nodes),
        "globals": Value::Object(serde_json::Map::new())
    })
}

pub async fn run_email_workflow_yaml_with_custom_worker_and_events(
    workflow: &YamlWorkflow,
    email_text: &str,
    executor: &dyn YamlWorkflowLlmExecutor,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
    event_sink: Option<&dyn YamlWorkflowEventSink>,
) -> Result<YamlWorkflowRunOutput, YamlWorkflowRunError> {
    let workflow_input = json!({ "email_text": email_text });
    run_workflow_yaml_with_custom_worker_and_events(
        workflow,
        &workflow_input,
        executor,
        custom_worker,
        event_sink,
    )
    .await
}

fn evaluate_switch_condition(
    condition: &str,
    context: &Value,
) -> Result<bool, YamlWorkflowRunError> {
    let (left, right) =
        condition
            .split_once("==")
            .ok_or_else(|| YamlWorkflowRunError::UnsupportedCondition {
                condition: condition.to_string(),
            })?;

    let left_path = left.trim().trim_start_matches("$.");
    let right_literal = right.trim().trim_matches('"').trim_matches('\'');
    let left_value = resolve_path(context, left_path);
    Ok(left_value
        .and_then(Value::as_str)
        .map(|value| value == right_literal)
        .unwrap_or(false))
}

fn parse_messages_from_context(path: &str, context: &Value) -> Result<Vec<Message>, String> {
    let normalized_path = path.trim().trim_start_matches("$.");
    let value = resolve_path(context, normalized_path)
        .ok_or_else(|| format!("messages_path not found: {path}"))?;
    let list: Vec<WorkflowMessage> = serde_json::from_value(value.clone()).map_err(|err| {
        format!("messages_path must resolve to a list of messages: {path}; {err}")
    })?;
    if list.is_empty() {
        return Err(format!(
            "messages_path must not resolve to an empty list: {path}"
        ));
    }

    let mut messages = Vec::with_capacity(list.len());
    for (index, item) in list.into_iter().enumerate() {
        let mut message = match item.role {
            Role::System => Message::system(item.content),
            Role::User => Message::user(item.content),
            Role::Assistant => Message::assistant(item.content),
            Role::Tool => {
                let tool_call_id = item
                    .tool_call_id
                    .ok_or_else(|| format!("tool message at index {index} missing tool_call_id"))?;
                Message::tool(item.content, tool_call_id)
            }
        };

        if let Some(name) = item.name {
            message = message.with_name(name);
        }

        messages.push(message);
    }

    Ok(messages)
}

pub fn verify_yaml_workflow(workflow: &YamlWorkflow) -> Vec<YamlWorkflowDiagnostic> {
    let mut diagnostics = Vec::new();
    let known_ids: HashMap<&str, &YamlNode> = workflow
        .nodes
        .iter()
        .map(|node| (node.id.as_str(), node))
        .collect();

    if !known_ids.contains_key(workflow.entry_node.as_str()) {
        diagnostics.push(YamlWorkflowDiagnostic {
            node_id: None,
            code: "missing_entry".to_string(),
            severity: YamlWorkflowDiagnosticSeverity::Error,
            message: format!("entry node '{}' does not exist", workflow.entry_node),
        });
    }

    for edge in &workflow.edges {
        if !known_ids.contains_key(edge.from.as_str()) {
            diagnostics.push(YamlWorkflowDiagnostic {
                node_id: Some(edge.from.clone()),
                code: "unknown_edge_from".to_string(),
                severity: YamlWorkflowDiagnosticSeverity::Error,
                message: format!("edge.from '{}' does not exist", edge.from),
            });
        }
        if !known_ids.contains_key(edge.to.as_str()) {
            diagnostics.push(YamlWorkflowDiagnostic {
                node_id: Some(edge.to.clone()),
                code: "unknown_edge_to".to_string(),
                severity: YamlWorkflowDiagnosticSeverity::Error,
                message: format!("edge.to '{}' does not exist", edge.to),
            });
        }
    }

    for node in &workflow.nodes {
        if let Some(llm) = &node.node_type.llm_call {
            if llm.model.trim().is_empty() {
                diagnostics.push(YamlWorkflowDiagnostic {
                    node_id: Some(node.id.clone()),
                    code: "empty_model".to_string(),
                    severity: YamlWorkflowDiagnosticSeverity::Error,
                    message: "llm_call.model must not be empty".to_string(),
                });
            }
            if llm.stream.unwrap_or(false) && llm.heal.unwrap_or(false) {
                diagnostics.push(YamlWorkflowDiagnostic {
                    node_id: Some(node.id.clone()),
                    code: "stream_heal_conflict".to_string(),
                    severity: YamlWorkflowDiagnosticSeverity::Warning,
                    message:
                        "llm_call.stream=true with heal=true is not streamable; runtime will disable streaming"
                            .to_string(),
                });
            }

            if llm.max_tool_roundtrips.unwrap_or(1) == 0 {
                diagnostics.push(YamlWorkflowDiagnostic {
                    node_id: Some(node.id.clone()),
                    code: "invalid_max_tool_roundtrips".to_string(),
                    severity: YamlWorkflowDiagnosticSeverity::Error,
                    message: "llm_call.max_tool_roundtrips must be >= 1".to_string(),
                });
            }

            if let Some(global_key) = llm.tool_calls_global_key.as_ref() {
                if global_key.trim().is_empty() {
                    diagnostics.push(YamlWorkflowDiagnostic {
                        node_id: Some(node.id.clone()),
                        code: "empty_tool_calls_global_key".to_string(),
                        severity: YamlWorkflowDiagnosticSeverity::Error,
                        message: "llm_call.tool_calls_global_key must not be empty".to_string(),
                    });
                }
            }

            match normalize_tool_choice(llm.tool_choice.clone()) {
                Ok(choice) => {
                    if let Some(ToolChoice::Tool(choice_tool)) = choice.as_ref() {
                        if !llm.tools.iter().any(|tool| match (llm.tools_format, tool) {
                            (YamlToolFormat::Openai, YamlToolDeclaration::OpenAi(openai)) => {
                                openai.function.name == choice_tool.function.name
                            }
                            (
                                YamlToolFormat::Simplified,
                                YamlToolDeclaration::Simplified(simple),
                            ) => simple.name == choice_tool.function.name,
                            _ => false,
                        }) {
                            diagnostics.push(YamlWorkflowDiagnostic {
                                node_id: Some(node.id.clone()),
                                code: "unknown_tool_choice_function".to_string(),
                                severity: YamlWorkflowDiagnosticSeverity::Error,
                                message: format!(
                                    "llm_call.tool_choice references unknown function '{}'",
                                    choice_tool.function.name
                                ),
                            });
                        }
                    }
                }
                Err(message) => {
                    diagnostics.push(YamlWorkflowDiagnostic {
                        node_id: Some(node.id.clone()),
                        code: "invalid_tool_choice".to_string(),
                        severity: YamlWorkflowDiagnosticSeverity::Error,
                        message,
                    });
                }
            }

            let normalized_tools = match normalize_llm_tools(llm) {
                Ok(tools) => tools,
                Err(message) => {
                    diagnostics.push(YamlWorkflowDiagnostic {
                        node_id: Some(node.id.clone()),
                        code: "invalid_tools_format".to_string(),
                        severity: YamlWorkflowDiagnosticSeverity::Error,
                        message,
                    });
                    Vec::new()
                }
            };

            let mut seen_tool_names = HashSet::new();
            for tool in &normalized_tools {
                let name = tool.definition.function.name.trim();
                if name.is_empty() {
                    diagnostics.push(YamlWorkflowDiagnostic {
                        node_id: Some(node.id.clone()),
                        code: "empty_tool_name".to_string(),
                        severity: YamlWorkflowDiagnosticSeverity::Error,
                        message: "tool function name must not be empty".to_string(),
                    });
                }
                if !seen_tool_names.insert(tool.definition.function.name.clone()) {
                    diagnostics.push(YamlWorkflowDiagnostic {
                        node_id: Some(node.id.clone()),
                        code: "duplicate_tool_name".to_string(),
                        severity: YamlWorkflowDiagnosticSeverity::Error,
                        message: format!(
                            "duplicate tool function name '{}' in node",
                            tool.definition.function.name
                        ),
                    });
                }

                let schema = tool
                    .definition
                    .function
                    .parameters
                    .clone()
                    .unwrap_or(Value::Null);
                if schema.is_null() {
                    diagnostics.push(YamlWorkflowDiagnostic {
                        node_id: Some(node.id.clone()),
                        code: "missing_tool_input_schema".to_string(),
                        severity: YamlWorkflowDiagnosticSeverity::Error,
                        message: format!(
                            "tool '{}' is missing input schema",
                            tool.definition.function.name
                        ),
                    });
                } else if let Err(message) = validate_json_schema(&schema) {
                    diagnostics.push(YamlWorkflowDiagnostic {
                        node_id: Some(node.id.clone()),
                        code: "invalid_tool_input_schema".to_string(),
                        severity: YamlWorkflowDiagnosticSeverity::Error,
                        message: format!(
                            "tool '{}' has invalid input schema: {}",
                            tool.definition.function.name, message
                        ),
                    });
                }

                if let Some(output_schema) = tool.output_schema.as_ref() {
                    if let Err(message) = validate_json_schema(output_schema) {
                        diagnostics.push(YamlWorkflowDiagnostic {
                            node_id: Some(node.id.clone()),
                            code: "invalid_tool_output_schema".to_string(),
                            severity: YamlWorkflowDiagnosticSeverity::Error,
                            message: format!(
                                "tool '{}' has invalid output schema: {}",
                                tool.definition.function.name, message
                            ),
                        });
                    }
                }
            }
        }

        if let Some(switch) = &node.node_type.switch {
            for branch in &switch.branches {
                if !known_ids.contains_key(branch.target.as_str()) {
                    diagnostics.push(YamlWorkflowDiagnostic {
                        node_id: Some(node.id.clone()),
                        code: "unknown_switch_target".to_string(),
                        severity: YamlWorkflowDiagnosticSeverity::Error,
                        message: format!("switch branch target '{}' does not exist", branch.target),
                    });
                }
            }
            if !known_ids.contains_key(switch.default.as_str()) {
                diagnostics.push(YamlWorkflowDiagnostic {
                    node_id: Some(node.id.clone()),
                    code: "unknown_switch_default".to_string(),
                    severity: YamlWorkflowDiagnosticSeverity::Error,
                    message: format!("switch default target '{}' does not exist", switch.default),
                });
            }
        }

        if let Some(config) = node.config.as_ref() {
            if let Some(update_globals) = config.update_globals.as_ref() {
                for (key, update) in update_globals {
                    let is_valid_op =
                        matches!(update.op.as_str(), "set" | "append" | "increment" | "merge");
                    if !is_valid_op {
                        diagnostics.push(YamlWorkflowDiagnostic {
                            node_id: Some(node.id.clone()),
                            code: "unknown_update_op".to_string(),
                            severity: YamlWorkflowDiagnosticSeverity::Error,
                            message: format!(
                                "update_globals key '{}' has unknown op '{}'; expected set|append|increment|merge",
                                key, update.op
                            ),
                        });
                    }

                    if update.op != "increment" && update.from.is_none() {
                        diagnostics.push(YamlWorkflowDiagnostic {
                            node_id: Some(node.id.clone()),
                            code: "missing_update_from".to_string(),
                            severity: YamlWorkflowDiagnosticSeverity::Error,
                            message: format!(
                                "update_globals key '{}' with op '{}' requires 'from'",
                                key, update.op
                            ),
                        });
                    }
                }
            }
        }
    }

    diagnostics
}

fn resolve_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
    path.split('.')
        .filter(|segment| !segment.is_empty())
        .try_fold(value, |current, segment| {
            if let Ok(index) = segment.parse::<usize>() {
                return current.get(index);
            }
            current.get(segment)
        })
}

fn interpolate_template(template: &str, context: &Value) -> String {
    let mut out = String::with_capacity(template.len());
    let mut rest = template;

    loop {
        let Some(start) = rest.find("{{") else {
            out.push_str(rest);
            break;
        };

        out.push_str(&rest[..start]);
        let after_start = &rest[start + 2..];
        let Some(end) = after_start.find("}}") else {
            out.push_str(&rest[start..]);
            break;
        };

        let expr = after_start[..end].trim();
        let source_path = expr.trim_start_matches("$.");
        let replacement = resolve_path(context, source_path)
            .map(value_to_template_string)
            .unwrap_or_default();
        out.push_str(replacement.as_str());

        rest = &after_start[end + 2..];
    }

    out
}

fn collect_template_bindings(template: &str, context: &Value) -> Vec<YamlTemplateBinding> {
    let mut bindings = Vec::new();
    let mut rest = template;

    loop {
        let Some(start) = rest.find("{{") else {
            break;
        };

        let after_start = &rest[start + 2..];
        let Some(end) = after_start.find("}}") else {
            break;
        };

        let expr = after_start[..end].trim();
        let source_path = expr.trim_start_matches("$.").to_string();
        let resolved = resolve_path(context, source_path.as_str()).cloned();
        let missing = resolved.is_none();
        let resolved_value = resolved.unwrap_or(Value::Null);
        bindings.push(YamlTemplateBinding {
            index: bindings.len(),
            expression: expr.to_string(),
            source_path,
            resolved_type: json_type_name(&resolved_value).to_string(),
            missing,
            resolved: resolved_value,
        });

        rest = &after_start[end + 2..];
    }

    bindings
}

fn json_type_name(value: &Value) -> &'static str {
    match value {
        Value::Null => "null",
        Value::Bool(_) => "bool",
        Value::Number(_) => "number",
        Value::String(_) => "string",
        Value::Array(_) => "array",
        Value::Object(_) => "object",
    }
}

fn value_to_template_string(value: &Value) -> String {
    match value {
        Value::Null => String::new(),
        Value::Bool(v) => v.to_string(),
        Value::Number(v) => v.to_string(),
        Value::String(v) => v.clone(),
        Value::Array(_) | Value::Object(_) => serde_json::to_string(value).unwrap_or_default(),
    }
}

fn apply_set_globals(
    node: &YamlNode,
    outputs: &BTreeMap<String, Value>,
    workflow_input: &Value,
    globals: &mut serde_json::Map<String, Value>,
) {
    let Some(config) = node.config.as_ref() else {
        return;
    };
    let Some(set_globals) = config.set_globals.as_ref() else {
        return;
    };

    let context = json!({
        "input": workflow_input,
        "nodes": outputs,
        "globals": Value::Object(globals.clone())
    });

    for (key, expr) in set_globals {
        let value = resolve_path(&context, expr.as_str())
            .cloned()
            .unwrap_or(Value::Null);
        globals.insert(key.clone(), value);
    }
}

fn apply_update_globals(
    node: &YamlNode,
    outputs: &BTreeMap<String, Value>,
    workflow_input: &Value,
    globals: &mut serde_json::Map<String, Value>,
) {
    let Some(config) = node.config.as_ref() else {
        return;
    };
    let Some(update_globals) = config.update_globals.as_ref() else {
        return;
    };

    let context = json!({
        "input": workflow_input,
        "nodes": outputs,
        "globals": Value::Object(globals.clone())
    });

    for (key, update) in update_globals {
        match update.op.as_str() {
            "set" => {
                if let Some(path) = update.from.as_ref() {
                    let value = resolve_path(&context, path.as_str())
                        .cloned()
                        .unwrap_or(Value::Null);
                    globals.insert(key.clone(), value);
                }
            }
            "append" => {
                if let Some(path) = update.from.as_ref() {
                    let value = resolve_path(&context, path.as_str())
                        .cloned()
                        .unwrap_or(Value::Null);
                    let entry = globals
                        .entry(key.clone())
                        .or_insert_with(|| Value::Array(Vec::new()));
                    match entry {
                        Value::Array(items) => items.push(value),
                        other => {
                            let existing = other.clone();
                            *other = Value::Array(vec![existing, value]);
                        }
                    }
                }
            }
            "increment" => {
                let by = update.by.unwrap_or(1.0);
                let current = globals
                    .get(key.as_str())
                    .and_then(Value::as_f64)
                    .unwrap_or(0.0);
                if let Some(next) = serde_json::Number::from_f64(current + by) {
                    globals.insert(key.clone(), Value::Number(next));
                }
            }
            "merge" => {
                if let Some(path) = update.from.as_ref() {
                    let source = resolve_path(&context, path.as_str())
                        .cloned()
                        .unwrap_or(Value::Null);
                    if let Value::Object(source_map) = source {
                        let target = globals
                            .entry(key.clone())
                            .or_insert_with(|| Value::Object(serde_json::Map::new()));
                        if let Value::Object(target_map) = target {
                            target_map.extend(source_map);
                        } else {
                            *target = Value::Object(source_map);
                        }
                    }
                }
            }
            _ => {}
        }
    }
}

fn llm_output_schema_for_node(node: &YamlNode) -> Value {
    if let Some(schema) = node
        .config
        .as_ref()
        .and_then(|cfg| cfg.output_schema.clone())
    {
        return schema;
    }

    default_llm_output_schema()
}

fn normalize_tool_choice(
    config: Option<YamlToolChoiceConfig>,
) -> Result<Option<ToolChoice>, String> {
    let Some(config) = config else {
        return Ok(None);
    };

    let choice = match config {
        YamlToolChoiceConfig::Mode(mode) => ToolChoice::Mode(mode),
        YamlToolChoiceConfig::Function { function } => ToolChoice::Tool(ToolChoiceTool {
            tool_type: ToolType::Function,
            function: ToolChoiceFunction { name: function },
        }),
        YamlToolChoiceConfig::OpenAi(tool) => ToolChoice::Tool(tool),
    };

    Ok(Some(choice))
}

fn normalize_llm_tools(llm: &YamlLlmCall) -> Result<Vec<YamlResolvedTool>, String> {
    llm.tools
        .iter()
        .map(|tool| match (llm.tools_format, tool) {
            (YamlToolFormat::Openai, YamlToolDeclaration::OpenAi(openai)) => {
                let definition = ToolDefinition {
                    tool_type: openai.tool_type.unwrap_or(ToolType::Function),
                    function: ToolFunction {
                        name: openai.function.name.clone(),
                        description: openai.function.description.clone(),
                        parameters: openai.function.parameters.clone(),
                    },
                };
                Ok(YamlResolvedTool {
                    definition,
                    output_schema: openai.function.output_schema.clone(),
                })
            }
            (YamlToolFormat::Simplified, YamlToolDeclaration::Simplified(simple)) => {
                let definition = ToolDefinition {
                    tool_type: ToolType::Function,
                    function: ToolFunction {
                        name: simple.name.clone(),
                        description: simple.description.clone(),
                        parameters: Some(simple.input_schema.clone()),
                    },
                };
                Ok(YamlResolvedTool {
                    definition,
                    output_schema: simple.output_schema.clone(),
                })
            }
            (YamlToolFormat::Openai, _) => {
                Err("tools_format=openai requires OpenAI-style tool declarations".to_string())
            }
            (YamlToolFormat::Simplified, _) => {
                Err("tools_format=simplified requires simplified tool declarations".to_string())
            }
        })
        .collect()
}

fn default_llm_output_schema() -> Value {
    json!({
        "type": "object",
        "additionalProperties": true
    })
}

fn mock_rag(topic: &str) -> Value {
    let (kb_source, playbook) = match topic {
        "probation" => (
            "hr_policy/probation.md",
            "Collect manager review, performance evidence, and probation timeline.",
        ),
        "leave_request" => (
            "hr_policy/leave.md",
            "Validate leave balance, manager approval, and blackout dates.",
        ),
        "supply_chain_order_assessment" => (
            "supply_chain/order_assessment.md",
            "Review order specs, inventory risk, and vendor lead-time guidance.",
        ),
        "supply_chain_order_replacement" => (
            "supply_chain/order_replacement.md",
            "Collect order id, damage proof, and replacement SLA policy.",
        ),
        "termination_first_time_offense" => (
            "hr_policy/termination_first_offense.md",
            "Validate first-incident criteria and route to HRBP review.",
        ),
        "termination_repeated_offense" => (
            "hr_policy/termination_repeated_offense.md",
            "Collect prior warnings and escalation approvals before final action.",
        ),
        _ => (
            "shared/request_clarification.md",
            "Request clarifying details before routing.",
        ),
    };

    json!({
        "kb_source": kb_source,
        "playbook": playbook,
    })
}

async fn execute_subworkflow_tool_call(
    payload: &Value,
    context: &Value,
    client: &SimpleAgentsClient,
    custom_worker: Option<&dyn YamlWorkflowCustomWorkerExecutor>,
    parent_options: &YamlWorkflowRunOptions,
    parent_trace_context: Option<&TraceContext>,
    resolved_trace_id: Option<&str>,
) -> Result<Value, String> {
    let workflow_id = payload
        .get("workflow_id")
        .and_then(Value::as_str)
        .ok_or_else(|| "run_workflow_graph requires payload.workflow_id".to_string())?;

    let input_context = context
        .get("input")
        .and_then(Value::as_object)
        .ok_or_else(|| "run_workflow_graph requires context.input".to_string())?;

    let registry = input_context
        .get("workflow_registry")
        .and_then(Value::as_object)
        .ok_or_else(|| {
            "run_workflow_graph requires input.workflow_registry map of workflow_id -> yaml_path"
                .to_string()
        })?;

    let workflow_path = registry
        .get(workflow_id)
        .and_then(Value::as_str)
        .ok_or_else(|| {
            format!(
                "workflow_registry has no entry for workflow_id '{}'",
                workflow_id
            )
        })?;

    let parent_depth = input_context
        .get("__subgraph_depth")
        .and_then(Value::as_u64)
        .unwrap_or(0);
    let max_depth = input_context
        .get("__subgraph_max_depth")
        .and_then(Value::as_u64)
        .unwrap_or(3);

    if parent_depth >= max_depth {
        return Err(format!(
            "run_workflow_graph depth limit reached (depth={}, max={})",
            parent_depth, max_depth
        ));
    }

    let mut subworkflow_input = payload
        .get("input")
        .and_then(Value::as_object)
        .cloned()
        .unwrap_or_default();

    if !subworkflow_input.contains_key("messages") {
        if let Some(messages) = input_context.get("messages") {
            subworkflow_input.insert("messages".to_string(), messages.clone());
        }
    }

    if !subworkflow_input.contains_key("email_text") {
        if let Some(email_text) = input_context.get("email_text") {
            subworkflow_input.insert("email_text".to_string(), email_text.clone());
        }
    }

    if !subworkflow_input.contains_key("workflow_registry") {
        subworkflow_input.insert(
            "workflow_registry".to_string(),
            Value::Object(registry.clone()),
        );
    }

    subworkflow_input.insert(
        "__subgraph_depth".to_string(),
        Value::Number(serde_json::Number::from(parent_depth + 1)),
    );
    subworkflow_input.insert(
        "__subgraph_max_depth".to_string(),
        Value::Number(serde_json::Number::from(max_depth)),
    );

    let subworkflow_options =
        build_subworkflow_options(parent_options, parent_trace_context, resolved_trace_id);

    let output = run_workflow_yaml_file_with_client_and_custom_worker_and_events_and_options(
        Path::new(workflow_path),
        &Value::Object(subworkflow_input),
        client,
        custom_worker,
        None,
        &subworkflow_options,
    )
    .await
    .map_err(|error| format!("subworkflow '{}' failed: {}", workflow_id, error))?;

    Ok(json!({
        "workflow_id": workflow_id,
        "workflow_path": workflow_path,
        "terminal_node": output.terminal_node,
        "terminal_output": output.terminal_output,
        "trace": output.trace,
    }))
}

fn build_subworkflow_options(
    parent_options: &YamlWorkflowRunOptions,
    parent_trace_context: Option<&TraceContext>,
    resolved_trace_id: Option<&str>,
) -> YamlWorkflowRunOptions {
    let mut subworkflow_options = parent_options.clone();
    if parent_trace_context.is_some() || resolved_trace_id.is_some() {
        let trace_context = YamlWorkflowTraceContextInput {
            trace_id: resolved_trace_id
                .map(|value| value.to_string())
                .or_else(|| parent_trace_context.and_then(|ctx| ctx.trace_id.clone())),
            span_id: parent_trace_context.and_then(|ctx| ctx.span_id.clone()),
            parent_span_id: parent_trace_context.and_then(|ctx| ctx.parent_span_id.clone()),
            traceparent: parent_trace_context.and_then(|ctx| ctx.traceparent.clone()),
            tracestate: parent_trace_context.and_then(|ctx| ctx.tracestate.clone()),
            baggage: parent_trace_context
                .map(|ctx| ctx.baggage.clone())
                .unwrap_or_default(),
        };
        subworkflow_options.trace.context = Some(trace_context);
    }
    subworkflow_options
}

fn mock_custom_worker_output(
    handler: &str,
    payload: &Value,
) -> Result<Value, YamlWorkflowRunError> {
    if handler == "get_employee_record" {
        let employee_name = payload
            .get("employee_name")
            .and_then(Value::as_str)
            .unwrap_or("Unknown Employee")
            .trim();
        let normalized_name = if employee_name.is_empty() {
            "Unknown Employee"
        } else {
            employee_name
        };

        let (employee_id, location) = match normalized_name.to_ascii_lowercase().as_str() {
            "alex johnson" => ("EMP-2041", "Austin"),
            "priya sharma" => ("EMP-3378", "Bengaluru"),
            "marcus lee" => ("EMP-1196", "Singapore"),
            "sarah chen" => ("EMP-4450", "Toronto"),
            _ => ("EMP-0000", "Unassigned"),
        };

        return Ok(json!({
            "employee_name": normalized_name,
            "employee_id": employee_id,
            "location": location,
        }));
    }

    if let Some(topic) = payload.get("topic").and_then(Value::as_str) {
        let mut value = mock_rag(topic);
        if let Value::Object(object) = &mut value {
            object.insert("handler".to_string(), Value::String(handler.to_string()));
        }
        return Ok(value);
    }

    Err(YamlWorkflowRunError::UnsupportedCustomHandler {
        handler: handler.to_string(),
    })
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlWorkflow {
    pub id: String,
    pub entry_node: String,
    #[serde(default)]
    pub nodes: Vec<YamlNode>,
    #[serde(default)]
    pub edges: Vec<YamlEdge>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlNode {
    pub id: String,
    pub node_type: YamlNodeType,
    pub config: Option<YamlNodeConfig>,
}

impl YamlNode {
    fn kind_name(&self) -> &'static str {
        if self.node_type.llm_call.is_some() {
            "llm_call"
        } else if self.node_type.switch.is_some() {
            "switch"
        } else if self.node_type.custom_worker.is_some() {
            "custom_worker"
        } else {
            "unknown"
        }
    }
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlNodeType {
    pub llm_call: Option<YamlLlmCall>,
    pub switch: Option<YamlSwitch>,
    pub custom_worker: Option<YamlCustomWorker>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlLlmCall {
    pub model: String,
    pub stream: Option<bool>,
    pub stream_json_as_text: Option<bool>,
    pub heal: Option<bool>,
    pub messages_path: Option<String>,
    pub append_prompt_as_user: Option<bool>,
    #[serde(default)]
    pub tools_format: YamlToolFormat,
    #[serde(default)]
    pub tools: Vec<YamlToolDeclaration>,
    pub tool_choice: Option<YamlToolChoiceConfig>,
    pub max_tool_roundtrips: Option<u8>,
    pub tool_calls_global_key: Option<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum YamlToolFormat {
    #[default]
    Openai,
    Simplified,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub enum YamlToolDeclaration {
    OpenAi(YamlOpenAiToolDeclaration),
    Simplified(YamlSimplifiedToolDeclaration),
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlOpenAiToolDeclaration {
    #[serde(rename = "type")]
    pub tool_type: Option<ToolType>,
    pub function: YamlOpenAiToolFunction,
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlOpenAiToolFunction {
    pub name: String,
    pub description: Option<String>,
    pub parameters: Option<Value>,
    pub output_schema: Option<Value>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlSimplifiedToolDeclaration {
    pub name: String,
    pub description: Option<String>,
    pub input_schema: Value,
    pub output_schema: Option<Value>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub enum YamlToolChoiceConfig {
    Mode(ToolChoiceMode),
    Function { function: String },
    OpenAi(ToolChoiceTool),
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlSwitch {
    #[serde(default)]
    pub branches: Vec<YamlSwitchBranch>,
    pub default: String,
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlSwitchBranch {
    pub condition: String,
    pub target: String,
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlCustomWorker {
    pub handler: String,
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlNodeConfig {
    pub prompt: Option<String>,
    #[serde(default, alias = "schema")]
    pub output_schema: Option<Value>,
    pub payload: Option<Value>,
    pub set_globals: Option<HashMap<String, String>>,
    pub update_globals: Option<HashMap<String, YamlGlobalUpdate>>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlGlobalUpdate {
    pub op: String,
    pub from: Option<String>,
    pub by: Option<f64>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct YamlEdge {
    pub from: String,
    pub to: String,
}

#[cfg(test)]
mod tests {
    use super::*;
    use simple_agent_type::provider::{Provider, ProviderRequest, ProviderResponse};
    use simple_agent_type::response::{CompletionChoice, CompletionResponse, Usage};
    use simple_agent_type::tool::{ToolCallFunction, ToolType};
    use simple_agent_type::{Result as SaResult, SimpleAgentsError};
    use simple_agents_core::SimpleAgentsClientBuilder;
    use std::fs;
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
    use std::sync::{Arc, Mutex, OnceLock};

    fn stream_debug_env_lock() -> &'static Mutex<()> {
        static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
        LOCK.get_or_init(|| Mutex::new(()))
    }

    struct MockExecutor;

    struct RecordingSink {
        events: Mutex<Vec<YamlWorkflowEvent>>,
    }

    struct CancelAfterFirstEventSink {
        cancelled: AtomicBool,
    }

    impl YamlWorkflowEventSink for CancelAfterFirstEventSink {
        fn emit(&self, _event: &YamlWorkflowEvent) {
            self.cancelled.store(true, Ordering::SeqCst);
        }

        fn is_cancelled(&self) -> bool {
            self.cancelled.load(Ordering::SeqCst)
        }
    }

    struct CountingExecutor {
        call_count: AtomicUsize,
    }

    struct CapturingWorker {
        context: Mutex<Option<Value>>,
    }

    struct ToolLoopProvider;

    struct UnknownToolProvider;

    struct ReasoningUsageProvider;

    struct ToolLoopReasoningProvider;

    #[async_trait]
    impl Provider for ToolLoopProvider {
        fn name(&self) -> &str {
            "openai"
        }

        fn transform_request(&self, req: &CompletionRequest) -> SaResult<ProviderRequest> {
            let body = serde_json::to_value(req).map_err(SimpleAgentsError::from)?;
            Ok(ProviderRequest::new("mock://tool-loop").with_body(body))
        }

        async fn execute(&self, req: ProviderRequest) -> SaResult<ProviderResponse> {
            let request: CompletionRequest =
                serde_json::from_value(req.body).map_err(SimpleAgentsError::from)?;

            let has_tools = request
                .tools
                .as_ref()
                .is_some_and(|tools| !tools.is_empty());
            let has_tool_result = request.messages.iter().any(|m| m.role == Role::Tool);

            let response = if has_tools && !has_tool_result {
                CompletionResponse {
                    id: "resp_tool_1".to_string(),
                    model: request.model.clone(),
                    choices: vec![CompletionChoice {
                        index: 0,
                        message: Message::assistant("").with_tool_calls(vec![ToolCall {
                            id: "call_get_context".to_string(),
                            tool_type: ToolType::Function,
                            function: ToolCallFunction {
                                name: "get_customer_context".to_string(),
                                arguments: "{\"order_id\":\"123\"}".to_string(),
                            },
                        }]),
                        finish_reason: FinishReason::ToolCalls,
                        logprobs: None,
                    }],
                    usage: Usage::new(10, 5),
                    created: None,
                    provider: Some(self.name().to_string()),
                    healing_metadata: None,
                }
            } else if has_tools && has_tool_result {
                CompletionResponse {
                    id: "resp_tool_2".to_string(),
                    model: request.model.clone(),
                    choices: vec![CompletionChoice {
                        index: 0,
                        message: Message::assistant("{\"state\":\"done\"}"),
                        finish_reason: FinishReason::Stop,
                        logprobs: None,
                    }],
                    usage: Usage::new(12, 6),
                    created: None,
                    provider: Some(self.name().to_string()),
                    healing_metadata: None,
                }
            } else {
                let prompt = request
                    .messages
                    .iter()
                    .rev()
                    .find(|m| m.role == Role::User)
                    .map(|m| m.content.clone())
                    .unwrap_or_default();
                let payload = json!({
                    "subject": "ok",
                    "body": prompt,
                })
                .to_string();
                CompletionResponse {
                    id: "resp_final".to_string(),
                    model: request.model.clone(),
                    choices: vec![CompletionChoice {
                        index: 0,
                        message: Message::assistant(payload),
                        finish_reason: FinishReason::Stop,
                        logprobs: None,
                    }],
                    usage: Usage::new(8, 4),
                    created: None,
                    provider: Some(self.name().to_string()),
                    healing_metadata: None,
                }
            };

            let body = serde_json::to_value(response).map_err(SimpleAgentsError::from)?;
            Ok(ProviderResponse::new(200, body))
        }

        fn transform_response(&self, resp: ProviderResponse) -> SaResult<CompletionResponse> {
            serde_json::from_value(resp.body).map_err(SimpleAgentsError::from)
        }
    }

    #[async_trait]
    impl Provider for UnknownToolProvider {
        fn name(&self) -> &str {
            "openai"
        }

        fn transform_request(&self, req: &CompletionRequest) -> SaResult<ProviderRequest> {
            let body = serde_json::to_value(req).map_err(SimpleAgentsError::from)?;
            Ok(ProviderRequest::new("mock://unknown-tool").with_body(body))
        }

        async fn execute(&self, req: ProviderRequest) -> SaResult<ProviderResponse> {
            let request: CompletionRequest =
                serde_json::from_value(req.body).map_err(SimpleAgentsError::from)?;

            let response = CompletionResponse {
                id: "resp_unknown_tool".to_string(),
                model: request.model,
                choices: vec![CompletionChoice {
                    index: 0,
                    message: Message::assistant("").with_tool_calls(vec![ToolCall {
                        id: "call_unknown".to_string(),
                        tool_type: ToolType::Function,
                        function: ToolCallFunction {
                            name: "unknown_tool".to_string(),
                            arguments: "{}".to_string(),
                        },
                    }]),
                    finish_reason: FinishReason::ToolCalls,
                    logprobs: None,
                }],
                usage: Usage::new(5, 2),
                created: None,
                provider: Some(self.name().to_string()),
                healing_metadata: None,
            };

            let body = serde_json::to_value(response).map_err(SimpleAgentsError::from)?;
            Ok(ProviderResponse::new(200, body))
        }

        fn transform_response(&self, resp: ProviderResponse) -> SaResult<CompletionResponse> {
            serde_json::from_value(resp.body).map_err(SimpleAgentsError::from)
        }
    }

    #[async_trait]
    impl Provider for ReasoningUsageProvider {
        fn name(&self) -> &str {
            "openai"
        }

        fn transform_request(&self, req: &CompletionRequest) -> SaResult<ProviderRequest> {
            let body = serde_json::to_value(req).map_err(SimpleAgentsError::from)?;
            Ok(ProviderRequest::new("mock://reasoning-usage").with_body(body))
        }

        async fn execute(&self, req: ProviderRequest) -> SaResult<ProviderResponse> {
            let request: CompletionRequest =
                serde_json::from_value(req.body).map_err(SimpleAgentsError::from)?;

            let mut usage = Usage::new(9, 5);
            usage.reasoning_tokens = Some(4);
            let response = CompletionResponse {
                id: "resp_reasoning".to_string(),
                model: request.model,
                choices: vec![CompletionChoice {
                    index: 0,
                    message: Message::assistant("{\"state\":\"ok\"}"),
                    finish_reason: FinishReason::Stop,
                    logprobs: None,
                }],
                usage,
                created: None,
                provider: Some(self.name().to_string()),
                healing_metadata: None,
            };

            let body = serde_json::to_value(response).map_err(SimpleAgentsError::from)?;
            Ok(ProviderResponse::new(200, body))
        }

        fn transform_response(&self, resp: ProviderResponse) -> SaResult<CompletionResponse> {
            serde_json::from_value(resp.body).map_err(SimpleAgentsError::from)
        }
    }

    #[async_trait]
    impl Provider for ToolLoopReasoningProvider {
        fn name(&self) -> &str {
            "openai"
        }

        fn transform_request(&self, req: &CompletionRequest) -> SaResult<ProviderRequest> {
            let body = serde_json::to_value(req).map_err(SimpleAgentsError::from)?;
            Ok(ProviderRequest::new("mock://tool-loop-reasoning").with_body(body))
        }

        async fn execute(&self, req: ProviderRequest) -> SaResult<ProviderResponse> {
            let request: CompletionRequest =
                serde_json::from_value(req.body).map_err(SimpleAgentsError::from)?;

            let has_tools = request
                .tools
                .as_ref()
                .is_some_and(|tools| !tools.is_empty());
            let has_tool_result = request.messages.iter().any(|m| m.role == Role::Tool);

            let response = if has_tools && !has_tool_result {
                let mut usage = Usage::new(10, 5);
                usage.reasoning_tokens = Some(2);
                CompletionResponse {
                    id: "resp_tool_reasoning_1".to_string(),
                    model: request.model.clone(),
                    choices: vec![CompletionChoice {
                        index: 0,
                        message: Message::assistant("").with_tool_calls(vec![ToolCall {
                            id: "call_get_context".to_string(),
                            tool_type: ToolType::Function,
                            function: ToolCallFunction {
                                name: "get_customer_context".to_string(),
                                arguments: "{\"order_id\":\"123\"}".to_string(),
                            },
                        }]),
                        finish_reason: FinishReason::ToolCalls,
                        logprobs: None,
                    }],
                    usage,
                    created: None,
                    provider: Some(self.name().to_string()),
                    healing_metadata: None,
                }
            } else {
                let mut usage = Usage::new(12, 6);
                usage.reasoning_tokens = Some(3);
                CompletionResponse {
                    id: "resp_tool_reasoning_2".to_string(),
                    model: request.model,
                    choices: vec![CompletionChoice {
                        index: 0,
                        message: Message::assistant("{\"state\":\"done\"}"),
                        finish_reason: FinishReason::Stop,
                        logprobs: None,
                    }],
                    usage,
                    created: None,
                    provider: Some(self.name().to_string()),
                    healing_metadata: None,
                }
            };

            let body = serde_json::to_value(response).map_err(SimpleAgentsError::from)?;
            Ok(ProviderResponse::new(200, body))
        }

        fn transform_response(&self, resp: ProviderResponse) -> SaResult<CompletionResponse> {
            serde_json::from_value(resp.body).map_err(SimpleAgentsError::from)
        }
    }

    struct FixedToolWorker {
        payload: Value,
    }

    struct CountingToolWorker {
        execute_calls: AtomicUsize,
    }

    #[async_trait]
    impl YamlWorkflowCustomWorkerExecutor for FixedToolWorker {
        async fn execute(
            &self,
            _handler: &str,
            _payload: &Value,
            _email_text: &str,
            _context: &Value,
        ) -> Result<Value, String> {
            Ok(self.payload.clone())
        }
    }

    #[async_trait]
    impl YamlWorkflowCustomWorkerExecutor for CountingToolWorker {
        async fn execute(
            &self,
            _handler: &str,
            _payload: &Value,
            _email_text: &str,
            _context: &Value,
        ) -> Result<Value, String> {
            self.execute_calls.fetch_add(1, Ordering::SeqCst);
            Ok(json!({"ok": true}))
        }
    }

    #[async_trait]
    impl YamlWorkflowLlmExecutor for CountingExecutor {
        async fn complete_structured(
            &self,
            _request: YamlLlmExecutionRequest,
            _event_sink: Option<&dyn YamlWorkflowEventSink>,
        ) -> Result<YamlLlmExecutionResult, String> {
            self.call_count.fetch_add(1, Ordering::SeqCst);
            Ok(YamlLlmExecutionResult {
                payload: json!({"state":"ok"}),
                usage: None,
                ttft_ms: None,
                tool_calls: Vec::new(),
            })
        }
    }

    impl YamlWorkflowEventSink for RecordingSink {
        fn emit(&self, event: &YamlWorkflowEvent) {
            self.events
                .lock()
                .expect("recording sink lock should not be poisoned")
                .push(event.clone());
        }
    }

    #[async_trait]
    impl YamlWorkflowCustomWorkerExecutor for CapturingWorker {
        async fn execute(
            &self,
            _handler: &str,
            _payload: &Value,
            _email_text: &str,
            context: &Value,
        ) -> Result<Value, String> {
            let mut guard = self
                .context
                .lock()
                .map_err(|_| "capturing worker lock should not be poisoned".to_string())?;
            *guard = Some(context.clone());
            Ok(json!({"ok": true}))
        }
    }

    #[async_trait]
    impl YamlWorkflowLlmExecutor for MockExecutor {
        async fn complete_structured(
            &self,
            request: YamlLlmExecutionRequest,
            _event_sink: Option<&dyn YamlWorkflowEventSink>,
        ) -> Result<YamlLlmExecutionResult, String> {
            let prompt = request.prompt;
            if prompt.contains("exactly one category") {
                return Ok(YamlLlmExecutionResult {
                    payload: json!({"category":"termination","reason":"mock"}),
                    usage: Some(YamlLlmTokenUsage {
                        prompt_tokens: 10,
                        completion_tokens: 5,
                        total_tokens: 15,
                        reasoning_tokens: None,
                    }),
                    ttft_ms: None,
                    tool_calls: Vec::new(),
                });
            }
            if prompt.contains("Determine termination subtype") {
                return Ok(YamlLlmExecutionResult {
                    payload: json!({"subtype":"repeated_offense","reason":"mock"}),
                    usage: Some(YamlLlmTokenUsage {
                        prompt_tokens: 12,
                        completion_tokens: 6,
                        total_tokens: 18,
                        reasoning_tokens: None,
                    }),
                    ttft_ms: None,
                    tool_calls: Vec::new(),
                });
            }
            if prompt.contains("Determine supply chain subtype") {
                return Ok(YamlLlmExecutionResult {
                    payload: json!({"subtype":"order_replacement","reason":"mock"}),
                    usage: Some(YamlLlmTokenUsage {
                        prompt_tokens: 11,
                        completion_tokens: 4,
                        total_tokens: 15,
                        reasoning_tokens: None,
                    }),
                    ttft_ms: None,
                    tool_calls: Vec::new(),
                });
            }
            Err("unexpected prompt".to_string())
        }
    }

    #[tokio::test]
    async fn runs_yaml_workflow_and_returns_step_timings() {
        let yaml = r#"
id: email-intake-classification
entry_node: classify_top_level
nodes:
  - id: classify_top_level
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
  - id: route_top_level
    node_type:
      switch:
        branches:
          - condition: '$.nodes.classify_top_level.output.category == "termination"'
            target: classify_termination_subtype
        default: rag_clarification
  - id: classify_termination_subtype
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Determine termination subtype:
        {{ input.email_text }}
  - id: route_termination_subtype
    node_type:
      switch:
        branches:
          - condition: '$.nodes.classify_termination_subtype.output.subtype == "repeated_offense"'
            target: rag_termination_repeated_offense
        default: rag_clarification
  - id: rag_termination_repeated_offense
    node_type:
      custom_worker:
        handler: GetRagData
    config:
      payload:
        topic: termination_repeated_offense
  - id: rag_clarification
    node_type:
      custom_worker:
        handler: GetRagData
    config:
      payload:
        topic: clarification
edges:
  - from: classify_top_level
    to: route_top_level
  - from: classify_termination_subtype
    to: route_termination_subtype
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let output = run_email_workflow_yaml(&workflow, "test", &MockExecutor)
            .await
            .expect("yaml workflow should execute");

        assert_eq!(output.workflow_id, "email-intake-classification");
        assert_eq!(output.terminal_node, "rag_termination_repeated_offense");
        assert!(!output.step_timings.is_empty());
        assert_eq!(output.step_timings.len(), output.trace.len());
        assert!(output
            .outputs
            .contains_key("rag_termination_repeated_offense"));
        assert_eq!(output.total_input_tokens, 22);
        assert_eq!(output.total_output_tokens, 11);
        assert_eq!(output.total_tokens, 33);
    }

    #[tokio::test]
    async fn emits_resolved_llm_input_event_with_bindings() {
        let yaml = r#"
id: email-intake-classification
entry_node: classify_top_level
nodes:
  - id: classify_top_level
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let sink = RecordingSink {
            events: Mutex::new(Vec::new()),
        };

        let output = run_email_workflow_yaml_with_custom_worker_and_events(
            &workflow,
            "Need help with termination",
            &MockExecutor,
            None,
            Some(&sink),
        )
        .await
        .expect("yaml workflow should execute");

        assert_eq!(output.terminal_node, "classify_top_level");

        let events = sink
            .events
            .lock()
            .expect("recording sink lock should not be poisoned");
        let llm_event = events
            .iter()
            .find(|event| event.event_type == "node_llm_input_resolved")
            .expect("expected llm input telemetry event");

        let metadata = llm_event
            .metadata
            .as_ref()
            .expect("llm input event must include metadata");
        assert_eq!(metadata["model"], Value::String("gpt-4.1".to_string()));
        assert_eq!(metadata["stream_requested"], Value::Bool(false));
        assert_eq!(metadata["heal_requested"], Value::Bool(false));
        assert!(metadata["prompt"]
            .as_str()
            .expect("prompt should be a string")
            .contains("Need help with termination"));

        let bindings = metadata["bindings"]
            .as_array()
            .expect("bindings should be an array");
        assert_eq!(bindings.len(), 1);
        assert_eq!(
            bindings[0]["source_path"],
            Value::String("input.email_text".to_string())
        );
        assert_eq!(
            bindings[0]["resolved"],
            Value::String("Need help with termination".to_string())
        );
        assert_eq!(bindings[0]["missing"], Value::Bool(false));
        assert_eq!(
            bindings[0]["resolved_type"],
            Value::String("string".to_string())
        );
    }

    #[tokio::test]
    async fn workflow_completed_event_includes_nerdstats_by_default() {
        let yaml = r#"
id: nerdstats-default
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let sink = RecordingSink {
            events: Mutex::new(Vec::new()),
        };

        let output = run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &MockExecutor,
            None,
            Some(&sink),
            &YamlWorkflowRunOptions::default(),
        )
        .await
        .expect("workflow should execute");

        let events = sink
            .events
            .lock()
            .expect("recording sink lock should not be poisoned");
        let completed = events
            .iter()
            .find(|event| event.event_type == "workflow_completed")
            .expect("expected workflow_completed event");
        let metadata = completed
            .metadata
            .as_ref()
            .expect("workflow_completed should include metadata by default");
        let nerdstats = metadata
            .get("nerdstats")
            .expect("nerdstats should be present by default");

        assert_eq!(nerdstats["workflow_id"], Value::String(output.workflow_id));
        assert_eq!(
            nerdstats["terminal_node"],
            Value::String(output.terminal_node)
        );
        assert_eq!(
            nerdstats["total_tokens"],
            Value::Number(output.total_tokens.into())
        );
        assert_eq!(nerdstats["token_metrics_available"], Value::Bool(true));
        assert_eq!(
            nerdstats["token_metrics_source"],
            Value::String("provider_usage".to_string())
        );
        assert!(nerdstats.get("step_timings").is_none());
        assert!(nerdstats.get("llm_node_metrics").is_none());
        assert!(nerdstats.get("step_details").is_some());
        assert!(nerdstats.get("llm_node_models").is_none());
        assert_eq!(
            nerdstats["step_details"][0]["model_name"],
            Value::String("gpt-4.1".to_string())
        );
        assert_eq!(
            nerdstats["step_details"][0]["node_id"],
            Value::String("classify".to_string())
        );
        assert_eq!(nerdstats["ttft_ms"], Value::Null);
    }

    #[tokio::test]
    async fn workflow_completed_event_omits_nerdstats_when_disabled() {
        let yaml = r#"
id: nerdstats-disabled
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let sink = RecordingSink {
            events: Mutex::new(Vec::new()),
        };
        let options = YamlWorkflowRunOptions {
            telemetry: YamlWorkflowTelemetryConfig {
                nerdstats: false,
                ..YamlWorkflowTelemetryConfig::default()
            },
            ..YamlWorkflowRunOptions::default()
        };

        run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &MockExecutor,
            None,
            Some(&sink),
            &options,
        )
        .await
        .expect("workflow should execute");

        let events = sink
            .events
            .lock()
            .expect("recording sink lock should not be poisoned");
        let completed = events
            .iter()
            .find(|event| event.event_type == "workflow_completed")
            .expect("expected workflow_completed event");
        assert!(completed.metadata.is_none());
    }

    struct StreamAwareMockExecutor;

    #[async_trait]
    impl YamlWorkflowLlmExecutor for StreamAwareMockExecutor {
        async fn complete_structured(
            &self,
            request: YamlLlmExecutionRequest,
            _event_sink: Option<&dyn YamlWorkflowEventSink>,
        ) -> Result<YamlLlmExecutionResult, String> {
            Ok(YamlLlmExecutionResult {
                payload: json!({"state":"ok"}),
                usage: Some(YamlLlmTokenUsage {
                    prompt_tokens: 20,
                    completion_tokens: 10,
                    total_tokens: 30,
                    reasoning_tokens: None,
                }),
                ttft_ms: if request.stream { Some(12) } else { None },
                tool_calls: Vec::new(),
            })
        }
    }

    #[tokio::test]
    async fn workflow_completed_event_includes_nerdstats_for_streaming_nodes() {
        let yaml = r#"
id: nerdstats-streaming
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
        stream: true
    config:
      prompt: |
        Return JSON only:
        {"state":"ok"}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let sink = RecordingSink {
            events: Mutex::new(Vec::new()),
        };

        run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &StreamAwareMockExecutor,
            None,
            Some(&sink),
            &YamlWorkflowRunOptions::default(),
        )
        .await
        .expect("workflow should execute");

        let events = sink
            .events
            .lock()
            .expect("recording sink lock should not be poisoned");
        let completed = events
            .iter()
            .find(|event| event.event_type == "workflow_completed")
            .expect("expected workflow_completed event");
        let metadata = completed
            .metadata
            .as_ref()
            .expect("workflow_completed should include metadata by default");
        let nerdstats = metadata
            .get("nerdstats")
            .expect("nerdstats should be present by default");

        assert_eq!(nerdstats["token_metrics_available"], Value::Bool(true));
        assert_eq!(nerdstats["total_tokens"], Value::Number(30u64.into()));
        assert_eq!(nerdstats["ttft_ms"], Value::Number(12u64.into()));
        assert!(nerdstats.get("step_timings").is_none());
        assert!(nerdstats.get("llm_node_metrics").is_none());
        assert_eq!(
            nerdstats["step_details"][0]["model_name"],
            Value::String("gpt-4.1".to_string())
        );
        assert_eq!(
            nerdstats["step_details"][0]["node_id"],
            Value::String("classify".to_string())
        );
    }

    #[test]
    fn workflow_nerdstats_marks_stream_token_metrics_unavailable() {
        let output = YamlWorkflowRunOutput {
            workflow_id: "workflow".to_string(),
            entry_node: "start".to_string(),
            email_text: "hello".to_string(),
            trace: vec!["llm_node".to_string()],
            outputs: BTreeMap::new(),
            terminal_node: "llm_node".to_string(),
            terminal_output: None,
            step_timings: vec![YamlStepTiming {
                node_id: "llm_node".to_string(),
                node_kind: "llm_call".to_string(),
                model_name: Some("gpt-4.1".to_string()),
                elapsed_ms: 100,
                prompt_tokens: None,
                completion_tokens: None,
                total_tokens: None,
                reasoning_tokens: None,
                tokens_per_second: None,
            }],
            llm_node_metrics: BTreeMap::new(),
            llm_node_models: BTreeMap::new(),
            total_elapsed_ms: 100,
            ttft_ms: None,
            total_input_tokens: 0,
            total_output_tokens: 0,
            total_tokens: 0,
            total_reasoning_tokens: None,
            tokens_per_second: 0.0,
            trace_id: Some("trace-1".to_string()),
            metadata: None,
        };

        let nerdstats = workflow_nerdstats(&output);
        assert_eq!(nerdstats["token_metrics_available"], Value::Bool(false));
        assert_eq!(
            nerdstats["token_metrics_source"],
            Value::String("provider_stream_usage_unavailable".to_string())
        );
        assert_eq!(nerdstats["total_tokens"], Value::Null);
        assert_eq!(nerdstats["ttft_ms"], Value::Null);
        assert!(nerdstats.get("step_timings").is_none());
        assert!(nerdstats.get("llm_node_metrics").is_none());
        assert_eq!(
            nerdstats["step_details"][0]["node_id"],
            Value::String("llm_node".to_string())
        );
        assert_eq!(
            nerdstats["step_details"][0]["model_name"],
            Value::String("gpt-4.1".to_string())
        );
    }

    #[test]
    fn workflow_nerdstats_includes_ttft_when_available() {
        let output = YamlWorkflowRunOutput {
            workflow_id: "workflow".to_string(),
            entry_node: "start".to_string(),
            email_text: "hello".to_string(),
            trace: vec!["llm_node".to_string()],
            outputs: BTreeMap::new(),
            terminal_node: "llm_node".to_string(),
            terminal_output: None,
            step_timings: vec![YamlStepTiming {
                node_id: "llm_node".to_string(),
                node_kind: "llm_call".to_string(),
                model_name: Some("gpt-4.1".to_string()),
                elapsed_ms: 100,
                prompt_tokens: Some(10),
                completion_tokens: Some(15),
                total_tokens: Some(25),
                reasoning_tokens: None,
                tokens_per_second: Some(150.0),
            }],
            llm_node_metrics: BTreeMap::new(),
            llm_node_models: BTreeMap::new(),
            total_elapsed_ms: 100,
            ttft_ms: Some(42),
            total_input_tokens: 10,
            total_output_tokens: 15,
            total_tokens: 25,
            total_reasoning_tokens: None,
            tokens_per_second: 150.0,
            trace_id: Some("trace-2".to_string()),
            metadata: None,
        };

        let nerdstats = workflow_nerdstats(&output);
        assert_eq!(nerdstats["ttft_ms"], Value::Number(42u64.into()));
        assert!(nerdstats.get("step_timings").is_none());
        assert!(nerdstats.get("llm_node_metrics").is_none());
        assert_eq!(
            nerdstats["step_details"][0]["node_id"],
            Value::String("llm_node".to_string())
        );
        assert_eq!(
            nerdstats["step_details"][0]["model_name"],
            Value::String("gpt-4.1".to_string())
        );
    }

    #[test]
    fn workflow_nerdstats_schema_contract_is_stable() {
        let output = YamlWorkflowRunOutput {
            workflow_id: "schema-workflow".to_string(),
            entry_node: "start".to_string(),
            email_text: "hello".to_string(),
            trace: vec!["classify".to_string(), "route".to_string()],
            outputs: BTreeMap::new(),
            terminal_node: "route".to_string(),
            terminal_output: None,
            step_timings: vec![
                YamlStepTiming {
                    node_id: "classify".to_string(),
                    node_kind: "llm_call".to_string(),
                    model_name: Some("gpt-4.1".to_string()),
                    elapsed_ms: 100,
                    prompt_tokens: Some(11),
                    completion_tokens: Some(22),
                    total_tokens: Some(33),
                    reasoning_tokens: Some(7),
                    tokens_per_second: Some(220.0),
                },
                YamlStepTiming {
                    node_id: "route".to_string(),
                    node_kind: "switch".to_string(),
                    model_name: None,
                    elapsed_ms: 0,
                    prompt_tokens: None,
                    completion_tokens: None,
                    total_tokens: None,
                    reasoning_tokens: None,
                    tokens_per_second: None,
                },
            ],
            llm_node_metrics: BTreeMap::new(),
            llm_node_models: BTreeMap::new(),
            total_elapsed_ms: 100,
            ttft_ms: Some(9),
            total_input_tokens: 11,
            total_output_tokens: 22,
            total_tokens: 33,
            total_reasoning_tokens: Some(7),
            tokens_per_second: 220.0,
            trace_id: Some("trace-schema".to_string()),
            metadata: None,
        };

        let nerdstats = workflow_nerdstats(&output);
        let expected = json!({
            "workflow_id": "schema-workflow",
            "terminal_node": "route",
            "total_elapsed_ms": 100,
            "ttft_ms": 9,
            "step_details": [
                {
                    "node_id": "classify",
                    "node_kind": "llm_call",
                    "model_name": "gpt-4.1",
                    "elapsed_ms": 100,
                    "prompt_tokens": 11,
                    "completion_tokens": 22,
                    "total_tokens": 33,
                    "reasoning_tokens": 7,
                    "tokens_per_second": 220.0
                },
                {
                    "node_id": "route",
                    "node_kind": "switch",
                    "elapsed_ms": 0
                }
            ],
            "total_input_tokens": 11,
            "total_output_tokens": 22,
            "total_tokens": 33,
            "total_reasoning_tokens": 7,
            "tokens_per_second": 220.0,
            "trace_id": "trace-schema",
            "token_metrics_available": true,
            "token_metrics_source": "provider_usage",
            "llm_nodes_without_usage": []
        });

        assert_eq!(nerdstats, expected);
    }

    struct MessageHistoryExecutor;

    #[async_trait]
    impl YamlWorkflowLlmExecutor for MessageHistoryExecutor {
        async fn complete_structured(
            &self,
            request: YamlLlmExecutionRequest,
            _event_sink: Option<&dyn YamlWorkflowEventSink>,
        ) -> Result<YamlLlmExecutionResult, String> {
            let messages = request
                .messages
                .ok_or_else(|| "expected messages in request".to_string())?;
            if messages.len() != 2 {
                return Err(format!("expected 2 messages, got {}", messages.len()));
            }
            Ok(YamlLlmExecutionResult {
                payload: json!({"category":"termination","reason":"history"}),
                usage: Some(YamlLlmTokenUsage {
                    prompt_tokens: 7,
                    completion_tokens: 3,
                    total_tokens: 10,
                    reasoning_tokens: None,
                }),
                ttft_ms: None,
                tool_calls: Vec::new(),
            })
        }
    }

    #[tokio::test]
    async fn supports_messages_path_in_workflow_input() {
        let yaml = r#"
id: email-intake-classification
entry_node: classify_top_level
nodes:
  - id: classify_top_level
    node_type:
      llm_call:
        model: gpt-4.1
        messages_path: input.messages
        append_prompt_as_user: false
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let input = json!({
            "email_text": "ignored",
            "messages": [
                {"role": "system", "content": "You are a classifier"},
                {"role": "user", "content": "Please classify this"}
            ]
        });

        let output = run_workflow_yaml(&workflow, &input, &MessageHistoryExecutor)
            .await
            .expect("workflow should use chat history from input");

        assert_eq!(output.terminal_node, "classify_top_level");
        assert_eq!(
            output.outputs["classify_top_level"]["output"]["reason"],
            Value::String("history".to_string())
        );
    }

    #[tokio::test]
    async fn wrapper_entrypoints_produce_equivalent_outputs() {
        let yaml = r#"
id: wrapper-equivalence
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let input = json!({"email_text":"hello"});

        let a = run_workflow_yaml(&workflow, &input, &MockExecutor)
            .await
            .expect("base entrypoint should execute");
        let b = run_workflow_yaml_with_custom_worker(&workflow, &input, &MockExecutor, None)
            .await
            .expect("custom worker wrapper should execute");
        let c = run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &input,
            &MockExecutor,
            None,
            None,
            &YamlWorkflowRunOptions::default(),
        )
        .await
        .expect("events/options wrapper should execute");

        assert_eq!(a.workflow_id, b.workflow_id);
        assert_eq!(a.workflow_id, c.workflow_id);
        assert_eq!(a.terminal_node, b.terminal_node);
        assert_eq!(a.terminal_node, c.terminal_node);
        assert_eq!(a.outputs, b.outputs);
        assert_eq!(a.outputs, c.outputs);
        assert_eq!(a.total_tokens, b.total_tokens);
        assert_eq!(a.total_tokens, c.total_tokens);
    }

    #[tokio::test]
    async fn yaml_llm_tool_calling_captures_traces_and_supports_globals_reference() {
        let yaml = r#"
id: tool-calling-workflow
entry_node: generate_with_tool
nodes:
  - id: generate_with_tool
    node_type:
      llm_call:
        model: gpt-4.1
        tools_format: simplified
        max_tool_roundtrips: 1
        tool_calls_global_key: audit
        tools:
          - name: get_customer_context
            description: Fetch customer context
            input_schema:
              type: object
              properties:
                order_id: { type: string }
              required: [order_id]
              additionalProperties: false
            output_schema:
              type: object
              properties:
                customer_name: { type: string }
              required: [customer_name]
              additionalProperties: false
    config:
      output_schema:
        type: object
        properties:
          state: { type: string }
        required: [state]
  - id: personalize
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Write an email greeting for {{ globals.audit.0.output.customer_name }}.
      output_schema:
        type: object
        properties:
          subject: { type: string }
          body: { type: string }
        required: [subject, body]
edges:
  - from: generate_with_tool
    to: personalize
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let client = SimpleAgentsClientBuilder::new()
            .with_provider(Arc::new(ToolLoopProvider))
            .build()
            .expect("client should build");
        let worker = FixedToolWorker {
            payload: json!({"customer_name": "Ava"}),
        };

        let output = run_workflow_yaml_with_client_and_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &client,
            Some(&worker),
            None,
            &YamlWorkflowRunOptions::default(),
        )
        .await
        .expect("workflow should execute");

        assert_eq!(output.trace, vec!["generate_with_tool", "personalize"]);
        assert_eq!(
            output.outputs["generate_with_tool"]["tool_calls"][0]["output"]["customer_name"],
            Value::String("Ava".to_string())
        );
        let body = output.outputs["personalize"]["output"]["body"]
            .as_str()
            .expect("body should be string");
        assert!(body.contains("Ava"));
    }

    #[tokio::test]
    async fn workflow_with_client_preserves_reasoning_tokens_in_output_and_nerdstats() {
        let yaml = r#"
id: reasoning-usage-workflow
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Return JSON only:
        {"state":"ok"}
      output_schema:
        type: object
        properties:
          state: { type: string }
        required: [state]
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let client = SimpleAgentsClientBuilder::new()
            .with_provider(Arc::new(ReasoningUsageProvider))
            .build()
            .expect("client should build");

        let output = run_workflow_yaml_with_client_and_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &client,
            None,
            None,
            &YamlWorkflowRunOptions::default(),
        )
        .await
        .expect("workflow should execute");

        assert_eq!(output.total_reasoning_tokens, Some(4));
        assert_eq!(output.step_timings.len(), 1);
        assert_eq!(output.step_timings[0].reasoning_tokens, Some(4));

        let nerdstats = workflow_nerdstats(&output);
        assert_eq!(
            nerdstats["total_reasoning_tokens"],
            Value::Number(4u64.into())
        );
        assert_eq!(
            nerdstats["step_details"][0]["reasoning_tokens"],
            Value::Number(4u64.into())
        );
    }

    #[tokio::test]
    async fn workflow_with_tools_accumulates_reasoning_tokens_across_roundtrips() {
        let yaml = r#"
id: tool-reasoning-workflow
entry_node: generate_with_tool
nodes:
  - id: generate_with_tool
    node_type:
      llm_call:
        model: gpt-4.1
        tools_format: simplified
        max_tool_roundtrips: 1
        tools:
          - name: get_customer_context
            input_schema:
              type: object
              properties:
                order_id: { type: string }
              required: [order_id]
              additionalProperties: false
            output_schema:
              type: object
              properties:
                customer_name: { type: string }
              required: [customer_name]
              additionalProperties: false
    config:
      output_schema:
        type: object
        properties:
          state: { type: string }
        required: [state]
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let client = SimpleAgentsClientBuilder::new()
            .with_provider(Arc::new(ToolLoopReasoningProvider))
            .build()
            .expect("client should build");
        let worker = FixedToolWorker {
            payload: json!({"customer_name": "Ava"}),
        };

        let output = run_workflow_yaml_with_client_and_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &client,
            Some(&worker),
            None,
            &YamlWorkflowRunOptions::default(),
        )
        .await
        .expect("workflow should execute");

        assert_eq!(output.total_reasoning_tokens, Some(5));
        assert_eq!(output.step_timings.len(), 1);
        assert_eq!(output.step_timings[0].reasoning_tokens, Some(5));

        let nerdstats = workflow_nerdstats(&output);
        assert_eq!(
            nerdstats["total_reasoning_tokens"],
            Value::Number(5u64.into())
        );
        assert_eq!(
            nerdstats["step_details"][0]["reasoning_tokens"],
            Value::Number(5u64.into())
        );
    }

    #[tokio::test]
    async fn yaml_llm_tool_output_schema_mismatch_hard_fails_node() {
        let yaml = r#"
id: tool-calling-schema-fail
entry_node: generate_with_tool
nodes:
  - id: generate_with_tool
    node_type:
      llm_call:
        model: gpt-4.1
        tools_format: simplified
        max_tool_roundtrips: 1
        tools:
          - name: get_customer_context
            input_schema:
              type: object
              properties:
                order_id: { type: string }
              required: [order_id]
              additionalProperties: false
            output_schema:
              type: object
              properties:
                customer_name: { type: string }
              required: [customer_name]
              additionalProperties: false
    config:
      output_schema:
        type: object
        properties:
          state: { type: string }
        required: [state]
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let client = SimpleAgentsClientBuilder::new()
            .with_provider(Arc::new(ToolLoopProvider))
            .build()
            .expect("client should build");
        let worker = FixedToolWorker {
            payload: json!({"unexpected": "shape"}),
        };

        let error = run_workflow_yaml_with_client_and_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &client,
            Some(&worker),
            None,
            &YamlWorkflowRunOptions::default(),
        )
        .await
        .expect_err("workflow should hard-fail on schema mismatch");

        match error {
            YamlWorkflowRunError::Llm { message, .. } => {
                assert!(message.contains("output failed schema validation"));
            }
            other => panic!("expected llm error, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn yaml_llm_unknown_tool_is_rejected_before_custom_worker_execution() {
        let yaml = r#"
id: unknown-tool-rejected
entry_node: generate_with_tool
nodes:
  - id: generate_with_tool
    node_type:
      llm_call:
        model: gpt-4.1
        tools_format: simplified
        max_tool_roundtrips: 1
        tools:
          - name: get_customer_context
            input_schema:
              type: object
              properties:
                order_id: { type: string }
              required: [order_id]
    config:
      output_schema:
        type: object
        properties:
          state: { type: string }
        required: [state]
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let client = SimpleAgentsClientBuilder::new()
            .with_provider(Arc::new(UnknownToolProvider))
            .build()
            .expect("client should build");
        let worker = CountingToolWorker {
            execute_calls: AtomicUsize::new(0),
        };

        let error = run_workflow_yaml_with_client_and_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &client,
            Some(&worker),
            None,
            &YamlWorkflowRunOptions::default(),
        )
        .await
        .expect_err("workflow should reject unknown tool before executing worker");

        match error {
            YamlWorkflowRunError::Llm { message, .. } => {
                assert!(message.contains("model requested unknown tool 'unknown_tool'"));
            }
            other => panic!("expected llm error, got {other:?}"),
        }

        assert_eq!(worker.execute_calls.load(Ordering::SeqCst), 0);
    }

    #[test]
    fn validates_tools_format_mismatch() {
        let yaml = r#"
id: mismatch
entry_node: generate
nodes:
  - id: generate
    node_type:
      llm_call:
        model: gpt-4.1
        tools_format: openai
        tools:
          - name: get_customer_context
            input_schema:
              type: object
              properties:
                order_id: { type: string }
              required: [order_id]
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let diagnostics = verify_yaml_workflow(&workflow);
        assert!(diagnostics
            .iter()
            .any(|diagnostic| diagnostic.code == "invalid_tools_format"));
    }

    #[test]
    fn mock_custom_worker_supports_get_employee_record() {
        let result = mock_custom_worker_output(
            "get_employee_record",
            &json!({"employee_name": "Alex Johnson"}),
        )
        .expect("mock tool should resolve employee record");

        assert_eq!(result["employee_id"], Value::String("EMP-2041".to_string()));
        assert_eq!(result["location"], Value::String("Austin".to_string()));
    }

    #[tokio::test]
    async fn custom_worker_receives_trace_context_block() {
        let yaml = r#"
id: custom-worker-trace-context
entry_node: lookup
nodes:
  - id: lookup
    node_type:
      custom_worker:
        handler: GetRagData
    config:
      payload:
        topic: demo
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let worker = CapturingWorker {
            context: Mutex::new(None),
        };
        let options = YamlWorkflowRunOptions {
            trace: YamlWorkflowTraceOptions {
                context: Some(YamlWorkflowTraceContextInput {
                    trace_id: Some("trace-fixed-ctx".to_string()),
                    traceparent: Some("00-trace-fixed-ctx-span-fixed-01".to_string()),
                    ..YamlWorkflowTraceContextInput::default()
                }),
                tenant: YamlWorkflowTraceTenantContext {
                    conversation_id: Some("7fd67af3-c67d-46cb-95af-08f8ed7b06a5".to_string()),
                    request_id: Some("turn-7".to_string()),
                    ..YamlWorkflowTraceTenantContext::default()
                },
            },
            ..YamlWorkflowRunOptions::default()
        };

        run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &MockExecutor,
            Some(&worker),
            None,
            &options,
        )
        .await
        .expect("workflow should execute");

        let captured_context = worker
            .context
            .lock()
            .expect("capturing worker lock should not be poisoned")
            .clone()
            .expect("custom worker should receive context");

        assert_eq!(
            captured_context
                .get("trace")
                .and_then(|trace| trace.get("context"))
                .and_then(|context| context.get("trace_id"))
                .and_then(Value::as_str),
            Some("trace-fixed-ctx")
        );
        assert_eq!(
            captured_context
                .get("trace")
                .and_then(|trace| trace.get("context"))
                .and_then(|context| context.get("traceparent"))
                .and_then(Value::as_str),
            Some("00-trace-fixed-ctx-span-fixed-01")
        );
        assert_eq!(
            captured_context
                .get("trace")
                .and_then(|trace| trace.get("tenant"))
                .and_then(|tenant| tenant.get("conversation_id"))
                .and_then(Value::as_str),
            Some("7fd67af3-c67d-46cb-95af-08f8ed7b06a5")
        );
    }

    #[tokio::test]
    async fn event_sink_cancellation_stops_workflow_before_llm_execution() {
        let yaml = r#"
id: cancellation-test
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let executor = CountingExecutor {
            call_count: AtomicUsize::new(0),
        };
        let sink = CancelAfterFirstEventSink {
            cancelled: AtomicBool::new(false),
        };

        let err = run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &executor,
            None,
            Some(&sink),
            &YamlWorkflowRunOptions::default(),
        )
        .await
        .expect_err("workflow should stop when sink signals cancellation");

        assert!(matches!(
            err,
            YamlWorkflowRunError::EventSinkCancelled { .. }
        ));
        assert_eq!(executor.call_count.load(Ordering::SeqCst), 0);
    }

    #[tokio::test]
    async fn rejects_invalid_messages_path_shape() {
        let yaml = r#"
id: email-intake-classification
entry_node: classify_top_level
nodes:
  - id: classify_top_level
    node_type:
      llm_call:
        model: gpt-4.1
        messages_path: input.messages
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let input = json!({
            "email_text": "ignored",
            "messages": "not-a-list"
        });

        let err = run_workflow_yaml(&workflow, &input, &MessageHistoryExecutor)
            .await
            .expect_err("workflow should fail for invalid messages shape");

        assert!(matches!(err, YamlWorkflowRunError::Llm { .. }));
    }

    #[test]
    fn renders_yaml_workflow_to_mermaid_with_switch_labels() {
        let yaml = r#"
id: chat-workflow
entry_node: decide
nodes:
  - id: decide
    node_type:
      switch:
        branches:
          - condition: '$.input.mode == "draft"'
            target: draft
        default: ask
  - id: draft
    node_type:
      llm_call:
        model: gpt-4.1
  - id: ask
    node_type:
      llm_call:
        model: gpt-4.1
edges:
  - from: draft
    to: ask
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let mermaid = yaml_workflow_to_mermaid(&workflow);

        assert!(mermaid.contains("flowchart TD"));
        assert!(mermaid.contains("decide -- \"route1\" --> draft"));
        assert!(mermaid.contains("decide -- \"default\" --> ask"));
        assert!(mermaid.contains("draft --> ask"));
    }

    #[test]
    fn renders_yaml_workflow_tools_as_colored_tool_nodes() {
        let yaml = r#"
id: tool-graph
entry_node: chat
nodes:
  - id: chat
    node_type:
      llm_call:
        model: gemini-3-flash
        tools_format: simplified
        tools:
          - name: run_workflow_graph
            input_schema:
              type: object
              properties:
                workflow_id: { type: string }
              required: [workflow_id]
edges: []
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let mermaid = yaml_workflow_to_mermaid(&workflow);

        assert!(mermaid.contains("chat__tool_0"));
        assert!(mermaid.contains("chat -.-> chat__tool_0"));
        assert!(mermaid.contains("classDef toolNode"));
        assert!(mermaid.contains("class chat__tool_0 toolNode;"));
    }

    #[test]
    fn renders_yaml_workflow_file_to_mermaid_with_subgraph_cluster_when_present() {
        let base_dir = std::env::temp_dir().join(format!(
            "simple_agents_mermaid_subgraph_{}",
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .expect("unix epoch")
                .as_nanos()
        ));
        fs::create_dir_all(&base_dir).expect("temp dir should be created");

        let orchestrator_path = base_dir.join("email-chat-orchestrator-with-subgraph-tool.yaml");
        let subgraph_path = base_dir.join("hr-warning-email-subgraph.yaml");

        let orchestrator_yaml = r#"
id: email-chat-orchestrator-with-subgraph-tool
entry_node: respond_casual
nodes:
  - id: respond_casual
    node_type:
      llm_call:
        model: gemini-3-flash
        tools_format: simplified
        tools:
          - name: run_workflow_graph
            input_schema:
              type: object
              properties:
                workflow_id: { type: string }
              required: [workflow_id]
    config:
      prompt: |
        Call with:
        {
          "workflow_id": "hr_warning_email_subgraph"
        }
edges: []
"#;

        let subgraph_yaml = r#"
id: hr-warning-email-subgraph
entry_node: draft_hr_warning_email
nodes:
  - id: draft_hr_warning_email
    node_type:
      llm_call:
        model: gemini-3-flash
edges: []
"#;

        fs::write(&orchestrator_path, orchestrator_yaml).expect("orchestrator yaml written");
        fs::write(&subgraph_path, subgraph_yaml).expect("subgraph yaml written");

        let mermaid =
            yaml_workflow_file_to_mermaid(&orchestrator_path).expect("mermaid should render");

        assert!(mermaid.contains("Main: email-chat-orchestrator-with-subgraph-tool"));
        assert!(mermaid.contains("Subgraph: hr_warning_email_subgraph"));
        assert!(mermaid.contains("calls hr_warning_email_subgraph"));
        assert!(mermaid.contains("subgraph_1__draft_hr_warning_email"));

        fs::remove_dir_all(base_dir).expect("temp dir removed");
    }

    #[test]
    fn converts_yaml_workflow_to_ir_definition() {
        let yaml = r#"
id: chat-workflow
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        classify
  - id: route
    node_type:
      switch:
        branches:
          - condition: '$.nodes.classify.output.kind == "x"'
            target: done
        default: done
  - id: done
    node_type:
      custom_worker:
        handler: GetRagData
    config:
      payload:
        topic: test
edges:
  - from: classify
    to: route
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let ir = yaml_workflow_to_ir(&workflow).expect("yaml should convert to ir");

        assert_eq!(ir.name, "chat-workflow");
        assert!(ir.nodes.iter().any(|n| n.id == "__yaml_start"));
        assert!(ir.nodes.iter().any(|n| n.id == "classify"));
        assert!(ir.nodes.iter().any(|n| n.id == "route"));
        assert!(ir.nodes.iter().any(|n| n.id == "done"));
    }

    #[test]
    fn supports_yaml_to_ir_when_messages_path_is_used() {
        let yaml = r#"
id: chat-workflow
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
        messages_path: input.messages
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let ir =
            yaml_workflow_to_ir(&workflow).expect("messages_path should convert to tool-based IR");
        assert!(ir.nodes.iter().any(|node| matches!(
            node.kind,
            crate::ir::NodeKind::Tool { ref tool, .. } if tool == "__yaml_llm_call"
        )));
    }

    #[tokio::test]
    async fn workflow_output_contains_trace_id_in_both_locations() {
        let yaml = r#"
id: trace-test
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let output = run_workflow_yaml(&workflow, &json!({"email_text":"hello"}), &MockExecutor)
            .await
            .expect("workflow should execute");

        let trace_id = output
            .trace_id
            .as_deref()
            .expect("trace_id should be present");
        assert!(!trace_id.is_empty());
        assert_eq!(
            output.metadata.as_ref().and_then(|value| {
                value
                    .get("telemetry")
                    .and_then(|telemetry| telemetry.get("trace_id"))
                    .and_then(Value::as_str)
            }),
            Some(trace_id)
        );
    }

    #[tokio::test]
    async fn workflow_run_options_sample_rate_zero_marks_trace_unsampled() {
        let yaml = r#"
id: sample-rate-zero
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let options = YamlWorkflowRunOptions {
            telemetry: YamlWorkflowTelemetryConfig {
                sample_rate: 0.0,
                ..YamlWorkflowTelemetryConfig::default()
            },
            trace: YamlWorkflowTraceOptions {
                context: Some(YamlWorkflowTraceContextInput {
                    trace_id: Some("trace-sample-zero".to_string()),
                    ..YamlWorkflowTraceContextInput::default()
                }),
                tenant: YamlWorkflowTraceTenantContext::default(),
            },
            model: None,
        };

        let output = run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &MockExecutor,
            None,
            None,
            &options,
        )
        .await
        .expect("workflow should execute");

        assert_eq!(output.trace_id.as_deref(), Some("trace-sample-zero"));
        assert_eq!(
            output
                .metadata
                .as_ref()
                .and_then(|value| value.get("telemetry"))
                .and_then(|telemetry| telemetry.get("sampled"))
                .and_then(Value::as_bool),
            Some(false)
        );
    }

    #[test]
    fn trace_id_from_traceparent_parses_w3c_header() {
        let traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
        assert_eq!(
            trace_id_from_traceparent(traceparent).as_deref(),
            Some("4bf92f3577b34da6a3ce929d0e0e4736")
        );
    }

    #[test]
    fn resolve_telemetry_context_marks_traceparent_source() {
        let mut options = YamlWorkflowRunOptions::default();
        options.trace.context = Some(YamlWorkflowTraceContextInput {
            traceparent: Some(
                "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01".to_string(),
            ),
            ..YamlWorkflowTraceContextInput::default()
        });

        let context = resolve_telemetry_context(&options, None);

        assert_eq!(context.trace_id_source, TraceIdSource::Traceparent);
        assert_eq!(
            context.trace_id.as_deref(),
            Some("4bf92f3577b34da6a3ce929d0e0e4736")
        );
    }

    #[tokio::test]
    async fn workflow_run_options_derives_trace_id_from_traceparent_when_trace_id_missing() {
        let yaml = r#"
id: traceparent-derived
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let traceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01";
        let options = YamlWorkflowRunOptions {
            telemetry: YamlWorkflowTelemetryConfig {
                sample_rate: 0.0,
                ..YamlWorkflowTelemetryConfig::default()
            },
            trace: YamlWorkflowTraceOptions {
                context: Some(YamlWorkflowTraceContextInput {
                    traceparent: Some(traceparent.to_string()),
                    ..YamlWorkflowTraceContextInput::default()
                }),
                tenant: YamlWorkflowTraceTenantContext::default(),
            },
            model: None,
        };

        let output = run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &MockExecutor,
            None,
            None,
            &options,
        )
        .await
        .expect("workflow should execute");

        assert_eq!(
            output.trace_id.as_deref(),
            Some("4bf92f3577b34da6a3ce929d0e0e4736")
        );
        assert_eq!(
            output
                .metadata
                .as_ref()
                .and_then(|value| value.get("telemetry"))
                .and_then(|telemetry| telemetry.get("sampled"))
                .and_then(Value::as_bool),
            Some(false)
        );
    }

    #[tokio::test]
    async fn workflow_run_options_sample_rate_one_marks_trace_sampled() {
        let yaml = r#"
id: sample-rate-one
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let options = YamlWorkflowRunOptions {
            telemetry: YamlWorkflowTelemetryConfig {
                sample_rate: 1.0,
                ..YamlWorkflowTelemetryConfig::default()
            },
            trace: YamlWorkflowTraceOptions {
                context: Some(YamlWorkflowTraceContextInput {
                    trace_id: Some("trace-sample-one".to_string()),
                    ..YamlWorkflowTraceContextInput::default()
                }),
                tenant: YamlWorkflowTraceTenantContext::default(),
            },
            model: None,
        };

        let output = run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &MockExecutor,
            None,
            None,
            &options,
        )
        .await
        .expect("workflow should execute");

        assert_eq!(output.trace_id.as_deref(), Some("trace-sample-one"));
        assert_eq!(
            output
                .metadata
                .as_ref()
                .and_then(|value| value.get("telemetry"))
                .and_then(|telemetry| telemetry.get("sampled"))
                .and_then(Value::as_bool),
            Some(true)
        );
    }

    #[tokio::test]
    async fn workflow_run_options_sampling_is_deterministic_for_fixed_trace_id() {
        let yaml = r#"
id: sample-rate-deterministic
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let options = YamlWorkflowRunOptions {
            telemetry: YamlWorkflowTelemetryConfig {
                sample_rate: 0.5,
                ..YamlWorkflowTelemetryConfig::default()
            },
            trace: YamlWorkflowTraceOptions {
                context: Some(YamlWorkflowTraceContextInput {
                    trace_id: Some("trace-sample-deterministic".to_string()),
                    ..YamlWorkflowTraceContextInput::default()
                }),
                tenant: YamlWorkflowTraceTenantContext::default(),
            },
            model: None,
        };

        let mut sampled_values = Vec::new();
        for _ in 0..3 {
            let output = run_workflow_yaml_with_custom_worker_and_events_and_options(
                &workflow,
                &json!({"email_text":"hello"}),
                &MockExecutor,
                None,
                None,
                &options,
            )
            .await
            .expect("workflow should execute");

            let sampled = output
                .metadata
                .as_ref()
                .and_then(|value| value.get("telemetry"))
                .and_then(|telemetry| telemetry.get("sampled"))
                .and_then(Value::as_bool)
                .expect("sampled flag should be present");
            sampled_values.push(sampled);
        }

        assert!(sampled_values
            .iter()
            .all(|value| *value == sampled_values[0]));
    }

    #[tokio::test]
    async fn workflow_run_options_reject_invalid_sample_rate() {
        let yaml = r#"
id: sample-rate-invalid
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let options = YamlWorkflowRunOptions {
            telemetry: YamlWorkflowTelemetryConfig {
                sample_rate: 1.1,
                ..YamlWorkflowTelemetryConfig::default()
            },
            ..YamlWorkflowRunOptions::default()
        };

        let err = run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &MockExecutor,
            None,
            None,
            &options,
        )
        .await
        .expect_err("invalid sample_rate should fail");

        match err {
            YamlWorkflowRunError::InvalidInput { message } => {
                assert!(message.contains("telemetry.sample_rate"));
            }
            _ => panic!("expected invalid input error for sample_rate"),
        }
    }

    #[tokio::test]
    async fn workflow_run_options_reject_nan_sample_rate() {
        let yaml = r#"
id: sample-rate-invalid
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let options = YamlWorkflowRunOptions {
            telemetry: YamlWorkflowTelemetryConfig {
                sample_rate: f32::NAN,
                ..YamlWorkflowTelemetryConfig::default()
            },
            ..YamlWorkflowRunOptions::default()
        };

        let err = run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &MockExecutor,
            None,
            None,
            &options,
        )
        .await
        .expect_err("nan sample_rate should fail");

        assert!(matches!(err, YamlWorkflowRunError::InvalidInput { .. }));
    }

    #[test]
    fn subworkflow_options_inherit_parent_telemetry_configuration() {
        let mut parent_options = YamlWorkflowRunOptions::default();
        parent_options.telemetry.sample_rate = 0.0;
        parent_options.telemetry.payload_mode = YamlWorkflowPayloadMode::RedactedPayload;
        parent_options.telemetry.tool_trace_mode = YamlToolTraceMode::Redacted;
        parent_options.trace.tenant.conversation_id = Some("conv-123".to_string());

        let parent_context = TraceContext {
            trace_id: Some("trace-parent".to_string()),
            span_id: Some("span-parent".to_string()),
            parent_span_id: None,
            traceparent: Some("00-trace-parent-span-parent-01".to_string()),
            tracestate: None,
            baggage: BTreeMap::new(),
        };

        let subworkflow_options =
            build_subworkflow_options(&parent_options, Some(&parent_context), None);

        assert_eq!(subworkflow_options.telemetry.sample_rate, 0.0);
        assert_eq!(
            subworkflow_options.telemetry.payload_mode,
            YamlWorkflowPayloadMode::RedactedPayload
        );
        assert_eq!(
            subworkflow_options.telemetry.tool_trace_mode,
            YamlToolTraceMode::Redacted
        );
        assert_eq!(
            subworkflow_options.trace.tenant.conversation_id.as_deref(),
            Some("conv-123")
        );
        assert_eq!(
            subworkflow_options
                .trace
                .context
                .as_ref()
                .and_then(|ctx| ctx.trace_id.as_deref()),
            Some("trace-parent")
        );
    }

    #[tokio::test]
    async fn workflow_run_options_use_explicit_trace_id_and_payload_mode() {
        let yaml = r#"
id: trace-options-test
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        Classify this email into exactly one category:
        {{ input.email_text }}
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let options = YamlWorkflowRunOptions {
            telemetry: YamlWorkflowTelemetryConfig {
                payload_mode: YamlWorkflowPayloadMode::RedactedPayload,
                ..YamlWorkflowTelemetryConfig::default()
            },
            trace: YamlWorkflowTraceOptions {
                context: Some(YamlWorkflowTraceContextInput {
                    trace_id: Some("trace-fixed-123".to_string()),
                    traceparent: Some("00-trace-fixed-123-span-1-01".to_string()),
                    ..YamlWorkflowTraceContextInput::default()
                }),
                tenant: YamlWorkflowTraceTenantContext {
                    conversation_id: Some("6e6d3125-b9f1-4af2-af1f-7cca024a2c42".to_string()),
                    ..YamlWorkflowTraceTenantContext::default()
                },
            },
            model: None,
        };

        let output = run_workflow_yaml_with_custom_worker_and_events_and_options(
            &workflow,
            &json!({"email_text":"hello"}),
            &MockExecutor,
            None,
            None,
            &options,
        )
        .await
        .expect("workflow should execute");

        assert_eq!(output.trace_id.as_deref(), Some("trace-fixed-123"));
        assert_eq!(
            output
                .metadata
                .as_ref()
                .and_then(|value| value.get("telemetry"))
                .and_then(|telemetry| telemetry.get("payload_mode"))
                .and_then(Value::as_str),
            Some("redacted_payload")
        );
        assert_eq!(
            output
                .metadata
                .as_ref()
                .and_then(|value| value.get("trace"))
                .and_then(|trace| trace.get("tenant"))
                .and_then(|tenant| tenant.get("conversation_id"))
                .and_then(Value::as_str),
            Some("6e6d3125-b9f1-4af2-af1f-7cca024a2c42")
        );
    }

    #[test]
    fn streamed_payload_parser_extracts_last_json_object() {
        let raw = r#"{"state":"missing_scenario","reason":"ok"}

extra reasoning text

{"state":"ready","reason":"final"}"#;

        let resolved = parse_streamed_structured_payload(raw, false)
            .expect("parser should extract final JSON object");
        assert_eq!(resolved.payload["state"], "ready");
        assert!(resolved.heal_confidence.is_none());
    }

    #[test]
    fn streamed_payload_parser_handles_unbalanced_reasoning_before_json() {
        let raw = "reasoning text with unmatched { braces and thoughts\n{\"state\":\"ready\",\"reason\":\"final\"}";

        let resolved = parse_streamed_structured_payload(raw, false)
            .expect("parser should recover final structured JSON object");
        assert_eq!(resolved.payload["state"], "ready");
    }

    #[test]
    fn streamed_payload_parser_handles_markdown_with_heal() {
        let raw = r#"Some preface
```json
{
  "state": "missing_scenario",
  "reason": "Need more details"
}
```
Some trailing explanation"#;

        let resolved = parse_streamed_structured_payload(raw, true)
            .expect("heal path should parse JSON block");
        assert_eq!(resolved.payload["state"], "missing_scenario");
        assert!(resolved.heal_confidence.is_some());
    }

    #[test]
    fn streamed_payload_parser_errors_when_no_json_candidate_exists() {
        let raw = "No JSON in this streamed output";
        let error = parse_streamed_structured_payload(raw, false)
            .expect_err("strict stream parse should fail without JSON candidate");
        assert!(error.contains("no JSON object candidate found"));
    }

    #[test]
    fn include_raw_stream_debug_events_defaults_to_false() {
        let _guard = stream_debug_env_lock()
            .lock()
            .expect("stream debug env lock should not be poisoned");
        std::env::remove_var("SIMPLE_AGENTS_WORKFLOW_STREAM_INCLUDE_RAW");
        assert!(!include_raw_stream_debug_events());
    }

    #[test]
    fn include_raw_stream_debug_events_accepts_truthy_values() {
        let _guard = stream_debug_env_lock()
            .lock()
            .expect("stream debug env lock should not be poisoned");
        std::env::set_var("SIMPLE_AGENTS_WORKFLOW_STREAM_INCLUDE_RAW", "true");
        assert!(include_raw_stream_debug_events());
        std::env::remove_var("SIMPLE_AGENTS_WORKFLOW_STREAM_INCLUDE_RAW");
    }

    #[test]
    fn structured_json_delta_filter_strips_reasoning_prefix_and_suffix() {
        let mut filter = StructuredJsonDeltaFilter::default();
        let chunks = vec![
            "I will think first... ",
            "{\"state\":\"missing_scenario\",",
            "\"reason\":\"Need more details\"}",
            " additional commentary",
        ];

        let filtered = chunks
            .into_iter()
            .filter_map(|chunk| filter.split(chunk).0)
            .collect::<String>();

        assert_eq!(
            filtered,
            "{\"state\":\"missing_scenario\",\"reason\":\"Need more details\"}"
        );
    }

    #[test]
    fn structured_json_delta_filter_handles_braces_inside_strings() {
        let mut filter = StructuredJsonDeltaFilter::default();
        let chunks = vec![
            "preface ",
            "{\"reason\":\"brace } in text\",\"state\":\"ok\"}",
            " trailing",
        ];

        let filtered = chunks
            .into_iter()
            .filter_map(|chunk| filter.split(chunk).0)
            .collect::<String>();

        assert_eq!(
            filtered,
            "{\"reason\":\"brace } in text\",\"state\":\"ok\"}"
        );
    }

    #[test]
    fn render_json_object_as_text_converts_top_level_fields() {
        let rendered =
            render_json_object_as_text(r#"{"question":"q","confidence":0.8,"nested":{"a":1}}"#);
        let lines: std::collections::HashSet<&str> = rendered.lines().collect();

        assert_eq!(lines.len(), 3);
        assert!(lines.contains("question: q"));
        assert!(lines.contains("confidence: 0.8"));
        assert!(lines.contains("nested: {\"a\":1}"));
    }

    #[test]
    fn stream_json_as_text_formatter_emits_once_when_complete() {
        let mut formatter = StreamJsonAsTextFormatter::default();
        formatter.push("{\"question\":\"hello\"}");

        let first = formatter.emit_if_ready(true);
        let second = formatter.emit_if_ready(true);

        assert_eq!(first, Some("question: hello".to_string()));
        assert_eq!(second, None);
    }

    #[test]
    fn rewrite_yaml_condition_preserves_output_prefix_in_field_names() {
        let expr = "$.nodes.classify.output.output_total == 1";
        let rewritten = rewrite_yaml_condition_to_ir(expr);
        assert_eq!(rewritten, "$.node_outputs.classify.output_total == 1");
    }

    #[tokio::test]
    async fn validates_workflow_input_before_ir_runtime_path() {
        let yaml = r#"
id: chat-workflow
entry_node: classify
nodes:
  - id: classify
    node_type:
      llm_call:
        model: gpt-4.1
    config:
      prompt: |
        classify
"#;

        let workflow: YamlWorkflow = serde_yaml::from_str(yaml).expect("yaml should parse");
        let err = run_workflow_yaml(&workflow, &json!("not-an-object"), &MockExecutor)
            .await
            .expect_err("non-object input should fail before execution");

        assert!(matches!(err, YamlWorkflowRunError::InvalidInput { .. }));
    }

    #[test]
    fn interpolate_template_supports_dollar_prefixed_paths() {
        let context = json!({
            "input": {
                "email_text": "hello"
            }
        });

        let rendered = interpolate_template("value={{ $.input.email_text }}", &context);
        assert_eq!(rendered, "value=hello");
    }
}