klieo-ops 3.5.0

//! [`ToolInvoker`] wrapper that runs the registered [`Gate`] stack before
//! every tool invocation and emits [`OpsEvent::GateDecision`] on every
//! decision.

use crate::emit::emit_ops_event;
use crate::escalation::Severity;
use crate::gates::{
    ApprovalError, ApprovalOutcome, ApprovalTimeoutPolicy, Gate, GateDecision, GateRequest,
};
use crate::ops_event::OpsEvent;
use crate::tenant::current_tenant;
use async_trait::async_trait;
use klieo_core::error::ToolError;
use klieo_core::ids::RunId;
use klieo_core::llm::ToolDef;
use klieo_core::memory::EpisodicMemory;
use klieo_core::tool::{ToolCtx, ToolInvoker};
use std::sync::Arc;
use std::time::Duration;

/// Default wait for human approval before the gate times out the tool call.
/// 10 minutes per spec § 2.3.
const DEFAULT_APPROVAL_TIMEOUT: Duration = Duration::from_secs(600);

/// Wraps an inner [`ToolInvoker`]. Consults every registered [`Gate`] before
/// delegating; first non-`Allow` outcome short-circuits with
/// [`ToolError::Permanent`]. Every decision is emitted into the provided
/// [`EpisodicMemory`] as [`OpsEvent::GateDecision`].
///
/// On approval timeout, the configured [`ApprovalTimeoutPolicy`] determines
/// behaviour. The default is `RequeueAndEscalate { max_requeues: 3 }` per
/// spec § 2.3.
pub struct GatedToolInvoker {
    inner: Arc<dyn ToolInvoker>,
    gates: Vec<Arc<dyn Gate>>,
    episodic: Arc<dyn EpisodicMemory>,
    run_id: RunId,
    approval_timeout: Duration,
    approval_timeout_policy: ApprovalTimeoutPolicy,
}

impl GatedToolInvoker {
    /// Construct a gated wrapper around `inner` with the default approval
    /// timeout (10 minutes) and default policy (`RequeueAndEscalate { max_requeues: 3 }`).
    #[must_use]
    pub fn new(
        inner: Arc<dyn ToolInvoker>,
        gates: Vec<Arc<dyn Gate>>,
        episodic: Arc<dyn EpisodicMemory>,
        run_id: RunId,
    ) -> Self {
        Self {
            inner,
            gates,
            episodic,
            run_id,
            approval_timeout: DEFAULT_APPROVAL_TIMEOUT,
            approval_timeout_policy: ApprovalTimeoutPolicy::default(),
        }
    }

    /// Override the approval timeout (used by tests and by the builder
    /// `approval_timeout` setter).
    #[must_use]
    pub fn with_approval_timeout(mut self, timeout: Duration) -> Self {
        self.approval_timeout = timeout;
        self
    }

    /// Override the approval timeout policy.
    #[must_use]
    pub fn with_approval_timeout_policy(mut self, policy: ApprovalTimeoutPolicy) -> Self {
        self.approval_timeout_policy = policy;
        self
    }

    async fn audit(&self, tool: &str, gate: &str, decision: &str, reason: Option<String>) {
        if let Err(err) = emit_ops_event(
            &*self.episodic,
            self.run_id,
            OpsEvent::GateDecision {
                tenant: current_tenant(),
                tool: tool.into(),
                decision: decision.into(),
                gate: gate.into(),
                policy_ref: None,
                reason,
            },
        )
        .await
        {
            tracing::warn!(
                target: "klieo.ops.audit",
                error = %err,
                "audit emit failed; episode not recorded"
            );
        }
    }

    async fn emit_escalation_state_changed(
        &self,
        ticket: &str,
        from: &str,
        to: &str,
        reason: &str,
    ) {
        if let Err(err) = emit_ops_event(
            &*self.episodic,
            self.run_id,
            OpsEvent::EscalationStateChanged {
                tenant: current_tenant(),
                ticket: ticket.into(),
                from: from.into(),
                to: to.into(),
                reason: Some(reason.into()),
            },
        )
        .await
        {
            tracing::warn!(
                target: "klieo.ops.audit",
                error = %err,
                "audit emit failed (EscalationStateChanged)"
            );
        }
    }

