use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use crate::crypto::{hash, Hash, PublicKey};
use crate::error::{Error, Result};
use crate::event::EventId;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TraceId(pub [u8; 16]);
impl TraceId {
pub fn generate() -> Self {
use rand::RngCore;
let mut bytes = [0u8; 16];
rand::thread_rng().fill_bytes(&mut bytes);
Self(bytes)
}
pub fn from_bytes(bytes: [u8; 16]) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8; 16] {
&self.0
}
pub fn to_hex(&self) -> String {
hex::encode(self.0)
}
pub fn from_hex(s: &str) -> Result<Self> {
let bytes = hex::decode(s).map_err(|_| Error::invalid_input("invalid hex"))?;
if bytes.len() != 16 {
return Err(Error::invalid_input("trace ID must be 16 bytes"));
}
let mut arr = [0u8; 16];
arr.copy_from_slice(&bytes);
Ok(Self(arr))
}
}
impl std::fmt::Display for TraceId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.to_hex())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Priority {
Low,
Normal,
High,
Critical,
}
impl Default for Priority {
fn default() -> Self {
Self::Normal
}
}
impl std::fmt::Display for Priority {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Priority::Low => write!(f, "low"),
Priority::Normal => write!(f, "normal"),
Priority::High => write!(f, "high"),
Priority::Critical => write!(f, "critical"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum GoalSource {
UserInstruction { event_id: EventId },
Derived { parent_goal_description: String },
System { policy: String },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Goal {
description: String,
structured: Option<serde_json::Value>,
source: GoalSource,
priority: Priority,
}
impl Goal {
pub fn from_user(description: impl Into<String>, event_id: EventId) -> Self {
Self {
description: description.into(),
structured: None,
source: GoalSource::UserInstruction { event_id },
priority: Priority::Normal,
}
}
pub fn derived(description: impl Into<String>, parent_description: impl Into<String>) -> Self {
Self {
description: description.into(),
structured: None,
source: GoalSource::Derived {
parent_goal_description: parent_description.into(),
},
priority: Priority::Normal,
}
}
pub fn system(description: impl Into<String>, policy: impl Into<String>) -> Self {
Self {
description: description.into(),
structured: None,
source: GoalSource::System {
policy: policy.into(),
},
priority: Priority::Normal,
}
}
pub fn with_structured(mut self, structured: serde_json::Value) -> Self {
self.structured = Some(structured);
self
}
pub fn with_priority(mut self, priority: Priority) -> Self {
self.priority = priority;
self
}
pub fn description(&self) -> &str {
&self.description
}
pub fn structured(&self) -> Option<&serde_json::Value> {
self.structured.as_ref()
}
pub fn source(&self) -> &GoalSource {
&self.source
}
pub fn priority(&self) -> Priority {
self.priority
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum StepAction {
Retrieve { query: String, source: String },
Analyze { subject: String, method: String },
ToolCall { tool: String, input_hash: Hash },
Delegate { agent: PublicKey, task: String },
Decide { decision: String },
}
impl StepAction {
pub fn retrieve(query: impl Into<String>, source: impl Into<String>) -> Self {
Self::Retrieve {
query: query.into(),
source: source.into(),
}
}
pub fn analyze(subject: impl Into<String>, method: impl Into<String>) -> Self {
Self::Analyze {
subject: subject.into(),
method: method.into(),
}
}
pub fn tool_call(tool: impl Into<String>, input_hash: Hash) -> Self {
Self::ToolCall {
tool: tool.into(),
input_hash,
}
}
pub fn delegate(agent: PublicKey, task: impl Into<String>) -> Self {
Self::Delegate {
agent,
task: task.into(),
}
}
pub fn decide(decision: impl Into<String>) -> Self {
Self::Decide {
decision: decision.into(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReasoningStep {
sequence: u32,
thought: String,
action: Option<StepAction>,
observation: Option<String>,
timestamp: i64,
}
impl ReasoningStep {
pub fn new(sequence: u32, thought: impl Into<String>) -> Self {
Self {
sequence,
thought: thought.into(),
action: None,
observation: None,
timestamp: chrono::Utc::now().timestamp_millis(),
}
}
pub fn with_action(mut self, action: StepAction) -> Self {
self.action = Some(action);
self
}
pub fn with_observation(mut self, observation: impl Into<String>) -> Self {
self.observation = Some(observation.into());
self
}
pub fn with_timestamp(mut self, timestamp: i64) -> Self {
self.timestamp = timestamp;
self
}
pub fn sequence(&self) -> u32 {
self.sequence
}
pub fn thought(&self) -> &str {
&self.thought
}
pub fn action(&self) -> Option<&StepAction> {
self.action.as_ref()
}
pub fn observation(&self) -> Option<&str> {
self.observation.as_deref()
}
pub fn timestamp(&self) -> i64 {
self.timestamp
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Decision {
action: String,
rationale: String,
expected_outcome: String,
success_criteria: Vec<String>,
}
impl Decision {
pub fn new(
action: impl Into<String>,
rationale: impl Into<String>,
expected_outcome: impl Into<String>,
) -> Self {
Self {
action: action.into(),
rationale: rationale.into(),
expected_outcome: expected_outcome.into(),
success_criteria: Vec::new(),
}
}
pub fn with_criterion(mut self, criterion: impl Into<String>) -> Self {
self.success_criteria.push(criterion.into());
self
}
pub fn with_criteria(mut self, criteria: Vec<String>) -> Self {
self.success_criteria = criteria;
self
}
pub fn action(&self) -> &str {
&self.action
}
pub fn rationale(&self) -> &str {
&self.rationale
}
pub fn expected_outcome(&self) -> &str {
&self.expected_outcome
}
pub fn success_criteria(&self) -> &[String] {
&self.success_criteria
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Confidence {
score: f64,
breakdown: HashMap<String, f64>,
uncertainties: Vec<String>,
would_help: Vec<String>,
}
impl Confidence {
pub fn new(score: f64) -> Self {
Self {
score: score.clamp(0.0, 1.0),
breakdown: HashMap::new(),
uncertainties: Vec::new(),
would_help: Vec::new(),
}
}
pub fn high() -> Self {
Self::new(0.9)
}
pub fn medium() -> Self {
Self::new(0.7)
}
pub fn low() -> Self {
Self::new(0.4)
}
pub fn with_factor(mut self, factor: impl Into<String>, score: f64) -> Self {
self.breakdown.insert(factor.into(), score.clamp(0.0, 1.0));
self
}
pub fn with_uncertainty(mut self, uncertainty: impl Into<String>) -> Self {
self.uncertainties.push(uncertainty.into());
self
}
pub fn with_would_help(mut self, help: impl Into<String>) -> Self {
self.would_help.push(help.into());
self
}
pub fn score(&self) -> f64 {
self.score
}
pub fn breakdown(&self) -> &HashMap<String, f64> {
&self.breakdown
}
pub fn uncertainties(&self) -> &[String] {
&self.uncertainties
}
pub fn would_help(&self) -> &[String] {
&self.would_help
}
pub fn should_reject(&self) -> bool {
self.score < 0.3
}
pub fn requires_approval(&self) -> bool {
self.score < 0.5
}
pub fn should_warn(&self) -> bool {
self.score < 0.7
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Alternative {
description: String,
rejection_reason: String,
estimated_outcome: String,
confidence: f64,
}
impl Alternative {
pub fn new(
description: impl Into<String>,
rejection_reason: impl Into<String>,
estimated_outcome: impl Into<String>,
confidence: f64,
) -> Self {
Self {
description: description.into(),
rejection_reason: rejection_reason.into(),
estimated_outcome: estimated_outcome.into(),
confidence: confidence.clamp(0.0, 1.0),
}
}
pub fn description(&self) -> &str {
&self.description
}
pub fn rejection_reason(&self) -> &str {
&self.rejection_reason
}
pub fn estimated_outcome(&self) -> &str {
&self.estimated_outcome
}
pub fn confidence(&self) -> f64 {
self.confidence
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Factor {
description: String,
influence: f64,
evidence: Vec<String>,
}
impl Factor {
pub fn new(description: impl Into<String>, influence: f64) -> Self {
Self {
description: description.into(),
influence: influence.clamp(-1.0, 1.0),
evidence: Vec::new(),
}
}
pub fn positive(description: impl Into<String>, influence: f64) -> Self {
Self::new(description, influence.abs())
}
pub fn negative(description: impl Into<String>, influence: f64) -> Self {
Self::new(description, -influence.abs())
}
pub fn with_evidence(mut self, evidence: impl Into<String>) -> Self {
self.evidence.push(evidence.into());
self
}
pub fn description(&self) -> &str {
&self.description
}
pub fn influence(&self) -> f64 {
self.influence
}
pub fn evidence(&self) -> &[String] {
&self.evidence
}
pub fn is_supportive(&self) -> bool {
self.influence > 0.0
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReasoningTrace {
id: TraceId,
goal: Goal,
steps: Vec<ReasoningStep>,
decision: Decision,
confidence: Confidence,
alternatives: Vec<Alternative>,
factors: Vec<Factor>,
trace_hash: Hash,
}
impl ReasoningTrace {
pub fn builder() -> ReasoningTraceBuilder {
ReasoningTraceBuilder::new()
}
pub fn id(&self) -> TraceId {
self.id
}
pub fn goal(&self) -> &Goal {
&self.goal
}
pub fn steps(&self) -> &[ReasoningStep] {
&self.steps
}
pub fn decision(&self) -> &Decision {
&self.decision
}
pub fn confidence(&self) -> &Confidence {
&self.confidence
}
pub fn alternatives(&self) -> &[Alternative] {
&self.alternatives
}
pub fn factors(&self) -> &[Factor] {
&self.factors
}
pub fn trace_hash(&self) -> Hash {
self.trace_hash
}
pub fn canonical_bytes(&self) -> Vec<u8> {
let mut data = Vec::new();
data.extend_from_slice(&self.id.0);
let goal_json = serde_json::to_vec(&self.goal).unwrap_or_default();
data.extend_from_slice(&goal_json);
for step in &self.steps {
let step_json = serde_json::to_vec(step).unwrap_or_default();
data.extend_from_slice(&step_json);
}
let decision_json = serde_json::to_vec(&self.decision).unwrap_or_default();
data.extend_from_slice(&decision_json);
data.extend_from_slice(&self.confidence.score.to_le_bytes());
for alt in &self.alternatives {
let alt_json = serde_json::to_vec(alt).unwrap_or_default();
data.extend_from_slice(&alt_json);
}
for factor in &self.factors {
let factor_json = serde_json::to_vec(factor).unwrap_or_default();
data.extend_from_slice(&factor_json);
}
data
}
pub fn verify_integrity(&self) -> bool {
let computed = hash(&self.canonical_bytes());
computed == self.trace_hash
}
pub fn is_complete(&self) -> bool {
if self.goal.description.is_empty() {
return false;
}
if self.steps.is_empty() {
return false;
}
if self.decision.action.is_empty() {
return false;
}
if !(0.0..=1.0).contains(&self.confidence.score) {
return false;
}
true
}
}
#[derive(Debug, Default)]
pub struct ReasoningTraceBuilder {
id: Option<TraceId>,
goal: Option<Goal>,
steps: Vec<ReasoningStep>,
decision: Option<Decision>,
confidence: Option<Confidence>,
alternatives: Vec<Alternative>,
factors: Vec<Factor>,
}
impl ReasoningTraceBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn id(mut self, id: TraceId) -> Self {
self.id = Some(id);
self
}
pub fn goal(mut self, goal: Goal) -> Self {
self.goal = Some(goal);
self
}
pub fn step(mut self, step: ReasoningStep) -> Self {
self.steps.push(step);
self
}
pub fn steps(mut self, steps: Vec<ReasoningStep>) -> Self {
self.steps = steps;
self
}
pub fn decision(mut self, decision: Decision) -> Self {
self.decision = Some(decision);
self
}
pub fn confidence(mut self, confidence: Confidence) -> Self {
self.confidence = Some(confidence);
self
}
pub fn alternative(mut self, alternative: Alternative) -> Self {
self.alternatives.push(alternative);
self
}
pub fn factor(mut self, factor: Factor) -> Self {
self.factors.push(factor);
self
}
pub fn build(self) -> Result<ReasoningTrace> {
let id = self.id.unwrap_or_else(TraceId::generate);
let goal = self
.goal
.ok_or_else(|| Error::invalid_input("goal is required"))?;
if self.steps.is_empty() {
return Err(Error::invalid_input(
"at least one reasoning step is required",
));
}
let decision = self
.decision
.ok_or_else(|| Error::invalid_input("decision is required"))?;
let confidence = self.confidence.unwrap_or_else(Confidence::medium);
let mut trace = ReasoningTrace {
id,
goal,
steps: self.steps,
decision,
confidence,
alternatives: self.alternatives,
factors: self.factors,
trace_hash: Hash::from_bytes([0u8; 32]), };
trace.trace_hash = hash(&trace.canonical_bytes());
Ok(trace)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn test_event_id() -> EventId {
EventId(hash(b"test-event"))
}
#[test]
fn trace_id_generates_unique() {
let id1 = TraceId::generate();
let id2 = TraceId::generate();
assert_ne!(id1, id2);
}
#[test]
fn trace_id_hex_roundtrip() {
let id = TraceId::generate();
let hex = id.to_hex();
let restored = TraceId::from_hex(&hex).unwrap();
assert_eq!(id, restored);
}
#[test]
fn priority_ordering() {
assert!(Priority::Low < Priority::Normal);
assert!(Priority::Normal < Priority::High);
assert!(Priority::High < Priority::Critical);
}
#[test]
fn goal_from_user() {
let goal = Goal::from_user("Complete the task", test_event_id());
assert_eq!(goal.description(), "Complete the task");
assert!(matches!(goal.source(), GoalSource::UserInstruction { .. }));
}
#[test]
fn goal_derived() {
let goal = Goal::derived("Sub-task", "Parent task");
assert_eq!(goal.description(), "Sub-task");
assert!(matches!(goal.source(), GoalSource::Derived { .. }));
}
#[test]
fn goal_with_priority() {
let goal =
Goal::from_user("Urgent task", test_event_id()).with_priority(Priority::Critical);
assert_eq!(goal.priority(), Priority::Critical);
}
#[test]
fn reasoning_step_basic() {
let step = ReasoningStep::new(1, "Analyzing the problem");
assert_eq!(step.sequence(), 1);
assert_eq!(step.thought(), "Analyzing the problem");
assert!(step.action().is_none());
}
#[test]
fn reasoning_step_with_action() {
let step = ReasoningStep::new(1, "Looking up information")
.with_action(StepAction::retrieve("user data", "database"));
assert!(step.action().is_some());
}
#[test]
fn reasoning_step_with_observation() {
let step = ReasoningStep::new(1, "Checking status").with_observation("Status is active");
assert_eq!(step.observation(), Some("Status is active"));
}
#[test]
fn decision_basic() {
let decision = Decision::new(
"Proceed with option A",
"It has the highest success probability",
"Task completed successfully",
);
assert_eq!(decision.action(), "Proceed with option A");
assert_eq!(
decision.rationale(),
"It has the highest success probability"
);
}
#[test]
fn decision_with_criteria() {
let decision = Decision::new("Execute", "Best option", "Success")
.with_criterion("Output matches expected format")
.with_criterion("No errors in logs");
assert_eq!(decision.success_criteria().len(), 2);
}
#[test]
fn confidence_clamped() {
let conf = Confidence::new(1.5);
assert_eq!(conf.score(), 1.0);
let conf2 = Confidence::new(-0.5);
assert_eq!(conf2.score(), 0.0);
}
#[test]
fn confidence_thresholds() {
let low = Confidence::new(0.2);
assert!(low.should_reject());
assert!(low.requires_approval());
assert!(low.should_warn());
let medium = Confidence::new(0.6);
assert!(!medium.should_reject());
assert!(!medium.requires_approval());
assert!(medium.should_warn());
let high = Confidence::new(0.8);
assert!(!high.should_reject());
assert!(!high.requires_approval());
assert!(!high.should_warn());
}
#[test]
fn confidence_with_breakdown() {
let conf = Confidence::medium()
.with_factor("data_quality", 0.8)
.with_factor("model_accuracy", 0.6)
.with_uncertainty("Limited training data");
assert_eq!(conf.breakdown().len(), 2);
assert_eq!(conf.uncertainties().len(), 1);
}
#[test]
fn alternative_basic() {
let alt = Alternative::new(
"Option B",
"Higher cost",
"Same result but more expensive",
0.6,
);
assert_eq!(alt.description(), "Option B");
assert_eq!(alt.rejection_reason(), "Higher cost");
}
#[test]
fn factor_positive() {
let factor = Factor::positive("Strong evidence", 0.8);
assert!(factor.is_supportive());
assert_eq!(factor.influence(), 0.8);
}
#[test]
fn factor_negative() {
let factor = Factor::negative("Risk of failure", 0.3);
assert!(!factor.is_supportive());
assert_eq!(factor.influence(), -0.3);
}
#[test]
fn factor_with_evidence() {
let factor = Factor::positive("Proven approach", 0.9)
.with_evidence("Study A shows 95% success rate")
.with_evidence("Historical data confirms");
assert_eq!(factor.evidence().len(), 2);
}
#[test]
fn reasoning_trace_requires_goal() {
let result = ReasoningTrace::builder()
.step(ReasoningStep::new(1, "Thinking"))
.decision(Decision::new("Do it", "Because", "Success"))
.build();
assert!(result.is_err());
}
#[test]
fn reasoning_trace_requires_steps() {
let result = ReasoningTrace::builder()
.goal(Goal::from_user("Task", test_event_id()))
.decision(Decision::new("Do it", "Because", "Success"))
.build();
assert!(result.is_err());
}
#[test]
fn reasoning_trace_requires_decision() {
let result = ReasoningTrace::builder()
.goal(Goal::from_user("Task", test_event_id()))
.step(ReasoningStep::new(1, "Thinking"))
.build();
assert!(result.is_err());
}
#[test]
fn reasoning_trace_complete() {
let trace = ReasoningTrace::builder()
.goal(Goal::from_user("Complete the task", test_event_id()))
.step(ReasoningStep::new(1, "Analyzing requirements"))
.step(ReasoningStep::new(2, "Evaluating options"))
.decision(Decision::new(
"Use approach A",
"Most efficient",
"Task completed",
))
.confidence(Confidence::high())
.alternative(Alternative::new("Approach B", "Slower", "Same result", 0.7))
.factor(Factor::positive("Clear requirements", 0.9))
.build()
.unwrap();
assert!(trace.is_complete());
assert_eq!(trace.steps().len(), 2);
assert_eq!(trace.alternatives().len(), 1);
assert_eq!(trace.factors().len(), 1);
}
#[test]
fn reasoning_trace_hash_computed() {
let trace = ReasoningTrace::builder()
.goal(Goal::from_user("Task", test_event_id()))
.step(ReasoningStep::new(1, "Thinking"))
.decision(Decision::new("Do it", "Because", "Success"))
.build()
.unwrap();
assert_ne!(trace.trace_hash().as_bytes(), &[0u8; 32]);
}
#[test]
fn reasoning_trace_integrity_verification() {
let trace = ReasoningTrace::builder()
.goal(Goal::from_user("Task", test_event_id()))
.step(ReasoningStep::new(1, "Thinking"))
.decision(Decision::new("Do it", "Because", "Success"))
.build()
.unwrap();
assert!(trace.verify_integrity());
}
#[test]
fn reasoning_trace_tamper_detected() {
let mut trace = ReasoningTrace::builder()
.goal(Goal::from_user("Task", test_event_id()))
.step(ReasoningStep::new(1, "Thinking"))
.decision(Decision::new("Do it", "Because", "Success"))
.build()
.unwrap();
trace.decision = Decision::new("Do something else", "Changed", "Different");
assert!(!trace.verify_integrity());
}
#[test]
fn step_action_variants() {
let retrieve = StepAction::retrieve("query", "db");
assert!(matches!(retrieve, StepAction::Retrieve { .. }));
let analyze = StepAction::analyze("data", "statistical");
assert!(matches!(analyze, StepAction::Analyze { .. }));
let tool = StepAction::tool_call("bash", hash(b"input"));
assert!(matches!(tool, StepAction::ToolCall { .. }));
let decide = StepAction::decide("Go ahead");
assert!(matches!(decide, StepAction::Decide { .. }));
}
}