use std::sync::Arc;
use crate::agent::ToolRegistry;
use crate::config::{AppConfig, ConfigError};
use crate::health::HealthRegistry;
use crate::module::{Module, ModuleBuilder, ModuleError};
use crate::thingd::{MemoryThingdBackend, ThingdBackend};
#[derive(Clone)]
pub struct AppState {
pub config: AppConfig,
pub storage: Arc<dyn ThingdBackend>,
pub tool_registry: Arc<ToolRegistry>,
pub storage_mode: String,
pub thingd_ready: bool,
pub health_registry: Option<Arc<HealthRegistry>>,
}
impl AppState {
pub fn builder() -> AppStateBuilder {
AppStateBuilder::new()
}
}
pub struct AppStateBuilder {
config: Option<AppConfig>,
storage: Option<Arc<dyn ThingdBackend>>,
tool_registry: Option<Arc<ToolRegistry>>,
storage_mode: Option<String>,
thingd_ready: Option<bool>,
health_registry: Option<Arc<HealthRegistry>>,
}
impl AppStateBuilder {
pub fn new() -> Self {
Self {
config: None,
storage: None,
tool_registry: None,
storage_mode: None,
thingd_ready: None,
health_registry: None,
}
}
pub fn with_config(mut self, config: AppConfig) -> Self {
self.config = Some(config);
self
}
pub fn with_storage(mut self, storage: Arc<dyn ThingdBackend>) -> Self {
self.storage = Some(storage);
self
}
pub fn with_tool_registry(mut self, registry: ToolRegistry) -> Self {
self.tool_registry = Some(Arc::new(registry));
self
}
pub fn with_storage_mode(mut self, mode: impl Into<String>) -> Self {
self.storage_mode = Some(mode.into());
self
}
pub fn with_thingd_ready(mut self, ready: bool) -> Self {
self.thingd_ready = Some(ready);
self
}
pub fn with_health_registry(mut self, registry: HealthRegistry) -> Self {
self.health_registry = Some(Arc::new(registry));
self
}
pub fn with_health_registry_arc(mut self, registry: Arc<HealthRegistry>) -> Self {
self.health_registry = Some(registry);
self
}
pub fn with_modules<M: Module + 'static>(
mut self,
modules: Vec<M>,
) -> Result<Self, ModuleError> {
let mut module_builder = ModuleBuilder::new();
for module in modules {
module_builder = module_builder.register(module);
}
module_builder.validate()?;
let mut tools = ToolRegistry::new(
&format!("{}-app", env!("CARGO_PKG_NAME")),
env!("CARGO_PKG_VERSION"),
"An Arqen application",
"memory",
);
let mut health = HealthRegistry::new();
module_builder.register_all(&mut tools, &mut health)?;
self.tool_registry = Some(Arc::new(tools));
self.health_registry = Some(Arc::new(health));
Ok(self)
}
pub fn build(self) -> Result<AppState, ConfigError> {
let config = self.config.unwrap_or_default();
let storage = self
.storage
.unwrap_or_else(|| Arc::new(MemoryThingdBackend::new()));
let tool_registry = self.tool_registry.unwrap_or_else(|| {
Arc::new(ToolRegistry::new(
&format!("{}-app", env!("CARGO_PKG_NAME")),
env!("CARGO_PKG_VERSION"),
"An Arqen application",
&format!("{:?}", config.storage.mode),
))
});
let storage_mode = self
.storage_mode
.unwrap_or_else(|| format!("{:?}", config.storage.mode).to_lowercase());
let thingd_ready = self.thingd_ready.unwrap_or(true);
let health_registry = self.health_registry;
Ok(AppState {
config,
storage,
tool_registry,
storage_mode,
thingd_ready,
health_registry,
})
}
}
impl Default for AppStateBuilder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::StorageMode;
#[test]
fn test_app_state_builder_defaults() {
let state = AppState::builder().build().unwrap();
assert_eq!(state.config.server.port, 8888);
assert_eq!(state.config.storage.mode, StorageMode::Memory);
}
#[test]
fn test_app_state_builder_with_config() {
let config = AppConfig {
server: crate::config::ServerConfig {
port: 8080,
..Default::default()
},
..Default::default()
};
let state = AppState::builder().with_config(config).build().unwrap();
assert_eq!(state.config.server.port, 8080);
}
#[tokio::test]
async fn test_app_state_builder_with_storage() {
let storage = Arc::new(MemoryThingdBackend::new());
let state = AppState::builder().with_storage(storage).build().unwrap();
assert!(state.storage.count_objects("test").await.is_ok());
}
#[tokio::test]
async fn test_app_state_builder_with_registry() {
let registry = ToolRegistry::new("test-app", "1.0.0", "Test", "memory");
let state = AppState::builder()
.with_tool_registry(registry)
.build()
.unwrap();
assert_eq!(state.tool_registry.generate_manifest().name, "test-app");
}
#[test]
fn test_app_state_builder_with_modules() {
use crate::agent::{ToolEffect, ToolMetadata};
use crate::module::{ModuleContext, ModuleHealth};
struct UsersModule;
#[async_trait::async_trait]
impl crate::module::Module for UsersModule {
fn name(&self) -> &str {
"users"
}
fn register(&self, ctx: &mut ModuleContext<'_>) -> Result<(), crate::core::AppError> {
ctx.tools.register_tool(ToolMetadata {
name: "get_user".to_string(),
description: "Get a user by ID".to_string(),
input: serde_json::json!({"type": "object"}),
output: serde_json::json!({"type": "object"}),
scopes: vec!["read:users".to_string()],
effect: ToolEffect::Read,
idempotent: true,
enqueues_job: None,
timeout: None,
});
Ok(())
}
async fn health_check(&self) -> ModuleHealth {
ModuleHealth::Healthy
}
}
let state = AppState::builder()
.with_modules(vec![UsersModule])
.unwrap()
.build()
.unwrap();
assert!(state.tool_registry.get_tool("get_user").is_some());
let health = state.health_registry.unwrap();
let rt = tokio::runtime::Runtime::new().unwrap();
let report = rt.block_on(health.check_liveness());
assert_eq!(report.checks.len(), 1);
assert_eq!(report.checks[0].name, "users");
}
#[test]
fn test_app_state_builder_with_modules_validation_error() {
struct DepModule;
#[async_trait::async_trait]
impl crate::module::Module for DepModule {
fn name(&self) -> &str {
"app"
}
fn dependencies(&self) -> Vec<&str> {
vec!["nonexistent"]
}
}
let result = AppState::builder().with_modules(vec![DepModule]);
assert!(result.is_err());
}
}