Skip to main content

lean_ctx/core/context_kernel/
bench.rs

1#[cfg(test)]
2mod golden_workloads {
3    use std::collections::HashMap;
4
5    use super::super::orchestrator::ContextKernel;
6    use super::super::types::{
7        CandidateProvider, ContextObjectKind, ContextObjectV1, Freshness, RetrievalContext,
8        SensitivityLevel, SideEffectPolicy,
9    };
10    use crate::core::context_field::{ContextItemId, Provenance, TokenBudget, ViewCosts};
11
12    fn test_object(
13        id: &str,
14        kind: ContextObjectKind,
15        title: &str,
16        confidence: f32,
17        tokens: usize,
18    ) -> ContextObjectV1 {
19        ContextObjectV1 {
20            id: ContextItemId::from_provider("golden", id),
21            kind,
22            source: "test".to_owned(),
23            content_ref: format!("ref:{id}"),
24            title: title.to_owned(),
25            content: Some(format!("Content for {title}")),
26            freshness: Freshness::default(),
27            confidence,
28            sensitivity: SensitivityLevel::Internal,
29            token_estimate: tokens,
30            view_costs: ViewCosts::default(),
31            provenance: Provenance::default(),
32            semantic_fingerprint: None,
33            metadata: HashMap::new(),
34        }
35    }
36
37    struct GoldenProvider {
38        name: String,
39        candidates: Vec<ContextObjectV1>,
40    }
41
42    impl CandidateProvider for GoldenProvider {
43        fn provider_id(&self) -> &str {
44            &self.name
45        }
46
47        fn candidates(&self, _ctx: &RetrievalContext) -> Vec<ContextObjectV1> {
48            self.candidates.clone()
49        }
50
51        fn side_effect_policy(&self) -> SideEffectPolicy {
52            SideEffectPolicy::ReadOnly
53        }
54    }
55
56    fn retrieval_ctx(query: &str, budget: usize) -> RetrievalContext {
57        RetrievalContext {
58            query: query.to_owned(),
59            task: Some(query.to_owned()),
60            project_root: "/tmp/golden-test".to_owned(),
61            budget: TokenBudget {
62                total: budget,
63                used: 0,
64            },
65            max_candidates: 20,
66        }
67    }
68
69    fn object_id(id: &str) -> String {
70        ContextItemId::from_provider("golden", id).to_string()
71    }
72
73    #[test]
74    fn golden_auth_bug_hunt() {
75        let knowledge = GoldenProvider {
76            name: "golden.knowledge".to_owned(),
77            candidates: vec![
78                test_object(
79                    "auth-middleware",
80                    ContextObjectKind::Fact,
81                    "auth middleware uses JWT",
82                    0.95,
83                    100,
84                ),
85                test_object(
86                    "rate-limiter",
87                    ContextObjectKind::Fact,
88                    "rate limiter config",
89                    0.10,
90                    100,
91                ),
92            ],
93        };
94        let memory = GoldenProvider {
95            name: "golden.memory".to_owned(),
96            candidates: vec![
97                test_object(
98                    "auth-bypass",
99                    ContextObjectKind::Episode,
100                    "fixed auth bypass in v2.3",
101                    0.75,
102                    100,
103                ),
104                test_object(
105                    "database-migration",
106                    ContextObjectKind::Episode,
107                    "database migration v1.0",
108                    0.10,
109                    100,
110                ),
111                test_object(
112                    "auth-debug-sequence",
113                    ContextObjectKind::Procedure,
114                    "auth debug sequence",
115                    0.95,
116                    100,
117                ),
118            ],
119        };
120        let kernel = ContextKernel::new(vec![Box::new(knowledge), Box::new(memory)]);
121        let plan = kernel.plan(&retrieval_ctx("auth middleware bypass debug sequence", 300));
122
123        assert!(
124            plan.selected
125                .iter()
126                .any(|entry| entry.object_id == object_id("auth-middleware"))
127        );
128        assert!(
129            plan.selected
130                .iter()
131                .any(|entry| entry.object_id == object_id("auth-bypass"))
132        );
133        assert!(
134            plan.excluded
135                .iter()
136                .any(|entry| entry.object_id == object_id("rate-limiter"))
137        );
138        assert!(
139            plan.excluded
140                .iter()
141                .any(|entry| entry.object_id == object_id("database-migration"))
142        );
143        assert!(plan.budget.used_tokens <= plan.budget.total_tokens);
144    }
145
146    #[test]
147    fn golden_caching_feature() {
148        let provider = GoldenProvider {
149            name: "golden.caching".to_owned(),
150            candidates: vec![
151                test_object(
152                    "cache-invalidation",
153                    ContextObjectKind::Procedure,
154                    "cache invalidation pattern",
155                    0.95,
156                    100,
157                ),
158                test_object(
159                    "redis-pool",
160                    ContextObjectKind::Fact,
161                    "Redis connection pool config",
162                    0.75,
163                    100,
164                ),
165                test_object(
166                    "current-cache-task",
167                    ContextObjectKind::SessionItem,
168                    "current task: add caching",
169                    0.95,
170                    100,
171                ),
172                test_object(
173                    "unrelated-refactoring",
174                    ContextObjectKind::Episode,
175                    "unrelated refactoring",
176                    0.10,
177                    100,
178                ),
179            ],
180        };
181        let kernel = ContextKernel::new(vec![Box::new(provider)]);
182        let plan = kernel.plan(&retrieval_ctx(
183            "cache invalidation Redis connection pool current task add caching",
184            300,
185        ));
186
187        for id in ["cache-invalidation", "redis-pool", "current-cache-task"] {
188            assert!(
189                plan.selected
190                    .iter()
191                    .any(|entry| entry.object_id == object_id(id)),
192                "expected {id} to be selected"
193            );
194        }
195        assert!(
196            plan.excluded
197                .iter()
198                .any(|entry| entry.object_id == object_id("unrelated-refactoring"))
199        );
200        assert!(plan.budget.used_tokens <= plan.budget.total_tokens);
201    }
202
203    #[test]
204    fn golden_budget_constraint() {
205        let provider = GoldenProvider {
206            name: "golden.budget".to_owned(),
207            candidates: vec![
208                test_object("item-a", ContextObjectKind::Fact, "item A", 0.9, 100),
209                test_object("item-b", ContextObjectKind::Fact, "item B", 0.8, 100),
210                test_object("item-c", ContextObjectKind::Fact, "item C", 0.7, 100),
211            ],
212        };
213        let kernel = ContextKernel::new(vec![Box::new(provider)]);
214        let plan = kernel.plan(&retrieval_ctx("item", 200));
215
216        assert_eq!(plan.selected.len(), 2);
217        assert!(
218            plan.selected
219                .iter()
220                .any(|entry| entry.object_id == object_id("item-a"))
221        );
222        assert!(
223            plan.selected
224                .iter()
225                .any(|entry| entry.object_id == object_id("item-b"))
226        );
227        assert!(
228            plan.deferred
229                .iter()
230                .any(|entry| entry.object_id == object_id("item-c"))
231                || plan
232                    .excluded
233                    .iter()
234                    .any(|entry| entry.object_id == object_id("item-c"))
235        );
236        assert!(plan.budget.used_tokens <= 200);
237    }
238}