use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use entelix_core::error::Result;
use entelix_core::{TenantId, ThreadKey};
#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct CheckpointId(uuid::Uuid);
impl CheckpointId {
pub fn new() -> Self {
Self(uuid::Uuid::now_v7())
}
pub const fn from_uuid(uuid: uuid::Uuid) -> Self {
Self(uuid)
}
pub const fn as_uuid(&self) -> &uuid::Uuid {
&self.0
}
pub fn to_hyphenated_string(&self) -> String {
self.0.to_string()
}
}
impl Default for CheckpointId {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct Checkpoint<S>
where
S: Clone + Send + Sync + 'static,
{
pub id: CheckpointId,
pub tenant_id: TenantId,
pub thread_id: String,
pub parent_id: Option<CheckpointId>,
pub step: usize,
pub state: S,
pub next_node: Option<String>,
pub timestamp: chrono::DateTime<chrono::Utc>,
}
impl<S> Checkpoint<S>
where
S: Clone + Send + Sync + 'static,
{
#[must_use]
pub fn new(key: &ThreadKey, step: usize, state: S, next_node: Option<String>) -> Self {
Self {
id: CheckpointId::new(),
tenant_id: key.tenant_id().clone(),
thread_id: key.thread_id().to_owned(),
parent_id: None,
step,
state,
next_node,
timestamp: chrono::Utc::now(),
}
}
#[must_use]
pub const fn with_parent(mut self, parent_id: CheckpointId) -> Self {
self.parent_id = Some(parent_id);
self
}
#[must_use]
#[allow(clippy::too_many_arguments)]
pub fn from_parts(
id: CheckpointId,
key: &ThreadKey,
parent_id: Option<CheckpointId>,
step: usize,
state: S,
next_node: Option<String>,
timestamp: chrono::DateTime<chrono::Utc>,
) -> Self {
Self {
id,
tenant_id: key.tenant_id().to_owned(),
thread_id: key.thread_id().to_owned(),
parent_id,
step,
state,
next_node,
timestamp,
}
}
#[must_use]
pub fn key(&self) -> ThreadKey {
ThreadKey::new(self.tenant_id.clone(), self.thread_id.clone())
}
}
#[async_trait]
pub trait Checkpointer<S>: Send + Sync + 'static
where
S: Clone + Send + Sync + 'static,
{
async fn put(&self, checkpoint: Checkpoint<S>) -> Result<()>;
async fn get_latest(&self, key: &ThreadKey) -> Result<Option<Checkpoint<S>>>;
async fn get_by_id(&self, key: &ThreadKey, id: &CheckpointId) -> Result<Option<Checkpoint<S>>>;
async fn list_history(&self, key: &ThreadKey, limit: usize) -> Result<Vec<Checkpoint<S>>>;
async fn update_state(
&self,
key: &ThreadKey,
parent_id: &CheckpointId,
new_state: S,
) -> Result<CheckpointId>;
}