klieo-core 0.41.2

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

/// 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>>,
}

/// 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>,
}

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
    }

    /// Consume the builder and produce an [`AgentContext`].
    ///
    /// Returns [`AgentContextBuilderError::MissingField`] when any
    /// required field was not set.
    pub fn build(self) -> Result<AgentContext, AgentContextBuilderError> {
        Ok(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"))?,
        ))
    }
}

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,
        }
    }

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

    /// 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(),
        }
    }
}

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

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

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

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

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