Skip to main content

context_governor/
high_roi.rs

1use crate::hash_text;
2use serde::{Deserialize, Serialize};
3use std::cmp::Ordering;
4use std::collections::BTreeSet;
5
6#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Default)]
7#[serde(rename_all = "snake_case")]
8pub enum RiskLevel {
9    #[default]
10    Low,
11    Medium,
12    High,
13    Critical,
14}
15
16#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
17#[serde(rename_all = "snake_case")]
18pub enum GovernanceFailureMode {
19    UnauthorizedLeakage,
20    StalePropagation,
21    ContradictionPersistence,
22    ProvenanceCollapse,
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
26pub struct GovernanceCase {
27    pub case_id: String,
28    pub mode: GovernanceFailureMode,
29    pub passed: bool,
30}
31
32impl GovernanceCase {
33    pub fn new(case_id: impl Into<String>, mode: GovernanceFailureMode, passed: bool) -> Self {
34        Self {
35            case_id: case_id.into(),
36            mode,
37            passed,
38        }
39    }
40}
41
42#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
43pub struct GovernedMemoryHarnessReceiptV1 {
44    pub schema: String,
45    pub harness_id: String,
46    pub total_cases: usize,
47    pub passed_cases: usize,
48    pub failed_cases: usize,
49    pub failure_modes: Vec<GovernanceFailureMode>,
50    pub certified: bool,
51    pub reasons: Vec<String>,
52}
53
54pub fn evaluate_governed_memory(
55    harness_id: impl Into<String>,
56    cases: &[GovernanceCase],
57) -> GovernedMemoryHarnessReceiptV1 {
58    let mut failure_modes = cases
59        .iter()
60        .filter(|case| !case.passed)
61        .map(|case| case.mode)
62        .collect::<BTreeSet<_>>()
63        .into_iter()
64        .collect::<Vec<_>>();
65    failure_modes.sort();
66    let passed_cases = cases.iter().filter(|case| case.passed).count();
67    let failed_cases = cases.len().saturating_sub(passed_cases);
68    let mut reasons = Vec::new();
69    for mode in &failure_modes {
70        reasons.push(format!("failed-{mode:?}").to_lowercase());
71    }
72    GovernedMemoryHarnessReceiptV1 {
73        schema: "GovernedMemoryHarnessReceiptV1".to_string(),
74        harness_id: harness_id.into(),
75        total_cases: cases.len(),
76        passed_cases,
77        failed_cases,
78        failure_modes,
79        certified: failed_cases == 0 && !cases.is_empty(),
80        reasons,
81    }
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
85pub struct ToolManifestEntry {
86    pub name: String,
87    pub description: String,
88}
89
90impl ToolManifestEntry {
91    pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
92        Self {
93            name: name.into(),
94            description: description.into(),
95        }
96    }
97}
98
99#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
100pub struct ToolSurfaceFinding {
101    pub reason: String,
102    pub risk: RiskLevel,
103    pub affected_tools: Vec<String>,
104}
105
106#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
107pub struct McpToolSurfaceAuditV1 {
108    pub schema: String,
109    pub tool_count: usize,
110    pub combined_surface_risk: RiskLevel,
111    pub findings: Vec<ToolSurfaceFinding>,
112}
113
114pub fn audit_mcp_tool_surface(tools: &[ToolManifestEntry]) -> McpToolSurfaceAuditV1 {
115    let mut findings = Vec::new();
116    for (left_idx, left) in tools.iter().enumerate() {
117        for right in tools.iter().skip(left_idx + 1) {
118            let combined = format!("{} {}", left.description, right.description).to_lowercase();
119            let left_l = left.description.to_lowercase();
120            let right_l = right.description.to_lowercase();
121            let split_payload = contains_any(&combined, &["ignore previous", "instructions"])
122                && contains_any(&combined, &["share a", "share b", "threshold", "combined"])
123                && left_l != right_l
124                && (contains_any(&left_l, &["ignore previous", "share a", "threshold"])
125                    || contains_any(&right_l, &["ignore previous", "share a", "threshold"]))
126                && (contains_any(&left_l, &["instructions", "share b", "combined"])
127                    || contains_any(&right_l, &["instructions", "share b", "combined"]));
128            if split_payload {
129                findings.push(ToolSurfaceFinding {
130                    reason: "split-instruction-fragments-across-tools".to_string(),
131                    risk: RiskLevel::High,
132                    affected_tools: vec![left.name.clone(), right.name.clone()],
133                });
134            }
135        }
136    }
137    let combined_surface_risk = findings
138        .iter()
139        .map(|finding| finding.risk)
140        .max()
141        .unwrap_or(RiskLevel::Low);
142    McpToolSurfaceAuditV1 {
143        schema: "McpToolSurfaceAuditV1".to_string(),
144        tool_count: tools.len(),
145        combined_surface_risk,
146        findings,
147    }
148}
149
150#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
151pub struct BoundaryFinding {
152    pub reason: String,
153    pub risk: RiskLevel,
154    pub location: String,
155    pub snippet: String,
156}
157
158#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
159pub struct CompressionBoundaryAuditV1 {
160    pub schema: String,
161    pub source_findings: Vec<BoundaryFinding>,
162    pub summary_findings: Vec<BoundaryFinding>,
163    pub relinking_risk: RiskLevel,
164    pub safe_to_reinject: bool,
165}
166
167pub fn audit_compression_boundary(
168    source_fragments: &[String],
169    compressed_summary: &str,
170) -> CompressionBoundaryAuditV1 {
171    let source_findings = source_fragments
172        .iter()
173        .enumerate()
174        .flat_map(|(idx, fragment)| {
175            scan_instruction_boundary(fragment, false, format!("source[{idx}]"))
176        })
177        .collect::<Vec<_>>();
178    let summary_findings =
179        scan_instruction_boundary(compressed_summary, true, "compressed_summary".to_string());
180    let relinking_risk = if summary_findings
181        .iter()
182        .any(|finding| finding.risk >= RiskLevel::High)
183    {
184        RiskLevel::High
185    } else if !source_findings.is_empty() {
186        RiskLevel::Medium
187    } else {
188        RiskLevel::Low
189    };
190    CompressionBoundaryAuditV1 {
191        schema: "CompressionBoundaryAuditV1".to_string(),
192        source_findings,
193        summary_findings,
194        relinking_risk,
195        safe_to_reinject: relinking_risk < RiskLevel::High,
196    }
197}
198
199fn scan_instruction_boundary(
200    text: &str,
201    post_compression: bool,
202    location: String,
203) -> Vec<BoundaryFinding> {
204    let lower = text.to_lowercase();
205    let mut findings = Vec::new();
206    if contains_any(
207        &lower,
208        &[
209            "ignore previous",
210            "disregard previous",
211            "run the release command",
212            "execute the command",
213        ],
214    ) {
215        findings.push(BoundaryFinding {
216            reason: if post_compression {
217                "post-compression-instruction".to_string()
218            } else {
219                "source-instruction-fragment".to_string()
220            },
221            risk: if post_compression {
222                RiskLevel::High
223            } else {
224                RiskLevel::Medium
225            },
226            location,
227            snippet: preview(text, 160),
228        });
229    }
230    findings
231}
232
233#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
234pub struct RagEvalInput {
235    pub task_id: String,
236    pub closed_book_correct: bool,
237    pub retrieved_answer_correct: bool,
238    pub retrieval_used: bool,
239}
240
241#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
242pub struct RagLeakageReceiptV1 {
243    pub schema: String,
244    pub task_id: String,
245    pub certified_retrieval_gain: bool,
246    pub reasons: Vec<String>,
247}
248
249pub fn evaluate_leakage_free_rag(input: RagEvalInput) -> RagLeakageReceiptV1 {
250    let mut reasons = Vec::new();
251    if input.closed_book_correct {
252        reasons.push("closed-book-solved-task".to_string());
253    }
254    if !input.retrieval_used {
255        reasons.push("retrieval-not-used".to_string());
256    }
257    if !input.retrieved_answer_correct {
258        reasons.push("retrieved-answer-incorrect".to_string());
259    }
260    let certified_retrieval_gain =
261        !input.closed_book_correct && input.retrieval_used && input.retrieved_answer_correct;
262    if certified_retrieval_gain {
263        reasons.push("retrieval-improved-unsolved-task".to_string());
264    }
265    RagLeakageReceiptV1 {
266        schema: "RagLeakageReceiptV1".to_string(),
267        task_id: input.task_id,
268        certified_retrieval_gain,
269        reasons,
270    }
271}
272
273#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
274pub struct EvidenceClaim {
275    pub id: String,
276    pub text: String,
277}
278
279impl EvidenceClaim {
280    pub fn new(id: impl Into<String>, text: impl Into<String>) -> Self {
281        Self {
282            id: id.into(),
283            text: text.into(),
284        }
285    }
286}
287
288#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
289pub struct ConflictFinding {
290    pub claim_a_id: String,
291    pub claim_b_id: String,
292    pub reason: String,
293}
294
295#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
296pub struct ConflictScreenReportV1 {
297    pub schema: String,
298    pub needs_expensive_review: bool,
299    pub conflicts: Vec<ConflictFinding>,
300}
301
302pub fn screen_knowledge_conflicts(claims: &[EvidenceClaim]) -> ConflictScreenReportV1 {
303    let mut conflicts = Vec::new();
304    for (left_idx, left) in claims.iter().enumerate() {
305        for right in claims.iter().skip(left_idx + 1) {
306            let left_nums = numbers(&left.text);
307            let right_nums = numbers(&right.text);
308            if !left_nums.is_empty() && !right_nums.is_empty() && left_nums != right_nums {
309                conflicts.push(ConflictFinding {
310                    claim_a_id: left.id.clone(),
311                    claim_b_id: right.id.clone(),
312                    reason: "numeric-disagreement".to_string(),
313                });
314            }
315            if negation_conflict(&left.text, &right.text) {
316                conflicts.push(ConflictFinding {
317                    claim_a_id: left.id.clone(),
318                    claim_b_id: right.id.clone(),
319                    reason: "negation-disagreement".to_string(),
320                });
321            }
322        }
323    }
324    ConflictScreenReportV1 {
325        schema: "ConflictScreenReportV1".to_string(),
326        needs_expensive_review: !conflicts.is_empty(),
327        conflicts,
328    }
329}
330
331#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
332#[serde(rename_all = "snake_case")]
333pub enum RetrievalRoute {
334    FlatSearchOnly,
335    GraphAssisted,
336    ConflictAware,
337    Synthesis,
338    Temporal,
339}
340
341#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
342pub struct RetrievalRouteDecisionV1 {
343    pub schema: String,
344    pub route: RetrievalRoute,
345    pub reasons: Vec<String>,
346}
347
348pub fn select_retrieval_route(query: &str) -> RetrievalRouteDecisionV1 {
349    let lower = query.to_lowercase();
350    let (route, reason) = if contains_any(
351        &lower,
352        &[
353            "contradict",
354            "conflict",
355            "vs",
356            "is it true",
357            "but ",
358            "compared",
359        ],
360    ) {
361        (RetrievalRoute::ConflictAware, "conflict-sensitive")
362    } else if contains_any(
363        &lower,
364        &[
365            "summarize",
366            "overview",
367            "all themes",
368            "landscape",
369            "everything",
370        ],
371    ) {
372        (RetrievalRoute::Synthesis, "synthesis")
373    } else if contains_any(
374        &lower,
375        &["changed", "after", "before", "when", "current vs", "latest"],
376    ) {
377        (RetrievalRoute::Temporal, "temporal")
378    } else if contains_any(
379        &lower,
380        &["connect", "relate", "depends", "between", "lead to"],
381    ) {
382        (RetrievalRoute::GraphAssisted, "multi-hop")
383    } else {
384        (RetrievalRoute::FlatSearchOnly, "simple-lookup")
385    };
386    RetrievalRouteDecisionV1 {
387        schema: "RetrievalRouteDecisionV1".to_string(),
388        route,
389        reasons: vec![reason.to_string()],
390    }
391}
392
393#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
394#[serde(rename_all = "snake_case")]
395pub enum MemoryModule {
396    Representation,
397    Organization,
398    RetrievalUpdate,
399    LifecycleGovernance,
400}
401
402#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
403pub struct MemoryModuleMetric {
404    pub module: MemoryModule,
405    pub score: f64,
406    pub observations: usize,
407}
408
409impl MemoryModuleMetric {
410    pub fn new(module: MemoryModule, score: f64, observations: usize) -> Self {
411        Self {
412            module,
413            score,
414            observations,
415        }
416    }
417}
418
419#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
420pub struct AgentMemoryModuleReportV1 {
421    pub schema: String,
422    pub ready_for_public_claims: bool,
423    pub missing_modules: Vec<MemoryModule>,
424    pub min_score: f64,
425    pub total_observations: usize,
426}
427
428pub fn evaluate_agent_memory_modules(metrics: &[MemoryModuleMetric]) -> AgentMemoryModuleReportV1 {
429    let present = metrics
430        .iter()
431        .map(|metric| metric.module)
432        .collect::<BTreeSet<_>>();
433    let required = [
434        MemoryModule::Representation,
435        MemoryModule::Organization,
436        MemoryModule::RetrievalUpdate,
437        MemoryModule::LifecycleGovernance,
438    ];
439    let missing_modules = required
440        .into_iter()
441        .filter(|module| !present.contains(module))
442        .collect::<Vec<_>>();
443    let min_score = metrics
444        .iter()
445        .map(|metric| metric.score)
446        .fold(1.0_f64, f64::min);
447    let total_observations = metrics.iter().map(|metric| metric.observations).sum();
448    AgentMemoryModuleReportV1 {
449        schema: "AgentMemoryModuleReportV1".to_string(),
450        ready_for_public_claims: missing_modules.is_empty()
451            && min_score >= 0.5
452            && total_observations >= required.len(),
453        missing_modules,
454        min_score,
455        total_observations,
456    }
457}
458
459#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
460#[serde(rename_all = "snake_case")]
461pub enum InferenceSurface {
462    HostedApi,
463    LocalInference,
464}
465
466#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
467pub struct TokenImportance {
468    pub token_id: String,
469    pub importance: f64,
470}
471
472impl TokenImportance {
473    pub fn new(token_id: impl Into<String>, importance: f64) -> Self {
474        Self {
475            token_id: token_id.into(),
476            importance,
477        }
478    }
479}
480
481#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
482pub struct SemanticKvRetentionPlanV1 {
483    pub schema: String,
484    pub surface: InferenceSurface,
485    pub retained_token_ids: Vec<String>,
486    pub boundary_notes: Vec<String>,
487}
488
489pub fn plan_semantic_kv_retention(
490    tokens: &[TokenImportance],
491    retain_count: usize,
492    surface: InferenceSurface,
493) -> SemanticKvRetentionPlanV1 {
494    let mut ranked = tokens.to_vec();
495    ranked.sort_by(|a, b| {
496        b.importance
497            .partial_cmp(&a.importance)
498            .unwrap_or(Ordering::Equal)
499            .then_with(|| a.token_id.cmp(&b.token_id))
500    });
501    let retained_token_ids = ranked
502        .into_iter()
503        .take(retain_count)
504        .map(|token| token.token_id)
505        .collect::<Vec<_>>();
506    let mut boundary_notes = Vec::new();
507    if surface == InferenceSurface::HostedApi {
508        boundary_notes.push(
509            "hosted APIs do not expose KV cache; use prompt/retrieval compaction instead"
510                .to_string(),
511        );
512    }
513    SemanticKvRetentionPlanV1 {
514        schema: "SemanticKvRetentionPlanV1".to_string(),
515        surface,
516        retained_token_ids,
517        boundary_notes,
518    }
519}
520
521#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
522#[serde(rename_all = "snake_case")]
523pub enum ProjectionKind {
524    TemporalTimeline,
525    EntityView,
526    CommunitySummary,
527}
528
529#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
530#[serde(rename_all = "snake_case")]
531pub enum ProjectionFreshness {
532    Fresh,
533    Stale,
534    Unknown,
535}
536
537#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
538pub struct ProjectionReceiptV1 {
539    pub schema: String,
540    pub projection_id: String,
541    pub kind: ProjectionKind,
542    pub source_ids: Vec<String>,
543    pub derivation_blake3: String,
544    pub freshness: ProjectionFreshness,
545}
546
547pub fn build_projection_receipt(
548    kind: ProjectionKind,
549    source_ids: &[String],
550    projection_content: &str,
551    freshness: ProjectionFreshness,
552) -> ProjectionReceiptV1 {
553    let mut hash_material = source_ids.join("\n");
554    hash_material.push('\n');
555    hash_material.push_str(projection_content);
556    let derivation_blake3 = hash_text(&hash_material);
557    ProjectionReceiptV1 {
558        schema: "ProjectionReceiptV1".to_string(),
559        projection_id: format!("projection:{}", &derivation_blake3[..16]),
560        kind,
561        source_ids: source_ids.to_vec(),
562        derivation_blake3,
563        freshness,
564    }
565}
566
567fn contains_any(text: &str, needles: &[&str]) -> bool {
568    needles.iter().any(|needle| text.contains(needle))
569}
570
571fn preview(text: &str, max_chars: usize) -> String {
572    text.split_whitespace()
573        .collect::<Vec<_>>()
574        .join(" ")
575        .chars()
576        .take(max_chars)
577        .collect()
578}
579
580fn numbers(text: &str) -> Vec<i64> {
581    let mut out = Vec::new();
582    let mut cur = String::new();
583    for ch in text.chars() {
584        if ch.is_ascii_digit() {
585            cur.push(ch);
586        } else if !cur.is_empty() {
587            if let Ok(n) = cur.parse::<i64>() {
588                out.push(n);
589            }
590            cur.clear();
591        }
592    }
593    if !cur.is_empty() {
594        if let Ok(n) = cur.parse::<i64>() {
595            out.push(n);
596        }
597    }
598    out
599}
600
601fn negation_conflict(left: &str, right: &str) -> bool {
602    let normalize = |text: &str| {
603        text.to_lowercase()
604            .replace(" not ", " ")
605            .replace(" no ", " ")
606            .replace(" unsupported", " supported")
607            .split_whitespace()
608            .collect::<Vec<_>>()
609            .join(" ")
610    };
611    let left_l = left.to_lowercase();
612    let right_l = right.to_lowercase();
613    let left_neg = contains_any(&left_l, &[" not ", " no ", "unsupported"]);
614    let right_neg = contains_any(&right_l, &[" not ", " no ", "unsupported"]);
615    left_neg != right_neg && normalize(left) == normalize(right)
616}