klieo-core 2.2.0

//! `Agent` trait and `AgentContext`.

use crate::bus::{JobQueue, KvStore, Pubsub, RequestReply};
use crate::ids::RunId;
use crate::llm::{LlmClient, ToolDef};
use crate::memory::{EpisodicMemory, LongTermMemory, ShortTermMemory};
use crate::redact::AuditRedactor;
use crate::tool::ToolInvoker;
use async_trait::async_trait;
use serde::de::DeserializeOwned;
use serde::Serialize;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;

/// Step-level event emitted by the runtime during [`Agent::run`].
/// Wire shape is transport-agnostic; MCP HTTP maps each variant
/// to a `notifications/progress` JSON-RPC notification.
///
/// Default [`AgentContext::progress`] is `None`, so emission is a
/// no-op for callers that don't opt in. Transports opt in by
/// passing a `broadcast::Sender` when constructing the context.
#[derive(Clone, Debug, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
#[non_exhaustive]
pub enum AgentEvent {
    /// LLM call about to be issued.
    LlmCallStarted,
    /// LLM call returned a response. `tokens` is `prompt + completion`
    /// when the provider streams a usage payload on the final chunk;
    /// `0` when the provider does not emit usage in stream mode.
    LlmCallCompleted {
        /// Total token count (`prompt_tokens + completion_tokens`).
        /// Zero when the provider does not surface usage in stream mode.
        tokens: u32,
        /// Wall-clock duration in milliseconds.
        latency_ms: u64,
    },
    /// Tool dispatch begun for `name`.
    ToolCallStarted {
        /// Name of the tool being dispatched.
        name: String,
    },
    /// Tool dispatch returned. `ok = false` means the tool
    /// errored; the wire-level redaction policy decides what to
    /// surface.
    ToolCallCompleted {
        /// Name of the tool that was dispatched.
        name: String,
        /// `true` if the tool call succeeded.
        ok: bool,
    },
    /// Agent reached its final assistant response.
    Completed,
    /// Agent failed terminally. `reason` is a sanitised summary;
    /// inner error chain is logged server-side by the runtime's
    /// existing `tracing::error!` path on the `Err` return.
    Failed {
        /// Sanitised failure reason suitable for wire transmission.
        reason: String,
    },
    /// Run suspended at a step awaiting human approval (ADR-045).
    Suspended {
        /// Reason supplied by the `ReviewPolicy`.
        reason: String,
    },
}

