1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
//! `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;
/// 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,
}
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(),
}
}
}
/// 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>;
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_utils::{
noop_bus, FakeLlmClient, FakeToolInvoker, InMemoryEpisodic, InMemoryLongTerm,
InMemoryShortTerm,
};
/// 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 {
let (pubsub, request_reply, kv, jobs) = noop_bus();
AgentContext {
llm: Arc::new(FakeLlmClient::new("fake")),
short_term: Arc::new(InMemoryShortTerm::default()),
long_term: Arc::new(InMemoryLongTerm::default()),
episodic: Arc::new(InMemoryEpisodic::default()),
pubsub,
kv,
request_reply,
jobs,
tools: Arc::new(FakeToolInvoker::new()),
run_id: RunId::new(),
cancel: CancellationToken::new(),
agent_name: "parent".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);
}
#[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));
}
}