klieo-core 2.2.0

//! Shared tool-dispatch helper used by both the blocking and streaming
//! runtime drivers.
//!
//! Both [`super::run_steps`] (blocking) and the streaming driver's loop
//! (`drive_streaming_loop` inside `super`) previously carried
//! byte-identical tool-dispatch arms. Extracting the shared loop here
//! removes the dedup risk: a bugfix lands in one place.

use crate::agent::{AgentContext, AgentEvent};
use crate::error::Error;
use crate::ids::ThreadId;
use crate::llm::{Message, Role, ToolCall};
use crate::memory::{Episode, ToolResult};
use crate::redact::{opaque_digest, AuditRedactor};
use tracing::debug;

use super::emit;

/// Marker recorded for a redacted error message when no
/// [`AuditRedactor`] is wired — the raw message may carry PII, so it is
/// dropped entirely rather than digested.
const REDACTED_ERROR_MARKER: &str = "[redacted error]";

/// Invoke every tool in `tool_calls` against `ctx.tools`, recording the
/// episodic [`Episode::ToolCall`] and appending a `Role::Tool` message
/// to short-term memory per call.
///
/// On a permanent (non-retryable) tool failure, returns `Err(Error::Tool)`
/// — the runtime aborts the run. Retryable failures are swallowed at
/// this layer (the error message has already been written to short-term
/// memory so the LLM observes it on the next cycle); upstream
/// backoff/retry policy belongs to wrappers (e.g. tool-level retry
/// adapters in `klieo-tools-mcp`).
///
/// `driver` is included in the trace span so log readers can tell the
/// blocking and streaming paths apart in a single combined log.
pub(crate) async fn dispatch_tool_calls(
    ctx: &AgentContext,
    thread: &ThreadId,
    tool_calls: &[ToolCall],
    driver: &'static str,
) -> Result<(), Error> {
    for call in tool_calls {
        let tool_ctx =
            crate::tool::ToolCtx::new(ctx.pubsub.clone(), ctx.kv.clone(), ctx.jobs.clone())
                .with_cancel(ctx.cancel.clone());
        emit(
            ctx,
            AgentEvent::ToolCallStarted {
                name: call.name.clone(),
            },
        );
        let outcome = ctx
            .tools
            .invoke(&call.name, call.args.clone(), tool_ctx)
            .await;
        emit(
            ctx,
            AgentEvent::ToolCallCompleted {
                name: call.name.clone(),
                ok: outcome.is_ok(),
            },
        );
        let raw_result = match &outcome {
            Ok(v) => ToolResult::Ok { value: v.clone() },
            Err(e) => ToolResult::Err {
                message: e.to_string(),
            },
        };
        // Short-term memory feeds the next LLM cycle and must stay raw so
        // the agent observes the real result; only the durable episodic
        // record is redacted when the tool handles PII.
        let tool_msg_content = match &outcome {
            Ok(v) => v.to_string(),
            Err(e) => format!("error: {e}"),
        };
        let (episode_args, episode_result) = if ctx.tools.tool_redacts_audit(&call.name) {
            redacted_args_and_result(ctx.audit_redactor.as_deref(), &call.args, raw_result)
        } else {
            (call.args.clone(), raw_result)
        };
        ctx.episodic
            .record(
                ctx.run_id,
                Episode::ToolCall {
                    name: call.name.clone(),
                    args: episode_args,
                    result: episode_result,
                },
            )
            .await?;
        ctx.short_term
            .append(
                thread.clone(),
                Message {
                    role: Role::Tool,
                    content: tool_msg_content,
                    tool_calls: vec![],
                    tool_call_id: Some(call.id.clone()),
                },
            )
            .await?;
        if let Err(e) = outcome {
            if !e.retryable() {
                return Err(Error::Tool(e));
            }
        }
    }
    debug!(n = tool_calls.len(), driver, "dispatched tools");
    Ok(())
}

