context_governor/
lib.rs

1//! Crate-agnostic governed context compaction for AI agents.
2//!
3//! `context-governor` does not call an LLM and does not pretend to extend a
4//! provider's hard context window. It allocates a limited prompt budget across
5//! agent conversation messages, preserves high-risk context verbatim, demotes or
6//! summarizes low-risk context, emits receipts, and keeps exact fallback ranges
7//! for later expansion.
8
9use chrono::{DateTime, Utc};
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12use std::collections::{BTreeMap, BTreeSet};
13use thiserror::Error;
14use uuid::Uuid;
15
16pub mod high_roi;
17pub use high_roi::*;
18
19const CHARS_PER_TOKEN: usize = 4;
20const SUMMARY_PREFIX: &str = "[CONTEXT COMPACTION — RECEIPT-BACKED REFERENCE ONLY]";
21
22#[derive(Debug, Error)]
23pub enum ContextGovernorError {
24    #[error("cannot compact an empty message list")]
25    EmptyMessages,
26    #[error("serialization failed: {0}")]
27    Serialization(#[from] serde_json::Error),
28    #[error("io failed: {0}")]
29    Io(#[from] std::io::Error),
30    #[error("receipt not found: {0}")]
31    ReceiptNotFound(String),
32    #[error("context budget exceeded: target {target_tokens} tokens, minimum required {minimum_required_tokens} tokens")]
33    BudgetExceeded {
34        target_tokens: usize,
35        minimum_required_tokens: usize,
36    },
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
40pub struct Message {
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub id: Option<String>,
43    pub role: String,
44    pub content: String,
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub name: Option<String>,
47    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
48    pub metadata: BTreeMap<String, Value>,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
52#[serde(rename_all = "snake_case")]
53pub enum ItemType {
54    LatestUserMessage,
55    ActiveInstruction,
56    AcceptanceGate,
57    ToolCall,
58    ToolResult,
59    ErrorOutput,
60    FilePathContext,
61    Decision,
62    UnresolvedQuestion,
63    SourceEvidence,
64    DurableFactCandidate,
65    ProjectStateCandidate,
66    ArtifactBoilerplate,
67    StalePlan,
68    DuplicateContext,
69    LowRiskNarrative,
70    #[default]
71    Unknown,
72}
73
74#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
75#[serde(rename_all = "snake_case")]
76pub enum AuthorityClass {
77    MustPreserveExact,
78    EvidenceCritical,
79    ActiveTask,
80    VerifiedToolReceipt,
81    DurableMemoryCandidate,
82    #[default]
83    SummaryOk,
84    ArchiveOk,
85    Discardable,
86    Quarantine,
87}
88
89#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
90#[serde(rename_all = "snake_case")]
91pub enum PreservationPolicy {
92    KeepVerbatim,
93    #[default]
94    ExtractiveSummary,
95    AbstractiveSummary,
96    SemanticMemoryArchive,
97    ReceiptOnly,
98    OmitDuplicate,
99    Quarantine,
100}
101
102#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
103#[serde(rename_all = "snake_case")]
104pub enum ContentKind {
105    PlainText,
106    Json,
107    Diff,
108    Rust,
109    Markdown,
110    CargoOutput,
111    ShellLog,
112    SearchResults,
113    #[default]
114    Unknown,
115}
116
117#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
118#[serde(rename_all = "snake_case")]
119pub enum BudgetMode {
120    #[default]
121    SoftWarn,
122    HardCascade,
123    FailClosed,
124}
125
126#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
127#[serde(rename_all = "snake_case")]
128pub enum TokenCounterKind {
129    #[default]
130    ApproxChars,
131    ApproxWords,
132    ProviderChatApprox,
133    /// OpenAI cl100k-compatible tokenizer surface. In the default build this
134    /// deliberately falls back to the provider chat approximation and emits a
135    /// warning instead of silently pretending native tokenization happened.
136    TiktokenCl100k,
137}
138
139#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
140#[serde(rename_all = "snake_case")]
141pub enum AllocatorMode {
142    #[default]
143    DeterministicV1,
144    AggressiveV1,
145}
146
147impl AllocatorMode {
148    fn as_str(&self) -> &'static str {
149        match self {
150            Self::DeterministicV1 => "deterministic_v1",
151            Self::AggressiveV1 => "aggressive_v1",
152        }
153    }
154}
155#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
156pub struct ContextItemV1 {
157    pub schema: String,
158    pub item_id: String,
159    pub session_id: String,
160    pub start_index: usize,
161    pub end_index: usize,
162    pub role_set: Vec<String>,
163    pub char_count: usize,
164    pub approx_tokens: usize,
165    pub content_blake3: String,
166    #[serde(default)]
167    pub content_kind: ContentKind,
168    pub item_type: ItemType,
169    pub authority_class: AuthorityClass,
170    pub preservation_policy: PreservationPolicy,
171    pub risk_reasons: Vec<String>,
172    pub source_message_ids: Vec<String>,
173    pub priority_score: i32,
174}
175
176#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
177pub struct ContextAllocationPlanV1 {
178    pub schema: String,
179    pub plan_id: String,
180    pub session_id: String,
181    pub created_utc: DateTime<Utc>,
182    pub context_budget_tokens: usize,
183    pub target_output_tokens: usize,
184    pub allocator: String,
185    pub items: Vec<ContextItemV1>,
186    pub kept_item_ids: Vec<String>,
187    pub summarized_item_ids: Vec<String>,
188    pub archived_item_ids: Vec<String>,
189    pub omitted_item_ids: Vec<String>,
190    pub quarantined_item_ids: Vec<String>,
191}
192
193#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
194pub struct ExactFallbackRefV1 {
195    pub item_id: String,
196    pub start_index: usize,
197    pub end_index: usize,
198    pub content_blake3: String,
199    pub approx_tokens: usize,
200}
201
202#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
203pub struct ExactStoredItemV1 {
204    pub item_id: String,
205    pub source_indices: Vec<usize>,
206    pub content: String,
207    pub content_blake3: String,
208}
209
210#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
211pub struct StructuredContextSummaryV1 {
212    pub active_task: Option<String>,
213    pub acceptance_gates: Vec<String>,
214    pub files: Vec<String>,
215    pub commands: Vec<String>,
216    pub errors: Vec<String>,
217    pub decisions: Vec<String>,
218    pub unresolved_questions: Vec<String>,
219    pub fallback_item_ids: Vec<String>,
220}
221
222#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
223pub struct SummaryLossReportV1 {
224    pub schema: String,
225    pub preserved_claims: Vec<String>,
226    pub omitted_claims: Vec<String>,
227    pub evidence_lost: Vec<String>,
228    pub uncertainty_introduced: Vec<String>,
229    pub exact_recovery_available: bool,
230    pub high_risk_omissions: Vec<String>,
231    #[serde(default)]
232    pub structured_summary: StructuredContextSummaryV1,
233}
234
235#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
236pub struct ContextCompactionReceiptV1 {
237    pub schema: String,
238    pub receipt_id: String,
239    pub session_id: String,
240    pub parent_session_id: Option<String>,
241    pub created_utc: DateTime<Utc>,
242    pub engine: String,
243    pub engine_version: String,
244    pub original_message_count: usize,
245    pub compacted_message_count: usize,
246    pub original_approx_tokens: usize,
247    pub compacted_approx_tokens: usize,
248    pub token_savings_estimate: isize,
249    pub token_counter: TokenCounterKind,
250    pub original_transcript_blake3: String,
251    pub compacted_transcript_blake3: String,
252    pub allocation_plan_id: String,
253    pub semantic_memory_fact_ids: Vec<String>,
254    pub semantic_memory_document_ids: Vec<String>,
255    pub exact_fallback_refs: Vec<ExactFallbackRefV1>,
256    pub summary_loss_report: SummaryLossReportV1,
257    pub warnings: Vec<String>,
258}
259
260#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
261pub struct CompactRequest {
262    pub session_id: String,
263    pub messages: Vec<Message>,
264    #[serde(default)]
265    pub policy: CompactionPolicy,
266    #[serde(default)]
267    pub focus: Option<String>,
268}
269
270#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
271pub struct CompactionPolicy {
272    pub target_tokens: usize,
273    pub protect_first_n: usize,
274    pub protect_last_n: usize,
275    pub summary_max_chars: usize,
276    pub allocator: String,
277    #[serde(default)]
278    pub semantic_memory_enabled: bool,
279    #[serde(default)]
280    pub archive_memory_enabled: bool,
281    #[serde(default)]
282    pub budget_mode: BudgetMode,
283    #[serde(default)]
284    pub token_counter: TokenCounterKind,
285}
286
287impl Default for CompactionPolicy {
288    fn default() -> Self {
289        Self {
290            target_tokens: 8_000,
291            protect_first_n: 3,
292            protect_last_n: 8,
293            summary_max_chars: 8_000,
294            allocator: "deterministic_v1".to_string(),
295            semantic_memory_enabled: false,
296            archive_memory_enabled: false,
297            budget_mode: BudgetMode::SoftWarn,
298            token_counter: TokenCounterKind::ApproxChars,
299        }
300    }
301}
302
303#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
304pub struct CompactResponse {
305    pub receipt: ContextCompactionReceiptV1,
306    pub allocation_plan: ContextAllocationPlanV1,
307    pub compacted_messages: Vec<Message>,
308    pub exact_store: Vec<ExactStoredItemV1>,
309}
310
311#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
312pub struct RecallCandidate {
313    pub source: Option<String>,
314    pub content: String,
315    pub score: Option<f64>,
316}
317
318#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
319#[serde(rename_all = "snake_case")]
320pub enum RecallDecision {
321    AdmitAuthoritative,
322    AdmitBackground,
323    DemoteArtifact,
324    QuarantineSpeculative,
325    RejectNoise,
326}
327
328#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
329pub struct RecallFilterResult {
330    pub decision: RecallDecision,
331    pub reasons: Vec<String>,
332}
333
334pub fn compact_context(request: CompactRequest) -> Result<CompactResponse, ContextGovernorError> {
335    compact_context_with_memory_sink(request, None)
336}
337
338pub fn compact_context_with_memory_sink(
339    request: CompactRequest,
340    memory_sink: Option<&mut dyn MemorySink>,
341) -> Result<CompactResponse, ContextGovernorError> {
342    if request.messages.is_empty() {
343        return Err(ContextGovernorError::EmptyMessages);
344    }
345
346    let original_tokens = count_tokens_messages(&request.messages, &request.policy);
347    let original_hash = hash_messages(&request.messages)?;
348    let mut items = classify_messages(&request.session_id, &request.messages, &request.policy);
349    score_items(&mut items, request.focus.as_deref());
350    let allocator_mode = resolve_allocator(&request.policy);
351    let allocator_unknown = {
352        let raw = request.policy.allocator.trim().to_lowercase();
353        !raw.is_empty()
354            && !matches!(
355                raw.as_str(),
356                "deterministic" | "deterministic_v1" | "aggressive" | "aggressive_v1"
357            )
358    };
359    let mut warning_from_allocator = Vec::new();
360    if allocator_unknown {
361        warning_from_allocator.push(format!(
362            "allocator '{}' is unknown; normalized to '{}'",
363            request.policy.allocator,
364            allocator_mode.as_str(),
365        ));
366    }
367    let plan = allocate_items(&request.session_id, items, &request.policy, allocator_mode);
368    let mut warnings = warning_from_allocator;
369    if matches!(
370        request.policy.token_counter,
371        TokenCounterKind::ApproxChars
372            | TokenCounterKind::ApproxWords
373            | TokenCounterKind::ProviderChatApprox
374    ) {
375        warnings.push(
376            "token_counter is approximate; provider-native token budgets may differ".to_string(),
377        );
378    }
379    if matches!(
380        request.policy.token_counter,
381        TokenCounterKind::ProviderChatApprox | TokenCounterKind::TiktokenCl100k
382    ) {
383        warnings.push(
384            "provider_chat_approx includes chat-role overhead but is not a native tokenizer"
385                .to_string(),
386        );
387    }
388    if matches!(
389        request.policy.token_counter,
390        TokenCounterKind::TiktokenCl100k
391    ) {
392        warnings.push(
393            "tiktoken_cl100k requested but native tokenizer feature is not compiled; using provider_chat_approx fallback".to_string(),
394        );
395    }
396
397    let exact_store = build_exact_store(&request.messages, &plan);
398    let receipt_id = format!("ctxr_{}", Uuid::new_v4().simple());
399    let structured_summary = build_structured_summary(&request.messages, &plan);
400    let summary = build_summary(
401        &request.messages,
402        &plan,
403        &request.policy,
404        &structured_summary,
405        &receipt_id,
406    );
407    let mut compacted_messages =
408        assemble_compacted_messages(&request.messages, &plan, &summary, &request.policy);
409    if matches!(request.policy.budget_mode, BudgetMode::HardCascade) {
410        warnings.push("hard cascade budget mode active".to_string());
411    }
412
413    if matches!(
414        request.policy.budget_mode,
415        BudgetMode::HardCascade | BudgetMode::FailClosed
416    ) {
417        compacted_messages = enforce_budget(compacted_messages, &request.policy, &mut warnings)?;
418    }
419
420    let compacted_tokens = count_tokens_messages(&compacted_messages, &request.policy);
421    let compacted_hash = hash_messages(&compacted_messages)?;
422    let exact_fallback_refs = exact_store
423        .iter()
424        .filter_map(|stored| {
425            plan.items
426                .iter()
427                .find(|i| i.item_id == stored.item_id)
428                .map(|item| ExactFallbackRefV1 {
429                    item_id: item.item_id.clone(),
430                    start_index: item.start_index,
431                    end_index: item.end_index,
432                    content_blake3: item.content_blake3.clone(),
433                    approx_tokens: item.approx_tokens,
434                })
435        })
436        .collect::<Vec<_>>();
437
438    if compacted_tokens > request.policy.target_tokens {
439        warnings.push(format!(
440            "compacted output still exceeds target budget: {} > {} tokens",
441            compacted_tokens, request.policy.target_tokens
442        ));
443    }
444
445    let mut receipt = ContextCompactionReceiptV1 {
446        schema: "ContextCompactionReceiptV1".to_string(),
447        receipt_id: receipt_id.clone(),
448        session_id: request.session_id.clone(),
449        parent_session_id: None,
450        created_utc: Utc::now(),
451        engine: "context-governor".to_string(),
452        engine_version: env!("CARGO_PKG_VERSION").to_string(),
453        original_message_count: request.messages.len(),
454        compacted_message_count: compacted_messages.len(),
455        original_approx_tokens: original_tokens,
456        compacted_approx_tokens: compacted_tokens,
457        token_savings_estimate: original_tokens as isize - compacted_tokens as isize,
458        token_counter: request.policy.token_counter.clone(),
459        original_transcript_blake3: original_hash,
460        compacted_transcript_blake3: compacted_hash,
461        allocation_plan_id: plan.plan_id.clone(),
462        semantic_memory_fact_ids: Vec::new(),
463        semantic_memory_document_ids: Vec::new(),
464        exact_fallback_refs,
465        summary_loss_report: build_loss_report(&plan, structured_summary),
466        warnings,
467    };
468
469    if request.policy.archive_memory_enabled || request.policy.semantic_memory_enabled {
470        if let Some(sink) = memory_sink {
471            let compact_response = CompactResponse {
472                receipt: receipt.clone(),
473                allocation_plan: plan.clone(),
474                compacted_messages: compacted_messages.clone(),
475                exact_store: exact_store.clone(),
476            };
477            let outcome = archive_response_to_memory(&compact_response, sink)?;
478            receipt.semantic_memory_fact_ids = outcome.record_ids;
479            if outcome.records_attempted == 0 {
480                receipt.warnings.push(
481                    "archive_memory_enabled true but no matching durable/archive candidates were found".to_string(),
482                );
483            }
484            if request.policy.semantic_memory_enabled && !request.policy.archive_memory_enabled {
485                receipt
486                    .warnings
487                    .push("semantic_memory_enabled=true but archive_memory_enabled=false; semantic IDs may be partial".to_string());
488            }
489        } else {
490            receipt.warnings.push(
491                "archive requested but no memory sink was supplied; semantic-memory IDs are intentionally empty".to_string(),
492            );
493        }
494    }
495
496    Ok(CompactResponse {
497        receipt,
498        allocation_plan: plan,
499        compacted_messages,
500        exact_store,
501    })
502}
503
504pub fn filter_recall_candidate(candidate: &RecallCandidate, query: &str) -> RecallFilterResult {
505    let content_l = candidate.content.to_lowercase();
506    let source_l = candidate.source.clone().unwrap_or_default().to_lowercase();
507    let query_l = query.to_lowercase();
508    let mut reasons = Vec::new();
509
510    let speculative_markers = [
511        "likely",
512        "would likely",
513        "potentially",
514        "may indicate",
515        "could imply",
516        "logically connect",
517        "intended to",
518        "seems to",
519    ];
520    if speculative_markers.iter().any(|m| content_l.contains(m)) {
521        reasons.push("speculative-language-without-evidence".to_string());
522        return RecallFilterResult {
523            decision: RecallDecision::QuarantineSpeculative,
524            reasons,
525        };
526    }
527
528    let artifact_markers = [
529        "_spec_pack.zip",
530        "_bundle.zip",
531        "pack_manifest",
532        "mvp_implementation_plan",
533        "rollback_and_quarantine_plan",
534        "master_codex_implementation_prompt",
535        "migration from prior drafts",
536        "profile_i",
537    ];
538    let is_artifact = artifact_markers
539        .iter()
540        .any(|m| content_l.contains(m) || source_l.contains(m));
541    if is_artifact {
542        let explicit = [
543            "aicc",
544            "context pack",
545            "pack manifest",
546            "zpy",
547            "gpt-pro-synthesis",
548        ]
549        .iter()
550        .any(|m| query_l.contains(m));
551        reasons.push("artifact-pack-marker".to_string());
552        return RecallFilterResult {
553            decision: if explicit {
554                RecallDecision::AdmitBackground
555            } else {
556                RecallDecision::RejectNoise
557            },
558            reasons,
559        };
560    }
561
562    if source_l.contains("tool-receipts")
563        && !query_l.contains("tool")
564        && !query_l.contains("receipt")
565    {
566        reasons.push("tool-receipt-not-query-matched".to_string());
567        return RecallFilterResult {
568            decision: RecallDecision::DemoteArtifact,
569            reasons,
570        };
571    }
572
573    if content_l.contains("verified")
574        || content_l.contains("passed")
575        || content_l.contains("source:")
576    {
577        reasons.push("verification-signal".to_string());
578        RecallFilterResult {
579            decision: RecallDecision::AdmitAuthoritative,
580            reasons,
581        }
582    } else {
583        RecallFilterResult {
584            decision: RecallDecision::AdmitBackground,
585            reasons,
586        }
587    }
588}
589
590pub fn approx_tokens_text(text: &str) -> usize {
591    (text.chars().count() / CHARS_PER_TOKEN).max(1)
592}
593
594fn approx_word_tokens(text: &str) -> usize {
595    text.split_whitespace()
596        .map(|token| {
597            if token.ends_with('.')
598                || token.ends_with(',')
599                || token.ends_with(';')
600                || token.ends_with('!')
601                || token.ends_with('?')
602            {
603                (token.chars().count() / CHARS_PER_TOKEN) + 1
604            } else {
605                token.chars().count() / CHARS_PER_TOKEN
606            }
607        })
608        .sum::<usize>()
609        .max(1)
610}
611
612pub fn approx_tokens_messages(messages: &[Message]) -> usize {
613    messages
614        .iter()
615        .map(|m| approx_tokens_text(&m.content) + 4)
616        .sum()
617}
618
619fn count_tokens_text(text: &str, policy: &CompactionPolicy) -> usize {
620    match policy.token_counter {
621        TokenCounterKind::ApproxChars => approx_tokens_text(text),
622        TokenCounterKind::ApproxWords => approx_word_tokens(text),
623        TokenCounterKind::ProviderChatApprox | TokenCounterKind::TiktokenCl100k => {
624            provider_chat_approx_tokens(text)
625        }
626    }
627}
628
629fn count_tokens_messages(messages: &[Message], policy: &CompactionPolicy) -> usize {
630    messages
631        .iter()
632        .map(|m| {
633            count_tokens_text(&m.content, policy)
634                + message_overhead_tokens(policy, &m.role, m.name.as_deref())
635        })
636        .sum()
637}
638
639fn message_overhead_tokens(policy: &CompactionPolicy, role: &str, name: Option<&str>) -> usize {
640    match policy.token_counter {
641        TokenCounterKind::ProviderChatApprox | TokenCounterKind::TiktokenCl100k => {
642            let role_overhead = match role {
643                "system" => 5,
644                "user" => 4,
645                "assistant" => 4,
646                "tool" => 7,
647                _ => 4,
648            };
649            role_overhead + usize::from(name.is_some())
650        }
651        TokenCounterKind::ApproxChars | TokenCounterKind::ApproxWords => 4,
652    }
653}
654
655fn provider_chat_approx_tokens(text: &str) -> usize {
656    let chars = text.chars().count();
657    let whitespace_tokens = text.split_whitespace().count();
658    let punctuation_tokens = text
659        .chars()
660        .filter(|ch| {
661            matches!(
662                ch,
663                '{' | '}' | '[' | ']' | ':' | ',' | '"' | '`' | '/' | '\\'
664            )
665        })
666        .count()
667        / 4;
668    let char_estimate = (chars / 4).max(1);
669    char_estimate
670        .max(whitespace_tokens + punctuation_tokens)
671        .max(1)
672}
673
674pub fn hash_messages(messages: &[Message]) -> Result<String, ContextGovernorError> {
675    let rendered = serde_json::to_string(messages)?;
676    Ok(hash_text(&rendered))
677}
678
679pub fn hash_text(text: &str) -> String {
680    blake3::hash(text.as_bytes()).to_hex().to_string()
681}
682
683fn classify_messages(
684    session_id: &str,
685    messages: &[Message],
686    policy: &CompactionPolicy,
687) -> Vec<ContextItemV1> {
688    let latest_user = messages
689        .iter()
690        .rposition(|m| m.role == "user")
691        .unwrap_or(messages.len() - 1);
692    let mut seen_hashes = BTreeSet::new();
693    messages
694        .iter()
695        .enumerate()
696        .map(|(idx, msg)| {
697            let h = hash_text(&msg.content);
698            let duplicate = !seen_hashes.insert(h.clone());
699            classify_message(session_id, idx, latest_user, msg, h, duplicate, policy)
700        })
701        .collect()
702}
703
704fn classify_message(
705    session_id: &str,
706    idx: usize,
707    latest_user: usize,
708    msg: &Message,
709    content_hash: String,
710    duplicate: bool,
711    policy: &CompactionPolicy,
712) -> ContextItemV1 {
713    let content_l = msg.content.to_lowercase();
714    let aggressive = is_aggressive_allocator(policy);
715    let counted_tokens = count_tokens_text(&msg.content, policy);
716    let msg_chars = msg.content.chars().count();
717    let long_message = msg_chars > 600;
718    let mut item_type = ItemType::LowRiskNarrative;
719    let mut authority = AuthorityClass::SummaryOk;
720    let mut policy = PreservationPolicy::ExtractiveSummary;
721    let mut reasons = Vec::new();
722
723    if idx == latest_user && msg.role == "user" {
724        item_type = ItemType::LatestUserMessage;
725        authority = AuthorityClass::ActiveTask;
726        policy = PreservationPolicy::KeepVerbatim;
727        reasons.push("latest-user-message".to_string());
728    } else if duplicate {
729        item_type = ItemType::DuplicateContext;
730        authority = AuthorityClass::Discardable;
731        policy = PreservationPolicy::OmitDuplicate;
732        reasons.push("duplicate-content".to_string());
733    } else if contains_any(
734        &content_l,
735        &[
736            "acceptance gate",
737            "must pass",
738            "required",
739            "requirement",
740            "do not",
741            "never ",
742        ],
743    ) {
744        item_type = ItemType::AcceptanceGate;
745        authority = AuthorityClass::MustPreserveExact;
746        policy = if aggressive && long_message {
747            PreservationPolicy::ReceiptOnly
748        } else {
749            PreservationPolicy::KeepVerbatim
750        };
751        reasons.push("acceptance-or-instruction".to_string());
752    } else if contains_any(
753        &content_l,
754        &[
755            "error:",
756            "error[",
757            "traceback",
758            "panic",
759            "failed",
760            "compilation failed",
761            "exit_code\":1",
762            "exit code 1",
763        ],
764    ) {
765        item_type = ItemType::ErrorOutput;
766        authority = AuthorityClass::EvidenceCritical;
767        policy = if msg.role == "tool" && msg.content.chars().count() > 600 {
768            PreservationPolicy::ReceiptOnly
769        } else {
770            PreservationPolicy::KeepVerbatim
771        };
772        reasons.push("error-output".to_string());
773    } else if contains_path_signal(&msg.content) {
774        item_type = ItemType::FilePathContext;
775        authority = AuthorityClass::EvidenceCritical;
776        policy = if aggressive && long_message {
777            PreservationPolicy::ReceiptOnly
778        } else {
779            PreservationPolicy::KeepVerbatim
780        };
781        reasons.push("path-signal".to_string());
782    } else if msg.role == "tool" {
783        item_type = ItemType::ToolResult;
784        authority = AuthorityClass::VerifiedToolReceipt;
785        policy = if msg.content.chars().count() > 600 {
786            PreservationPolicy::ReceiptOnly
787        } else {
788            PreservationPolicy::KeepVerbatim
789        };
790        reasons.push("tool-result".to_string());
791    } else if contains_any(
792        &content_l,
793        &["decided", "decision", "architecture", "verdict"],
794    ) {
795        item_type = ItemType::Decision;
796        authority = AuthorityClass::DurableMemoryCandidate;
797        policy = PreservationPolicy::SemanticMemoryArchive;
798        reasons.push("decision-signal".to_string());
799    } else if contains_any(&content_l, &["source:", "evidence", "verified", "receipt"]) {
800        item_type = ItemType::SourceEvidence;
801        authority = AuthorityClass::EvidenceCritical;
802        policy = if aggressive && long_message {
803            PreservationPolicy::ReceiptOnly
804        } else {
805            PreservationPolicy::KeepVerbatim
806        };
807        reasons.push("evidence-signal".to_string());
808    }
809
810    if contains_any(
811        &content_l,
812        &[
813            "likely",
814            "potentially",
815            "would likely",
816            "may indicate",
817            "logically connect",
818        ],
819    ) {
820        item_type = ItemType::ArtifactBoilerplate;
821        authority = AuthorityClass::Quarantine;
822        policy = PreservationPolicy::Quarantine;
823        reasons.push("speculative-language".to_string());
824    }
825
826    ContextItemV1 {
827        schema: "ContextItemV1".to_string(),
828        item_id: format!("ctxi_{idx:04}_{}", &content_hash[..12]),
829        session_id: session_id.to_string(),
830        start_index: idx,
831        end_index: idx,
832        role_set: vec![msg.role.clone()],
833        char_count: msg.content.chars().count(),
834        approx_tokens: counted_tokens,
835        content_blake3: content_hash,
836        content_kind: detect_content_kind(&msg.role, &msg.content),
837        item_type,
838        authority_class: authority,
839        preservation_policy: policy,
840        risk_reasons: reasons,
841        source_message_ids: msg.id.clone().into_iter().collect(),
842        priority_score: 0,
843    }
844}
845
846fn resolve_allocator(policy: &CompactionPolicy) -> AllocatorMode {
847    if policy.allocator.trim().is_empty() {
848        return AllocatorMode::DeterministicV1;
849    }
850    match policy.allocator.trim().to_lowercase().as_str() {
851        "deterministic" | "deterministic_v1" => AllocatorMode::DeterministicV1,
852        "aggressive" | "aggressive_v1" => AllocatorMode::AggressiveV1,
853        _ => AllocatorMode::DeterministicV1,
854    }
855}
856
857fn is_aggressive_allocator(policy: &CompactionPolicy) -> bool {
858    matches!(resolve_allocator(policy), AllocatorMode::AggressiveV1)
859}
860fn is_aggressive_allocator_mode(mode: &AllocatorMode) -> bool {
861    matches!(mode, AllocatorMode::AggressiveV1)
862}
863
864fn score_items(items: &mut [ContextItemV1], focus: Option<&str>) {
865    for item in items {
866        let mut score = match item.item_type {
867            ItemType::LatestUserMessage => 1000,
868            ItemType::ActiveInstruction => 900,
869            ItemType::AcceptanceGate => 850,
870            ItemType::ErrorOutput => 800,
871            ItemType::SourceEvidence => 750,
872            ItemType::ToolResult => 650,
873            ItemType::FilePathContext => 650,
874            ItemType::Decision => 600,
875            ItemType::UnresolvedQuestion => 575,
876            ItemType::DurableFactCandidate | ItemType::ProjectStateCandidate => 500,
877            ItemType::LowRiskNarrative => 100,
878            ItemType::ArtifactBoilerplate => -300,
879            ItemType::StalePlan => -350,
880            ItemType::DuplicateContext => -500,
881            ItemType::Unknown | ItemType::ToolCall => 50,
882        };
883        if matches!(
884            item.authority_class,
885            AuthorityClass::MustPreserveExact | AuthorityClass::ActiveTask
886        ) {
887            score += 300;
888        }
889        if focus.is_some() && !matches!(item.preservation_policy, PreservationPolicy::Quarantine) {
890            score += 25;
891        }
892        item.priority_score = score;
893    }
894}
895
896fn allocate_items(
897    session_id: &str,
898    items: Vec<ContextItemV1>,
899    policy: &CompactionPolicy,
900    allocator: AllocatorMode,
901) -> ContextAllocationPlanV1 {
902    let mut used_tokens = 0usize;
903    let mut kept = Vec::new();
904    let mut summarized = Vec::new();
905    let mut archived = Vec::new();
906    let mut omitted = Vec::new();
907    let mut quarantined = Vec::new();
908
909    let len = items.len();
910    let aggressive = is_aggressive_allocator_mode(&allocator);
911    for item in &items {
912        let in_head = item.start_index < policy.protect_first_n;
913        let in_tail = item.start_index + policy.protect_last_n >= len;
914        let protected_head = in_head && (!aggressive || item.approx_tokens <= 1_000);
915        let protected_tail = in_tail
916            && (!aggressive
917                || item.approx_tokens <= 300
918                || matches!(item.item_type, ItemType::LatestUserMessage));
919        match item.preservation_policy {
920            PreservationPolicy::Quarantine => quarantined.push(item.item_id.clone()),
921            PreservationPolicy::OmitDuplicate => omitted.push(item.item_id.clone()),
922            PreservationPolicy::SemanticMemoryArchive => {
923                archived.push(item.item_id.clone());
924                if protected_tail || item.priority_score >= 650 {
925                    kept.push(item.item_id.clone());
926                    used_tokens = used_tokens.saturating_add(item.approx_tokens);
927                } else {
928                    summarized.push(item.item_id.clone());
929                }
930            }
931            PreservationPolicy::KeepVerbatim => {
932                let must_keep = matches!(item.item_type, ItemType::LatestUserMessage)
933                    || protected_head
934                    || (!aggressive && protected_tail);
935                let fits_aggressive_budget = !aggressive
936                    || used_tokens + item.approx_tokens <= policy.target_tokens.saturating_div(2);
937                if must_keep || fits_aggressive_budget {
938                    kept.push(item.item_id.clone());
939                    used_tokens = used_tokens.saturating_add(item.approx_tokens);
940                } else {
941                    summarized.push(item.item_id.clone());
942                }
943            }
944            PreservationPolicy::ReceiptOnly => {
945                if protected_head
946                    || protected_tail
947                    || used_tokens + item.approx_tokens <= policy.target_tokens
948                {
949                    summarized.push(item.item_id.clone());
950                } else {
951                    omitted.push(item.item_id.clone());
952                }
953            }
954            PreservationPolicy::ExtractiveSummary | PreservationPolicy::AbstractiveSummary => {
955                let must_keep = matches!(item.item_type, ItemType::LatestUserMessage)
956                    || protected_head
957                    || protected_tail;
958                if must_keep
959                    || (!aggressive && used_tokens + item.approx_tokens <= policy.target_tokens / 2)
960                {
961                    kept.push(item.item_id.clone());
962                    used_tokens = used_tokens.saturating_add(item.approx_tokens);
963                } else {
964                    summarized.push(item.item_id.clone());
965                }
966            }
967        }
968    }
969
970    ContextAllocationPlanV1 {
971        schema: "ContextAllocationPlanV1".to_string(),
972        plan_id: format!("ctxp_{}", Uuid::new_v4().simple()),
973        session_id: session_id.to_string(),
974        created_utc: Utc::now(),
975        context_budget_tokens: policy.target_tokens,
976        target_output_tokens: policy.target_tokens,
977        allocator: allocator.as_str().to_string(),
978        items,
979        kept_item_ids: kept,
980        summarized_item_ids: summarized,
981        archived_item_ids: archived,
982        omitted_item_ids: omitted,
983        quarantined_item_ids: quarantined,
984    }
985}
986
987fn build_exact_store(
988    messages: &[Message],
989    plan: &ContextAllocationPlanV1,
990) -> Vec<ExactStoredItemV1> {
991    plan.items
992        .iter()
993        .filter(|item| {
994            plan.summarized_item_ids.contains(&item.item_id)
995                || plan.archived_item_ids.contains(&item.item_id)
996                || plan.omitted_item_ids.contains(&item.item_id)
997                || plan.quarantined_item_ids.contains(&item.item_id)
998                || matches!(item.preservation_policy, PreservationPolicy::ReceiptOnly)
999        })
1000        .filter_map(|item| {
1001            messages.get(item.start_index).map(|m| ExactStoredItemV1 {
1002                item_id: item.item_id.clone(),
1003                source_indices: vec![item.start_index],
1004                content: m.content.clone(),
1005                content_blake3: item.content_blake3.clone(),
1006            })
1007        })
1008        .collect()
1009}
1010
1011fn build_structured_summary(
1012    messages: &[Message],
1013    plan: &ContextAllocationPlanV1,
1014) -> StructuredContextSummaryV1 {
1015    let mut out = StructuredContextSummaryV1 {
1016        active_task: messages
1017            .iter()
1018            .rev()
1019            .find(|m| m.role == "user")
1020            .map(|m| compact_preview(&m.content, 240)),
1021        fallback_item_ids: plan
1022            .summarized_item_ids
1023            .iter()
1024            .chain(plan.omitted_item_ids.iter())
1025            .chain(plan.quarantined_item_ids.iter())
1026            .cloned()
1027            .collect(),
1028        ..Default::default()
1029    };
1030
1031    for msg in messages {
1032        let content_l = msg.content.to_lowercase();
1033        if content_l.contains("acceptance gate") || content_l.contains("must pass") {
1034            push_unique(
1035                &mut out.acceptance_gates,
1036                compact_preview(&msg.content, 240),
1037            );
1038        }
1039        if content_l.contains("decision:") || content_l.contains("decided") {
1040            push_unique(&mut out.decisions, compact_preview(&msg.content, 240));
1041        }
1042        if content_l.contains("unresolved question") || content_l.contains('?') {
1043            push_unique(
1044                &mut out.unresolved_questions,
1045                compact_preview(&msg.content, 240),
1046            );
1047        }
1048        if content_l.contains("error")
1049            || content_l.contains("traceback")
1050            || content_l.contains("failed")
1051        {
1052            push_unique(
1053                &mut out.errors,
1054                important_lines(&msg.content, &["error", "failed", "traceback"], 240),
1055            );
1056        }
1057        for file in extract_file_like_tokens(&msg.content) {
1058            push_unique(&mut out.files, file);
1059        }
1060        for command in extract_command_like_lines(&msg.content) {
1061            push_unique(&mut out.commands, command);
1062        }
1063    }
1064    out
1065}
1066
1067fn push_summary_list(lines: &mut Vec<String>, label: &str, values: &[String]) {
1068    if values.is_empty() {
1069        return;
1070    }
1071    lines.push(format!("- {label}:"));
1072    for value in values.iter().take(5) {
1073        lines.push(format!("  - {}", compact_preview(value, 220)));
1074    }
1075}
1076
1077fn push_unique(values: &mut Vec<String>, value: String) {
1078    if !value.is_empty() && !values.contains(&value) {
1079        values.push(value);
1080    }
1081}
1082
1083fn extract_file_like_tokens(text: &str) -> Vec<String> {
1084    text.split_whitespace()
1085        .map(|token| {
1086            token.trim_matches(|c: char| {
1087                c == ',' || c == ':' || c == ';' || c == ')' || c == '(' || c == '`' || c == '"'
1088            })
1089        })
1090        .filter(|token| {
1091            token.contains("/home/")
1092                || token.contains("/usr/")
1093                || token.contains("/etc/")
1094                || token.contains("/var/")
1095                || token.contains("/tmp/")
1096                || token.starts_with("~/")
1097                || token.starts_with("/")
1098                || token.ends_with(".rs")
1099                || token.ends_with(".py")
1100                || token.ends_with(".ts")
1101                || token.ends_with(".tsx")
1102                || token.ends_with(".js")
1103                || token.ends_with(".md")
1104                || token.ends_with(".toml")
1105                || token.ends_with(".yaml")
1106                || token.ends_with(".yml")
1107                || token.ends_with(".json")
1108                || token.ends_with(".go")
1109                || token.ends_with(".c")
1110                || token.ends_with(".h")
1111        })
1112        .map(|token| token.to_string())
1113        .collect()
1114}
1115
1116fn extract_command_like_lines(text: &str) -> Vec<String> {
1117    text.lines()
1118        .map(str::trim)
1119        .filter(|line| {
1120            line.starts_with("running:")
1121                || line.starts_with('$')
1122                || line.starts_with("cargo ")
1123                || line.starts_with("npm ")
1124                || line.starts_with("python ")
1125                || line.starts_with("git ")
1126        })
1127        .map(|line| compact_preview(line, 240))
1128        .collect()
1129}
1130
1131fn important_lines(text: &str, needles: &[&str], max_chars: usize) -> String {
1132    let mut out = Vec::new();
1133    for line in text.lines() {
1134        let line_l = line.to_lowercase();
1135        if needles.iter().any(|needle| line_l.contains(needle)) {
1136            out.push(line.trim());
1137        }
1138        if out.len() >= 5 {
1139            break;
1140        }
1141    }
1142    if out.is_empty() {
1143        compact_preview(text, max_chars)
1144    } else {
1145        compact_preview(&out.join(" | "), max_chars)
1146    }
1147}
1148
1149fn build_summary(
1150    messages: &[Message],
1151    plan: &ContextAllocationPlanV1,
1152    policy: &CompactionPolicy,
1153    structured: &StructuredContextSummaryV1,
1154    receipt_id: &str,
1155) -> String {
1156    let mut lines = vec![
1157        format!("{SUMMARY_PREFIX}"),
1158        "Earlier turns were compacted. This summary is background, not an active task.".to_string(),
1159        "Only the latest user message after this summary is active.".to_string(),
1160        format!("Use context_expand(receipt_id=\"{receipt_id}\", item_id=...) to recover exact omitted text when needed."),
1161        String::new(),
1162        "## Structured anchors".to_string(),
1163    ];
1164    if let Some(active_task) = &structured.active_task {
1165        lines.push(format!(
1166            "- active_task: {}",
1167            compact_preview(active_task, 220)
1168        ));
1169    }
1170    push_summary_list(&mut lines, "acceptance_gates", &structured.acceptance_gates);
1171    push_summary_list(&mut lines, "commands", &structured.commands);
1172    push_summary_list(&mut lines, "errors", &structured.errors);
1173    push_summary_list(&mut lines, "files", &structured.files);
1174    push_summary_list(&mut lines, "decisions", &structured.decisions);
1175    push_summary_list(
1176        &mut lines,
1177        "unresolved_questions",
1178        &structured.unresolved_questions,
1179    );
1180    if !structured.fallback_item_ids.is_empty() {
1181        let shown = structured
1182            .fallback_item_ids
1183            .iter()
1184            .take(30)
1185            .cloned()
1186            .collect::<Vec<_>>();
1187        let hidden = structured
1188            .fallback_item_ids
1189            .len()
1190            .saturating_sub(shown.len());
1191        let suffix = if hidden > 0 {
1192            format!(" (+{hidden} more in receipt)")
1193        } else {
1194            String::new()
1195        };
1196        lines.push(format!(
1197            "- fallback_item_ids: {}{}",
1198            shown.join(", "),
1199            suffix
1200        ));
1201    }
1202
1203    lines.push(String::new());
1204    lines.push("## Preserved / summarized context".to_string());
1205    for item in &plan.items {
1206        if plan.summarized_item_ids.contains(&item.item_id)
1207            || plan.archived_item_ids.contains(&item.item_id)
1208            || matches!(item.preservation_policy, PreservationPolicy::ReceiptOnly)
1209        {
1210            if let Some(msg) = messages.get(item.start_index) {
1211                let preview = content_aware_preview(&item.content_kind, &msg.content, 220);
1212                lines.push(format!(
1213                    "- {} {:?}/{:?}/{:?}: {}",
1214                    item.item_id,
1215                    item.item_type,
1216                    item.content_kind,
1217                    item.preservation_policy,
1218                    preview
1219                ));
1220            }
1221        }
1222    }
1223    if !plan.quarantined_item_ids.is_empty() {
1224        lines.push(String::new());
1225        lines.push(format!(
1226            "## Quarantined context\n{} items were excluded from authoritative context due to speculative/artifact signals; exact fallback remains available.",
1227            plan.quarantined_item_ids.len()
1228        ));
1229    }
1230    let mut summary = lines.join("\n");
1231    if summary.chars().count() > policy.summary_max_chars {
1232        summary = summary
1233            .chars()
1234            .take(policy.summary_max_chars)
1235            .collect::<String>();
1236        summary.push_str("\n[summary truncated by context-governor budget]");
1237    }
1238    summary
1239}
1240
1241fn assemble_compacted_messages(
1242    messages: &[Message],
1243    plan: &ContextAllocationPlanV1,
1244    summary: &str,
1245    policy: &CompactionPolicy,
1246) -> Vec<Message> {
1247    let mut out = Vec::new();
1248    let kept_set = plan.kept_item_ids.iter().cloned().collect::<BTreeSet<_>>();
1249    let aggressive = is_aggressive_allocator(policy);
1250    let item_by_index = plan
1251        .items
1252        .iter()
1253        .map(|i| (i.start_index, i))
1254        .collect::<BTreeMap<_, _>>();
1255
1256    let tail_start = messages.len().saturating_sub(policy.protect_last_n);
1257
1258    for (idx, msg) in messages.iter().enumerate() {
1259        if idx >= tail_start {
1260            continue;
1261        }
1262        if idx < policy.protect_first_n {
1263            let Some(item) = item_by_index.get(&idx) else {
1264                out.push(msg.clone());
1265                continue;
1266            };
1267            let keep_head = matches!(item.item_type, ItemType::LatestUserMessage)
1268                || (kept_set.contains(&item.item_id)
1269                    && (!aggressive || item.approx_tokens <= 1_000))
1270                || (matches!(
1271                    item.authority_class,
1272                    AuthorityClass::MustPreserveExact | AuthorityClass::ActiveTask
1273                ) && (!aggressive || item.approx_tokens <= 1_000));
1274            if keep_head {
1275                out.push(msg.clone());
1276            }
1277        } else if let Some(item) = item_by_index.get(&idx) {
1278            if kept_set.contains(&item.item_id)
1279                && !out
1280                    .iter()
1281                    .any(|m: &Message| m.content == msg.content && m.role == msg.role)
1282            {
1283                out.push(msg.clone());
1284            }
1285        }
1286    }
1287
1288    let summary_msg = Message {
1289        id: Some(format!("summary_{}", plan.plan_id)),
1290        role: choose_summary_role(out.last().map(|m| m.role.as_str())).to_string(),
1291        content: summary.to_string(),
1292        name: Some("context_governor".to_string()),
1293        metadata: BTreeMap::from([("compressed_summary".to_string(), Value::Bool(true))]),
1294    };
1295    out.push(summary_msg);
1296
1297    for (offset, msg) in messages[tail_start..].iter().enumerate() {
1298        let idx = tail_start + offset;
1299        let _ = item_by_index.get(&idx); // item exists but we always keep tail
1300                                         // Always include tail messages within protect_last_n, regardless of
1301                                         // kept_set. The built-in compressor always includes the last N messages;
1302                                         // excluding them because the allocator summarized them causes data loss
1303                                         // of recent context the model needs.
1304        let should_keep_tail = !out
1305            .iter()
1306            .any(|m: &Message| m.content == msg.content && m.role == msg.role);
1307        if should_keep_tail {
1308            out.push(msg.clone());
1309        }
1310    }
1311    out
1312}
1313
1314fn choose_summary_role(previous: Option<&str>) -> &'static str {
1315    match previous {
1316        Some("assistant") => "user",
1317        _ => "assistant",
1318    }
1319}
1320
1321fn build_loss_report(
1322    plan: &ContextAllocationPlanV1,
1323    structured_summary: StructuredContextSummaryV1,
1324) -> SummaryLossReportV1 {
1325    SummaryLossReportV1 {
1326        schema: "SummaryLossReportV1".to_string(),
1327        preserved_claims: plan.kept_item_ids.clone(),
1328        omitted_claims: plan.omitted_item_ids.clone(),
1329        evidence_lost: plan
1330            .omitted_item_ids
1331            .iter()
1332            .chain(plan.summarized_item_ids.iter())
1333            .cloned()
1334            .collect(),
1335        uncertainty_introduced: plan.quarantined_item_ids.clone(),
1336        exact_recovery_available: true,
1337        high_risk_omissions: plan
1338            .items
1339            .iter()
1340            .filter(|i| {
1341                plan.omitted_item_ids.contains(&i.item_id)
1342                    && matches!(
1343                        i.authority_class,
1344                        AuthorityClass::MustPreserveExact | AuthorityClass::EvidenceCritical
1345                    )
1346            })
1347            .map(|i| i.item_id.clone())
1348            .collect(),
1349        structured_summary,
1350    }
1351}
1352
1353fn compact_preview(text: &str, max_chars: usize) -> String {
1354    let clean = text.split_whitespace().collect::<Vec<_>>().join(" ");
1355    if clean.chars().count() <= max_chars {
1356        clean
1357    } else {
1358        let mut s = clean.chars().take(max_chars).collect::<String>();
1359        s.push_str("...");
1360        s
1361    }
1362}
1363
1364fn content_aware_preview(kind: &ContentKind, text: &str, max_chars: usize) -> String {
1365    match kind {
1366        ContentKind::Json => json_preview(text, max_chars),
1367        ContentKind::Diff => important_lines(text, &["diff --git", "@@", "+++", "---"], max_chars),
1368        ContentKind::CargoOutput => important_lines(
1369            text,
1370            &["error", "warning", "test result", "failed", "finished"],
1371            max_chars,
1372        ),
1373        ContentKind::ShellLog => {
1374            important_lines(text, &["error", "failed", "exit", "$"], max_chars)
1375        }
1376        ContentKind::SearchResults => important_lines(
1377            text,
1378            &["/home/", ":", "LINE", "match", ".rs", ".py", ".md"],
1379            max_chars,
1380        ),
1381        ContentKind::Rust => important_lines(
1382            text,
1383            &["pub ", "fn ", "struct ", "impl ", "use ", "error", "TODO"],
1384            max_chars,
1385        ),
1386        ContentKind::Markdown => important_lines(text, &["#", "##", "- ", "```"], max_chars),
1387        _ => compact_preview(text, max_chars),
1388    }
1389}
1390
1391fn json_preview(text: &str, max_chars: usize) -> String {
1392    if let Ok(value) = serde_json::from_str::<Value>(text) {
1393        if let Some(obj) = value.as_object() {
1394            let keys = obj.keys().take(12).cloned().collect::<Vec<_>>().join(", ");
1395            return compact_preview(&format!("json keys: {keys}"), max_chars);
1396        }
1397    }
1398    compact_preview(text, max_chars)
1399}
1400
1401fn detect_content_kind(role: &str, content: &str) -> ContentKind {
1402    let trimmed = content.trim();
1403    let lower = trimmed.to_lowercase();
1404    if serde_json::from_str::<Value>(trimmed).is_ok() {
1405        return ContentKind::Json;
1406    }
1407    if lower.starts_with("diff --git") || lower.contains("\n@@ ") {
1408        return ContentKind::Diff;
1409    }
1410    if lower.contains("error[")
1411        || lower.contains("test result:")
1412        || lower.contains("warning:")
1413        || lower.contains("compiling ")
1414    {
1415        return ContentKind::CargoOutput;
1416    }
1417    if trimmed.starts_with("# ") || trimmed.contains("\n## ") || trimmed.contains("\n- ") {
1418        return ContentKind::Markdown;
1419    }
1420    if lower.contains("fn main") || lower.contains("pub struct") || lower.contains("impl ") {
1421        return ContentKind::Rust;
1422    }
1423    if looks_like_search_results(trimmed) {
1424        return ContentKind::SearchResults;
1425    }
1426    if role == "tool" {
1427        return ContentKind::ShellLog;
1428    }
1429    ContentKind::PlainText
1430}
1431
1432fn looks_like_search_results(text: &str) -> bool {
1433    let mut path_line_count = 0usize;
1434    for line in text.lines().take(1_000) {
1435        let lower = line.to_lowercase();
1436        let has_path = lower.contains("/home/")
1437            || lower.contains("/tmp/")
1438            || lower.contains(".rs:")
1439            || lower.contains(".py:")
1440            || lower.contains(".md:")
1441            || lower.contains(".toml:")
1442            || lower.contains(".yaml:")
1443            || lower.contains(".json:");
1444        let has_location_separator = line.matches(':').count() >= 2;
1445        if has_path && has_location_separator {
1446            path_line_count += 1;
1447        }
1448    }
1449    path_line_count >= 1
1450}
1451
1452fn enforce_budget(
1453    mut messages: Vec<Message>,
1454    policy: &CompactionPolicy,
1455    warnings: &mut Vec<String>,
1456) -> Result<Vec<Message>, ContextGovernorError> {
1457    let target = policy.target_tokens;
1458    if count_tokens_messages(&messages, policy) <= target {
1459        return Ok(messages);
1460    }
1461    let Some(summary_idx) = messages
1462        .iter()
1463        .position(|m| m.content.starts_with(SUMMARY_PREFIX))
1464    else {
1465        let minimum = count_tokens_messages(&messages, policy);
1466        return Err(ContextGovernorError::BudgetExceeded {
1467            target_tokens: target,
1468            minimum_required_tokens: minimum,
1469        });
1470    };
1471
1472    let mut without_summary = messages.clone();
1473    without_summary.remove(summary_idx);
1474    let required = count_tokens_messages(&without_summary, policy);
1475    if required > target {
1476        if matches!(policy.budget_mode, BudgetMode::HardCascade) {
1477            warnings.push(format!(
1478                "hard cascade removed summary but exact protected minimum still exceeds target budget: {required} > {target} tokens"
1479            ));
1480            return Ok(without_summary);
1481        }
1482        return Err(ContextGovernorError::BudgetExceeded {
1483            target_tokens: target,
1484            minimum_required_tokens: required,
1485        });
1486    }
1487
1488    let available_chars = target
1489        .saturating_sub(required)
1490        .saturating_mul(CHARS_PER_TOKEN);
1491    let minimal_summary = minimal_recovery_summary(&messages[summary_idx].content);
1492    let minimal_tokens = count_tokens_text(&minimal_summary, policy) + 4;
1493    if available_chars < 80 {
1494        if required + minimal_tokens <= target {
1495            messages[summary_idx].content = minimal_summary;
1496            warnings.push("hard cascade reduced summary to minimal recovery pointer".to_string());
1497        } else {
1498            messages.remove(summary_idx);
1499            warnings.push("hard cascade removed summary to satisfy target budget".to_string());
1500        }
1501        return Ok(messages);
1502    }
1503
1504    let summary = &mut messages[summary_idx];
1505    if summary.content.chars().count() > available_chars {
1506        let truncated_body = summary
1507            .content
1508            .chars()
1509            .take(available_chars.saturating_sub(96))
1510            .collect::<String>();
1511        summary.content = format!(
1512            "{}\nhard cascade summary truncated; exact fallback remains available.\n{}",
1513            SUMMARY_PREFIX, truncated_body
1514        );
1515        if count_tokens_messages(&messages, policy) > target && required + minimal_tokens <= target
1516        {
1517            messages[summary_idx].content = minimal_summary;
1518        }
1519        warnings.push("hard cascade truncated summary to satisfy target budget".to_string());
1520    }
1521    if count_tokens_messages(&messages, policy) > target {
1522        messages.remove(summary_idx);
1523        warnings.push("hard cascade removed summary after truncation did not fit".to_string());
1524    }
1525    Ok(messages)
1526}
1527
1528fn minimal_recovery_summary(summary: &str) -> String {
1529    let mut lines = vec![
1530        SUMMARY_PREFIX.to_string(),
1531        "Hard cascade kept only the recovery pointer; exact fallback remains available."
1532            .to_string(),
1533    ];
1534    for line in summary.lines() {
1535        if line.contains("context_expand(") || line.starts_with("- fallback_item_ids:") {
1536            lines.push(compact_preview(line, 600));
1537        }
1538    }
1539    lines.join("\n")
1540}
1541
1542fn contains_any(text: &str, needles: &[&str]) -> bool {
1543    needles.iter().any(|needle| text.contains(needle))
1544}
1545
1546fn contains_path_signal(text: &str) -> bool {
1547    text.contains("/home/")
1548        || text.contains("~/")
1549        || text.contains("Cargo.toml")
1550        || text.contains("pyproject.toml")
1551        || text.contains(".rs")
1552        || text.contains(".py")
1553        || text.contains(".ts")
1554        || text.contains(".tsx")
1555        || text.contains(".md")
1556        || text.contains(".yaml")
1557}
1558
1559#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
1560#[serde(rename_all = "snake_case")]
1561pub enum SearchScope {
1562    All,
1563    Summary,
1564    ExactStore,
1565    Receipt,
1566}
1567
1568#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
1569pub struct ContextSearchHit {
1570    pub source: String,
1571    pub item_id: Option<String>,
1572    pub snippet: String,
1573    pub content_blake3: Option<String>,
1574}
1575
1576#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
1577pub struct ContextExpandResult {
1578    pub item_id: String,
1579    pub content: String,
1580    pub content_blake3: String,
1581    pub truncated: bool,
1582}
1583
1584#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
1585pub struct ContextDiffSummary {
1586    pub kept_count: usize,
1587    pub summarized_count: usize,
1588    pub archived_count: usize,
1589    pub omitted_count: usize,
1590    pub quarantined_count: usize,
1591    pub original_approx_tokens: usize,
1592    pub compacted_approx_tokens: usize,
1593    pub token_savings_estimate: isize,
1594    pub warnings: Vec<String>,
1595    #[serde(default)]
1596    pub content_kind_reductions: BTreeMap<String, usize>,
1597}
1598
1599pub fn context_expand(
1600    response: &CompactResponse,
1601    item_id: &str,
1602    max_chars: usize,
1603) -> Option<ContextExpandResult> {
1604    response
1605        .exact_store
1606        .iter()
1607        .find(|item| item.item_id == item_id)
1608        .map(|item| {
1609            let total_chars = item.content.chars().count();
1610            let truncated = total_chars > max_chars;
1611            let content = if truncated {
1612                item.content.chars().take(max_chars).collect::<String>()
1613            } else {
1614                item.content.clone()
1615            };
1616            ContextExpandResult {
1617                item_id: item.item_id.clone(),
1618                content,
1619                content_blake3: item.content_blake3.clone(),
1620                truncated,
1621            }
1622        })
1623}
1624
1625pub fn context_search(
1626    response: &CompactResponse,
1627    query: &str,
1628    top_k: usize,
1629    scope: SearchScope,
1630) -> Vec<ContextSearchHit> {
1631    let query_l = query.to_lowercase();
1632    let mut scored_hits = Vec::new();
1633
1634    let mut add_hit = |source: &'static str,
1635                       rank: usize,
1636                       item_id: Option<String>,
1637                       rendered: &str,
1638                       hash: Option<String>| {
1639        if rendered.to_lowercase().contains(&query_l) {
1640            scored_hits.push((
1641                rank,
1642                ContextSearchHit {
1643                    source: source.to_string(),
1644                    item_id,
1645                    snippet: snippet_around(rendered, query, 320),
1646                    content_blake3: hash,
1647                },
1648            ));
1649        }
1650    };
1651
1652    if matches!(scope, SearchScope::All | SearchScope::ExactStore) {
1653        for item in &response.exact_store {
1654            add_hit(
1655                "exact_store",
1656                90,
1657                Some(item.item_id.clone()),
1658                &item.content,
1659                Some(item.content_blake3.clone()),
1660            );
1661        }
1662    }
1663
1664    if matches!(scope, SearchScope::All | SearchScope::Summary) {
1665        for msg in &response.compacted_messages {
1666            add_hit(
1667                "compacted_messages",
1668                60,
1669                msg.id.clone(),
1670                &msg.content,
1671                Some(hash_text(&msg.content)),
1672            );
1673        }
1674    }
1675
1676    if matches!(scope, SearchScope::All | SearchScope::Receipt) {
1677        if let Ok(rendered) = serde_json::to_string(&response.receipt) {
1678            add_hit(
1679                "receipt",
1680                40,
1681                Some(response.receipt.receipt_id.clone()),
1682                &rendered,
1683                Some(hash_text(&rendered)),
1684            );
1685        }
1686    }
1687
1688    scored_hits.sort_by(|a, b| {
1689        b.0.cmp(&a.0)
1690            .then_with(|| a.1.source.cmp(&b.1.source))
1691            .then_with(|| a.1.snippet.cmp(&b.1.snippet))
1692    });
1693
1694    let mut seen = std::collections::HashSet::new();
1695    let mut hits = Vec::new();
1696    for (_, hit) in scored_hits {
1697        let key = format!(
1698            "{}:{}",
1699            hit.source,
1700            hit.item_id.clone().unwrap_or_else(|| "receipt".to_string())
1701        );
1702        if seen.insert(key) {
1703            hits.push(hit);
1704            if hits.len() >= top_k {
1705                break;
1706            }
1707        }
1708    }
1709
1710    hits
1711}
1712
1713pub fn context_diff(response: &CompactResponse) -> ContextDiffSummary {
1714    let mut content_kind_reductions: BTreeMap<String, usize> = BTreeMap::new();
1715    let reduced_ids = response
1716        .allocation_plan
1717        .summarized_item_ids
1718        .iter()
1719        .chain(response.allocation_plan.archived_item_ids.iter())
1720        .chain(response.allocation_plan.omitted_item_ids.iter())
1721        .chain(response.allocation_plan.quarantined_item_ids.iter())
1722        .cloned()
1723        .collect::<BTreeSet<_>>();
1724    for item in &response.allocation_plan.items {
1725        if reduced_ids.contains(&item.item_id) {
1726            *content_kind_reductions
1727                .entry(format!("{:?}", item.content_kind))
1728                .or_insert(0) += 1;
1729        }
1730    }
1731    ContextDiffSummary {
1732        kept_count: response.allocation_plan.kept_item_ids.len(),
1733        summarized_count: response.allocation_plan.summarized_item_ids.len(),
1734        archived_count: response.allocation_plan.archived_item_ids.len(),
1735        omitted_count: response.allocation_plan.omitted_item_ids.len(),
1736        quarantined_count: response.allocation_plan.quarantined_item_ids.len(),
1737        original_approx_tokens: response.receipt.original_approx_tokens,
1738        compacted_approx_tokens: response.receipt.compacted_approx_tokens,
1739        token_savings_estimate: response.receipt.token_savings_estimate,
1740        warnings: response.receipt.warnings.clone(),
1741        content_kind_reductions,
1742    }
1743}
1744
1745fn snippet_around(content: &str, query: &str, max_chars: usize) -> String {
1746    let content_l = content.to_lowercase();
1747    let query_l = query.to_lowercase();
1748    let start_byte = content_l
1749        .find(&query_l)
1750        .map(|idx| floor_char_boundary(content, idx.saturating_sub(max_chars / 3)))
1751        .unwrap_or(0);
1752    let mut snippet = content[start_byte..]
1753        .chars()
1754        .take(max_chars)
1755        .collect::<String>();
1756    if snippet.chars().count() == max_chars {
1757        snippet.push_str("...");
1758    }
1759    snippet
1760}
1761
1762fn floor_char_boundary(text: &str, mut idx: usize) -> usize {
1763    idx = idx.min(text.len());
1764    while idx > 0 && !text.is_char_boundary(idx) {
1765        idx -= 1;
1766    }
1767    idx
1768}
1769
1770#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
1771pub struct ReplayProbe {
1772    pub category: String,
1773    pub needle: String,
1774}
1775
1776#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1777pub struct ReplayBaselineScore {
1778    pub name: String,
1779    pub visible_probes: usize,
1780    pub recoverable_probes: usize,
1781    pub total_probes: usize,
1782    pub visible_rate: f64,
1783    pub recoverable_rate: f64,
1784    pub tokens: usize,
1785    pub active_task_visible: bool,
1786}
1787
1788#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1789pub struct ReplayFixtureReport {
1790    pub fixture_id: String,
1791    pub original_message_count: usize,
1792    pub compacted_message_count: usize,
1793    pub probes: Vec<ReplayProbe>,
1794    pub baselines: Vec<ReplayBaselineScore>,
1795    pub receipt_id: String,
1796    pub warnings: Vec<String>,
1797}
1798
1799#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
1800pub struct ReplayAnswerabilityQuestion {
1801    pub question: String,
1802    pub expected_terms: Vec<String>,
1803    pub forbidden_terms: Vec<String>,
1804}
1805
1806#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1807pub struct ReplayAnswerabilityBaselineScore {
1808    pub name: String,
1809    pub answerable_questions: usize,
1810    pub total_questions: usize,
1811    pub answerability_rate: f64,
1812    pub incorrect_action_risk: usize,
1813    pub tokens: usize,
1814    pub active_task_visible: bool,
1815}
1816
1817#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1818pub struct ReplayAnswerabilityReport {
1819    pub fixture_id: String,
1820    pub baselines: Vec<ReplayAnswerabilityBaselineScore>,
1821    pub receipt_id: String,
1822    pub warnings: Vec<String>,
1823}
1824
1825pub fn build_replay_probes(messages: &[Message], max_probes: usize) -> Vec<ReplayProbe> {
1826    let mut probes = Vec::new();
1827    if let Some(latest_user) = messages.iter().rev().find(|message| message.role == "user") {
1828        push_probe(
1829            &mut probes,
1830            "active_task",
1831            literal_probe_fragment(&latest_user.content, 120),
1832            max_probes,
1833        );
1834    }
1835
1836    for message in messages {
1837        let content_l = message.content.to_lowercase();
1838        if content_l.contains("acceptance gate") || content_l.contains("must pass") {
1839            push_probe(
1840                &mut probes,
1841                "acceptance_gate",
1842                important_lines(&message.content, &["acceptance gate", "must pass"], 160),
1843                max_probes,
1844            );
1845        }
1846        if content_l.contains("decision:") || content_l.contains("decided") {
1847            push_probe(
1848                &mut probes,
1849                "decision",
1850                important_lines(&message.content, &["decision", "decided"], 160),
1851                max_probes,
1852            );
1853        }
1854        if content_l.contains("error")
1855            || content_l.contains("traceback")
1856            || content_l.contains("failed")
1857        {
1858            push_probe(
1859                &mut probes,
1860                "error",
1861                important_lines(&message.content, &["error", "traceback", "failed"], 160),
1862                max_probes,
1863            );
1864        }
1865        for file in extract_file_like_tokens(&message.content) {
1866            push_probe(&mut probes, "file", file, max_probes);
1867        }
1868        if probes.len() >= max_probes {
1869            break;
1870        }
1871    }
1872    probes
1873}
1874
1875pub fn evaluate_replay_answerability(
1876    fixture_id: impl Into<String>,
1877    request: CompactRequest,
1878    questions: Vec<ReplayAnswerabilityQuestion>,
1879) -> Result<ReplayAnswerabilityReport, ContextGovernorError> {
1880    let fixture_id = fixture_id.into();
1881    let active_task = request
1882        .messages
1883        .iter()
1884        .rev()
1885        .find(|message| message.role == "user")
1886        .map(|message| literal_probe_fragment(&message.content, 120))
1887        .unwrap_or_default();
1888
1889    let full_text = join_message_content(&request.messages);
1890    let full_tokens = approx_tokens_messages(&request.messages);
1891    let head_tail_messages = head_tail_messages(&request.messages, 1, 1);
1892    let head_tail_text = join_message_content(&head_tail_messages);
1893    let head_tail_tokens = approx_tokens_messages(&head_tail_messages);
1894
1895    let response = compact_context(request)?;
1896    let compacted_text = join_message_content(&response.compacted_messages);
1897    let compacted_tokens = response.receipt.compacted_approx_tokens;
1898
1899    let baselines = vec![
1900        score_answerability_text_baseline(
1901            "full",
1902            &full_text,
1903            full_tokens,
1904            &questions,
1905            &active_task,
1906        ),
1907        score_answerability_text_baseline(
1908            "head_tail",
1909            &head_tail_text,
1910            head_tail_tokens,
1911            &questions,
1912            &active_task,
1913        ),
1914        score_answerability_context_governor(
1915            &response,
1916            &compacted_text,
1917            compacted_tokens,
1918            &questions,
1919            &active_task,
1920        ),
1921    ];
1922
1923    Ok(ReplayAnswerabilityReport {
1924        fixture_id,
1925        baselines,
1926        receipt_id: response.receipt.receipt_id,
1927        warnings: response.receipt.warnings,
1928    })
1929}
1930
1931pub fn evaluate_replay_fixture(
1932    fixture_id: impl Into<String>,
1933    request: CompactRequest,
1934    max_probes: usize,
1935) -> Result<ReplayFixtureReport, ContextGovernorError> {
1936    let fixture_id = fixture_id.into();
1937    let original_message_count = request.messages.len();
1938    let probes = build_replay_probes(&request.messages, max_probes);
1939    let active_task = probes
1940        .iter()
1941        .find(|probe| probe.category == "active_task")
1942        .map(|probe| probe.needle.as_str())
1943        .unwrap_or("");
1944
1945    let full_text = join_message_content(&request.messages);
1946    let full_tokens = approx_tokens_messages(&request.messages);
1947    let head_tail_messages = head_tail_messages(&request.messages, 1, 1);
1948    let head_tail_text = join_message_content(&head_tail_messages);
1949    let head_tail_tokens = approx_tokens_messages(&head_tail_messages);
1950
1951    let response = compact_context(request)?;
1952    let compacted_text = join_message_content(&response.compacted_messages);
1953    let compacted_tokens = response.receipt.compacted_approx_tokens;
1954
1955    let baselines = vec![
1956        score_text_baseline("full", &full_text, full_tokens, &probes, active_task),
1957        score_text_baseline(
1958            "head_tail",
1959            &head_tail_text,
1960            head_tail_tokens,
1961            &probes,
1962            active_task,
1963        ),
1964        score_context_governor(
1965            &response,
1966            &compacted_text,
1967            compacted_tokens,
1968            &probes,
1969            active_task,
1970        ),
1971    ];
1972
1973    Ok(ReplayFixtureReport {
1974        fixture_id,
1975        original_message_count,
1976        compacted_message_count: response.compacted_messages.len(),
1977        probes,
1978        baselines,
1979        receipt_id: response.receipt.receipt_id,
1980        warnings: response.receipt.warnings,
1981    })
1982}
1983
1984fn literal_probe_fragment(text: &str, max_chars: usize) -> String {
1985    text.split_whitespace()
1986        .collect::<Vec<_>>()
1987        .join(" ")
1988        .chars()
1989        .take(max_chars)
1990        .collect()
1991}
1992
1993fn push_probe(probes: &mut Vec<ReplayProbe>, category: &str, needle: String, max_probes: usize) {
1994    let mut needle = compact_preview(needle.trim(), 180);
1995    if let Some(stripped) = needle.strip_suffix("...") {
1996        needle = stripped.to_string();
1997    }
1998    if needle.len() < 4 || probes.len() >= max_probes {
1999        return;
2000    }
2001    if probes.iter().any(|probe| probe.needle == needle) {
2002        return;
2003    }
2004    probes.push(ReplayProbe {
2005        category: category.to_string(),
2006        needle,
2007    });
2008}
2009
2010fn join_message_content(messages: &[Message]) -> String {
2011    messages
2012        .iter()
2013        .map(|message| message.content.as_str())
2014        .collect::<Vec<_>>()
2015        .join("\n")
2016}
2017
2018fn head_tail_messages(messages: &[Message], head: usize, tail: usize) -> Vec<Message> {
2019    if messages.len() <= head + tail {
2020        return messages.to_vec();
2021    }
2022    messages
2023        .iter()
2024        .take(head)
2025        .chain(messages.iter().skip(messages.len().saturating_sub(tail)))
2026        .cloned()
2027        .collect()
2028}
2029
2030fn text_contains_probe(text: &str, needle: &str) -> bool {
2031    text.contains(needle) || compact_preview(text, usize::MAX / 4).contains(needle)
2032}
2033
2034fn score_text_baseline(
2035    name: &str,
2036    text: &str,
2037    tokens: usize,
2038    probes: &[ReplayProbe],
2039    active_task: &str,
2040) -> ReplayBaselineScore {
2041    let visible = probes
2042        .iter()
2043        .filter(|probe| text_contains_probe(text, &probe.needle))
2044        .count();
2045    build_score(
2046        name,
2047        visible,
2048        visible,
2049        probes.len(),
2050        tokens,
2051        text_contains_probe(text, active_task),
2052    )
2053}
2054
2055fn score_answerability_text_baseline(
2056    name: &str,
2057    text: &str,
2058    tokens: usize,
2059    questions: &[ReplayAnswerabilityQuestion],
2060    active_task: &str,
2061) -> ReplayAnswerabilityBaselineScore {
2062    let answerable = questions
2063        .iter()
2064        .filter(|question| answerability_terms_present(text, question))
2065        .count();
2066    let incorrect_action_risk = questions
2067        .iter()
2068        .filter(|question| forbidden_terms_present(text, question))
2069        .count();
2070    build_answerability_score(
2071        name,
2072        answerable,
2073        questions.len(),
2074        incorrect_action_risk,
2075        tokens,
2076        text_contains_probe(text, active_task),
2077    )
2078}
2079
2080fn score_answerability_context_governor(
2081    response: &CompactResponse,
2082    compacted_text: &str,
2083    tokens: usize,
2084    questions: &[ReplayAnswerabilityQuestion],
2085    active_task: &str,
2086) -> ReplayAnswerabilityBaselineScore {
2087    let answerable = questions
2088        .iter()
2089        .filter(|question| {
2090            question.expected_terms.iter().all(|term| {
2091                text_contains_probe(compacted_text, term)
2092                    || !context_search(response, term, 1, SearchScope::All).is_empty()
2093            })
2094        })
2095        .count();
2096    let incorrect_action_risk = questions
2097        .iter()
2098        .filter(|question| forbidden_terms_present(compacted_text, question))
2099        .count();
2100    build_answerability_score(
2101        "context_governor",
2102        answerable,
2103        questions.len(),
2104        incorrect_action_risk,
2105        tokens,
2106        text_contains_probe(compacted_text, active_task),
2107    )
2108}
2109
2110fn answerability_terms_present(text: &str, question: &ReplayAnswerabilityQuestion) -> bool {
2111    question
2112        .expected_terms
2113        .iter()
2114        .all(|term| text_contains_probe(text, term))
2115}
2116
2117fn forbidden_terms_present(text: &str, question: &ReplayAnswerabilityQuestion) -> bool {
2118    let text_l = text.to_lowercase();
2119    question.forbidden_terms.iter().any(|term| {
2120        let term_l = term.to_lowercase();
2121        let negated = text_l.contains(&format!("not {term_l}"))
2122            || text_l.contains(&format!("no {term_l}"))
2123            || text_l.contains(&format!("avoid {term_l}"));
2124        !negated && text_contains_probe(text, term)
2125    })
2126}
2127
2128fn build_answerability_score(
2129    name: &str,
2130    answerable_questions: usize,
2131    total_questions: usize,
2132    incorrect_action_risk: usize,
2133    tokens: usize,
2134    active_task_visible: bool,
2135) -> ReplayAnswerabilityBaselineScore {
2136    let denominator = total_questions.max(1) as f64;
2137    ReplayAnswerabilityBaselineScore {
2138        name: name.to_string(),
2139        answerable_questions,
2140        total_questions,
2141        answerability_rate: answerable_questions as f64 / denominator,
2142        incorrect_action_risk,
2143        tokens,
2144        active_task_visible,
2145    }
2146}
2147
2148fn score_context_governor(
2149    response: &CompactResponse,
2150    compacted_text: &str,
2151    tokens: usize,
2152    probes: &[ReplayProbe],
2153    active_task: &str,
2154) -> ReplayBaselineScore {
2155    let visible = probes
2156        .iter()
2157        .filter(|probe| text_contains_probe(compacted_text, &probe.needle))
2158        .count();
2159    let recoverable = probes
2160        .iter()
2161        .filter(|probe| !context_search(response, &probe.needle, 1, SearchScope::All).is_empty())
2162        .count();
2163    build_score(
2164        "context_governor",
2165        visible,
2166        recoverable,
2167        probes.len(),
2168        tokens,
2169        text_contains_probe(compacted_text, active_task),
2170    )
2171}
2172
2173fn build_score(
2174    name: &str,
2175    visible_probes: usize,
2176    recoverable_probes: usize,
2177    total_probes: usize,
2178    tokens: usize,
2179    active_task_visible: bool,
2180) -> ReplayBaselineScore {
2181    let denominator = total_probes.max(1) as f64;
2182    ReplayBaselineScore {
2183        name: name.to_string(),
2184        visible_probes,
2185        recoverable_probes,
2186        total_probes,
2187        visible_rate: visible_probes as f64 / denominator,
2188        recoverable_rate: recoverable_probes as f64 / denominator,
2189        tokens,
2190        active_task_visible,
2191    }
2192}
2193
2194#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
2195pub struct MemoryArchiveRecordV1 {
2196    pub schema: String,
2197    pub receipt_id: String,
2198    pub item_id: String,
2199    pub content_blake3: String,
2200    pub content_kind: ContentKind,
2201    pub archive_reason: String,
2202    pub sensitivity: String,
2203    pub preview: String,
2204}
2205
2206#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
2207pub struct MemoryArchiveOutcome {
2208    pub record_ids: Vec<String>,
2209    pub records_attempted: usize,
2210}
2211
2212pub trait MemorySink {
2213    fn archive_context_record(
2214        &mut self,
2215        record: &MemoryArchiveRecordV1,
2216    ) -> Result<Option<String>, ContextGovernorError>;
2217}
2218
2219pub fn archive_response_to_memory(
2220    response: &CompactResponse,
2221    sink: &mut dyn MemorySink,
2222) -> Result<MemoryArchiveOutcome, ContextGovernorError> {
2223    let mut outcome = MemoryArchiveOutcome::default();
2224    for item in &response.allocation_plan.items {
2225        let should_archive = response
2226            .allocation_plan
2227            .archived_item_ids
2228            .contains(&item.item_id)
2229            || matches!(
2230                item.authority_class,
2231                AuthorityClass::DurableMemoryCandidate | AuthorityClass::EvidenceCritical
2232            );
2233        if !should_archive {
2234            continue;
2235        }
2236        let Some(stored) = response
2237            .exact_store
2238            .iter()
2239            .find(|stored| stored.item_id == item.item_id)
2240        else {
2241            continue;
2242        };
2243        let record = MemoryArchiveRecordV1 {
2244            schema: "MemoryArchiveRecordV1".to_string(),
2245            receipt_id: response.receipt.receipt_id.clone(),
2246            item_id: item.item_id.clone(),
2247            content_blake3: stored.content_blake3.clone(),
2248            content_kind: item.content_kind.clone(),
2249            archive_reason: format!("semantic archive {:?}", item.item_type),
2250            sensitivity: "internal".to_string(),
2251            preview: content_aware_preview(&item.content_kind, &stored.content, 280),
2252        };
2253        outcome.records_attempted += 1;
2254        if let Some(id) = sink.archive_context_record(&record)? {
2255            outcome.record_ids.push(id);
2256        }
2257    }
2258    Ok(outcome)
2259}
2260
2261#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
2262pub struct StoredContextSearchHit {
2263    pub receipt_id: String,
2264    pub hit: ContextSearchHit,
2265}
2266
2267#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
2268pub struct FileContextStoreStatusV1 {
2269    pub schema: String,
2270    pub root: String,
2271    pub receipt_count: usize,
2272    pub total_bytes: u64,
2273    pub stale_tmp_files_removed: usize,
2274    pub index_built: bool,
2275    pub searchable: bool,
2276    pub last_receipt: Option<String>,
2277}
2278
2279#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
2280pub struct FileContextStorePruneResultV1 {
2281    pub schema: String,
2282    pub kept_receipts: usize,
2283    pub removed_receipts: usize,
2284    pub total_bytes: u64,
2285    pub index_built: bool,
2286    pub last_receipt: Option<String>,
2287}
2288
2289#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
2290struct PersistedReceiptIndexV1 {
2291    schema: String,
2292    receipt_ids: Vec<String>,
2293    token_to_receipts: BTreeMap<String, Vec<String>>,
2294}
2295
2296#[derive(Debug, Clone)]
2297struct ReceiptFileInfo {
2298    receipt_id: String,
2299    path: std::path::PathBuf,
2300    created_utc: Option<DateTime<Utc>>,
2301}
2302
2303#[derive(Debug, Clone)]
2304pub struct FileContextStore {
2305    root: std::path::PathBuf,
2306    // In-memory inverted index: token -> set of receipt_ids that contain it.
2307    // Built lazily on first search() call and invalidated on save().
2308    index: std::cell::RefCell<
2309        Option<std::collections::HashMap<String, std::collections::HashSet<String>>>,
2310    >,
2311}
2312
2313impl FileContextStore {
2314    pub fn new(root: impl AsRef<std::path::Path>) -> Self {
2315        Self {
2316            root: root.as_ref().to_path_buf(),
2317            index: std::cell::RefCell::new(None),
2318        }
2319    }
2320
2321    pub fn save(
2322        &self,
2323        response: &CompactResponse,
2324    ) -> Result<std::path::PathBuf, ContextGovernorError> {
2325        std::fs::create_dir_all(&self.root)?;
2326        let path = self.path_for_receipt(&response.receipt.receipt_id);
2327        let tmp_path = path.with_extension("json.tmp");
2328        let json = serde_json::to_string_pretty(response)?;
2329        std::fs::write(&tmp_path, json)?;
2330        std::fs::rename(&tmp_path, &path)?;
2331        // Invalidate the in-memory index so the next search rebuilds it
2332        *self.index.borrow_mut() = None;
2333        self.rebuild_and_persist_index()?;
2334        Ok(path)
2335    }
2336
2337    pub fn load(&self, receipt_id: &str) -> Result<CompactResponse, ContextGovernorError> {
2338        let path = self.path_for_receipt(receipt_id);
2339        if !path.exists() {
2340            return Err(ContextGovernorError::ReceiptNotFound(
2341                receipt_id.to_string(),
2342            ));
2343        }
2344        let json = std::fs::read_to_string(path)?;
2345        Ok(serde_json::from_str(&json)?)
2346    }
2347
2348    pub fn list_receipts(&self) -> Result<Vec<String>, ContextGovernorError> {
2349        Ok(self
2350            .scan_receipts()?
2351            .0
2352            .into_iter()
2353            .map(|info| info.receipt_id)
2354            .collect())
2355    }
2356
2357    pub fn status(&self) -> Result<FileContextStoreStatusV1, ContextGovernorError> {
2358        let (receipts, total_bytes, stale_tmp_files_removed) = self.scan_receipts()?;
2359        let index_built =
2360            self.index.borrow().is_some() || self.persisted_index_matches(&receipts)?;
2361        Ok(FileContextStoreStatusV1 {
2362            schema: "FileContextStoreStatusV1".to_string(),
2363            root: self.root.display().to_string(),
2364            receipt_count: receipts.len(),
2365            total_bytes,
2366            stale_tmp_files_removed,
2367            index_built,
2368            searchable: index_built || receipts.is_empty(),
2369            last_receipt: receipts.last().map(|info| info.receipt_id.clone()),
2370        })
2371    }
2372
2373    pub fn prune_receipts_keep_last(
2374        &self,
2375        keep_last: usize,
2376    ) -> Result<FileContextStorePruneResultV1, ContextGovernorError> {
2377        let (receipts, _, _) = self.scan_receipts()?;
2378        let remove_count = receipts.len().saturating_sub(keep_last);
2379        for info in receipts.iter().take(remove_count) {
2380            std::fs::remove_file(&info.path)?;
2381        }
2382        *self.index.borrow_mut() = None;
2383        self.rebuild_and_persist_index()?;
2384        let status = self.status()?;
2385        Ok(FileContextStorePruneResultV1 {
2386            schema: "FileContextStorePruneResultV1".to_string(),
2387            kept_receipts: status.receipt_count,
2388            removed_receipts: remove_count,
2389            total_bytes: status.total_bytes,
2390            index_built: status.index_built,
2391            last_receipt: status.last_receipt,
2392        })
2393    }
2394
2395    fn scan_receipts(&self) -> Result<(Vec<ReceiptFileInfo>, u64, usize), ContextGovernorError> {
2396        if !self.root.exists() {
2397            return Ok((Vec::new(), 0, 0));
2398        }
2399        let mut ids = Vec::new();
2400        let mut total_bytes = 0u64;
2401        let mut stale_tmp_files_removed = 0usize;
2402        for entry in std::fs::read_dir(&self.root)? {
2403            let entry = entry?;
2404            let path = entry.path();
2405            // Clean up stale .tmp files from interrupted save() calls
2406            if path.extension().and_then(|ext| ext.to_str()) == Some("tmp") {
2407                if std::fs::remove_file(&path).is_ok() {
2408                    stale_tmp_files_removed += 1;
2409                }
2410                continue;
2411            }
2412            if path.extension().and_then(|ext| ext.to_str()) != Some("json") {
2413                continue;
2414            }
2415            if path.file_name().and_then(|name| name.to_str()) == Some(".receipt-index.json") {
2416                continue;
2417            }
2418            let Some(stem) = path.file_stem().and_then(|stem| stem.to_str()) else {
2419                continue;
2420            };
2421            if let Ok(metadata) = entry.metadata() {
2422                total_bytes = total_bytes.saturating_add(metadata.len());
2423                let created_utc = std::fs::read_to_string(&path)
2424                    .ok()
2425                    .and_then(|json| serde_json::from_str::<CompactResponse>(&json).ok())
2426                    .map(|response| response.receipt.created_utc);
2427                ids.push(ReceiptFileInfo {
2428                    receipt_id: stem.to_string(),
2429                    path,
2430                    created_utc,
2431                });
2432            }
2433        }
2434        ids.sort_by(|a, b| {
2435            a.created_utc
2436                .cmp(&b.created_utc)
2437                .then_with(|| a.receipt_id.cmp(&b.receipt_id))
2438        });
2439        Ok((ids, total_bytes, stale_tmp_files_removed))
2440    }
2441
2442    pub fn expand(
2443        &self,
2444        receipt_id: &str,
2445        item_id: &str,
2446        max_chars: usize,
2447    ) -> Result<ContextExpandResult, ContextGovernorError> {
2448        let response = self.load(receipt_id)?;
2449        context_expand(&response, item_id, max_chars)
2450            .ok_or_else(|| ContextGovernorError::ReceiptNotFound(item_id.to_string()))
2451    }
2452
2453    /// Build an in-memory inverted index from all stored receipts.
2454    /// Tokenizes content on whitespace and stores receipt_id sets per token.
2455    fn build_index(
2456        &self,
2457    ) -> Result<
2458        std::collections::HashMap<String, std::collections::HashSet<String>>,
2459        ContextGovernorError,
2460    > {
2461        let mut idx: std::collections::HashMap<String, std::collections::HashSet<String>> =
2462            std::collections::HashMap::new();
2463        for receipt_id in self.list_receipts()? {
2464            let response = self.load(&receipt_id)?;
2465            // Index exact_store content
2466            for item in &response.exact_store {
2467                for token in item.content.split_whitespace() {
2468                    let token_l = token.to_lowercase();
2469                    idx.entry(token_l).or_default().insert(receipt_id.clone());
2470                }
2471            }
2472            // Index compacted_messages content
2473            for msg in &response.compacted_messages {
2474                for token in msg.content.split_whitespace() {
2475                    let token_l = token.to_lowercase();
2476                    idx.entry(token_l).or_default().insert(receipt_id.clone());
2477                }
2478            }
2479        }
2480        Ok(idx)
2481    }
2482
2483    fn rebuild_and_persist_index(&self) -> Result<(), ContextGovernorError> {
2484        let index = self.build_index()?;
2485        self.persist_index(&index)?;
2486        *self.index.borrow_mut() = Some(index);
2487        Ok(())
2488    }
2489
2490    fn persisted_index_matches(
2491        &self,
2492        receipts: &[ReceiptFileInfo],
2493    ) -> Result<bool, ContextGovernorError> {
2494        let Some(index) = self.load_persisted_index()? else {
2495            return Ok(false);
2496        };
2497        let receipt_ids = receipts
2498            .iter()
2499            .map(|info| info.receipt_id.clone())
2500            .collect::<Vec<_>>();
2501        Ok(index.receipt_ids == receipt_ids)
2502    }
2503
2504    fn load_persisted_index(
2505        &self,
2506    ) -> Result<Option<PersistedReceiptIndexV1>, ContextGovernorError> {
2507        let path = self.index_path();
2508        if !path.exists() {
2509            return Ok(None);
2510        }
2511        let json = std::fs::read_to_string(path)?;
2512        Ok(Some(serde_json::from_str(&json)?))
2513    }
2514
2515    fn persist_index(
2516        &self,
2517        index: &std::collections::HashMap<String, std::collections::HashSet<String>>,
2518    ) -> Result<(), ContextGovernorError> {
2519        std::fs::create_dir_all(&self.root)?;
2520        let receipt_ids = self.list_receipts()?;
2521        let token_to_receipts = index
2522            .iter()
2523            .map(|(token, receipt_set)| {
2524                let mut receipts = receipt_set.iter().cloned().collect::<Vec<_>>();
2525                receipts.sort();
2526                (token.clone(), receipts)
2527            })
2528            .collect::<BTreeMap<_, _>>();
2529        let persisted = PersistedReceiptIndexV1 {
2530            schema: "PersistedReceiptIndexV1".to_string(),
2531            receipt_ids,
2532            token_to_receipts,
2533        };
2534        let path = self.index_path();
2535        let tmp_path = path.with_extension("json.tmp");
2536        std::fs::write(&tmp_path, serde_json::to_string_pretty(&persisted)?)?;
2537        std::fs::rename(tmp_path, path)?;
2538        Ok(())
2539    }
2540
2541    fn index_from_persisted(
2542        persisted: PersistedReceiptIndexV1,
2543    ) -> std::collections::HashMap<String, std::collections::HashSet<String>> {
2544        persisted
2545            .token_to_receipts
2546            .into_iter()
2547            .map(|(token, receipts)| (token, receipts.into_iter().collect()))
2548            .collect()
2549    }
2550
2551    /// Ensure the index is built and return a reference to it.
2552    fn with_index<F, R>(&self, f: F) -> Result<R, ContextGovernorError>
2553    where
2554        F: FnOnce(&std::collections::HashMap<String, std::collections::HashSet<String>>) -> R,
2555    {
2556        let mut guard = self.index.borrow_mut();
2557        if guard.is_none() {
2558            let receipts = self.scan_receipts()?.0;
2559            if let Some(persisted) = self.load_persisted_index()? {
2560                let receipt_ids = receipts
2561                    .iter()
2562                    .map(|info| info.receipt_id.clone())
2563                    .collect::<Vec<_>>();
2564                if persisted.receipt_ids == receipt_ids {
2565                    *guard = Some(Self::index_from_persisted(persisted));
2566                }
2567            }
2568            if guard.is_none() {
2569                let index = self.build_index()?;
2570                self.persist_index(&index)?;
2571                *guard = Some(index);
2572            }
2573        }
2574        Ok(f(guard.as_ref().unwrap()))
2575    }
2576
2577    pub fn search(
2578        &self,
2579        query: &str,
2580        top_k: usize,
2581        scope: SearchScope,
2582    ) -> Result<Vec<StoredContextSearchHit>, ContextGovernorError> {
2583        // Use the inverted index to find candidate receipts that contain
2584        // any query token, then do full search only on those candidates.
2585        // If the index returns no candidates (e.g. the query is a substring
2586        // of a token, not a full token), fall back to searching all receipts.
2587        let query_tokens: Vec<String> =
2588            query.split_whitespace().map(|t| t.to_lowercase()).collect();
2589        let candidate_ids: Vec<String> = if query_tokens.is_empty() {
2590            self.list_receipts()?
2591        } else {
2592            let candidates = self.with_index(|idx| {
2593                let mut candidates: std::collections::HashSet<String> =
2594                    std::collections::HashSet::new();
2595                for token in &query_tokens {
2596                    if let Some(receipt_ids) = idx.get(token) {
2597                        for rid in receipt_ids {
2598                            candidates.insert(rid.clone());
2599                        }
2600                    }
2601                }
2602                let mut sorted: Vec<String> = candidates.into_iter().collect();
2603                sorted.sort();
2604                sorted
2605            })?;
2606            if candidates.is_empty() {
2607                // No exact token matches — fall back to full scan so that
2608                // substring queries (e.g. "NEEDLE" in "NEEDLE_PARSER") still work.
2609                self.list_receipts()?
2610            } else {
2611                candidates
2612            }
2613        };
2614
2615        let mut out = Vec::new();
2616        for receipt_id in candidate_ids {
2617            let response = self.load(&receipt_id)?;
2618            for hit in context_search(&response, query, top_k, scope) {
2619                out.push(StoredContextSearchHit {
2620                    receipt_id: receipt_id.clone(),
2621                    hit,
2622                });
2623                if out.len() >= top_k {
2624                    return Ok(out);
2625                }
2626            }
2627        }
2628        Ok(out)
2629    }
2630
2631    fn path_for_receipt(&self, receipt_id: &str) -> std::path::PathBuf {
2632        let safe = receipt_id
2633            .chars()
2634            .map(|ch| {
2635                if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' {
2636                    ch
2637                } else {
2638                    '_'
2639                }
2640            })
2641            .collect::<String>();
2642        self.root.join(format!("{safe}.json"))
2643    }
2644
2645    fn index_path(&self) -> std::path::PathBuf {
2646        self.root.join(".receipt-index.json")
2647    }
2648}
context_governor/lib.rs

context_governor/
lib.rs
