klieo_core/

agent.rs

//! `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 }
    }

    /// Replace the episodic store, returning a new context with all other
    /// fields preserved (e.g. to inject a per-run scoped or durable backend).
    pub fn with_episodic(self, episodic: Arc<dyn EpisodicMemory>) -> Self {
        Self { episodic, ..self }
    }

    /// Replace the short-term memory store, returning a new context with all
    /// other fields preserved (e.g. to inject a thread-isolated backend).
    pub fn with_short_term(self, short_term: Arc<dyn ShortTermMemory>) -> Self {
        Self { short_term, ..self }
    }

    /// Replace the long-term memory store, returning a new context with all
    /// other fields preserved (e.g. to swap vector backends per tenant).
    pub fn with_long_term(self, long_term: Arc<dyn LongTermMemory>) -> Self {
        Self { long_term, ..self }
    }

    /// Replace the run id, returning a new context with all other fields
    /// preserved (e.g. to resume a run or stitch episodes into a known chain).
    pub fn with_run_id(self, run_id: RunId) -> Self {
        Self { run_id, ..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
        }
    }

    /// Publish `payload` on `subject`, threading this run's id into the
    /// [`crate::bus::CAUSATION_HEADER`] so receivers can attribute causation,
    /// and record an [`crate::Episode::BusPublish`]. The publisher identity is
    /// the recording run, so the episode itself needs no extra field.
    ///
    /// Each dotted `subject` segment is validated via
    /// [`crate::bus::validate_subject_token`] before publishing (CWE-74 subject
    /// injection guard); a forbidden segment returns `BusError::Invalid` and
    /// nothing is published. Episode recording is best-effort and logged on
    /// failure — the message is already on the wire at that point.
    pub async fn publish(
        &self,
        subject: &str,
        payload: bytes::Bytes,
    ) -> Result<(), crate::error::BusError> {
        for segment in subject.split('.') {
            crate::bus::validate_subject_token(segment)?;
        }
        let mut headers = crate::bus::Headers::new();
        headers.insert(
            crate::bus::CAUSATION_HEADER.to_string(),
            self.run_id.to_string(),
        );
        self.pubsub.publish(subject, payload, headers).await?;
        if let Err(error) = self
            .episodic
            .record(
                self.run_id,
                crate::memory::Episode::BusPublish {
                    subject: subject.to_string(),
                },
            )
            .await
        {
            tracing::warn!(%error, run_id = %self.run_id, subject, "failed to record BusPublish episode");
        }
        Ok(())
    }

    /// Record receipt of a bus `msg` on `subject`: always records a
    /// [`crate::Episode::BusReceive`], and additionally records a
    /// [`crate::Episode::BusCausalLink`] when the message carries a
    /// [`crate::bus::CAUSATION_HEADER`] identifying the publishing run.
    ///
    /// Episode recording is best-effort; failures are logged and not
    /// propagated — the message has already been delivered at call time.
    pub async fn record_received(&self, subject: &str, msg: &crate::bus::Msg) {
        if let Err(error) = self
            .episodic
            .record(
                self.run_id,
                crate::memory::Episode::BusReceive {
                    subject: subject.to_string(),
                },
            )
            .await
        {
            tracing::warn!(%error, run_id = %self.run_id, subject, "failed to record BusReceive episode");
        }
        if let Some(caused_by_run) = msg.headers.get(crate::bus::CAUSATION_HEADER) {
            // The header is publisher-supplied and untrusted. A well-formed
            // value is the canonical `RunId` string `ctx.publish` stamps; a
            // non-ULID is malformed (or injected) and dropped — attribution is
            // best-effort metadata, mirroring the nats-jobs claim path.
            if ulid::Ulid::from_string(caused_by_run).is_err() {
                tracing::warn!(
                    run_id = %self.run_id,
                    subject,
                    "dropping malformed (non-ULID) causation header; no BusCausalLink recorded"
                );
            } else if let Err(error) = self
                .episodic
                .record(
                    self.run_id,
                    crate::memory::Episode::BusCausalLink {
                        subject: subject.to_string(),
                        caused_by_run: caused_by_run.clone(),
                    },
                )
                .await
            {
                tracing::warn!(%error, run_id = %self.run_id, subject, "failed to record BusCausalLink episode");
            }
        }
    }

    /// Enqueue `job` on `queue`, stamping this run as the causer
    /// (`job.causation_run_id`) so a worker that claims it can record a
    /// causation link via [`Self::record_claimed`]. Each dotted `queue`
    /// segment is validated ([`crate::bus::validate_subject_token`], CWE-74).
    pub async fn enqueue(
        &self,
        queue: &str,
        mut job: crate::bus::Job,
    ) -> Result<crate::ids::JobId, crate::error::BusError> {
        for segment in queue.split('.') {
            crate::bus::validate_subject_token(segment)?;
        }
        job.causation_run_id = Some(self.run_id);
        self.jobs.enqueue(queue, job).await
    }

    /// Record receipt of a claimed `job` from `queue`: always records a
    /// [`crate::Episode::BusReceive`], and additionally records a
    /// [`crate::Episode::BusCausalLink`] when the job carries a causing run.
    ///
    /// Episode recording is best-effort; failures are logged and not
    /// propagated — the job has already been claimed at call time.
    pub async fn record_claimed(&self, queue: &str, job: &crate::bus::ClaimedJob) {
        if let Err(error) = self
            .episodic
            .record(
                self.run_id,
                crate::memory::Episode::BusReceive {
                    subject: queue.to_string(),
                },
            )
            .await
        {
            tracing::warn!(%error, run_id = %self.run_id, queue, "failed to record BusReceive episode");
        }
        if let Some(caused_by_run) = job.causation_run_id {
            if let Err(error) = self
                .episodic
                .record(
                    self.run_id,
                    crate::memory::Episode::BusCausalLink {
                        subject: queue.to_string(),
                        caused_by_run: caused_by_run.to_string(),
                    },
                )
                .await
            {
                tracing::warn!(%error, run_id = %self.run_id, queue, "failed to record BusCausalLink episode");
            }
        }
    }

    /// CAS-write `value` at `key` in `bucket` and, on success, stamp this run as
    /// the value's causer ([`crate::bus::KvStore::put_causer`]) so a later reader
    /// can attribute the write via [`Self::record_kv_read`]. A CAS conflict
    /// (e.g. dedup redelivery) leaves any existing causer intact —
    /// first-writer-wins. The causer write is best-effort (logged, not
    /// propagated). Returns the value revision.
    pub async fn kv_cas_caused_by(
        &self,
        bucket: &str,
        key: &str,
        value: bytes::Bytes,
        expected: Option<crate::bus::Revision>,
    ) -> Result<crate::bus::Revision, crate::error::BusError> {
        let revision = self.kv.cas(bucket, key, value, expected).await?;
        if let Err(error) = self.kv.put_causer(bucket, key, self.run_id).await {
            tracing::warn!(%error, run_id = %self.run_id, bucket, key, "failed to write KV causer");
        }
        Ok(revision)
    }

    /// Record reading `key` from `bucket`: always a
    /// [`crate::Episode::BusReceive`], plus a [`crate::Episode::BusCausalLink`]
    /// to the run that wrote the value when one was stamped (via
    /// [`Self::kv_cas_caused_by`]). Best-effort; failures are logged, not
    /// propagated.
    pub async fn record_kv_read(&self, bucket: &str, key: &str) {
        if let Err(error) = self
            .episodic
            .record(
                self.run_id,
                crate::memory::Episode::BusReceive {
                    subject: bucket.to_string(),
                },
            )
            .await
        {
            tracing::warn!(%error, run_id = %self.run_id, bucket, key, "failed to record BusReceive episode");
        }
        match self.kv.causer_of(bucket, key).await {
            Ok(Some(caused_by_run)) => {
                if let Err(error) = self
                    .episodic
                    .record(
                        self.run_id,
                        crate::memory::Episode::BusCausalLink {
                            subject: bucket.to_string(),
                            caused_by_run: caused_by_run.to_string(),
                        },
                    )
                    .await
                {
                    tracing::warn!(%error, run_id = %self.run_id, bucket, key, "failed to record BusCausalLink episode");
                }
            }
            Ok(None) => {}
            Err(error) => {
                tracing::warn!(%error, run_id = %self.run_id, bucket, key, "failed to read KV causer")
            }
        }
    }

    /// 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"));
    }

    /// Minimal pubsub that records every published message for test inspection.
    struct CapturingPubsub {
        published: std::sync::Mutex<Vec<crate::bus::Msg>>,
    }

    impl CapturingPubsub {
        fn new() -> Self {
            Self {
                published: std::sync::Mutex::new(Vec::new()),
            }
        }

        fn messages(&self) -> Vec<crate::bus::Msg> {
            self.published
                .lock()
                .unwrap()
                .iter()
                .map(|m| crate::bus::Msg {
                    subject: m.subject.clone(),
                    payload: m.payload.clone(),
                    headers: m.headers.clone(),
                    ack: crate::bus::AckHandle::new(Box::new(NoopAckImpl)),
                })
                .collect()
        }
    }

    struct NoopAckImpl;

    #[async_trait]
    impl crate::bus::AckHandleImpl for NoopAckImpl {
        async fn ack(self: Box<Self>) -> Result<(), crate::error::BusError> {
            Ok(())
        }
        async fn nak(
            self: Box<Self>,
            _delay: std::time::Duration,
        ) -> Result<(), crate::error::BusError> {
            Ok(())
        }
        async fn term(self: Box<Self>) -> Result<(), crate::error::BusError> {
            Ok(())
        }
    }

    #[async_trait]
    impl crate::bus::Pubsub for CapturingPubsub {
        async fn publish(
            &self,
            subject: &str,
            payload: bytes::Bytes,
            headers: crate::bus::Headers,
        ) -> Result<(), crate::error::BusError> {
            self.published.lock().unwrap().push(crate::bus::Msg {
                subject: subject.to_string(),
                payload,
                headers,
                ack: crate::bus::AckHandle::new(Box::new(NoopAckImpl)),
            });
            Ok(())
        }

        async fn subscribe(
            &self,
            _subject: &str,
            _durable: crate::ids::DurableName,
        ) -> Result<crate::bus::MsgStream, crate::error::BusError> {
            Ok(Box::pin(tokio_stream::empty()))
        }
    }

    #[tokio::test]
    async fn publish_injects_causation_header_and_records_bus_publish() {
        use crate::bus::CAUSATION_HEADER;

        let capturing = Arc::new(CapturingPubsub::new());
        let mut ctx = fake_context("planner");
        ctx.pubsub = capturing.clone();
        let run = ctx.run_id;

        ctx.publish("jobs.research", bytes::Bytes::from_static(b"x"))
            .await
            .unwrap();

        let messages = capturing.messages();
        assert_eq!(messages.len(), 1, "expected exactly one published message");
        let msg = &messages[0];
        assert_eq!(msg.subject, "jobs.research");
        assert_eq!(
            msg.headers.get(CAUSATION_HEADER).map(String::as_str),
            Some(run.to_string().as_str()),
            "causation header must carry the publisher's run id"
        );

        let episodes = ctx.episodic.replay(run).await.unwrap();
        assert!(
            episodes
                .iter()
                .any(|e| matches!(e, crate::Episode::BusPublish { subject } if subject == "jobs.research")),
            "expected BusPublish episode for jobs.research; got {episodes:?}"
        );
    }

    #[tokio::test]
    async fn publish_rejects_wildcard_subject_segment() {
        let capturing = Arc::new(CapturingPubsub::new());
        let mut ctx = fake_context("planner");
        ctx.pubsub = capturing.clone();

        let result = ctx.publish("jobs.>", bytes::Bytes::from_static(b"x")).await;

        assert!(
            result.is_err(),
            "a wildcard subject segment must be rejected"
        );
        assert!(
            capturing.messages().is_empty(),
            "nothing is published when validation fails"
        );
    }

    #[tokio::test]
    async fn record_received_links_to_publisher_run() {
        use crate::bus::{AckHandle, Headers, Msg, CAUSATION_HEADER};

        let receiver = fake_context("synth");
        let recv_run = receiver.run_id;
        let pub_run = crate::ids::RunId::new();

        let mut headers = Headers::new();
        headers.insert(CAUSATION_HEADER.to_string(), pub_run.to_string());
        let msg = Msg {
            subject: "jobs.research".into(),
            payload: bytes::Bytes::from_static(b"x"),
            headers,
            ack: AckHandle::new(Box::new(NoopAckImpl)),
        };

        receiver.record_received("jobs.research", &msg).await;

        let eps = receiver.episodic.replay(recv_run).await.unwrap();
        assert!(
            eps.iter().any(|e| matches!(e,
                crate::Episode::BusCausalLink { subject, caused_by_run }
                if subject == "jobs.research" && *caused_by_run == pub_run.to_string())),
            "records BusCausalLink at the publisher run; got {eps:?}"
        );
        assert!(
            eps.iter().any(|e| matches!(e, crate::Episode::BusReceive { subject } if subject == "jobs.research")),
            "still records BusReceive"
        );
    }

    #[tokio::test]
    async fn record_received_with_malformed_causation_header_records_no_link() {
        use crate::bus::{AckHandle, Headers, Msg, CAUSATION_HEADER};

        let receiver = fake_context("synth");
        let recv_run = receiver.run_id;

        let mut headers = Headers::new();
        headers.insert(CAUSATION_HEADER.to_string(), "not-a-ulid".to_string());
        let msg = Msg {
            subject: "jobs.research".into(),
            payload: bytes::Bytes::from_static(b"x"),
            headers,
            ack: AckHandle::new(Box::new(NoopAckImpl)),
        };

        receiver.record_received("jobs.research", &msg).await;

        let eps = receiver.episodic.replay(recv_run).await.unwrap();
        assert!(
            !eps.iter()
                .any(|e| matches!(e, crate::Episode::BusCausalLink { .. })),
            "malformed (non-ULID) header -> no link; got {eps:?}"
        );
        assert!(
            eps.iter()
                .any(|e| matches!(e, crate::Episode::BusReceive { subject } if subject == "jobs.research")),
            "BusReceive still recorded unconditionally"
        );
    }

    #[tokio::test]
    async fn record_received_without_header_records_only_receive() {
        use crate::bus::{AckHandle, Headers, Msg};

        let receiver = fake_context("synth");
        let recv_run = receiver.run_id;
        let msg = Msg {
            subject: "jobs.research".into(),
            payload: bytes::Bytes::new(),
            headers: Headers::new(),
            ack: AckHandle::new(Box::new(NoopAckImpl)),
        };

        receiver.record_received("jobs.research", &msg).await;

        let eps = receiver.episodic.replay(recv_run).await.unwrap();
        assert!(
            !eps.iter()
                .any(|e| matches!(e, crate::Episode::BusCausalLink { .. })),
            "no header -> no link"
        );
        assert!(
            eps.iter()
                .any(|e| matches!(e, crate::Episode::BusReceive { .. })),
            "receive still recorded"
        );
    }

    #[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");
    }

    /// Minimal JobQueue that records every enqueued job for test inspection.
    struct CapturingJobQueue {
        enqueued: std::sync::Mutex<Vec<(String, crate::bus::Job)>>,
        counter: std::sync::atomic::AtomicU32,
    }

    impl CapturingJobQueue {
        fn new() -> Self {
            Self {
                enqueued: std::sync::Mutex::new(Vec::new()),
                counter: std::sync::atomic::AtomicU32::new(0),
            }
        }

        fn jobs(&self) -> Vec<(String, crate::bus::Job)> {
            self.enqueued.lock().unwrap().clone()
        }
    }

    #[async_trait]
    impl crate::bus::JobQueue for CapturingJobQueue {
        async fn enqueue(
            &self,
            queue: &str,
            job: crate::bus::Job,
        ) -> Result<crate::ids::JobId, crate::error::BusError> {
            let n = self
                .counter
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            self.enqueued.lock().unwrap().push((queue.to_string(), job));
            Ok(crate::ids::JobId(format!("test-{n}")))
        }

        async fn claim(
            &self,
            _queue: &str,
            _worker_id: &str,
            _lease_ttl: std::time::Duration,
        ) -> Result<Option<crate::bus::ClaimedJob>, crate::error::BusError> {
            Ok(None)
        }
    }

    struct NoopLeaseImpl;
    #[async_trait]
    impl crate::bus::LeaseImpl for NoopLeaseImpl {
        async fn heartbeat(&self) -> Result<(), crate::error::BusError> {
            Ok(())
        }
    }

    struct NoopClaimImpl;
    #[async_trait]
    impl crate::bus::ClaimHandleImpl for NoopClaimImpl {
        async fn ack(self: Box<Self>) -> Result<(), crate::error::BusError> {
            Ok(())
        }
        async fn nak(
            self: Box<Self>,
            _delay: std::time::Duration,
        ) -> Result<(), crate::error::BusError> {
            Ok(())
        }
        async fn dead_letter(self: Box<Self>, _reason: &str) -> Result<(), crate::error::BusError> {
            Ok(())
        }
    }

    fn noop_claimed_job(causation: Option<crate::ids::RunId>) -> crate::bus::ClaimedJob {
        crate::bus::ClaimedJob::new(
            crate::ids::JobId("test-0".into()),
            bytes::Bytes::from_static(b"payload"),
            crate::bus::Lease::new(Box::new(NoopLeaseImpl)),
            crate::bus::ClaimHandle::new(Box::new(NoopClaimImpl)),
        )
        .with_causation(causation)
    }

    #[tokio::test]
    async fn enqueue_sets_causation_to_this_run() {
        let capturing = Arc::new(CapturingJobQueue::new());
        let mut ctx = fake_context("enqueuer");
        ctx.jobs = capturing.clone();
        let run = ctx.run_id;

        let job = crate::bus::Job::new(bytes::Bytes::from_static(b"work"));
        ctx.enqueue("work.items", job).await.unwrap();

        let jobs = capturing.jobs();
        assert_eq!(jobs.len(), 1, "expected exactly one enqueued job");
        let (queue, enqueued_job) = &jobs[0];
        assert_eq!(queue, "work.items");
        assert_eq!(
            enqueued_job.causation_run_id,
            Some(run),
            "enqueue must stamp this run's id as causation_run_id"
        );
    }

    #[tokio::test]
    async fn enqueue_rejects_wildcard_queue_segment() {
        let capturing = Arc::new(CapturingJobQueue::new());
        let mut ctx = fake_context("enqueuer");
        ctx.jobs = capturing.clone();

        let job = crate::bus::Job::new(bytes::Bytes::from_static(b"work"));
        let result = ctx.enqueue("work.>", job).await;

        assert!(result.is_err(), "a wildcard queue segment must be rejected");
        assert!(
            capturing.jobs().is_empty(),
            "nothing is enqueued when validation fails"
        );
    }

    #[tokio::test]
    async fn record_claimed_links_to_causing_run() {
        let ctx = fake_context("worker");
        let worker_run = ctx.run_id;
        let causing_run = crate::ids::RunId::new();

        let claimed = noop_claimed_job(Some(causing_run));
        ctx.record_claimed("work.items", &claimed).await;

        let eps = ctx.episodic.replay(worker_run).await.unwrap();
        assert!(
            eps.iter().any(|e| matches!(e,
                crate::Episode::BusCausalLink { subject, caused_by_run }
                if subject == "work.items" && *caused_by_run == causing_run.to_string()
            )),
            "record_claimed must emit BusCausalLink for the causing run; got {eps:?}"
        );
        assert!(
            eps.iter().any(
                |e| matches!(e, crate::Episode::BusReceive { subject } if subject == "work.items")
            ),
            "record_claimed must also emit BusReceive; got {eps:?}"
        );
    }

    #[tokio::test]
    async fn record_claimed_without_causation_records_only_receive() {
        let ctx = fake_context("worker-no-cause");
        let worker_run = ctx.run_id;

        let claimed = noop_claimed_job(None);
        ctx.record_claimed("work.items", &claimed).await;

        let eps = ctx.episodic.replay(worker_run).await.unwrap();
        assert!(
            !eps.iter()
                .any(|e| matches!(e, crate::Episode::BusCausalLink { .. })),
            "no causation_run_id -> no BusCausalLink; got {eps:?}"
        );
        assert!(
            eps.iter()
                .any(|e| matches!(e, crate::Episode::BusReceive { .. })),
            "BusReceive must still be recorded; got {eps:?}"
        );
    }

    #[tokio::test]
    async fn kv_cas_caused_by_then_record_kv_read_links_to_writer() {
        use crate::test_utils::FakeKvStore;
        let kv: Arc<dyn KvStore> = Arc::new(FakeKvStore::default());

        let mut writer = fake_context("worker");
        writer.kv = kv.clone();
        let writer_run = writer.run_id;
        writer
            .kv_cas_caused_by("results", "k", bytes::Bytes::from_static(b"v"), None)
            .await
            .unwrap();

        let mut reader = fake_context("synth");
        reader.kv = kv.clone();
        let reader_run = reader.run_id;
        reader.record_kv_read("results", "k").await;

        let eps = reader.episodic.replay(reader_run).await.unwrap();
        assert!(
            eps.iter().any(
                |e| matches!(e, crate::Episode::BusReceive { subject } if subject == "results")
            ),
            "reader records BusReceive; got {eps:?}"
        );
        assert!(
            eps.iter().any(
                |e| matches!(e, crate::Episode::BusCausalLink { subject, caused_by_run }
                if subject == "results" && *caused_by_run == writer_run.to_string())
            ),
            "reader links to the writer run; got {eps:?}"
        );
    }

    #[tokio::test]
    async fn record_kv_read_unstamped_records_receive_only() {
        use crate::test_utils::FakeKvStore;
        let mut reader = fake_context("synth-no-cause");
        reader.kv = Arc::new(FakeKvStore::default());
        let run = reader.run_id;
        reader
            .kv
            .put("results", "k", bytes::Bytes::from_static(b"v"))
            .await
            .unwrap();
        reader.record_kv_read("results", "k").await;

        let eps = reader.episodic.replay(run).await.unwrap();
        assert!(
            eps.iter().any(
                |e| matches!(e, crate::Episode::BusReceive { subject } if subject == "results")
            ),
            "unstamped read still records BusReceive on the bucket; got {eps:?}"
        );
        assert!(
            !eps.iter()
                .any(|e| matches!(e, crate::Episode::BusCausalLink { .. })),
            "unstamped read -> no BusCausalLink; got {eps:?}"
        );
    }

    #[test]
    fn with_episodic_replaces_episodic_store() {
        use crate::test_utils::InMemoryEpisodic;
        let base = fake_context("with-episodic-test");
        let stub: Arc<dyn crate::memory::EpisodicMemory> = Arc::new(InMemoryEpisodic::default());
        // Confirm the stub is pointer-distinct from the default installed by fake_context.
        assert!(
            !Arc::ptr_eq(&base.episodic, &stub,),
            "precondition: stub must be a different allocation"
        );
        let ctx = base.with_episodic(Arc::clone(&stub));
        assert!(
            Arc::ptr_eq(&ctx.episodic, &stub),
            "with_episodic must install the supplied store"
        );
    }

    #[test]
    fn with_short_term_replaces_short_term_store() {
        use crate::test_utils::InMemoryShortTerm;
        let base = fake_context("with-short-term-test");
        let stub: Arc<dyn crate::memory::ShortTermMemory> = Arc::new(InMemoryShortTerm::default());
        assert!(
            !Arc::ptr_eq(&base.short_term, &stub),
            "precondition: stub must be a different allocation"
        );
        let ctx = base.with_short_term(Arc::clone(&stub));
        assert!(
            Arc::ptr_eq(&ctx.short_term, &stub),
            "with_short_term must install the supplied store"
        );
    }

    #[test]
    fn with_long_term_replaces_long_term_store() {
        use crate::test_utils::InMemoryLongTerm;
        let base = fake_context("with-long-term-test");
        let stub: Arc<dyn crate::memory::LongTermMemory> = Arc::new(InMemoryLongTerm::default());
        assert!(
            !Arc::ptr_eq(&base.long_term, &stub),
            "precondition: stub must be a different allocation"
        );
        let ctx = base.with_long_term(Arc::clone(&stub));
        assert!(
            Arc::ptr_eq(&ctx.long_term, &stub),
            "with_long_term must install the supplied store"
        );
    }

    #[test]
    fn with_run_id_installs_known_run_id() {
        let base = fake_context("with-run-id-test");
        let known_id = RunId::new();
        // Ensure we have a distinct id to install.
        assert_ne!(
            base.run_id, known_id,
            "precondition: RunId::new() produces distinct ids"
        );
        let ctx = base.with_run_id(known_id);
        assert_eq!(
            ctx.run_id, known_id,
            "with_run_id must install the given id"
        );
    }
}