lean_ctx/core/
firewall.rs

1//! Context firewall: replace large tool outputs with a compact digest + retrieval ref.
2//!
3//! When ephemeral mode is active (`[archive].ephemeral`, default on), genuinely large
4//! tool results are stored out-of-band via [`crate::core::archive`] and only a
5//! deterministic digest — a head/tail excerpt, size stats, and `ctx_expand` drilldown
6//! instructions — is returned inline. This keeps the agent's context window small while
7//! preserving full, slice-addressable access to the raw output.
8//!
9//! Scope: tool *outputs* (`ctx_shell`, `ctx_execute`, `ctx_search`, `ctx_tree`). Explicit
10//! file reads keep their own read-mode system and are never firewalled.
11
12use crate::core::config::Config;
13
14const HEAD_LINES: usize = 20;
15const TAIL_LINES: usize = 8;
16const LONG_LINE_HEAD_CHARS: usize = 800;
17const LONG_LINE_TAIL_CHARS: usize = 300;
18
19/// Tools whose large outputs are eligible for the firewall. Explicit file reads are
20/// intentionally excluded — they have their own read-mode (`lines:`, `signatures`, …).
21pub fn is_firewallable_tool(name: &str) -> bool {
22    matches!(
23        name,
24        "ctx_shell" | "ctx_execute" | "ctx_search" | "ctx_tree"
25    )
26}
27
28/// Explicit file-read tools whose result *is* the file content the agent reads and
29/// edits against. They must always return that content inline — never a head/tail
30/// digest (firewall) nor a stored-reference stub (`reference_results`) — regardless
31/// of output size or config. This is the single source of truth for "an explicit
32/// read always returns content"; both the firewall and the reference-results path
33/// honour it so a `ctx_read` can never degrade to a preview the agent can't edit.
34pub fn is_protected_read(name: &str) -> bool {
35    matches!(name, "ctx_read" | "ctx_multi_read" | "ctx_smart_read")
36}
37
38/// Effective minimum token count before firewalling (config + env override).
39pub fn min_tokens(config: &Config) -> usize {
40    config.archive.ephemeral_min_tokens_effective()
41}
42
43/// Whether a result of `output_tokens` from `tool` should be firewalled.
44pub fn should_firewall(tool: &str, output_tokens: usize, config: &Config) -> bool {
45    config.archive.ephemeral_effective()
46        && is_firewallable_tool(tool)
47        && output_tokens >= min_tokens(config)
48}
49
50/// Build the inline digest that replaces a firewalled output. Deterministic (no LLM):
51/// a head/tail excerpt for multi-line output, or a char-bounded excerpt for output with
52/// few but very long lines (e.g. a single giant JSON line), followed by drilldown
53/// instructions keyed on `archive_id`.
54pub fn summarize(full: &str, archive_id: &str, tool: &str, output_tokens: usize) -> String {
55    let chars = full.len();
56    let lines: Vec<&str> = full.lines().collect();
57    let line_count = lines.len();
58
59    let mut out = String::new();
60    out.push_str(&format!(
61        "[Firewalled {tool} output — {chars} chars, {output_tokens} tok, {line_count} lines stored out-of-band]\n"
62    ));
63
64    if line_count > HEAD_LINES + TAIL_LINES + 1 {
65        out.push_str("--- head ---\n");
66        out.push_str(&lines[..HEAD_LINES].join("\n"));
67        out.push_str(&format!(
68            "\n--- … {} lines omitted … ---\n",
69            line_count - HEAD_LINES - TAIL_LINES
70        ));
71        out.push_str("--- tail ---\n");
72        out.push_str(&lines[line_count - TAIL_LINES..].join("\n"));
73        out.push('\n');
74    } else {
75        // Few lines but large (e.g. one giant minified JSON line): char-bounded excerpt.
76        let head_end = full.floor_char_boundary(LONG_LINE_HEAD_CHARS.min(chars));
77        out.push_str(&full[..head_end]);
78        if chars > LONG_LINE_HEAD_CHARS + LONG_LINE_TAIL_CHARS {
79            out.push_str("\n… (truncated) …\n");
80            let tail_start = full.floor_char_boundary(chars - LONG_LINE_TAIL_CHARS);
81            out.push_str(&full[tail_start..]);
82            out.push('\n');
83        }
84    }
85
86    out.push_str("--- retrieve full output ---\n");
87    // Non-MCP route first: the verbatim blob is a real file any tool can read,
88    // for agents/orgs where MCP is unavailable or forbidden.
89    out.push_str(&format!(
90        "Direct:  read {} directly (no MCP)\n",
91        crate::core::archive::content_path_str(archive_id)
92    ));
93    out.push_str(&format!("Full:    ctx_expand(id=\"{archive_id}\")\n"));
94    out.push_str(&format!(
95        "Range:   ctx_expand(id=\"{archive_id}\", start_line=1, end_line=80)\n"
96    ));
97    out.push_str(&format!(
98        "Head:    ctx_expand(id=\"{archive_id}\", head=120)\n"
99    ));
100    out.push_str(&format!(
101        "Search:  ctx_expand(id=\"{archive_id}\", search=\"ERROR\")\n"
102    ));
103    out.push_str(&format!(
104        "JSON:    ctx_expand(id=\"{archive_id}\", json_keys=true)"
105    ));
106    out
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn firewallable_tools_are_outputs_not_reads() {
115        assert!(is_firewallable_tool("ctx_shell"));
116        assert!(is_firewallable_tool("ctx_search"));
117        assert!(is_firewallable_tool("ctx_tree"));
118        assert!(is_firewallable_tool("ctx_execute"));
119        assert!(!is_firewallable_tool("ctx_read"));
120        assert!(!is_firewallable_tool("ctx_multi_read"));
121        assert!(!is_firewallable_tool("ctx_knowledge"));
122    }
123
124    #[test]
125    fn protected_reads_are_file_readers_and_never_firewallable() {
126        // Explicit reads must always return content (no firewall digest, no
127        // reference stub) so the agent can edit against the lines.
128        for read in ["ctx_read", "ctx_multi_read", "ctx_smart_read"] {
129            assert!(is_protected_read(read), "{read} must be a protected read");
130            assert!(
131                !is_firewallable_tool(read),
132                "{read} must never be firewallable"
133            );
134        }
135        assert!(!is_protected_read("ctx_shell"));
136        assert!(!is_protected_read("ctx_search"));
137    }
138
139    #[test]
140    fn should_firewall_respects_tool_and_threshold() {
141        let mut cfg = Config::default();
142        cfg.archive.enabled = true;
143        cfg.archive.ephemeral = true;
144        cfg.archive.ephemeral_min_tokens = 2000;
145        // Env can override ephemeral; clear it for a deterministic test.
146        crate::test_env::remove_var("LEAN_CTX_EPHEMERAL");
147        crate::test_env::remove_var("LEAN_CTX_EPHEMERAL_MIN_TOKENS");
148
149        assert!(should_firewall("ctx_shell", 5000, &cfg));
150        assert!(!should_firewall("ctx_shell", 1000, &cfg)); // below threshold
151        assert!(!should_firewall("ctx_read", 5000, &cfg)); // not firewallable
152    }
153
154    #[test]
155    fn summarize_includes_excerpt_stats_and_ref() {
156        let full = (1..=200)
157            .map(|i| format!("line {i}"))
158            .collect::<Vec<_>>()
159            .join("\n");
160        let digest = summarize(&full, "abc123", "ctx_shell", 1234);
161        assert!(digest.contains("Firewalled ctx_shell output"));
162        assert!(digest.contains("1234 tok"));
163        assert!(digest.contains("line 1")); // head
164        assert!(digest.contains("line 200")); // tail
165        assert!(digest.contains("lines omitted"));
166        assert!(digest.contains("ctx_expand(id=\"abc123\")"));
167        assert!(digest.contains("json_keys=true"));
168        // The digest must be far smaller than the original.
169        assert!(digest.len() < full.len());
170    }
171
172    #[test]
173    fn summarize_handles_single_giant_line() {
174        let full = "x".repeat(5000);
175        let digest = summarize(&full, "id9", "ctx_search", 1300);
176        assert!(digest.contains("Firewalled ctx_search output"));
177        assert!(digest.contains("truncated"));
178        assert!(digest.len() < full.len());
179    }
180}
lean_ctx/core/firewall.rs

lean_ctx/core/
firewall.rs
