diff --git a/crates/monty-python/python/pydantic_monty/_monty.pyi b/crates/monty-python/python/pydantic_monty/_monty.pyi index ae80a2bf5..14b1b8706 100644 --- a/crates/monty-python/python/pydantic_monty/_monty.pyi +++ b/crates/monty-python/python/pydantic_monty/_monty.pyi @@ -47,18 +47,27 @@ NOT_HANDLED = object() @final class CollectStreams: - """Collect printed output as `(stream, text)` tuples.""" + """Collect printed output as `(stream, text)` tuples. - def __new__(cls) -> CollectStreams: ... + Defaults to a 10 MiB cap. Pass `max_bytes=None` to disable (trusted hosts). + Exceeding the cap raises `MemoryError`. Not covered by `ResourceLimits.max_memory`. + The cap includes a fixed per-entry overhead (many tiny fragments). + """ + + def __new__(cls, max_bytes: int | None = 10 * 1024 * 1024) -> CollectStreams: ... @property def output(self) -> list[tuple[Literal['stdout', 'stderr'], str]]: """Collected output so far.""" @final class CollectString: - """Collect printed output as one concatenated string.""" + """Collect printed output as one concatenated string. + + Defaults to a 10 MiB cap. Pass `max_bytes=None` to disable (trusted hosts). + Exceeding the cap raises `MemoryError`. Not covered by `ResourceLimits.max_memory`. + """ - def __new__(cls) -> CollectString: ... + def __new__(cls, max_bytes: int | None = 10 * 1024 * 1024) -> CollectString: ... @property def output(self) -> str: """Collected output so far.""" diff --git a/crates/monty-python/src/print_target.rs b/crates/monty-python/src/print_target.rs index cda0f21d5..480146e8b 100644 --- a/crates/monty-python/src/print_target.rs +++ b/crates/monty-python/src/print_target.rs @@ -10,13 +10,18 @@ //! - A `CollectString()` instance — fragments accumulate into a shared flat //! `String` exposed via `CollectString.output`. //! +//! Both collectors default to [`DEFAULT_MAX_PRINT_COLLECT_BYTES`]; pass +//! `max_bytes=None` to disable. The cap is enforced here in +//! [`PrintTarget::write_event`] (pool host path) — not by worker +//! `ResourceLimits.max_memory`. +//! //! This module encapsulates that dispatch. The rest of the bindings thread a //! [`PrintTarget`] value through `feed_run`, while the collector objects //! themselves remain the single public place that exposes the captured output. use std::sync::{Arc, Mutex, PoisonError}; -use monty::{MontyException, PrintStream}; +use monty::{DEFAULT_MAX_PRINT_COLLECT_BYTES, MontyException, PrintStream, check_print_collect_limit}; use monty_proto::python::exc_py_to_monty; use pyo3::{ PyRef, @@ -26,18 +31,41 @@ use pyo3::{ types::{PyList, PyString}, }; +/// Host bytes charged per retained `(stream, text)` entry beyond the payload. +/// +/// `String` / `Vec` bookkeeping is not free: many tiny prints can exhaust the +/// host long before payload bytes hit the cap. Charged toward `max_bytes`. +const COLLECT_STREAMS_ENTRY_OVERHEAD: usize = 64; + +/// Shared collect-streams state: labelled fragments plus optional byte cap. +#[derive(Debug)] +pub(crate) struct CollectStreamsState { + output: Vec<(PrintStream, String)>, + /// Running charge: payload UTF-8 bytes plus [`COLLECT_STREAMS_ENTRY_OVERHEAD`] + /// per retained entry. Avoids O(n) rescans on each print event. + collected_bytes: usize, + max_bytes: Option, +} + /// Shared buffer for the `CollectStreams` mode. /// /// The `Arc>` wrapper lets a single collector keep accumulating /// across `start()` / `resume()` / async / snapshot-load boundaries while still /// allowing read access from Python between transitions. -type CollectStreamsBuffer = Arc>>; +type CollectStreamsBuffer = Arc>; + +/// Shared collect-string state: flat text plus optional byte cap. +#[derive(Debug)] +pub(crate) struct CollectStringState { + output: String, + max_bytes: Option, +} /// Shared buffer for the `CollectString` mode. /// /// This follows the same sharing scheme as [`CollectStreamsBuffer`], but stores /// a flat concatenated string instead of labelled stream fragments. -type CollectStringBuffer = Arc>; +type CollectStringBuffer = Arc>; /// Python collector that records printed fragments as `(stream, text)` tuples. /// @@ -45,18 +73,30 @@ type CollectStringBuffer = Arc>; /// entire run or snapshot chain. Reading `.output` clones the current buffer /// without draining it, so callers can inspect intermediate state and continue /// accumulating into the same collector. +/// +/// Defaults to a [`DEFAULT_MAX_PRINT_COLLECT_BYTES`] cap; `max_bytes=None` +/// disables the limit (trusted hosts only). The cap includes a fixed per-entry +/// overhead so many tiny fragments cannot exhaust the host before payload bytes +/// hit the limit. #[pyclass(name = "CollectStreams", module = "pydantic_monty", frozen)] -#[derive(Debug, Default)] +#[derive(Debug)] pub struct PyCollectStreams { buffer: CollectStreamsBuffer, } #[pymethods] impl PyCollectStreams { - /// Creates an empty stream collector. + /// Creates an empty stream collector with an optional byte cap. #[new] - fn new() -> Self { - Self::default() + #[pyo3(signature = (max_bytes=DEFAULT_MAX_PRINT_COLLECT_BYTES))] + fn new(max_bytes: Option) -> Self { + Self { + buffer: Arc::new(Mutex::new(CollectStreamsState { + output: Vec::new(), + collected_bytes: 0, + max_bytes, + })), + } } /// Returns the collected `(stream, text)` tuples so far. @@ -67,6 +107,7 @@ impl PyCollectStreams { self.buffer .lock() .unwrap_or_else(PoisonError::into_inner) + .output .iter() .map(|(stream, text)| { let label = match stream { @@ -95,25 +136,34 @@ impl PyCollectStreams { /// Pass `CollectString()` as `print_callback` to accumulate raw printed text /// while still letting the corresponding run or snapshot return its ordinary /// execution value. +/// +/// Defaults to a [`DEFAULT_MAX_PRINT_COLLECT_BYTES`] cap; `max_bytes=None` +/// disables the limit (trusted hosts only). #[pyclass(name = "CollectString", module = "pydantic_monty", frozen)] -#[derive(Debug, Default)] +#[derive(Debug)] pub struct PyCollectString { buffer: CollectStringBuffer, } #[pymethods] impl PyCollectString { - /// Creates an empty string collector. + /// Creates an empty string collector with an optional byte cap. #[new] - fn new() -> Self { - Self::default() + #[pyo3(signature = (max_bytes=DEFAULT_MAX_PRINT_COLLECT_BYTES))] + fn new(max_bytes: Option) -> Self { + Self { + buffer: Arc::new(Mutex::new(CollectStringState { + output: String::new(), + max_bytes, + })), + } } /// Returns the collected text so far. #[getter] fn output<'py>(&self, py: Python<'py>) -> Bound<'py, PyString> { let guard = self.buffer.lock().unwrap_or_else(PoisonError::into_inner); - PyString::new(py, guard.as_str()) + PyString::new(py, guard.output.as_str()) } fn __repr__(&self, py: Python<'_>) -> PyResult { @@ -131,9 +181,8 @@ impl PyCollectString { /// Destination for Monty `print()` output. /// /// The variant is chosen once from the Python `print_callback` argument (via -/// [`PrintTarget::from_py`]) and threaded through the execution chain. It is -/// not invoked directly — call [`PrintTarget::with_writer`] to build a -/// `PrintWriter` on demand for each VM transition. +/// [`PrintTarget::from_py`]) and threaded through the execution chain. Pool +/// sessions deliver pre-rendered fragments via [`PrintTarget::write_event`]. /// /// # Foot-guns /// @@ -237,13 +286,17 @@ impl PrintTarget { }) .map_err(|e| Python::attach(|py| exc_py_to_monty(py, &e))), Self::CollectStreams(buf) => { - buf.lock() - .unwrap_or_else(PoisonError::into_inner) - .push((stream, text.to_owned())); + let mut state = buf.lock().unwrap_or_else(PoisonError::into_inner); + let add = text.len().saturating_add(COLLECT_STREAMS_ENTRY_OVERHEAD); + check_print_collect_limit(state.collected_bytes, add, state.max_bytes)?; + state.output.push((stream, text.to_owned())); + state.collected_bytes = state.collected_bytes.saturating_add(add); Ok(()) } Self::CollectString(buf) => { - buf.lock().unwrap_or_else(PoisonError::into_inner).push_str(text); + let mut state = buf.lock().unwrap_or_else(PoisonError::into_inner); + check_print_collect_limit(state.output.len(), text.len(), state.max_bytes)?; + state.output.push_str(text); Ok(()) } } diff --git a/crates/monty-python/tests/test_print.py b/crates/monty-python/tests/test_print.py index 11ab9d5bb..a8104dd11 100644 --- a/crates/monty-python/tests/test_print.py +++ b/crates/monty-python/tests/test_print.py @@ -343,3 +343,26 @@ def test_collectors_are_valid_print_callback_values(monty_run: RunMonty) -> None assert str(exc_info.value) == snapshot( 'print_callback must be a callable, CollectStreams(), CollectString(), or None' ) + + +def test_collect_string_max_bytes_raises(monty_run: RunMonty) -> None: + """Host CollectString cap is independent of ResourceLimits.max_memory (issue #464).""" + collector = CollectString(max_bytes=100) + with pytest.raises(MontyRuntimeError) as exc_info: + monty_run("print('x' * 200)", print_callback=collector) + inner = exc_info.value.exception() + assert isinstance(inner, MemoryError) + assert str(inner) == snapshot('memory limit exceeded: 201 bytes > 100 bytes') + assert collector.output == snapshot('') + + +def test_collect_streams_max_bytes_raises(monty_run: RunMonty) -> None: + """Host CollectStreams cap charges payload + per-entry overhead (issue #464).""" + collector = CollectStreams(max_bytes=100) + with pytest.raises(MontyRuntimeError) as exc_info: + monty_run("print('x' * 200)", print_callback=collector) + inner = exc_info.value.exception() + assert isinstance(inner, MemoryError) + # 201 payload bytes + 64 entry overhead + assert str(inner) == snapshot('memory limit exceeded: 265 bytes > 100 bytes') + assert collector.output == snapshot([]) diff --git a/crates/monty/src/io.rs b/crates/monty/src/io.rs index 47e66d26a..aca396680 100644 --- a/crates/monty/src/io.rs +++ b/crates/monty/src/io.rs @@ -1,6 +1,14 @@ use std::borrow::Cow; -use crate::exception_public::MontyException; +use crate::{exception_private::ExcType, exception_public::MontyException, resource::ResourceError}; + +/// Default cap for [`PrintWriter::CollectString`] / [`PrintWriter::CollectStreams`] +/// and the matching Python collectors. +/// +/// Host-side print buffers sit outside [`crate::ResourceLimits::max_memory`]; +/// without a cap, a print loop can OOM the host while sandbox limits stay green. +/// Pass `max_bytes: None` to opt out on trusted hosts. +pub const DEFAULT_MAX_PRINT_COLLECT_BYTES: usize = 10 * 1024 * 1024; /// Identifies the output stream for a single print fragment. /// @@ -27,10 +35,13 @@ pub enum PrintStream { /// - `Stdout` — writes to standard output (the default behavior). /// - `CollectString` — accumulates output into a target `String` for programmatic access. /// No stream labels are preserved; every fragment is appended in the order it was emitted. +/// The `Option` is an optional byte cap (`None` = unlimited); constructors +/// [`collect_string`](Self::collect_string) / default Python collectors use +/// [`DEFAULT_MAX_PRINT_COLLECT_BYTES`]. /// - `CollectStreams` — accumulates output as `(stream, text)` pairs, merging consecutive /// same-stream fragments into one tuple. Each write to the same stream extends the /// trailing entry rather than producing a new one; a new tuple is only pushed when -/// the stream changes. +/// the stream changes. Same optional byte cap as `CollectString`. /// - `Callback` — delegates to a user-provided [`PrintWriterCallback`] implementation. pub enum PrintWriter<'a> { /// Silently discard all output. @@ -38,7 +49,10 @@ pub enum PrintWriter<'a> { /// Write to standard output. Stdout, /// Collect all output into a single `String`, in emit order, with no stream labels. - CollectString(&'a mut String), + /// + /// Second field: max collected bytes (`None` = unlimited). Exceeding raises + /// `MemoryError` with the same message as [`ResourceError::Memory`]. + CollectString(&'a mut String, Option), /// Collect all output as `(stream, text)` tuples. /// /// The builtin `print()` implementation calls `stdout_write` for each argument @@ -49,12 +63,24 @@ pub enum PrintWriter<'a> { /// `print()` only writes to stdout), a single `print(a, b)` call produces one /// `(Stdout, "a b\n")` entry — and consecutive prints with `end=''` likewise /// merge into a single trailing entry. - CollectStreams(&'a mut Vec<(PrintStream, String)>), + /// + /// Second field: max collected bytes across all tuples (`None` = unlimited). + CollectStreams(&'a mut Vec<(PrintStream, String)>, Option), /// Delegate to a custom callback. Callback(&'a mut dyn PrintWriterCallback), } impl PrintWriter<'_> { + /// Collect into `buf` with the default [`DEFAULT_MAX_PRINT_COLLECT_BYTES`] cap. + pub fn collect_string(buf: &mut String) -> PrintWriter<'_> { + PrintWriter::CollectString(buf, Some(DEFAULT_MAX_PRINT_COLLECT_BYTES)) + } + + /// Collect into `buf` with the default [`DEFAULT_MAX_PRINT_COLLECT_BYTES`] cap. + pub fn collect_streams(buf: &mut Vec<(PrintStream, String)>) -> PrintWriter<'_> { + PrintWriter::CollectStreams(buf, Some(DEFAULT_MAX_PRINT_COLLECT_BYTES)) + } + /// Creates a new `PrintWriter` that reborrows the same underlying target. /// /// This is useful in iterative execution (`start`/`resume` loops) where each @@ -65,8 +91,8 @@ impl PrintWriter<'_> { match self { Self::Disabled => PrintWriter::Disabled, Self::Stdout => PrintWriter::Stdout, - Self::CollectString(buf) => PrintWriter::CollectString(buf), - Self::CollectStreams(buf) => PrintWriter::CollectStreams(buf), + Self::CollectString(buf, max) => PrintWriter::CollectString(buf, *max), + Self::CollectStreams(buf, max) => PrintWriter::CollectStreams(buf, *max), Self::Callback(cb) => PrintWriter::Callback(&mut **cb), } } @@ -83,14 +109,12 @@ impl PrintWriter<'_> { print!("{output}"); Ok(()) } - Self::CollectString(buf) => { + Self::CollectString(buf, max_bytes) => { + check_print_collect_limit(buf.len(), output.len(), *max_bytes)?; buf.push_str(&output); Ok(()) } - Self::CollectStreams(buf) => { - append_streams_str(buf, PrintStream::Stdout, &output); - Ok(()) - } + Self::CollectStreams(buf, max_bytes) => append_streams_str(buf, PrintStream::Stdout, &output, *max_bytes), Self::Callback(cb) => cb.stdout_write(output), } } @@ -106,35 +130,75 @@ impl PrintWriter<'_> { print!("{end}"); Ok(()) } - Self::CollectString(buf) => { + Self::CollectString(buf, max_bytes) => { + check_print_collect_limit(buf.len(), end.len_utf8(), *max_bytes)?; buf.push(end); Ok(()) } - Self::CollectStreams(buf) => { - append_streams_char(buf, PrintStream::Stdout, end); - Ok(()) - } + Self::CollectStreams(buf, max_bytes) => append_streams_char(buf, PrintStream::Stdout, end, *max_bytes), Self::Callback(cb) => cb.stdout_push(end), } } } +/// Rejects a collect-buffer growth that would exceed `max_bytes`. +/// +/// `None` means unlimited. On overflow, returns the same `MemoryError` message +/// as [`ResourceError::Memory`] so hosts see one familiar limit string. +pub fn check_print_collect_limit( + current_len: usize, + add: usize, + max_bytes: Option, +) -> Result<(), MontyException> { + let Some(limit) = max_bytes else { + return Ok(()); + }; + let used = current_len.saturating_add(add); + if used > limit { + Err(MontyException::new( + ExcType::MemoryError, + Some(ResourceError::Memory { limit, used }.to_string()), + )) + } else { + Ok(()) + } +} + +/// Total UTF-8 bytes across all collect-streams tuples. +fn streams_byte_len(buf: &[(PrintStream, String)]) -> usize { + buf.iter().map(|(_, s)| s.len()).sum() +} + /// Appends a string fragment to the collect-streams buffer, merging into the /// trailing tuple when the stream matches. -fn append_streams_str(buf: &mut Vec<(PrintStream, String)>, stream: PrintStream, text: &str) { +fn append_streams_str( + buf: &mut Vec<(PrintStream, String)>, + stream: PrintStream, + text: &str, + max_bytes: Option, +) -> Result<(), MontyException> { + check_print_collect_limit(streams_byte_len(buf), text.len(), max_bytes)?; match buf.last_mut() { Some((s, existing)) if *s == stream => existing.push_str(text), _ => buf.push((stream, text.to_owned())), } + Ok(()) } /// Appends a single character to the collect-streams buffer, merging into the /// trailing tuple when the stream matches. -fn append_streams_char(buf: &mut Vec<(PrintStream, String)>, stream: PrintStream, ch: char) { +fn append_streams_char( + buf: &mut Vec<(PrintStream, String)>, + stream: PrintStream, + ch: char, + max_bytes: Option, +) -> Result<(), MontyException> { + check_print_collect_limit(streams_byte_len(buf), ch.len_utf8(), max_bytes)?; match buf.last_mut() { Some((s, existing)) if *s == stream => existing.push(ch), _ => buf.push((stream, String::from(ch))), } + Ok(()) } /// Trait for custom output handling from the `print()` builtin function. diff --git a/crates/monty/src/lib.rs b/crates/monty/src/lib.rs index b260955f1..2a0f3d44e 100644 --- a/crates/monty/src/lib.rs +++ b/crates/monty/src/lib.rs @@ -41,7 +41,7 @@ pub use crate::{ exception_public::{ CodeLoc, ExcData, JsonErrorData, MontyException, StackFrame, UnicodeErrorData, UnicodeErrorObject, }, - io::{PrintStream, PrintWriter, PrintWriterCallback}, + io::{DEFAULT_MAX_PRINT_COLLECT_BYTES, PrintStream, PrintWriter, PrintWriterCallback, check_print_collect_limit}, object::{ DictPairs, InvalidInputError, MontyDate, MontyDateTime, MontyFileHandle, MontyObject, MontyTimeDelta, MontyTimeZone, MontyType, diff --git a/crates/monty/tests/print_collect_limits.rs b/crates/monty/tests/print_collect_limits.rs new file mode 100644 index 000000000..0ca2e629d --- /dev/null +++ b/crates/monty/tests/print_collect_limits.rs @@ -0,0 +1,174 @@ +//! Byte caps on `PrintWriter::CollectString` / `CollectStreams`. +//! +//! Host-side collectors sit outside `ResourceLimits::max_memory` (Monty heap +//! only). These tests lock the optional `max_bytes` check: capped runs raise +//! `MemoryError` without growing past the limit; `None` opts out; `Disabled` +//! under a tight heap limit still succeeds. +//! +//! Loops stay at ~256 KiB — safe, not a real OOM. + +use monty::{ExcType, LimitedTracker, MontyRun, NoLimitTracker, PrintStream, PrintWriter, ResourceLimits}; + +/// One KiB payload reused across prints so heap growth stays small. +const CHUNK: &str = "A"; +const CHUNK_REPS: usize = 1024; +const PRINTS: usize = 256; +/// Host collector target ≈ 256 KiB with `end=''`. +const EXPECTED_MIN_BYTES: usize = CHUNK_REPS * PRINTS; +/// Cap / heap limit well below collected output. +const LIMIT_BYTES: usize = 64 * 1024; + +fn print_loop_code() -> String { + format!("s = '{CHUNK}' * {CHUNK_REPS}\nfor _ in range({PRINTS}):\n print(s, end='')\n") +} + +#[test] +fn collect_string_respects_max_bytes() { + let code = print_loop_code(); + let ex = MontyRun::new(code, "test.py", vec![]).unwrap(); + let mut output = String::new(); + + let err = ex + .run( + vec![], + NoLimitTracker, + PrintWriter::CollectString(&mut output, Some(LIMIT_BYTES)), + ) + .expect_err("expected MemoryError when collect buffer exceeds max_bytes"); + + assert_eq!(err.exc_type(), ExcType::MemoryError); + let expected = format!( + "memory limit exceeded: {} bytes > {LIMIT_BYTES} bytes", + // first write that would cross the limit: LIMIT + one chunk + LIMIT_BYTES + CHUNK_REPS + ); + assert_eq!(err.message(), Some(expected.as_str())); + assert!( + output.len() <= LIMIT_BYTES, + "buffer must stay at or under cap, got {}", + output.len() + ); + // Filled up to the last successful chunk boundary (exact multiple of CHUNK_REPS). + assert_eq!(output.len() % CHUNK_REPS, 0); + assert_eq!(output.len(), LIMIT_BYTES); +} + +#[test] +fn collect_streams_respects_max_bytes() { + let code = print_loop_code(); + let ex = MontyRun::new(code, "test.py", vec![]).unwrap(); + let mut streams: Vec<(PrintStream, String)> = Vec::new(); + + let err = ex + .run( + vec![], + NoLimitTracker, + PrintWriter::CollectStreams(&mut streams, Some(LIMIT_BYTES)), + ) + .expect_err("expected MemoryError when collect buffer exceeds max_bytes"); + + let total: usize = streams.iter().map(|(_, s)| s.len()).sum(); + assert_eq!(err.exc_type(), ExcType::MemoryError); + let expected = format!( + "memory limit exceeded: {} bytes > {LIMIT_BYTES} bytes", + LIMIT_BYTES + CHUNK_REPS + ); + assert_eq!(err.message(), Some(expected.as_str())); + assert!(total <= LIMIT_BYTES, "buffer must stay at or under cap, got {total}"); + assert_eq!(total, LIMIT_BYTES); +} + +/// Control: same loop under tight `max_memory` with `Disabled` still succeeds — +/// proves the print loop itself fits the heap budget; the Collect* failures above +/// are from the host buffer cap, not sandbox memory. +#[test] +fn disabled_print_stays_under_max_memory() { + let code = print_loop_code(); + let ex = MontyRun::new(code, "test.py", vec![]).unwrap(); + let limits = ResourceLimits::new().max_memory(LIMIT_BYTES); + + let result = ex.run(vec![], LimitedTracker::new(limits), PrintWriter::Disabled); + assert!(result.is_ok(), "control failed: {result:?}"); +} + +/// Opt-out: `max_bytes=None` still allows growth past a 64 KiB would-be cap. +#[test] +fn collect_string_unlimited_allows_growth_past_64kib() { + let code = print_loop_code(); + let ex = MontyRun::new(code, "test.py", vec![]).unwrap(); + let mut output = String::new(); + + ex.run(vec![], NoLimitTracker, PrintWriter::CollectString(&mut output, None)) + .expect("unlimited collect should succeed"); + + assert!( + output.len() >= EXPECTED_MIN_BYTES, + "expected >= {EXPECTED_MIN_BYTES} bytes, got {}", + output.len() + ); + assert!( + output.len() > LIMIT_BYTES, + "opt-out not shown: collected {} did not exceed {LIMIT_BYTES}", + output.len() + ); +} + +/// Covers `PrintWriter::collect_streams` and `stdout_push` → `append_streams_char` +/// (the `end=''` loop tests never push a terminator). +#[test] +fn collect_streams_helper_merges_newline_push() { + let ex = MontyRun::new("print('hi')".to_owned(), "test.py", vec![]).unwrap(); + let mut streams: Vec<(PrintStream, String)> = Vec::new(); + + ex.run(vec![], NoLimitTracker, PrintWriter::collect_streams(&mut streams)) + .expect("default-capped collect_streams should accept a short print"); + + assert_eq!(streams, vec![(PrintStream::Stdout, "hi\n".to_owned())]); +} + +/// Bare `print()` only `stdout_push`es `'\n'` — exercises the empty-buffer branch +/// of `append_streams_char`. +#[test] +fn collect_streams_empty_print_pushes_newline_entry() { + let ex = MontyRun::new("print()".to_owned(), "test.py", vec![]).unwrap(); + let mut streams: Vec<(PrintStream, String)> = Vec::new(); + + ex.run(vec![], NoLimitTracker, PrintWriter::collect_streams(&mut streams)) + .expect("empty print should succeed"); + + assert_eq!(streams, vec![(PrintStream::Stdout, "\n".to_owned())]); +} + +/// Cap of 1 byte: `print('a')` writes `'a'` then fails on the newline push. +#[test] +fn collect_string_max_bytes_rejects_newline_push() { + let ex = MontyRun::new("print('a')".to_owned(), "test.py", vec![]).unwrap(); + let mut output = String::new(); + + let err = ex + .run(vec![], NoLimitTracker, PrintWriter::CollectString(&mut output, Some(1))) + .expect_err("expected MemoryError on newline push past max_bytes"); + + assert_eq!(err.exc_type(), ExcType::MemoryError); + assert_eq!(err.message(), Some("memory limit exceeded: 2 bytes > 1 bytes")); + assert_eq!(output, "a"); +} + +/// Same as the string case, but through CollectStreams' char-append path. +#[test] +fn collect_streams_max_bytes_rejects_newline_push() { + let ex = MontyRun::new("print('a')".to_owned(), "test.py", vec![]).unwrap(); + let mut streams: Vec<(PrintStream, String)> = Vec::new(); + + let err = ex + .run( + vec![], + NoLimitTracker, + PrintWriter::CollectStreams(&mut streams, Some(1)), + ) + .expect_err("expected MemoryError on newline push past max_bytes"); + + assert_eq!(err.exc_type(), ExcType::MemoryError); + assert_eq!(err.message(), Some("memory limit exceeded: 2 bytes > 1 bytes")); + assert_eq!(streams, vec![(PrintStream::Stdout, "a".to_owned())]); +} diff --git a/crates/monty/tests/print_writer.rs b/crates/monty/tests/print_writer.rs index 9e9936d05..d057c2fcd 100644 --- a/crates/monty/tests/print_writer.rs +++ b/crates/monty/tests/print_writer.rs @@ -17,7 +17,7 @@ use monty::{MontyRun, NoLimitTracker, PrintWriter}; fn run_and_capture(code: &str) -> String { let ex = MontyRun::new(code.to_owned(), "test.py", vec![]).unwrap(); let mut output = String::new(); - ex.run(vec![], NoLimitTracker, PrintWriter::CollectString(&mut output)) + ex.run(vec![], NoLimitTracker, PrintWriter::collect_string(&mut output)) .unwrap(); output } @@ -103,11 +103,11 @@ fn writer_reuse_accumulates() { let mut output = String::new(); let ex1 = MontyRun::new("print('first')".to_owned(), "test.py", vec![]).unwrap(); - ex1.run(vec![], NoLimitTracker, PrintWriter::CollectString(&mut output)) + ex1.run(vec![], NoLimitTracker, PrintWriter::collect_string(&mut output)) .unwrap(); let ex2 = MontyRun::new("print('second')".to_owned(), "test.py", vec![]).unwrap(); - ex2.run(vec![], NoLimitTracker, PrintWriter::CollectString(&mut output)) + ex2.run(vec![], NoLimitTracker, PrintWriter::collect_string(&mut output)) .unwrap(); assert_snapshot!(output, @r" diff --git a/crates/monty/tests/repl.rs b/crates/monty/tests/repl.rs index a517c9a5a..c4de5b512 100644 --- a/crates/monty/tests/repl.rs +++ b/crates/monty/tests/repl.rs @@ -710,7 +710,7 @@ fn call_function_captures_print() { .call_function( "say_hello", vec![MontyObject::String("world".to_owned())], - PrintWriter::CollectString(&mut output), + PrintWriter::collect_string(&mut output), ) .unwrap(); assert_eq!(result, MontyObject::None); diff --git a/limitations/print.md b/limitations/print.md index 2a9faec19..fb980254b 100644 --- a/limitations/print.md +++ b/limitations/print.md @@ -30,3 +30,22 @@ underneath (see [sys.md](sys.md)). 8 KiB; a single `print()` call can therefore arrive in more than one callback. There is no atomicity guarantee across multiple `print()` calls if the host interleaves with other output. + +## CollectString / CollectStreams caps + +`CollectString` and `CollectStreams` (Rust `PrintWriter` variants and the +matching `pydantic_monty` collectors) accumulate print output in **host-side** +buffers. That growth is **not** covered by +`ResourceLimits.max_memory` (heap-only, and in the pool only on the worker). + +- Default cap: **10 MiB** (`DEFAULT_MAX_PRINT_COLLECT_BYTES`). +- Exceeding the cap raises host-visible `MemoryError` with + `memory limit exceeded: {used} bytes > {limit} bytes` (same wording as + heap `ResourceError::Memory`). +- Pass `max_bytes=None` to disable the cap (trusted hosts only). +- Python `CollectStreams` also charges a fixed per-entry overhead toward the + cap (many tiny fragments would otherwise OOM the host before payload bytes + hit the limit). Rust `PrintWriter::CollectStreams` merges consecutive + same-stream fragments, so entry count stays small for normal `print()`. +- `Stdout` / `Disabled` / `Callback` are unchanged — `Callback` hosts can + already self-limit by returning an error.