Skip to main content

kaish_kernel/
output_limit.rs

1//! Configurable output size limits for agent safety.
2//!
3//! When output exceeds the threshold the result is capped and `ExecResult.out`
4//! is replaced with a head+tail preview. Two strategies, selected at runtime by
5//! [`SpillMode`]:
6//! - [`SpillMode::Disk`] (default): the full output is written to a spill file
7//!   on the real filesystem and the preview points at it. The agent can then
8//!   selectively read the file.
9//! - [`SpillMode::Memory`]: the output is truncated in memory only — no disk
10//!   I/O, no recoverable file. For runtime read-only kernels (e.g. kaibo) that
11//!   must not touch the host filesystem even when `localfs` is compiled in.
12//!   Memory stays bounded regardless of how much the command produces.
13//!
14//! Either way the exit code is remapped to 3 (`did_spill`) so callers can tell
15//! the output was capped.
16//!
17//! Per-mode defaults: sandboxed-agent kernels get an 8KB limit, REPL/test
18//! kernels are unlimited. Runtime-switchable via the `kaish-output-limit` builtin.
19
20use std::path::PathBuf;
21
22use crate::interpreter::ExecResult;
23#[cfg(feature = "localfs")]
24use crate::paths;
25
26/// Default output limit for the sandboxed-agent preset (8KB).
27const DEFAULT_AGENT_LIMIT: usize = 8 * 1024;
28
29/// Default head preview size (bytes of output start to keep).
30const DEFAULT_HEAD_BYTES: usize = 1024;
31
32/// Default tail preview size (bytes of output end to keep).
33const DEFAULT_TAIL_BYTES: usize = 512;
34
35/// Where overflow output goes when it exceeds the limit.
36///
37/// This is a *runtime* choice, distinct from the compile-time `localfs`
38/// feature: a `localfs`-built kernel can still be told to truncate in memory.
39/// A build without `localfs` always behaves as [`SpillMode::Memory`] regardless
40/// of this setting, since disk I/O is unavailable.
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
42pub enum SpillMode {
43    /// Write overflow to a disk spill file under `paths::spill_dir()` and keep a
44    /// head+tail preview in the result (the message carries the file path).
45    /// Requires the `localfs` feature. This is the default.
46    ///
47    /// Auto-overridden to [`Memory`](Self::Memory) at kernel construction when
48    /// the VFS mount is `NoLocal` (memory-only) — such a kernel has no host
49    /// filesystem to spill to. See `Kernel::assemble`.
50    #[default]
51    Disk,
52    /// Truncate in memory to head+tail only — no disk I/O, no recoverable file.
53    /// For runtime read-only kernels (e.g. kaibo) that must never touch the host
54    /// filesystem even when `localfs` is compiled in.
55    Memory,
56}
57
58/// Configurable output size limit.
59///
60/// Threaded through `KernelConfig` → `ExecContext` → kernel pipeline execution.
61/// Runtime-mutable via the `kaish-output-limit` builtin.
62#[derive(Debug, Clone)]
63pub struct OutputLimitConfig {
64    max_bytes: Option<usize>,
65    head_bytes: usize,
66    tail_bytes: usize,
67    spill_mode: SpillMode,
68}
69
70impl OutputLimitConfig {
71    /// No limiting — REPL/embedded/test default.
72    pub fn none() -> Self {
73        Self {
74            max_bytes: None,
75            head_bytes: DEFAULT_HEAD_BYTES,
76            tail_bytes: DEFAULT_TAIL_BYTES,
77            spill_mode: SpillMode::Disk,
78        }
79    }
80
81    /// Default limit used by `on` subcommand and `set -o output-limit`.
82    pub fn default_limit() -> usize {
83        DEFAULT_AGENT_LIMIT
84    }
85
86    /// Sandboxed-agent defaults: 8KB limit, 1KB head, 512B tail, disk spill.
87    pub fn agent() -> Self {
88        Self {
89            max_bytes: Some(DEFAULT_AGENT_LIMIT),
90            head_bytes: DEFAULT_HEAD_BYTES,
91            tail_bytes: DEFAULT_TAIL_BYTES,
92            spill_mode: SpillMode::Disk,
93        }
94    }
95
96    /// Switch to in-memory truncation — no disk spill, no host filesystem
97    /// writes. For runtime read-only kernels (e.g. kaibo). Builder form of
98    /// [`set_spill_mode`](Self::set_spill_mode).
99    ///
100    /// Note: a `NoLocal` VFS mount forces this mode automatically at kernel
101    /// construction, so an embedder only needs this for a `localfs`-mounted
102    /// kernel it nonetheless wants to keep off the host disk.
103    pub fn in_memory(mut self) -> Self {
104        self.spill_mode = SpillMode::Memory;
105        self
106    }
107
108    /// Whether output limiting is enabled.
109    pub fn is_enabled(&self) -> bool {
110        self.max_bytes.is_some()
111    }
112
113    /// The spill mode (disk vs in-memory truncation).
114    pub fn spill_mode(&self) -> SpillMode {
115        self.spill_mode
116    }
117
118    /// Set the spill mode.
119    pub fn set_spill_mode(&mut self, mode: SpillMode) {
120        self.spill_mode = mode;
121    }
122
123    /// The maximum output size in bytes, if set.
124    pub fn max_bytes(&self) -> Option<usize> {
125        self.max_bytes
126    }
127
128    /// Bytes of output head to preserve in truncated result.
129    pub fn head_bytes(&self) -> usize {
130        self.head_bytes
131    }
132
133    /// Bytes of output tail to preserve in truncated result.
134    pub fn tail_bytes(&self) -> usize {
135        self.tail_bytes
136    }
137
138    /// Set the output limit. `None` disables limiting.
139    pub fn set_limit(&mut self, max: Option<usize>) {
140        self.max_bytes = max;
141    }
142
143    /// Set the head preview size.
144    pub fn set_head_bytes(&mut self, bytes: usize) {
145        self.head_bytes = bytes;
146    }
147
148    /// Set the tail preview size.
149    pub fn set_tail_bytes(&mut self, bytes: usize) {
150        self.tail_bytes = bytes;
151    }
152}
153
154/// Result of a spill operation.
155pub struct SpillResult {
156    pub path: PathBuf,
157    pub total_bytes: usize,
158}
159
160/// Check if the result output exceeds the limit and spill to disk if so.
161///
162/// Mutates `result.out` in place: replaces with head+tail+pointer message.
163/// Returns `Some(SpillResult)` if a spill file was written, `None` otherwise.
164///
165/// If the filesystem write fails, the result is replaced with an error.
166/// Fail fast: truncating output silently could corrupt structured data
167/// that an agent acts on. An explicit error is safer.
168///
169/// In [`SpillMode::Memory`], or in any build without the `localfs` feature,
170/// performs in-memory head+tail truncation (no disk I/O) instead.
171pub async fn spill_if_needed(
172    result: &mut ExecResult,
173    config: &OutputLimitConfig,
174) -> Option<SpillResult> {
175    let max = config.max_bytes?;
176
177    // Binary payloads are measured and spilled by RAW bytes. text_out() would
178    // lossy-decode a Bytes result — corrupting the spill file and mis-measuring
179    // the size (U+FFFD is 3 bytes per invalid byte). Handle them up front.
180    if let Some(total) = result.out_bytes().map(|b| b.len()) {
181        if total <= max {
182            return None;
183        }
184        #[cfg(feature = "localfs")]
185        if config.spill_mode == SpillMode::Disk {
186            let bytes = result.out_bytes().unwrap_or_default().to_vec();
187            return match write_spill_file(&bytes).await {
188                Ok((path, written)) => {
189                    result.set_out(format!(
190                        "[binary output: {total} bytes spilled to {} — read it with `cat {}`]",
191                        path.display(),
192                        path.display()
193                    ));
194                    result.did_spill = true;
195                    Some(SpillResult { path, total_bytes: written })
196                }
197                Err(e) => {
198                    tracing::error!("binary output spill failed: {}", e);
199                    *result = ExecResult::failure(
200                        1,
201                        format!(
202                            "binary output exceeded {max} byte limit ({total} bytes) and spill \
203                             to disk failed: {e}"
204                        ),
205                    );
206                    None
207                }
208            };
209        }
210        // Memory mode (or no localfs): bounded head+tail of the raw bytes. The
211        // result stays binary, just truncated.
212        let bytes = result.out_bytes().unwrap_or_default().to_vec();
213        let head_n = config.head_bytes.min(bytes.len());
214        let tail_n = config.tail_bytes.min(bytes.len().saturating_sub(head_n));
215        let mut truncated = bytes[..head_n].to_vec();
216        truncated.extend_from_slice(&bytes[bytes.len() - tail_n..]);
217        result.set_out_bytes(truncated);
218        result.did_spill = true;
219        return None;
220    }
221
222    // Disk spill requires `localfs` AND the caller selecting it. Memory mode
223    // (or a build without `localfs`) falls through to in-memory truncation.
224    #[cfg(feature = "localfs")]
225    if config.spill_mode == SpillMode::Disk {
226        // If result.out is already populated (external commands), check it directly
227        if !result.text_out().is_empty() && !result.has_output() {
228            let total = result.text_out().len();
229            if total <= max {
230                return None;
231            }
232            return spill_string(result, config, max).await;
233        }
234
235        // If we have structured OutputData, estimate size before materializing
236        if let Some(output) = result.output() {
237            let estimate = output.estimated_byte_size();
238            if estimate <= max {
239                // Small enough — materialize normally
240                result.materialize();
241                // Re-check actual size (estimate is a lower bound)
242                if result.text_out().len() <= max {
243                    return None;
244                }
245                return spill_string(result, config, max).await;
246            }
247
248            // Large — stream directly to spill file, never holding full String
249            return spill_output_data(result, config, max).await;
250        }
251
252        return None;
253    }
254
255    // In-memory head+tail truncation (Memory mode or no `localfs`): no disk I/O.
256    truncate_in_memory(result, config, max)
257}
258
259/// Truncate output in memory to head+tail, with no disk I/O.
260///
261/// Sets `did_spill = true` so the kernel remaps the exit code to 3 — the same
262/// "output was capped" signal as a disk spill — but the message carries no file
263/// path because there is no recoverable file. Returns `None` (no `SpillResult`,
264/// since nothing was written); the caller distinguishes truncation via
265/// `result.did_spill`.
266///
267/// Memory is bounded: large structured `OutputData` is streamed through a byte
268/// budget rather than materialized into a full `String`, so a builtin emitting
269/// a huge tree (e.g. a recursive `ls` of a giant directory) cannot OOM a
270/// read-only kernel.
271fn truncate_in_memory(
272    result: &mut ExecResult,
273    config: &OutputLimitConfig,
274    max: usize,
275) -> Option<SpillResult> {
276    // Structured OutputData: estimate first. If it would clearly overflow,
277    // render only a bounded head prefix via `write_canonical` rather than
278    // materializing the whole thing.
279    if let Some(output) = result.output() {
280        let estimate = output.estimated_byte_size();
281        if estimate > max {
282            // Render a bounded head prefix only — no full materialization.
283            let mut buf = Vec::with_capacity(config.head_bytes + 64);
284            // write_canonical stops shortly after the budget; ignore the count.
285            let _ = output.write_canonical(&mut buf, Some(config.head_bytes));
286            let s = String::from_utf8_lossy(&buf);
287            let head = truncate_to_char_boundary(&s, config.head_bytes);
288            let truncated = format!(
289                "{}\n...\n[output truncated in memory: ~{} bytes (exceeds {} byte limit) — head only, no spill file]",
290                head, estimate, max
291            );
292            result.set_out(truncated);
293            result.did_spill = true;
294            return None;
295        }
296        // Small enough to materialize safely.
297        result.materialize();
298    }
299
300    let total = result.text_out().len();
301    if total <= max {
302        return None;
303    }
304
305    // Already-materialized text fits in memory (it was produced into RAM
306    // regardless) — give a precise head+tail+total.
307    let text = result.text_out().into_owned();
308    let head = truncate_to_char_boundary(&text, config.head_bytes);
309    let tail = tail_from_str(&text, config.tail_bytes);
310    let truncated = format!(
311        "{}\n...\n{}\n[output truncated in memory: {} bytes total — no spill file]",
312        head, tail, total
313    );
314    result.set_out(truncated);
315    result.did_spill = true;
316    None
317}
318
319/// Spill an already-materialized string in result.out.
320#[cfg(feature = "localfs")]
321async fn spill_string(
322    result: &mut ExecResult,
323    config: &OutputLimitConfig,
324    max: usize,
325) -> Option<SpillResult> {
326    let total = result.text_out().len();
327    match write_spill_file(result.text_out().as_bytes()).await {
328        Ok((path, written)) => {
329            let truncated = build_truncated_output(&result.text_out(), config, &path, total);
330            result.set_out(truncated);
331            result.did_spill = true;
332            Some(SpillResult {
333                path,
334                total_bytes: written,
335            })
336        }
337        Err(e) => {
338            tracing::error!("output spill failed: {}", e);
339            *result = ExecResult::failure(1, format!(
340                "output exceeded {} byte limit ({} bytes) and spill to disk failed: {}",
341                max, total, e
342            ));
343            None
344        }
345    }
346}
347
348/// Stream OutputData directly to a spill file without materializing the full String.
349#[cfg(feature = "localfs")]
350async fn spill_output_data(
351    result: &mut ExecResult,
352    config: &OutputLimitConfig,
353    max: usize,
354) -> Option<SpillResult> {
355    let output = result.output()?;
356
357    let dir = paths::spill_dir();
358    if let Err(e) = tokio::fs::create_dir_all(&dir).await {
359        tracing::error!("output spill dir creation failed: {}", e);
360        *result = ExecResult::failure(1, format!(
361            "output exceeded {} byte limit and spill dir creation failed: {}", max, e
362        ));
363        return None;
364    }
365
366    let filename = generate_spill_filename();
367    let path = dir.join(&filename);
368
369    // Write OutputData directly to file via write_canonical
370    let total = match std::fs::File::create(&path) {
371        Ok(mut file) => {
372            match output.write_canonical(&mut file, None) {
373                Ok(n) => n,
374                Err(e) => {
375                    tracing::error!("output spill write failed: {}", e);
376                    *result = ExecResult::failure(1, format!(
377                        "output exceeded {} byte limit and spill to disk failed: {}", max, e
378                    ));
379                    return None;
380                }
381            }
382        }
383        Err(e) => {
384            tracing::error!("output spill file creation failed: {}", e);
385            *result = ExecResult::failure(1, format!(
386                "output exceeded {} byte limit and spill to disk failed: {}", max, e
387            ));
388            return None;
389        }
390    };
391
392    // Read head and tail from the spill file for the truncated preview
393    let head = read_head_from_file(&path, config.head_bytes).await.unwrap_or_default();
394    let tail = read_tail_from_file(&path, config.tail_bytes).await.unwrap_or_default();
395    let path_str = path.to_string_lossy();
396
397    result.set_out(format!(
398        "{}\n...\n{}\n[output truncated: {} bytes total — full output at {}]",
399        head, tail, total, path_str
400    ));
401    result.did_spill = true;
402
403    Some(SpillResult {
404        path,
405        total_bytes: total,
406    })
407}
408
409/// Write output bytes to a new spill file. Returns (path, bytes_written).
410#[cfg(feature = "localfs")]
411async fn write_spill_file(data: &[u8]) -> Result<(PathBuf, usize), std::io::Error> {
412    let dir = paths::spill_dir();
413    tokio::fs::create_dir_all(&dir).await?;
414
415    let filename = generate_spill_filename();
416    let path = dir.join(filename);
417    tokio::fs::write(&path, data).await?;
418    Ok((path, data.len()))
419}
420
421/// Build the truncated output string with head, tail, and pointer.
422#[cfg(feature = "localfs")]
423fn build_truncated_output(
424    full: &str,
425    config: &OutputLimitConfig,
426    spill_path: &std::path::Path,
427    total_bytes: usize,
428) -> String {
429    let head = truncate_to_char_boundary(full, config.head_bytes);
430    let tail = tail_from_str(full, config.tail_bytes);
431    let path_str = spill_path.to_string_lossy();
432    format!(
433        "{}\n...\n{}\n[output truncated: {} bytes total — full output at {}]",
434        head, tail, total_bytes, path_str
435    )
436}
437
438/// Truncate a string to at most `max_bytes`, respecting UTF-8 char boundaries.
439fn truncate_to_char_boundary(s: &str, max_bytes: usize) -> &str {
440    if s.len() <= max_bytes {
441        return s;
442    }
443    // Find the last char boundary at or before max_bytes
444    let mut end = max_bytes;
445    while end > 0 && !s.is_char_boundary(end) {
446        end -= 1;
447    }
448    &s[..end]
449}
450
451/// Get the last `max_bytes` of a string, respecting UTF-8 char boundaries.
452fn tail_from_str(s: &str, max_bytes: usize) -> &str {
453    if s.len() <= max_bytes {
454        return s;
455    }
456    let start = s.len() - max_bytes;
457    let mut adjusted = start;
458    while adjusted < s.len() && !s.is_char_boundary(adjusted) {
459        adjusted += 1;
460    }
461    &s[adjusted..]
462}
463
464/// Read the first N bytes from a file for head preview.
465#[cfg(feature = "localfs")]
466async fn read_head_from_file(path: &std::path::Path, max_bytes: usize) -> Result<String, std::io::Error> {
467    use tokio::io::AsyncReadExt;
468
469    let mut file = tokio::fs::File::open(path).await?;
470    let mut buf = vec![0u8; max_bytes];
471    let n = file.read(&mut buf).await?;
472    buf.truncate(n);
473
474    let s = String::from_utf8_lossy(&buf);
475    // Truncate to char boundary
476    let result = truncate_to_char_boundary(&s, max_bytes);
477    Ok(result.to_string())
478}
479
480/// Read the last N bytes from a file for tail preview.
481#[cfg(feature = "localfs")]
482async fn read_tail_from_file(path: &std::path::Path, max_bytes: usize) -> Result<String, std::io::Error> {
483    use tokio::io::{AsyncReadExt, AsyncSeekExt};
484
485    let mut file = tokio::fs::File::open(path).await?;
486    let metadata = file.metadata().await?;
487    let len = metadata.len() as usize;
488
489    if len <= max_bytes {
490        let mut buf = Vec::new();
491        file.read_to_end(&mut buf).await?;
492        return Ok(String::from_utf8_lossy(&buf).into_owned());
493    }
494
495    let offset = len - max_bytes;
496    file.seek(std::io::SeekFrom::Start(offset as u64)).await?;
497    let mut buf = vec![0u8; max_bytes];
498    let n = file.read(&mut buf).await?;
499    buf.truncate(n);
500
501    // Adjust to char boundary
502    let s = String::from_utf8_lossy(&buf);
503    Ok(s.into_owned())
504}
505
506/// Generate a unique spill filename using timestamp, PID, and monotonic counter.
507#[cfg(feature = "localfs")]
508fn generate_spill_filename() -> String {
509    use std::sync::atomic::{AtomicUsize, Ordering};
510    use std::time::SystemTime;
511
512    static COUNTER: AtomicUsize = AtomicUsize::new(0);
513    let seq = COUNTER.fetch_add(1, Ordering::Relaxed);
514    let ts = SystemTime::now()
515        .duration_since(SystemTime::UNIX_EPOCH)
516        .unwrap_or_default();
517    let pid = std::process::id();
518    format!("spill-{}.{}-{}-{}.txt", ts.as_secs(), ts.subsec_nanos(), pid, seq)
519}
520
521/// Parse a size string with optional K/M suffix into bytes.
522///
523/// Accepts: "64K", "64k", "1M", "1m", "65536" (raw bytes).
524pub fn parse_size(s: &str) -> Result<usize, String> {
525    let s = s.trim();
526    if s.is_empty() {
527        return Err("empty size string".to_string());
528    }
529
530    let (num_str, multiplier) = if let Some(n) = s.strip_suffix('K').or_else(|| s.strip_suffix('k')) {
531        (n, 1024)
532    } else if let Some(n) = s.strip_suffix('M').or_else(|| s.strip_suffix('m')) {
533        (n, 1024 * 1024)
534    } else {
535        (s, 1)
536    };
537
538    let num: usize = num_str
539        .parse()
540        .map_err(|_| format!("invalid size: {}", s))?;
541
542    Ok(num * multiplier)
543}
544
545#[cfg(all(test, feature = "localfs"))]
546mod tests {
547    use super::*;
548
549    #[test]
550    fn test_none_is_disabled() {
551        let config = OutputLimitConfig::none();
552        assert!(!config.is_enabled());
553        assert_eq!(config.max_bytes(), None);
554    }
555
556    #[test]
557    fn test_agent_preset_is_enabled() {
558        let config = OutputLimitConfig::agent();
559        assert!(config.is_enabled());
560        assert_eq!(config.max_bytes(), Some(8 * 1024));
561        assert_eq!(config.head_bytes(), 1024);
562        assert_eq!(config.tail_bytes(), 512);
563    }
564
565    #[test]
566    fn test_set_limit() {
567        let mut config = OutputLimitConfig::none();
568        assert!(!config.is_enabled());
569
570        config.set_limit(Some(1024));
571        assert!(config.is_enabled());
572        assert_eq!(config.max_bytes(), Some(1024));
573
574        config.set_limit(None);
575        assert!(!config.is_enabled());
576    }
577
578    #[test]
579    fn test_set_head_tail() {
580        let mut config = OutputLimitConfig::agent();
581        config.set_head_bytes(2048);
582        config.set_tail_bytes(1024);
583        assert_eq!(config.head_bytes(), 2048);
584        assert_eq!(config.tail_bytes(), 1024);
585    }
586
587    #[test]
588    fn test_parse_size() {
589        assert_eq!(parse_size("64K").unwrap(), 64 * 1024);
590        assert_eq!(parse_size("64k").unwrap(), 64 * 1024);
591        assert_eq!(parse_size("1M").unwrap(), 1024 * 1024);
592        assert_eq!(parse_size("1m").unwrap(), 1024 * 1024);
593        assert_eq!(parse_size("65536").unwrap(), 65536);
594        assert!(parse_size("").is_err());
595        assert!(parse_size("abc").is_err());
596    }
597
598    #[test]
599    fn test_truncate_to_char_boundary() {
600        assert_eq!(truncate_to_char_boundary("hello", 10), "hello");
601        assert_eq!(truncate_to_char_boundary("hello", 3), "hel");
602        // Multi-byte: "日" is 3 bytes
603        assert_eq!(truncate_to_char_boundary("日本語", 3), "日");
604        assert_eq!(truncate_to_char_boundary("日本語", 4), "日");
605        assert_eq!(truncate_to_char_boundary("日本語", 6), "日本");
606    }
607
608    #[test]
609    fn test_tail_from_str() {
610        assert_eq!(tail_from_str("hello", 10), "hello");
611        assert_eq!(tail_from_str("hello", 3), "llo");
612        // Multi-byte
613        assert_eq!(tail_from_str("日本語", 3), "語");
614        assert_eq!(tail_from_str("日本語", 6), "本語");
615    }
616
617    #[test]
618    fn test_generate_spill_filename() {
619        let name = generate_spill_filename();
620        assert!(name.starts_with("spill-"));
621        assert!(name.ends_with(".txt"));
622    }
623
624    #[tokio::test]
625    async fn test_spill_if_needed_under_limit() {
626        let config = OutputLimitConfig::agent();
627        let mut result = ExecResult::success("short output");
628        let spill = spill_if_needed(&mut result, &config).await;
629        assert!(spill.is_none());
630        assert_eq!(&*result.text_out(), "short output");
631        assert!(!result.did_spill);
632    }
633
634    #[tokio::test]
635    async fn test_spill_if_needed_over_limit() {
636        let config = OutputLimitConfig {
637            max_bytes: Some(100),
638            head_bytes: 20,
639            tail_bytes: 10,
640            spill_mode: SpillMode::Disk,
641        };
642        let big_output = "x".repeat(200);
643        let mut result = ExecResult::success(big_output);
644        let spill = spill_if_needed(&mut result, &config).await;
645        assert!(spill.is_some());
646        assert!(result.did_spill);
647
648        let spill = spill.unwrap();
649        assert_eq!(spill.total_bytes, 200);
650        assert!(spill.path.exists());
651
652        // Verify truncated output
653        assert!(result.text_out().contains("..."));
654        assert!(result.text_out().contains("[output truncated: 200 bytes total"));
655        assert!(result.text_out().contains(&spill.path.to_string_lossy().to_string()));
656
657        // Verify head (first 20 bytes)
658        assert!(result.text_out().starts_with(&"x".repeat(20)));
659
660        // Verify spill file has full content
661        let spill_content = tokio::fs::read_to_string(&spill.path).await.unwrap();
662        assert_eq!(spill_content.len(), 200);
663
664        // Clean up
665        let _ = tokio::fs::remove_file(&spill.path).await;
666    }
667
668    #[tokio::test]
669    async fn test_spill_if_needed_disabled() {
670        let config = OutputLimitConfig::none();
671        let big_output = "x".repeat(200);
672        let mut result = ExecResult::success(big_output.clone());
673        let spill = spill_if_needed(&mut result, &config).await;
674        assert!(spill.is_none());
675        assert_eq!(&*result.text_out(), big_output);
676        assert!(!result.did_spill);
677    }
678
679    #[test]
680    fn test_build_truncated_output() {
681        let config = OutputLimitConfig {
682            max_bytes: Some(100),
683            head_bytes: 5,
684            tail_bytes: 3,
685            spill_mode: SpillMode::Disk,
686        };
687        let full = "abcdefghijklmnop";
688        let path = PathBuf::from("/tmp/test-spill.txt");
689        let result = build_truncated_output(full, &config, &path, 16);
690        assert!(result.starts_with("abcde"));
691        assert!(result.contains("..."));
692        assert!(result.contains("nop"));
693        assert!(result.contains("[output truncated: 16 bytes total — full output at /tmp/test-spill.txt]"));
694    }
695
696    #[tokio::test]
697    async fn test_kernel_agent_truncates_large_output() {
698        use crate::kernel::{Kernel, KernelConfig};
699
700        // agent preset has 8K limit by default — use a smaller limit for testing
701        let config = KernelConfig::agent()
702            .with_output_limit(OutputLimitConfig {
703                max_bytes: Some(200),
704                head_bytes: 50,
705                tail_bytes: 30,
706                spill_mode: SpillMode::Disk,
707            });
708        let kernel = Kernel::new(config).expect("kernel creation");
709
710        // seq 1 10000 produces lots of output
711        let result = kernel.execute("seq 1 10000").await.expect("execute");
712        assert!(result.text_out().contains("[output truncated:"));
713        assert!(result.text_out().contains("full output at"));
714        // Head should contain the first numbers
715        assert!(result.text_out().starts_with("1\n"));
716    }
717
718    /// GH #177: every other Disk-mode test in this module re-specifies
719    /// `spill_mode: SpillMode::Disk` explicitly, so none of them actually pin
720    /// that a host-backed (`Sandboxed`) kernel's *untouched* default — no
721    /// `.with_output_limit()` override at all — is `Disk`. This is the literal
722    /// "`Kernel::new` without `.in_memory()`" scenario the issue names. The
723    /// forcing logic in `Kernel::assemble` (`no_host_side_channel`) must leave
724    /// a `Sandboxed` kernel's config alone; only `NoLocal`/`with_backend`
725    /// override it to `Memory`.
726    #[tokio::test]
727    async fn test_agent_kernel_unmodified_default_spills_to_disk() {
728        use crate::kernel::{Kernel, KernelConfig};
729
730        // Untouched: OutputLimitConfig::agent() — 8K limit, SpillMode::Disk.
731        let config = KernelConfig::agent();
732        assert_eq!(config.output_limit.spill_mode(), SpillMode::Disk);
733        let kernel = Kernel::new(config).expect("kernel creation");
734
735        let big = "x".repeat(8 * 1024 + 200);
736        let result = kernel.execute(&format!("echo '{}'", big)).await.expect("execute");
737        assert_eq!(result.code, 3, "default 8K agent limit should trip the spill");
738        assert!(
739            result.text_out().contains("full output at"),
740            "an unmodified agent() default must spill to a real file, not truncate in \
741             memory: {}",
742            result.text_out()
743        );
744    }
745
746    #[tokio::test]
747    async fn test_spill_exits_3() {
748        use crate::kernel::{Kernel, KernelConfig};
749
750        let config = KernelConfig::agent()
751            .with_output_limit(OutputLimitConfig {
752                max_bytes: Some(100),
753                head_bytes: 30,
754                tail_bytes: 20,
755                spill_mode: SpillMode::Disk,
756            });
757        let kernel = Kernel::new(config).expect("kernel creation");
758
759        let big = "x".repeat(200);
760        let result = kernel.execute(&format!("echo '{}'", big)).await.expect("execute");
761        assert_eq!(result.code, 3, "spill should always exit 3");
762        assert_eq!(result.original_code, Some(0), "original command exit code preserved");
763        assert!(result.text_out().contains("[output truncated:"));
764    }
765
766    #[tokio::test]
767    async fn test_kernel_repl_no_truncation() {
768        use crate::kernel::{Kernel, KernelConfig};
769
770        // REPL has no limit
771        let config = KernelConfig::repl();
772        let kernel = Kernel::new(config).expect("kernel creation");
773
774        let result = kernel.execute("seq 1 100").await.expect("execute");
775        assert!(!result.text_out().contains("[output truncated:"));
776        assert!(result.text_out().contains("100"));
777    }
778
779    #[tokio::test]
780    async fn test_kernel_builtin_truncation() {
781        use crate::kernel::{Kernel, KernelConfig};
782
783        // Builtins go through post-hoc spill check
784        let config = KernelConfig::agent()
785            .with_output_limit(OutputLimitConfig {
786                max_bytes: Some(100),
787                head_bytes: 30,
788                tail_bytes: 20,
789                spill_mode: SpillMode::Disk,
790            });
791        let kernel = Kernel::new(config).expect("kernel creation");
792
793        // echo with a large string
794        let big = "x".repeat(200);
795        let result = kernel.execute(&format!("echo '{}'", big)).await.expect("execute");
796        assert!(result.text_out().contains("[output truncated:"));
797    }
798
799    // ── OutputData estimation and streaming tests ──
800
801    #[test]
802    fn test_estimated_byte_size_text() {
803        use crate::interpreter::OutputData;
804        let data = OutputData::text("hello world");
805        assert_eq!(data.estimated_byte_size(), 11);
806    }
807
808    #[test]
809    fn test_estimated_byte_size_table() {
810        use crate::interpreter::{OutputData, OutputNode};
811        let data = OutputData::table(
812            vec!["NAME".into(), "SIZE".into()],
813            vec![
814                OutputNode::new("foo").with_cells(vec!["123".into()]),
815                OutputNode::new("bar").with_cells(vec!["456".into()]),
816            ],
817        );
818        // "foo\t123\nbar\t456" = 3+1+3 + 1 + 3+1+3 = 15
819        assert_eq!(data.estimated_byte_size(), 15);
820    }
821
822    #[test]
823    fn test_estimated_byte_size_tree() {
824        use crate::interpreter::{OutputData, OutputNode};
825        let data = OutputData::nodes(vec![
826            OutputNode::new("src").with_children(vec![
827                OutputNode::new("main.rs"),
828                OutputNode::new("lib.rs"),
829            ]),
830        ]);
831        // "src/{main.rs,lib.rs}" = 3 + 2 + 7 + 1 + 6 + 1 = 20
832        assert_eq!(data.estimated_byte_size(), 20);
833    }
834
835    #[test]
836    fn test_write_canonical_matches_to_canonical_string() {
837        use crate::interpreter::{OutputData, OutputNode};
838
839        let cases: Vec<OutputData> = vec![
840            OutputData::text("hello world"),
841            OutputData::nodes(vec![
842                OutputNode::new("file1"),
843                OutputNode::new("file2"),
844            ]),
845            OutputData::table(
846                vec!["NAME".into(), "SIZE".into()],
847                vec![
848                    OutputNode::new("foo").with_cells(vec!["123".into()]),
849                    OutputNode::new("bar").with_cells(vec!["456".into()]),
850                ],
851            ),
852            OutputData::nodes(vec![
853                OutputNode::new("src").with_children(vec![
854                    OutputNode::new("main.rs"),
855                    OutputNode::new("lib.rs"),
856                ]),
857            ]),
858        ];
859
860        for data in cases {
861            let expected = data.to_canonical_string();
862            let mut buf = Vec::new();
863            let written = data.write_canonical(&mut buf, None).unwrap();
864            let got = String::from_utf8(buf).unwrap();
865            assert_eq!(got, expected, "write_canonical mismatch for {:?}", data);
866            assert_eq!(written, expected.len(), "byte count mismatch");
867        }
868    }
869
870    #[test]
871    fn test_write_canonical_budget_stops_early() {
872        use crate::interpreter::{OutputData, OutputNode};
873
874        let data = OutputData::nodes(
875            (0..1000).map(|i| OutputNode::new(format!("file_{:04}", i))).collect()
876        );
877        let mut buf = Vec::new();
878        let written = data.write_canonical(&mut buf, Some(100)).unwrap();
879        // Should have stopped shortly after 100 bytes
880        assert!(written > 100, "should exceed budget slightly");
881        assert!(written < 500, "should stop soon after budget: got {}", written);
882    }
883
884    #[tokio::test]
885    async fn test_spill_if_needed_large_output_data_no_oom() {
886        use crate::interpreter::{OutputData, OutputNode};
887
888        let config = OutputLimitConfig {
889            max_bytes: Some(1024),
890            head_bytes: 100,
891            tail_bytes: 50,
892            spill_mode: SpillMode::Disk,
893        };
894
895        // 100K nodes — large enough to detect OOM if materialized carelessly,
896        // but small enough to not slow down the test
897        let nodes: Vec<OutputNode> = (0..100_000)
898            .map(|i| OutputNode::new(format!("node_{:06}", i)))
899            .collect();
900        let data = OutputData::nodes(nodes);
901        let mut result = ExecResult::with_output(data);
902
903        let spill = spill_if_needed(&mut result, &config).await;
904        assert!(spill.is_some(), "should have spilled");
905        assert!(result.did_spill);
906        assert!(result.text_out().contains("[output truncated:"));
907
908        // Clean up
909        if let Some(s) = spill {
910            let _ = tokio::fs::remove_file(&s.path).await;
911        }
912    }
913
914    // ── In-memory spill mode (SpillMode::Memory) ──
915
916    #[test]
917    fn test_in_memory_builder_and_default() {
918        assert_eq!(OutputLimitConfig::agent().spill_mode(), SpillMode::Disk);
919        assert_eq!(OutputLimitConfig::agent().in_memory().spill_mode(), SpillMode::Memory);
920
921        let mut config = OutputLimitConfig::none();
922        config.set_spill_mode(SpillMode::Memory);
923        assert_eq!(config.spill_mode(), SpillMode::Memory);
924    }
925
926    #[tokio::test]
927    async fn test_memory_mode_truncates_string_without_disk() {
928        let config = OutputLimitConfig {
929            max_bytes: Some(100),
930            head_bytes: 20,
931            tail_bytes: 10,
932            spill_mode: SpillMode::Memory,
933        };
934        let mut result = ExecResult::success("x".repeat(200));
935        let spill = spill_if_needed(&mut result, &config).await;
936
937        // No SpillResult (no file written) but did_spill flags the truncation.
938        assert!(spill.is_none(), "memory mode must not write a spill file");
939        assert!(result.did_spill, "memory truncation must set did_spill for the exit-3 remap");
940
941        let out = result.text_out();
942        assert!(out.contains("truncated in memory"), "got: {}", out);
943        assert!(out.contains("200 bytes total"), "got: {}", out);
944        assert!(!out.contains("full output at"), "memory mode must not point at a file: {}", out);
945        assert!(out.starts_with(&"x".repeat(20)), "head preserved");
946    }
947
948    #[tokio::test]
949    async fn test_memory_mode_under_limit_untouched() {
950        let config = OutputLimitConfig {
951            max_bytes: Some(100),
952            head_bytes: 20,
953            tail_bytes: 10,
954            spill_mode: SpillMode::Memory,
955        };
956        let mut result = ExecResult::success("short");
957        let spill = spill_if_needed(&mut result, &config).await;
958        assert!(spill.is_none());
959        assert!(!result.did_spill);
960        assert_eq!(&*result.text_out(), "short");
961    }
962
963    #[tokio::test]
964    async fn test_memory_mode_large_output_data_bounded() {
965        use crate::interpreter::{OutputData, OutputNode};
966
967        let config = OutputLimitConfig {
968            max_bytes: Some(1024),
969            head_bytes: 100,
970            tail_bytes: 50,
971            spill_mode: SpillMode::Memory,
972        };
973
974        // 100K nodes — would be a huge String if fully materialized.
975        let nodes: Vec<OutputNode> = (0..100_000)
976            .map(|i| OutputNode::new(format!("node_{:06}", i)))
977            .collect();
978        let mut result = ExecResult::with_output(OutputData::nodes(nodes));
979
980        let spill = spill_if_needed(&mut result, &config).await;
981        assert!(spill.is_none(), "memory mode writes no file");
982        assert!(result.did_spill);
983        let out = result.text_out();
984        assert!(out.contains("truncated in memory"), "got: {}", out);
985        assert!(out.starts_with("node_000000"), "head rendered: {}", out);
986        // Head-only path for oversized structured data: no tail section echoed.
987        assert!(out.contains("head only"), "got: {}", out);
988    }
989
990    #[tokio::test]
991    async fn test_kernel_memory_mode_exits_3_preserves_original() {
992        use crate::kernel::{Kernel, KernelConfig};
993
994        let config = KernelConfig::agent().with_output_limit(OutputLimitConfig {
995            max_bytes: Some(100),
996            head_bytes: 30,
997            tail_bytes: 20,
998            spill_mode: SpillMode::Memory,
999        });
1000        let kernel = Kernel::new(config).expect("kernel creation");
1001
1002        let big = "x".repeat(200);
1003        let result = kernel.execute(&format!("echo '{}'", big)).await.expect("execute");
1004        assert_eq!(result.code, 3, "memory truncation still signals via exit 3");
1005        assert_eq!(result.original_code, Some(0), "original exit code preserved");
1006        assert!(result.text_out().contains("truncated in memory"));
1007        assert!(!result.text_out().contains("full output at"));
1008    }
1009
1010    #[tokio::test]
1011    async fn test_nolocal_kernel_forces_memory_spill() {
1012        use crate::kernel::{Kernel, KernelConfig, VfsMountMode};
1013
1014        // NoLocal mount + an explicit Disk spill mode: the kernel must override
1015        // to Memory so nothing is written to a host spill file, even though
1016        // `localfs` is compiled in.
1017        let config = KernelConfig::agent()
1018            .with_vfs_mode(VfsMountMode::NoLocal)
1019            .with_output_limit(OutputLimitConfig {
1020                max_bytes: Some(100),
1021                head_bytes: 30,
1022                tail_bytes: 20,
1023                spill_mode: SpillMode::Disk,
1024            });
1025        let kernel = Kernel::new(config).expect("kernel creation");
1026
1027        let big = "x".repeat(200);
1028        let result = kernel.execute(&format!("echo '{}'", big)).await.expect("execute");
1029        assert_eq!(result.code, 3, "still signals truncation via exit 3");
1030        assert!(result.text_out().contains("truncated in memory"), "got: {}", result.text_out());
1031        assert!(
1032            !result.text_out().contains("full output at"),
1033            "NoLocal kernel must not write a host spill file: {}",
1034            result.text_out()
1035        );
1036    }
1037
1038}