Ellipsize long outputs from commands

Signed-off-by: Nikola Forró <[email protected]>

Author: nforro

agents/tests/unit/test_tools.py

@@ -85,6 +85,28 @@ async def test_run_shell_command(command, exit_code, stdout, stderr):
     assert result.stderr == stderr
 
 
+@pytest.mark.parametrize(
+    "full_output",
+    [False, True],
+)
+@pytest.mark.asyncio
+async def test_run_shell_command_huge_output(full_output):
+    command = "printf 'Line\n%.0s' {1..1000}"
+    tool = RunShellCommandTool()
+    output = await tool.run(input=RunShellCommandToolInput(command=command, full_output=full_output)).middleware(
+        GlobalTrajectoryMiddleware(pretty=True)
+    )
+    result = output.to_json_safe()
+    assert result.exit_code == 0
+    assert result.stderr is None
+    if full_output:
+        assert len(result.stdout.splitlines()) == 1000
+        assert "[...]" not in result.stdout.splitlines()
+    else:
+        assert len(result.stdout.splitlines()) == 200
+        assert "[...]" in result.stdout.splitlines()
+
+
 @pytest.mark.asyncio
 async def test_add_changelog_entry(minimal_spec):
     content = ["- some change", "  second line"]

agents/tools/commands.py

@@ -1,4 +1,5 @@
 import asyncio
+import math
 from typing import Any
 
 from pydantic import BaseModel, Field
@@ -10,10 +11,19 @@
 from utils import run_subprocess
 
 TIMEOUT = 10 * 60  # seconds
+ELLIPSIZED_LINES = 200
 
 
 class RunShellCommandToolInput(BaseModel):
     command: str = Field(description="Command to run")
+    full_output: bool = Field(
+        default=False,
+        description=(
+            "Whether the content of stdout and stderr should be included in full. "
+            f"Only approximately {ELLIPSIZED_LINES // 2} lines from the beginning "
+            "and the end are included by default."
+        ),
+    )
 
 
 class RunShellCommandToolResult(BaseModel):
@@ -50,9 +60,24 @@ async def _run(
             )
         except TimeoutError as e:
             raise ToolError(f"The specified command timed out after {TIMEOUT} seconds") from e
+
+        def ellipsize(output):
+            if output is None:
+                return None
+            if tool_input.full_output:
+                return output
+            lines = output.splitlines(keepends=True)
+            if len(lines) <= ELLIPSIZED_LINES:
+                return output
+            return "".join(
+                lines[: math.floor((ELLIPSIZED_LINES - 1) / 2)]
+                + ["[...]\n"]
+                + lines[-math.ceil((ELLIPSIZED_LINES - 1) / 2) :]
+            )
+
         result = {
             "exit_code": exit_code,
-            "stdout": stdout,
-            "stderr": stderr,
+            "stdout": ellipsize(stdout),
+            "stderr": ellipsize(stderr),
         }
         return RunShellCommandToolOutput(RunShellCommandToolResult.model_validate(result))