agent_runtime_memory_journal/
memory.rs

1use std::collections::BTreeMap;
2
3use async_trait::async_trait;
4use runtime_api_contract::{
5    EventPayload, ExecutionEvent, ExecutionFailure, ExecutionOutcome, ExecutionState,
6};
7use runtime_link_api::{
8    ExecutionJournal, JournalClaim, JournalDelegation, JournalError, JournalExecution,
9    JournalInputReceipt, JournalInteraction, JournalMutation, JournalReservation,
10    NewJournalExecution,
11};
12use runtime_types::{DelegationLeaseRef, ExecutionId, OperationId};
13use tokio::sync::Mutex;
14
15struct State {
16    executions: BTreeMap<ExecutionId, JournalExecution>,
17    idempotency: BTreeMap<String, ExecutionId>,
18    events: BTreeMap<ExecutionId, Vec<ExecutionEvent>>,
19}
20
21/// Deterministic bounded journal for local development and contract tests.
22pub struct MemoryExecutionJournal {
23    state: Mutex<State>,
24    max_events_per_execution: usize,
25}
26
27impl MemoryExecutionJournal {
28    pub fn new(max_events_per_execution: usize) -> Self {
29        Self {
30            state: Mutex::new(State {
31                executions: BTreeMap::new(),
32                idempotency: BTreeMap::new(),
33                events: BTreeMap::new(),
34            }),
35            max_events_per_execution,
36        }
37    }
38}
39
40#[async_trait]
41impl ExecutionJournal for MemoryExecutionJournal {
42    async fn reserve(
43        &self,
44        execution: NewJournalExecution,
45    ) -> Result<JournalReservation, JournalError> {
46        let mut state = self.state.lock().await;
47        if let Some(id) = state.idempotency.get(&execution.idempotency_key) {
48            let existing = state
49                .executions
50                .get(id)
51                .ok_or_else(|| JournalError::Unavailable("idempotency index is corrupt".into()))?;
52            if existing.request != execution.request {
53                return Err(JournalError::Conflict);
54            }
55            return Ok(JournalReservation::Existing(existing.clone()));
56        }
57        if state.executions.contains_key(&execution.view.id) {
58            return Err(JournalError::Conflict);
59        }
60        let stored = JournalExecution {
61            idempotency_key: execution.idempotency_key.clone(),
62            caller: execution.caller,
63            request: execution.request,
64            view: execution.view,
65            version: 1,
66            pending_interaction: None,
67            input_receipts: BTreeMap::new(),
68            delegation: None,
69            claim: None,
70        };
71        state
72            .idempotency
73            .insert(execution.idempotency_key, stored.view.id.clone());
74        state.events.insert(stored.view.id.clone(), Vec::new());
75        state
76            .executions
77            .insert(stored.view.id.clone(), stored.clone());
78        Ok(JournalReservation::Created(stored))
79    }
80
81    async fn get(&self, id: &ExecutionId) -> Result<JournalExecution, JournalError> {
82        self.state
83            .lock()
84            .await
85            .executions
86            .get(id)
87            .cloned()
88            .ok_or(JournalError::NotFound)
89    }
90
91    async fn claim(
92        &self,
93        id: &ExecutionId,
94        worker_id: &str,
95        now_ms: i64,
96        expires_at_ms: i64,
97    ) -> Result<JournalExecution, JournalError> {
98        if worker_id.is_empty() || expires_at_ms <= now_ms {
99            return Err(JournalError::Conflict);
100        }
101        let mut state = self.state.lock().await;
102        let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
103        if execution
104            .claim
105            .as_ref()
106            .is_some_and(|claim| claim.expires_at_ms > now_ms && claim.worker_id != worker_id)
107        {
108            return Err(JournalError::Conflict);
109        }
110        execution.claim = Some(JournalClaim {
111            worker_id: worker_id.into(),
112            expires_at_ms,
113        });
114        execution.version = execution.version.saturating_add(1);
115        Ok(execution.clone())
116    }
117
118    async fn recoverable(
119        &self,
120        now_ms: i64,
121        limit: usize,
122    ) -> Result<Vec<JournalExecution>, JournalError> {
123        if limit == 0 || limit > 10_000 {
124            return Err(JournalError::Conflict);
125        }
126        Ok(self
127            .state
128            .lock()
129            .await
130            .executions
131            .values()
132            .filter(|record| {
133                let pending = !record.view.state.is_terminal()
134                    || record
135                        .delegation
136                        .as_ref()
137                        .is_some_and(|delegation| !delegation.cleanup_complete);
138                pending
139                    && record
140                        .claim
141                        .as_ref()
142                        .is_none_or(|claim| claim.expires_at_ms <= now_ms)
143            })
144            .take(limit)
145            .cloned()
146            .collect())
147    }
148
149    async fn transition(
150        &self,
151        id: &ExecutionId,
152        expected: &[ExecutionState],
153        next: ExecutionState,
154    ) -> Result<JournalExecution, JournalError> {
155        self.finish(id, expected, next, None, None).await
156    }
157
158    async fn attach_delegation(
159        &self,
160        id: &ExecutionId,
161        delegation: JournalDelegation,
162    ) -> Result<JournalExecution, JournalError> {
163        let mut state = self.state.lock().await;
164        let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
165        if execution.view.state != ExecutionState::Queued {
166            return Err(JournalError::Conflict);
167        }
168        if let Some(existing) = &execution.delegation {
169            return if existing.lease_ref == delegation.lease_ref
170                && delegation.revision >= existing.revision
171            {
172                execution.delegation = Some(delegation);
173                execution.version = execution.version.saturating_add(1);
174                Ok(execution.clone())
175            } else {
176                Err(JournalError::Conflict)
177            };
178        }
179        execution.delegation = Some(delegation);
180        execution.version = execution.version.saturating_add(1);
181        Ok(execution.clone())
182    }
183
184    async fn complete_delegation_cleanup(
185        &self,
186        id: &ExecutionId,
187        lease_ref: &DelegationLeaseRef,
188    ) -> Result<JournalExecution, JournalError> {
189        let mut state = self.state.lock().await;
190        let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
191        if execution.view.state != ExecutionState::Finalizing && !execution.view.state.is_terminal()
192        {
193            return Err(JournalError::Conflict);
194        }
195        let delegation = execution
196            .delegation
197            .as_mut()
198            .ok_or(JournalError::Conflict)?;
199        if &delegation.lease_ref != lease_ref {
200            return Err(JournalError::Conflict);
201        }
202        delegation.cleanup_complete = true;
203        execution.version = execution.version.saturating_add(1);
204        Ok(execution.clone())
205    }
206
207    async fn finish(
208        &self,
209        id: &ExecutionId,
210        expected: &[ExecutionState],
211        next: ExecutionState,
212        outcome: Option<ExecutionOutcome>,
213        failure: Option<ExecutionFailure>,
214    ) -> Result<JournalExecution, JournalError> {
215        let mut state = self.state.lock().await;
216        let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
217        if !expected.contains(&execution.view.state) {
218            return Err(JournalError::Conflict);
219        }
220        validate_transition(execution.view.state, next)?;
221        execution.view.state = next;
222        execution.view.outcome = outcome;
223        execution.view.failure = failure;
224        execution.view.updated_at_ms = chrono::Utc::now().timestamp_millis();
225        execution.version = execution.version.saturating_add(1);
226        Ok(execution.clone())
227    }
228
229    async fn transition_with_event(
230        &self,
231        id: &ExecutionId,
232        expected: &[ExecutionState],
233        next: ExecutionState,
234        outcome: Option<ExecutionOutcome>,
235        failure: Option<ExecutionFailure>,
236        payload: EventPayload,
237    ) -> Result<JournalMutation, JournalError> {
238        let mut state = self.state.lock().await;
239        ensure_event_capacity(&state, id, self.max_events_per_execution)?;
240        let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
241        if !expected.contains(&execution.view.state) {
242            return Err(JournalError::Conflict);
243        }
244        validate_transition(execution.view.state, next)?;
245        execution.view.state = next;
246        execution.view.outcome = outcome;
247        execution.view.failure = failure;
248        execution.view.updated_at_ms = chrono::Utc::now().timestamp_millis();
249        execution.version = execution.version.saturating_add(1);
250        let execution = execution.clone();
251        let event = push_event(&mut state, id, payload)?;
252        Ok(JournalMutation {
253            execution,
254            event,
255            replayed: false,
256        })
257    }
258
259    async fn append_event(
260        &self,
261        id: &ExecutionId,
262        payload: EventPayload,
263    ) -> Result<ExecutionEvent, JournalError> {
264        let mut state = self.state.lock().await;
265        if !state.executions.contains_key(id) {
266            return Err(JournalError::NotFound);
267        }
268        let events = state
269            .events
270            .get_mut(id)
271            .ok_or_else(|| JournalError::Unavailable("event partition is missing".into()))?;
272        if events.len() >= self.max_events_per_execution {
273            return Err(JournalError::Unavailable(
274                "event retention limit exceeded".into(),
275            ));
276        }
277        let event = ExecutionEvent {
278            execution_id: id.clone(),
279            sequence: events.len() as u64 + 1,
280            created_at_ms: chrono::Utc::now().timestamp_millis(),
281            payload,
282        };
283        events.push(event.clone());
284        Ok(event)
285    }
286
287    async fn begin_interaction(
288        &self,
289        id: &ExecutionId,
290        request_id: &str,
291        prompt: &str,
292    ) -> Result<JournalMutation, JournalError> {
293        let mut state = self.state.lock().await;
294        ensure_event_capacity(&state, id, self.max_events_per_execution)?;
295        let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
296        if execution.view.state != ExecutionState::Running
297            || execution.pending_interaction.is_some()
298        {
299            return Err(JournalError::Conflict);
300        }
301        execution.view.state = ExecutionState::WaitingForInput;
302        execution.view.updated_at_ms = chrono::Utc::now().timestamp_millis();
303        execution.version = execution.version.saturating_add(1);
304        execution.pending_interaction = Some(JournalInteraction {
305            request_id: request_id.into(),
306            prompt: prompt.into(),
307        });
308        let execution = execution.clone();
309        let event = push_event(
310            &mut state,
311            id,
312            EventPayload::InteractionRequired {
313                request_id: request_id.into(),
314                prompt: prompt.into(),
315            },
316        )?;
317        Ok(JournalMutation {
318            execution,
319            event,
320            replayed: false,
321        })
322    }
323
324    async fn commit_interaction_input(
325        &self,
326        id: &ExecutionId,
327        operation_id: &OperationId,
328        request_id: &str,
329        response_text: &str,
330    ) -> Result<JournalMutation, JournalError> {
331        let mut state = self.state.lock().await;
332        let existing = state
333            .executions
334            .get(id)
335            .ok_or(JournalError::NotFound)?
336            .input_receipts
337            .get(operation_id)
338            .cloned();
339        if let Some(receipt) = existing {
340            if receipt.request_id != request_id || receipt.response_text != response_text {
341                return Err(JournalError::Conflict);
342            }
343            let event = state
344                .events
345                .get(id)
346                .and_then(|events| {
347                    events
348                        .iter()
349                        .find(|event| event.sequence == receipt.event_sequence)
350                })
351                .cloned()
352                .ok_or_else(|| {
353                    JournalError::Unavailable("input receipt event is missing".into())
354                })?;
355            return Ok(JournalMutation {
356                execution: state
357                    .executions
358                    .get(id)
359                    .ok_or(JournalError::NotFound)?
360                    .clone(),
361                event,
362                replayed: true,
363            });
364        }
365        ensure_event_capacity(&state, id, self.max_events_per_execution)?;
366        let event_sequence = state
367            .events
368            .get(id)
369            .map_or(1, |events| events.len() as u64 + 1);
370        let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
371        let matches_pending = execution
372            .pending_interaction
373            .as_ref()
374            .is_some_and(|pending| pending.request_id == request_id);
375        if execution.view.state != ExecutionState::WaitingForInput || !matches_pending {
376            return Err(JournalError::Conflict);
377        }
378        execution.input_receipts.insert(
379            operation_id.clone(),
380            JournalInputReceipt {
381                request_id: request_id.into(),
382                response_text: response_text.into(),
383                event_sequence,
384            },
385        );
386        execution.view.updated_at_ms = chrono::Utc::now().timestamp_millis();
387        execution.version = execution.version.saturating_add(1);
388        let execution = execution.clone();
389        let event = push_event(
390            &mut state,
391            id,
392            EventPayload::InteractionReceived {
393                request_id: request_id.into(),
394            },
395        )?;
396        Ok(JournalMutation {
397            execution,
398            event,
399            replayed: false,
400        })
401    }
402
403    async fn complete_interaction(
404        &self,
405        id: &ExecutionId,
406        request_id: &str,
407    ) -> Result<JournalExecution, JournalError> {
408        let mut state = self.state.lock().await;
409        let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
410        let matches_pending = execution
411            .pending_interaction
412            .as_ref()
413            .is_some_and(|pending| pending.request_id == request_id);
414        let has_response = execution
415            .input_receipts
416            .values()
417            .any(|receipt| receipt.request_id == request_id);
418        if execution.view.state != ExecutionState::WaitingForInput
419            || !matches_pending
420            || !has_response
421        {
422            return Err(JournalError::Conflict);
423        }
424        execution.pending_interaction = None;
425        execution.view.state = ExecutionState::Running;
426        execution.view.updated_at_ms = chrono::Utc::now().timestamp_millis();
427        execution.version = execution.version.saturating_add(1);
428        Ok(execution.clone())
429    }
430
431    async fn events(
432        &self,
433        id: &ExecutionId,
434        after: Option<u64>,
435        limit: usize,
436    ) -> Result<Vec<ExecutionEvent>, JournalError> {
437        let state = self.state.lock().await;
438        if !state.executions.contains_key(id) {
439            return Err(JournalError::NotFound);
440        }
441        let after = after.unwrap_or(0);
442        Ok(state
443            .events
444            .get(id)
445            .into_iter()
446            .flatten()
447            .filter(|event| event.sequence > after)
448            .take(limit.min(1000))
449            .cloned()
450            .collect())
451    }
452}
453
454fn ensure_event_capacity(
455    state: &State,
456    id: &ExecutionId,
457    limit: usize,
458) -> Result<(), JournalError> {
459    let events = state.events.get(id).ok_or(JournalError::NotFound)?;
460    if events.len() >= limit {
461        Err(JournalError::Unavailable(
462            "event retention limit exceeded".into(),
463        ))
464    } else {
465        Ok(())
466    }
467}
468
469fn push_event(
470    state: &mut State,
471    id: &ExecutionId,
472    payload: EventPayload,
473) -> Result<ExecutionEvent, JournalError> {
474    let events = state.events.get_mut(id).ok_or(JournalError::NotFound)?;
475    let event = ExecutionEvent {
476        execution_id: id.clone(),
477        sequence: events.len() as u64 + 1,
478        created_at_ms: chrono::Utc::now().timestamp_millis(),
479        payload,
480    };
481    events.push(event.clone());
482    Ok(event)
483}
484
485fn validate_transition(current: ExecutionState, next: ExecutionState) -> Result<(), JournalError> {
486    let valid = matches!(
487        (current, next),
488        (
489            ExecutionState::Queued,
490            ExecutionState::Running
491                | ExecutionState::Finalizing
492                | ExecutionState::Canceled
493                | ExecutionState::Failed
494        ) | (
495            ExecutionState::Running,
496            ExecutionState::WaitingForInput
497                | ExecutionState::Finalizing
498                | ExecutionState::Failed
499                | ExecutionState::Canceled
500        ) | (
501            ExecutionState::WaitingForInput,
502            ExecutionState::Running
503                | ExecutionState::Finalizing
504                | ExecutionState::Failed
505                | ExecutionState::Canceled
506        ) | (
507            ExecutionState::Finalizing,
508            ExecutionState::Completed | ExecutionState::Failed | ExecutionState::Canceled
509        )
510    );
511    if valid {
512        Ok(())
513    } else {
514        Err(JournalError::Conflict)
515    }
516}
517
518#[cfg(test)]
519mod tests {
520    use super::*;
521    use runtime_api_contract::{
522        CreateExecutionRequest, ExecutionOptions, ExecutionView, RuntimeInput,
523    };
524    use runtime_types::{CallerScope, RuntimeInstanceId};
525
526    fn new_execution(key: &str) -> NewJournalExecution {
527        let id = ExecutionId::random();
528        let request = CreateExecutionRequest {
529            runtime_instance_id: RuntimeInstanceId::new("runtime-1").unwrap(),
530            conversation_id: None,
531            input: RuntimeInput::UserMessage {
532                text: "hello".into(),
533            },
534            workspace_id: None,
535            model: None,
536            instructions: None,
537            metadata: Default::default(),
538            generation: Default::default(),
539            options: ExecutionOptions::default(),
540        };
541        NewJournalExecution {
542            idempotency_key: key.into(),
543            caller: CallerScope {
544                subject: "user-1".into(),
545                tenant_id: "tenant-1".into(),
546                project_id: "project-1".into(),
547                capabilities: Vec::new(),
548            },
549            request: request.clone(),
550            view: ExecutionView {
551                id,
552                runtime_instance_id: request.runtime_instance_id.clone(),
553                conversation_id: None,
554                workspace_id: None,
555                model: None,
556                metadata: Default::default(),
557                state: ExecutionState::Queued,
558                outcome: None,
559                failure: None,
560                created_at_ms: 1,
561                updated_at_ms: 1,
562            },
563        }
564    }
565
566    #[tokio::test]
567    async fn reservation_is_idempotent_and_transitions_are_guarded() {
568        let journal = MemoryExecutionJournal::new(10);
569        let created = journal.reserve(new_execution("key")).await.unwrap();
570        let id = match created {
571            JournalReservation::Created(value) => value.view.id,
572            _ => panic!("created"),
573        };
574        assert!(matches!(
575            journal.reserve(new_execution("key")).await,
576            Ok(JournalReservation::Existing(_))
577        ));
578        let mut conflicting = new_execution("key");
579        conflicting.request.input = RuntimeInput::UserMessage {
580            text: "different".into(),
581        };
582        assert!(matches!(
583            journal.reserve(conflicting).await,
584            Err(JournalError::Conflict)
585        ));
586        journal
587            .transition(&id, &[ExecutionState::Queued], ExecutionState::Running)
588            .await
589            .unwrap();
590        assert!(matches!(
591            journal
592                .transition(&id, &[ExecutionState::Queued], ExecutionState::Running)
593                .await,
594            Err(JournalError::Conflict)
595        ));
596    }
597
598    #[tokio::test]
599    async fn claim_expiry_allows_exactly_one_recovery_worker() {
600        let journal = MemoryExecutionJournal::new(10);
601        let id = match journal.reserve(new_execution("claim-key")).await.unwrap() {
602            JournalReservation::Created(value) => value.view.id,
603            _ => panic!("created"),
604        };
605        journal.claim(&id, "worker-a", 100, 200).await.unwrap();
606        assert!(matches!(
607            journal.claim(&id, "worker-b", 150, 250).await,
608            Err(JournalError::Conflict)
609        ));
610        assert!(journal.recoverable(199, 10).await.unwrap().is_empty());
611        assert_eq!(journal.recoverable(200, 10).await.unwrap().len(), 1);
612        let claimed = journal.claim(&id, "worker-b", 200, 300).await.unwrap();
613        assert_eq!(claimed.claim.unwrap().worker_id, "worker-b");
614    }
615
616    #[tokio::test]
617    async fn interaction_input_is_committed_before_resume_and_deduplicated() {
618        let journal = MemoryExecutionJournal::new(10);
619        let id = match journal.reserve(new_execution("interaction")).await.unwrap() {
620            JournalReservation::Created(value) => value.view.id,
621            _ => panic!("created"),
622        };
623        journal
624            .transition(&id, &[ExecutionState::Queued], ExecutionState::Running)
625            .await
626            .unwrap();
627        let required = journal
628            .begin_interaction(&id, "request-1", "approve?")
629            .await
630            .unwrap();
631        assert_eq!(
632            required.execution.view.state,
633            ExecutionState::WaitingForInput
634        );
635
636        let operation_id = OperationId::new("input-1").unwrap();
637        let committed = journal
638            .commit_interaction_input(&id, &operation_id, "request-1", "yes")
639            .await
640            .unwrap();
641        assert!(!committed.replayed);
642        assert_eq!(
643            committed.execution.view.state,
644            ExecutionState::WaitingForInput
645        );
646        let replay = journal
647            .commit_interaction_input(&id, &operation_id, "request-1", "yes")
648            .await
649            .unwrap();
650        assert!(replay.replayed);
651        assert_eq!(replay.event.sequence, committed.event.sequence);
652        assert!(matches!(
653            journal
654                .commit_interaction_input(&id, &operation_id, "request-1", "no")
655                .await,
656            Err(JournalError::Conflict)
657        ));
658        let resumed = journal
659            .complete_interaction(&id, "request-1")
660            .await
661            .unwrap();
662        assert_eq!(resumed.view.state, ExecutionState::Running);
663    }
664}
agent_runtime_memory_journal/memory.rs

agent_runtime_memory_journal/
memory.rs
