Skip to content

datarobot_genai.drtools.core.sandbox.utils

utils

Sandbox-backed code execution.

:func:execute_code is a plain async function for now — deliberately not registered as an MCP tool, so it doesn't collide with the execute tool that FastMCP CodeMode exposes (see #376). A later PR will add the MCP tool layer that calls this function, gated by the ENABLE_MCP_SANDBOX DR entitlement (PR datarobot/DataRobot#154256; pattern mirrors datarobot/global-mcp#120).

execute_code async

execute_code(code: str, *, inputs: dict[str, Any] | None = None, timeout_s: float = 30.0, image: str = DEFAULT_SANDBOX_IMAGE) -> dict[str, Any]

Execute Python code in an isolated, short-lived DataRobot workload container.

The container runs with a locked-down securityContext (read-only rootfs, drop ALL caps, no privilege escalation) and is force-deleted in a finally block on success, failure, timeout, or cancellation. Returns a dict with stdout/stderr, an optional _return value the code may assign, duration, exit code, and the image used. The default image is built by DRUM at datarobotdev/datarobot-user-models:public_dropin_environments_dr_mcp_execute_sandbox_minimal_latest and ships polars, pyarrow, datarobot, requests.

This is a plain function — a later PR adds the MCP tool layer that calls it (gated by ENABLE_MCP_SANDBOX), kept separate so it doesn't override FastMCP CodeMode's execute tool.

Source code in datarobot_genai/drtools/core/sandbox/utils.py
async def execute_code(
    code: Annotated[str, "Python source to execute. Assign to `_return` to return a value."],
    *,
    inputs: Annotated[
        dict[str, Any] | None,
        "Optional JSON-serializable mapping bound as `inputs` in the sandbox.",
    ] = None,
    timeout_s: Annotated[float, "Max wall-clock time before forced teardown."] = 30.0,
    image: Annotated[
        str, "Container image URI. Defaults to the minimal sandbox runner."
    ] = DEFAULT_SANDBOX_IMAGE,
) -> dict[str, Any]:
    """Execute Python code in an isolated, short-lived DataRobot workload container.

    The container runs with a locked-down securityContext (read-only rootfs,
    drop ALL caps, no privilege escalation) and is force-deleted in a finally
    block on success, failure, timeout, or cancellation. Returns a dict with
    stdout/stderr, an optional ``_return`` value the code may assign, duration,
    exit code, and the image used. The default image is built by DRUM at
    ``datarobotdev/datarobot-user-models:public_dropin_environments_dr_mcp_execute_sandbox_minimal_latest``
    and ships polars, pyarrow, datarobot, requests.

    This is a plain function — a later PR adds the MCP tool layer that calls it
    (gated by ``ENABLE_MCP_SANDBOX``), kept separate so it doesn't override FastMCP
    CodeMode's ``execute`` tool.
    """
    # Derive credentials the same way the rest of drtools / MCP does, rather
    # than reading DATAROBOT_ENDPOINT / DATAROBOT_API_TOKEN off os.environ: the
    # requesting user's token comes from the request headers (falling back to
    # the application token only in non-HTTP contexts), and the endpoint comes
    # from configured credentials. `get_datarobot_access_token` raises
    # ToolError(AUTHENTICATION) when no token is available.
    token = get_datarobot_access_token()
    endpoint = get_credentials().datarobot.datarobot_endpoint

    # The wrap is unconditional: without a bootstrapped MeterProvider the
    # instruments record into OTel's no-op provider.
    sandbox = InstrumentedSandbox(
        DataRobotWorkloadSandbox(
            image=image,
            datarobot_endpoint=endpoint,
            datarobot_api_token=token,
            security_context=_resolve_security_context(),
        )
    )
    try:
        result = await sandbox.run(code, inputs=inputs, timeout_s=timeout_s)
    except SandboxTimeout as exc:
        raise ToolError(
            f"Sandbox execution timed out after {timeout_s}s: {exc}",
            kind=ToolErrorKind.UPSTREAM,
        ) from exc
    except SandboxError as exc:
        raise ToolError(
            f"Sandbox execution failed: {exc}",
            kind=ToolErrorKind.UPSTREAM,
        ) from exc

    payload = dataclasses.asdict(result)
    payload["image"] = image
    return payload