Skip to main content

assay_runner_schema/
correlation.rs

1use serde::{Deserialize, Serialize};
2use thiserror::Error;
3
4pub const CORRELATION_REPORT_SCHEMA: &str = "assay.runner.correlation_report.v0";
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
7#[serde(rename_all = "snake_case")]
8pub enum CorrelationStatus {
9    Clean,
10    Partial,
11    Failed,
12}
13
14#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
15pub struct BindingWindow {
16    pub start: String,
17    pub end: String,
18}
19
20#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21pub struct CorrelationBinding {
22    pub tool_call_id: String,
23    pub policy_decision: Option<String>,
24    pub kernel_event_count: u64,
25    pub window: BindingWindow,
26}
27
28#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
29pub struct CorrelationReport {
30    pub schema: String,
31    pub run_id: String,
32    pub status: CorrelationStatus,
33    pub bindings: Vec<CorrelationBinding>,
34    pub ambiguities: Vec<String>,
35}
36
37#[derive(Debug, Clone, PartialEq, Eq, Error)]
38pub enum CorrelationReportError {
39    #[error("correlation report schema must be {CORRELATION_REPORT_SCHEMA}")]
40    InvalidSchema,
41    #[error("run_id must not be empty")]
42    EmptyRunId,
43}
44
45impl CorrelationReport {
46    pub fn clean(run_id: impl Into<String>) -> Self {
47        Self {
48            schema: CORRELATION_REPORT_SCHEMA.to_string(),
49            run_id: run_id.into(),
50            status: CorrelationStatus::Clean,
51            bindings: Vec::new(),
52            ambiguities: Vec::new(),
53        }
54    }
55
56    pub fn add_binding(&mut self, binding: CorrelationBinding) {
57        self.bindings.push(binding);
58    }
59
60    pub fn mark_partial(&mut self, ambiguity: impl Into<String>) {
61        if self.status == CorrelationStatus::Clean {
62            self.status = CorrelationStatus::Partial;
63        }
64        self.ambiguities.push(ambiguity.into());
65    }
66
67    pub fn mark_failed(&mut self, ambiguity: impl Into<String>) {
68        self.status = CorrelationStatus::Failed;
69        self.ambiguities.push(ambiguity.into());
70    }
71
72    pub fn validate(&self) -> Result<(), CorrelationReportError> {
73        if self.schema != CORRELATION_REPORT_SCHEMA {
74            return Err(CorrelationReportError::InvalidSchema);
75        }
76        if self.run_id.is_empty() {
77            return Err(CorrelationReportError::EmptyRunId);
78        }
79        Ok(())
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn adding_partial_ambiguity_changes_clean_report_to_partial() {
89        let mut report = CorrelationReport::clean("run_001");
90
91        report.mark_partial("sdk_layer_absent");
92
93        assert_eq!(report.status, CorrelationStatus::Partial);
94        assert_eq!(report.ambiguities, vec!["sdk_layer_absent"]);
95    }
96
97    #[test]
98    fn failed_status_is_sticky() {
99        let mut report = CorrelationReport::clean("run_001");
100
101        report.mark_failed("cgroup_correlation_failed");
102        report.mark_partial("sdk_layer_absent");
103
104        assert_eq!(report.status, CorrelationStatus::Failed);
105        assert_eq!(
106            report.ambiguities,
107            vec!["cgroup_correlation_failed", "sdk_layer_absent"]
108        );
109    }
110
111    #[test]
112    fn validate_rejects_unexpected_schema() {
113        let mut report = CorrelationReport::clean("run_001");
114        report.schema = "assay.runner.correlation_report.v_future".to_string();
115
116        assert_eq!(
117            report.validate(),
118            Err(CorrelationReportError::InvalidSchema)
119        );
120    }
121}