Skip to main content

walrus_runtime/
lib.rs

1//! Walrus runtime: agent registry and hook orchestration.
2//!
3//! The [`Runtime`] holds agents behind a `RwLock` and provides take/put
4//! semantics for callers to drive execution directly via `Agent::run()`
5//! or `Agent::run_stream()`.
6
7pub use hook::Hook;
8pub use memory::{InMemory, Memory, NoEmbedder};
9pub use wcore::AgentConfig;
10pub use wcore::model::{Message, Request, Response, Role, StreamChunk, Tool};
11
12use anyhow::Result;
13use compact_str::CompactString;
14use std::{collections::BTreeMap, future::Future, sync::Arc};
15use tokio::sync::RwLock;
16
17pub mod hook;
18
19/// Re-exports of the most commonly used types.
20pub mod prelude {
21    pub use crate::{
22        AgentConfig, Hook, InMemory, Message, Request, Response, Role, Runtime, StreamChunk, Tool,
23    };
24}
25
26/// A type-erased async tool handler.
27pub type Handler =
28    Arc<dyn Fn(String) -> std::pin::Pin<Box<dyn Future<Output = String> + Send>> + Send + Sync>;
29
30/// Thin wrapper that implements wcore's `Dispatcher` by forwarding to Hook.
31pub struct AgentDispatcher<'a, H: Hook> {
32    /// The hook backend.
33    pub hook: &'a H,
34    /// The agent name for scoped dispatch.
35    pub agent: &'a str,
36}
37
38impl<H: Hook> wcore::Dispatcher for AgentDispatcher<'_, H> {
39    fn dispatch(&self, calls: &[(&str, &str)]) -> impl Future<Output = Vec<Result<String>>> + Send {
40        self.hook.dispatch(self.agent, calls)
41    }
42
43    fn tools(&self) -> Vec<Tool> {
44        self.hook.tools(self.agent)
45    }
46}
47
48/// The walrus runtime — agent registry and hook orchestration.
49///
50/// Generic over `H: Hook` for the runtime backend. Stores agents (each
51/// holding its own model clone, config, and history). Callers take agents
52/// out for execution and put them back when done.
53pub struct Runtime<H: Hook> {
54    hook: Arc<H>,
55    agents: RwLock<BTreeMap<CompactString, wcore::Agent<H::Model>>>,
56}
57
58impl<H: Hook + 'static> Runtime<H> {
59    /// Create a new runtime with the given hook backend.
60    pub fn new(hook: Arc<H>) -> Self {
61        Self {
62            hook,
63            agents: RwLock::new(BTreeMap::new()),
64        }
65    }
66
67    /// Access the hook backend.
68    pub fn hook(&self) -> &H {
69        &self.hook
70    }
71
72    /// Register an agent from its configuration.
73    ///
74    /// Clones the Hook's model into the agent so it can drive LLM calls.
75    pub async fn add_agent(&self, config: AgentConfig) {
76        let name = config.name.clone();
77        let agent = wcore::AgentBuilder::new(self.hook.model().clone())
78            .config(config)
79            .build();
80        self.agents.write().await.insert(name, agent);
81    }
82
83    /// Get a registered agent's config by name (cloned).
84    pub async fn agent(&self, name: &str) -> Option<AgentConfig> {
85        self.agents.read().await.get(name).map(|a| a.config.clone())
86    }
87
88    /// Get all registered agent configs (cloned, alphabetical order).
89    pub async fn agents(&self) -> Vec<AgentConfig> {
90        self.agents
91            .read()
92            .await
93            .values()
94            .map(|a| a.config.clone())
95            .collect()
96    }
97
98    /// Take an agent out of the registry for execution.
99    ///
100    /// The agent is removed from the map. Caller must call [`put_agent`]
101    /// to re-insert it after execution completes.
102    pub async fn take_agent(&self, name: &str) -> Option<wcore::Agent<H::Model>> {
103        self.agents.write().await.remove(name)
104    }
105
106    /// Put an agent back into the registry after execution.
107    pub async fn put_agent(&self, agent: wcore::Agent<H::Model>) {
108        let name = agent.config.name.clone();
109        self.agents.write().await.insert(name, agent);
110    }
111
112    /// Clear the conversation history for a named agent.
113    pub async fn clear_session(&self, agent: &str) {
114        if let Some(a) = self.agents.write().await.get_mut(agent) {
115            a.clear_history();
116        }
117    }
118}