Cap CollectString/CollectStreams so print loops can't OOM the host

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."""

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,37 +31,72 @@ 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<usize>,
+}
+
 /// Shared buffer for the `CollectStreams` mode.
 ///
 /// The `Arc<Mutex<..>>` 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<Mutex<Vec<(PrintStream, String)>>>;
+type CollectStreamsBuffer = Arc<Mutex<CollectStreamsState>>;
+
+/// Shared collect-string state: flat text plus optional byte cap.
+#[derive(Debug)]
+pub(crate) struct CollectStringState {
+    output: String,
+    max_bytes: Option<usize>,
+}
 
 /// 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<Mutex<String>>;
+type CollectStringBuffer = Arc<Mutex<CollectStringState>>;
 
 /// Python collector that records printed fragments as `(stream, text)` tuples.
 ///
 /// Pass `CollectStreams()` as `print_callback` to share one collector across an
 /// 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<usize>) -> 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<usize>) -> 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<String> {
@@ -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(())
             }
         }

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([])