/// Borrow-free agent execution context. Holds `Arc<dyn …>` so it can be
/// cloned freely across `tokio::spawn` boundaries (`'static` requirement).
#[derive(Clone)]
#[non_exhaustive]
pub struct AgentContext {
    /// LLM provider.
    pub llm: Arc<dyn LlmClient>,
    /// Short-term conversation memory.
    pub short_term: Arc<dyn ShortTermMemory>,
    /// Long-term semantic memory.
    pub long_term: Arc<dyn LongTermMemory>,
    /// Episodic event log.
    pub episodic: Arc<dyn EpisodicMemory>,
    /// Pub/sub bus.
    pub pubsub: Arc<dyn Pubsub>,
    /// KV store.
    pub kv: Arc<dyn KvStore>,
    /// Synchronous request/reply.
    pub request_reply: Arc<dyn RequestReply>,
    /// Job queue.
    pub jobs: Arc<dyn JobQueue>,
    /// Tool dispatcher.
    pub tools: Arc<dyn ToolInvoker>,
    /// Stable id for this run.
    pub run_id: RunId,
    /// Cooperative cancellation token. Runtime checks between steps.
    pub cancel: CancellationToken,
    /// Agent name; recorded in episodic events. Caller must set this
    /// before invoking the runtime — typically from `Agent::name()`.
    pub agent_name: String,
    /// Optional fan-out channel for step-level events. When `Some`,
    /// the runtime emits one [`AgentEvent`] per LLM call, tool call,
    /// and terminal transition. Caller (e.g. MCP HTTP transport)
    /// owns the receiver and serialises events to the wire.
    ///
    /// Default `None` — existing single-shot callers see no
    /// behaviour change. Best-effort send; dropped receivers are
    /// silently ignored.
    pub progress: Option<tokio::sync::broadcast::Sender<AgentEvent>>,
    /// Redactor applied to audit-flagged tools' recorded args/results.
    /// Non-public so adding it stays SemVer-minor on the
    /// `#[non_exhaustive]`-less struct; set only via
    /// [`AgentContextBuilder::audit_redactor`]. `None` ⇒ dispatch uses
    /// the fail-closed [`crate::opaque_digest`] fallback for flagged
    /// tools.
    pub(crate) audit_redactor: Option<Arc<dyn AuditRedactor>>,
    /// Derived non-PII attribution label for the caller driving this
    /// run.
    ///
    /// `Some` ⇒ runtime records an [`crate::Episode::RunAttributed`]
    /// next to [`crate::Episode::Started`] so the audit trail can
    /// attribute the run to its driving tenant. The raw principal is
    /// never stored here — that stays in server-side tracing —
    /// and the label is never appended to short-term memory or any
    /// LLM-visible message.
    ///
    /// Non-public; install via [`Self::with_tenant_label`] or
    /// [`AgentContextBuilder::tenant_label`].
    pub(crate) tenant_label: Option<String>,
    /// Opaque cross-hop provenance anchor — the caller's own
    /// provenance chain-entry id, supplied by an authenticated external
    /// caller over MCP.
    ///
    /// `Some` ⇒ runtime records an [`crate::Episode::RunOrigin`] next to
    /// [`crate::Episode::Started`] so klieo→klieo composition can be
    /// stitched across deployments. Recorded verbatim and treated as a
    /// caller-asserted, unverified claim; never appended to short-term
    /// memory or any LLM-visible message. Inherited by child runs so the
    /// whole local run tree shares one external origin.
    ///
    /// Non-public; install via [`Self::with_parent_anchor`].
    pub(crate) parent_anchor: Option<String>,
}

