klieo-workflow 3.8.1

//! Node wrappers: turn each declarative node into an `Arc<dyn Flow>` that
//! reads/writes named fields on a shared JSON envelope.

use crate::error::CompileError;
use crate::registry::Registry;
use async_trait::async_trait;
use klieo_core::agent::AgentContext;
use klieo_core::llm::LlmClient;
use klieo_core::tool::ToolCtx;
use klieo_core::SimpleAgent;
use klieo_flows::{AgentFlow, Flow, FlowError};
use serde_json::Value;
use std::sync::Arc;

/// Merge `output` into `envelope` under `field`, returning the envelope.
fn write_field(mut envelope: Value, field: &str, output: Value) -> Value {
    if !envelope.is_object() {
        envelope = Value::Object(serde_json::Map::new());
    }
    envelope
        .as_object_mut()
        .expect("envelope forced to object above")
        .insert(field.to_string(), output);
    envelope
}

/// Read `field` from `envelope`, or `Value::Null` if absent.
fn read_field(envelope: &Value, field: &str) -> Value {
    envelope.get(field).cloned().unwrap_or(Value::Null)
}

/// Message surfaced when a compiled workflow is run on a non-object input.
pub(crate) const NON_OBJECT_ENVELOPE: &str = "workflow envelope must be a JSON object";

/// Enforces the envelope contract at the top-level flow boundary: the run
/// input must be a JSON object. Rejects a non-object input with a typed
/// error rather than silently coercing it (which would discard the caller's
/// upstream data when this flow is composed inside another).
pub(crate) struct EnvelopeGuard {
    pub inner: Arc<dyn Flow>,
}

#[async_trait]
impl Flow for EnvelopeGuard {
    fn name(&self) -> &str {
        self.inner.name()
    }
    async fn run(&self, ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
        if !input.is_object() {
            return Err(FlowError::Agent(NON_OBJECT_ENVELOPE.to_string()));
        }
        self.inner.run(ctx, input).await
    }
}

/// An agent node: swaps the context LLM to the registry-resolved model,
/// runs a `SimpleAgent` on the `input_from` field (as a string), writes
/// the string result to `output_to`.
pub(crate) struct AgentNode {
    pub name: String,
    pub llm: Arc<dyn LlmClient>,
    pub agent: Arc<SimpleAgent>,
    pub input_from: String,
    pub output_to: String,
}

#[async_trait]
impl Flow for AgentNode {
    fn name(&self) -> &str {
        &self.name
    }
    async fn run(&self, mut ctx: AgentContext, envelope: Value) -> Result<Value, FlowError> {
        ctx.llm = self.llm.clone();
        let prompt = match read_field(&envelope, &self.input_from) {
            Value::String(s) => s,
            other => other.to_string(),
        };
        let flow = AgentFlow::from_arc(self.agent.clone());
        let out = flow.run(ctx, Value::String(prompt)).await?;
        Ok(write_field(envelope, &self.output_to, out))
    }
}

/// A tool node: invokes a registered tool by name with the `input_from`
/// field as its args, writes the result to `output_to`.
pub(crate) struct ToolNode {
    pub name: String,
    pub tool: String,
    pub input_from: String,
    pub output_to: String,
}

#[async_trait]
impl Flow for ToolNode {
    fn name(&self) -> &str {
        &self.name
    }
    async fn run(&self, ctx: AgentContext, envelope: Value) -> Result<Value, FlowError> {
        let args = read_field(&envelope, &self.input_from);
        let tool_ctx = ToolCtx::new(ctx.pubsub.clone(), ctx.kv.clone(), ctx.jobs.clone());
        let out = ctx
            .tools
            .invoke(&self.tool, args, tool_ctx)
            .await
            .map_err(|e| {
                tracing::error!(tool = %self.tool, error = %e, "workflow tool node failed");
                FlowError::Tool(e)
            })?;
        Ok(write_field(envelope, &self.output_to, out))
    }
}

