use std::collections::BTreeMap;
use async_trait::async_trait;
use runtime_api_contract::{
EventPayload, ExecutionEvent, ExecutionFailure, ExecutionOutcome, ExecutionState,
};
use runtime_link_api::{
ExecutionJournal, JournalClaim, JournalDelegation, JournalError, JournalExecution,
JournalInputReceipt, JournalInteraction, JournalMutation, JournalReservation,
NewJournalExecution,
};
use runtime_types::{DelegationLeaseRef, ExecutionId, OperationId};
use tokio::sync::Mutex;
struct State {
executions: BTreeMap<ExecutionId, JournalExecution>,
idempotency: BTreeMap<String, ExecutionId>,
events: BTreeMap<ExecutionId, Vec<ExecutionEvent>>,
}
pub struct MemoryExecutionJournal {
state: Mutex<State>,
max_events_per_execution: usize,
}
impl MemoryExecutionJournal {
pub fn new(max_events_per_execution: usize) -> Self {
Self {
state: Mutex::new(State {
executions: BTreeMap::new(),
idempotency: BTreeMap::new(),
events: BTreeMap::new(),
}),
max_events_per_execution,
}
}
}
#[async_trait]
impl ExecutionJournal for MemoryExecutionJournal {
async fn reserve(
&self,
execution: NewJournalExecution,
) -> Result<JournalReservation, JournalError> {
let mut state = self.state.lock().await;
if let Some(id) = state.idempotency.get(&execution.idempotency_key) {
let existing = state
.executions
.get(id)
.ok_or_else(|| JournalError::Unavailable("idempotency index is corrupt".into()))?;
if existing.request != execution.request {
return Err(JournalError::Conflict);
}
return Ok(JournalReservation::Existing(existing.clone()));
}
if state.executions.contains_key(&execution.view.id) {
return Err(JournalError::Conflict);
}
let stored = JournalExecution {
idempotency_key: execution.idempotency_key.clone(),
caller: execution.caller,
request: execution.request,
view: execution.view,
version: 1,
pending_interaction: None,
input_receipts: BTreeMap::new(),
delegation: None,
claim: None,
};
state
.idempotency
.insert(execution.idempotency_key, stored.view.id.clone());
state.events.insert(stored.view.id.clone(), Vec::new());
state
.executions
.insert(stored.view.id.clone(), stored.clone());
Ok(JournalReservation::Created(stored))
}
async fn get(&self, id: &ExecutionId) -> Result<JournalExecution, JournalError> {
self.state
.lock()
.await
.executions
.get(id)
.cloned()
.ok_or(JournalError::NotFound)
}
async fn claim(
&self,
id: &ExecutionId,
worker_id: &str,
now_ms: i64,
expires_at_ms: i64,
) -> Result<JournalExecution, JournalError> {
if worker_id.is_empty() || expires_at_ms <= now_ms {
return Err(JournalError::Conflict);
}
let mut state = self.state.lock().await;
let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
if execution
.claim
.as_ref()
.is_some_and(|claim| claim.expires_at_ms > now_ms && claim.worker_id != worker_id)
{
return Err(JournalError::Conflict);
}
execution.claim = Some(JournalClaim {
worker_id: worker_id.into(),
expires_at_ms,
});
execution.version = execution.version.saturating_add(1);
Ok(execution.clone())
}
async fn recoverable(
&self,
now_ms: i64,
limit: usize,
) -> Result<Vec<JournalExecution>, JournalError> {
if limit == 0 || limit > 10_000 {
return Err(JournalError::Conflict);
}
Ok(self
.state
.lock()
.await
.executions
.values()
.filter(|record| {
let pending = !record.view.state.is_terminal()
|| record
.delegation
.as_ref()
.is_some_and(|delegation| !delegation.cleanup_complete);
pending
&& record
.claim
.as_ref()
.is_none_or(|claim| claim.expires_at_ms <= now_ms)
})
.take(limit)
.cloned()
.collect())
}
async fn transition(
&self,
id: &ExecutionId,
expected: &[ExecutionState],
next: ExecutionState,
) -> Result<JournalExecution, JournalError> {
self.finish(id, expected, next, None, None).await
}
async fn attach_delegation(
&self,
id: &ExecutionId,
delegation: JournalDelegation,
) -> Result<JournalExecution, JournalError> {
let mut state = self.state.lock().await;
let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
if execution.view.state != ExecutionState::Queued {
return Err(JournalError::Conflict);
}
if let Some(existing) = &execution.delegation {
return if existing.lease_ref == delegation.lease_ref
&& delegation.revision >= existing.revision
{
execution.delegation = Some(delegation);
execution.version = execution.version.saturating_add(1);
Ok(execution.clone())
} else {
Err(JournalError::Conflict)
};
}
execution.delegation = Some(delegation);
execution.version = execution.version.saturating_add(1);
Ok(execution.clone())
}
async fn complete_delegation_cleanup(
&self,
id: &ExecutionId,
lease_ref: &DelegationLeaseRef,
) -> Result<JournalExecution, JournalError> {
let mut state = self.state.lock().await;
let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
if execution.view.state != ExecutionState::Finalizing && !execution.view.state.is_terminal()
{
return Err(JournalError::Conflict);
}
let delegation = execution
.delegation
.as_mut()
.ok_or(JournalError::Conflict)?;
if &delegation.lease_ref != lease_ref {
return Err(JournalError::Conflict);
}
delegation.cleanup_complete = true;
execution.version = execution.version.saturating_add(1);
Ok(execution.clone())
}
async fn finish(
&self,
id: &ExecutionId,
expected: &[ExecutionState],
next: ExecutionState,
outcome: Option<ExecutionOutcome>,
failure: Option<ExecutionFailure>,
) -> Result<JournalExecution, JournalError> {
let mut state = self.state.lock().await;
let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
if !expected.contains(&execution.view.state) {
return Err(JournalError::Conflict);
}
validate_transition(execution.view.state, next)?;
execution.view.state = next;
execution.view.outcome = outcome;
execution.view.failure = failure;
execution.view.updated_at_ms = chrono::Utc::now().timestamp_millis();
execution.version = execution.version.saturating_add(1);
Ok(execution.clone())
}
async fn transition_with_event(
&self,
id: &ExecutionId,
expected: &[ExecutionState],
next: ExecutionState,
outcome: Option<ExecutionOutcome>,
failure: Option<ExecutionFailure>,
payload: EventPayload,
) -> Result<JournalMutation, JournalError> {
let mut state = self.state.lock().await;
ensure_event_capacity(&state, id, self.max_events_per_execution)?;
let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
if !expected.contains(&execution.view.state) {
return Err(JournalError::Conflict);
}
validate_transition(execution.view.state, next)?;
execution.view.state = next;
execution.view.outcome = outcome;
execution.view.failure = failure;
execution.view.updated_at_ms = chrono::Utc::now().timestamp_millis();
execution.version = execution.version.saturating_add(1);
let execution = execution.clone();
let event = push_event(&mut state, id, payload)?;
Ok(JournalMutation {
execution,
event,
replayed: false,
})
}
async fn append_event(
&self,
id: &ExecutionId,
payload: EventPayload,
) -> Result<ExecutionEvent, JournalError> {
let mut state = self.state.lock().await;
if !state.executions.contains_key(id) {
return Err(JournalError::NotFound);
}
let events = state
.events
.get_mut(id)
.ok_or_else(|| JournalError::Unavailable("event partition is missing".into()))?;
if events.len() >= self.max_events_per_execution {
return Err(JournalError::Unavailable(
"event retention limit exceeded".into(),
));
}
let event = ExecutionEvent {
execution_id: id.clone(),
sequence: events.len() as u64 + 1,
created_at_ms: chrono::Utc::now().timestamp_millis(),
payload,
};
events.push(event.clone());
Ok(event)
}
async fn begin_interaction(
&self,
id: &ExecutionId,
request_id: &str,
prompt: &str,
) -> Result<JournalMutation, JournalError> {
let mut state = self.state.lock().await;
ensure_event_capacity(&state, id, self.max_events_per_execution)?;
let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
if execution.view.state != ExecutionState::Running
|| execution.pending_interaction.is_some()
{
return Err(JournalError::Conflict);
}
execution.view.state = ExecutionState::WaitingForInput;
execution.view.updated_at_ms = chrono::Utc::now().timestamp_millis();
execution.version = execution.version.saturating_add(1);
execution.pending_interaction = Some(JournalInteraction {
request_id: request_id.into(),
prompt: prompt.into(),
});
let execution = execution.clone();
let event = push_event(
&mut state,
id,
EventPayload::InteractionRequired {
request_id: request_id.into(),
prompt: prompt.into(),
},
)?;
Ok(JournalMutation {
execution,
event,
replayed: false,
})
}
async fn commit_interaction_input(
&self,
id: &ExecutionId,
operation_id: &OperationId,
request_id: &str,
response_text: &str,
) -> Result<JournalMutation, JournalError> {
let mut state = self.state.lock().await;
let existing = state
.executions
.get(id)
.ok_or(JournalError::NotFound)?
.input_receipts
.get(operation_id)
.cloned();
if let Some(receipt) = existing {
if receipt.request_id != request_id || receipt.response_text != response_text {
return Err(JournalError::Conflict);
}
let event = state
.events
.get(id)
.and_then(|events| {
events
.iter()
.find(|event| event.sequence == receipt.event_sequence)
})
.cloned()
.ok_or_else(|| {
JournalError::Unavailable("input receipt event is missing".into())
})?;
return Ok(JournalMutation {
execution: state
.executions
.get(id)
.ok_or(JournalError::NotFound)?
.clone(),
event,
replayed: true,
});
}
ensure_event_capacity(&state, id, self.max_events_per_execution)?;
let event_sequence = state
.events
.get(id)
.map_or(1, |events| events.len() as u64 + 1);
let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
let matches_pending = execution
.pending_interaction
.as_ref()
.is_some_and(|pending| pending.request_id == request_id);
if execution.view.state != ExecutionState::WaitingForInput || !matches_pending {
return Err(JournalError::Conflict);
}
execution.input_receipts.insert(
operation_id.clone(),
JournalInputReceipt {
request_id: request_id.into(),
response_text: response_text.into(),
event_sequence,
},
);
execution.view.updated_at_ms = chrono::Utc::now().timestamp_millis();
execution.version = execution.version.saturating_add(1);
let execution = execution.clone();
let event = push_event(
&mut state,
id,
EventPayload::InteractionReceived {
request_id: request_id.into(),
},
)?;
Ok(JournalMutation {
execution,
event,
replayed: false,
})
}
async fn complete_interaction(
&self,
id: &ExecutionId,
request_id: &str,
) -> Result<JournalExecution, JournalError> {
let mut state = self.state.lock().await;
let execution = state.executions.get_mut(id).ok_or(JournalError::NotFound)?;
let matches_pending = execution
.pending_interaction
.as_ref()
.is_some_and(|pending| pending.request_id == request_id);
let has_response = execution
.input_receipts
.values()
.any(|receipt| receipt.request_id == request_id);
if execution.view.state != ExecutionState::WaitingForInput
|| !matches_pending
|| !has_response
{
return Err(JournalError::Conflict);
}
execution.pending_interaction = None;
execution.view.state = ExecutionState::Running;
execution.view.updated_at_ms = chrono::Utc::now().timestamp_millis();
execution.version = execution.version.saturating_add(1);
Ok(execution.clone())
}
async fn events(
&self,
id: &ExecutionId,
after: Option<u64>,
limit: usize,
) -> Result<Vec<ExecutionEvent>, JournalError> {
let state = self.state.lock().await;
if !state.executions.contains_key(id) {
return Err(JournalError::NotFound);
}
let after = after.unwrap_or(0);
Ok(state
.events
.get(id)
.into_iter()
.flatten()
.filter(|event| event.sequence > after)
.take(limit.min(1000))
.cloned()
.collect())
}
}
fn ensure_event_capacity(
state: &State,
id: &ExecutionId,
limit: usize,
) -> Result<(), JournalError> {
let events = state.events.get(id).ok_or(JournalError::NotFound)?;
if events.len() >= limit {
Err(JournalError::Unavailable(
"event retention limit exceeded".into(),
))
} else {
Ok(())
}
}
fn push_event(
state: &mut State,
id: &ExecutionId,
payload: EventPayload,
) -> Result<ExecutionEvent, JournalError> {
let events = state.events.get_mut(id).ok_or(JournalError::NotFound)?;
let event = ExecutionEvent {
execution_id: id.clone(),
sequence: events.len() as u64 + 1,
created_at_ms: chrono::Utc::now().timestamp_millis(),
payload,
};
events.push(event.clone());
Ok(event)
}
fn validate_transition(current: ExecutionState, next: ExecutionState) -> Result<(), JournalError> {
let valid = matches!(
(current, next),
(
ExecutionState::Queued,
ExecutionState::Running
| ExecutionState::Finalizing
| ExecutionState::Canceled
| ExecutionState::Failed
) | (
ExecutionState::Running,
ExecutionState::WaitingForInput
| ExecutionState::Finalizing
| ExecutionState::Failed
| ExecutionState::Canceled
) | (
ExecutionState::WaitingForInput,
ExecutionState::Running
| ExecutionState::Finalizing
| ExecutionState::Failed
| ExecutionState::Canceled
) | (
ExecutionState::Finalizing,
ExecutionState::Completed | ExecutionState::Failed | ExecutionState::Canceled
)
);
if valid {
Ok(())
} else {
Err(JournalError::Conflict)
}
}
#[cfg(test)]
mod tests {
use super::*;
use runtime_api_contract::{
CreateExecutionRequest, ExecutionOptions, ExecutionView, RuntimeInput,
};
use runtime_types::{CallerScope, RuntimeInstanceId};
fn new_execution(key: &str) -> NewJournalExecution {
let id = ExecutionId::random();
let request = CreateExecutionRequest {
runtime_instance_id: RuntimeInstanceId::new("runtime-1").unwrap(),
conversation_id: None,
input: RuntimeInput::UserMessage {
text: "hello".into(),
},
workspace_id: None,
model: None,
instructions: None,
metadata: Default::default(),
generation: Default::default(),
options: ExecutionOptions::default(),
};
NewJournalExecution {
idempotency_key: key.into(),
caller: CallerScope {
subject: "user-1".into(),
tenant_id: "tenant-1".into(),
project_id: "project-1".into(),
capabilities: Vec::new(),
},
request: request.clone(),
view: ExecutionView {
id,
runtime_instance_id: request.runtime_instance_id.clone(),
conversation_id: None,
workspace_id: None,
model: None,
metadata: Default::default(),
state: ExecutionState::Queued,
outcome: None,
failure: None,
created_at_ms: 1,
updated_at_ms: 1,
},
}
}
#[tokio::test]
async fn reservation_is_idempotent_and_transitions_are_guarded() {
let journal = MemoryExecutionJournal::new(10);
let created = journal.reserve(new_execution("key")).await.unwrap();
let id = match created {
JournalReservation::Created(value) => value.view.id,
_ => panic!("created"),
};
assert!(matches!(
journal.reserve(new_execution("key")).await,
Ok(JournalReservation::Existing(_))
));
let mut conflicting = new_execution("key");
conflicting.request.input = RuntimeInput::UserMessage {
text: "different".into(),
};
assert!(matches!(
journal.reserve(conflicting).await,
Err(JournalError::Conflict)
));
journal
.transition(&id, &[ExecutionState::Queued], ExecutionState::Running)
.await
.unwrap();
assert!(matches!(
journal
.transition(&id, &[ExecutionState::Queued], ExecutionState::Running)
.await,
Err(JournalError::Conflict)
));
}
#[tokio::test]
async fn claim_expiry_allows_exactly_one_recovery_worker() {
let journal = MemoryExecutionJournal::new(10);
let id = match journal.reserve(new_execution("claim-key")).await.unwrap() {
JournalReservation::Created(value) => value.view.id,
_ => panic!("created"),
};
journal.claim(&id, "worker-a", 100, 200).await.unwrap();
assert!(matches!(
journal.claim(&id, "worker-b", 150, 250).await,
Err(JournalError::Conflict)
));
assert!(journal.recoverable(199, 10).await.unwrap().is_empty());
assert_eq!(journal.recoverable(200, 10).await.unwrap().len(), 1);
let claimed = journal.claim(&id, "worker-b", 200, 300).await.unwrap();
assert_eq!(claimed.claim.unwrap().worker_id, "worker-b");
}
#[tokio::test]
async fn interaction_input_is_committed_before_resume_and_deduplicated() {
let journal = MemoryExecutionJournal::new(10);
let id = match journal.reserve(new_execution("interaction")).await.unwrap() {
JournalReservation::Created(value) => value.view.id,
_ => panic!("created"),
};
journal
.transition(&id, &[ExecutionState::Queued], ExecutionState::Running)
.await
.unwrap();
let required = journal
.begin_interaction(&id, "request-1", "approve?")
.await
.unwrap();
assert_eq!(
required.execution.view.state,
ExecutionState::WaitingForInput
);
let operation_id = OperationId::new("input-1").unwrap();
let committed = journal
.commit_interaction_input(&id, &operation_id, "request-1", "yes")
.await
.unwrap();
assert!(!committed.replayed);
assert_eq!(
committed.execution.view.state,
ExecutionState::WaitingForInput
);
let replay = journal
.commit_interaction_input(&id, &operation_id, "request-1", "yes")
.await
.unwrap();
assert!(replay.replayed);
assert_eq!(replay.event.sequence, committed.event.sequence);
assert!(matches!(
journal
.commit_interaction_input(&id, &operation_id, "request-1", "no")
.await,
Err(JournalError::Conflict)
));
let resumed = journal
.complete_interaction(&id, "request-1")
.await
.unwrap();
assert_eq!(resumed.view.state, ExecutionState::Running);
}
}