From 0c6deef5c3509ebc2437085919d8575f42215c46 Mon Sep 17 00:00:00 2001 From: HsingTang <32690368+HsingTang@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:31:50 +0000 Subject: [PATCH] feat: add Databricks sandbox backend for the native executor Provide a native Environment backed by a Databricks sandbox for the v1.0 (codex-harness) SDK. A localhost WSS gateway relays the SDK executor connection to the sandbox exec-server over the databricks-sdk streaming command API; databricks_sandbox_environment provisions the gateway and yields a configured Environment. Add a databricks optional-dependency extra, unit tests for the relay and bearer auth, a runnable example, and a package README. The README flags two prerequisites to confirm against a live sandbox: exec-server stdio mode in the image, and bidirectional streaming via execute_command. Co-authored-by: Isaac --- examples/databricks/runner.py | 53 ++++ pyproject.toml | 12 + src/agents/extensions/__init__.py | 1 + src/agents/extensions/sandbox/__init__.py | 6 + .../extensions/sandbox/databricks/README.md | 106 ++++++++ .../extensions/sandbox/databricks/__init__.py | 17 ++ .../sandbox/databricks/environment.py | 135 ++++++++++ .../extensions/sandbox/databricks/gateway.py | 239 ++++++++++++++++++ .../databricks/test_databricks_environment.py | 164 ++++++++++++ .../databricks/test_databricks_gateway.py | 203 +++++++++++++++ 10 files changed, 936 insertions(+) create mode 100644 examples/databricks/runner.py create mode 100644 src/agents/extensions/__init__.py create mode 100644 src/agents/extensions/sandbox/__init__.py create mode 100644 src/agents/extensions/sandbox/databricks/README.md create mode 100644 src/agents/extensions/sandbox/databricks/__init__.py create mode 100644 src/agents/extensions/sandbox/databricks/environment.py create mode 100644 src/agents/extensions/sandbox/databricks/gateway.py create mode 100644 tests/extensions/sandbox/databricks/test_databricks_environment.py create mode 100644 tests/extensions/sandbox/databricks/test_databricks_gateway.py diff --git a/examples/databricks/runner.py b/examples/databricks/runner.py new file mode 100644 index 0000000000..5ee6c1837f --- /dev/null +++ b/examples/databricks/runner.py @@ -0,0 +1,53 @@ +"""Run an agent against a Databricks sandbox through the native executor. + +Prerequisites (see src/agents/extensions/sandbox/databricks/README.md): + - the SDK preview wheel plus the ``databricks`` extra + - a running Databricks sandbox whose image carries the ``codex`` exec-server + - Databricks auth (a CLI profile) and model-provider auth + +Environment variables: + DATABRICKS_CONFIG_PROFILE Databricks CLI profile (default: DEFAULT) + SANDBOX_NAME the target sandbox name/id (required) + SANDBOX_CWD remote workspace path (default: /home/sandbox-agent/work) + OPENAI_API_KEY model-provider key for NativeProvider (required) + OPENAI_BASE_URL optional model-provider base URL + OPENAI_MODEL model name (default: gpt-5.5) + +Run: + python -m examples.databricks.runner "Inspect the project and list next steps" +""" + +import asyncio +import os +import sys + +from databricks.sdk import WorkspaceClient + +from agents import Agent, NativeProvider +from agents.extensions.sandbox.databricks import databricks_sandbox_environment + + +async def main(task: str) -> None: + workspace_client = WorkspaceClient(profile=os.getenv("DATABRICKS_CONFIG_PROFILE", "DEFAULT")) + async with databricks_sandbox_environment( + workspace_client=workspace_client, + sandbox_name=os.environ["SANDBOX_NAME"], + cwd=os.getenv("SANDBOX_CWD", "/home/sandbox-agent/work"), + ) as environment: + async with NativeProvider( + base_url=os.getenv("OPENAI_BASE_URL"), + api_key=os.environ["OPENAI_API_KEY"], + supports_websockets=False, + ) as provider: + agent = Agent( + instructions="Inspect the assigned workspace and report findings.", + model=os.getenv("OPENAI_MODEL", "gpt-5.5"), + provider=provider, + ) + async with agent.session(environment=environment) as session: + result = await (await session.submit(task)).get_result() + print(result.final_output) + + +if __name__ == "__main__": + asyncio.run(main(" ".join(sys.argv[1:]) or "Inspect the assigned workspace")) diff --git a/pyproject.toml b/pyproject.toml index f636d111ac..bbe08ac9d6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,14 @@ classifiers = [ "License :: OSI Approved :: MIT License", ] +[project.optional-dependencies] +databricks = [ + "databricks-sdk[sandbox]>=0.153.0", + "websockets>=13", + "cryptography>=42", + "certifi", +] + [project.urls] Homepage = "https://github.com/openai/openai-agents-python/tree/feat/codex-harness" Repository = "https://github.com/openai/openai-agents-python" @@ -140,6 +148,10 @@ disallow_incomplete_defs = false disallow_untyped_defs = false disallow_untyped_calls = false +[[tool.mypy.overrides]] +module = ["databricks.*"] +ignore_missing_imports = true + [tool.coverage.run] source = ["src/agents"] omit = ["tests/*"] diff --git a/src/agents/extensions/__init__.py b/src/agents/extensions/__init__.py new file mode 100644 index 0000000000..2a109d6445 --- /dev/null +++ b/src/agents/extensions/__init__.py @@ -0,0 +1 @@ +"""Optional, separately installed integrations for the Agents SDK.""" diff --git a/src/agents/extensions/sandbox/__init__.py b/src/agents/extensions/sandbox/__init__.py new file mode 100644 index 0000000000..d42d7c7190 --- /dev/null +++ b/src/agents/extensions/sandbox/__init__.py @@ -0,0 +1,6 @@ +"""Sandbox integrations that provision executors for the native ``Environment``. + +Each provider lives in its own subpackage and depends on an extra, so importing +``agents`` never requires a provider SDK. Import the provider explicitly, e.g. +``from agents.extensions.sandbox.databricks import databricks_sandbox_environment``. +""" diff --git a/src/agents/extensions/sandbox/databricks/README.md b/src/agents/extensions/sandbox/databricks/README.md new file mode 100644 index 0000000000..efd7cbf083 --- /dev/null +++ b/src/agents/extensions/sandbox/databricks/README.md @@ -0,0 +1,106 @@ +# Databricks sandbox backend + +Back a native `Environment` with a Databricks sandbox, so an agent driven by +this SDK executes `exec_command` / `write_stdin` / `apply_patch` inside the +sandbox. + +## How it works + +The native `Environment` dials an authenticated `wss://` executor and relays the +exec-server protocol over it. A Databricks sandbox runs that exec-server in +stdio mode; the `databricks-sdk` streaming command API +(`WorkspaceClient.sandbox.execute_command`) gives bidirectional stdio to it. +Nothing in the SDK speaks stdio and nothing in the sandbox speaks WebSocket, so +`SandboxExecServerGateway` bridges them: it terminates the WebSocket, +authenticates the bearer token, opens the exec-server's stdio, and relays bytes +both ways. It parses neither protocol. + +Because `execute_command` reaches the sandbox outbound through the Databricks +control plane, the gateway runs on `localhost` next to the application and the +sandbox needs no inbound ingress. The agent loop runs in your process (via +`NativeProvider`), which is what makes the `Environment` approval / transform +callbacks and `NetworkPolicy` available - the trade against the managed +Agents API path, where OpenAI runs the loop and the sandbox exec-server dials +out instead. + +## Install + +The v1.0 SDK is a GitHub preview wheel (off PyPI). Install it, then this extra: + +```sh +pip install "" +pip install "openai-agents[databricks]" # from this source tree +``` + +The extra pulls `databricks-sdk[sandbox]`, `websockets`, `cryptography`, and +`certifi`. + +## Usage + +```python +from databricks.sdk import WorkspaceClient +from agents import Agent, NativeProvider +from agents.extensions.sandbox.databricks import databricks_sandbox_environment + +async with databricks_sandbox_environment( + workspace_client=WorkspaceClient(profile="my-workspace"), + sandbox_name="openai-agent-sandbox", + cwd="/home/sandbox-agent/work", +) as environment: + async with NativeProvider(api_key=..., supports_websockets=False) as provider: + agent = Agent(model="gpt-5.5", provider=provider, instructions="...") + async with agent.session(environment=environment) as session: + print((await (await session.submit("...")).get_result()).final_output) +``` + +A runnable version is in `examples/databricks/runner.py`. + +## Provisioning the sandbox + +This package wires an `Environment` to an *existing* sandbox; it does not create +one, to avoid depending on sandbox-spec classes it cannot pin. Provision with +the usual API, then ensure the exec-server is reachable, e.g.: + +```python +from databricks.sdk import WorkspaceClient + +w = WorkspaceClient(profile="my-workspace") +w.sandbox.create_sandbox(..., sandbox_id="openai-agent-sandbox") +``` + +`exec-server` is killed when the sandbox stops, so a new connection after a +stop relaunches it. + +## TLS + +`Environment` requires `wss`, so the gateway needs TLS even on loopback. By +default `databricks_sandbox_environment` generates an ephemeral self-signed +certificate and, for the session, points `SSL_CERT_FILE` at a bundle of the +public CAs (`certifi`) plus that certificate, then restores it on exit - so the +native runtime trusts the gateway without breaking model-provider TLS. Pass your +own `ssl_context` and set `trust_gateway_ca=False` to manage trust yourself. + +## Unverified prerequisites (draft) + +Two assumptions this draft is built on still need confirmation against a live +sandbox before it can be relied on: + +1. **exec-server stdio mode in the image.** The default command is + `("codex", "exec-server", "stdio")`. The SDK's own `offline_docker.py` + example proves `exec-server stdio` is valid; the Databricks cookbook only + exercised the outbound `codex exec-server --remote` mode. Confirm with + `codex exec-server --help` in a sandbox, and override `exec_server_command` + if the binary path or subcommand differ. +2. **Streaming `execute_command` suitability.** The relay assumes + `CommandProcess` gives a usable bidirectional binary stdio stream (`stdin` + writes, `stdout.read1`) for the exec-server protocol. The v0.x provider used + only the buffered `execute_command_sync`, so this streaming path has not been + exercised for this purpose. Validate interactive `write_stdin` under load and + how `CommandProcess` reconnect interacts with the exec-server session. +3. **WebSocket framing.** The gateway relays an opaque byte stream: it forwards + each inbound frame to stdin as-is and sends arbitrary `read1` chunks of stdout + back as frames. This is correct only if the native client reassembles the + byte stream rather than requiring one protocol message per frame (the + reference `offline_docker.py` bridges the socket straight to byte pipes, which + indicates byte-stream semantics). If the executor is message-framed, the relay + must preserve message boundaries instead. diff --git a/src/agents/extensions/sandbox/databricks/__init__.py b/src/agents/extensions/sandbox/databricks/__init__.py new file mode 100644 index 0000000000..c958873118 --- /dev/null +++ b/src/agents/extensions/sandbox/databricks/__init__.py @@ -0,0 +1,17 @@ +"""Databricks sandbox backend for the native ``Environment``. + +Install with the ``databricks`` extra. The application provisions a Databricks +sandbox whose image carries the ``codex`` exec-server; this package runs a +localhost WSS gateway that relays that exec-server's stdio to the SDK, and +builds the ``Environment`` pointing at it. +""" + +from .environment import databricks_sandbox_environment +from .gateway import GatewayHandle, SandboxExecServerGateway, self_signed_server_context + +__all__ = [ + "databricks_sandbox_environment", + "GatewayHandle", + "SandboxExecServerGateway", + "self_signed_server_context", +] diff --git a/src/agents/extensions/sandbox/databricks/environment.py b/src/agents/extensions/sandbox/databricks/environment.py new file mode 100644 index 0000000000..ee790601cf --- /dev/null +++ b/src/agents/extensions/sandbox/databricks/environment.py @@ -0,0 +1,135 @@ +"""Provision a native ``Environment`` backed by a Databricks sandbox. + +This async context manager starts a :class:`SandboxExecServerGateway` for an +existing sandbox and yields an ``Environment`` wired to it. The SDK's +``EnvironmentSource`` accepts an async context manager, so the whole +start-gateway / tear-down lifecycle fits one ``async with``: + + async with databricks_sandbox_environment( + workspace_client=WorkspaceClient(profile="my-workspace"), + sandbox_name="openai-agent-sandbox", + cwd="/home/sandbox-agent/work", + ) as environment: + async with agent.session(environment=environment) as session: + ... + +Provisioning the sandbox itself (``WorkspaceClient.sandbox.create_sandbox``) and +ensuring the ``codex`` exec-server is present are left to the caller, so this +code takes no dependency on sandbox-spec classes it cannot pin. See the package +README for the create/ensure pattern. +""" + +from __future__ import annotations + +import os +import secrets +import ssl +import tempfile +from collections.abc import AsyncIterator, Callable, Sequence +from contextlib import asynccontextmanager +from typing import TYPE_CHECKING + +from ....executor import Environment, ExecutorToolName +from .gateway import SandboxExecServerGateway, self_signed_server_context + +if TYPE_CHECKING: + from databricks.sdk import WorkspaceClient + +_DEFAULT_TOOLS: tuple[ExecutorToolName, ...] = ("exec_command", "write_stdin", "apply_patch") + + +@asynccontextmanager +async def databricks_sandbox_environment( + *, + workspace_client: WorkspaceClient, + sandbox_name: str, + cwd: str, + tools: Sequence[ExecutorToolName] = _DEFAULT_TOOLS, + exec_server_command: Sequence[str] = ("codex", "exec-server", "stdio"), + token: str | None = None, + ssl_context: ssl.SSLContext | None = None, + trust_gateway_ca: bool = True, +) -> AsyncIterator[Environment]: + """Yield an ``Environment`` connected to ``sandbox_name`` through a gateway. + + Args: + workspace_client: An authenticated ``databricks.sdk.WorkspaceClient``. + sandbox_name: An existing, running sandbox with the exec-server available. + cwd: Absolute remote workspace path the executor operates in. + tools: Native executor tools to expose. ``write_stdin`` requires + ``exec_command``. + exec_server_command: Argv that starts the exec-server in stdio mode in + the sandbox. The default assumes the image's ``codex`` binary serves + ``exec-server stdio`` - confirm this before relying on it. + token: Bearer secret shared between the gateway and the SDK. A random + value is generated when omitted. + ssl_context: Server TLS context for the gateway. A self-signed loopback + certificate is generated when omitted. + trust_gateway_ca: When a self-signed context is generated, add its CA to + a combined bundle and point ``SSL_CERT_FILE`` at it for the session, + so the native runtime trusts the gateway while keeping public CAs + (via ``certifi``) for the model provider. Restored on exit. Disable + if you manage trust yourself. + """ + secret = token or secrets.token_urlsafe(32) + generated_ca: str | None = None + if ssl_context is None: + ssl_context, generated_ca = self_signed_server_context() + + gateway = SandboxExecServerGateway( + workspace_client=workspace_client, + sandbox_name=sandbox_name, + token=secret, + ssl_context=ssl_context, + exec_server_command=exec_server_command, + ca_cert_path=generated_ca, + ) + handle = await gateway.start() + + restore: Callable[[], None] | None = None + if trust_gateway_ca and handle.ca_cert_path is not None: + restore = _install_trust_bundle(handle.ca_cert_path) + try: + yield Environment( + url=handle.url, + cwd=cwd, + tools=tuple(tools), + headers={"Authorization": f"Bearer {secret}"}, + ) + finally: + if restore is not None: + restore() + await gateway.aclose() + + +def _install_trust_bundle(ca_cert_path: str) -> Callable[[], None]: + """Point ``SSL_CERT_FILE`` at ``certifi`` plus the gateway CA; return a restorer. + + The native runtime shares one trust store for the model provider and the + executor connection, so the gateway CA is appended to the public bundle + rather than replacing it. + """ + import certifi + + with open(certifi.where(), "rb") as public_bundle: + public = public_bundle.read() + with open(ca_cert_path, "rb") as gateway_ca: + extra = gateway_ca.read() + + combined = tempfile.NamedTemporaryFile(prefix="dbx-trust-", suffix=".pem", delete=False) + with combined: + combined.write(public) + if not public.endswith(b"\n"): + combined.write(b"\n") + combined.write(extra) + + previous = os.environ.get("SSL_CERT_FILE") + os.environ["SSL_CERT_FILE"] = combined.name + + def restore() -> None: + if previous is None: + os.environ.pop("SSL_CERT_FILE", None) + else: + os.environ["SSL_CERT_FILE"] = previous + + return restore diff --git a/src/agents/extensions/sandbox/databricks/gateway.py b/src/agents/extensions/sandbox/databricks/gateway.py new file mode 100644 index 0000000000..d720cd68a3 --- /dev/null +++ b/src/agents/extensions/sandbox/databricks/gateway.py @@ -0,0 +1,239 @@ +"""A localhost WSS gateway that bridges the SDK executor to a sandbox exec-server. + +The native ``Environment`` dials an authenticated ``wss://`` endpoint and relays +the exec-server wire protocol over it. A Databricks sandbox runs that +exec-server in stdio mode, reached through the ``databricks-sdk`` streaming +command API (``WorkspaceClient.sandbox.execute_command``). Nothing in the SDK +speaks stdio and nothing in the sandbox speaks WebSocket, so this gateway sits +between them: it terminates the WebSocket, authenticates the bearer token, +opens the exec-server's stdio stream, and relays bytes in both directions. + +The gateway parses neither protocol. It moves bytes between the WebSocket frames +and the exec-server's stdin/stdout, so it stays valid as the exec-server +protocol evolves. This assumes the native WebSocket client treats the +connection as an opaque byte stream and reassembles messages itself, rather than +requiring one protocol message per WebSocket frame; confirm this framing against +the executor before relying on it (see the package README). Because +``execute_command`` reaches the sandbox outbound +through the Databricks control plane, the gateway can run on ``localhost`` next +to the application; the sandbox needs no inbound ingress. +""" + +from __future__ import annotations + +import asyncio +import hmac +import http +import logging +import ssl +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from databricks.sdk import WorkspaceClient + from websockets.asyncio.server import ServerConnection + from websockets.http11 import Request, Response + +logger = logging.getLogger("agents.extensions.sandbox.databricks") + +_READ_CHUNK = 64 * 1024 +_DEFAULT_EXEC_SERVER_COMMAND: tuple[str, ...] = ("codex", "exec-server", "stdio") + + +@dataclass(frozen=True, kw_only=True) +class GatewayHandle: + """A running gateway. + + Attributes: + url: The ``wss://`` endpoint to pass as ``Environment.url``. + ca_cert_path: Path to the CA that signed the gateway certificate, or + None when the caller supplied its own already-trusted context. Add + it to the trust store the native runtime uses (see + ``environment.databricks_sandbox_environment``). + """ + + url: str + ca_cert_path: str | None + + +class SandboxExecServerGateway: + """Serve one sandbox's exec-server over an authenticated localhost WebSocket. + + Args: + workspace_client: An authenticated ``databricks.sdk.WorkspaceClient``. + sandbox_name: The target sandbox name/id. + token: Shared secret the SDK must present as ``Authorization: Bearer``. + The provisioner puts the same value in ``Environment.headers``. + ssl_context: Server-side TLS context. ``Environment`` requires ``wss``, + so TLS is mandatory even on loopback. Use + ``self_signed_server_context`` for local runs or pass your own. + exec_server_command: Argv that starts the exec-server in stdio mode + inside the sandbox. PREREQUISITE: the sandbox image's ``codex`` + binary must support this; confirm with ``codex exec-server --help`` + in a live sandbox. Override if the binary path or subcommand differ. + host: Loopback interface to bind. + ca_cert_path: Path to the CA for ``ssl_context`` when self-signed, echoed + back on the handle so callers can trust it. + """ + + def __init__( + self, + *, + workspace_client: WorkspaceClient, + sandbox_name: str, + token: str, + ssl_context: ssl.SSLContext, + exec_server_command: Sequence[str] = _DEFAULT_EXEC_SERVER_COMMAND, + host: str = "127.0.0.1", + ca_cert_path: str | None = None, + ) -> None: + if not token: + raise ValueError("Gateway token must be a non-empty secret") + command = tuple(exec_server_command) + if not command: + raise ValueError("exec_server_command must name the exec-server program") + self._w = workspace_client + self._sandbox_name = sandbox_name + self._token = token + self._ssl_context = ssl_context + self._command = command + self._host = host + self._ca_cert_path = ca_cert_path + self._server: Any = None + + async def start(self) -> GatewayHandle: + """Bind the WebSocket server on an ephemeral port and start serving.""" + from websockets.asyncio.server import serve + + self._server = await serve( + self._handle, + self._host, + 0, + ssl=self._ssl_context, + process_request=self._process_request, + ) + port = self._server.sockets[0].getsockname()[1] + return GatewayHandle(url=f"wss://{self._host}:{port}", ca_cert_path=self._ca_cert_path) + + async def aclose(self) -> None: + """Stop accepting connections and wait for the server to close.""" + if self._server is not None: + self._server.close() + await self._server.wait_closed() + self._server = None + + def _process_request(self, connection: ServerConnection, request: Request) -> Response | None: + """Reject an unauthenticated client at the HTTP handshake with 401.""" + presented = request.headers.get("Authorization", "") + if hmac.compare_digest(presented, f"Bearer {self._token}"): + return None + return connection.respond(http.HTTPStatus.UNAUTHORIZED, "unauthorized\n") + + async def _handle(self, connection: ServerConnection) -> None: + proc = await asyncio.to_thread(self._open_exec_server) + inbound = asyncio.create_task(self._ws_to_stdin(connection, proc)) + outbound = asyncio.create_task(self._stdout_to_ws(proc, connection)) + try: + await asyncio.wait({inbound, outbound}, return_when=asyncio.FIRST_COMPLETED) + finally: + # Closing the process unblocks the stdout read thread so both tasks finish. + await asyncio.to_thread(self._close_quietly, proc) + inbound.cancel() + outbound.cancel() + await asyncio.gather(inbound, outbound, return_exceptions=True) + + def _open_exec_server(self) -> Any: + return self._w.sandbox.execute_command( + self._sandbox_name, + self._command[0], + args=list(self._command[1:]), + text=False, + ) + + @staticmethod + def _close_quietly(proc: Any) -> None: + try: + proc.close() + except Exception: # teardown must not mask the real error + logger.debug("exec-server process close raised during teardown", exc_info=True) + + async def _ws_to_stdin(self, connection: ServerConnection, proc: Any) -> None: + async for message in connection: + data = message if isinstance(message, bytes) else message.encode() + await asyncio.to_thread(proc.stdin.write, data) + flush = getattr(proc.stdin, "flush", None) + if flush is not None: + await asyncio.to_thread(flush) + await asyncio.to_thread(self._close_stdin_quietly, proc) + + @staticmethod + def _close_stdin_quietly(proc: Any) -> None: + try: + proc.stdin.close() + except Exception: # half-closed stream during teardown + logger.debug("exec-server stdin close raised", exc_info=True) + + async def _stdout_to_ws(self, proc: Any, connection: ServerConnection) -> None: + while True: + chunk = await asyncio.to_thread(proc.stdout.read1, _READ_CHUNK) + if not chunk: + break + await connection.send(chunk) + + +def self_signed_server_context(host: str = "127.0.0.1") -> tuple[ssl.SSLContext, str]: + """Build a TLS server context with an ephemeral self-signed certificate. + + Returns the server context and the path to the generated CA certificate so + the client side can trust it. Intended for local runs; supply your own + context and CA in production. Requires the ``cryptography`` dependency from + the ``databricks`` extra. + + Args: + host: The loopback host the gateway binds, embedded as a SAN. + """ + import datetime + import ipaddress + import tempfile + + from cryptography import x509 + from cryptography.hazmat.primitives import hashes, serialization + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.x509.oid import NameOID + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, host)]) + try: + san: x509.GeneralName = x509.IPAddress(ipaddress.ip_address(host)) + except ValueError: + san = x509.DNSName(host) + now = datetime.datetime.now(datetime.timezone.utc) + certificate = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(minutes=5)) + .not_valid_after(now + datetime.timedelta(days=1)) + .add_extension(x509.SubjectAlternativeName([san]), critical=False) + .sign(key, hashes.SHA256()) + ) + + cert_pem = certificate.public_bytes(serialization.Encoding.PEM) + key_pem = key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.TraditionalOpenSSL, + serialization.NoEncryption(), + ) + cert_file = tempfile.NamedTemporaryFile(prefix="dbx-gw-", suffix=".crt", delete=False) + key_file = tempfile.NamedTemporaryFile(prefix="dbx-gw-", suffix=".key", delete=False) + with cert_file: + cert_file.write(cert_pem) + with key_file: + key_file.write(key_pem) + + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(cert_file.name, key_file.name) + return context, cert_file.name diff --git a/tests/extensions/sandbox/databricks/test_databricks_environment.py b/tests/extensions/sandbox/databricks/test_databricks_environment.py new file mode 100644 index 0000000000..eb76458f6e --- /dev/null +++ b/tests/extensions/sandbox/databricks/test_databricks_environment.py @@ -0,0 +1,164 @@ +"""Tests for the Databricks sandbox environment context manager. + +Cover the public entry point (Environment wiring, teardown, trust handling) and +the ``SSL_CERT_FILE`` bundle helper. Require the ``databricks`` extra's +``websockets``, ``cryptography`` and ``certifi``. +""" + +from __future__ import annotations + +import os +import queue +import ssl +from pathlib import Path + +import pytest + +pytest.importorskip("websockets") +pytest.importorskip("cryptography") +pytest.importorskip("certifi") + +from websockets.asyncio.client import connect # noqa: E402 + +from agents.extensions.sandbox.databricks.environment import ( # noqa: E402 + _install_trust_bundle, + databricks_sandbox_environment, +) + +_FAKE_CA = b"-----BEGIN CERTIFICATE-----\nFAKE\n-----END CERTIFICATE-----\n" + + +class _FakeStdin: + def __init__(self) -> None: + self.written = bytearray() + + def write(self, data: bytes) -> None: + self.written.extend(data) + + def flush(self) -> None: + pass + + def close(self) -> None: + pass + + +class _FakeStdout: + def __init__(self) -> None: + self._chunks: queue.Queue[bytes | None] = queue.Queue() + + def eof(self) -> None: + self._chunks.put(None) + + def read1(self, _size: int) -> bytes: + item = self._chunks.get() + return b"" if item is None else item + + +class _FakeProcess: + def __init__(self) -> None: + self.stdin = _FakeStdin() + self.stdout = _FakeStdout() + + def close(self) -> None: + self.stdout.eof() + + +class _FakeSandbox: + def __init__(self, proc: _FakeProcess) -> None: + self._proc = proc + + def execute_command(self, name: str, program: str, args: list[str], text: bool) -> _FakeProcess: + return self._proc + + +class _FakeWorkspaceClient: + def __init__(self, proc: _FakeProcess) -> None: + self.sandbox = _FakeSandbox(proc) + + +async def test_environment_wiring(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("SSL_CERT_FILE", raising=False) + client = _FakeWorkspaceClient(_FakeProcess()) + async with databricks_sandbox_environment( + workspace_client=client, + sandbox_name="sbx", + cwd="/home/sandbox-agent/work", + token="tok-123", + ) as environment: + assert environment.url.startswith("wss://127.0.0.1:") + assert environment.cwd == "/home/sandbox-agent/work" + # Environment lowercases header names. + assert environment.headers["authorization"] == "Bearer tok-123" + assert tuple(environment.tools) == ("exec_command", "write_stdin", "apply_patch") + assert "SSL_CERT_FILE" in os.environ + assert "SSL_CERT_FILE" not in os.environ + + +async def test_trust_gateway_ca_false_leaves_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("SSL_CERT_FILE", raising=False) + client = _FakeWorkspaceClient(_FakeProcess()) + async with databricks_sandbox_environment( + workspace_client=client, + sandbox_name="sbx", + cwd="/w", + token="t", + trust_gateway_ca=False, + ): + assert "SSL_CERT_FILE" not in os.environ + assert "SSL_CERT_FILE" not in os.environ + + +async def test_restores_and_closes_on_exception(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SSL_CERT_FILE", "/orig") + client = _FakeWorkspaceClient(_FakeProcess()) + captured: dict[str, str] = {} + with pytest.raises(RuntimeError, match="boom"): + async with databricks_sandbox_environment( + workspace_client=client, + sandbox_name="sbx", + cwd="/w", + token="t", + ) as environment: + captured["url"] = environment.url + raise RuntimeError("boom") + + assert os.environ["SSL_CERT_FILE"] == "/orig" + + context = ssl.create_default_context() + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + with pytest.raises(OSError): + async with connect(captured["url"], ssl=context): + pass + + +def test_install_trust_bundle_sets_and_restores( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + import certifi + + ca = tmp_path / "ca.pem" + ca.write_bytes(_FAKE_CA) + monkeypatch.delenv("SSL_CERT_FILE", raising=False) + + restore = _install_trust_bundle(str(ca)) + bundle = Path(os.environ["SSL_CERT_FILE"]).read_bytes() + assert Path(certifi.where()).read_bytes() in bundle + assert _FAKE_CA in bundle + + restore() + assert "SSL_CERT_FILE" not in os.environ + + +def test_install_trust_bundle_restores_previous( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + ca = tmp_path / "ca.pem" + ca.write_bytes(_FAKE_CA) + monkeypatch.setenv("SSL_CERT_FILE", "/x") + + restore = _install_trust_bundle(str(ca)) + assert os.environ["SSL_CERT_FILE"] != "/x" + + restore() + assert os.environ["SSL_CERT_FILE"] == "/x" diff --git a/tests/extensions/sandbox/databricks/test_databricks_gateway.py b/tests/extensions/sandbox/databricks/test_databricks_gateway.py new file mode 100644 index 0000000000..c1208283e6 --- /dev/null +++ b/tests/extensions/sandbox/databricks/test_databricks_gateway.py @@ -0,0 +1,203 @@ +"""Relay and auth tests for the Databricks sandbox gateway. + +These exercise the certain core - byte relay and bearer auth - with a fake +``CommandProcess`` and a real TLS WebSocket client. They require the +``databricks`` extra's ``websockets``, ``cryptography`` and ``certifi``. +""" + +from __future__ import annotations + +import queue +import ssl +from typing import Any + +import pytest + +pytest.importorskip("websockets") +pytest.importorskip("cryptography") + +from websockets.asyncio.client import connect # noqa: E402 +from websockets.exceptions import ConnectionClosed, InvalidStatus # noqa: E402 + +from agents.extensions.sandbox.databricks.gateway import ( # noqa: E402 + SandboxExecServerGateway, + self_signed_server_context, +) + + +class _FakeStdin: + def __init__(self) -> None: + self.written = bytearray() + self.closed = False + + def write(self, data: bytes) -> None: + self.written.extend(data) + + def flush(self) -> None: + pass + + def close(self) -> None: + self.closed = True + + +class _FakeStdout: + """Blocking stdout fed from the test; ``read1`` returns b'' after EOF.""" + + def __init__(self) -> None: + self._chunks: queue.Queue[bytes | None] = queue.Queue() + + def feed(self, data: bytes) -> None: + self._chunks.put(data) + + def eof(self) -> None: + self._chunks.put(None) + + def read1(self, _size: int) -> bytes: + item = self._chunks.get() + return b"" if item is None else item + + +class _FakeProcess: + def __init__(self) -> None: + self.stdin = _FakeStdin() + self.stdout = _FakeStdout() + self.closed = False + + def close(self) -> None: + self.closed = True + self.stdout.eof() + + +class _FakeSandbox: + def __init__(self, proc: _FakeProcess) -> None: + self._proc = proc + self.calls: list[tuple[str, str, list[str]]] = [] + + def execute_command(self, name: str, program: str, args: list[str], text: bool) -> _FakeProcess: + assert text is False + self.calls.append((name, program, args)) + return self._proc + + +class _FakeWorkspaceClient: + def __init__(self, proc: _FakeProcess) -> None: + self.sandbox = _FakeSandbox(proc) + + +def _client_ssl(ca_path: str) -> ssl.SSLContext: + context = ssl.create_default_context(cafile=ca_path) + return context + + +async def _start(proc: _FakeProcess, token: str = "s3cret") -> tuple[SandboxExecServerGateway, Any]: + server_ctx, ca_path = self_signed_server_context("127.0.0.1") + gateway = SandboxExecServerGateway( + workspace_client=_FakeWorkspaceClient(proc), + sandbox_name="sbx", + token=token, + ssl_context=server_ctx, + exec_server_command=("codex", "exec-server", "stdio"), + ca_cert_path=ca_path, + ) + handle = await gateway.start() + return gateway, handle + + +async def test_relay_both_directions() -> None: + proc = _FakeProcess() + client = _FakeWorkspaceClient(proc) + server_ctx, ca_path = self_signed_server_context("127.0.0.1") + gateway = SandboxExecServerGateway( + workspace_client=client, + sandbox_name="sbx", + token="s3cret", + ssl_context=server_ctx, + exec_server_command=("codex", "exec-server", "stdio"), + ca_cert_path=ca_path, + ) + handle = await gateway.start() + try: + async with connect( + handle.url, + ssl=_client_ssl(ca_path), + additional_headers={"Authorization": "Bearer s3cret"}, + ) as ws: + await ws.send(b"hello-stdin") + proc.stdout.feed(b"hello-stdout") + received = await ws.recv() + assert received == b"hello-stdout" + finally: + await gateway.aclose() + + assert proc.stdin.written == b"hello-stdin" + assert proc.closed is True + assert proc.stdin.closed is True + assert client.sandbox.calls[0] == ("sbx", "codex", ["exec-server", "stdio"]) + + +async def test_rejects_wrong_token() -> None: + proc = _FakeProcess() + gateway, handle = await _start(proc) + try: + with pytest.raises(InvalidStatus): + async with connect( + handle.url, + ssl=_client_ssl(handle.ca_cert_path or ""), + additional_headers={"Authorization": "Bearer wrong"}, + ): + pass + finally: + await gateway.aclose() + + assert proc.stdin.written == b"" + + +class _RaisingSandbox: + def execute_command(self, name: str, program: str, args: list[str], text: bool) -> Any: + raise RuntimeError("exec-server unavailable") + + +class _RaisingWorkspaceClient: + def __init__(self) -> None: + self.sandbox = _RaisingSandbox() + + +async def test_closes_websocket_when_exec_server_exits() -> None: + proc = _FakeProcess() + gateway, handle = await _start(proc) + try: + async with connect( + handle.url, + ssl=_client_ssl(handle.ca_cert_path or ""), + additional_headers={"Authorization": "Bearer s3cret"}, + ) as ws: + proc.stdout.eof() # exec-server exits first + with pytest.raises(ConnectionClosed): + await ws.recv() + finally: + await gateway.aclose() + + assert proc.closed is True + + +async def test_connection_fails_when_exec_server_cannot_start() -> None: + server_ctx, ca_path = self_signed_server_context("127.0.0.1") + gateway = SandboxExecServerGateway( + workspace_client=_RaisingWorkspaceClient(), + sandbox_name="sbx", + token="s3cret", + ssl_context=server_ctx, + ca_cert_path=ca_path, + ) + handle = await gateway.start() + try: + for _ in range(2): # the gateway keeps serving after a failed connection + async with connect( + handle.url, + ssl=_client_ssl(ca_path), + additional_headers={"Authorization": "Bearer s3cret"}, + ) as ws: + with pytest.raises(ConnectionClosed): + await ws.recv() + finally: + await gateway.aclose()