diff --git a/src/agents/sandbox/util/token_truncation.py b/src/agents/sandbox/util/token_truncation.py index 41440b33af..c3996d6079 100644 --- a/src/agents/sandbox/util/token_truncation.py +++ b/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 "" diff --git a/tests/sandbox/capabilities/test_shell_capability.py b/tests/sandbox/capabilities/test_shell_capability.py index 533ff62f1f..9eea4e5fee 100644 --- a/tests/sandbox/capabilities/test_shell_capability.py +++ b/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 diff --git a/tests/sandbox/test_memory.py b/tests/sandbox/test_memory.py index c917dacd6d..3359359f15 100644 --- a/tests/sandbox/test_memory.py +++ b/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) diff --git a/tests/sandbox/test_token_truncation.py b/tests/sandbox/test_token_truncation.py index fdd0f0627c..a63bfbce78 100644 --- a/tests/sandbox/test_token_truncation.py +++ b/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,14 +58,26 @@ 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: @@ -63,6 +85,22 @@ def test_truncate_with_token_budget_handles_empty_and_short_content() -> 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))