/// Build the agent node's `SimpleAgent` from config + resolve its model.
/// `node_id` is used only for error messages.
pub(crate) fn build_agent_node(
    node_id: &str,
    cfg: &crate::def::AgentConfig,
    input_from: &str,
    output_to: &str,
    registry: &Registry,
) -> Result<AgentNode, CompileError> {
    let llm = registry
        .model(&cfg.model)
        .ok_or_else(|| CompileError::UnknownRef {
            kind: "model",
            node: node_id.to_string(),
            id: cfg.model.clone(),
        })?;
    for tool_id in &cfg.tools {
        if !registry.allows_tool(tool_id) {
            return Err(CompileError::UnknownRef {
                kind: "tool",
                node: node_id.to_string(),
                id: tool_id.clone(),
            });
        }
    }
    // The agent runs with an empty ToolDef catalogue; tool dispatch is
    // resolved at runtime through ctx.tools. Surfacing the registry's
    // ToolDefs to the agent is not yet wired.
    let agent = SimpleAgent::new(node_id, &cfg.system_prompt, Vec::new());
    Ok(AgentNode {
        name: node_id.to_string(),
        llm,
        agent: Arc::new(agent),
        input_from: input_from.to_string(),
        output_to: output_to.to_string(),
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use klieo_core::error::ToolError;
    use klieo_core::test_utils::FakeToolInvoker;

    fn tool_node(tool: &str) -> ToolNode {
        ToolNode {
            name: "t".into(),
            tool: tool.into(),
            input_from: "in".into(),
            output_to: "out".into(),
        }
    }

    #[tokio::test]
    async fn tool_node_run_writes_tool_output_to_output_field() {
        let mut ctx = crate::test_support::test_ctx();
        ctx.tools = Arc::new(FakeToolInvoker::new().with_tool("echo", "echoes its args", Ok));
        let envelope = serde_json::json!({"in": {"x": 1}});

        let result = tool_node("echo").run(ctx, envelope).await.unwrap();

        assert_eq!(result["out"], serde_json::json!({"x": 1}));
    }

    #[tokio::test]
    async fn tool_node_run_maps_invoke_failure_to_flow_error_tool() {
        let mut ctx = crate::test_support::test_ctx();
        ctx.tools = Arc::new(
            FakeToolInvoker::new().with_tool("boom", "always fails", |_args| {
                Err(ToolError::Permanent("kaboom".into()))
            }),
        );
        let envelope = serde_json::json!({"in": {}});

        match tool_node("boom").run(ctx, envelope).await {
            Ok(_) => panic!("expected the tool failure to surface"),
            Err(FlowError::Tool(ToolError::Permanent(msg))) => assert_eq!(msg, "kaboom"),
            Err(other) => panic!("expected FlowError::Tool(Permanent), got {other:?}"),
        }
    }

    #[test]
    fn write_field_creates_object_from_null() {
        let out = write_field(Value::Null, "a", serde_json::json!(1));
        assert_eq!(out, serde_json::json!({"a": 1}));
    }

    #[test]
    fn write_field_preserves_existing_keys() {
        let env = serde_json::json!({"x": 1});
        let out = write_field(env, "y", serde_json::json!(2));
        assert_eq!(out, serde_json::json!({"x": 1, "y": 2}));
    }

    #[test]
    fn write_field_overwrites_existing_key() {
        let env = serde_json::json!({"a": 1});
        let out = write_field(env, "a", serde_json::json!(2));
        assert_eq!(out, serde_json::json!({"a": 2}));
    }

    #[test]
    fn read_absent_field_is_null() {
        assert_eq!(read_field(&serde_json::json!({}), "nope"), Value::Null);
    }

    #[test]
    fn build_agent_node_rejects_unregistered_model() {
        let cfg = crate::def::AgentConfig {
            model: "ghost".into(),
            system_prompt: "hi".into(),
            tools: vec![],
        };
        // `AgentNode` holds `Arc<dyn LlmClient>` / `Arc<SimpleAgent>`, neither
        // `Debug`, so `Result::unwrap_err` (which requires the `Ok` side to
        // be `Debug`) isn't available here — match instead.
        match build_agent_node("n", &cfg, "in", "out", &Registry::new()) {
            Ok(_) => panic!("expected an UnknownRef model error"),
            Err(e) => assert!(
                matches!(e, CompileError::UnknownRef { kind: "model", .. }),
                "got: {e}"
            ),
        }
    }

    #[test]
    fn build_agent_node_rejects_unregistered_tool() {
        let cfg = crate::def::AgentConfig {
            model: "m".into(),
            system_prompt: "hi".into(),
            tools: vec!["danger".into()],
        };
        let registry = Registry::new().with_model("m", crate::test_support::dummy_llm());
        match build_agent_node("n", &cfg, "in", "out", &registry) {
            Ok(_) => panic!("expected an UnknownRef tool error"),
            Err(e) => assert!(
                matches!(e, CompileError::UnknownRef { kind: "tool", .. }),
                "got: {e}"
            ),
        }
    }
}