inspect-ai 0.3.92__py3-none-any.whl → 0.3.93__py3-none-any.whl

inspect_ai/tool/_tools/_bash_session.py

@@ -1,20 +1,27 @@
-from pydantic import BaseModel, Field, RootModel
+from textwrap import dedent
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Discriminator, Field, RootModel
+from semver import Version
 from shortuuid import uuid
 
+from inspect_ai._util.error import PrerequisiteError
 from inspect_ai.tool import ToolResult
-from inspect_ai.tool._tool_support_helpers import (
-    exec_model_request,
-    tool_container_sandbox,
-)
 from inspect_ai.util import StoreModel, store_as
+from inspect_ai.util._sandbox.environment import SandboxEnvironment
 
 from .._tool import Tool, ToolParsingError, tool
-from .._tool_call import ToolCall, ToolCallContent, ToolCallView, ToolCallViewer
-
+from .._tool_support_helpers import (
+    exec_model_request,
+    exec_scalar_request,
+    tool_support_sandbox,
+)
 
 # These models are cloned from the container code. If/when we decide to create
 # a package that is shared between the inspect and tool-container codebases, we'll
 # just have to live with it.
+
+
 class NewSessionResult(BaseModel):
     session_name: str
 
@@ -23,106 +30,221 @@ class BashRestartResult(BaseModel):
     pass
 
 
-class BashCommandResult(BaseModel):
-    status: int
-    stdout: str
-    stderr: str
+class BashSessionStore(StoreModel):
+    session_id: str = Field(default_factory=str)
+    sandbox: SandboxEnvironment | None = Field(default=None)
 
 
-class BashResult(RootModel[BashRestartResult | BashCommandResult]):
-    pass
+# Action-specific parameter models
 
 
-class BashSessionStore(StoreModel):
-    session_id: str = Field(default_factory=str)
+class TypeParams(BaseModel):
+    action: Literal["type"] = "type"
+    input: str
 
 
-# custom viewer for bash
-def code_viewer(language: str, code_param: str) -> ToolCallViewer:
-    def viewer(tool_call: ToolCall) -> ToolCallView:
-        code = tool_call.arguments.get(code_param, None)
-        code = (code or tool_call.function).strip()
-        call = ToolCallContent(
-            title=language,
-            format="markdown",
-            content=f"```{language}\n" + code + "\n```\n",
-        )
-        return ToolCallView(call=call)
+class TypeSubmitParams(BaseModel):
+    action: Literal["type_submit"] = "type_submit"
+    input: str
+
+
+class RestartParams(BaseModel):
+    action: Literal["restart"] = "restart"
+
 
-    return viewer
+class ReadParams(BaseModel):
+    action: Literal["read"] = "read"
 
 
-@tool(viewer=code_viewer("bash", "command"))
-def bash_session(*, timeout: int | None = None, instance: str | None = uuid()) -> Tool:
-    """Bash shell session command execution tool.
+class InterruptParams(BaseModel):
+    action: Literal["interrupt"] = "interrupt"
 
-    Execute bash shell commands in a long running session using a sandbox environment (e.g. "docker").
+
+class BashSessionParams(
+    RootModel[
+        TypeParams | TypeSubmitParams | RestartParams | ReadParams | InterruptParams
+    ]
+):
+    root: Annotated[
+        TypeParams | TypeSubmitParams | RestartParams | ReadParams | InterruptParams,
+        Discriminator("action"),
+    ]
+
+
+DEFAULT_WAIT_FOR_OUTPUT = 30
+DEFAULT_IDLE_TIME = 0.5
+# this is how long we're willing to wait for the basic RPC call overhead.
+TRANSPORT_TIMEOUT = 5
+
+
+@tool()
+def bash_session(
+    *,
+    timeout: int | None = None,  # default is max_wait + 5 seconds
+    wait_for_output: int | None = None,  # default is 30 seconds
+    instance: str | None = uuid(),
+) -> Tool:
+    """Interactive bash shell session tool.
+
+    Interact with a bash shell in a long running session using a sandbox
+    environment (e.g. "docker"). This tool allows sending text to the shell,
+    which could be a command followed by a newline character or any other input
+    text such as the response to a password prompt.
 
     By default, a separate bash process is created within the sandbox for each
