autoagents_core/
environment.rs

1use crate::error::Error;
2use crate::protocol::{Event, RuntimeID};
3use crate::runtime::manager::RuntimeManager;
4use crate::runtime::{Runtime, RuntimeError};
5use crate::utils::BoxEventStream;
6use std::path::PathBuf;
7use std::sync::Arc;
8use tokio::task::JoinHandle;
9
10#[derive(Debug, thiserror::Error)]
11pub enum EnvironmentError {
12    #[error("Runtime not found: {0}")]
13    RuntimeNotFound(RuntimeID),
14
15    #[error("Runtime error: {0}")]
16    RuntimeError(#[from] RuntimeError),
17
18    #[error("Error when consuming receiver")]
19    EventError,
20}
21
22#[derive(Clone)]
23pub struct EnvironmentConfig {
24    pub working_dir: PathBuf,
25}
26
27impl Default for EnvironmentConfig {
28    fn default() -> Self {
29        Self {
30            working_dir: std::env::current_dir().unwrap_or_default(),
31        }
32    }
33}
34
35pub struct Environment {
36    config: EnvironmentConfig,
37    runtime_manager: Arc<RuntimeManager>,
38    default_runtime: Option<RuntimeID>,
39    handle: Option<JoinHandle<Result<(), RuntimeError>>>,
40}
41
42impl Environment {
43    pub fn new(config: Option<EnvironmentConfig>) -> Self {
44        let config = config.unwrap_or_default();
45        let runtime_manager = Arc::new(RuntimeManager::new());
46
47        Self {
48            config,
49            runtime_manager,
50            default_runtime: None,
51            handle: None,
52        }
53    }
54
55    pub async fn register_runtime(&mut self, runtime: Arc<dyn Runtime>) -> Result<(), Error> {
56        self.runtime_manager
57            .register_runtime(runtime.clone())
58            .await?;
59        if self.default_runtime.is_none() {
60            self.default_runtime = Some(runtime.id());
61        }
62        Ok(())
63    }
64
65    pub fn config(&self) -> &EnvironmentConfig {
66        &self.config
67    }
68
69    pub async fn get_runtime(&self, runtime_id: &RuntimeID) -> Option<Arc<dyn Runtime>> {
70        self.runtime_manager.get_runtime(runtime_id).await
71    }
72
73    pub async fn get_runtime_or_default(
74        &self,
75        runtime_id: Option<RuntimeID>,
76    ) -> Result<Arc<dyn Runtime>, Error> {
77        let rid = runtime_id.unwrap_or(self.default_runtime.unwrap());
78        self.get_runtime(&rid)
79            .await
80            .ok_or_else(|| EnvironmentError::RuntimeNotFound(rid).into())
81    }
82
83    pub fn run(&mut self) -> JoinHandle<Result<(), RuntimeError>> {
84        let manager = self.runtime_manager.clone();
85        // Spawn background task to run the runtimes.
86        let handle = tokio::spawn(async move { manager.run().await });
87        handle
88    }
89
90    pub async fn take_event_receiver(
91        &mut self,
92        runtime_id: Option<RuntimeID>,
93    ) -> Result<BoxEventStream<Event>, EnvironmentError> {
94        if let Ok(runtime) = self.get_runtime_or_default(runtime_id).await {
95            runtime
96                .take_event_receiver()
97                .await
98                .ok_or_else(|| EnvironmentError::EventError)
99        } else {
100            Err(EnvironmentError::RuntimeNotFound(runtime_id.unwrap()))
101        }
102    }
103
104    pub async fn shutdown(&mut self) {
105        let _ = self.runtime_manager.stop().await;
106
107        if let Some(handle) = self.handle.take() {
108            let _ = handle.await;
109        }
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use crate::runtime::SingleThreadedRuntime;
116
117    use super::*;
118    use uuid::Uuid;
119
120    #[test]
121    fn test_environment_config_default() {
122        let config = EnvironmentConfig::default();
123        assert_eq!(
124            config.working_dir,
125            std::env::current_dir().unwrap_or_default()
126        );
127    }
128
129    #[test]
130    fn test_environment_config_custom() {
131        let config = EnvironmentConfig {
132            working_dir: std::path::PathBuf::from("/tmp"),
133        };
134        assert_eq!(config.working_dir, std::path::PathBuf::from("/tmp"));
135    }
136
137    #[tokio::test]
138    async fn test_environment_get_runtime() {
139        let mut env = Environment::new(None);
140        let runtime = SingleThreadedRuntime::new(None);
141        let runtime_id = runtime.id;
142        env.register_runtime(runtime).await.unwrap();
143
144        // Test getting default runtime
145        let runtime = env.get_runtime(&runtime_id).await;
146
147        assert!(runtime.is_some());
148
149        // Test getting non-existent runtime
150        let non_existent_id = Uuid::new_v4();
151        let runtime = env.get_runtime(&non_existent_id).await;
152        assert!(runtime.is_none());
153    }
154
155    #[tokio::test]
156    async fn test_environment_take_event_receiver() {
157        let mut env = Environment::new(None);
158        let runtime = SingleThreadedRuntime::new(None);
159        let _ = runtime.id;
160        env.register_runtime(runtime).await.unwrap();
161        let receiver = env.take_event_receiver(None).await;
162        assert!(receiver.is_ok());
163
164        // Second call should return None
165        let receiver2 = env.take_event_receiver(None).await;
166        assert!(receiver2.is_err());
167    }
168
169    #[tokio::test]
170    async fn test_environment_shutdown() {
171        let mut env = Environment::new(None);
172        env.shutdown().await;
173        // Should not panic
174    }
175
176    #[tokio::test]
177    async fn test_environment_error_runtime_not_found() {
178        let mut env = Environment::new(None);
179        let runtime = SingleThreadedRuntime::new(None);
180        let _ = runtime.id;
181        env.register_runtime(runtime).await.unwrap();
182        let non_existent_id = Uuid::new_v4();
183
184        let result = env.get_runtime_or_default(Some(non_existent_id)).await;
185        assert!(result.is_err());
186
187        assert!(result.is_err());
188        // Just test that it's an error, not the specific variant
189        assert!(result.is_err());
190    }
191
192    #[test]
193    fn test_environment_error_display() {
194        let runtime_id = Uuid::new_v4();
195        let error = EnvironmentError::RuntimeNotFound(runtime_id);
196        assert!(error.to_string().contains("Runtime not found"));
197        assert!(error.to_string().contains(&runtime_id.to_string()));
198    }
199}
autoagents_core/environment.rs

autoagents_core/
environment.rs
