aep_platform/
store.rs

1use std::{
2    collections::BTreeMap,
3    sync::{Arc, Mutex},
4    time::Duration,
5};
6
7use async_trait::async_trait;
8use time::OffsetDateTime;
9
10use crate::{
11    Clock, IdempotencyInput, IdempotencyOperation, IdempotencyResult, IdempotencyStore,
12    IdentityCreation, IdentityListQuery, IdentityListResult, IdentityRecord, IdentityStore,
13    LifecyclePolicy, ManagedAgentStatus, PlatformError, ReplayStore, RequestContext,
14    StoredResponse, validate_identity_record,
15};
16
17type ScopeLock = Arc<futures::lock::Mutex<()>>;
18
19/// How long a replayable idempotency result stays available (PLT-IDP-004).
20const RETENTION: Duration = Duration::from_secs(3600);
21type IdempotencyKey = (String, String);
22
23#[derive(Default)]
24struct IdentityRecords {
25    by_agent_did: BTreeMap<String, String>,
26    by_agent_did_id: BTreeMap<String, String>,
27    by_scope: BTreeMap<(String, String), String>,
28    records: BTreeMap<String, IdentityRecord>,
29}
30
31#[derive(Default)]
32pub struct MemoryIdentityStore {
33    locks: Mutex<BTreeMap<(String, String), ScopeLock>>,
34    records: Mutex<IdentityRecords>,
35}
36
37impl MemoryIdentityStore {
38    fn scope_lock(&self, key: &(String, String)) -> Result<ScopeLock, PlatformError> {
39        Ok(self
40            .locks
41            .lock()
42            .map_err(lock_error)?
43            .entry(key.clone())
44            .or_default()
45            .clone())
46    }
47
48    fn find_indexed(
49        &self,
50        key: &str,
51        index: impl FnOnce(&IdentityRecords) -> &BTreeMap<String, String>,
52    ) -> Result<Option<IdentityRecord>, PlatformError> {
53        let records = self.records.lock().map_err(lock_error)?;
54        Ok(index(&records)
55            .get(key)
56            .and_then(|identity_id| records.records.get(identity_id))
57            .cloned())
58    }
59
60    fn release_scope_lock(
61        &self,
62        key: &(String, String),
63        scope_lock: &ScopeLock,
64    ) -> Result<(), PlatformError> {
65        let mut locks = self.locks.lock().map_err(lock_error)?;
66        if locks
67            .get(key)
68            .is_some_and(|current| Arc::ptr_eq(current, scope_lock))
69            && Arc::strong_count(scope_lock) == 2
70        {
71            locks.remove(key);
72        }
73        Ok(())
74    }
75}
76
77#[async_trait]
78impl IdentityStore for MemoryIdentityStore {
79    async fn find_or_create(
80        &self,
81        principal: &str,
82        service_did: &str,
83        create: IdentityCreation,
84    ) -> Result<IdentityRecord, PlatformError> {
85        if principal.is_empty() || service_did.is_empty() {
86            return Err(PlatformError::Store(
87                "AEP Platform identity scope must not be empty".to_owned(),
88            ));
89        }
90        let key = (principal.to_owned(), service_did.to_owned());
91        let scope_lock = self.scope_lock(&key)?;
92        let result = async {
93            let _guard = scope_lock.lock().await;
94            let existing = {
95                let records = self.records.lock().map_err(lock_error)?;
96                records
97                    .by_scope
98                    .get(&key)
99                    .and_then(|identity_id| records.records.get(identity_id))
100                    .cloned()
101            };
102            if let Some(existing) = existing {
103                Ok(existing)
104            } else {
105                let identity = create().await?;
106                validate_identity_record(&identity)?;
107                if identity.principal != principal || identity.service_did != service_did {
108                    return Err(PlatformError::Store(
109                        "AEP Platform identity does not match its requested scope".to_owned(),
110                    ));
111                }
112                let mut records = self.records.lock().map_err(lock_error)?;
113                if records.records.contains_key(&identity.agent_identity_id)
114                    || records.by_agent_did.contains_key(&identity.agent_did)
115                    || records.by_agent_did_id.contains_key(&identity.agent_did_id)
116                {
117                    return Err(PlatformError::Store(
118                        "AEP Platform identity material must be unique".to_owned(),
119                    ));
120                }
121                records
122                    .by_scope
123                    .insert(key.clone(), identity.agent_identity_id.clone());
124                records.by_agent_did.insert(
125                    identity.agent_did.clone(),
126                    identity.agent_identity_id.clone(),
127                );
128                records.by_agent_did_id.insert(
129                    identity.agent_did_id.clone(),
130                    identity.agent_identity_id.clone(),
131                );
132                records
133                    .records
134                    .insert(identity.agent_identity_id.clone(), identity.clone());
135                Ok(identity)
136            }
137        }
138        .await;
139        self.release_scope_lock(&key, &scope_lock)?;
140        result
141    }
142
143    async fn find_by_agent_did(
144        &self,
145        agent_did: &str,
146    ) -> Result<Option<IdentityRecord>, PlatformError> {
147        self.find_indexed(agent_did, |records| &records.by_agent_did)
148    }
149
150    async fn find_by_agent_did_id(
151        &self,
152        agent_did_id: &str,
153    ) -> Result<Option<IdentityRecord>, PlatformError> {
154        self.find_indexed(agent_did_id, |records| &records.by_agent_did_id)
155    }
156
157    async fn get(&self, agent_identity_id: &str) -> Result<Option<IdentityRecord>, PlatformError> {
158        Ok(self
159            .records
160            .lock()
161            .map_err(lock_error)?
162            .records
163            .get(agent_identity_id)
164            .cloned())
165    }
166
167    async fn list(
168        &self,
169        principal: &str,
170        query: &IdentityListQuery,
171    ) -> Result<IdentityListResult, PlatformError> {
172        if principal.is_empty() {
173            return Err(PlatformError::Store(
174                "AEP Platform identity-list principal must not be empty".to_owned(),
175            ));
176        }
177        let mut identities = self
178            .records
179            .lock()
180            .map_err(lock_error)?
181            .records
182            .values()
183            .filter(|identity| {
184                identity.principal == principal
185                    && query
186                        .service_did
187                        .as_ref()
188                        .is_none_or(|service_did| &identity.service_did == service_did)
189                    && query.status.is_none_or(|status| identity.status == status)
190            })
191            .cloned()
192            .collect::<Vec<_>>();
193        identities.sort_by(|left, right| {
194            left.created_at
195                .cmp(&right.created_at)
196                .then_with(|| left.agent_identity_id.cmp(&right.agent_identity_id))
197        });
198        if query.descending {
199            identities.reverse();
200        }
201        let total = identities.len();
202        let identities = identities
203            .into_iter()
204            .skip(query.offset)
205            .take(query.limit)
206            .collect();
207        Ok(IdentityListResult { identities, total })
208    }
209
210    async fn update_status(
211        &self,
212        agent_identity_id: &str,
213        status: ManagedAgentStatus,
214        updated_at: OffsetDateTime,
215    ) -> Result<Option<IdentityRecord>, PlatformError> {
216        let mut records = self.records.lock().map_err(lock_error)?;
217        let Some(identity) = records.records.get_mut(agent_identity_id) else {
218            return Ok(None);
219        };
220        identity.status = status;
221        identity.updated_at = updated_at;
222        Ok(Some(identity.clone()))
223    }
224}
225
226#[derive(Clone, Debug)]
227struct IdempotencyRecord {
228    input: IdempotencyInput,
229    response: StoredResponse,
230    // Retention starts at completion, independently of caller-supplied response metadata.
231    stored_at: OffsetDateTime,
232}
233
234pub struct MemoryIdempotencyStore {
235    clock: Arc<dyn Clock>,
236    locks: Mutex<BTreeMap<IdempotencyKey, ScopeLock>>,
237    records: Mutex<BTreeMap<IdempotencyKey, IdempotencyRecord>>,
238}
239
240impl MemoryIdempotencyStore {
241    pub fn new(clock: Arc<dyn Clock>) -> Self {
242        Self {
243            clock,
244            locks: Mutex::new(BTreeMap::new()),
245            records: Mutex::new(BTreeMap::new()),
246        }
247    }
248
249    fn operation_lock(&self, key: &IdempotencyKey) -> Result<ScopeLock, PlatformError> {
250        Ok(self
251            .locks
252            .lock()
253            .map_err(lock_error)?
254            .entry(key.clone())
255            .or_default()
256            .clone())
257    }
258
259    fn release_operation_lock(
260        &self,
261        key: &IdempotencyKey,
262        operation_lock: &ScopeLock,
263    ) -> Result<(), PlatformError> {
264        let mut locks = self.locks.lock().map_err(lock_error)?;
265        if locks
266            .get(key)
267            .is_some_and(|current| Arc::ptr_eq(current, operation_lock))
268            && Arc::strong_count(operation_lock) == 2
269        {
270            locks.remove(key);
271        }
272        Ok(())
273    }
274}
275
276#[async_trait]
277impl IdempotencyStore for MemoryIdempotencyStore {
278    async fn execute(
279        &self,
280        input: IdempotencyInput,
281        operation: IdempotencyOperation,
282    ) -> Result<IdempotencyResult, PlatformError> {
283        if input.principal.is_empty()
284            || input.idempotency_key.is_empty()
285            || input.request_hash.is_empty()
286        {
287            return Err(PlatformError::Store(
288                "AEP Platform idempotency input is invalid".to_owned(),
289            ));
290        }
291        let key = (input.principal.clone(), input.idempotency_key.clone());
292        let operation_lock = self.operation_lock(&key)?;
293        let result = async {
294            let _guard = operation_lock.lock().await;
295            let now = self.clock.now();
296            let existing = {
297                let mut records = self.records.lock().map_err(lock_error)?;
298                // PLT-IDP-004 retains the replayable result for at least one hour, so a record
299                // is still usable at exactly one hour old.
300                records.retain(|_, record| record.stored_at + RETENTION >= now);
301                records.get(&key).cloned()
302            };
303            if let Some(existing) = existing {
304                if existing.input.operation == input.operation
305                    && existing.input.request_hash == input.request_hash
306                {
307                    Ok(IdempotencyResult::Replayed(existing.response))
308                } else {
309                    Ok(IdempotencyResult::Conflict)
310                }
311            } else {
312                let mut response = operation().await?;
313                let stored_at = self.clock.now();
314                if response.created_at == OffsetDateTime::UNIX_EPOCH {
315                    response.created_at = stored_at;
316                }
317                self.records.lock().map_err(lock_error)?.insert(
318                    key.clone(),
319                    IdempotencyRecord {
320                        input,
321                        response: response.clone(),
322                        stored_at,
323                    },
324                );
325                Ok(IdempotencyResult::Created(response))
326            }
327        }
328        .await;
329        self.release_operation_lock(&key, &operation_lock)?;
330        result
331    }
332}
333
334impl Default for MemoryIdempotencyStore {
335    fn default() -> Self {
336        Self::new(Arc::new(SystemClock))
337    }
338}
339
340#[derive(Default)]
341pub struct MemoryReplayStore {
342    records: Mutex<BTreeMap<String, OffsetDateTime>>,
343}
344
345#[async_trait]
346impl ReplayStore for MemoryReplayStore {
347    async fn consume(
348        &self,
349        key: &str,
350        expires_at: OffsetDateTime,
351        now: OffsetDateTime,
352    ) -> Result<bool, PlatformError> {
353        if key.is_empty() || expires_at <= now {
354            return Ok(false);
355        }
356        let mut records = self.records.lock().map_err(lock_error)?;
357        records.retain(|_, expiry| *expiry > now);
358        if records.contains_key(key) {
359            return Ok(false);
360        }
361        records.insert(key.to_owned(), expires_at);
362        Ok(true)
363    }
364}
365
366#[derive(Default)]
367pub struct SystemClock;
368
369impl Clock for SystemClock {
370    fn now(&self) -> OffsetDateTime {
371        OffsetDateTime::now_utc()
372    }
373}
374
375#[derive(Default)]
376pub struct DefaultLifecyclePolicy;
377
378#[async_trait]
379impl LifecyclePolicy for DefaultLifecyclePolicy {
380    async fn can_sign(
381        &self,
382        identity: &IdentityRecord,
383        _context: &RequestContext,
384    ) -> Result<bool, PlatformError> {
385        Ok(identity.status == ManagedAgentStatus::Active)
386    }
387
388    async fn can_transition(
389        &self,
390        _identity: &IdentityRecord,
391        _status: ManagedAgentStatus,
392        _context: &RequestContext,
393    ) -> Result<bool, PlatformError> {
394        Ok(true)
395    }
396
397    async fn can_verify(
398        &self,
399        identity: &IdentityRecord,
400        _context: &RequestContext,
401    ) -> Result<bool, PlatformError> {
402        Ok(identity.status == ManagedAgentStatus::Active)
403    }
404}
405
406fn lock_error<T>(_error: std::sync::PoisonError<T>) -> PlatformError {
407    PlatformError::Store("AEP Platform memory store lock is poisoned".to_owned())
408}
aep_platform/store.rs

aep_platform/
store.rs
