1use std::sync::Arc;
20
21use chrono::Utc;
22use tokio::task;
23
24use crate::{
25 application::{
26 audit::AuditUseCase,
27 dto::{
28 AuditQueryRequest, AuditQueryResponse, IngestClaimRequest, IngestClaimResponse,
29 QueryHistoryRequest, QueryHistoryResponse, QueryMemoryRequest, QueryMemoryResponse,
30 QuerySubjectRequest, QuerySubjectResponse,
31 ReconcileRequest, ReconcileResponse,
32 },
33 ingest_claim::IngestClaimUseCase,
34 query_history::QueryHistoryUseCase,
35 query_memory::QueryMemoryUseCase,
36 query_subject::QuerySubjectUseCase,
37 reconcile::ReconcileUseCase,
38 submit_adjudication::SubmitAdjudicationUseCase,
39 sweep_adjudications::SweepAdjudicationsUseCase,
40 },
41 concurrency::agent_lock::AgentWriteLockMap,
42 config::EngineConfig,
43 error::MemError,
44 ports::{OraclePort, PendingAdjudicationPort, PersistencePort, VectorPort},
45};
46
47#[allow(missing_docs)]
57pub trait ErasedPendingStore: Send + Sync + 'static {
58 fn insert_pending_erased(
59 &self,
60 row: &crate::ports::pending_adjudication::PendingAdjudicationRow,
61 ) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>>;
62
63 fn get_pending_erased(
64 &self,
65 handle_id: uuid::Uuid,
66 ) -> Result<Option<crate::ports::pending_adjudication::PendingAdjudicationRow>, Box<dyn std::error::Error + Send + Sync + 'static>>;
67
68 fn list_pending_erased(
69 &self,
70 agent_id: Option<&mempill_types::AgentId>,
71 ) -> Result<Vec<crate::ports::pending_adjudication::PendingAdjudicationRow>, Box<dyn std::error::Error + Send + Sync + 'static>>;
72
73 fn list_expired_erased(
74 &self,
75 now: chrono::DateTime<chrono::Utc>,
76 ) -> Result<Vec<crate::ports::pending_adjudication::PendingAdjudicationRow>, Box<dyn std::error::Error + Send + Sync + 'static>>;
77
78 fn mark_resolved_erased(
79 &self,
80 handle_id: uuid::Uuid,
81 ) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>>;
82
83 fn mark_expired_erased(
84 &self,
85 handle_id: uuid::Uuid,
86 ) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>>;
87
88 fn list_queued_orphan_claims_erased(
89 &self,
90 ) -> Result<Vec<crate::ports::pending_adjudication::OrphanedQueuedClaim>, Box<dyn std::error::Error + Send + Sync + 'static>>;
91}
92
93pub struct ErasedPendingStoreAdapter<S: PendingAdjudicationPort> {
95 inner: S,
96}
97
98impl<S: PendingAdjudicationPort> ErasedPendingStoreAdapter<S> {
99 pub fn new(inner: S) -> Self {
101 Self { inner }
102 }
103}
104
105impl<S: PendingAdjudicationPort> ErasedPendingStore for ErasedPendingStoreAdapter<S> {
106 fn insert_pending_erased(
107 &self,
108 row: &crate::ports::pending_adjudication::PendingAdjudicationRow,
109 ) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
110 self.inner.insert_pending(row).map_err(|e| Box::new(e) as _)
111 }
112
113 fn get_pending_erased(
114 &self,
115 handle_id: uuid::Uuid,
116 ) -> Result<Option<crate::ports::pending_adjudication::PendingAdjudicationRow>, Box<dyn std::error::Error + Send + Sync + 'static>> {
117 self.inner.get_pending(handle_id).map_err(|e| Box::new(e) as _)
118 }
119
120 fn list_pending_erased(
121 &self,
122 agent_id: Option<&mempill_types::AgentId>,
123 ) -> Result<Vec<crate::ports::pending_adjudication::PendingAdjudicationRow>, Box<dyn std::error::Error + Send + Sync + 'static>> {
124 self.inner.list_pending(agent_id).map_err(|e| Box::new(e) as _)
125 }
126
127 fn list_expired_erased(
128 &self,
129 now: chrono::DateTime<chrono::Utc>,
130 ) -> Result<Vec<crate::ports::pending_adjudication::PendingAdjudicationRow>, Box<dyn std::error::Error + Send + Sync + 'static>> {
131 self.inner.list_expired(now).map_err(|e| Box::new(e) as _)
132 }
133
134 fn mark_resolved_erased(
135 &self,
136 handle_id: uuid::Uuid,
137 ) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
138 self.inner.mark_resolved(handle_id).map_err(|e| Box::new(e) as _)
139 }
140
141 fn mark_expired_erased(
142 &self,
143 handle_id: uuid::Uuid,
144 ) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
145 self.inner.mark_expired(handle_id).map_err(|e| Box::new(e) as _)
146 }
147
148 fn list_queued_orphan_claims_erased(
149 &self,
150 ) -> Result<Vec<crate::ports::pending_adjudication::OrphanedQueuedClaim>, Box<dyn std::error::Error + Send + Sync + 'static>> {
151 self.inner.list_queued_orphan_claims().map_err(|e| Box::new(e) as _)
152 }
153}
154
155pub struct EngineHandle<P, O, V>
162where
163 P: PersistencePort + Send + Sync + 'static,
164 O: OraclePort + Send + Sync + 'static,
165 V: VectorPort + Send + Sync + 'static,
166{
167 persistence: Arc<P>,
168 oracle: Option<Arc<O>>,
169 vector: Option<Arc<V>>,
170 pending_store: Option<Arc<dyn ErasedPendingStore>>,
172 config: EngineConfig,
173 write_locks: AgentWriteLockMap,
174 store_write_lock: Arc<tokio::sync::Mutex<()>>,
178}
179
180impl<P, O, V> EngineHandle<P, O, V>
181where
182 P: PersistencePort + Send + Sync + 'static,
183 O: OraclePort + Send + Sync + 'static,
184 V: VectorPort + Send + Sync + 'static,
185{
186 pub fn new(
192 persistence: Arc<P>,
193 oracle: Option<Arc<O>>,
194 vector: Option<Arc<V>>,
195 config: EngineConfig,
196 ) -> Self {
197 Self {
198 persistence,
199 oracle,
200 vector,
201 pending_store: None,
202 config,
203 write_locks: AgentWriteLockMap::new(),
204 store_write_lock: Arc::new(tokio::sync::Mutex::new(())),
205 }
206 }
207
208 pub fn new_with_pending_store<S>(
224 persistence: Arc<P>,
225 oracle: Option<Arc<O>>,
226 vector: Option<Arc<V>>,
227 pending_store: Arc<dyn ErasedPendingStore>,
228 config: EngineConfig,
229 ) -> Self {
230 Self {
231 persistence,
232 oracle,
233 vector,
234 pending_store: Some(pending_store),
235 config,
236 write_locks: AgentWriteLockMap::new(),
237 store_write_lock: Arc::new(tokio::sync::Mutex::new(())),
238 }
239 }
240
241 pub async fn ingest_claim(
249 &self,
250 req: IngestClaimRequest,
251 ) -> Result<IngestClaimResponse, MemError> {
252 let now = Utc::now(); let _store_lock = if self.persistence.requires_global_write_serialization() {
255 Some(self.store_write_lock.lock().await)
256 } else {
257 None
258 };
259 let _guard = self.write_locks.acquire(&req.agent_id).await;
260 let uc = IngestClaimUseCase::new(
261 Arc::clone(&self.persistence),
262 self.oracle.clone(),
263 self.pending_store.clone(),
264 self.config.clone(),
265 );
266 task::spawn_blocking(move || uc.execute_with_time(req, now))
267 .await
268 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })?
269 }
270
271 pub async fn query_memory(
275 &self,
276 req: QueryMemoryRequest,
277 ) -> Result<QueryMemoryResponse, MemError> {
278 let now = Utc::now();
279 let uc = QueryMemoryUseCase::new(
280 Arc::clone(&self.persistence),
281 self.vector.clone(),
282 self.config.clone(),
283 );
284 task::spawn_blocking(move || uc.execute_with_time(req, now))
285 .await
286 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })?
287 }
288
289 pub async fn query_subject(
294 &self,
295 req: QuerySubjectRequest,
296 ) -> Result<QuerySubjectResponse, MemError> {
297 let now = Utc::now();
298 let uc = QuerySubjectUseCase::new(
299 Arc::clone(&self.persistence),
300 self.vector.clone(),
301 self.config.clone(),
302 );
303 task::spawn_blocking(move || uc.execute_with_time(req, now))
304 .await
305 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })?
306 }
307
308 pub async fn query_history(
316 &self,
317 req: QueryHistoryRequest,
318 ) -> Result<QueryHistoryResponse, MemError> {
319 let now = Utc::now();
320 let uc = QueryHistoryUseCase::new(
321 Arc::clone(&self.persistence),
322 self.vector.clone(),
323 self.config.clone(),
324 );
325 task::spawn_blocking(move || uc.execute_with_time(req, now))
326 .await
327 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })?
328 }
329
330 pub async fn reconcile(
334 &self,
335 req: ReconcileRequest,
336 ) -> Result<ReconcileResponse, MemError> {
337 let _store_lock = if self.persistence.requires_global_write_serialization() {
339 Some(self.store_write_lock.lock().await)
340 } else {
341 None
342 };
343 let _guard = self.write_locks.acquire(&req.agent_id).await;
344 let uc = ReconcileUseCase::new(
345 Arc::clone(&self.persistence),
346 self.oracle.clone(),
347 self.config.clone(),
348 );
349 task::spawn_blocking(move || uc.execute(req))
350 .await
351 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })?
352 }
353
354 pub async fn query_audit(
356 &self,
357 req: AuditQueryRequest,
358 ) -> Result<AuditQueryResponse, MemError> {
359 let uc = AuditUseCase::new(Arc::clone(&self.persistence));
360 task::spawn_blocking(move || uc.execute(req))
361 .await
362 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })?
363 }
364
365 pub async fn submit_adjudication(
394 &self,
395 handle_id: uuid::Uuid,
396 response: mempill_types::AdjudicationResponse,
397 ) -> Result<mempill_types::AdjudicationOutcome, MemError> {
398 let now = Utc::now(); let pending_store = self.pending_store.as_ref()
408 .ok_or(MemError::AdjudicationHandleNotFound { handle_id })?;
409 let pending_store_arc = Arc::clone(pending_store);
410
411 let resolve_result = task::spawn_blocking(move || {
412 let row = pending_store_arc
413 .get_pending_erased(handle_id)
414 .map_err(|e| MemError::PendingStore { source: e })?
415 .ok_or(MemError::AdjudicationHandleNotFound { handle_id })?;
416 Ok::<_, MemError>(row)
417 })
418 .await
419 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })??;
420
421 let row = resolve_result;
422
423 let agent_id = row.agent_id.clone();
427
428 let _store_lock = if self.persistence.requires_global_write_serialization() {
430 Some(self.store_write_lock.lock().await)
431 } else {
432 None
433 };
434 let _guard = self.write_locks.acquire(&agent_id).await;
435
436 let pending_store_arc2 = Arc::clone(pending_store);
438 let uc = SubmitAdjudicationUseCase::new(
439 Arc::clone(&self.persistence),
440 pending_store_arc2,
441 );
442 task::spawn_blocking(move || uc.execute(handle_id, response, now))
443 .await
444 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })?
445 }
446
447 pub async fn list_pending_adjudications(
455 &self,
456 agent_id: Option<mempill_types::AgentId>,
457 ) -> Result<Vec<crate::ports::pending_adjudication::PendingAdjudicationRow>, MemError> {
458 let pending_store = match &self.pending_store {
459 Some(ps) => Arc::clone(ps),
460 None => return Ok(vec![]),
461 };
462
463 task::spawn_blocking(move || {
464 pending_store
465 .list_pending_erased(agent_id.as_ref())
466 .map_err(|e| MemError::PendingStore { source: e })
467 })
468 .await
469 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })?
470 }
471
472 pub async fn sweep_expired_adjudications(&self) -> Result<usize, MemError> {
494 let now = Utc::now();
495
496 let pending_store = match &self.pending_store {
497 Some(ps) => Arc::clone(ps),
498 None => return Ok(0),
499 };
500
501 let ps_for_list = Arc::clone(&pending_store);
503 let expired_rows = task::spawn_blocking(move || {
504 ps_for_list
505 .list_expired_erased(now)
506 .map_err(|e| MemError::PendingStore { source: e })
507 })
508 .await
509 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })??;
510
511 let mut swept = 0usize;
512
513 for row in expired_rows {
514 let agent_id = row.agent_id.clone();
515
516 let _store_lock = if self.persistence.requires_global_write_serialization() {
517 Some(self.store_write_lock.lock().await)
518 } else {
519 None
520 };
521 let _guard = self.write_locks.acquire(&agent_id).await;
522
523 let persistence = Arc::clone(&self.persistence);
524 let ps = Arc::clone(&pending_store);
525 let row_clone = row.clone();
526
527 let result = task::spawn_blocking(move || {
528 let uc = SweepAdjudicationsUseCase::new(persistence, ps);
529 uc.revert_expired_row(&row_clone, now)
530 })
531 .await
532 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })??;
533
534 if result {
535 swept += 1;
536 }
537 }
538
539 let ps_for_orphans = Arc::clone(&pending_store);
542 let orphans = task::spawn_blocking(move || {
543 ps_for_orphans
544 .list_queued_orphan_claims_erased()
545 .map_err(|e| MemError::PendingStore { source: e })
546 })
547 .await
548 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })??;
549
550 for orphan in orphans {
551 let agent_id = orphan.agent_id.clone();
552
553 let _store_lock = if self.persistence.requires_global_write_serialization() {
554 Some(self.store_write_lock.lock().await)
555 } else {
556 None
557 };
558 let _guard = self.write_locks.acquire(&agent_id).await;
559
560 let persistence = Arc::clone(&self.persistence);
561 let ps = Arc::clone(&pending_store);
562 let orphan_clone = orphan.clone();
563
564 let result = task::spawn_blocking(move || {
565 let uc = SweepAdjudicationsUseCase::new(persistence, ps);
566 uc.revert_orphan(&orphan_clone, now)
567 })
568 .await
569 .map_err(|e| MemError::SpawnBlocking { reason: e.to_string() })??;
570
571 if result {
572 swept += 1;
573 }
574 }
575
576 Ok(swept)
577 }
578}
579
580impl<P, O, V> Clone for EngineHandle<P, O, V>
581where
582 P: PersistencePort + Send + Sync + 'static,
583 O: OraclePort + Send + Sync + 'static,
584 V: VectorPort + Send + Sync + 'static,
585{
586 fn clone(&self) -> Self {
587 Self {
588 persistence: Arc::clone(&self.persistence),
589 oracle: self.oracle.clone(),
590 vector: self.vector.clone(),
591 pending_store: self.pending_store.clone(),
592 config: self.config.clone(),
593 write_locks: self.write_locks.clone(),
594 store_write_lock: Arc::clone(&self.store_write_lock),
595 }
596 }
597}