use std::time::Duration;
use uuid::Uuid;
use crate::error::{GraphError, ObservedError};
use crate::state::{GraphResult, State};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TraceId(Uuid);
impl TraceId {
pub fn new() -> Self {
Self(Uuid::new_v4())
}
pub fn to_string(&self) -> String {
self.0.to_string()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SpanId(Uuid);
impl SpanId {
pub fn new() -> Self {
Self(Uuid::new_v4())
}
pub fn to_string(&self) -> String {
self.0.to_string()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct BarrierId {
pub node_id: String,
pub occurrence: u32,
}
impl BarrierId {
pub fn new(node_id: impl Into<String>, occurrence: u32) -> Self {
Self {
node_id: node_id.into(),
occurrence,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum EventLevel {
Graph,
Node,
Agent,
Debug,
}
#[derive(Debug)]
pub enum NodeEvent {
Agent(lellm_agent::AgentEvent),
Barrier(BarrierInnerEvent),
}
#[derive(Debug, Clone)]
pub enum BarrierInnerEvent {
StateChange { from: String, to: String },
}
#[derive(Debug, Clone)]
pub enum BarrierDecision {
Approve,
Reject { reason: String },
Modify { key: String, value: serde_json::Value },
Reroute { target: String },
}
#[derive(Debug)]
pub enum GraphEvent {
NodeStart {
node_name: String,
span_id: SpanId,
step: usize,
},
NodeEnd {
node_name: String,
span_id: SpanId,
success: bool,
duration: Duration,
},
Node {
span_id: SpanId,
node_name: String,
event: NodeEvent,
},
BarrierWaiting {
barrier_id: BarrierId,
node_name: String,
span_id: SpanId,
},
BarrierResolved {
barrier_id: BarrierId,
decision: BarrierDecision,
},
ObservedError {
error: ObservedError,
node_name: String,
},
GraphComplete {
result: GraphResult,
},
GraphError {
error: GraphError,
state: State,
},
}
pub type GraphStream = tokio::sync::mpsc::Receiver<GraphEvent>;
pub struct GraphExecution {
pub stream: GraphStream,
pub handle: GraphHandle,
}
pub struct GraphHandle {
decision_tx: tokio::sync::mpsc::Sender<BarrierDecisionMessage>,
cancel_tx: tokio::sync::mpsc::Sender<()>,
}
#[allow(dead_code)]
pub(crate) enum BarrierDecisionMessage {
Exact { barrier_id: BarrierId, decision: BarrierDecision },
Wildcard { node_id: String, decision: BarrierDecision },
}
impl GraphHandle {
pub(crate) fn new(
decision_tx: tokio::sync::mpsc::Sender<BarrierDecisionMessage>,
cancel_tx: tokio::sync::mpsc::Sender<()>,
) -> Self {
Self {
decision_tx,
cancel_tx,
}
}
pub async fn decide(
&self,
barrier_id: BarrierId,
decision: BarrierDecision,
) -> Result<(), GraphError> {
self.decision_tx
.send(BarrierDecisionMessage::Exact {
barrier_id,
decision,
})
.await
.map_err(|_| GraphError::Terminal(crate::error::TerminalError::BarrierCancelled {
node: "decision channel closed".into(),
}))
}
pub async fn decide_wildcard(
&self,
node_id: impl Into<String>,
decision: BarrierDecision,
) -> Result<(), GraphError> {
self.decision_tx
.send(BarrierDecisionMessage::Wildcard {
node_id: node_id.into(),
decision,
})
.await
.map_err(|_| GraphError::Terminal(crate::error::TerminalError::BarrierCancelled {
node: "decision channel closed".into(),
}))
}
pub fn cancel(&self) {
let _ = self.cancel_tx.try_send(());
}
}