fn redacted_args_and_result(
    redactor: Option<&dyn AuditRedactor>,
    args: &serde_json::Value,
    result: ToolResult,
) -> (serde_json::Value, ToolResult) {
    if let Some(redactor) = redactor {
        let redacted_result = match result {
            ToolResult::Ok { value } => ToolResult::Ok {
                value: redactor.redact(&value),
            },
            ToolResult::Err { message } => ToolResult::Err {
                message: redactor
                    .redact(&serde_json::Value::String(message))
                    .to_string(),
            },
        };
        return (redactor.redact(args), redacted_result);
    }
    let opaque_result = match result {
        ToolResult::Ok { value } => ToolResult::Ok {
            value: opaque_digest(&value),
        },
        ToolResult::Err { .. } => ToolResult::Err {
            message: REDACTED_ERROR_MARKER.to_string(),
        },
    };
    (opaque_digest(args), opaque_result)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agent::AgentContext;
    use crate::error::ToolError;
    use crate::ids::{RunId, ThreadId};
    use crate::memory::EpisodicMemory;
    use crate::test_utils::{fake_context, FakeToolInvoker, InMemoryEpisodic};
    use std::sync::Arc;

    fn make_ctx(tools: Arc<FakeToolInvoker>) -> (AgentContext, Arc<InMemoryEpisodic>) {
        let episodic = Arc::new(InMemoryEpisodic::default());
        let mut ctx = fake_context("dispatch-test");
        ctx.episodic = episodic.clone();
        ctx.tools = tools;
        (ctx, episodic)
    }

    struct MaskingRedactor;

    const MASK_TOKEN: &str = "[masked]";

    impl crate::redact::AuditRedactor for MaskingRedactor {
        fn redact(&self, value: &serde_json::Value) -> serde_json::Value {
            match value {
                serde_json::Value::Object(map) => serde_json::Value::Object(
                    map.iter()
                        .map(|(k, v)| (k.clone(), self.redact(v)))
                        .collect(),
                ),
                serde_json::Value::Array(items) => {
                    serde_json::Value::Array(items.iter().map(|v| self.redact(v)).collect())
                }
                serde_json::Value::String(_) => serde_json::json!(MASK_TOKEN),
                other => other.clone(),
            }
        }
    }

    async fn recorded_episode_string(episodic: &InMemoryEpisodic, run: RunId) -> String {
        let episodes = episodic.replay(run).await.unwrap();
        serde_json::to_string(&episodes).unwrap()
    }

    fn first_tool_result(episodes: &[Episode]) -> &ToolResult {
        episodes
            .iter()
            .find_map(|e| match e {
                Episode::ToolCall { result, .. } => Some(result),
                _ => None,
            })
            .expect("a ToolCall episode was recorded")
    }

    const PII_IBAN: &str = "DE89370400440532013000";
    const PII_EMAIL: &str = "alice@example.com";

    fn call_with_args(id: &str, name: &str, args: serde_json::Value) -> ToolCall {
        ToolCall {
            id: id.into(),
            name: name.into(),
            args,
        }
    }

    fn call(id: &str, name: &str) -> ToolCall {
        ToolCall {
            id: id.into(),
            name: name.into(),
            args: serde_json::json!({}),
        }
    }

    #[tokio::test]
    async fn records_one_episode_and_one_short_term_message_per_successful_call() {
        let tools = Arc::new(
            FakeToolInvoker::new()
                .with_tool("ok_a", "", |_| Ok(serde_json::json!("a")))
                .with_tool("ok_b", "", |_| Ok(serde_json::json!("b")))
                .with_tool("ok_c", "", |_| Ok(serde_json::json!("c"))),
        );
        let (ctx, episodic) = make_ctx(tools);
        let thread = ThreadId::new("t-1");
        let calls = vec![
            call("c-1", "ok_a"),
            call("c-2", "ok_b"),
            call("c-3", "ok_c"),
        ];

        dispatch_tool_calls(&ctx, &thread, &calls, "blocking")
            .await
            .expect("all-Ok dispatch must succeed");

        let episodes = episodic.replay(ctx.run_id).await.unwrap();
        let tool_call_count = episodes
            .iter()
            .filter(|e| matches!(e, Episode::ToolCall { .. }))
            .count();
        assert_eq!(tool_call_count, 3);
        let history = ctx.short_term.load(thread, 1024).await.unwrap();
        let tool_msgs: Vec<_> = history.iter().filter(|m| m.role == Role::Tool).collect();
        assert_eq!(tool_msgs.len(), 3);
        let ids: Vec<&str> = tool_msgs
            .iter()
            .filter_map(|m| m.tool_call_id.as_deref())
            .collect();
        assert_eq!(ids, vec!["c-1", "c-2", "c-3"]);
    }

    #[tokio::test]
    async fn non_retryable_tool_error_aborts_after_persisting_that_call() {
        let tools = Arc::new(
            FakeToolInvoker::new()
                .with_tool("ok", "", |_| Ok(serde_json::json!("ok-value")))
                .with_tool("boom", "", |_| Err(ToolError::Permanent("nope".into())))
                .with_tool("never_runs", "", |_| Ok(serde_json::json!("x"))),
        );
        let (ctx, episodic) = make_ctx(tools);
        let thread = ThreadId::new("t-abort");
        let calls = vec![
            call("c-1", "ok"),
            call("c-2", "boom"),
            call("c-3", "never_runs"),
        ];

        let err = dispatch_tool_calls(&ctx, &thread, &calls, "blocking")
            .await
            .expect_err("non-retryable must abort");
        assert!(
            matches!(&err, Error::Tool(ToolError::Permanent(m)) if m == "nope"),
            "expected Tool(Permanent), got {err:?}"
        );
        let episodes = episodic.replay(ctx.run_id).await.unwrap();
        let tool_episode_count = episodes
            .iter()
            .filter(|e| matches!(e, Episode::ToolCall { .. }))
            .count();
        assert_eq!(
            tool_episode_count, 2,
            "only ok + boom should have been recorded; never_runs aborts"
        );
        let history = ctx.short_term.load(thread, 1024).await.unwrap();
        let tool_msgs: Vec<_> = history.iter().filter(|m| m.role == Role::Tool).collect();
        assert_eq!(tool_msgs.len(), 2);
        assert!(tool_msgs[1].content.contains("error: permanent: nope"));
    }

    #[tokio::test]
    async fn retryable_tool_error_is_swallowed_and_loop_continues() {
        let tools = Arc::new(
            FakeToolInvoker::new()
                .with_tool("flaky", "", |_| {
                    Err(ToolError::Retryable {
                        message: "transient".into(),
                        retry_after_secs: 1,
                    })
                })
                .with_tool("after", "", |_| Ok(serde_json::json!("after-value"))),
        );
        let (ctx, episodic) = make_ctx(tools);
        let thread = ThreadId::new("t-retry");
        let calls = vec![call("c-1", "flaky"), call("c-2", "after")];

        dispatch_tool_calls(&ctx, &thread, &calls, "streaming")
            .await
            .expect("retryable error must not abort the loop");

        let episodes = episodic.replay(ctx.run_id).await.unwrap();
        let tool_episode_count = episodes
            .iter()
            .filter(|e| matches!(e, Episode::ToolCall { .. }))
            .count();
        assert_eq!(tool_episode_count, 2);
        let history = ctx.short_term.load(thread, 1024).await.unwrap();
        let tool_msgs: Vec<_> = history.iter().filter(|m| m.role == Role::Tool).collect();
        assert_eq!(tool_msgs.len(), 2);
        assert!(tool_msgs[0].content.contains("error: retryable"));
        assert_eq!(tool_msgs[1].content, "\"after-value\"");
    }

    #[tokio::test]
    async fn flagged_tool_with_redactor_records_masked_shape_never_raw_pii() {
        let pii = serde_json::json!({"iban": PII_IBAN, "email": PII_EMAIL});
        let tools = Arc::new(
            FakeToolInvoker::new().with_redacting_tool("claim", "", move |_| Ok(pii.clone())),
        );
        let (mut ctx, episodic) = make_ctx(tools);
        ctx.audit_redactor = Some(Arc::new(MaskingRedactor));
        let calls = vec![call_with_args(
            "c-1",
            "claim",
            serde_json::json!({"applicant": PII_EMAIL}),
        )];

        dispatch_tool_calls(&ctx, &ThreadId::new("t"), &calls, "blocking")
            .await
            .unwrap();

        let recorded = recorded_episode_string(&episodic, ctx.run_id).await;
        assert!(
            !recorded.contains(PII_IBAN) && !recorded.contains(PII_EMAIL),
            "no raw PII may survive in the episodic record; got {recorded}"
        );
        assert!(
            recorded.contains(MASK_TOKEN),
            "the masked shape must be present, proving structure is kept"
        );
    }

    #[tokio::test]
    async fn flagged_tool_without_redactor_records_opaque_digest_fail_closed() {
        let pii = serde_json::json!({"iban": PII_IBAN});
        let tools = Arc::new(
            FakeToolInvoker::new().with_redacting_tool("claim", "", move |_| Ok(pii.clone())),
        );
        let (ctx, episodic) = make_ctx(tools);
        assert!(ctx.audit_redactor.is_none(), "no redactor wired");
        let calls = vec![call_with_args(
            "c-1",
            "claim",
            serde_json::json!({"iban": PII_IBAN}),
        )];

        dispatch_tool_calls(&ctx, &ThreadId::new("t"), &calls, "blocking")
            .await
            .unwrap();

        let recorded = recorded_episode_string(&episodic, ctx.run_id).await;
        assert!(
            !recorded.contains(PII_IBAN),
            "fail-closed must not leak the secret; got {recorded}"
        );
        assert!(
            !recorded.contains("[REDACTED:"),
            "fail-closed must not emit a partial-redaction token"
        );
        assert!(
            recorded.contains("sha256"),
            "fail-closed must record the opaque-digest marker"
        );
    }

    #[tokio::test]
    async fn non_flagged_tool_records_args_and_result_raw() {
        let marker = "RAW-MARKER-1234";
        let value = serde_json::json!({"echo": marker});
        let tools =
            Arc::new(FakeToolInvoker::new().with_tool("plain", "", move |_| Ok(value.clone())));
        let (ctx, episodic) = make_ctx(tools);
        let calls = vec![call_with_args(
            "c-1",
            "plain",
            serde_json::json!({"in": marker}),
        )];

        dispatch_tool_calls(&ctx, &ThreadId::new("t"), &calls, "blocking")
            .await
            .unwrap();

        let recorded = recorded_episode_string(&episodic, ctx.run_id).await;
        assert!(
            recorded.contains(marker),
            "a non-flagged tool must persist the value verbatim; got {recorded}"
        );
    }

    #[tokio::test]
    async fn flagged_tool_error_arm_redacts_message() {
        let tools = Arc::new(
            FakeToolInvoker::new().with_redacting_tool("claim", "", |_| {
                Err(ToolError::Retryable {
                    message: format!("lookup failed for {PII_EMAIL}"),
                    retry_after_secs: 1,
                })
            }),
        );
        let (ctx, episodic) = make_ctx(tools);
        let calls = vec![call_with_args("c-1", "claim", serde_json::json!({}))];

        dispatch_tool_calls(&ctx, &ThreadId::new("t"), &calls, "blocking")
            .await
            .unwrap();

        let episodes = episodic.replay(ctx.run_id).await.unwrap();
        match first_tool_result(&episodes) {
            ToolResult::Err { message } => assert!(
                !message.contains(PII_EMAIL),
                "the error message must not carry PII; got {message}"
            ),
            ToolResult::Ok { .. } => panic!("expected an Err result"),
        }
    }

    #[test]
    fn redacted_helper_without_redactor_digests_args_and_marks_error() {
        let (args, result) = redacted_args_and_result(
            None,
            &serde_json::json!({"iban": PII_IBAN}),
            ToolResult::Err {
                message: format!("failed {PII_EMAIL}"),
            },
        );
        assert_eq!(args["redacted"], serde_json::json!("sha256"));
        match result {
            ToolResult::Err { message } => assert_eq!(message, REDACTED_ERROR_MARKER),
            ToolResult::Ok { .. } => panic!("expected Err"),
        }
    }

    #[test]
    fn redacted_helper_with_redactor_masks_ok_value() {
        let redactor = MaskingRedactor;
        let (args, result) = redacted_args_and_result(
            Some(&redactor),
            &serde_json::json!({"k": PII_IBAN}),
            ToolResult::Ok {
                value: serde_json::json!({"out": PII_EMAIL}),
            },
        );
        assert_eq!(args["k"], serde_json::json!(MASK_TOKEN));
        match result {
            ToolResult::Ok { value } => assert_eq!(value["out"], serde_json::json!(MASK_TOKEN)),
            ToolResult::Err { .. } => panic!("expected Ok"),
        }
    }
}