klieo-core 0.36.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::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")]
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)]
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>>,
}

impl AgentContext {
    /// 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"#));
    }
}