use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RuleSource {
Local,
Team,
Global,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct RecalledVerdict {
pub id: String,
pub title: String,
pub similarity: f32,
pub excerpt: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum TrajectoryStep {
ChunksRetrieved {
count: usize,
symbols: Vec<String>,
similarity_scores: Vec<f32>,
},
RulesApplied {
rule_ids: Vec<String>,
source: RuleSource,
},
LlmCall {
perspective: String,
input_tokens: u32,
output_tokens: u32,
#[serde(skip_serializing_if = "Option::is_none", default)]
raw_output: Option<String>,
},
PastVerdictsRecalled {
count: usize,
top_similarities: Vec<f32>,
#[serde(default)]
recalled_items: Vec<RecalledVerdict>,
},
SelfCheck {
keep_count: u32,
drop_count: u32,
avg_confidence: f32,
},
SignatureConfidenceAdjust {
accepted_bumps: u32,
rejected_bumps: u32,
},
FinalDecision { issue_ids_emitted: Vec<String> },
McpResponseSize {
tool: String,
total_tokens: usize,
rules_injected: usize,
},
RuleHitByOrigin {
manual: u32,
conversation: u32,
pr_review: u32,
extracted: u32,
cloud: u32,
},
RetrievalFilter { before: u32, after: u32 },
HybridFusion {
fts_hits: u32,
emb_hits: u32,
overlap: u32,
},
AnnRecall {
used: bool,
index_size: u32,
candidates: u32,
},
}
#[derive(Debug, Clone, Default)]
pub struct TrajectoryBuilder {
steps: Vec<TrajectoryStep>,
}
impl TrajectoryBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn push(&mut self, step: TrajectoryStep) {
self.steps.push(step);
}
pub const fn len(&self) -> usize {
self.steps.len()
}
pub const fn is_empty(&self) -> bool {
self.steps.is_empty()
}
pub fn steps(&self) -> &[TrajectoryStep] {
&self.steps
}
pub fn into_json(self) -> serde_json::Value {
serde_json::to_value(self.steps).unwrap_or(serde_json::Value::Array(vec![]))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn serialize_matches_ts_shape_exactly() {
let mut b = TrajectoryBuilder::new();
b.push(TrajectoryStep::ChunksRetrieved {
count: 2,
symbols: vec!["foo".into()],
similarity_scores: vec![0.91],
});
b.push(TrajectoryStep::SelfCheck {
keep_count: 3,
drop_count: 1,
avg_confidence: 0.82,
});
let value = b.into_json();
let arr = value.as_array().expect("top-level must be array");
assert_eq!(arr.len(), 2);
let first = arr[0].as_object().unwrap();
assert_eq!(
first.get("kind").and_then(|v| v.as_str()),
Some("chunks_retrieved")
);
assert_eq!(
first.get("count").and_then(serde_json::Value::as_u64),
Some(2)
);
assert!(first.contains_key("symbols"));
assert!(first.contains_key("similarity_scores"));
let second = arr[1].as_object().unwrap();
assert_eq!(
second.get("kind").and_then(|v| v.as_str()),
Some("self_check")
);
assert_eq!(
second.get("keep_count").and_then(serde_json::Value::as_u64),
Some(3)
);
assert_eq!(
second.get("drop_count").and_then(serde_json::Value::as_u64),
Some(1)
);
assert!(
(second
.get("avg_confidence")
.and_then(serde_json::Value::as_f64)
.unwrap()
- 0.82)
.abs()
< 1e-6
);
}
#[test]
fn llm_call_omits_raw_output_when_absent() {
let mut b = TrajectoryBuilder::new();
b.push(TrajectoryStep::LlmCall {
perspective: "safety".into(),
input_tokens: 123,
output_tokens: 45,
raw_output: None,
});
let value = b.into_json();
let obj = value.as_array().unwrap()[0].as_object().unwrap();
assert_eq!(
obj.get("perspective").and_then(|v| v.as_str()),
Some("safety")
);
assert!(!obj.contains_key("raw_output"));
}
#[test]
fn full_pipeline_shape_matches_plan_capture_points() {
let mut b = TrajectoryBuilder::new();
b.push(TrajectoryStep::ChunksRetrieved {
count: 4,
symbols: vec!["foo".into()],
similarity_scores: vec![],
});
b.push(TrajectoryStep::RulesApplied {
rule_ids: vec!["r1".into(), "r2".into()],
source: RuleSource::Team,
});
b.push(TrajectoryStep::PastVerdictsRecalled {
count: 2,
top_similarities: vec![],
recalled_items: vec![],
});
for p in ["safety", "performance", "style", "docs", "api_design"] {
b.push(TrajectoryStep::LlmCall {
perspective: p.to_owned(),
input_tokens: 200,
output_tokens: 0,
raw_output: None,
});
}
b.push(TrajectoryStep::SelfCheck {
keep_count: 3,
drop_count: 1,
avg_confidence: 0.87,
});
b.push(TrajectoryStep::FinalDecision {
issue_ids_emitted: vec!["issue-1".into(), "issue-2".into(), "issue-3".into()],
});
assert_eq!(b.len(), 1 + 1 + 1 + 5 + 1 + 1);
let kinds: Vec<&str> = b
.steps()
.iter()
.map(|s| match s {
TrajectoryStep::ChunksRetrieved { .. } => "chunks_retrieved",
TrajectoryStep::RulesApplied { .. } => "rules_applied",
TrajectoryStep::PastVerdictsRecalled { .. } => "past_verdicts_recalled",
TrajectoryStep::LlmCall { .. } => "llm_call",
TrajectoryStep::SelfCheck { .. } => "self_check",
TrajectoryStep::SignatureConfidenceAdjust { .. } => "signature_confidence_adjust",
TrajectoryStep::FinalDecision { .. } => "final_decision",
TrajectoryStep::McpResponseSize { .. } => "mcp_response_size",
TrajectoryStep::RuleHitByOrigin { .. } => "rule_hit_by_origin",
TrajectoryStep::RetrievalFilter { .. } => "retrieval_filter",
TrajectoryStep::HybridFusion { .. } => "hybrid_fusion",
TrajectoryStep::AnnRecall { .. } => "ann_recall",
})
.collect();
assert_eq!(
kinds,
vec![
"chunks_retrieved",
"rules_applied",
"past_verdicts_recalled",
"llm_call",
"llm_call",
"llm_call",
"llm_call",
"llm_call",
"self_check",
"final_decision",
]
);
}
#[test]
fn mcp_response_size_and_rule_hit_by_origin_serialize() {
let mut b = TrajectoryBuilder::new();
b.push(TrajectoryStep::McpResponseSize {
tool: "search_rules".into(),
total_tokens: 1234,
rules_injected: 3,
});
b.push(TrajectoryStep::RuleHitByOrigin {
manual: 1,
conversation: 2,
pr_review: 0,
extracted: 1,
cloud: 0,
});
let value = b.clone().into_json();
let arr = value.as_array().expect("top-level array");
assert_eq!(arr[0]["kind"], "mcp_response_size");
assert_eq!(arr[0]["tool"], "search_rules");
assert_eq!(arr[0]["total_tokens"], 1234);
assert_eq!(arr[0]["rules_injected"], 3);
assert_eq!(arr[1]["kind"], "rule_hit_by_origin");
assert_eq!(arr[1]["manual"], 1);
assert_eq!(arr[1]["conversation"], 2);
assert_eq!(arr[1]["pr_review"], 0);
assert_eq!(arr[1]["extracted"], 1);
assert_eq!(arr[1]["cloud"], 0);
let text = serde_json::to_string(&value).unwrap();
let parsed: Vec<TrajectoryStep> = serde_json::from_str(&text).unwrap();
assert_eq!(parsed, b.steps().to_vec());
}
#[test]
fn round_trip_deserialize_via_serde_json() {
let mut b = TrajectoryBuilder::new();
b.push(TrajectoryStep::PastVerdictsRecalled {
count: 4,
top_similarities: vec![0.95, 0.88, 0.80, 0.72],
recalled_items: vec![RecalledVerdict {
id: "verdict-1".into(),
title: "avoid unwrap in request handlers".into(),
similarity: 0.95,
excerpt: "fn handler() { ... .unwrap() ... }".into(),
}],
});
b.push(TrajectoryStep::RulesApplied {
rule_ids: vec!["r1".into(), "r2".into()],
source: RuleSource::Global,
});
b.push(TrajectoryStep::FinalDecision {
issue_ids_emitted: vec!["issue-1".into()],
});
let value = b.clone().into_json();
let text = serde_json::to_string(&value).unwrap();
let parsed: Vec<TrajectoryStep> = serde_json::from_str(&text).unwrap();
assert_eq!(parsed, b.steps().to_vec());
}
}