    async fn emit_escalation_resolved(&self, ticket: &str, outcome: &str, reason: &str) {
        if let Err(err) = emit_ops_event(
            &*self.episodic,
            self.run_id,
            OpsEvent::EscalationResolved {
                tenant: current_tenant(),
                ticket: ticket.into(),
                outcome: outcome.into(),
                reason: Some(reason.into()),
            },
        )
        .await
        {
            tracing::warn!(
                target: "klieo.ops.audit",
                error = %err,
                "audit emit failed (EscalationResolved)"
            );
        }
    }

    /// Drive the requeue loop for a single gate's `RequireApproval`.
    ///
    /// Returns `Ok(())` when approval is eventually granted, or
    /// `Err(ToolError::Permanent(...))` on exhaustion / explicit denial /
    /// other terminal error.
    async fn wait_with_requeue(
        &self,
        name: &str,
        gate: &Arc<dyn Gate>,
        ticket: String,
    ) -> Result<(), ToolError> {
        let max_requeues = match self.approval_timeout_policy {
            ApprovalTimeoutPolicy::Deny => {
                return self.wait_once(name, gate, &ticket).await.map(|_| ());
            }
            ApprovalTimeoutPolicy::RequeueAndEscalate { max_requeues } => max_requeues,
        };

        let mut attempts: u8 = 0;
        // Track severity so each requeue bumps it one level.
        let mut severity = Severity::High;

        loop {
            match gate
                .wait_for_approval(ticket.clone(), self.approval_timeout)
                .await
            {
                Ok(ApprovalOutcome::Allow) => {
                    self.audit(
                        name,
                        gate.name(),
                        "approved",
                        Some(format!("ticket={ticket}")),
                    )
                    .await;
                    return Ok(());
                }
                Ok(ApprovalOutcome::Deny { code, reason }) => {
                    self.audit(name, gate.name(), "approval_denied", Some(reason.clone()))
                        .await;
                    return Err(ToolError::Permanent(format!(
                        "gate `{}` approval denied: code={code} reason={reason}",
                        gate.name()
                    )));
                }
                Err(ApprovalError::TimedOut { millis }) => {
                    attempts += 1;
                    if attempts > max_requeues {
                        self.emit_escalation_resolved(
                            &ticket,
                            "timed_out",
                            "max requeues exhausted",
                        )
                        .await;
                        return Err(ToolError::Permanent(format!(
                            "gate `{}` approval timed out after {millis}ms; \
                             max requeues ({max_requeues}) exhausted",
                            gate.name()
                        )));
                    }
                    let from_label = severity_label(severity);
                    severity = severity.escalate_one_level();
                    let to_label = severity_label(severity);
                    self.emit_escalation_state_changed(
                        &ticket,
                        from_label,
                        to_label,
                        &format!("timeout requeue {attempts}/{max_requeues}"),
                    )
                    .await;
                    self.audit(
                        name,
                        gate.name(),
                        "approval_timed_out_requeue",
                        Some(format!(
                            "attempt {attempts}/{max_requeues} after {millis}ms"
                        )),
                    )
                    .await;
                }
                Err(err) => return Err(self.terminal_approval_error(gate.name(), err)),
            }
        }
    }

    /// Wait once (no retry). Used for the `Deny` policy and for non-timeout
    /// error paths.
    async fn wait_once(
        &self,
        name: &str,
        gate: &Arc<dyn Gate>,
        ticket: &str,
    ) -> Result<ApprovalOutcome, ToolError> {
        match gate
            .wait_for_approval(ticket.to_string(), self.approval_timeout)
            .await
        {
            Ok(outcome @ ApprovalOutcome::Allow) => {
                self.audit(
                    name,
                    gate.name(),
                    "approved",
                    Some(format!("ticket={ticket}")),
                )
                .await;
                Ok(outcome)
            }
            Ok(ApprovalOutcome::Deny { code, reason }) => {
                self.audit(name, gate.name(), "approval_denied", Some(reason.clone()))
                    .await;
                Err(ToolError::Permanent(format!(
                    "gate `{}` approval denied: code={code} reason={reason}",
                    gate.name()
                )))
            }
            Err(ApprovalError::TimedOut { millis }) => {
                self.audit(
                    name,
                    gate.name(),
                    "approval_timed_out",
                    Some(format!("after {millis}ms")),
                )
                .await;
                Err(ToolError::Permanent(format!(
                    "gate `{}` approval timed out after {millis}ms",
                    gate.name()
                )))
            }
            Err(err) => Err(self.terminal_approval_error(gate.name(), err)),
        }
    }

