fix(sandbox): enforce token output budgets

src/agents/sandbox/util/token_truncation.py

@@ -40,6 +40,9 @@ def formatted_truncate_text(content: str, policy: TruncationPolicy) -> str:
     if _byte_len(content) <= policy.byte_budget():
         return content
     total_lines = len(content.splitlines())
+    if policy.mode == "tokens":
+        prefix = f"Total output lines: {total_lines}\n\n"
+        return _truncate_token_output(content, policy, prefix=prefix)
     result = truncate_text(content, policy)
     return f"Total output lines: {total_lines}\n\n{result}"
 
@@ -61,9 +64,10 @@ def formatted_truncate_text_with_token_count(
     if _byte_len(content) <= policy.byte_budget():
         return content, None
 
-    truncated, original_token_count = truncate_with_token_budget(content, policy)
     total_lines = len(content.splitlines())
-    return f"Total output lines: {total_lines}\n\n{truncated}", original_token_count
+    prefix = f"Total output lines: {total_lines}\n\n"
+    truncated = _truncate_token_output(content, policy, prefix=prefix)
+    return truncated, approx_token_count(content)
 
 
 def truncate_with_token_budget(s: str, policy: TruncationPolicy) -> tuple[str, int | None]:
@@ -75,13 +79,44 @@ def truncate_with_token_budget(s: str, policy: TruncationPolicy) -> tuple[str, i
     if max_tokens > 0 and byte_len <= approx_bytes_for_tokens(max_tokens):
         return s, None
 
-    truncated = truncate_with_byte_estimate(s, policy)
     approx_total = approx_token_count(s)
+    truncated = _truncate_token_output(s, policy)
     if truncated == s:
         return truncated, None
     return truncated, approx_total
 
 
+def _truncate_token_output(content: str, policy: TruncationPolicy, *, prefix: str = "") -> str:
+    max_bytes = policy.byte_budget()
+    if max_bytes == 0:
+        return ""
+
+    source_bytes = content.encode("utf-8")
+    marker = format_truncation_marker(policy, approx_token_count(content))
+
+    if _byte_len(prefix) + _byte_len(marker) > max_bytes:
+        prefix = ""
+
+    content_budget = max_bytes - _byte_len(prefix) - _byte_len(marker)
+    if content_budget < 0:
+        return _truncate_utf8(marker, max_bytes)
+
+    left_budget, right_budget = split_budget(content_budget)
+    _, left, right = split_string(content, left_budget, right_budget)
+    retained_bytes = _byte_len(left) + _byte_len(right)
+    removed_bytes = len(source_bytes) - retained_bytes
+    removed_chars = len(content) - len(left) - len(right)
+    marker = format_truncation_marker(
+        policy,
+        removed_units_for_source(policy, removed_bytes, removed_chars),
+    )
+    return assemble_truncated_output(f"{prefix}{left}", right, marker)
+
+
+def _truncate_utf8(text: str, max_bytes: int) -> str:
+    return text.encode("utf-8")[:max_bytes].decode("utf-8", errors="ignore")
+
+
 def truncate_with_byte_estimate(s: str, policy: TruncationPolicy) -> str:
     if s == "":
         return ""

tests/sandbox/capabilities/test_shell_capability.py

@@ -469,8 +469,7 @@ async def test_exec_command_tool_includes_original_token_count_when_truncating(
             "Process exited with code 7\n"
             "Original token count: 6\n"
             "Output:\n"
-            "Total output lines: 2\n\n"
-            "stdo…4 tokens truncated… pwd"
+            "…6 tok"
         )
 
     @pytest.mark.asyncio

tests/sandbox/test_memory.py

@@ -817,11 +817,11 @@ async def test_memory_capability_injects_truncated_memory_summary(
 
     try:
         async with session:
-            monkeypatch.setattr(memory_module, "_MEMORY_SUMMARY_MAX_TOKENS", 1)
+            monkeypatch.setattr(memory_module, "_MEMORY_SUMMARY_MAX_TOKENS", 8)
             await session.mkdir("memories", parents=True)
             await session.write(
                 Path("memories/memory_summary.md"),
-                io.BytesIO(b"abcdefg"),
+                io.BytesIO(b"abcdefghijklmnopqrstuvwxyz" * 2),
             )
             capability.bind(session)
 

tests/sandbox/test_token_truncation.py

@@ -40,6 +40,16 @@ def test_formatted_truncate_text_adds_line_count_when_truncated() -> None:
     assert "chars truncated" in result
 
 
+def test_formatted_truncate_text_keeps_token_metadata_within_budget() -> None:
+    content = "\n".join(f"line {index}: {('value ' * 8).strip()}" for index in range(20))
+
+    result = formatted_truncate_text(content, TruncationPolicy.tokens(32))
+
+    assert result.startswith("Total output lines: 20\n\n")
+    assert "tokens truncated" in result
+    assert approx_token_count(result) <= 32
+
+
 def test_formatted_truncate_text_with_token_count_handles_none_and_short_content() -> None:
     assert formatted_truncate_text_with_token_count("short", None) == ("short", None)
     assert formatted_truncate_text_with_token_count("short", 10) == ("short", None)
@@ -48,21 +58,49 @@ def test_formatted_truncate_text_with_token_count_handles_none_and_short_content
 def test_formatted_truncate_text_with_token_count_reports_original_count() -> None:
     result, original_token_count = formatted_truncate_text_with_token_count("abcdefghi", 1)
 
-    assert result.startswith("Total output lines: 1\n\n")
-    assert "tokens truncated" in result
+    assert approx_token_count(result) <= 1
     assert original_token_count == approx_token_count("abcdefghi")
 
 
+def test_formatted_truncate_text_with_token_count_keeps_metadata_within_budget() -> None:
+    content = "\n".join(f"line {index}: {('value ' * 8).strip()}" for index in range(20))
+
+    result, original_token_count = formatted_truncate_text_with_token_count(content, 32)
+
+    assert result.startswith("Total output lines: 20\n\n")
+    assert "tokens truncated" in result
+    assert approx_token_count(result) <= 32
+    assert original_token_count == approx_token_count(content)
+
+
 def test_truncate_text_dispatches_byte_and_token_modes() -> None:
     assert truncate_text("abcdef", TruncationPolicy.bytes(4)).startswith("a")
-    assert "tokens truncated" in truncate_text("abcdefghi", TruncationPolicy.tokens(1))
+    token_result = truncate_text("abcdefghijklmnopqrstuvwxyz" * 2, TruncationPolicy.tokens(8))
+    assert "tokens truncated" in token_result
+    assert approx_token_count(token_result) <= 8
 
 
 def test_truncate_with_token_budget_handles_empty_and_short_content() -> None:
     assert truncate_with_token_budget("", TruncationPolicy.tokens(1)) == ("", None)
     assert truncate_with_token_budget("abc", TruncationPolicy.tokens(1)) == ("abc", None)
 
 
+def test_truncate_with_token_budget_includes_marker_within_budget() -> None:
+    content = "abcdefghijklmnopqrstuvwxyz" * 2
+    result, original_token_count = truncate_with_token_budget(content, TruncationPolicy.tokens(8))
+
+    assert "tokens truncated" in result
+    assert approx_token_count(result) <= 8
+    assert original_token_count == approx_token_count(content)
+
+
+def test_formatted_truncate_text_with_zero_token_budget_returns_empty_payload() -> None:
+    result, original_token_count = formatted_truncate_text_with_token_count("content", 0)
+
+    assert result == ""
+    assert original_token_count == approx_token_count("content")
+
+
 def test_truncate_with_byte_estimate_handles_empty_zero_and_short_content() -> None:
     assert truncate_with_byte_estimate("", TruncationPolicy.bytes(0)) == ""
     assert "chars truncated" in truncate_with_byte_estimate("abc", TruncationPolicy.bytes(0))