mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-21 06:11:31 +08:00
Piotr Mlocek
34fd3cbf64
Closes #67 ## Summary Implements transparent inference interception and routing for sandboxed AI agents. The sandbox proxy intercepts outbound AI SDK calls (OpenAI, Anthropic) and reroutes them through the gateway to policy-controlled backends — enabling organizations to redirect inference traffic to local or self-hosted models without modifying agent code. **Decision model** — a tri-state OPA evaluation for every CONNECT request: 1. Binary + endpoint explicitly allowed in `network_policies` → **allow** (pass through) 2. Not explicitly allowed + `inference.allowed_routes` configured → **inspect for inference** (TLS intercept, detect API patterns, route through gateway) 3. Otherwise → **deny** No endpoint declarations or binary lists needed for inference routing. Just configure `inference.allowed_routes`. ## Key Changes ### Sandbox (interception) - **OPA policy**: New `network_action` Rego rule with three outcomes (`allow`, `inspect_for_inference`, `deny`). New `NetworkAction` enum replaces `PolicyDecision.allowed` bool for the proxy's main decision path. - **Proxy**: New `InspectForInference` path — TLS-terminates client, parses HTTP, detects inference API patterns (`POST /v1/chat/completions`, `/v1/completions`, `/v1/messages`), strips auth headers, forwards via gRPC. - **New module**: `l7/inference.rs` — `InferenceApiPattern`, `detect_inference_pattern()`, HTTP request/response parsing. - **gRPC client**: New `proxy_inference()` for sandbox→gateway forwarding. - **Sandbox init**: Creates OPA engine when inference is configured, even without `network_policies`. ### Gateway (dispatch) - **InferenceService**: `ProxyInference` RPC loads sandbox policy, resolves allowed routes, dispatches to router. Full CRUD for inference routes. - **Proto**: `InferenceRoute`, `InferenceRouteSpec`, `ProxyInferenceRequest/Response`, Inference gRPC service. ### Router (backend proxying) - **New crate**: `navigator-router` with `Router`, `proxy_with_candidates()`, protocol-based route selection, backend HTTP proxying with auth header rewriting. - **Mock support** for testing (`mock://` scheme). ### CLI - `nav inference create/update/delete/list` commands for route management. ### Python SDK - Updated protobuf bindings. Removed old `inference.py` client (replaced by transparent interception). ### Documentation - New `architecture/inference-routing.md` — end-to-end system documentation. - Updated `architecture/sandbox.md` — proxy, OPA, and source index sections. - Updated `architecture/README.md` — new subsystem overview and diagram. ## Addendum: Chunked Transfer Compatibility This branch now also fixes intercepted SDK requests that send chunked request bodies: - `inspect_for_inference` now accepts `Transfer-Encoding: chunked` and decodes chunked request bodies before forwarding to the gateway - Removed the prior `411 Length Required` response for chunked intercepted requests - Added request/response header sanitization for framing and hop-by-hop headers (`content-length`, `transfer-encoding`, `connection`, etc.) to keep forwarded requests and returned responses valid - Added unit tests for chunked parsing and header sanitization Note: this improves compatibility for streaming-style SDK request patterns; true token-by-token passthrough response streaming is still a separate follow-up. ## Minimal Policy for Inference Routing ```yaml inference: allowed_routes: - local ``` Any outgoing connection from a binary not explicitly allowed in `network_policies` will be intercepted and checked for inference API patterns. ## Test Plan - [x] `cargo test --workspace` — all tests pass - [x] `mise run pre-commit` — all checks pass - [x] E2E: OpenAI chat completions routed through gateway - [x] E2E: Anthropic messages routed through gateway - [x] E2E: Python OpenAI SDK from sandbox (`examples/inference/inference.py`) - [x] E2E test: `e2e/python/test_inference_routing.py`
107 lines
3.0 KiB
Python
107 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import time
|
|
from typing import TYPE_CHECKING
|
|
|
|
import grpc
|
|
import pytest
|
|
|
|
from navigator import InferenceRouteClient, Sandbox, SandboxClient
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable, Iterator
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def cluster_name() -> str | None:
|
|
return os.environ.get("NAVIGATOR_CLUSTER")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def sandbox_client(cluster_name: str | None) -> Iterator[SandboxClient]:
|
|
with SandboxClient.from_active_cluster(cluster=cluster_name) as client:
|
|
yield client
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def ensure_sandbox_persistence_ready(sandbox_client: SandboxClient) -> None:
|
|
for _ in range(60):
|
|
try:
|
|
sandbox_client.list_ids(limit=1)
|
|
return
|
|
except grpc.RpcError as exc:
|
|
details = exc.details() or ""
|
|
if exc.code() == grpc.StatusCode.UNAVAILABLE:
|
|
time.sleep(2)
|
|
continue
|
|
if (
|
|
exc.code() == grpc.StatusCode.INTERNAL
|
|
and "no such table: objects" in details
|
|
):
|
|
time.sleep(1)
|
|
continue
|
|
raise
|
|
|
|
pytest.fail(
|
|
"navigator-server persistence is not initialized (missing sqlite objects table); "
|
|
"redeploy the active cluster and rerun e2e sandbox tests"
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def sandbox(cluster_name: str | None) -> Callable[..., Sandbox]:
|
|
def _create(*, spec: object | None = None, delete_on_exit: bool = True) -> Sandbox:
|
|
return Sandbox(
|
|
cluster=cluster_name,
|
|
spec=spec,
|
|
delete_on_exit=delete_on_exit,
|
|
# The sandbox image is large (Python, Node.js, coding agents) so the
|
|
# first pod in the cluster may need extra time for the image pull.
|
|
ready_timeout_seconds=300.0,
|
|
)
|
|
|
|
return _create
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def inference_client(sandbox_client: SandboxClient) -> InferenceRouteClient:
|
|
return InferenceRouteClient.from_sandbox_client(sandbox_client)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def mock_inference_route(
|
|
inference_client: InferenceRouteClient,
|
|
) -> Iterator[str]:
|
|
name = "e2e-mock-local"
|
|
routing_hint = "e2e_mock_local"
|
|
# Clean up any leftover route from a previous run.
|
|
try:
|
|
inference_client.delete(name)
|
|
except grpc.RpcError:
|
|
pass
|
|
|
|
inference_client.create(
|
|
name=name,
|
|
routing_hint=routing_hint,
|
|
base_url="mock://e2e",
|
|
protocols=["openai_chat_completions"],
|
|
api_key="mock",
|
|
model_id="mock/test-model",
|
|
enabled=True,
|
|
)
|
|
yield name
|
|
try:
|
|
inference_client.delete(name)
|
|
except grpc.RpcError:
|
|
pass
|
|
|
|
|
|
@pytest.fixture
|
|
def run_python() -> Callable[[Sandbox, str], tuple[int, str, str]]:
|
|
def _run(sandbox: Sandbox, code: str) -> tuple[int, str, str]:
|
|
result = sandbox.exec(["python", "-c", code], timeout_seconds=20)
|
|
return result.exit_code, result.stdout, result.stderr
|
|
|
|
return _run
|