    fn terminal_approval_error(&self, gate_name: &str, err: ApprovalError) -> ToolError {
        match err {
            ApprovalError::Denied(reason) => {
                ToolError::Permanent(format!("gate `{gate_name}` approval denied: {reason}"))
            }
            ApprovalError::Halted => {
                ToolError::Permanent(format!("gate `{gate_name}` approval halted (kill-switch)"))
            }
            ApprovalError::NotSupported => ToolError::Permanent(format!(
                "gate `{gate_name}` requires approval but does not support waiting; \
                 register a FourEyesGate or remove the dual_control classifier"
            )),
            ApprovalError::VerificationFailed(msg) => ToolError::Permanent(format!(
                "gate `{gate_name}` approver verification failed: {msg}"
            )),
            ApprovalError::TimedOut { millis } => ToolError::Permanent(format!(
                "gate `{gate_name}` approval timed out after {millis}ms"
            )),
        }
    }
}

#[async_trait]
impl ToolInvoker for GatedToolInvoker {
    async fn invoke(
        &self,
        name: &str,
        args: serde_json::Value,
        ctx: ToolCtx,
    ) -> Result<serde_json::Value, ToolError> {
        let req = GateRequest::new(name, args.clone());

        for gate in &self.gates {
            match gate.evaluate(req.clone()).await {
                GateDecision::Allow => {
                    self.audit(name, gate.name(), "allow", None).await;
                }
                GateDecision::Deny { code, reason } => {
                    self.audit(name, gate.name(), "deny", Some(reason.clone()))
                        .await;
                    return Err(ToolError::Permanent(format!(
                        "gate `{}` denied: code={code} reason={reason}",
                        gate.name()
                    )));
                }
                GateDecision::RequireApproval { ticket, quorum } => {
                    self.audit(
                        name,
                        gate.name(),
                        "require_approval",
                        Some(format!("ticket={ticket} quorum={quorum}")),
                    )
                    .await;
                    self.wait_with_requeue(name, gate, ticket).await?;
                    // Approval granted — continue evaluating remaining gates.
                }
            }
        }

        self.inner.invoke(name, args, ctx).await
    }

    fn catalogue(&self) -> Vec<ToolDef> {
        self.inner.catalogue()
    }

    fn is_tool_idempotent(&self, name: &str) -> bool {
        self.inner.is_tool_idempotent(name)
    }

    fn tool_redacts_audit(&self, name: &str) -> bool {
        self.inner.tool_redacts_audit(name)
    }
}

fn severity_label(s: Severity) -> &'static str {
    match s {
        Severity::Low => "low",
        Severity::Medium => "medium",
        Severity::High => "high",
        Severity::Critical => "critical",
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use klieo_core::test_utils::{FakeToolInvoker, InMemoryEpisodic};

    fn gated_over(inner: FakeToolInvoker) -> GatedToolInvoker {
        GatedToolInvoker::new(
            Arc::new(inner),
            Vec::new(),
            Arc::new(InMemoryEpisodic::default()),
            RunId::new(),
        )
    }

    #[test]
    fn forwards_tool_redacts_audit_for_flagged_inner_tool() {
        let inner =
            FakeToolInvoker::new().with_redacting_tool("claimant_lookup", "handles PII", Ok);
        let gated = gated_over(inner);
        assert!(
            gated.tool_redacts_audit("claimant_lookup"),
            "a PII-flagged tool behind the gate must not fail open to raw audit"
        );
    }

    #[test]
    fn does_not_report_redacts_audit_for_unflagged_inner_tool() {
        let inner = FakeToolInvoker::new().with_tool("echo", "plain", Ok);
        let gated = gated_over(inner);
        assert!(!gated.tool_redacts_audit("echo"));
    }

    #[test]
    fn forwards_is_tool_idempotent_for_flagged_inner_tool() {
        let inner = FakeToolInvoker::new().with_idempotent_tool("read_only", "safe to replay", Ok);
        let gated = gated_over(inner);
        assert!(gated.is_tool_idempotent("read_only"));
    }

    #[test]
    fn does_not_report_idempotent_for_unflagged_inner_tool() {
        let inner = FakeToolInvoker::new().with_tool("mutate", "side effect", Ok);
        let gated = gated_over(inner);
        assert!(!gated.is_tool_idempotent("mutate"));
    }
}