Skip to main content

lean_ctx/core/context_kernel/
conformance.rs

1#[cfg(test)]
2mod tests {
3    use std::collections::HashMap;
4
5    use crate::core::context_field::{ContextItemId, TokenBudget};
6    use crate::core::context_kernel::attribution::{AttributionReport, compute_attribution};
7    use crate::core::context_kernel::learning::{OutcomeLearner, WeightUpdate};
8    use crate::core::context_kernel::orchestrator::ContextKernel;
9    use crate::core::context_kernel::policy::{ContextPolicy, PolicyFilter};
10    use crate::core::context_kernel::types::{
11        ContextObjectKind, ContextObjectV1, ContextPlanV1, ContextReceiptV1, PlanBudget, PlanEntry,
12        QualitySignal, ReceiptOutcome, RetrievalContext, SensitivityLevel,
13    };
14
15    fn test_candidate(
16        source: &str,
17        sensitivity: SensitivityLevel,
18        tokens: usize,
19    ) -> ContextObjectV1 {
20        ContextObjectV1 {
21            id: ContextItemId(format!("test:{source}")),
22            kind: ContextObjectKind::Fact,
23            source: source.to_owned(),
24            sensitivity,
25            token_estimate: tokens,
26            ..ContextObjectV1::default()
27        }
28    }
29
30    #[test]
31    fn plan_receipt_roundtrip() {
32        let project_root = std::env::temp_dir().join("lean-ctx-kernel-conformance");
33        let project_root_text = project_root.to_string_lossy();
34        let kernel = ContextKernel::for_project(project_root_text.as_ref());
35        let context = RetrievalContext {
36            query: "context kernel conformance".to_owned(),
37            task: Some("verify plan receipt roundtrip".to_owned()),
38            project_root: project_root_text.into_owned(),
39            budget: TokenBudget {
40                total: 1_000,
41                used: 0,
42            },
43            max_candidates: 10,
44        };
45
46        let plan = kernel.plan(&context);
47        let receipt = kernel.record_receipt(&plan, 64, ReceiptOutcome::Accepted);
48
49        assert_eq!(receipt.plan_id, plan.plan_id);
50        assert!(receipt.delivered_tokens > 0);
51    }
52
53    #[test]
54    fn policy_filters_sensitive_candidates() {
55        let candidates: Vec<ContextObjectV1> = vec![
56            test_candidate("public", SensitivityLevel::Public, 20),
57            test_candidate("restricted", SensitivityLevel::Restricted, 20),
58        ];
59        let policy = ContextPolicy {
60            max_sensitivity: SensitivityLevel::Internal,
61            allowed_sources: None,
62            blocked_sources: Vec::new(),
63            budget_cap_tokens: None,
64            retention_days: None,
65        };
66        let filter = PolicyFilter::new(policy);
67
68        let filtered = filter.apply(candidates);
69
70        assert_eq!(filtered.len(), 1);
71        assert_eq!(filtered[0].source, "public");
72        assert_eq!(filtered[0].sensitivity, SensitivityLevel::Public);
73    }
74
75    #[test]
76    fn attribution_no_double_counting() {
77        let plan = ContextPlanV1 {
78            plan_id: "plan:conformance".to_owned(),
79            intent: "verify provider attribution".to_owned(),
80            budget: PlanBudget {
81                total_tokens: 1_000,
82                used_tokens: 700,
83                remaining_tokens: 300,
84            },
85            selected: vec![
86                plan_entry("file:a", "files", 400),
87                plan_entry("file:b", "files", 100),
88                plan_entry("fact:c", "knowledge", 200),
89            ],
90            excluded: Vec::new(),
91            deferred: Vec::new(),
92            provider_stats: HashMap::new(),
93        };
94        let receipt = ContextReceiptV1 {
95            receipt_id: "receipt:conformance".to_owned(),
96            plan_id: plan.plan_id.clone(),
97            delivered_tokens: 500,
98            cache_hits: 1,
99            cache_misses: 2,
100            outcome: ReceiptOutcome::Accepted,
101            quality_signals: vec![QualitySignal {
102                signal_type: "outcome".to_owned(),
103                value: 1.0,
104            }],
105            feedback_attribution: HashMap::new(),
106        };
107
108        let report: AttributionReport = compute_attribution(&plan, &receipt);
109        let summed_savings: usize = report.entries.iter().map(|entry| entry.tokens_saved).sum();
110        let mut provider_occurrences: HashMap<&str, usize> = HashMap::new();
111        for entry in &report.entries {
112            *provider_occurrences
113                .entry(entry.provider.as_str())
114                .or_insert(0) += 1;
115        }
116
117        assert!(summed_savings <= report.total_tokens_saved);
118        assert_eq!(provider_occurrences.get("files"), Some(&1));
119        assert_eq!(provider_occurrences.get("knowledge"), Some(&1));
120        assert!(provider_occurrences.values().all(|count| *count == 1));
121    }
122
123    #[test]
124    fn learning_updates_provider_weights() {
125        let initial_weights: HashMap<String, f64> =
126            HashMap::from([("files".to_owned(), 0.5), ("knowledge".to_owned(), 0.5)]);
127        let receipt = ContextReceiptV1 {
128            receipt_id: "receipt:learning".to_owned(),
129            plan_id: "plan:learning".to_owned(),
130            delivered_tokens: 300,
131            cache_hits: 0,
132            cache_misses: 0,
133            outcome: ReceiptOutcome::Accepted,
134            quality_signals: Vec::new(),
135            feedback_attribution: HashMap::from([
136                ("files".to_owned(), 0.6),
137                ("knowledge".to_owned(), 0.4),
138            ]),
139        };
140        let learner = OutcomeLearner::default_learner();
141
142        let updates: Vec<WeightUpdate> = learner.learn_from_receipt(&receipt, &initial_weights);
143
144        assert_eq!(updates.len(), initial_weights.len());
145        for update in updates {
146            let old_weight = initial_weights
147                .get(&update.provider)
148                .copied()
149                .expect("attributed provider has an initial weight");
150            assert_eq!(update.old_weight, old_weight);
151            assert_ne!(update.new_weight, old_weight);
152            assert!(update.new_weight >= old_weight);
153        }
154    }
155
156    fn plan_entry(object_id: &str, provider: &str, tokens: usize) -> PlanEntry {
157        PlanEntry {
158            object_id: object_id.to_owned(),
159            provider: provider.to_owned(),
160            view: "summary".to_owned(),
161            tokens,
162            phi: 0.8,
163            reason: "selected for conformance".to_owned(),
164        }
165    }
166}