Preserve truncated output recovery details

interpreter/core/utils/truncate_output.py

@@ -1,3 +1,30 @@
+import re
+import tempfile
+from pathlib import Path
+
+
+ANSI_ESCAPE_RE = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")
+TRUNCATION_PREFIX = "Output truncated ("
+
+
+def _strip_ansi_codes(data):
+    return ANSI_ESCAPE_RE.sub("", data)
+
+
+def _latest_output_path():
+    return Path(tempfile.gettempdir()) / "oi-output-latest.txt"
+
+
+def _build_truncation_message(total_chars, chars_per_end, output_path):
+    return (
+        f"Output truncated ({total_chars:,} visible characters total). "
+        f"Showing {chars_per_end:,} visible characters from start/end. "
+        f"Full output saved to `{output_path}`. "
+        "To inspect it, use tools like `head`, `tail`, or `grep`, "
+        "or rerun the command in smaller steps.\n\n"
+    )
+
+
 def truncate_output(data, max_output_chars=2800, add_scrollbars=False):
     # if "@@@DO_NOT_TRUNCATE@@@" in data:
     #     return data
@@ -6,33 +33,37 @@ def truncate_output(data, max_output_chars=2800, add_scrollbars=False):
 
     # Calculate how much to show from start and end
     chars_per_end = max_output_chars // 2
-
-    message = (f"Output truncated ({len(data):,} characters total). "
-               f"Showing {chars_per_end:,} characters from start/end. "
-               "To handle large outputs, store result in python var first "
-               "`result = command()` then `computer.ai.summarize(result)` for "
-               "a summary, search with `result.find('text')`, "
-               "repeat shell commands with wc/grep/sed, etc. or break it down "
-               "into smaller steps.\n\n")
+    visible_data = _strip_ansi_codes(data)
 
     # This won't work because truncated code is stored in interpreter.messages :/
     # If the full code was stored, we could do this:
     if add_scrollbars:
-        message = (
-            message.strip()
-            + f" Run `get_last_output()[0:{max_output_chars}]` to see the first page.\n\n"
+        extra_scrollbar_message = (
+            f" Run `get_last_output()[0:{max_output_chars}]` to see the first page."
         )
+    else:
+        extra_scrollbar_message = ""
     # Then we have code in `terminal.py` which makes that function work. It should be a computer tool though to just access messages IMO. Or like, self.messages.
 
     # Remove previous truncation message if it exists
-    if data.startswith(message):
-        data = data[len(message) :]
+    if data.startswith(TRUNCATION_PREFIX):
+        _, separator, remainder = data.partition("\n\n")
+        if separator:
+            data = remainder
+            visible_data = _strip_ansi_codes(data)
         needs_truncation = True
 
     # If data exceeds max length, truncate it and add message
-    if len(data) > max_output_chars or needs_truncation:
-        first_part = data[:chars_per_end]
-        last_part = data[-chars_per_end:]
+    if len(visible_data) > max_output_chars or needs_truncation:
+        output_path = _latest_output_path()
+        output_path.write_text(data)
+        message = _build_truncation_message(
+            len(visible_data), chars_per_end, output_path
+        )
+        if extra_scrollbar_message:
+            message = message.strip() + extra_scrollbar_message + "\n\n"
+        first_part = visible_data[:chars_per_end]
+        last_part = visible_data[-chars_per_end:]
         data = message + first_part + "\n[...]\n" + last_part
 
     return data

tests/core/utils/test_truncate_output.py

@@ -0,0 +1,43 @@
+import importlib.util
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+
+
+MODULE_PATH = Path(__file__).resolve().parents[3] / "interpreter/core/utils/truncate_output.py"
+SPEC = importlib.util.spec_from_file_location("truncate_output_module", MODULE_PATH)
+truncate_output_module = importlib.util.module_from_spec(SPEC)
+assert SPEC.loader is not None
+SPEC.loader.exec_module(truncate_output_module)
+truncate_output = truncate_output_module.truncate_output
+
+
+class TestTruncateOutput(unittest.TestCase):
+    def test_does_not_truncate_when_only_ansi_codes_exceed_limit(self):
+        data = "\x1b[31mhello\x1b[0m"
+
+        result = truncate_output(data, max_output_chars=5)
+
+        self.assertEqual(result, data)
+
+    def test_truncation_message_uses_real_recovery_paths(self):
+        data = "abcdefghijklmnopqrstuvwxyz"
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            with mock.patch.object(
+                truncate_output_module.tempfile,
+                "gettempdir",
+                return_value=temp_dir,
+            ):
+                result = truncate_output(data, max_output_chars=10)
+
+            saved_output = Path(temp_dir) / "oi-output-latest.txt"
+            self.assertTrue(saved_output.exists())
+            self.assertEqual(saved_output.read_text(), data)
+
+        self.assertIn("head", result)
+        self.assertIn("tail", result)
+        self.assertIn("grep", result)
+        self.assertNotIn("computer.ai.summarize(result)", result)
+