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 and reroutes through the gateway's ProxyInference gRPC
endpoint, which forwards to the policy-allowed inference route (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_rerouted_through_gateway(
    sandbox: Callable[..., Sandbox],
    mock_inference_route: str,
) -> None:
    """Inference call to undeclared endpoint is intercepted and rerouted.

    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 through the gateway's ProxyInference RPC
    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()