chasm_cli/agency/
runtime.rs

1// Copyright (c) 2024-2026 Nervosys LLC
2// SPDX-License-Identifier: Apache-2.0
3//! Agency Runtime
4//!
5//! High-level API for running agents with automatic session management.
6
7#![allow(dead_code)]
8
9use crate::agency::agent::Agent;
10use crate::agency::error::{AgencyError, AgencyResult};
11use crate::agency::executor::{ExecutionContext, ExecutionResult, Executor};
12use crate::agency::models::AgencyEvent;
13use crate::agency::orchestrator::{Orchestrator, OrchestratorResult, Pipeline, Swarm};
14use crate::agency::session::{Session, SessionManager};
15use crate::agency::tools::ToolRegistry;
16use serde::{Deserialize, Serialize};
17use std::collections::HashMap;
18use std::path::PathBuf;
19use std::sync::Arc;
20use tokio::sync::mpsc;
21
22/// Runtime configuration
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct RuntimeConfig {
25    /// Database path for session storage
26    pub db_path: PathBuf,
27    /// Default model to use
28    pub default_model: String,
29    /// Maximum tool calls per turn
30    pub max_tool_calls: u32,
31    /// Request timeout in seconds
32    pub timeout_seconds: u64,
33    /// Enable streaming by default
34    pub streaming: bool,
35    /// API keys for providers (provider -> key)
36    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
37    pub api_keys: HashMap<String, String>,
38}
39
40impl Default for RuntimeConfig {
41    fn default() -> Self {
42        Self {
43            db_path: PathBuf::from("Agency_sessions.db"),
44            default_model: "gemini-2.5-flash".to_string(),
45            max_tool_calls: 10,
46            timeout_seconds: 120,
47            streaming: true,
48            api_keys: HashMap::new(),
49        }
50    }
51}
52
53/// The Agency Runtime - main entry point for running agents
54pub struct Runtime {
55    config: RuntimeConfig,
56    tool_registry: Arc<ToolRegistry>,
57    session_manager: Arc<SessionManager>,
58    executor: Arc<Executor>,
59    orchestrator: Orchestrator,
60    agents: HashMap<String, Arc<Agent>>,
61}
62
63impl Runtime {
64    /// Create a new runtime with default configuration
65    pub fn new() -> AgencyResult<Self> {
66        Self::with_config(RuntimeConfig::default())
67    }
68
69    /// Create a new runtime with custom configuration
70    pub fn with_config(config: RuntimeConfig) -> AgencyResult<Self> {
71        let tool_registry = Arc::new(ToolRegistry::with_builtins());
72        let session_manager = Arc::new(SessionManager::new(&config.db_path)?);
73        let executor = Arc::new(Executor::new(tool_registry.clone()));
74        let orchestrator = Orchestrator::new(executor.clone());
75
76        Ok(Self {
77            config,
78            tool_registry,
79            session_manager,
80            executor,
81            orchestrator,
82            agents: HashMap::new(),
83        })
84    }
85
86    /// Create an in-memory runtime (for testing)
87    pub fn in_memory() -> AgencyResult<Self> {
88        let tool_registry = Arc::new(ToolRegistry::with_builtins());
89        let session_manager = Arc::new(SessionManager::in_memory()?);
90        let executor = Arc::new(Executor::new(tool_registry.clone()));
91        let orchestrator = Orchestrator::new(executor.clone());
92
93        Ok(Self {
94            config: RuntimeConfig::default(),
95            tool_registry,
96            session_manager,
97            executor,
98            orchestrator,
99            agents: HashMap::new(),
100        })
101    }
102
103    /// Register an agent
104    pub fn register_agent(&mut self, agent: Agent) {
105        self.agents
106            .insert(agent.name().to_string(), Arc::new(agent));
107    }
108
109    /// Get a registered agent
110    pub fn get_agent(&self, name: &str) -> Option<&Arc<Agent>> {
111        self.agents.get(name)
112    }
113
114    /// List registered agents
115    pub fn list_agents(&self) -> Vec<&str> {
116        self.agents.keys().map(|s| s.as_str()).collect()
117    }
118
119    /// Get the tool registry
120    pub fn tools(&self) -> &ToolRegistry {
121        &self.tool_registry
122    }
123
124    /// Get the session manager
125    pub fn sessions(&self) -> &SessionManager {
126        &self.session_manager
127    }
128
129    /// Run an agent with a message
130    pub async fn run(
131        &self,
132        agent_name: &str,
133        message: &str,
134        options: Option<RunOptions>,
135    ) -> AgencyResult<ExecutionResult> {
136        let options = options.unwrap_or_default();
137
138        // Get or create agent
139        let agent_arc = self
140            .agents
141            .get(agent_name)
142            .ok_or_else(|| AgencyError::AgentNotFound(agent_name.to_string()))?;
143
144        // Get or create session
145        let mut session = if let Some(session_id) = &options.session_id {
146            self.session_manager
147                .get(session_id)?
148                .ok_or_else(|| AgencyError::SessionNotFound(session_id.clone()))?
149        } else {
150            self.session_manager
151                .create(agent_name, options.user_id.clone())?
152        };
153
154        // Create execution context
155        let mut ctx = ExecutionContext::new(&session);
156        ctx.user_id = options.user_id;
157        ctx.allow_tools = options.allow_tools;
158        ctx.max_tool_calls = options.max_tool_calls.unwrap_or(self.config.max_tool_calls);
159        ctx.event_sender = options.event_sender;
160
161        // Execute
162        let result = self
163            .executor
164            .execute(agent_arc.as_ref(), &mut session, message, &mut ctx)
165            .await?;
166
167        // Save session
168        self.session_manager.save(&session)?;
169
170        Ok(result)
171    }
172
173    /// Run an agent with streaming events
174    pub async fn run_stream(
175        &self,
176        agent_name: &str,
177        message: &str,
178        options: Option<RunOptions>,
179    ) -> AgencyResult<(ExecutionResult, mpsc::Receiver<AgencyEvent>)> {
180        let (tx, rx) = mpsc::channel(100);
181        let mut options = options.unwrap_or_default();
182        options.event_sender = Some(tx);
183
184        let result = self.run(agent_name, message, Some(options)).await?;
185        Ok((result, rx))
186    }
187
188    /// Run a pipeline
189    pub async fn run_pipeline(
190        &self,
191        pipeline: &Pipeline,
192        input: &str,
193        options: Option<RunOptions>,
194    ) -> AgencyResult<OrchestratorResult> {
195        let options = options.unwrap_or_default();
196
197        let session = Session::new(&pipeline.name, options.user_id.clone());
198        let mut ctx = ExecutionContext::new(&session);
199        ctx.user_id = options.user_id;
200        ctx.allow_tools = options.allow_tools;
201        ctx.event_sender = options.event_sender;
202
203        self.orchestrator
204            .run_pipeline(pipeline, input, &mut ctx)
205            .await
206    }
207
208    /// Run a swarm
209    pub async fn run_swarm(
210        &self,
211        swarm: &Swarm,
212        input: &str,
213        options: Option<RunOptions>,
214    ) -> AgencyResult<OrchestratorResult> {
215        let options = options.unwrap_or_default();
216
217        let session = Session::new(&swarm.name, options.user_id.clone());
218        let mut ctx = ExecutionContext::new(&session);
219        ctx.user_id = options.user_id;
220        ctx.allow_tools = options.allow_tools;
221        ctx.event_sender = options.event_sender;
222
223        self.orchestrator.run_swarm(swarm, input, &mut ctx).await
224    }
225
226    /// Create a new session for an agent
227    pub fn create_session(
228        &self,
229        agent_name: &str,
230        user_id: Option<String>,
231    ) -> AgencyResult<Session> {
232        self.session_manager.create(agent_name, user_id)
233    }
234
235    /// Get a session by ID
236    pub fn get_session(&self, session_id: &str) -> AgencyResult<Option<Session>> {
237        self.session_manager.get(session_id)
238    }
239
240    /// List sessions for an agent
241    pub fn list_sessions(
242        &self,
243        agent_name: &str,
244        limit: Option<u32>,
245    ) -> AgencyResult<Vec<Session>> {
246        self.session_manager.list_by_agent(agent_name, limit)
247    }
248
249    /// Delete a session
250    pub fn delete_session(&self, session_id: &str) -> AgencyResult<bool> {
251        self.session_manager.delete(session_id)
252    }
253}
254
255/// Options for running an agent
256#[derive(Debug, Clone, Default)]
257pub struct RunOptions {
258    /// Session ID to continue (creates new if not provided)
259    pub session_id: Option<String>,
260    /// User ID
261    pub user_id: Option<String>,
262    /// Allow tool execution
263    pub allow_tools: bool,
264    /// Maximum tool calls
265    pub max_tool_calls: Option<u32>,
266    /// Event sender for streaming
267    pub event_sender: Option<mpsc::Sender<AgencyEvent>>,
268}
269
270impl RunOptions {
271    pub fn new() -> Self {
272        Self {
273            allow_tools: true,
274            ..Default::default()
275        }
276    }
277
278    pub fn with_session(mut self, session_id: impl Into<String>) -> Self {
279        self.session_id = Some(session_id.into());
280        self
281    }
282
283    pub fn with_user(mut self, user_id: impl Into<String>) -> Self {
284        self.user_id = Some(user_id.into());
285        self
286    }
287
288    pub fn with_tools(mut self, allow: bool) -> Self {
289        self.allow_tools = allow;
290        self
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297    use crate::agency::agent::AgentBuilder;
298
299    #[tokio::test]
300    #[ignore = "Integration test - requires API credentials"]
301    async fn test_runtime() -> AgencyResult<()> {
302        let mut runtime = Runtime::in_memory()?;
303
304        let agent = AgentBuilder::new("assistant")
305            .description("A helpful assistant")
306            .instruction("You are a helpful AI assistant.")
307            .model("gemini-2.5-flash")
308            .build();
309
310        runtime.register_agent(agent);
311
312        let result = runtime.run("assistant", "Hello!", None).await?;
313        assert!(result.success);
314        assert!(!result.response.is_empty());
315
316        Ok(())
317    }
318
319    #[tokio::test]
320    #[ignore = "Integration test - requires API credentials"]
321    async fn test_runtime_sessions() -> AgencyResult<()> {
322        let mut runtime = Runtime::in_memory()?;
323
324        let agent = AgentBuilder::new("test_agent")
325            .instruction("You are a test agent.")
326            .build();
327
328        runtime.register_agent(agent);
329
330        // Create session
331        let session = runtime.create_session("test_agent", Some("user1".to_string()))?;
332        assert_eq!(session.agent_name, "test_agent");
333
334        // Run with session
335        let options = RunOptions::new()
336            .with_session(&session.id)
337            .with_user("user1");
338
339        let result = runtime
340            .run("test_agent", "Test message", Some(options))
341            .await?;
342        assert!(result.success);
343
344        // Verify session was updated
345        let updated_session = runtime.get_session(&session.id)?.unwrap();
346        assert!(!updated_session.messages.is_empty());
347
348        Ok(())
349    }
350}
chasm_cli/agency/runtime.rs

chasm_cli/agency/
runtime.rs
