diff --git a/navi/tools/code_exec.py b/navi/tools/code_exec.py index 78671ac..3b9c5ce 100644 --- a/navi/tools/code_exec.py +++ b/navi/tools/code_exec.py @@ -20,7 +20,20 @@ current_working_directory, ) -_TIMEOUT = 30 +_DEFAULT_TIMEOUT = 30 +_MAX_TIMEOUT = 300 +_MIN_TIMEOUT = 1 + + +def _resolve_timeout(value: object) -> int: + """Clamp a caller-supplied timeout to the allowed range, defaulting to 30s.""" + if value is None: + return _DEFAULT_TIMEOUT + try: + secs = int(value) + except (TypeError, ValueError): + return _DEFAULT_TIMEOUT + return max(_MIN_TIMEOUT, min(secs, _MAX_TIMEOUT)) def _resolve_working_dir( @@ -66,7 +79,8 @@ "Run Python code and return output. Use for calculations, data parsing, " "text processing, or anything that benefits from a script. " "Each call is a fresh interpreter — import everything you need, no state persists. " - "For shell-native tasks (pipes, system commands) prefer terminal instead." + "For shell-native tasks (pipes, system commands) prefer terminal instead. " + "Timeout defaults to 30s (max 300s) — set the `timeout` param for long runs." ) parameters = { "type": "object", @@ -79,6 +93,13 @@ "type": "string", "description": "Working directory for the script (optional).", }, + "timeout": { + "type": "number", + "description": ( + "Execution timeout in seconds (default 30, max 300). " + "Raise it for long computations or test suites." + ), + }, }, "required": ["code"], } @@ -93,6 +114,8 @@ params.get("working_dir"), user_id, role, cwd=ctx.cwd if ctx else None ) + timeout = _resolve_timeout(params.get("timeout")) + if user_id and role != "admin": # Write temp file inside the sandbox so file I/O in user code # is implicitly sandboxed unless they escape via absolute paths. @@ -112,13 +135,14 @@ cwd=str(cwd), ) try: - stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=_TIMEOUT) + stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout) except asyncio.TimeoutError: proc.kill() return ToolResult( success=False, - output=f"Code execution timed out after {_TIMEOUT}s", + output=f"Code execution timed out after {timeout}s", error="timeout", + metadata={"returncode": None, "language": "python", "timeout": timeout}, ) output_parts = [] @@ -131,7 +155,7 @@ return ToolResult( success=success, output="\n".join(output_parts) or "(no output)", - metadata={"returncode": proc.returncode}, + metadata={"returncode": proc.returncode, "language": "python"}, error=None if success else f"Exit code {proc.returncode}", ) except Exception as e: diff --git a/tests/unit/tools/test_code_exec.py b/tests/unit/tools/test_code_exec.py index 3f240bc..b303fa9 100644 --- a/tests/unit/tools/test_code_exec.py +++ b/tests/unit/tools/test_code_exec.py @@ -28,3 +28,24 @@ async def test_syntax_error(self, tool): result = await tool.execute({"code": "print("}) assert not result.success + + async def test_metadata_carries_returncode_and_language(self, tool): + result = await tool.execute({"code": "print('hi')"}) + assert result.success + assert result.metadata["returncode"] == 0 + assert result.metadata["language"] == "python" + + async def test_timeout_param_clamped_to_max(self, tool): + # 99999 -> clamped to _MAX_TIMEOUT (300). Fast code still succeeds; we + # only assert the clamp path doesn't error and metadata is well-formed. + result = await tool.execute({"code": "print('ok')", "timeout": 99999}) + assert result.success + assert result.metadata["language"] == "python" + + async def test_timeout_param_triggers_timeout(self, tool): + result = await tool.execute({"code": "import time; time.sleep(5)", "timeout": 1}) + assert not result.success + assert result.error == "timeout" + assert "timed out after 1s" in result.output + assert result.metadata["language"] == "python" + assert result.metadata["timeout"] == 1