1use std::{collections::BTreeMap, sync::Arc};
2
3use async_trait::async_trait;
4
5pub use loong_contracts::{
7 MemoryCoreOutcome, MemoryCoreRequest, MemoryExtensionOutcome, MemoryExtensionRequest,
8 MemoryTier,
9};
10
11use crate::errors::MemoryPlaneError;
12
13#[async_trait]
14pub trait CoreMemoryAdapter: Send + Sync {
15 fn name(&self) -> &str;
16
17 async fn execute_core_memory(
18 &self,
19 request: MemoryCoreRequest,
20 ) -> Result<MemoryCoreOutcome, MemoryPlaneError>;
21}
22
23#[async_trait]
24pub trait MemoryExtensionAdapter: Send + Sync {
25 fn name(&self) -> &str;
26
27 async fn execute_memory_extension(
28 &self,
29 request: MemoryExtensionRequest,
30 core: &(dyn CoreMemoryAdapter + Sync),
31 ) -> Result<MemoryExtensionOutcome, MemoryPlaneError>;
32}
33
34#[derive(Default)]
35pub struct MemoryPlane {
36 core_adapters: BTreeMap<String, Arc<dyn CoreMemoryAdapter>>,
37 extension_adapters: BTreeMap<String, Arc<dyn MemoryExtensionAdapter>>,
38 default_core_adapter: Option<String>,
39}
40
41impl MemoryPlane {
42 #[must_use]
43 pub fn new() -> Self {
44 Self {
45 core_adapters: BTreeMap::new(),
46 extension_adapters: BTreeMap::new(),
47 default_core_adapter: None,
48 }
49 }
50
51 pub fn register_core_adapter<A: CoreMemoryAdapter + 'static>(&mut self, adapter: A) {
52 let name = adapter.name().to_owned();
53 if self.default_core_adapter.is_none() {
54 self.default_core_adapter = Some(name.clone());
55 }
56 self.core_adapters.insert(name, Arc::new(adapter));
57 }
58
59 pub fn register_extension_adapter<A: MemoryExtensionAdapter + 'static>(&mut self, adapter: A) {
60 let name = adapter.name().to_owned();
61 self.extension_adapters.insert(name, Arc::new(adapter));
62 }
63
64 pub fn set_default_core_adapter(&mut self, name: &str) -> Result<(), MemoryPlaneError> {
65 if !self.core_adapters.contains_key(name) {
66 return Err(MemoryPlaneError::CoreAdapterNotFound(name.to_owned()));
67 }
68 self.default_core_adapter = Some(name.to_owned());
69 Ok(())
70 }
71
72 #[must_use]
73 pub fn default_core_adapter_name(&self) -> Option<&str> {
74 self.default_core_adapter.as_deref()
75 }
76
77 pub async fn execute_core(
78 &self,
79 core_name: Option<&str>,
80 request: MemoryCoreRequest,
81 ) -> Result<MemoryCoreOutcome, MemoryPlaneError> {
82 let resolved_name = if let Some(name) = core_name {
83 name
84 } else {
85 self.default_core_adapter
86 .as_deref()
87 .ok_or(MemoryPlaneError::NoDefaultCoreAdapter)?
88 };
89
90 let adapter = self
91 .core_adapters
92 .get(resolved_name)
93 .ok_or(MemoryPlaneError::CoreAdapterNotFound(
94 resolved_name.to_owned(),
95 ))?
96 .clone();
97
98 return adapter.execute_core_memory(request).await;
99 }
100
101 pub async fn execute_extension(
102 &self,
103 extension_name: &str,
104 core_name: Option<&str>,
105 request: MemoryExtensionRequest,
106 ) -> Result<MemoryExtensionOutcome, MemoryPlaneError> {
107 let extension = self
108 .extension_adapters
109 .get(extension_name)
110 .ok_or_else(|| MemoryPlaneError::ExtensionNotFound(extension_name.to_owned()))?
111 .clone();
112
113 let resolved_core_name = if let Some(name) = core_name {
114 name
115 } else {
116 self.default_core_adapter
117 .as_deref()
118 .ok_or(MemoryPlaneError::NoDefaultCoreAdapter)?
119 };
120
121 let core = self
122 .core_adapters
123 .get(resolved_core_name)
124 .ok_or(MemoryPlaneError::CoreAdapterNotFound(
125 resolved_core_name.to_owned(),
126 ))?
127 .clone();
128
129 return extension
130 .execute_memory_extension(request, core.as_ref())
131 .await;
132 }
133}