use klieo_core::{
Agent, AgentContext, AuditRedactor, BusHandles, EpisodicMemory, Error, LlmClient,
LongTermMemory, MemoryHandles, RunId, ShortTermMemory, ToolDef, ToolInvoker,
};
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
#[cfg(feature = "tools")]
use klieo_core::Tool;
#[cfg(feature = "graph-rag")]
use crate::recall_recorder::{RecallRecorder, MAX_RECALL_QUERY_CHARS};
#[cfg(feature = "graph-rag")]
use klieo_core::Scope;
#[cfg(feature = "graph-rag")]
use klieo_memory_graph::KnowledgeGraph;
#[cfg(feature = "graph-rag")]
use klieo_memory_graph_rag::GraphAwareLongTerm;
#[cfg(feature = "runlog-sqlite")]
use klieo_runlog::RunLogStore;
#[cfg(any(feature = "runlog-sqlite", feature = "provenance-sqlite"))]
use std::path::PathBuf;
#[cfg(any(
feature = "llm-openai",
feature = "llm-anthropic",
feature = "llm-gemini"
))]
use secrecy::SecretString;
#[derive(Clone)]
pub struct App {
llm: Arc<dyn LlmClient>,
memory: MemoryHandles,
bus: BusHandles,
tools: Arc<dyn ToolInvoker>,
catalogue: Vec<ToolDef>,
parent_cancel: CancellationToken,
audit_redactor: Option<Arc<dyn AuditRedactor>>,
#[cfg(feature = "runlog")]
capture_sink: Option<Arc<klieo_runlog::RunLogCaptureSink>>,
#[cfg(feature = "runlog-sqlite")]
runlog_store: Option<Arc<dyn RunLogStore>>,
#[cfg(feature = "graph-rag")]
graph_rag_metrics: Option<Arc<klieo_memory_graph::RecallMetrics>>,
#[cfg(feature = "graph-rag")]
graph_projector: Option<Arc<std::sync::Mutex<Option<tokio::task::JoinHandle<()>>>>>,
#[cfg(feature = "graph-rag")]
graph_projector_cancel: Option<CancellationToken>,
#[cfg(feature = "graph-rag")]
graph_explorer_graph: Option<Arc<dyn KnowledgeGraph>>,
#[cfg(feature = "graph-rag")]
graph_explorer_recall: Option<Arc<GraphAwareLongTerm>>,
#[cfg(feature = "graph-rag")]
graph_explorer_scope: Option<Scope>,
}
impl std::fmt::Debug for App {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("App")
.field("llm", &self.llm.name())
.field("tools", &self.catalogue.len())
.field("cancelled", &self.parent_cancel.is_cancelled())
.finish_non_exhaustive()
}
}
#[cfg(feature = "runlog-sqlite")]
const MAX_AUTO_PERSIST_EPISODES: usize = 10_000;
impl App {
#[cfg(all(
feature = "llm-ollama",
feature = "memory-sqlite",
feature = "bus-memory",
feature = "tools",
))]
pub fn local() -> AppBuilder {
AppBuilder::new()
.ollama("http://localhost:11434", "qwen2.5:14b")
.sqlite(":memory:")
.memory_bus()
}
#[cfg(all(
feature = "llm-ollama",
feature = "memory-sqlite",
feature = "bus-memory",
feature = "tools",
))]
pub fn local_at(db_path: impl Into<std::path::PathBuf>) -> AppBuilder {
AppBuilder::new()
.ollama("http://localhost:11434", "qwen2.5:14b")
.sqlite(db_path)
.memory_bus()
}
#[cfg(feature = "data-dir")]
pub fn local_at_data_dir() -> Result<AppBuilder, Error> {
let dir = dirs::data_dir().ok_or_else(|| {
Error::Config(klieo_core::error::ConfigError::MissingKey(
"platform data directory unavailable; call App::local_at with an explicit path"
.into(),
))
})?;
let path = dir.join("klieo").join("klieo.db");
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).map_err(|e| Error::wrap("create klieo data dir", e))?;
}
Ok(AppBuilder::new()
.ollama("http://localhost:11434", "qwen2.5:14b")
.sqlite(path)
.memory_bus())
}
pub fn builder() -> AppBuilder {
AppBuilder::new()
}
pub fn context(&self, agent_name: impl Into<String>) -> AgentContext {
let run_id = RunId::new();
let ctx = AgentContext::new(
self.llm.clone(),
self.memory.short_term.clone(),
self.long_term_for_run(run_id),
self.memory.episodic.clone(),
self.bus.pubsub.clone(),
self.bus.kv.clone(),
self.bus.request_reply.clone(),
self.bus.jobs.clone(),
self.tools.clone(),
run_id,
self.parent_cancel.child_token(),
agent_name,
);
match &self.audit_redactor {
Some(redactor) => ctx.with_audit_redactor(redactor.clone()),
None => ctx,
}
}
#[cfg(feature = "graph-rag")]
fn long_term_for_run(&self, run_id: RunId) -> Arc<dyn LongTermMemory> {
if self.graph_explorer_recall.is_none() {
return self.memory.long_term.clone();
}
Arc::new(RecallRecorder::new(
self.memory.long_term.clone(),
self.memory.episodic.clone(),
run_id,
self.recall_redactor(),
MAX_RECALL_QUERY_CHARS,
))
}
#[cfg(feature = "graph-rag")]
fn recall_redactor(&self) -> Arc<dyn AuditRedactor> {
self.audit_redactor
.clone()
.unwrap_or_else(|| Arc::new(PassthroughRedactor))
}
#[cfg(not(feature = "graph-rag"))]
fn long_term_for_run(&self, _run_id: RunId) -> Arc<dyn LongTermMemory> {
self.memory.long_term.clone()
}
pub fn llm(&self) -> Arc<dyn LlmClient> {
self.llm.clone()
}
pub fn memory(&self) -> &MemoryHandles {
&self.memory
}
pub fn bus(&self) -> &BusHandles {
&self.bus
}
pub fn tools_invoker(&self) -> Arc<dyn ToolInvoker> {
self.tools.clone()
}
pub fn tools_catalogue(&self) -> &[ToolDef] {
&self.catalogue
}
#[cfg(feature = "mcp-server")]
pub fn mcp_server(&self) -> Arc<klieo_mcp_server::McpServer> {
Arc::new(klieo_mcp_server::McpServer::expose_tools(
self.tools_invoker(),
))
}
#[cfg(feature = "mcp-server")]
pub async fn serve_mcp_stdio(&self) -> Result<(), Error> {
self.mcp_server()
.serve_stdio()
.await
.map_err(|e| Error::wrap("serve mcp stdio", e))
}
pub fn cancel_token(&self) -> CancellationToken {
self.parent_cancel.clone()
}
#[cfg(feature = "runlog")]
pub fn capture_sink(&self) -> Option<Arc<dyn klieo_core::runtime::CaptureSink>> {
self.capture_sink
.clone()
.map(|sink| sink as Arc<dyn klieo_core::runtime::CaptureSink>)
}
#[cfg(feature = "runlog")]
pub fn runlog_capture_sink(&self) -> Option<Arc<klieo_runlog::RunLogCaptureSink>> {
self.capture_sink.clone()
}
#[cfg(feature = "graph-rag")]
pub fn graph_rag_metrics(&self) -> Option<klieo_memory_graph::RecallSnapshot> {
self.graph_rag_metrics.as_ref().map(|m| m.snapshot())
}
#[cfg(feature = "graph-rag")]
pub fn graph_explorer(
&self,
) -> Option<(Arc<dyn KnowledgeGraph>, Arc<GraphAwareLongTerm>, Scope)> {
Some((
self.graph_explorer_graph.clone()?,
self.graph_explorer_recall.clone()?,
self.graph_explorer_scope.clone()?,
))
}
#[cfg(all(feature = "graph-rag", feature = "observability-viz"))]
pub fn wire_graph_explorer(
&self,
builder: klieo_ops_viz::VizRouterBuilder,
) -> klieo_ops_viz::VizRouterBuilder {
match self.graph_explorer() {
Some((graph, recall, scope)) => builder.with_graph_explorer(graph, recall, scope),
None => builder,
}
}
#[cfg(feature = "runlog-sqlite")]
pub async fn run<A: Agent>(&self, agent: &A, input: A::Input) -> Result<A::Output, A::Error> {
let ctx = self.context(agent.name());
let run_id = ctx.run_id;
let agent_name = agent.name().to_string();
let result = agent.run(ctx, input).await;
if self.runlog_store.is_some() {
self.persist_run_log(run_id, &agent_name).await;
}
result
}
#[cfg(not(feature = "runlog-sqlite"))]
pub async fn run<A: Agent>(&self, agent: &A, input: A::Input) -> Result<A::Output, A::Error> {
let ctx = self.context(agent.name());
agent.run(ctx, input).await
}
#[cfg(feature = "runlog-sqlite")]
async fn persist_run_log(&self, run_id: RunId, agent_name: &str) {
let Some(store) = &self.runlog_store else {
return;
};
let episodes = match self.memory.episodic.replay(run_id).await {
Ok(episodes) => episodes,
Err(err) => {
tracing::warn!(%err, %run_id, "runlog auto-persist: replay failed");
return;
}
};
if episodes.is_empty() {
return;
}
if episodes.len() > MAX_AUTO_PERSIST_EPISODES {
tracing::warn!(
%run_id,
episode_count = episodes.len(),
cap = MAX_AUTO_PERSIST_EPISODES,
"runlog auto-persist: skipped, run's episode count exceeds the cap"
);
return;
}
let log = klieo_runlog::project_with_price_table(
run_id,
agent_name,
&episodes,
&[],
&klieo_runlog::PriceTable::default(),
);
if let Err(err) = store.put(&log).await {
tracing::warn!(%err, %run_id, "runlog auto-persist: store put failed");
}
}
}
#[cfg(feature = "graph-rag")]
impl Drop for App {
fn drop(&mut self) {
if let Some(cancel) = &self.graph_projector_cancel {
cancel.cancel();
}
}
}
#[cfg(feature = "graph-rag")]
impl App {
pub async fn shutdown(&self) {
if let Some(cancel) = &self.graph_projector_cancel {
cancel.cancel();
}
let handle = self
.graph_projector
.as_ref()
.and_then(|projector| match projector.lock() {
Ok(mut guard) => guard.take(),
Err(_) => {
tracing::error!("graph-rag projector mutex poisoned; skipping drain");
None
}
});
if let Some(handle) = handle {
if let Err(error) = handle.await {
tracing::error!(%error, "graph-rag projector task panicked during shutdown");
}
}
}
}
enum LlmShortcut {
#[cfg(feature = "llm-ollama")]
Ollama { base_url: String, model: String },
#[cfg(feature = "llm-openai")]
OpenAi {
api_key: SecretString,
model: String,
},
#[cfg(feature = "llm-anthropic")]
Anthropic {
api_key: SecretString,
model: String,
},
#[cfg(feature = "llm-gemini")]
Gemini {
api_key: SecretString,
model: String,
},
}
enum BusShortcut {
#[cfg(feature = "bus-memory")]
Memory,
#[cfg(feature = "bus-nats")]
Nats(Box<klieo_bus_nats::NatsBusConfig>),
}
#[cfg(any(feature = "memory-qdrant", feature = "memory-neo4j"))]
enum MemoryShortcut {
#[cfg(feature = "memory-qdrant")]
Qdrant {
config: klieo_memory_qdrant::QdrantConfig,
embedder: Option<Arc<dyn klieo_embed_common::Embedder>>,
},
#[cfg(feature = "memory-neo4j")]
Neo4j(Box<klieo_memory_neo4j::Neo4jConfig>),
}
#[cfg(any(feature = "memory-qdrant", feature = "memory-neo4j"))]
enum ResolvedPorts {
#[cfg(feature = "memory-qdrant")]
LongTermOnly(Arc<dyn LongTermMemory>),
#[cfg(feature = "memory-neo4j")]
ShortTermAndEpisodic {
short_term: Arc<dyn ShortTermMemory>,
episodic: Arc<dyn EpisodicMemory>,
},
}
#[cfg(any(feature = "memory-qdrant", feature = "memory-neo4j"))]
fn apply_ports(resolved: ResolvedPorts, memory: &mut MemoryHandles) {
match resolved {
#[cfg(feature = "memory-qdrant")]
ResolvedPorts::LongTermOnly(long_term) => {
memory.long_term = long_term;
}
#[cfg(feature = "memory-neo4j")]
ResolvedPorts::ShortTermAndEpisodic {
short_term,
episodic,
} => {
memory.short_term = short_term;
memory.episodic = episodic;
}
}
}
#[cfg(any(feature = "memory-qdrant", feature = "memory-neo4j"))]
async fn resolve_memory_backend(shortcut: MemoryShortcut) -> Result<ResolvedPorts, Error> {
match shortcut {
#[cfg(feature = "memory-qdrant")]
MemoryShortcut::Qdrant { config, embedder } => {
let embedder = resolve_qdrant_embedder(embedder)?;
let mem = klieo_memory_qdrant::MemoryQdrant::new(config, embedder).await?;
Ok(ResolvedPorts::LongTermOnly(mem.long_term))
}
#[cfg(feature = "memory-neo4j")]
MemoryShortcut::Neo4j(config) => {
let mem = klieo_memory_neo4j::MemoryNeo4j::new(*config).await?;
Ok(ResolvedPorts::ShortTermAndEpisodic {
short_term: mem.short_term,
episodic: mem.episodic,
})
}
}
}
#[cfg(feature = "memory-qdrant")]
fn resolve_qdrant_embedder(
explicit: Option<Arc<dyn klieo_embed_common::Embedder>>,
) -> Result<Arc<dyn klieo_embed_common::Embedder>, Error> {
if let Some(embedder) = explicit {
return Ok(embedder);
}
#[cfg(feature = "embed-fastembed")]
{
let embedder = klieo_embed_common::FastEmbedEmbedder::new()?;
Ok(Arc::new(embedder))
}
#[cfg(not(feature = "embed-fastembed"))]
{
Err(Error::AppBuildError {
missing: "memory.long_term embedder (qdrant_with requires an explicit \
embedder when the embed-fastembed feature is off)",
})
}
}
#[cfg(feature = "mcp")]
struct McpStdioSpec {
id: String,
cmd: String,
args: Vec<String>,
}
#[cfg(feature = "http")]
struct McpHttpSpec {
id: String,
url: String,
opts: klieo_tools_mcp::HttpConnectOptions,
}
#[cfg(feature = "mcp")]
enum McpSpec {
Stdio(McpStdioSpec),
#[cfg(feature = "http")]
Http(McpHttpSpec),
}
#[derive(Default)]
pub struct AppBuilder {
llm: Option<Arc<dyn LlmClient>>,
memory: Option<MemoryHandles>,
bus: Option<BusHandles>,
tools_invoker: Option<Arc<dyn ToolInvoker>>,
#[cfg(feature = "tools")]
pending_tools: Vec<Arc<dyn Tool>>,
parent_cancel: Option<CancellationToken>,
audit_redactor: Option<Arc<dyn AuditRedactor>>,
#[cfg(feature = "runlog")]
runlog_capture_provider: Option<String>,
#[cfg(feature = "runlog-sqlite")]
runlog_sqlite_path: Option<PathBuf>,
#[cfg(feature = "provenance-sqlite")]
provenance_sqlite_path: Option<PathBuf>,
llm_shortcut: Option<LlmShortcut>,
bus_shortcut: Option<BusShortcut>,
#[cfg(any(feature = "memory-qdrant", feature = "memory-neo4j"))]
memory_shortcut: Option<MemoryShortcut>,
model_registry_path: Option<std::path::PathBuf>,
short_term: Option<Arc<dyn ShortTermMemory>>,
long_term: Option<Arc<dyn LongTermMemory>>,
episodic: Option<Arc<dyn EpisodicMemory>>,
#[cfg(feature = "memory-sqlite")]
sqlite_path: Option<std::path::PathBuf>,
#[cfg(feature = "graph-rag")]
graph_rag: Option<crate::graph_rag::GraphRagConfig>,
#[cfg(feature = "mcp")]
pending_mcp: Vec<McpSpec>,
}
impl AppBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn llm(mut self, llm: Arc<dyn LlmClient>) -> Self {
self.llm = Some(llm);
self
}
pub fn memory(mut self, memory: impl Into<MemoryHandles>) -> Self {
self.memory = Some(memory.into());
self
}
pub fn bus(mut self, bus: impl Into<BusHandles>) -> Self {
self.bus = Some(bus.into());
self
}
pub fn tools_invoker(mut self, tools: Arc<dyn ToolInvoker>) -> Self {
self.tools_invoker = Some(tools);
self
}
pub fn audit_redactor(mut self, redactor: Arc<dyn AuditRedactor>) -> Self {
self.audit_redactor = Some(redactor);
self
}
#[cfg(feature = "runlog")]
pub fn runlog_capture(mut self, provider: impl Into<String>) -> Self {
self.runlog_capture_provider = Some(provider.into());
self
}
#[cfg(feature = "runlog-sqlite")]
pub fn sqlite_runlog(mut self, path: impl Into<PathBuf>) -> Self {
self.runlog_sqlite_path = Some(path.into());
self
}
#[cfg(feature = "provenance-sqlite")]
pub fn sqlite_provenance(mut self, path: impl Into<PathBuf>) -> Self {
self.provenance_sqlite_path = Some(path.into());
self
}
#[cfg(feature = "tools")]
pub fn tool<T: Tool + 'static>(mut self, tool: T) -> Self {
self.pending_tools.push(Arc::new(tool));
self
}
pub fn cancel_token(mut self, token: CancellationToken) -> Self {
self.parent_cancel = Some(token);
self
}
pub fn short_term(mut self, short_term: Arc<dyn ShortTermMemory>) -> Self {
self.short_term = Some(short_term);
self
}
pub fn long_term(mut self, long_term: Arc<dyn LongTermMemory>) -> Self {
self.long_term = Some(long_term);
self
}
pub fn episodic(mut self, episodic: Arc<dyn EpisodicMemory>) -> Self {
self.episodic = Some(episodic);
self
}
#[cfg(feature = "llm-ollama")]
pub fn ollama(mut self, base_url: impl Into<String>, model: impl Into<String>) -> Self {
self.llm_shortcut = Some(LlmShortcut::Ollama {
base_url: base_url.into(),
model: model.into(),
});
self
}
#[cfg(feature = "llm-ollama")]
pub fn model(mut self, model: impl Into<String>) -> Self {
let base_url = match self.llm_shortcut.take() {
Some(LlmShortcut::Ollama { base_url, .. }) => base_url,
_ => "http://localhost:11434".to_string(),
};
self.llm_shortcut = Some(LlmShortcut::Ollama {
base_url,
model: model.into(),
});
self
}
#[cfg(feature = "llm-openai")]
pub fn openai(mut self, api_key: impl Into<SecretString>, model: impl Into<String>) -> Self {
self.llm_shortcut = Some(LlmShortcut::OpenAi {
api_key: api_key.into(),
model: model.into(),
});
self
}
#[cfg(feature = "llm-anthropic")]
pub fn anthropic(mut self, api_key: impl Into<SecretString>, model: impl Into<String>) -> Self {
self.llm_shortcut = Some(LlmShortcut::Anthropic {
api_key: api_key.into(),
model: model.into(),
});
self
}
#[cfg(feature = "llm-gemini")]
pub fn gemini(mut self, api_key: impl Into<SecretString>, model: impl Into<String>) -> Self {
self.llm_shortcut = Some(LlmShortcut::Gemini {
api_key: api_key.into(),
model: model.into(),
});
self
}
pub fn model_registry(mut self, path: impl Into<std::path::PathBuf>) -> Self {
self.model_registry_path = Some(path.into());
self
}
#[cfg(feature = "memory-sqlite")]
pub fn sqlite(mut self, path: impl Into<std::path::PathBuf>) -> Self {
self.sqlite_path = Some(path.into());
self
}
#[cfg(feature = "memory-qdrant")]
pub fn qdrant_with(
mut self,
config: klieo_memory_qdrant::QdrantConfig,
embedder: Arc<dyn klieo_embed_common::Embedder>,
) -> Self {
self.memory_shortcut = Some(MemoryShortcut::Qdrant {
config,
embedder: Some(embedder),
});
self
}
#[cfg(all(feature = "memory-qdrant", feature = "embed-fastembed"))]
pub fn qdrant(mut self, url: impl Into<String>) -> Self {
self.memory_shortcut = Some(MemoryShortcut::Qdrant {
config: klieo_memory_qdrant::QdrantConfig::new(url),
embedder: None,
});
self
}
#[cfg(feature = "memory-neo4j")]
pub fn neo4j_with(mut self, config: klieo_memory_neo4j::Neo4jConfig) -> Self {
self.memory_shortcut = Some(MemoryShortcut::Neo4j(Box::new(config)));
self
}
#[cfg(feature = "memory-neo4j")]
pub fn neo4j(
self,
uri: impl Into<String>,
user: impl Into<String>,
password: impl Into<secrecy::SecretString>,
) -> Self {
self.neo4j_with(klieo_memory_neo4j::Neo4jConfig::new(uri, user, password))
}
#[cfg(feature = "graph-rag")]
pub fn graph_rag(self) -> Self {
self.graph_rag_with(crate::graph_rag::GraphRagConfig::in_memory())
}
#[cfg(feature = "graph-rag")]
pub fn graph_rag_with(mut self, cfg: crate::graph_rag::GraphRagConfig) -> Self {
self.graph_rag = Some(cfg);
self
}
#[cfg(feature = "bus-memory")]
pub fn memory_bus(mut self) -> Self {
self.bus_shortcut = Some(BusShortcut::Memory);
self
}
#[cfg(feature = "bus-nats")]
pub fn nats(mut self, config: klieo_bus_nats::NatsBusConfig) -> Self {
self.bus_shortcut = Some(BusShortcut::Nats(Box::new(config)));
self
}
#[cfg(feature = "mcp")]
pub fn mcp_stdio(
mut self,
id: impl Into<String>,
cmd: impl Into<String>,
args: Vec<String>,
) -> Self {
self.pending_mcp.push(McpSpec::Stdio(McpStdioSpec {
id: id.into(),
cmd: cmd.into(),
args,
}));
self
}
#[cfg(feature = "http")]
pub fn mcp_http(
mut self,
id: impl Into<String>,
url: impl Into<String>,
opts: klieo_tools_mcp::HttpConnectOptions,
) -> Self {
self.pending_mcp.push(McpSpec::Http(McpHttpSpec {
id: id.into(),
url: url.into(),
opts,
}));
self
}
pub async fn build(mut self) -> Result<App, Error> {
self.enforce_model_registry()?;
self.reject_qdrant_graph_rag_conflict()?;
#[cfg(feature = "graph-rag")]
let graph_rag_sqlite_path = self.sqlite_path.clone();
self.resolve_shortcuts().await?;
let llm = self.llm.ok_or(Error::AppBuildError { missing: "llm" })?;
#[cfg_attr(
not(any(feature = "graph-rag", feature = "provenance-sqlite")),
allow(unused_mut)
)]
let mut memory = self
.memory
.ok_or(Error::AppBuildError { missing: "memory" })?;
let bus = self.bus.ok_or(Error::AppBuildError { missing: "bus" })?;
let parent_cancel = self.parent_cancel.unwrap_or_default();
#[cfg(feature = "graph-rag")]
let mut graph_explorer_graph: Option<Arc<dyn KnowledgeGraph>> = None;
#[cfg(feature = "graph-rag")]
let mut graph_explorer_recall: Option<Arc<GraphAwareLongTerm>> = None;
#[cfg(feature = "graph-rag")]
let mut graph_explorer_scope: Option<Scope> = None;
#[cfg(feature = "graph-rag")]
let mut graph_rag = match self.graph_rag.take() {
Some(cfg) => {
let wired = cfg.wire(llm.clone(), graph_rag_sqlite_path).await?;
memory.long_term = wired.long_term.clone();
graph_explorer_graph = Some(wired.graph.clone());
graph_explorer_recall = Some(wired.graph_aware.clone());
graph_explorer_scope = Some(wired.projection_scope.clone());
Some(install_graph_rag(wired, &mut memory, &parent_cancel))
}
None => None,
};
#[cfg(feature = "provenance-sqlite")]
install_sqlite_provenance(self.provenance_sqlite_path, &mut memory)?;
let (tools, catalogue) = resolve_tools(
self.tools_invoker,
#[cfg(feature = "tools")]
self.pending_tools,
)?;
#[cfg(feature = "runlog-sqlite")]
let runlog_store = resolve_runlog_store(self.runlog_sqlite_path).await?;
let audit_redactor = resolve_audit_redactor(self.audit_redactor);
#[cfg(feature = "runlog")]
let capture_sink = build_capture_sink(self.runlog_capture_provider, audit_redactor.clone());
#[cfg(feature = "graph-rag")]
if graph_explorer_recall.is_some() && audit_redactor.is_none() {
tracing::warn!(
"graph-rag recall queries will be recorded with the length bound only and \
no content redaction; enable the `ops` feature or set `audit_redactor` \
to scrub PII from recorded MemoryRecall queries",
);
}
Ok(App {
llm,
memory,
bus,
tools,
catalogue,
parent_cancel,
audit_redactor,
#[cfg(feature = "runlog")]
capture_sink,
#[cfg(feature = "runlog-sqlite")]
runlog_store,
#[cfg(feature = "graph-rag")]
graph_rag_metrics: graph_rag.as_ref().map(|g| g.metrics.clone()),
#[cfg(feature = "graph-rag")]
graph_projector: graph_rag
.as_mut()
.and_then(|g| g.projector.take())
.map(|handle| Arc::new(std::sync::Mutex::new(Some(handle)))),
#[cfg(feature = "graph-rag")]
graph_projector_cancel: graph_rag.as_ref().and_then(|g| g.cancel.clone()),
#[cfg(feature = "graph-rag")]
graph_explorer_graph,
#[cfg(feature = "graph-rag")]
graph_explorer_recall,
#[cfg(feature = "graph-rag")]
graph_explorer_scope,
})
}
#[cfg(all(feature = "memory-qdrant", feature = "graph-rag"))]
fn reject_qdrant_graph_rag_conflict(&self) -> Result<(), Error> {
let has_qdrant_shortcut =
matches!(self.memory_shortcut, Some(MemoryShortcut::Qdrant { .. }));
if has_qdrant_shortcut && self.graph_rag.is_some() {
return Err(Error::Config(
klieo_core::error::ConfigError::InvalidValue {
key: "memory".to_string(),
reason:
"graph_rag manages the long-term vector store; do not also call \
.qdrant() — pass the Qdrant URL via GraphRagConfig::remote(qdrant_url, ..)"
.to_string(),
},
));
}
Ok(())
}
#[cfg(not(all(feature = "memory-qdrant", feature = "graph-rag")))]
fn reject_qdrant_graph_rag_conflict(&self) -> Result<(), Error> {
Ok(())
}
fn enforce_model_registry(&self) -> Result<(), Error> {
let Some(path) = self.model_registry_path.clone() else {
return Ok(());
};
let registry = match klieo_model_registry::ModelRegistry::load_from_path(&path) {
Ok(r) => r,
Err(e) => {
tracing::error!(target: "governance", error = %e, "model registry load failed");
return Err(Error::AppBuildError {
missing: "model-registry",
});
}
};
match shortcut_model(self.llm_shortcut.as_ref()) {
Some(pin) => match registry.status(&pin) {
Some(klieo_model_registry::DeprecationStatus::Sunset(since)) => {
return Err(Error::ModelSunset {
model: pin.to_string(),
since: since.to_string(),
});
}
Some(klieo_model_registry::DeprecationStatus::DeprecatedSince(since)) => {
tracing::warn!(target: "governance", model = %pin, %since, "configured model is deprecated");
}
Some(klieo_model_registry::DeprecationStatus::Active) => {}
None => {
tracing::warn!(target: "governance", model = %pin, "configured model is unpinned/unknown in the registry")
}
Some(_) => {
tracing::warn!(target: "governance", model = %pin, "unrecognized model status; treating as non-fatal")
}
},
None => {
tracing::warn!(target: "governance", "model not observable (raw LlmClient injected); registry not applied")
}
}
Ok(())
}
async fn resolve_shortcuts(&mut self) -> Result<(), Error> {
self.resolve_llm_shortcut();
self.resolve_sqlite().await?;
self.resolve_memory_partials()?;
#[cfg(any(feature = "memory-qdrant", feature = "memory-neo4j"))]
self.resolve_memory_shortcut().await?;
self.resolve_bus_shortcut().await?;
#[cfg(feature = "mcp")]
self.resolve_mcp().await?;
Ok(())
}
fn resolve_llm_shortcut(&mut self) {
if self.llm.is_some() {
return;
}
match self.llm_shortcut.take() {
None => {}
#[cfg(feature = "llm-ollama")]
Some(LlmShortcut::Ollama { base_url, model }) => {
self.llm = Some(Arc::new(klieo_llm_ollama::OllamaClient::new(
base_url, model,
)));
}
#[cfg(feature = "llm-openai")]
Some(LlmShortcut::OpenAi { api_key, model }) => {
self.llm = Some(Arc::new(klieo_llm_openai::OpenAiClient::new(
api_key, model,
)));
}
#[cfg(feature = "llm-anthropic")]
Some(LlmShortcut::Anthropic { api_key, model }) => {
self.llm = Some(Arc::new(klieo_llm_anthropic::AnthropicClient::new(
api_key, model,
)));
}
#[cfg(feature = "llm-gemini")]
Some(LlmShortcut::Gemini { api_key, model }) => {
self.llm = Some(Arc::new(klieo_llm_gemini::GeminiClient::new(
api_key, model,
)));
}
}
}
#[cfg(feature = "memory-sqlite")]
async fn resolve_sqlite(&mut self) -> Result<(), Error> {
if self.memory.is_some() {
return Ok(());
}
if let Some(path) = self.sqlite_path.take() {
let mem = klieo_memory_sqlite::MemorySqlite::new(
path,
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await?;
self.memory = Some(mem.into());
}
Ok(())
}
#[cfg(not(feature = "memory-sqlite"))]
#[allow(clippy::unused_async)]
async fn resolve_sqlite(&mut self) -> Result<(), Error> {
Ok(())
}
fn resolve_memory_partials(&mut self) -> Result<(), Error> {
if self.memory.is_some() {
return Ok(());
}
let any_partial =
self.short_term.is_some() || self.long_term.is_some() || self.episodic.is_some();
if !any_partial {
return Ok(());
}
let short_term = self.short_term.take().ok_or(Error::AppBuildError {
missing: "memory.short_term",
})?;
let long_term = self.long_term.take().ok_or(Error::AppBuildError {
missing: "memory.long_term",
})?;
let episodic = self.episodic.take().ok_or(Error::AppBuildError {
missing: "memory.episodic",
})?;
self.memory = Some(MemoryHandles::new(short_term, long_term, episodic));
Ok(())
}
#[cfg(any(feature = "memory-qdrant", feature = "memory-neo4j"))]
async fn resolve_memory_shortcut(&mut self) -> Result<(), Error> {
let Some(shortcut) = self.memory_shortcut.take() else {
return Ok(());
};
let memory = self.memory.as_mut().ok_or(Error::AppBuildError {
missing: "memory (a backend shortcut composes onto a base triple; \
also set .sqlite() or .memory())",
})?;
let resolved = resolve_memory_backend(shortcut).await?;
apply_ports(resolved, memory);
Ok(())
}
async fn resolve_bus_shortcut(&mut self) -> Result<(), Error> {
if self.bus.is_some() {
return Ok(());
}
match self.bus_shortcut.take() {
None => {}
#[cfg(feature = "bus-memory")]
Some(BusShortcut::Memory) => {
self.bus = Some(klieo_bus_memory::MemoryBus::new().into());
}
#[cfg(feature = "bus-nats")]
Some(BusShortcut::Nats(config)) => {
let bus = klieo_bus_nats::NatsBus::connect(*config).await?;
self.bus = Some(bus.into());
}
}
Ok(())
}
#[cfg(feature = "mcp")]
async fn resolve_mcp(&mut self) -> Result<(), Error> {
if self.tools_invoker.is_some() && !self.pending_mcp.is_empty() {
tracing::error!(
op = "resolve_mcp",
"conflicting builder options: mcp_stdio()/mcp_http() combined with tools_invoker()"
);
return Err(Error::AppBuildError {
missing: "tools (cannot combine .mcp_stdio()/.mcp_http() with .tools_invoker())",
});
}
for spec in self.pending_mcp.drain(..).collect::<Vec<_>>() {
let tools = match spec {
McpSpec::Stdio(stdio) => connect_mcp_stdio(stdio).await?,
#[cfg(feature = "http")]
McpSpec::Http(http) => connect_mcp_http(http).await?,
};
for tool in tools {
self.pending_tools.push(tool);
}
}
Ok(())
}
}
#[cfg(feature = "mcp")]
async fn connect_mcp_stdio(spec: McpStdioSpec) -> Result<Vec<Arc<dyn Tool>>, Error> {
klieo_tools_mcp::McpToolset::connect_stdio(
klieo_tools_mcp::McpServerId::new(&spec.id),
&spec.cmd,
klieo_tools_mcp::StdioConnectOptions {
args: spec.args,
..Default::default()
},
)
.await
.map_err(|e| {
tracing::error!(op = "mcp_stdio", server_id = %spec.id, cmd = %spec.cmd, err = %e, "MCP connect failed");
Error::Tool(e)
})
}
#[cfg(feature = "http")]
async fn connect_mcp_http(spec: McpHttpSpec) -> Result<Vec<Arc<dyn Tool>>, Error> {
klieo_tools_mcp::McpToolset::connect_http(
klieo_tools_mcp::McpServerId::new(&spec.id),
&spec.url,
spec.opts,
)
.await
.map_err(|e| {
tracing::error!(op = "mcp_http", server_id = %spec.id, url = %spec.url, err = %e, "MCP connect failed");
Error::Tool(e)
})
}
fn shortcut_model(shortcut: Option<&LlmShortcut>) -> Option<klieo_model_registry::ModelPin> {
let (provider, model): (&str, &str) = match shortcut? {
#[cfg(feature = "llm-ollama")]
LlmShortcut::Ollama { model, .. } => ("ollama", model),
#[cfg(feature = "llm-openai")]
LlmShortcut::OpenAi { model, .. } => ("openai", model),
#[cfg(feature = "llm-anthropic")]
LlmShortcut::Anthropic { model, .. } => ("anthropic", model),
#[cfg(feature = "llm-gemini")]
LlmShortcut::Gemini { model, .. } => ("gemini", model),
#[cfg(not(any(
feature = "llm-ollama",
feature = "llm-openai",
feature = "llm-anthropic",
feature = "llm-gemini"
)))]
_ => return None,
};
#[cfg(any(
feature = "llm-ollama",
feature = "llm-openai",
feature = "llm-anthropic",
feature = "llm-gemini"
))]
Some(klieo_model_registry::ModelPin::new(provider, model))
}
#[cfg(feature = "graph-rag")]
struct InstalledGraphRag {
metrics: Arc<klieo_memory_graph::RecallMetrics>,
projector: Option<tokio::task::JoinHandle<()>>,
cancel: Option<CancellationToken>,
}
#[cfg(feature = "graph-rag")]
fn install_graph_rag(
wired: crate::graph_rag::WiredGraphRag,
memory: &mut MemoryHandles,
parent_cancel: &CancellationToken,
) -> InstalledGraphRag {
if !wired.projector_enabled {
return InstalledGraphRag {
metrics: wired.metrics,
projector: None,
cancel: None,
};
}
let child = parent_cancel.child_token();
let (projectable, rx) =
klieo_graph_projector::ProjectableEpisodic::new(memory.episodic.clone());
memory.episodic = Arc::new(projectable);
let projector = Arc::new(
klieo_graph_projector::EpisodeProjector::builder()
.metrics(wired.metrics.clone())
.build(wired.graph, wired.extractor, wired.projection_scope),
);
let handle = projector.spawn(rx, child.clone());
InstalledGraphRag {
metrics: wired.metrics,
projector: Some(handle),
cancel: Some(child),
}
}
fn resolve_audit_redactor(
explicit: Option<Arc<dyn AuditRedactor>>,
) -> Option<Arc<dyn AuditRedactor>> {
if explicit.is_some() {
return explicit;
}
#[cfg(feature = "ops")]
{
Some(Arc::new(klieo_ops::redactor::DefaultRedactor::new()))
}
#[cfg(not(feature = "ops"))]
{
None
}
}
#[cfg(feature = "graph-rag")]
struct PassthroughRedactor;
#[cfg(feature = "graph-rag")]
impl AuditRedactor for PassthroughRedactor {
fn redact(&self, value: &serde_json::Value) -> serde_json::Value {
value.clone()
}
}
#[cfg(feature = "runlog")]
fn build_capture_sink(
provider: Option<String>,
audit_redactor: Option<Arc<dyn AuditRedactor>>,
) -> Option<Arc<klieo_runlog::RunLogCaptureSink>> {
provider.map(|provider| {
let mut sink = klieo_runlog::RunLogCaptureSink::new().with_provider(provider);
match audit_redactor {
Some(redactor) => sink = sink.with_redactor(redactor),
None => tracing::warn!(
"runlog capture is enabled but no audit_redactor is configured; \
captured prompts and completions are recorded verbatim with no \
PII redaction — enable the `ops` feature or set `audit_redactor` \
to scrub captured LLM I/O"
),
}
Arc::new(sink)
})
}
#[cfg(feature = "provenance-sqlite")]
pub const PROVENANCE_SCOPE: &str = "klieo-episodic";
#[cfg(feature = "provenance-sqlite")]
pub const PROVENANCE_ACTOR: &str = "klieo-agent";
#[cfg(feature = "provenance-sqlite")]
fn install_sqlite_provenance(
path: Option<PathBuf>,
memory: &mut MemoryHandles,
) -> Result<(), Error> {
let Some(path) = path else {
return Ok(());
};
let repo = klieo_provenance::SqliteProvenanceRepository::open(path)
.map_err(|e| Error::wrap("open sqlite provenance repository", e))?;
memory.episodic = Arc::new(klieo_provenance::ProvenanceEpisodic::new(
memory.episodic.clone(),
Arc::new(repo),
PROVENANCE_SCOPE,
PROVENANCE_ACTOR,
));
Ok(())
}
#[cfg(feature = "runlog-sqlite")]
async fn resolve_runlog_store(
path: Option<PathBuf>,
) -> Result<Option<Arc<dyn RunLogStore>>, Error> {
let Some(path) = path else {
return Ok(None);
};
let store = klieo_runlog::SqliteRunLogStore::new(path)
.await
.map_err(|e| Error::wrap("open sqlite runlog store", e))?;
Ok(Some(Arc::new(store)))
}
fn resolve_tools(
invoker: Option<Arc<dyn ToolInvoker>>,
#[cfg(feature = "tools")] pending: Vec<Arc<dyn Tool>>,
) -> Result<(Arc<dyn ToolInvoker>, Vec<ToolDef>), Error> {
#[cfg(feature = "tools")]
{
let has_pending = !pending.is_empty();
match (invoker, has_pending) {
(Some(_), true) => Err(Error::AppBuildError {
missing: "tools (cannot combine .tool() with .tools_invoker())",
}),
(Some(inv), false) => {
let cat = inv.catalogue();
Ok((inv, cat))
}
(None, _) => {
let mut chained = klieo_tools::ChainedInvoker::new();
for t in pending {
chained
.add_tool(t)
.map_err(|e| Error::wrap("tool registration failed", e))?;
}
let cat = chained.catalogue();
Ok((Arc::new(chained), cat))
}
}
}
#[cfg(not(feature = "tools"))]
{
let inv = invoker.ok_or(Error::AppBuildError { missing: "tools" })?;
let cat = inv.catalogue();
Ok((inv, cat))
}
}
#[cfg(test)]
mod tests {
use super::*;
use klieo_core::test_utils::{FakeLlmClient, FakeLlmStep, FakeToolInvoker};
fn fake_llm() -> Arc<dyn LlmClient> {
Arc::new(FakeLlmClient::new("fake"))
}
#[cfg(feature = "llm-ollama")]
#[test]
fn shortcut_model_maps_ollama() {
let s = LlmShortcut::Ollama {
base_url: "x".into(),
model: "qwen2.5:14b".into(),
};
assert_eq!(
shortcut_model(Some(&s)),
Some(klieo_model_registry::ModelPin::new("ollama", "qwen2.5:14b"))
);
}
#[test]
fn shortcut_model_none_for_raw_llm() {
assert_eq!(shortcut_model(None), None);
}
fn fake_tools() -> Arc<dyn ToolInvoker> {
Arc::new(FakeToolInvoker::new())
}
#[cfg(feature = "memory-sqlite")]
async fn sqlite_triple() -> klieo_memory_sqlite::MemorySqlite {
klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap()
}
#[cfg(all(feature = "memory-sqlite", feature = "memory-qdrant"))]
#[tokio::test]
async fn apply_ports_long_term_only_leaves_others_unchanged() {
let prior = sqlite_triple().await;
let resolved = sqlite_triple().await;
let prior_short = prior.short_term.clone();
let prior_episodic = prior.episodic.clone();
let resolved_long = resolved.long_term.clone();
let mut handles: MemoryHandles = prior.into();
apply_ports(
ResolvedPorts::LongTermOnly(resolved_long.clone()),
&mut handles,
);
assert!(
Arc::ptr_eq(&handles.long_term, &resolved_long),
"Qdrant-resolved result must set long_term"
);
assert!(
Arc::ptr_eq(&handles.short_term, &prior_short),
"long-term-only resolution must leave short_term untouched"
);
assert!(
Arc::ptr_eq(&handles.episodic, &prior_episodic),
"long-term-only resolution must leave episodic untouched"
);
}
#[cfg(all(feature = "memory-sqlite", feature = "memory-neo4j"))]
#[tokio::test]
async fn apply_ports_short_term_and_episodic_leaves_long_term_unchanged() {
let prior = sqlite_triple().await;
let resolved = sqlite_triple().await;
let prior_long = prior.long_term.clone();
let resolved_short = resolved.short_term.clone();
let resolved_episodic = resolved.episodic.clone();
let mut handles: MemoryHandles = prior.into();
apply_ports(
ResolvedPorts::ShortTermAndEpisodic {
short_term: resolved_short.clone(),
episodic: resolved_episodic.clone(),
},
&mut handles,
);
assert!(
Arc::ptr_eq(&handles.short_term, &resolved_short),
"Neo4j-resolved result must set short_term"
);
assert!(
Arc::ptr_eq(&handles.episodic, &resolved_episodic),
"Neo4j-resolved result must set episodic"
);
assert!(
Arc::ptr_eq(&handles.long_term, &prior_long),
"Neo4j-resolved result must leave long_term at the prior default"
);
}
#[cfg(all(
feature = "memory-qdrant",
feature = "embed-fastembed",
feature = "graph-rag"
))]
#[tokio::test]
async fn qdrant_shortcut_with_graph_rag_is_config_conflict() {
let err = App::local()
.model("dummy")
.qdrant("http://127.0.0.1:6334")
.graph_rag()
.build()
.await
.unwrap_err();
assert!(
matches!(err, Error::Config(_)),
"expected Error::Config for qdrant+graph_rag conflict, got {err:?}"
);
}
#[tokio::test]
async fn build_missing_llm_errors() {
let err = App::builder()
.memory(
klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap(),
)
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(err, Error::AppBuildError { missing: "llm" }));
}
#[tokio::test]
async fn build_missing_memory_errors() {
let err = App::builder()
.llm(fake_llm())
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(err, Error::AppBuildError { missing: "memory" }));
}
#[tokio::test]
async fn build_missing_bus_errors() {
let err = App::builder()
.llm(fake_llm())
.memory(
klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap(),
)
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(err, Error::AppBuildError { missing: "bus" }));
}
#[cfg(all(feature = "provenance-sqlite", feature = "memory-sqlite"))]
#[tokio::test]
async fn sqlite_provenance_wraps_episodic_and_records_through_it() {
use klieo_core::{Episode, RunId};
let app = App::builder()
.llm(fake_llm())
.memory(sqlite_triple().await)
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.sqlite_provenance(":memory:")
.build()
.await
.unwrap();
let ctx = app.context("prov-agent");
let run = RunId::new();
ctx.episodic
.record(
run,
Episode::Started {
agent: "prov-agent".into(),
},
)
.await
.expect("record through the provenance-wrapped episodic port");
let replayed = ctx.episodic.replay(run).await.unwrap();
assert_eq!(
replayed.len(),
1,
"the decorated episodic port must still persist the episode"
);
}
#[cfg(all(feature = "provenance-sqlite", feature = "memory-sqlite"))]
#[tokio::test]
async fn sqlite_provenance_open_failure_surfaces_as_build_error() {
let err = App::builder()
.llm(fake_llm())
.memory(sqlite_triple().await)
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.sqlite_provenance("/nonexistent-dir-klieo-prov/chain.db")
.build()
.await
.unwrap_err();
assert!(
matches!(err, Error::Other { .. }),
"provenance open failure must surface as a wrapped build error, got {err:?}"
);
}
#[cfg(feature = "graph-rag")]
#[tokio::test]
async fn metrics_none_when_graph_rag_not_configured() {
let app = App::local().model("dummy").build().await.unwrap();
assert!(
app.graph_rag_metrics().is_none(),
"an App built without .graph_rag* must expose no recall metrics"
);
}
#[cfg(feature = "graph-rag")]
#[tokio::test]
async fn dropped_app_cancels_projector() {
use crate::embed_common::FakeEmbedder;
let cfg = crate::graph_rag::GraphRagConfig::in_memory()
.embedder(Arc::new(FakeEmbedder::new(384)), "test-fake");
let cancel = {
let app = App::local()
.model("dummy")
.graph_rag_with(cfg)
.build()
.await
.unwrap();
app.graph_projector_cancel.clone().unwrap()
};
assert!(
cancel.is_cancelled(),
"dropping the App must cancel the projector token"
);
}
#[cfg(feature = "graph-rag")]
#[tokio::test]
async fn graph_explorer_returns_handles_when_graph_rag_configured() {
use crate::embed_common::FakeEmbedder;
let cfg = crate::graph_rag::GraphRagConfig::in_memory()
.embedder(Arc::new(FakeEmbedder::new(384)), "test-fake");
let app = App::local()
.model("dummy")
.graph_rag_with(cfg)
.build()
.await
.unwrap();
assert!(
app.graph_explorer().is_some(),
"an App built with .graph_rag* must expose graph explorer handles"
);
}
#[cfg(feature = "graph-rag")]
#[tokio::test]
async fn graph_explorer_is_none_without_graph_rag() {
let app = App::local().model("dummy").build().await.unwrap();
assert!(
app.graph_explorer().is_none(),
"an App built without .graph_rag* must expose no graph explorer handles"
);
}
#[cfg(feature = "graph-rag")]
#[tokio::test]
async fn context_records_memory_recall_when_graph_rag_configured() {
use crate::embed_common::FakeEmbedder;
use klieo_core::Episode;
let cfg = crate::graph_rag::GraphRagConfig::in_memory()
.embedder(Arc::new(FakeEmbedder::new(384)), "test-fake");
let app = App::local()
.model("dummy")
.graph_rag_with(cfg)
.build()
.await
.unwrap();
let ctx = app.context("recall-test-agent");
ctx.long_term
.recall(
Scope::Workspace("recall-test".to_string()),
"hello world query",
3,
)
.await
.expect("recall through the per-run context's long-term port");
let episodes = app
.memory()
.episodic
.replay(ctx.run_id)
.await
.expect("replay this run's episodes");
let recalls: Vec<_> = episodes
.iter()
.filter(|e| matches!(e, Episode::MemoryRecall { .. }))
.collect();
assert_eq!(
recalls.len(),
1,
"context() must install a per-run RecallRecorder when graph-rag is configured"
);
match recalls[0] {
Episode::MemoryRecall { query, k, .. } => {
assert_eq!(
query, "hello world query",
"no PII pattern here — query stored as-is"
);
assert_eq!(*k, 3);
}
other => panic!("expected MemoryRecall, got {other:?}"),
}
}
#[cfg(feature = "graph-rag")]
#[tokio::test]
async fn context_records_redacted_query_with_configured_redactor() {
use crate::embed_common::FakeEmbedder;
use klieo_core::Episode;
const SECRET_TOKEN: &str = "AKIA1234567890";
let cfg = crate::graph_rag::GraphRagConfig::in_memory()
.embedder(Arc::new(FakeEmbedder::new(384)), "test-fake");
let app = App::local()
.model("dummy")
.audit_redactor(Arc::new(ConstRedactor))
.graph_rag_with(cfg)
.build()
.await
.unwrap();
let ctx = app.context("redaction-test-agent");
ctx.long_term
.recall(
Scope::Workspace("redaction-test".to_string()),
&format!("query containing {SECRET_TOKEN}"),
3,
)
.await
.expect("recall through the per-run context's long-term port");
let episodes = app
.memory()
.episodic
.replay(ctx.run_id)
.await
.expect("replay this run's episodes");
let recalls: Vec<_> = episodes
.iter()
.filter(|e| matches!(e, Episode::MemoryRecall { .. }))
.collect();
assert_eq!(
recalls.len(),
1,
"exactly one MemoryRecall must be recorded"
);
match recalls[0] {
Episode::MemoryRecall { query, .. } => {
assert!(
!query.contains(SECRET_TOKEN),
"the configured redactor must scrub the token from the recorded query"
);
assert_eq!(
query, "[const]",
"the configured ConstRedactor's fixed value must reach the recorded episode"
);
}
other => panic!("expected MemoryRecall, got {other:?}"),
}
}
#[tokio::test]
async fn build_local_with_only_model_succeeds() {
let app = App::local().model("dummy-model").build().await.unwrap();
assert!(app.tools_catalogue().is_empty());
let ctx = app.context("test-agent");
assert_eq!(ctx.agent_name, "test-agent");
}
#[tokio::test]
async fn context_mints_fresh_run_id_per_call() {
let app = App::local().model("dummy").build().await.unwrap();
let a = app.context("agent");
let b = app.context("agent");
assert_ne!(a.run_id, b.run_id);
}
#[tokio::test]
async fn cancel_token_propagates_to_minted_context() {
let app = App::local().model("dummy").build().await.unwrap();
let ctx = app.context("agent");
assert!(!ctx.cancel.is_cancelled());
app.cancel_token().cancel();
assert!(ctx.cancel.is_cancelled());
}
struct ConstRedactor;
impl AuditRedactor for ConstRedactor {
fn redact(&self, _value: &serde_json::Value) -> serde_json::Value {
serde_json::json!("[const]")
}
}
#[cfg(feature = "ops")]
#[tokio::test]
async fn ops_feature_wires_default_audit_redactor() {
let app = App::local().model("dummy").build().await.unwrap();
assert!(
app.audit_redactor.is_some(),
"the ops feature must default a redactor so PII-flagged tools \
are masked rather than falling back to the opaque digest"
);
}
#[tokio::test]
async fn explicit_audit_redactor_overrides_default() {
let app = App::local()
.model("dummy")
.audit_redactor(Arc::new(ConstRedactor))
.build()
.await
.unwrap();
let redactor = app
.audit_redactor
.as_ref()
.expect("explicit redactor must be retained");
assert_eq!(
redactor.redact(&serde_json::json!("anything")),
serde_json::json!("[const]")
);
}
#[cfg(feature = "data-dir")]
#[test]
fn local_at_data_dir_resolves_and_creates_parent_dir() {
let _builder =
App::local_at_data_dir().expect("dirs::data_dir resolvable on this platform");
let parent = dirs::data_dir()
.expect("data_dir already validated above")
.join("klieo");
assert!(
parent.is_dir(),
"expected klieo data dir to be created: {parent:?}",
);
let _builder_again = App::local_at_data_dir().expect("data_dir fn must be idempotent");
}
#[tokio::test]
async fn port_accessors_expose_configured_handles() {
let llm = fake_llm();
let app = App::builder()
.llm(Arc::clone(&llm))
.memory(
klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap(),
)
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap();
assert!(Arc::ptr_eq(&app.llm(), &llm));
let ctx = app.context("ports-test");
assert!(Arc::ptr_eq(&app.memory().short_term, &ctx.short_term));
assert!(Arc::ptr_eq(&app.memory().long_term, &ctx.long_term));
assert!(Arc::ptr_eq(&app.memory().episodic, &ctx.episodic));
assert!(Arc::ptr_eq(&app.bus().pubsub, &ctx.pubsub));
assert!(Arc::ptr_eq(&app.bus().kv, &ctx.kv));
assert!(Arc::ptr_eq(&app.bus().request_reply, &ctx.request_reply));
assert!(Arc::ptr_eq(&app.bus().jobs, &ctx.jobs));
assert!(Arc::ptr_eq(&app.tools_invoker(), &ctx.tools));
}
#[tokio::test]
async fn cannot_combine_tool_with_tools_invoker() {
let mut builder = App::local().model("dummy");
builder.pending_tools.push(Arc::new(EchoTool));
let err = builder
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(
err,
Error::AppBuildError {
missing: "tools (cannot combine .tool() with .tools_invoker())"
}
));
}
struct EchoTool;
#[async_trait::async_trait]
impl Tool for EchoTool {
fn name(&self) -> &str {
"echo"
}
fn description(&self) -> &str {
"echo"
}
fn json_schema(&self) -> &serde_json::Value {
static SCHEMA: std::sync::OnceLock<serde_json::Value> = std::sync::OnceLock::new();
SCHEMA.get_or_init(|| serde_json::json!({"type": "object"}))
}
async fn invoke(
&self,
_args: serde_json::Value,
_ctx: klieo_core::ToolCtx,
) -> Result<serde_json::Value, klieo_core::ToolError> {
Ok(serde_json::Value::Null)
}
}
#[tokio::test]
async fn partial_memory_composes_into_handles() {
let mem = klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap();
let app = App::builder()
.llm(fake_llm())
.short_term(mem.short_term.clone())
.long_term(mem.long_term.clone())
.episodic(mem.episodic.clone())
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap();
let ctx = app.context("partial");
assert!(Arc::ptr_eq(&ctx.short_term, &mem.short_term));
assert!(Arc::ptr_eq(&ctx.long_term, &mem.long_term));
assert!(Arc::ptr_eq(&ctx.episodic, &mem.episodic));
}
#[tokio::test]
async fn partial_memory_missing_long_term_errors() {
let mem = klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap();
let err = App::builder()
.llm(fake_llm())
.short_term(mem.short_term.clone())
.episodic(mem.episodic.clone())
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(
err,
Error::AppBuildError {
missing: "memory.long_term"
}
));
}
#[tokio::test]
async fn explicit_memory_overrides_partials() {
let primary = klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap();
let secondary = klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap();
let primary_st = primary.short_term.clone();
let app = App::builder()
.llm(fake_llm())
.memory(primary)
.short_term(secondary.short_term)
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap();
let ctx = app.context("override");
assert!(Arc::ptr_eq(&ctx.short_term, &primary_st));
}
#[cfg(feature = "llm-openai")]
#[tokio::test]
async fn openai_shortcut_sets_llm() {
let app = App::local()
.openai("test-key".to_string(), "gpt-4o-mini")
.build()
.await
.unwrap();
let ctx = app.context("a");
assert!(
ctx.llm.name().starts_with("openai"),
"got {}",
ctx.llm.name()
);
}
#[cfg(feature = "llm-anthropic")]
#[tokio::test]
async fn anthropic_shortcut_sets_llm() {
let app = App::local()
.anthropic("test-key".to_string(), "claude-3-5-haiku-latest")
.build()
.await
.unwrap();
let ctx = app.context("a");
assert!(
ctx.llm.name().starts_with("anthropic"),
"got {}",
ctx.llm.name()
);
}
#[cfg(feature = "llm-gemini")]
#[tokio::test]
async fn gemini_shortcut_sets_llm() {
let app = App::local()
.gemini("test-key".to_string(), "gemini-2.0-flash")
.build()
.await
.unwrap();
let ctx = app.context("a");
assert!(ctx.llm.name().contains("gemini"), "got {}", ctx.llm.name());
}
#[tokio::test]
async fn last_llm_shortcut_wins() {
let app = App::local()
.ollama("http://localhost:11434", "qwen")
.ollama("http://other:11434", "llama")
.build()
.await
.unwrap();
let ctx = app.context("a");
assert!(
ctx.llm.name().starts_with("ollama"),
"got {}",
ctx.llm.name()
);
}
#[cfg(feature = "mcp")]
#[tokio::test]
async fn mcp_stdio_registers_entry_verified_via_conflict_guard() {
struct DummyInvoker;
#[async_trait::async_trait]
impl klieo_core::ToolInvoker for DummyInvoker {
fn catalogue(&self) -> Vec<klieo_core::ToolDef> {
vec![]
}
async fn invoke(
&self,
_name: &str,
_args: serde_json::Value,
_ctx: klieo_core::ToolCtx,
) -> Result<serde_json::Value, klieo_core::error::ToolError> {
Ok(serde_json::Value::Null)
}
}
let err = App::local()
.mcp_stdio("test-srv", "/bin/echo", vec![])
.tools_invoker(Arc::new(DummyInvoker))
.build()
.await
.unwrap_err();
assert!(
matches!(err, klieo_core::Error::AppBuildError { .. }),
"expected AppBuildError for conflict, got {err:?}"
);
}
#[cfg(feature = "mcp")]
#[tokio::test]
async fn mcp_stdio_errors_on_missing_binary() {
let err = App::local()
.mcp_stdio("test-server", "/nonexistent/binary/path", vec![])
.build()
.await
.unwrap_err();
assert!(
matches!(err, klieo_core::Error::Tool(_)),
"expected Error::Tool, got {err:?}"
);
}
#[cfg(feature = "mcp")]
#[tokio::test]
async fn mcp_stdio_combined_with_tools_invoker_errors_at_build() {
struct DummyInvoker;
#[async_trait::async_trait]
impl klieo_core::ToolInvoker for DummyInvoker {
fn catalogue(&self) -> Vec<klieo_core::ToolDef> {
vec![]
}
async fn invoke(
&self,
_name: &str,
_args: serde_json::Value,
_ctx: klieo_core::ToolCtx,
) -> Result<serde_json::Value, klieo_core::error::ToolError> {
Ok(serde_json::Value::Null)
}
}
let err = App::local()
.tools_invoker(Arc::new(DummyInvoker))
.mcp_stdio("srv", "/bin/echo", vec![])
.build()
.await
.unwrap_err();
assert!(
matches!(err, klieo_core::Error::AppBuildError { .. }),
"expected AppBuildError, got {err:?}"
);
}
#[cfg(feature = "http")]
#[tokio::test]
async fn mcp_http_combined_with_tools_invoker_errors_at_build() {
struct DummyInvoker;
#[async_trait::async_trait]
impl klieo_core::ToolInvoker for DummyInvoker {
fn catalogue(&self) -> Vec<klieo_core::ToolDef> {
vec![]
}
async fn invoke(
&self,
_name: &str,
_args: serde_json::Value,
_ctx: klieo_core::ToolCtx,
) -> Result<serde_json::Value, klieo_core::error::ToolError> {
Ok(serde_json::Value::Null)
}
}
let err = App::local()
.mcp_http(
"remote",
"https://mcp.example.com/mcp",
klieo_tools_mcp::HttpConnectOptions::default(),
)
.tools_invoker(Arc::new(DummyInvoker))
.build()
.await
.unwrap_err();
assert!(
matches!(err, klieo_core::Error::AppBuildError { .. }),
"expected AppBuildError for conflict, got {err:?}"
);
}
#[cfg(feature = "http")]
#[tokio::test]
async fn mcp_http_rejects_plain_http_remote_url_at_build() {
let err = App::local()
.mcp_http(
"remote",
"http://mcp.example.com/mcp",
klieo_tools_mcp::HttpConnectOptions::default(),
)
.build()
.await
.unwrap_err();
assert!(
matches!(err, klieo_core::Error::Tool(_)),
"expected Error::Tool for plain-http rejection, got {err:?}"
);
}
#[cfg(feature = "bus-nats")]
#[tokio::test]
async fn nats_shortcut_errors_on_unreachable_server() {
let cfg = klieo_bus_nats::NatsBusConfig {
url: "nats://127.0.0.1:14223".to_string(),
..klieo_bus_nats::NatsBusConfig::default()
};
let err = App::builder()
.llm(fake_llm())
.memory(
klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap(),
)
.nats(cfg)
.tools_invoker(fake_tools())
.build()
.await
.unwrap_err();
assert!(matches!(err, Error::Bus(_)), "got {err:?}");
}
#[tokio::test]
async fn run_delegates_to_agent_with_fresh_context() {
let app = App::builder()
.llm(Arc::new(
FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]),
))
.memory(
klieo_memory_sqlite::MemorySqlite::new(
":memory:",
Arc::new(klieo_memory_sqlite::DummyEmbedder),
)
.await
.unwrap(),
)
.bus(klieo_bus_memory::MemoryBus::new())
.tools_invoker(fake_tools())
.build()
.await
.unwrap();
let agent =
klieo_core::SimpleAgent::new("greeter", "be brief", app.tools_catalogue().to_vec());
let out = app.run(&agent, "hi".into()).await.unwrap();
assert_eq!(out, "done");
}
#[cfg(all(
feature = "llm-ollama",
feature = "memory-sqlite",
feature = "bus-memory",
feature = "tools",
))]
#[tokio::test]
async fn build_local_at_file_path_succeeds() {
let db_path = std::env::temp_dir().join("klieo_test_local_at.db");
let _ = std::fs::remove_file(&db_path);
let app = App::local_at(&db_path)
.model("dummy-model")
.build()
.await
.unwrap();
assert!(app.tools_catalogue().is_empty());
let ctx = app.context("test-agent");
assert_eq!(ctx.agent_name, "test-agent");
assert!(
db_path.exists(),
"sqlite file should have been created at {db_path:?}"
);
let _ = std::fs::remove_file(&db_path);
}
}