-    call to `bash_session()`. You can modify this behavior by passing `instance=None`
-    (which will result in a single bash process for the entire sample) or use other
-    `instance` values that implement another scheme).
+    call to `bash_session()`. You can modify this behavior by passing
+    `instance=None` (which will result in a single bash process for the entire
+    sample) or use other `instance` values that implement another scheme).
 
     See complete documentation at <https://inspect.aisi.org.uk/tools-standard.html#sec-bash-session>.
 
     Args:
       timeout: Timeout (in seconds) for command.
+      wait_for_output: Maximum time (in seconds) to wait for output. If no
+          output is received within this period, the function will return an
+          empty string. The model may need to make multiple tool calls to obtain
+          all output from a given command.
       instance: Instance id (each unique instance id has its own bash process)
 
     Returns:
-      String with command output (stdout) or command error (stderr).
+      String with output from the shell.
     """
+    wait_for_output = wait_for_output or DEFAULT_WAIT_FOR_OUTPUT
+    min_timeout = wait_for_output + TRANSPORT_TIMEOUT
+    if timeout is None:
+        timeout = min_timeout
+    elif timeout < min_timeout:
+        raise ValueError(
+            f"Timeout must be at least {min_timeout} seconds, but got {timeout}."
+        )
 
     async def execute(
-        command: str | None = None,
-        restart: bool | None = None,
+        action: Literal["type", "type_submit", "restart", "read", "interrupt"],
+        input: str | None = None,
     ) -> ToolResult:
-        """
-        Use this function to execute bash commands.
+        r"""
+        Interact with a bash shell.
+
+        Interact with a bash shell by sending it input text and retrieving output
+        from it. There is no guarantee that all output will be returned in a
+        single call. Call this function multiple times to retrieve additional
+        output from the shell.
+
+        USAGE NOTES:
+        - Ensure that the shell is at a command prompt (typically when the
+          output ends in "$ " or "# ") before submitting a new command.
+        - Control characters must be sent as Unicode escape sequences (e.g., use
+          "\u0003" for Ctrl+C/ETX, "\u0004" for Ctrl+D/EOT). The literal string
+          "Ctrl+C" will not be interpreted as a control character.
+        - Use the "read" action to retrieve output from the shell without
+          sending any input. This is useful for long-running commands that
+          produce output over time. The "read" action will return any new output
+          since the last call.
+        - If a long-running command is in progress, additional input to execute
+          a new command will not be processed until the previous completes. To
+          abort a long-running command, use the "interrupt" action:
+          `bash_session(action="interrupt")`
+
+        Example use case:
+        - For a short-running command with a nominal amount of output, a single
+          call may suffice.
+          ```
+          bash_session(action="type_submit", input="echo foo") -> "foo\nuser@host:/# "
+          ```
+        - For a long-running command with output over time, multiple calls to are needed.
+          ```
+          bash_session(action="type_submit", input="tail -f /tmp/foo.log") -> <some output>
+          bash_session(action="read") -> <more output>
+          # Send interrupt (Ctrl+C)
+          bash_session(action="interrupt") -> "<final output>^Cuser@host:/# "
+          ```
+        - Interactive command awaiting more input from the user.
+          ```
+          bash_session(action="type_submit", input="ssh fred@foo.com") -> "foo.com's password: "
+          bash_session(action="type_submit", input="secret") -> "fred@foo.com:~$ "
+          ```
 
         Args:
-          command: The bash command to run. Required unless the tool is being restarted.
-          restart: Specifying true will restart this tool. Otherwise, leave this unspecified.
+          action: The action to execute:
+                - "type": Send input without a return key
+                - "type_submit": Send input followed by a return key
+                - "read": Read any new output without sending input
+                - "interrupt": Send a Ctrl+C (ETX character) to interrupt the current process
+                - "restart": Restart the bash session
+          input: The input to send to the shell.
+                Required for "type". Optional for "type_submit" actions. Must
+                not be provided for "restart", "read", or "interrupt" actions.
 
         Returns:
-          The output of the command.
+          The accumulated output of the shell.
         """