/// Error returned when a required field is missing from [`AgentContextBuilder`].
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum AgentContextBuilderError {
    /// A required builder field was not set before calling `build()`.
    #[error("required field missing: {0}")]
    MissingField(&'static str),
}

/// Fluent builder for [`AgentContext`].
///
/// All fields except `run_id` and `cancel` are required. Omitting
/// `run_id` defaults to [`RunId::new()`]; omitting `cancel` defaults
/// to a fresh [`CancellationToken`].
///
/// ```
/// # use klieo_core::agent::{AgentContext, AgentContextBuilderError};
/// # use klieo_core::test_utils::fake_context;
/// // A full build is exercised in tests; this snippet shows the chaining shape.
/// let result: Result<AgentContext, AgentContextBuilderError> = AgentContext::builder()
///     .agent_name("my-agent")
///     .build();
/// assert!(result.is_err()); // other required fields are missing
/// ```
#[derive(Default)]
pub struct AgentContextBuilder {
    llm: Option<Arc<dyn LlmClient>>,
    short_term: Option<Arc<dyn ShortTermMemory>>,
    long_term: Option<Arc<dyn LongTermMemory>>,
    episodic: Option<Arc<dyn EpisodicMemory>>,
    pubsub: Option<Arc<dyn Pubsub>>,
    kv: Option<Arc<dyn KvStore>>,
    request_reply: Option<Arc<dyn RequestReply>>,
    jobs: Option<Arc<dyn JobQueue>>,
    tools: Option<Arc<dyn ToolInvoker>>,
    run_id: Option<RunId>,
    cancel: Option<CancellationToken>,
    agent_name: Option<String>,
    audit_redactor: Option<Arc<dyn AuditRedactor>>,
    tenant_label: Option<String>,
}

impl AgentContextBuilder {
    /// Set the LLM client.
    pub fn llm(mut self, v: Arc<dyn LlmClient>) -> Self {
        self.llm = Some(v);
        self
    }

    /// Set the short-term memory store.
    pub fn short_term(mut self, v: Arc<dyn ShortTermMemory>) -> Self {
        self.short_term = Some(v);
        self
    }

    /// Set the long-term memory store.
    pub fn long_term(mut self, v: Arc<dyn LongTermMemory>) -> Self {
        self.long_term = Some(v);
        self
    }

    /// Set the episodic memory store.
    pub fn episodic(mut self, v: Arc<dyn EpisodicMemory>) -> Self {
        self.episodic = Some(v);
        self
    }

    /// Set the pub/sub bus.
    pub fn pubsub(mut self, v: Arc<dyn Pubsub>) -> Self {
        self.pubsub = Some(v);
        self
    }

    /// Set the KV store.
    pub fn kv(mut self, v: Arc<dyn KvStore>) -> Self {
        self.kv = Some(v);
        self
    }

    /// Set the request/reply bus.
    pub fn request_reply(mut self, v: Arc<dyn RequestReply>) -> Self {
        self.request_reply = Some(v);
        self
    }

    /// Set the job queue.
    pub fn jobs(mut self, v: Arc<dyn JobQueue>) -> Self {
        self.jobs = Some(v);
        self
    }

    /// Set the tool invoker.
    pub fn tools(mut self, v: Arc<dyn ToolInvoker>) -> Self {
        self.tools = Some(v);
        self
    }

    /// Override the run ID. Defaults to [`RunId::new()`] when omitted.
    pub fn run_id(mut self, v: RunId) -> Self {
        self.run_id = Some(v);
        self
    }

    /// Override the cancellation token. Defaults to a fresh token when omitted.
    pub fn cancel(mut self, v: CancellationToken) -> Self {
        self.cancel = Some(v);
        self
    }

    /// Set the agent name (required).
    pub fn agent_name(mut self, v: impl Into<String>) -> Self {
        self.agent_name = Some(v.into());
        self
    }

    /// Install the audit redactor applied to audit-flagged tools'
    /// recorded args/results. When omitted, dispatch falls back to the
    /// fail-closed [`crate::opaque_digest`] for flagged tools.
    pub fn audit_redactor(mut self, redactor: Arc<dyn AuditRedactor>) -> Self {
        self.audit_redactor = Some(redactor);
        self
    }

    /// Install the derived non-PII tenant attribution label. Runtime
    /// records [`crate::Episode::RunAttributed`] at run entry when set.
    /// Caller must pre-derive a non-PII label (e.g. via
    /// [`crate::principal_hash`]) — the value is recorded into
    /// episodic memory verbatim.
    pub fn tenant_label(mut self, label: String) -> Self {
        self.tenant_label = Some(label);
        self
    }

    /// Consume the builder and produce an [`AgentContext`].
    ///
    /// Returns [`AgentContextBuilderError::MissingField`] when any
    /// required field was not set.
    pub fn build(self) -> Result<AgentContext, AgentContextBuilderError> {
        let audit_redactor = self.audit_redactor.clone();
        let tenant_label = self.tenant_label.clone();
        let ctx = AgentContext::new(
            self.llm
                .ok_or(AgentContextBuilderError::MissingField("llm"))?,
            self.short_term
                .ok_or(AgentContextBuilderError::MissingField("short_term"))?,
            self.long_term
                .ok_or(AgentContextBuilderError::MissingField("long_term"))?,
            self.episodic
                .ok_or(AgentContextBuilderError::MissingField("episodic"))?,
            self.pubsub
                .ok_or(AgentContextBuilderError::MissingField("pubsub"))?,
            self.kv
                .ok_or(AgentContextBuilderError::MissingField("kv"))?,
            self.request_reply
                .ok_or(AgentContextBuilderError::MissingField("request_reply"))?,
            self.jobs
                .ok_or(AgentContextBuilderError::MissingField("jobs"))?,
            self.tools
                .ok_or(AgentContextBuilderError::MissingField("tools"))?,
            self.run_id.unwrap_or_default(),
            self.cancel.unwrap_or_default(),
            self.agent_name
                .ok_or(AgentContextBuilderError::MissingField("agent_name"))?,
        );
        Ok(AgentContext {
            audit_redactor,
            tenant_label,
            ..ctx
        })
    }
}

impl AgentContext {
    /// Fluent builder. See [`AgentContextBuilder`].
    pub fn builder() -> AgentContextBuilder {
        AgentContextBuilder::default()
    }

    /// Construct an `AgentContext` with all required fields. `progress` defaults to `None`.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        llm: Arc<dyn LlmClient>,
        short_term: Arc<dyn ShortTermMemory>,
        long_term: Arc<dyn LongTermMemory>,
        episodic: Arc<dyn EpisodicMemory>,
        pubsub: Arc<dyn Pubsub>,
        kv: Arc<dyn KvStore>,
        request_reply: Arc<dyn RequestReply>,
        jobs: Arc<dyn JobQueue>,
        tools: Arc<dyn ToolInvoker>,
        run_id: RunId,
        cancel: CancellationToken,
        agent_name: impl Into<String>,
    ) -> Self {
        Self {
            llm,
            short_term,
            long_term,
            episodic,
            pubsub,
            kv,
            request_reply,
            jobs,
            tools,
            run_id,
            cancel,
            agent_name: agent_name.into(),
            progress: None,
            audit_redactor: None,
            tenant_label: None,
            parent_anchor: None,
        }
    }

    /// Replace the LLM client, returning a new context with all other fields cloned.
    pub fn with_llm(self, llm: Arc<dyn LlmClient>) -> Self {
        Self { llm, ..self }
    }

    /// Replace the tool invoker, returning a new context with all other fields cloned.
    pub fn with_tools(self, tools: Arc<dyn ToolInvoker>) -> Self {
        Self { tools, ..self }
    }

    /// Install the audit redactor for PII-flagged tools, returning a new
    /// context with all other fields cloned. When unset, dispatch falls
    /// back to the fail-closed [`crate::opaque_digest`] for flagged
    /// tools. Symmetric with [`AgentContextBuilder::audit_redactor`] for
    /// callers that construct via [`AgentContext::new`].
    pub fn with_audit_redactor(self, audit_redactor: Arc<dyn AuditRedactor>) -> Self {
        Self {
            audit_redactor: Some(audit_redactor),
            ..self
        }
    }

    /// Install the derived non-PII tenant attribution label. When set,
    /// the runtime records [`crate::Episode::RunAttributed`] alongside
    /// [`crate::Episode::Started`] at run entry. Caller must pre-derive
    /// the label (e.g. via [`crate::principal_hash`]); raw principals
    /// must never reach this surface. The label is audit metadata only
    /// and is never appended to short-term memory or LLM-visible
    /// messages.
    pub fn with_tenant_label(self, label: String) -> Self {
        Self {
            tenant_label: Some(label),
            ..self
        }
    }

    /// Install the cross-hop provenance anchor (the caller's chain-entry
    /// id). When set, the runtime records [`crate::Episode::RunOrigin`]
    /// alongside [`crate::Episode::Started`] at run entry. The value is
    /// recorded verbatim, is a caller-asserted unverified claim, and is
    /// never appended to short-term memory or LLM-visible messages.
    pub fn with_parent_anchor(self, anchor: String) -> Self {
        Self {
            parent_anchor: Some(anchor),
            ..self
        }
    }

    /// Read the derived non-PII tenant attribution label installed via
    /// [`Self::with_tenant_label`] or [`AgentContextBuilder::tenant_label`].
    /// Returns `None` when the context has not been bound to a verified
    /// caller (e.g. an in-process supervisor run that elides the
    /// inbound-MCP tenant-binding hop). External crates use this to
    /// route per-tenant budget governance without touching the
    /// pub(crate) field directly.
    #[must_use]
    pub fn tenant_label(&self) -> Option<&str> {
        self.tenant_label.as_deref()
    }

    /// Spawn a child context for a sub-run. Clones every `Arc<dyn …>`
    /// handle, mints a fresh [`RunId`], sets `agent_name`, and inherits
    /// the parent's cancellation token (cancelling the parent cancels
    /// the child, but the child can also be cancelled independently).
    ///
    /// Used by composite agents (`klieo-flows`'s `SequentialAgent`,
    /// `ParallelAgent`, etc.) to build per-leg contexts without manual
    /// struct-spread boilerplate.
    pub fn child(&self, agent_name: impl Into<String>) -> Self {
        Self {
            llm: self.llm.clone(),
            short_term: self.short_term.clone(),
            long_term: self.long_term.clone(),
            episodic: self.episodic.clone(),
            pubsub: self.pubsub.clone(),
            kv: self.kv.clone(),
            request_reply: self.request_reply.clone(),
            jobs: self.jobs.clone(),
            tools: self.tools.clone(),
            run_id: RunId::new(),
            cancel: self.cancel.child_token(),
            agent_name: agent_name.into(),
            progress: self.progress.clone(),
            audit_redactor: self.audit_redactor.clone(),
            tenant_label: self.tenant_label.clone(),
            parent_anchor: self.parent_anchor.clone(),
        }
    }
}

