use crate::error::Error;
use crate::runtime::manager::RuntimeManager;
use crate::runtime::{Runtime, RuntimeError};
use crate::utils::BoxEventStream;
use autoagents_protocol::{Event, RuntimeID};
use futures_util::FutureExt;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::task::JoinHandle;
#[derive(Debug, thiserror::Error)]
pub enum EnvironmentError {
#[error("Runtime not found: {0}")]
RuntimeNotFound(RuntimeID),
#[error("No default runtime registered")]
NoDefaultRuntime,
#[error("Environment is already running")]
AlreadyRunning,
#[error("Runtime error: {0}")]
RuntimeError(#[from] Box<RuntimeError>),
#[error("Error when consuming receiver")]
EventError,
#[error("Run task join error: {0}")]
JoinError(#[from] tokio::task::JoinError),
#[error("Finished run task result was not yet available")]
RunResultNotReady,
}
#[derive(Clone)]
pub struct EnvironmentConfig {
pub working_dir: PathBuf,
}
impl Default for EnvironmentConfig {
fn default() -> Self {
Self {
working_dir: std::env::current_dir().unwrap_or_default(),
}
}
}
pub struct Environment {
config: EnvironmentConfig,
runtime_manager: Arc<RuntimeManager>,
default_runtime: Option<RuntimeID>,
handle: Option<JoinHandle<Result<(), RuntimeError>>>,
launch_state: RuntimeLaunchState,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
enum RuntimeLaunchState {
#[default]
Idle,
Managed,
Background,
}
impl Environment {
pub fn new(config: Option<EnvironmentConfig>) -> Self {
let config = config.unwrap_or_default();
let runtime_manager = Arc::new(RuntimeManager::new());
Self {
config,
runtime_manager,
default_runtime: None,
handle: None,
launch_state: RuntimeLaunchState::Idle,
}
}
pub async fn register_runtime(&mut self, runtime: Arc<dyn Runtime>) -> Result<(), Error> {
self.runtime_manager
.register_runtime(runtime.clone())
.await?;
if self.default_runtime.is_none() {
self.default_runtime = Some(runtime.id());
}
Ok(())
}
pub fn config(&self) -> &EnvironmentConfig {
&self.config
}
pub async fn get_runtime(&self, runtime_id: &RuntimeID) -> Option<Arc<dyn Runtime>> {
self.runtime_manager.get_runtime(runtime_id).await
}
pub async fn get_runtime_or_default(
&self,
runtime_id: Option<RuntimeID>,
) -> Result<Arc<dyn Runtime>, Error> {
let rid = match runtime_id {
Some(id) => id,
None => self
.default_runtime
.ok_or(EnvironmentError::NoDefaultRuntime)?,
};
self.get_runtime(&rid)
.await
.ok_or_else(|| EnvironmentError::RuntimeNotFound(rid).into())
}
#[allow(clippy::result_large_err)] pub fn run(&mut self) -> Result<(), EnvironmentError> {
self.reconcile_finished_managed_launch()?;
if self.launch_state == RuntimeLaunchState::Background {
return Err(EnvironmentError::AlreadyRunning);
}
if self.is_running() {
return Err(EnvironmentError::AlreadyRunning);
}
let manager = self.runtime_manager.clone();
let handle = tokio::spawn(async move { manager.run().await });
self.handle = Some(handle);
self.launch_state = RuntimeLaunchState::Managed;
Ok(())
}
pub async fn wait(&mut self) -> Result<Result<(), RuntimeError>, tokio::task::JoinError> {
let handle = match self.handle.take() {
Some(handle) => handle,
None => return Ok(Ok(())),
};
let mut guard = RestoreRunHandleOnDrop {
environment: self,
handle: Some(handle),
};
let join_result = guard.handle.as_mut().expect("handle was just set").await;
guard.handle = None;
guard.environment.launch_state = RuntimeLaunchState::Idle;
join_result
}
pub async fn run_background(&mut self) -> Result<(), EnvironmentError> {
self.reconcile_finished_managed_launch()?;
if self.launch_state != RuntimeLaunchState::Idle || self.is_running() {
return Err(EnvironmentError::AlreadyRunning);
}
let manager = self.runtime_manager.clone();
manager
.run_background()
.await
.map_err(|e| EnvironmentError::RuntimeError(Box::new(e)))?;
self.launch_state = RuntimeLaunchState::Background;
Ok(())
}
pub async fn take_event_receiver(
&mut self,
runtime_id: Option<RuntimeID>,
) -> Result<BoxEventStream<Event>, EnvironmentError> {
let runtime = self
.get_runtime_or_default(runtime_id)
.await
.map_err(|err| match err {
Error::EnvironmentError(env_err) => env_err,
_ => EnvironmentError::EventError,
})?;
runtime
.take_event_receiver()
.await
.ok_or(EnvironmentError::EventError)
}
pub async fn subscribe_events(
&self,
runtime_id: Option<RuntimeID>,
) -> Result<BoxEventStream<Event>, EnvironmentError> {
let runtime = self
.get_runtime_or_default(runtime_id)
.await
.map_err(|err| match err {
Error::EnvironmentError(env_err) => env_err,
_ => EnvironmentError::EventError,
})?;
Ok(runtime.subscribe_events().await)
}
pub async fn shutdown(&mut self) -> Result<(), EnvironmentError> {
let stop_result = self.runtime_manager.stop().await;
let join_result = if let Some(handle) = self.handle.take() {
Some(handle.await)
} else {
None
};
self.launch_state = RuntimeLaunchState::Idle;
if let Err(e) = stop_result {
return Err(EnvironmentError::RuntimeError(Box::new(e)));
}
match join_result {
None | Some(Ok(Ok(()))) => Ok(()),
Some(Ok(Err(e))) => Err(EnvironmentError::RuntimeError(Box::new(e))),
Some(Err(e)) => Err(EnvironmentError::JoinError(e)),
}
}
pub fn is_running(&self) -> bool {
match self.launch_state {
RuntimeLaunchState::Background => true,
RuntimeLaunchState::Managed => self
.handle
.as_ref()
.is_some_and(|handle| !handle.is_finished()),
RuntimeLaunchState::Idle => false,
}
}
#[allow(clippy::result_large_err)]
fn reconcile_finished_managed_launch(&mut self) -> Result<(), EnvironmentError> {
if self.launch_state != RuntimeLaunchState::Managed {
return Ok(());
}
let Some(handle) = self.handle.take() else {
self.launch_state = RuntimeLaunchState::Idle;
return Ok(());
};
if !handle.is_finished() {
self.handle = Some(handle);
return Ok(());
}
self.launch_state = RuntimeLaunchState::Idle;
Self::join_finished_handle(handle)
}
#[allow(clippy::result_large_err)]
fn join_finished_handle(
handle: JoinHandle<Result<(), RuntimeError>>,
) -> Result<(), EnvironmentError> {
debug_assert!(
handle.is_finished(),
"join_finished_handle requires a finished run task"
);
match handle.now_or_never() {
Some(Ok(Ok(()))) => Ok(()),
Some(Ok(Err(e))) => Err(EnvironmentError::RuntimeError(Box::new(e))),
Some(Err(e)) => Err(EnvironmentError::JoinError(e)),
None => Err(EnvironmentError::RunResultNotReady),
}
}
}
struct RestoreRunHandleOnDrop<'a> {
environment: &'a mut Environment,
handle: Option<JoinHandle<Result<(), RuntimeError>>>,
}
impl Drop for RestoreRunHandleOnDrop<'_> {
fn drop(&mut self) {
if let Some(handle) = self.handle.take() {
self.environment.handle = Some(handle);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::runtime::SingleThreadedRuntime;
use tempfile::tempdir;
use tokio::sync::mpsc;
use uuid::Uuid;
#[test]
fn test_environment_config_default() {
let config = EnvironmentConfig::default();
assert_eq!(
config.working_dir,
std::env::current_dir().unwrap_or_default()
);
}
#[test]
fn test_environment_config_custom() {
let dir = tempdir().expect("Unable to create temp dir");
let config = EnvironmentConfig {
working_dir: dir.path().to_path_buf(),
};
assert_eq!(config.working_dir, dir.path().to_path_buf());
}
#[tokio::test]
async fn test_environment_get_runtime() {
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
let runtime_id = runtime.id;
env.register_runtime(runtime).await.unwrap();
let runtime = env.get_runtime(&runtime_id).await;
assert!(runtime.is_some());
let non_existent_id = Uuid::new_v4();
let runtime = env.get_runtime(&non_existent_id).await;
assert!(runtime.is_none());
}
#[tokio::test]
async fn test_environment_take_event_receiver() {
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
let _ = runtime.id;
env.register_runtime(runtime).await.unwrap();
let receiver = env.take_event_receiver(None).await;
assert!(receiver.is_ok());
let receiver2 = env.take_event_receiver(None).await;
assert!(receiver2.is_err());
}
#[tokio::test]
async fn test_environment_shutdown() {
let mut env = Environment::new(None);
env.shutdown()
.await
.expect("shutdown should succeed when idle");
}
#[tokio::test]
async fn test_environment_error_runtime_not_found() {
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
let _ = runtime.id;
env.register_runtime(runtime).await.unwrap();
let non_existent_id = Uuid::new_v4();
let result = env.get_runtime_or_default(Some(non_existent_id)).await;
assert!(matches!(
result,
Err(Error::EnvironmentError(EnvironmentError::RuntimeNotFound(id)))
if id == non_existent_id
));
}
#[test]
fn test_environment_error_already_running_display() {
let error = EnvironmentError::AlreadyRunning;
assert!(error.to_string().contains("already running"));
}
#[test]
fn test_environment_error_run_result_not_ready_display() {
let error = EnvironmentError::RunResultNotReady;
assert!(error.to_string().contains("not yet available"));
}
#[test]
fn test_environment_error_display() {
let runtime_id = Uuid::new_v4();
let error = EnvironmentError::RuntimeNotFound(runtime_id);
assert!(error.to_string().contains("Runtime not found"));
assert!(error.to_string().contains(&runtime_id.to_string()));
}
#[test]
fn test_environment_error_no_default_display() {
let error = EnvironmentError::NoDefaultRuntime;
assert!(error.to_string().contains("No default runtime registered"));
}
#[tokio::test]
async fn test_get_runtime_or_default_no_default_runtime() {
let env = Environment::new(None);
let result = env.get_runtime_or_default(None).await;
assert!(matches!(
result,
Err(Error::EnvironmentError(EnvironmentError::NoDefaultRuntime))
));
}
#[tokio::test]
async fn test_take_event_receiver_no_default_runtime() {
let mut env = Environment::new(None);
let result = env.take_event_receiver(None).await;
assert!(matches!(result, Err(EnvironmentError::NoDefaultRuntime)));
}
#[tokio::test]
async fn test_subscribe_events_no_default_runtime() {
let env = Environment::new(None);
let result = env.subscribe_events(None).await;
assert!(matches!(result, Err(EnvironmentError::NoDefaultRuntime)));
}
#[tokio::test]
async fn test_environment_run_stores_handle() {
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
env.register_runtime(runtime).await.unwrap();
assert!(!env.is_running());
env.run().expect("run should succeed");
assert!(env.is_running());
assert!(matches!(env.run(), Err(EnvironmentError::AlreadyRunning)));
env.shutdown().await.expect("shutdown should succeed");
assert!(!env.is_running());
}
#[tokio::test]
async fn test_environment_run_can_restart_after_shutdown() {
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
env.register_runtime(runtime).await.unwrap();
env.run().expect("initial run should succeed");
env.shutdown().await.expect("shutdown should succeed");
env.run().expect("run after shutdown should succeed");
let run_result = env.wait().await.expect("wait should join run task");
assert!(run_result.is_err());
env.shutdown()
.await
.expect("shutdown should succeed when idle after failed run");
}
#[tokio::test]
async fn test_get_runtime_or_default_runtime_not_found_variant() {
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
env.register_runtime(runtime).await.unwrap();
let non_existent_id = Uuid::new_v4();
let result = env.get_runtime_or_default(Some(non_existent_id)).await;
assert!(matches!(
result,
Err(Error::EnvironmentError(EnvironmentError::RuntimeNotFound(id)))
if id == non_existent_id
));
}
#[tokio::test]
async fn test_take_event_receiver_runtime_not_found() {
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
env.register_runtime(runtime).await.unwrap();
let non_existent_id = Uuid::new_v4();
let result = env.take_event_receiver(Some(non_existent_id)).await;
assert!(matches!(
result,
Err(EnvironmentError::RuntimeNotFound(id)) if id == non_existent_id
));
}
#[tokio::test]
async fn test_environment_wait_restores_handle_when_cancelled() {
use tokio::time::{Duration, sleep};
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
env.register_runtime(runtime).await.unwrap();
env.run().expect("run should succeed");
assert!(env.is_running());
{
let wait_fut = env.wait();
tokio::pin!(wait_fut);
tokio::select! {
_ = &mut wait_fut => panic!("run task should not finish immediately"),
_ = sleep(Duration::from_millis(10)) => {}
}
}
assert!(
env.is_running(),
"cancelled wait should restore the join handle"
);
env.shutdown()
.await
.expect("shutdown should join the restored handle");
assert!(!env.is_running());
}
#[tokio::test]
async fn test_environment_wait_is_idempotent() {
use tokio::time::{Duration, timeout};
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
env.register_runtime(runtime).await.unwrap();
env.run().expect("run should succeed");
env.shutdown().await.expect("shutdown should succeed");
for _ in 0..2 {
let result = timeout(Duration::from_secs(1), env.wait())
.await
.expect("wait should not hang");
assert!(result.is_ok());
}
}
#[derive(Clone, Copy)]
enum ImmediateRuntimeBehavior {
Success,
Error,
}
struct ImmediateRuntime {
id: RuntimeID,
behavior: ImmediateRuntimeBehavior,
tx: mpsc::Sender<Event>,
}
#[async_trait::async_trait]
impl Runtime for ImmediateRuntime {
fn id(&self) -> RuntimeID {
self.id
}
async fn subscribe_any(
&self,
_topic_name: &str,
_topic_type: std::any::TypeId,
_actor: Arc<dyn crate::actor::AnyActor>,
) -> Result<(), RuntimeError> {
Ok(())
}
async fn publish_any(
&self,
_topic_name: &str,
_topic_type: std::any::TypeId,
_message: Arc<dyn std::any::Any + Send + Sync>,
) -> Result<(), RuntimeError> {
Ok(())
}
fn tx(&self) -> mpsc::Sender<Event> {
self.tx.clone()
}
async fn transport(&self) -> Arc<dyn crate::actor::Transport> {
Arc::new(crate::actor::LocalTransport)
}
async fn take_event_receiver(&self) -> Option<BoxEventStream<Event>> {
None
}
async fn subscribe_events(&self) -> BoxEventStream<Event> {
Box::pin(futures::stream::empty())
}
async fn run(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
match self.behavior {
ImmediateRuntimeBehavior::Success => Ok(()),
ImmediateRuntimeBehavior::Error => {
Err(std::io::Error::other("immediate runtime run failed").into())
}
}
}
async fn stop(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
Ok(())
}
}
#[tokio::test]
async fn test_environment_wait_awaits_completed_run_task() {
let mut env = Environment::new(None);
let (tx, _rx) = mpsc::channel(1);
let runtime = Arc::new(ImmediateRuntime {
id: RuntimeID::new_v4(),
behavior: ImmediateRuntimeBehavior::Success,
tx,
}) as Arc<dyn Runtime>;
env.register_runtime(runtime).await.unwrap();
env.run().expect("run should succeed");
let wait_result = env.wait().await.expect("wait join should succeed");
assert!(wait_result.is_ok());
}
#[tokio::test]
async fn test_environment_run_restarts_after_finished_handle() {
use tokio::time::{Duration, sleep};
let mut env = Environment::new(None);
let (tx, _rx) = mpsc::channel(1);
let runtime = Arc::new(ImmediateRuntime {
id: RuntimeID::new_v4(),
behavior: ImmediateRuntimeBehavior::Success,
tx,
}) as Arc<dyn Runtime>;
env.register_runtime(runtime).await.unwrap();
env.run().expect("initial run should succeed");
sleep(Duration::from_millis(20)).await;
assert!(!env.is_running());
env.run().expect("run after finished handle should succeed");
env.shutdown().await.expect("shutdown should succeed");
}
#[tokio::test]
async fn test_environment_run_background_starts_runtimes() {
let mut env = Environment::new(None);
let (tx, _rx) = mpsc::channel(1);
let runtime = Arc::new(ImmediateRuntime {
id: RuntimeID::new_v4(),
behavior: ImmediateRuntimeBehavior::Success,
tx,
}) as Arc<dyn Runtime>;
env.register_runtime(runtime).await.unwrap();
env.run_background()
.await
.expect("run_background should succeed");
assert!(
env.is_running(),
"run_background should mark the environment as running until shutdown"
);
env.shutdown().await.expect("shutdown should succeed");
}
#[tokio::test]
async fn test_environment_run_background_rejects_after_run() {
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
env.register_runtime(runtime).await.unwrap();
env.run().expect("run should succeed");
assert!(matches!(
env.run_background().await,
Err(EnvironmentError::AlreadyRunning)
));
env.shutdown().await.expect("shutdown should succeed");
}
#[tokio::test]
async fn test_environment_run_rejects_after_run_background() {
let mut env = Environment::new(None);
let (tx, _rx) = mpsc::channel(1);
let runtime = Arc::new(ImmediateRuntime {
id: RuntimeID::new_v4(),
behavior: ImmediateRuntimeBehavior::Success,
tx,
}) as Arc<dyn Runtime>;
env.register_runtime(runtime).await.unwrap();
env.run_background()
.await
.expect("run_background should succeed");
assert!(matches!(env.run(), Err(EnvironmentError::AlreadyRunning)));
env.shutdown().await.expect("shutdown should succeed");
}
#[tokio::test]
async fn test_environment_wait_propagates_run_failure() {
let mut env = Environment::new(None);
let (tx, _rx) = mpsc::channel(1);
let runtime = Arc::new(ImmediateRuntime {
id: RuntimeID::new_v4(),
behavior: ImmediateRuntimeBehavior::Error,
tx,
}) as Arc<dyn Runtime>;
env.register_runtime(runtime).await.unwrap();
env.run().expect("run should succeed");
let wait_result = env.wait().await.expect("wait join should succeed");
assert!(wait_result.is_err());
}
#[tokio::test]
async fn test_environment_run_surfaces_prior_failure_without_wait() {
use tokio::time::{Duration, sleep};
let mut env = Environment::new(None);
let (tx, _rx) = mpsc::channel(1);
let runtime = Arc::new(ImmediateRuntime {
id: RuntimeID::new_v4(),
behavior: ImmediateRuntimeBehavior::Error,
tx,
}) as Arc<dyn Runtime>;
env.register_runtime(runtime).await.unwrap();
env.run().expect("initial run should succeed");
sleep(Duration::from_millis(20)).await;
assert!(!env.is_running());
let err = env
.run()
.expect_err("restart should surface prior run failure");
assert!(matches!(err, EnvironmentError::RuntimeError(_)));
}
#[tokio::test]
async fn test_environment_run_background_surfaces_prior_failure_without_wait() {
use tokio::time::{Duration, sleep};
let mut env = Environment::new(None);
let (tx, _rx) = mpsc::channel(1);
let runtime = Arc::new(ImmediateRuntime {
id: RuntimeID::new_v4(),
behavior: ImmediateRuntimeBehavior::Error,
tx,
}) as Arc<dyn Runtime>;
env.register_runtime(runtime).await.unwrap();
env.run().expect("initial run should succeed");
sleep(Duration::from_millis(20)).await;
assert!(!env.is_running());
let err = env
.run_background()
.await
.expect_err("run_background should surface prior run failure");
assert!(matches!(err, EnvironmentError::RuntimeError(_)));
}
#[test]
fn test_environment_config_accessor() {
let dir = tempdir().expect("Unable to create temp dir");
let config = EnvironmentConfig {
working_dir: dir.path().to_path_buf(),
};
let env = Environment::new(Some(config.clone()));
assert_eq!(env.config().working_dir, config.working_dir);
}
#[tokio::test]
async fn test_subscribe_events_with_default_runtime() {
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
env.register_runtime(runtime).await.unwrap();
let stream = env.subscribe_events(None).await;
assert!(stream.is_ok());
}
#[tokio::test]
async fn test_get_runtime_or_default_uses_default_runtime() {
let mut env = Environment::new(None);
let runtime = SingleThreadedRuntime::new(None);
let runtime_id = runtime.id;
env.register_runtime(runtime).await.unwrap();
let resolved = env
.get_runtime_or_default(None)
.await
.expect("default runtime should resolve");
assert_eq!(resolved.id(), runtime_id);
}
}