use crate::events::{EventLog, EventSink, InMemoryEventLog, NoopEventSink};
use crate::in_memory::{InMemorySessionStorageStore, InMemorySessionStore};
use async_trait::async_trait;
use everruns_core::agent::Agent;
use everruns_core::error::Result;
use everruns_core::harness::Harness;
use everruns_core::in_memory::{InMemoryAgentStore, InMemoryHarnessStore, InMemoryProviderStore};
use everruns_core::session::Session;
use everruns_core::session_task::SessionTaskRegistry;
use everruns_core::traits::{
AgentStore, HarnessStore, ProviderStore, ResolvedModel, SessionMutator, SessionScheduleStore,
SessionStorageStore, SessionStore, UserConnectionResolver,
};
use everruns_core::typed_id::SessionId;
use everruns_platform::PlatformStore;
use std::sync::Arc;
pub type ScheduleStoreFactory = Arc<dyn Fn(i64) -> Arc<dyn SessionScheduleStore> + Send + Sync>;
pub type PlatformStoreFactory = Arc<dyn Fn(i64, SessionId) -> Arc<dyn PlatformStore> + Send + Sync>;
#[async_trait]
pub trait RuntimeAgentStore: AgentStore + Send + Sync {
async fn add_agent(&self, agent: Agent) -> Result<()>;
}
#[async_trait]
pub trait RuntimeHarnessStore: HarnessStore + Send + Sync {
async fn add_harness(&self, harness: Harness) -> Result<()>;
}
#[async_trait]
pub trait RuntimeSessionStore: SessionStore + SessionMutator + Send + Sync {
async fn add_session(&self, session: Session) -> Result<()>;
}
#[async_trait]
pub trait RuntimeProviderStore: ProviderStore + Send + Sync {
async fn set_default_model(&self, model: ResolvedModel) -> Result<()>;
}
#[derive(Clone)]
pub struct HostBackends {
pub harness_store: Arc<dyn RuntimeHarnessStore>,
pub agent_store: Arc<dyn RuntimeAgentStore>,
pub session_store: Arc<dyn RuntimeSessionStore>,
pub event_log: Arc<dyn EventLog>,
pub compaction_checkpoint_store: Arc<dyn everruns_core::CompactionCheckpointStore>,
pub provider_store: Arc<dyn RuntimeProviderStore>,
pub event_sink: Arc<dyn EventSink>,
pub storage_store: Arc<dyn SessionStorageStore>,
pub connection_resolver: Option<Arc<dyn UserConnectionResolver>>,
pub session_task_registry: Option<Arc<dyn SessionTaskRegistry>>,
pub schedule_store_factory: Option<ScheduleStoreFactory>,
pub platform_store_factory: Option<PlatformStoreFactory>,
}
impl HostBackends {
pub fn in_memory() -> Self {
Self {
harness_store: Arc::new(InMemoryHarnessStore::new()),
agent_store: Arc::new(InMemoryAgentStore::new()),
session_store: Arc::new(InMemorySessionStore::new()),
event_log: Arc::new(InMemoryEventLog::new()),
compaction_checkpoint_store: Arc::new(
everruns_core::InMemoryCompactionCheckpointStore::default(),
),
provider_store: Arc::new(InMemoryProviderStore::new()),
event_sink: Arc::new(NoopEventSink),
storage_store: Arc::new(InMemorySessionStorageStore::new()),
connection_resolver: None,
session_task_registry: None,
schedule_store_factory: None,
platform_store_factory: None,
}
}
pub fn with_harness_store(mut self, store: Arc<dyn RuntimeHarnessStore>) -> Self {
self.harness_store = store;
self
}
pub fn with_agent_store(mut self, store: Arc<dyn RuntimeAgentStore>) -> Self {
self.agent_store = store;
self
}
pub fn with_session_store(mut self, store: Arc<dyn RuntimeSessionStore>) -> Self {
self.session_store = store;
self
}
pub fn with_event_log(mut self, log: Arc<dyn EventLog>) -> Self {
self.event_log = log;
self
}
pub fn with_compaction_checkpoint_store(
mut self,
store: Arc<dyn everruns_core::CompactionCheckpointStore>,
) -> Self {
self.compaction_checkpoint_store = store;
self
}
pub fn with_provider_store(mut self, store: Arc<dyn RuntimeProviderStore>) -> Self {
self.provider_store = store;
self
}
pub fn with_event_sink(mut self, sink: Arc<dyn EventSink>) -> Self {
self.event_sink = sink;
self
}
pub fn with_storage_store(mut self, store: Arc<dyn SessionStorageStore>) -> Self {
self.storage_store = store;
self
}
pub fn with_connection_resolver(mut self, resolver: Arc<dyn UserConnectionResolver>) -> Self {
self.connection_resolver = Some(resolver);
self
}
pub fn with_session_task_registry(mut self, registry: Arc<dyn SessionTaskRegistry>) -> Self {
self.session_task_registry = Some(registry);
self
}
pub fn with_schedule_store_factory(mut self, factory: ScheduleStoreFactory) -> Self {
self.schedule_store_factory = Some(factory);
self
}
pub fn with_platform_store_factory(mut self, factory: PlatformStoreFactory) -> Self {
self.platform_store_factory = Some(factory);
self
}
}
#[async_trait]
impl RuntimeAgentStore for InMemoryAgentStore {
async fn add_agent(&self, agent: Agent) -> Result<()> {
InMemoryAgentStore::add_agent(self, agent).await;
Ok(())
}
}
#[async_trait]
impl RuntimeHarnessStore for InMemoryHarnessStore {
async fn add_harness(&self, harness: Harness) -> Result<()> {
InMemoryHarnessStore::add_harness(self, harness).await;
Ok(())
}
}
#[async_trait]
impl RuntimeSessionStore for InMemorySessionStore {
async fn add_session(&self, session: Session) -> Result<()> {
InMemorySessionStore::add_session(self, session).await;
Ok(())
}
}
#[async_trait]
impl RuntimeProviderStore for InMemoryProviderStore {
async fn set_default_model(&self, model: ResolvedModel) -> Result<()> {
InMemoryProviderStore::set_default_model(self, model).await;
Ok(())
}
}