/// One agent — a typed function from `Input` to `Output` plus prompt
/// configuration.
#[async_trait]
pub trait Agent: Send + Sync {
    /// Input payload type.
    type Input: DeserializeOwned + Send + 'static;
    /// Output payload type.
    type Output: Serialize + Send + 'static;
    /// Domain-specific error type. Wrap `crate::Error` if you don't need
    /// a custom one.
    type Error: std::error::Error + Send + Sync + 'static;

    /// Stable agent name (used in spans + episodic events).
    fn name(&self) -> &str;

    /// System prompt prepended to the conversation.
    fn system_prompt(&self) -> &str;

    /// Tool catalogue this agent advertises to the LLM.
    fn tools(&self) -> &[ToolDef];

    /// Run one turn. Runtime supplies `ctx`; agent owns the per-call shape.
    ///
    /// ```
    /// # tokio_test::block_on(async {
    /// use async_trait::async_trait;
    /// use klieo_core::{Agent, AgentContext, ToolDef};
    /// struct Echo;
    /// #[async_trait]
    /// impl Agent for Echo {
    ///     type Input = String;
    ///     type Output = String;
    ///     type Error = std::io::Error;
    ///     fn name(&self) -> &str { "echo" }
    ///     fn system_prompt(&self) -> &str { "" }
    ///     fn tools(&self) -> &[ToolDef] { &[] }
    ///     async fn run(&self, _ctx: AgentContext, input: String) -> Result<String, Self::Error> {
    ///         Ok(input)
    ///     }
    /// }
    /// let agent = Echo;
    /// assert_eq!(agent.name(), "echo");
    /// # });
    /// ```
    async fn run(&self, ctx: AgentContext, input: Self::Input)
        -> Result<Self::Output, Self::Error>;
}

