OpenShell/e2e/python/test_inference_routing.py

"""E2E tests for inference interception and routing.

When a process inside the sandbox makes an inference API call (e.g. POST
/v1/chat/completions) to an endpoint not explicitly allowed by network policy,
the proxy intercepts it, TLS-terminates the connection, detects the inference
API pattern, and the sandbox routes the request locally to the configured
backend (configured with `mock://` for testing).
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from navigator._proto import datamodel_pb2, sandbox_pb2

if TYPE_CHECKING:
    from collections.abc import Callable

    from navigator import Sandbox


# =============================================================================
# Policy helpers
# =============================================================================

_BASE_FILESYSTEM = sandbox_pb2.FilesystemPolicy(
    include_workdir=True,
    read_only=["/usr", "/lib", "/etc", "/app", "/var/log"],
    read_write=["/sandbox", "/tmp"],
)
_BASE_LANDLOCK = sandbox_pb2.LandlockPolicy(compatibility="best_effort")
_BASE_PROCESS = sandbox_pb2.ProcessPolicy(run_as_user="sandbox", run_as_group="sandbox")


def _inference_routing_policy() -> sandbox_pb2.SandboxPolicy:
    """Policy with inference routing enabled.

    No network_policies needed — any connection from any binary to an endpoint
    not in an explicit policy will be intercepted for inference when
    allowed_routes is non-empty.
    """
    return sandbox_pb2.SandboxPolicy(
        version=1,
        inference=sandbox_pb2.InferencePolicy(allowed_routes=["e2e_mock_local"]),
        filesystem=_BASE_FILESYSTEM,
        landlock=_BASE_LANDLOCK,
        process=_BASE_PROCESS,
    )


# =============================================================================
# Tests
# =============================================================================


def test_inference_call_routed_to_backend(
    sandbox: Callable[..., Sandbox],
    mock_inference_route: str,
) -> None:
    """Inference call to undeclared endpoint is intercepted and routed.

    A Python process inside the sandbox calls the OpenAI chat completions
    endpoint via raw urllib. Since api.openai.com is not in any network
    policy, but inference routing is configured, the proxy should:
    1. Detect no explicit policy match (inspect_for_inference)
    2. TLS-terminate the connection
    3. Detect the inference API pattern (POST /v1/chat/completions)
    4. Forward locally via sandbox router to the policy-allowed backend
    5. Return the mock response from the configured route
    """
    spec = datamodel_pb2.SandboxSpec(policy=_inference_routing_policy())

    def call_chat_completions() -> str:
        import json
        import ssl
        import urllib.request

        body = json.dumps(
            {
                "model": "test-model",
                "messages": [{"role": "user", "content": "hello"}],
            }
        ).encode()

        req = urllib.request.Request(
            "https://api.openai.com/v1/chat/completions",
            data=body,
            headers={
                "Content-Type": "application/json",
                "Authorization": "Bearer dummy-key",
            },
            method="POST",
        )
        # The proxy will TLS-terminate, so we need to accept its cert.
        ctx = ssl.create_default_context()
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE

        resp = urllib.request.urlopen(req, timeout=30, context=ctx)
        return resp.read().decode()

    with sandbox(spec=spec, delete_on_exit=True) as sb:
        result = sb.exec_python(call_chat_completions, timeout_seconds=60)
        assert result.exit_code == 0, f"stderr: {result.stderr}"
        output = result.stdout.strip()
        assert "Hello from navigator mock backend" in output
        assert "mock/test-model" in output


def test_non_inference_request_denied(
    sandbox: Callable[..., Sandbox],
    mock_inference_route: str,
) -> None:
    """Non-inference HTTP request on an intercepted connection is denied.

    A process making a non-inference request (e.g. GET /v1/models) to an
    undeclared endpoint should be denied with 403 when inference routing
    is configured — only recognized inference API patterns are routed.
    """
    spec = datamodel_pb2.SandboxSpec(policy=_inference_routing_policy())

    def make_non_inference_request() -> str:
        import ssl
        import urllib.request

        ctx = ssl.create_default_context()
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE

        try:
            req = urllib.request.Request("https://api.openai.com/v1/models")
            urllib.request.urlopen(req, timeout=10, context=ctx)
            return "unexpected_success"
        except urllib.error.HTTPError as e:
            return f"http_error_{e.code}"
        except Exception as e:
            return f"error: {e}"

    with sandbox(spec=spec, delete_on_exit=True) as sb:
        result = sb.exec_python(make_non_inference_request, timeout_seconds=30)
        assert result.exit_code == 0, f"stderr: {result.stderr}"
        assert "403" in result.stdout.strip()


def test_inference_anthropic_messages_protocol(
    sandbox: Callable[..., Sandbox],
    mock_anthropic_route: str,
) -> None:
    """Anthropic messages protocol (POST /v1/messages) is intercepted and routed.

    Verifies multi-protocol routing: a request using the Anthropic messages
    format is correctly detected and forwarded to a route configured with
    the anthropic_messages protocol.
    """
    policy = sandbox_pb2.SandboxPolicy(
        version=1,
        inference=sandbox_pb2.InferencePolicy(
            allowed_routes=["e2e_mock_anthropic"],
        ),
        filesystem=_BASE_FILESYSTEM,
        landlock=_BASE_LANDLOCK,
        process=_BASE_PROCESS,
    )
    spec = datamodel_pb2.SandboxSpec(policy=policy)

    def call_anthropic_messages() -> str:
        import json
        import ssl
        import urllib.request

        body = json.dumps(
            {
                "model": "claude-test",
                "max_tokens": 64,
                "messages": [{"role": "user", "content": "hello"}],
            }
        ).encode()

        req = urllib.request.Request(
            "https://api.anthropic.com/v1/messages",
            data=body,
            headers={
                "Content-Type": "application/json",
                "x-api-key": "dummy-key",
                "anthropic-version": "2023-06-01",
            },
            method="POST",
        )
        ctx = ssl.create_default_context()
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE

        resp = urllib.request.urlopen(req, timeout=30, context=ctx)
        return resp.read().decode()

    with sandbox(spec=spec, delete_on_exit=True) as sb:
        result = sb.exec_python(call_anthropic_messages, timeout_seconds=60)
        assert result.exit_code == 0, f"stderr: {result.stderr}"
        output = result.stdout.strip()
        assert "Hello from navigator mock backend" in output
        assert "mock/claude-test" in output


def test_inference_route_filtering_by_allowed_routes(
    sandbox: Callable[..., Sandbox],
    mock_inference_route: str,
    mock_disallowed_route: str,
) -> None:
    """Only routes in allowed_routes are available; others produce errors.

    Two routes exist (e2e_mock_local and e2e_mock_disallowed), but the
    policy only allows e2e_mock_local. A request that would match the
    allowed route should succeed, while inference requests that can't
    match any allowed route get an error from the sandbox router.
    """
    # Policy only allows e2e_mock_local, NOT e2e_mock_disallowed
    spec = datamodel_pb2.SandboxSpec(policy=_inference_routing_policy())

    def call_allowed_route() -> str:
        import json
        import ssl
        import urllib.request

        body = json.dumps(
            {
                "model": "test-model",
                "messages": [{"role": "user", "content": "hello"}],
            }
        ).encode()

        req = urllib.request.Request(
            "https://api.openai.com/v1/chat/completions",
            data=body,
            headers={
                "Content-Type": "application/json",
                "Authorization": "Bearer dummy",
            },
            method="POST",
        )
        ctx = ssl.create_default_context()
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE

        resp = urllib.request.urlopen(req, timeout=30, context=ctx)
        return resp.read().decode()

    with sandbox(spec=spec, delete_on_exit=True) as sb:
        result = sb.exec_python(call_allowed_route, timeout_seconds=60)
        assert result.exit_code == 0, f"stderr: {result.stderr}"
        output = result.stdout.strip()
        # The allowed route (e2e_mock_local) should serve the request
        assert "Hello from navigator mock backend" in output
        assert "mock/test-model" in output