-        if not ((command is None) ^ (restart is None)):
-            raise ToolParsingError(
-                "Either 'command' or 'restart' must be specified, but not both."
-            )
-        params: dict[str, object] = {"command": command, "restart": restart}
+        # Validate parameters based on action
+        match action:
+            case "type":
+                if input is None:
+                    raise ToolParsingError(
+                        f"'input' is required for '{action}' action."
+                    )
+            case "restart" | "read" | "interrupt":
+                if input is not None:
+                    raise ToolParsingError(
+                        f"Do not provide 'input' with '{action}' action."
+                    )
 
-        sandbox = await tool_container_sandbox("bash session")
         store = store_as(BashSessionStore, instance=instance)
+        sandbox = await _get_sandbox(store)
 
         if not store.session_id:
             store.session_id = (
                 await exec_model_request(
-                    sandbox=sandbox,
-                    method="bash_session_new_session",
-                    params={},
-                    result_type=NewSessionResult,
-                    timeout=timeout,
+                    sandbox,
+                    "bash_session_new_session",
+                    {},
+                    NewSessionResult,
+                    TRANSPORT_TIMEOUT,
                 )
             ).session_name
 
-        params["session_name"] = store.session_id
+        timing: dict[str, object] = {
+            "wait_for_output": wait_for_output,
+            "idle_timeout": DEFAULT_IDLE_TIME,
+        }
+        action_specific: dict[str, dict[str, object]] = {
+            "type": {"input": input, **timing},
+            "type_submit": {"input": f"{input}\n", **timing},
+            "interrupt": {"input": "\u0003", **timing},
+            "read": timing,
+            "restart": {"restart": True},
+        }
+
+        result = await exec_scalar_request(
+            sandbox,
+            "bash_session",
+            {"session_name": store.session_id, **(action_specific[action])},
+            str,
+            timeout,
+        )
 
-        result = (
-            await exec_model_request(
-                sandbox=sandbox,
-                method="bash_session",
-                params=params,
-                result_type=BashResult,
-                timeout=timeout,
-            )
-        ).root
+        # Return the appropriate response
+        return (
+            "Bash session restarted."
+            if isinstance(result, BashRestartResult)
+            else result
+        )
 
-        if isinstance(result, BashRestartResult):
-            return "Bash session restarted."
+    return execute
 
-        # return output (including stderr if any)
-        return f"{result.stderr}\n{result.stdout}" if result.stderr else result.stdout
 
-    return execute
+async def _get_sandbox(store: BashSessionStore) -> SandboxEnvironment:
+    if not store.sandbox:
+        (sandbox, sandbox_version) = await tool_support_sandbox("bash session")
+        required_version = Version.parse("1.0.0")
+        if sandbox_version < required_version:
+            raise PrerequisiteError(
+                dedent(f"""
+                    The 'inspect-tool-support' version in your container is '{sandbox_version}'. The 'bash_session' tool requires version '{required_version}' or newer. Please update your container image to the latest version of 'inspect-tool-support'.
+                    """).strip()
+            )
+        store.sandbox = sandbox
+
+    return store.sandbox

inspect_ai/tool/_tools/_computer/_computer.py

@@ -6,7 +6,31 @@ from inspect_ai.tool._tool import TOOL_INIT_MODEL_INPUT, ToolParsingError
 from inspect_ai.tool._tool_call import ToolCallModelInput, ToolCallModelInputHints
 
 from . import _common as common
-from ._resources.tool._constants import Action
+
+# this is duplicated from ._resources.tool._constants import Action
+# changes should be synchronized!
+
+Action = Literal[
+    "key",
+    "hold_key",
+    "type",
+    "cursor_position",
+    "mouse_move",
+    "left_mouse_down",
+    "left_mouse_up",
+    "left_click",
+    "left_click_drag",
+    "right_click",
+    "middle_click",
+    "back_click",
+    "forward_click",
+    "double_click",
+    "triple_click",
+    "scroll",
+    "wait",
+    "screenshot",
+]
+
 
 ActionFunction = Callable[[str], ToolResult | Awaitable[ToolResult]]
 

inspect_ai/tool/_tools/_text_editor.py