/// Canonical [`Agent`] implementation for the `String → String` case.
///
/// Wraps the boilerplate every example repeats: append the user
/// message to short-term memory, delegate to
/// [`crate::runtime::run_steps`] with the supplied system prompt.
///
/// Custom-typed agents (non-`String` input or output, alternative
/// turn shapes) still implement `Agent` by hand; `SimpleAgent` is
/// shortcut, not replacement.
///
/// ```
/// # tokio_test::block_on(async {
/// use klieo_core::{Agent, SimpleAgent};
/// let agent = SimpleAgent::new("hello", "Be brief.", vec![]);
/// assert_eq!(agent.name(), "hello");
/// assert_eq!(agent.system_prompt(), "Be brief.");
/// assert!(agent.tools().is_empty());
/// # });
/// ```
pub struct SimpleAgent {
    name: String,
    system_prompt: String,
    catalogue: Vec<crate::llm::ToolDef>,
    run_options: crate::runtime::RunOptions,
}

impl SimpleAgent {
    /// Build a `SimpleAgent` with the supplied name, system prompt,
    /// and tool catalogue. Uses [`crate::runtime::RunOptions::default`]
    /// — override via [`SimpleAgent::with_run_options`].
    pub fn new(
        name: impl Into<String>,
        system_prompt: impl Into<String>,
        catalogue: Vec<crate::llm::ToolDef>,
    ) -> Self {
        Self {
            name: name.into(),
            system_prompt: system_prompt.into(),
            catalogue,
            run_options: crate::runtime::RunOptions::default(),
        }
    }

