From 13822ece177d46bd54d364e9018790836dca4940 Mon Sep 17 00:00:00 2001 From: Jaideep <67646710+jdchawla29@users.noreply.github.com> Date: Thu, 23 Jul 2026 02:19:26 -0700 Subject: [PATCH] fix(openai): bound native shell output (HUD-2204) --- hud/agents/openai/tools/coding.py | 77 ++++++++--- .../tests/test_provider_native_tools.py | 122 +++++++++++++++++- 2 files changed, 181 insertions(+), 18 deletions(-) diff --git a/hud/agents/openai/tools/coding.py b/hud/agents/openai/tools/coding.py index 87f362e0a..bcf9d8761 100644 --- a/hud/agents/openai/tools/coding.py +++ b/hud/agents/openai/tools/coding.py @@ -7,7 +7,6 @@ import mcp.types as mcp_types from hud.agents.tools import SSHTool -from hud.agents.tools.base import result_text from hud.types import MCPToolResult from .base import OpenAIToolSpec @@ -17,6 +16,9 @@ api_name="shell", ) +MAX_SHELL_OUTPUT_LENGTH = 10 * 1024 * 1024 +TRUNCATION_MARKER = "[truncated]" + class OpenAIShellTool(SSHTool): name = "shell" @@ -32,14 +34,31 @@ def to_params(self) -> Any: return {"type": "shell", "environment": {"type": "local"}} async def execute(self, arguments: dict[str, Any]) -> MCPToolResult: + requested_limit = arguments.get("max_output_length") + if requested_limit is None: + max_output_length = MAX_SHELL_OUTPUT_LENGTH + elif ( + isinstance(requested_limit, bool) + or not isinstance(requested_limit, int) + or requested_limit <= 0 + ): + text = "max_output_length must be a positive integer" + return _shell_result( + [text], + is_error=True, + structured={"output": [shell_output("", text, 1)]}, + ) + else: + max_output_length = min(requested_limit, MAX_SHELL_OUTPUT_LENGTH) + def invalid_commands_result() -> MCPToolResult: text = "commands must be a list of strings" return _shell_result( - text, + [text], is_error=True, structured={ "output": [shell_output("", text, 1)], - "max_output_length": arguments.get("max_output_length"), + "max_output_length": max_output_length, }, ) @@ -54,7 +73,7 @@ def invalid_commands_result() -> MCPToolResult: command_list = cast("list[str]", raw_commands) outputs: list[dict[str, Any]] = [] - text_parts: list[str] = [] + display_outputs: list[str] = [] is_error = False env_arguments: dict[str, Any] = {} timeout_ms = arguments.get("timeout_ms") @@ -66,40 +85,64 @@ def invalid_commands_result() -> MCPToolResult: full_cmd = f"timeout {int(env_arguments['timeout_seconds'])} {command}" else: full_cmd = command - result = await self.bash(full_cmd) - text = result_text(result) - if result.isError: - outputs.append(shell_output("", text, 1)) - is_error = True - else: - outputs.append(shell_output(text, "", 0)) - if text: - text_parts.append(text) + completed = await self.client.conn.run(full_cmd, check=False) + stdout = completed.stdout if isinstance(completed.stdout, str) else "" + stderr = completed.stderr if isinstance(completed.stderr, str) else "" + exit_code = completed.exit_status if completed.exit_status is not None else 1 + stdout, stderr = _bound_output(stdout, stderr, max_output_length) + outputs.append(shell_output(stdout, stderr, exit_code)) + display_outputs.append(stdout + stderr) + is_error = is_error or bool(exit_code) return _shell_result( - "\n".join(text_parts), + display_outputs, is_error=is_error, structured={ "output": outputs, - "max_output_length": arguments.get("max_output_length"), + "max_output_length": max_output_length, }, ) def _shell_result( - text: str, + texts: list[str], *, is_error: bool = False, structured: dict[str, Any] | None = None, ) -> MCPToolResult: payload = {"provider_tool": "shell", **(structured or {})} return MCPToolResult( - content=[mcp_types.TextContent(type="text", text=text)] if text else [], + content=[mcp_types.TextContent(type="text", text=text) for text in texts if text], isError=is_error, structuredContent=payload, ) +def _bound_output(stdout: str, stderr: str, limit: int) -> tuple[str, str]: + total_length = len(stdout) + len(stderr) + if total_length <= limit: + return stdout, stderr + + marker = TRUNCATION_MARKER[:limit] + available = limit - len(marker) + prefix_length = (available + 1) // 2 + suffix_length = available // 2 + + stdout_prefix = stdout[:prefix_length] + stderr_prefix = stderr[: max(0, prefix_length - len(stdout))] + stderr_suffix = stderr[-suffix_length:] if suffix_length else "" + stdout_suffix_length = max(0, suffix_length - len(stderr)) + stdout_suffix = stdout[-stdout_suffix_length:] if stdout_suffix_length else "" + + if len(stdout_prefix) + len(stdout_suffix) < len(stdout): + stdout = stdout_prefix + marker + stdout_suffix + stderr = stderr_prefix + stderr_suffix + else: + stdout = stdout_prefix + stdout_suffix + stderr = stderr_prefix + marker + stderr_suffix + return stdout, stderr + + def shell_output(stdout: str, stderr: str, exit_code: int) -> dict[str, Any]: return { "stdout": stdout, diff --git a/hud/agents/tests/test_provider_native_tools.py b/hud/agents/tests/test_provider_native_tools.py index b5aa7c343..3711b0832 100644 --- a/hud/agents/tests/test_provider_native_tools.py +++ b/hud/agents/tests/test_provider_native_tools.py @@ -10,6 +10,7 @@ import shlex from typing import Any, cast +import mcp.types as mcp_types import pytest from hud.agents.claude.tools.coding import ClaudeBashTool, ClaudeTextEditorTool @@ -82,6 +83,7 @@ def __init__( self, *, stdout: str = "ok", + stderr: str = "", exit_status: int = 0, files: dict[str, bytes] | None = None, cwd: str | None = None, @@ -90,7 +92,13 @@ def __init__( params = {"cwd": cwd} if cwd else {} super().__init__( Capability(name="shell", protocol="ssh/2", url="ssh://localhost:22", params=params), - cast("Any", _Conn(_Completed(stdout=stdout, exit_status=exit_status), self.files)), + cast( + "Any", + _Conn( + _Completed(stdout=stdout, stderr=stderr, exit_status=exit_status), + self.files, + ), + ), ) @@ -130,6 +138,118 @@ async def test_openai_shell_runs_each_command_without_timeout() -> None: assert _commands(tool) == ["echo a", "echo b"] +async def test_openai_shell_bounds_large_output_in_every_result_field() -> None: + limit = 20_000 + stderr_prefix = "stderr-start\n" + stderr_suffix = "\nstderr-end" + stderr = ( + stderr_prefix + "x" * (10_523_560 - len(stderr_prefix) - len(stderr_suffix)) + stderr_suffix + ) + tool = OpenAIShellTool( + spec=OpenAIShellTool.default_spec("gpt-5.5"), + client=_ssh(stderr=stderr, exit_status=1), + ) + + result = await tool.execute({"commands": ["noisy-command"], "max_output_length": limit}) + + assert result.structuredContent is not None + assert result.structuredContent["max_output_length"] == limit + output = result.structuredContent["output"][0] + assert len(output["stdout"]) + len(output["stderr"]) == limit + assert output["stderr"].startswith(stderr_prefix) + assert output["stderr"].endswith(stderr_suffix) + assert "[truncated]" in output["stderr"] + text_blocks = [ + block.text for block in result.content if isinstance(block, mcp_types.TextContent) + ] + assert len(text_blocks) == 1 + assert len(text_blocks[0]) == limit + assert text_blocks[0] == output["stdout"] + output["stderr"] + assert "[truncated]" in text_blocks[0] + + +async def test_openai_shell_applies_limit_independently_to_each_command() -> None: + limit = 80 + tool = OpenAIShellTool( + spec=OpenAIShellTool.default_spec("gpt-5.5"), + client=_ssh(stdout="stdout-start-" + "a" * 100, stderr="b" * 100 + "-stderr-end"), + ) + + result = await tool.execute( + {"commands": ["first-command", "second-command"], "max_output_length": limit} + ) + + assert result.structuredContent is not None + outputs = result.structuredContent["output"] + assert len(outputs) == 2 + assert all(len(output["stdout"]) + len(output["stderr"]) == limit for output in outputs) + assert all(output["stdout"].startswith("stdout-start-") for output in outputs) + assert all(output["stderr"].endswith("-stderr-end") for output in outputs) + assert sum(len(output["stdout"]) + len(output["stderr"]) for output in outputs) == 2 * limit + text_blocks = [ + block.text for block in result.content if isinstance(block, mcp_types.TextContent) + ] + assert len(text_blocks) == 2 + assert all(len(text) == limit for text in text_blocks) + assert text_blocks == [output["stdout"] + output["stderr"] for output in outputs] + + +@pytest.mark.parametrize( + ("max_output_length", "expected_output"), + [(1, "["), (10, "[truncated")], +) +async def test_openai_shell_honors_small_positive_output_limits( + max_output_length: int, + expected_output: str, +) -> None: + tool = OpenAIShellTool( + spec=OpenAIShellTool.default_spec("gpt-5.5"), + client=_ssh(stdout="x" * 100), + ) + + result = await tool.execute( + {"commands": ["noisy-command"], "max_output_length": max_output_length} + ) + + assert _commands(tool) == ["noisy-command"] + assert result.structuredContent is not None + assert result.structuredContent["max_output_length"] == max_output_length + output = result.structuredContent["output"][0] + assert output["stdout"] == expected_output + assert output["stderr"] == "" + assert result_text(result) == expected_output + + +@pytest.mark.parametrize("max_output_length", [0, -1, "20000", 20_000.0, True]) +async def test_openai_shell_rejects_invalid_output_limits_without_running( + max_output_length: Any, +) -> None: + tool = OpenAIShellTool(spec=OpenAIShellTool.default_spec("gpt-5.5"), client=_ssh()) + + result = await tool.execute( + {"commands": ["echo should-not-run"], "max_output_length": max_output_length} + ) + + assert result.isError is True + assert _commands(tool) == [] + assert result.structuredContent is not None + assert "max_output_length" not in result.structuredContent + assert result_text(result) == "max_output_length must be a positive integer" + + +@pytest.mark.parametrize("max_output_length", [None, 20 * 1024 * 1024]) +async def test_openai_shell_uses_safe_effective_limit(max_output_length: int | None) -> None: + tool = OpenAIShellTool(spec=OpenAIShellTool.default_spec("gpt-5.5"), client=_ssh()) + arguments: dict[str, Any] = {"commands": ["echo ok"]} + if max_output_length is not None: + arguments["max_output_length"] = max_output_length + + result = await tool.execute(arguments) + + assert result.structuredContent is not None + assert result.structuredContent["max_output_length"] == 10 * 1024 * 1024 + + async def test_openai_shell_rejects_non_list_commands_without_running() -> None: tool = OpenAIShellTool(spec=OpenAIShellTool.default_spec("gpt-5.5"), client=_ssh())