use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use thiserror::Error;
use crate::execution::{
ConversationSurface, DurableEvent, InvocationContext, Message, QueuedRun, RunId, Scope,
ThreadId,
};
#[derive(Clone, Copy, Debug, Default, Eq, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
#[serde(transparent)]
pub struct Revision(pub u64);
impl Revision {
pub const INITIAL: Self = Self(0);
#[must_use]
pub fn next(self) -> Option<Self> {
self.0.checked_add(1).map(Self)
}
}
#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
pub struct ThreadKey {
pub scope: Scope,
pub surface_id: ConversationSurface,
pub thread_id: ThreadId,
}
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
pub struct Checkpoint {
pub format_version: u32,
pub bytes: Vec<u8>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct RuntimeSnapshot {
pub checkpoint: Checkpoint,
pub history: Vec<Message>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ThreadRecord {
pub revision: Revision,
pub checkpoint: Checkpoint,
pub history: Vec<Message>,
pub pending_event: Option<DurableEvent>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", content = "messages", rename_all = "snake_case")]
pub enum HistoryMutation {
Append(Vec<Message>),
Replace(Vec<Message>),
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ThreadCommit {
pub fencing_token: Option<u64>,
pub checkpoint: Checkpoint,
pub history: HistoryMutation,
pub acknowledge_event: Option<crate::execution::EventId>,
pub event: Option<DurableEvent>,
pub recovery_context: Option<InvocationContext>,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum CommitOutcome {
Applied {
revision: Revision,
},
Conflict {
actual: Revision,
},
StaleFence {
actual: u64,
},
}
#[derive(Clone, Debug, Error, Eq, PartialEq)]
pub enum StoreError {
#[error("invalid durable state transition: {message}")]
InvalidInput {
message: String,
},
#[error("corrupt durable state: {message}")]
Corrupt {
message: String,
},
#[error("durable state access denied")]
PermissionDenied,
#[error("durable state unavailable: {message}")]
Unavailable {
message: String,
},
#[error("durable state failure: {message}")]
Internal {
message: String,
},
}
#[async_trait]
pub trait ThreadStore: Send + Sync {
async fn load(&self, key: &ThreadKey) -> Result<Option<ThreadRecord>, StoreError>;
async fn commit(
&self,
key: &ThreadKey,
expected: Revision,
transition: ThreadCommit,
) -> Result<CommitOutcome, StoreError>;
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct QueueSnapshot {
pub revision: Revision,
pub items: Vec<QueuedRun>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ClaimNextOutcome {
Empty,
Claimed {
revision: Revision,
request: Box<QueuedRun>,
},
Conflict {
actual: Revision,
},
}
#[async_trait]
pub trait QueueStore: Send + Sync {
async fn load(&self, key: &ThreadKey) -> Result<QueueSnapshot, StoreError>;
async fn enqueue(
&self,
key: &ThreadKey,
request: QueuedRun,
) -> Result<QueueSnapshot, StoreError>;
async fn claim_next(
&self,
key: &ThreadKey,
expected: Revision,
claimant: &str,
) -> Result<ClaimNextOutcome, StoreError>;
async fn ack_started(
&self,
key: &ThreadKey,
run_id: &RunId,
) -> Result<QueueSnapshot, StoreError>;
async fn remove(&self, key: &ThreadKey, run_id: &RunId) -> Result<QueueSnapshot, StoreError>;
}
#[async_trait]
pub trait RunControl: Send + Sync {
async fn request_cancel(
&self,
key: &ThreadKey,
context: &InvocationContext,
) -> Result<(), StoreError>;
async fn is_cancel_requested(
&self,
key: &ThreadKey,
run_id: &RunId,
) -> Result<bool, StoreError>;
async fn wait_for_cancel(&self, key: &ThreadKey, run_id: &RunId) -> Result<(), StoreError>;
async fn clear_cancel(&self, key: &ThreadKey, run_id: &RunId) -> Result<(), StoreError>;
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct RunLease {
pub key: ThreadKey,
pub run_id: RunId,
pub fencing_token: u64,
pub expires_at_unix_ms: u64,
}
#[async_trait]
pub trait RunCoordinator: Send + Sync {
async fn acquire(
&self,
key: &ThreadKey,
run_id: &RunId,
lease_ms: u64,
) -> Result<Option<RunLease>, StoreError>;
async fn renew(&self, lease: &RunLease, lease_ms: u64) -> Result<Option<RunLease>, StoreError>;
async fn release(&self, lease: &RunLease) -> Result<(), StoreError>;
}