    /// Override the [`crate::runtime::RunOptions`] passed to
    /// [`crate::runtime::run_steps`].
    pub fn with_run_options(mut self, options: crate::runtime::RunOptions) -> Self {
        self.run_options = options;
        self
    }

    /// Install a review policy (the ADR-045 HITL suspend gate) on this agent's
    /// run options. Composes with [`SimpleAgent::with_run_options`]: every other
    /// option field keeps its configured value; only the review policy is
    /// replaced.
    pub fn with_review_policy(
        mut self,
        policy: std::sync::Arc<dyn crate::runtime::ReviewPolicy>,
    ) -> Self {
        self.run_options = self.run_options.with_review_policy(policy);
        self
    }
}

#[async_trait]
impl Agent for SimpleAgent {
    type Input = String;
    type Output = String;
    type Error = crate::error::Error;

    fn name(&self) -> &str {
        &self.name
    }

    fn system_prompt(&self) -> &str {
        &self.system_prompt
    }

    fn tools(&self) -> &[ToolDef] {
        &self.catalogue
    }

    async fn run(&self, ctx: AgentContext, input: String) -> Result<String, Self::Error> {
        let thread = crate::ids::ThreadId::new(&self.name);
        ctx.short_term
            .append(
                thread.clone(),
                crate::llm::Message {
                    role: crate::llm::Role::User,
                    content: input,
                    tool_calls: vec![],
                    tool_call_id: None,
                },
            )
            .await?;
        crate::runtime::run_steps(&ctx, &self.system_prompt, thread, self.run_options.clone()).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_utils::{fake_context, FakeLlmClient};

    /// Compile-time check that AgentContext is Send + Sync + 'static.
    fn _assert_ctx_send_sync_static() {
        fn check<T: Send + Sync + 'static>() {}
        check::<AgentContext>();
    }

    fn parent_ctx() -> AgentContext {
        fake_context("parent")
    }

    struct PauseAlways;
    #[async_trait]
    impl crate::runtime::ReviewPolicy for PauseAlways {
        async fn should_pause_for_approval(
            &self,
            _step: u32,
            _message: &crate::llm::Message,
        ) -> Result<Option<String>, crate::error::Error> {
            Ok(Some("always".into()))
        }
    }

    #[test]
    fn child_mints_fresh_run_id() {
        let p = parent_ctx();
        let c = p.child("child-agent");
        assert_ne!(c.run_id, p.run_id);
    }

    #[tokio::test]
    async fn simple_agent_with_review_policy_folds_into_run_options() {
        let plain = SimpleAgent::new("a", "p", vec![]);
        assert!(
            plain.run_options.review_policy.is_never(),
            "a fresh SimpleAgent defaults to NeverReview"
        );

        let gated =
            SimpleAgent::new("a", "p", vec![]).with_review_policy(std::sync::Arc::new(PauseAlways));

        // Drive the installed policy, not just the `is_never()` sentinel: prove the
        // PauseAlways the builder stored is the one the run loop would consult.
        let msg = crate::llm::Message {
            role: crate::llm::Role::Assistant,
            content: "hi".into(),
            tool_calls: vec![],
            tool_call_id: None,
        };
        assert_eq!(
            gated
                .run_options
                .review_policy
                .should_pause_for_approval(1, &msg)
                .await
                .unwrap(),
            Some("always".to_string()),
            "with_review_policy must install the supplied policy on run_options, not merely flip is_never()"
        );
    }

    #[test]
    fn with_review_policy_preserves_other_run_options() {
        let opts = crate::runtime::RunOptions::default().with_checkpoint_bucket("custom-bucket");
        let agent = SimpleAgent::new("a", "p", vec![])
            .with_run_options(opts)
            .with_review_policy(std::sync::Arc::new(PauseAlways));

        assert!(
            !agent.run_options.review_policy.is_never(),
            "the review policy is installed"
        );
        assert_eq!(
            agent.run_options.checkpoint_kv_bucket.as_deref(),
            Some("custom-bucket"),
            "with_review_policy replaces only the policy — a field set via with_run_options survives"
        );
    }

    #[test]
    fn child_sets_new_agent_name() {
        let p = parent_ctx();
        let c = p.child("child-agent");
        assert_eq!(c.agent_name, "child-agent");
        assert_eq!(p.agent_name, "parent");
    }

    #[test]
    fn child_inherits_cancellation_from_parent() {
        let p = parent_ctx();
        let c = p.child("child-agent");
        assert!(!c.cancel.is_cancelled());
        p.cancel.cancel();
        assert!(
            c.cancel.is_cancelled(),
            "cancelling parent must propagate to child"
        );
    }

    #[test]
    fn child_shares_arc_handles_with_parent() {
        let p = parent_ctx();
        let c = p.child("child-agent");
        assert!(Arc::ptr_eq(&p.llm, &c.llm));
        assert!(Arc::ptr_eq(&p.short_term, &c.short_term));
        assert!(Arc::ptr_eq(&p.long_term, &c.long_term));
        assert!(Arc::ptr_eq(&p.episodic, &c.episodic));
        assert!(Arc::ptr_eq(&p.pubsub, &c.pubsub));
        assert!(Arc::ptr_eq(&p.kv, &c.kv));
        assert!(Arc::ptr_eq(&p.request_reply, &c.request_reply));
        assert!(Arc::ptr_eq(&p.jobs, &c.jobs));
        assert!(Arc::ptr_eq(&p.tools, &c.tools));
    }

    #[test]
    fn agent_event_variants_serialize_to_snake_case() {
        let evt = AgentEvent::LlmCallCompleted {
            tokens: 42,
            latency_ms: 180,
        };
        let s = serde_json::to_string(&evt).unwrap();
        assert!(s.contains(r#""kind":"llm_call_completed""#), "got: {s}");
        assert!(s.contains(r#""tokens":42"#));
        assert!(s.contains(r#""latency_ms":180"#));
    }

    #[test]
    fn simple_agent_exposes_constructor_args() {
        let cat = vec![ToolDef {
            name: "echo".into(),
            description: "e".into(),
            json_schema: serde_json::json!({"type": "object"}),
        }];
        let agent = SimpleAgent::new("hello", "Be brief.", cat.clone());
        assert_eq!(agent.name(), "hello");
        assert_eq!(agent.system_prompt(), "Be brief.");
        assert_eq!(agent.tools().len(), 1);
        assert_eq!(agent.tools()[0].name, "echo");
    }

    #[test]
    fn simple_agent_with_run_options_swaps_in_place() {
        let opts = crate::runtime::RunOptions {
            max_steps: 3,
            ..crate::runtime::RunOptions::default()
        };
        let agent = SimpleAgent::new("a", "s", vec![]).with_run_options(opts);
        assert_eq!(agent.run_options.max_steps, 3);
    }

    #[tokio::test]
    async fn simple_agent_run_appends_user_then_returns_assistant_text() {
        use crate::test_utils::FakeLlmStep;
        let mut ctx = fake_context("simple-test");
        ctx.llm =
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]));
        let short_term = ctx.short_term.clone();
        let agent = SimpleAgent::new("simple-test", "be brief", vec![]);
        let out = agent.run(ctx, "hi".into()).await.unwrap();
        assert_eq!(out, "done");

        let thread = crate::ids::ThreadId::new("simple-test");
        let loaded = short_term.load(thread, 1000).await.unwrap();
        assert!(
            loaded
                .iter()
                .any(|m| matches!(m.role, crate::llm::Role::User) && m.content == "hi"),
            "user message must be persisted to short-term before run_steps; got {loaded:?}",
        );
    }

    #[test]
    fn agent_event_tool_call_completed_serialises_name_and_ok() {
        let evt = AgentEvent::ToolCallCompleted {
            name: "echo".into(),
            ok: true,
        };
        let s = serde_json::to_string(&evt).unwrap();
        assert!(s.contains(r#""kind":"tool_call_completed""#));
        assert!(s.contains(r#""name":"echo""#));
        assert!(s.contains(r#""ok":true"#));
    }

    #[test]
    fn builder_missing_required_field_returns_err() {
        let result = AgentContext::builder().agent_name("test").build();
        match result {
            Err(AgentContextBuilderError::MissingField(name)) => {
                assert!(!name.is_empty());
            }
            Ok(_) => panic!("expected error when required fields are absent"),
        }
    }

    struct ConstRedactor;

    impl AuditRedactor for ConstRedactor {
        fn redact(&self, _value: &serde_json::Value) -> serde_json::Value {
            serde_json::json!("[redacted]")
        }
    }

    #[test]
    fn with_audit_redactor_installs_redactor_on_cloned_context() {
        let base = fake_context("redactor-test");
        assert!(base.audit_redactor.is_none());
        let ctx = base.with_audit_redactor(Arc::new(ConstRedactor));
        let redactor = ctx
            .audit_redactor
            .as_ref()
            .expect("redactor must be installed");
        assert_eq!(
            redactor.redact(&serde_json::json!({"pii": "secret"})),
            serde_json::json!("[redacted]")
        );
    }

    #[test]
    fn tenant_label_defaults_to_none() {
        let ctx = fake_context("tenant-default");
        assert!(ctx.tenant_label.is_none());
    }

    #[test]
    fn with_tenant_label_installs_label_on_cloned_context() {
        let base = fake_context("tenant-set");
        assert!(base.tenant_label.is_none());
        let ctx = base.with_tenant_label("hashed-tenant-abc".into());
        assert_eq!(ctx.tenant_label.as_deref(), Some("hashed-tenant-abc"));
    }

    #[test]
    fn builder_threads_tenant_label_into_context() {
        let base = fake_context("tenant-builder");
        let ctx = AgentContext::builder()
            .llm(base.llm.clone())
            .short_term(base.short_term.clone())
            .long_term(base.long_term.clone())
            .episodic(base.episodic.clone())
            .pubsub(base.pubsub.clone())
            .kv(base.kv.clone())
            .request_reply(base.request_reply.clone())
            .jobs(base.jobs.clone())
            .tools(base.tools.clone())
            .agent_name("builder-tenant")
            .tenant_label("hashed-tenant-xyz".into())
            .build()
            .expect("all required fields are set");
        assert_eq!(ctx.tenant_label.as_deref(), Some("hashed-tenant-xyz"));
    }

    #[test]
    fn child_inherits_tenant_label_from_parent() {
        let parent = fake_context("tenant-parent").with_tenant_label("hashed".into());
        let child = parent.child("tenant-child");
        assert_eq!(child.tenant_label.as_deref(), Some("hashed"));
    }

    #[test]
    fn parent_anchor_defaults_to_none() {
        let ctx = fake_context("anchor-default");
        assert!(ctx.parent_anchor.is_none());
    }

    #[test]
    fn with_parent_anchor_installs_anchor_on_cloned_context() {
        let base = fake_context("anchor-set");
        assert!(base.parent_anchor.is_none());
        let ctx = base.with_parent_anchor("sha256:abc123".into());
        assert_eq!(ctx.parent_anchor.as_deref(), Some("sha256:abc123"));
    }

    #[test]
    fn child_inherits_parent_anchor_from_parent() {
        let parent = fake_context("anchor-parent").with_parent_anchor("sha256:abc123".into());
        let child = parent.child("anchor-child");
        assert_eq!(child.parent_anchor.as_deref(), Some("sha256:abc123"));
    }

    #[test]
    fn builder_produces_valid_context_when_all_fields_set() {
        let base = fake_context("base");
        let ctx = AgentContext::builder()
            .llm(base.llm.clone())
            .short_term(base.short_term.clone())
            .long_term(base.long_term.clone())
            .episodic(base.episodic.clone())
            .pubsub(base.pubsub.clone())
            .kv(base.kv.clone())
            .request_reply(base.request_reply.clone())
            .jobs(base.jobs.clone())
            .tools(base.tools.clone())
            .agent_name("builder-test")
            .build()
            .expect("all required fields are set");
        assert_eq!(ctx.agent_name, "builder-test");
    }
}