@@ -6,7 +6,7 @@ from pydantic import BaseModel, Discriminator, RootModel
 from inspect_ai.tool import ToolResult
 from inspect_ai.tool._tool_support_helpers import (
     exec_scalar_request,
-    tool_container_sandbox,
+    tool_support_sandbox,
 )
 
 from .._tool import Tool, tool
@@ -70,12 +70,13 @@ def text_editor(timeout: int | None = None, user: str | None = None) -> Tool:
     that a change made to a file by on Subtask will be visible to another Subtask.
 
     Args:
-      timeout: Timeout (in seconds) for command.
+      timeout: Timeout (in seconds) for command. Defaults to 180 if not provided.
       user: User to execute commands as.
 
     Returns:
       String with command output (stdout) or command error (stderr).
     """
+    timeout = timeout or 180
 
     async def execute(
         command: Literal["view", "create", "str_replace", "insert", "undo_edit"],
@@ -101,7 +102,7 @@ def text_editor(timeout: int | None = None, user: str | None = None) -> Tool:
         Returns:
           The output of the command.
         """
-        sandbox = await tool_container_sandbox("editor")
+        (sandbox, _) = await tool_support_sandbox("editor")
 
         # Create a dictionary of the parameters
         params = {

inspect_ai/tool/_tools/_web_browser/_web_browser.py

@@ -10,7 +10,7 @@ from inspect_ai.tool._tool_call import ToolCall, ToolCallContent, ToolCallView
 from inspect_ai.tool._tool_info import parse_tool_info
 from inspect_ai.tool._tool_support_helpers import (
     exec_model_request,
-    tool_container_sandbox,
+    tool_support_sandbox,
 )
 from inspect_ai.tool._tool_with import tool_with
 from inspect_ai.util._store_model import StoreModel, store_as
@@ -397,8 +397,10 @@ def web_browser_refresh(instance: str | None = None) -> Tool:
 async def _web_browser_cmd(
     tool_name: str, instance: str | None, params: dict[str, object]
 ) -> ToolResult:
+    # TODO: Is it worth it to plumb this down from the @tool?
+    timeout = 180
     try:
-        sandbox_env = await tool_container_sandbox("web browser")
+        (sandbox_env, _) = await tool_support_sandbox("web browser")
     except PrerequisiteError as e:
         # The user may have the old, incompatible, sandbox. If so, use that and
         # execute the old compatible code.
@@ -419,13 +421,18 @@ async def _web_browser_cmd(
                 method="web_new_session",
                 params={"headful": False},
                 result_type=NewSessionResult,
+                timeout=timeout,
             )
         ).session_name
 
     params["session_name"] = store.session_id
 
     crawler_result = await exec_model_request(
-        sandbox=sandbox_env, method=tool_name, params=params, result_type=CrawlerResult
+        sandbox=sandbox_env,
+        method=tool_name,
+        params=params,
+        result_type=CrawlerResult,
+        timeout=timeout,
     )
     if crawler_result.error and crawler_result.error.strip() != "":
         raise ToolError(crawler_result.error)

inspect_ai/util/__init__.py

@@ -1,5 +1,12 @@
 from inspect_ai._util.registry import RegistryType, registry_create
 from inspect_ai._util.trace import trace_action, trace_message
+from inspect_ai.util._limit import (
+    Limit,
+    LimitExceededError,
+    apply_limits,
+    message_limit,
+    token_limit,
+)
 
 from ._concurrency import concurrency
 from ._console import input_screen
@@ -31,6 +38,7 @@ from ._subtask import Subtask, subtask
 from ._throttle import throttle
 
 __all__ = [
+    "apply_limits",
     "ExecResult",
     "concurrency",
     "DisplayType",
@@ -42,9 +50,12 @@ __all__ = [
     "JSONType",
     "JSONSchema",
     "json_schema",
+    "Limit",
+    "message_limit",
     "OutputLimitExceededError",
     "resource",
     "subprocess",
+    "LimitExceededError",
     "SandboxEnvironment",
     "SandboxEnvironmentConfigType",
     "SandboxEnvironmentLimits",
@@ -63,6 +74,7 @@ __all__ = [
     "Subtask",
     "subtask",
     "throttle",
+    "token_limit",
     "trace_action",
     "trace_message",
     "RegistryType",