diff --git a/pyproject.toml b/pyproject.toml index 4d0122049f..c154dd407a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,7 @@ litellm = ["litellm>=1.83.0"] any-llm = ["any-llm-sdk>=1.11.0, <2; python_version >= '3.11'"] realtime = ["websockets>=15.0, <17"] sqlalchemy = ["SQLAlchemy>=2.0", "asyncpg>=0.29.0"] -encrypt = ["cryptography>=45.0, <46"] +encrypt = ["cryptography>=45.0, <46", "pynacl>=1.5, <1.6"] redis = ["redis>=7"] dapr = ["dapr>=1.16.0", "grpcio>=1.60.0"] mongodb = ["pymongo>=4.14"] @@ -86,6 +86,7 @@ dev = [ "fastapi >= 0.110.0, <1", "aiosqlite>=0.21.0", "cryptography>=45.0, <46", + "pynacl>=1.5, <1.6", "fakeredis>=2.31.3", "dapr>=1.14.0", "grpcio>=1.60.0", diff --git a/src/agents/__init__.py b/src/agents/__init__.py index e3bc96abff..b17b6dadc0 100644 --- a/src/agents/__init__.py +++ b/src/agents/__init__.py @@ -82,6 +82,7 @@ from .models.interface import Model, ModelProvider, ModelTracing from .models.multi_provider import MultiProvider from .models.openai_agent_registration import OpenAIAgentRegistrationConfig +from .models.openai_agent_runtime_auth import OpenAIAgentRuntimeAuthConfig from .models.openai_chatcompletions import OpenAIChatCompletionsModel from .models.openai_provider import OpenAIProvider from .models.openai_responses import ( @@ -352,6 +353,7 @@ def enable_verbose_stdout_logging(): "MultiProvider", "OpenAIProvider", "OpenAIAgentRegistrationConfig", + "OpenAIAgentRuntimeAuthConfig", "OpenAIResponsesModel", "OpenAIResponsesWSModel", "AgentOutputSchema", diff --git a/src/agents/models/openai_agent_runtime_auth.py b/src/agents/models/openai_agent_runtime_auth.py new file mode 100644 index 0000000000..6d96836039 --- /dev/null +++ b/src/agents/models/openai_agent_runtime_auth.py @@ -0,0 +1,380 @@ +from __future__ import annotations + +import asyncio +import base64 +import datetime as dt +import json +import os +from collections.abc import Sequence +from dataclasses import dataclass, replace +from typing import Any, Protocol + +from pydantic import BaseModel + +from ..exceptions import UserError +from ..model_settings import ModelSettings +from ..version import __version__ +from .openai_agent_registration import resolve_openai_harness_id_for_model_provider + +_ENV_RUNTIME_AUTH_ENABLED = "OPENAI_AGENT_RUNTIME_AUTH_ENABLED" +_ENV_AGENT_VERSION = "OPENAI_AGENT_VERSION" +_ENV_AGENT_RUNNING_LOCATION = "OPENAI_AGENT_RUNNING_LOCATION" +_ENV_AGENT_CAPABILITIES = "OPENAI_AGENT_CAPABILITIES" +_ENV_AGENT_RUNTIME_TTL = "OPENAI_AGENT_RUNTIME_TTL" +_DEFAULT_AGENT_HARNESS_ID = "agents-sdk-python" +_DEFAULT_RUNNING_LOCATION = "client" +_DEFAULT_CAPABILITIES = ("responsesapi",) +_AGENT_ASSERTION_SCHEME = "AgentAssertion" + + +@dataclass(frozen=True) +class OpenAIAgentRuntimeAuthConfig: + """Opt-in configuration for verified sandbox runtime attribution on Responses calls.""" + + agent_harness_id: str | None = None + """Stable registry or interface identifier for the agent runtime.""" + + agent_version: str | None = None + """Version of the running agent. Defaults to the installed Agents SDK version.""" + + running_location: str | None = None + """Logical location where the agent is running.""" + + capabilities: Sequence[str] | None = None + """Capabilities authorized for this agent runtime registration.""" + + ttl: int | None = None + """Optional runtime identity TTL in whole seconds.""" + + external_task_ref: str | None = None + """Optional caller-supplied task reference used to resolve a durable task id.""" + + enabled: bool = True + """Whether runtime auth should be enabled for this configuration.""" + + +@dataclass(frozen=True) +class ResolvedOpenAIAgentRuntimeAuthConfig: + agent_harness_id: str + agent_version: str + running_location: str + capabilities: tuple[str, ...] + ttl: int | None + external_task_ref: str | None + + +class _RegisterAgentResponse(BaseModel): + agent_runtime_id: str + + +class _RegisterTaskResponse(BaseModel): + encrypted_task_id: str + + +class _OpenAIClientWithPost(Protocol): + async def post( + self, + path: str, + *, + cast_to: type[Any], + body: object | None = None, + **kwargs: Any, + ) -> Any: ... + + +def resolve_openai_agent_runtime_auth_config( + config: OpenAIAgentRuntimeAuthConfig | None, + *, + model_provider: object, +) -> ResolvedOpenAIAgentRuntimeAuthConfig | None: + if config is None and not _env_bool(_ENV_RUNTIME_AUTH_ENABLED): + return None + if config is not None and not config.enabled: + return None + + agent_harness_id = _first_non_empty( + config.agent_harness_id if config is not None else None, + resolve_openai_harness_id_for_model_provider(model_provider), + _DEFAULT_AGENT_HARNESS_ID, + ) + agent_version = _first_non_empty( + config.agent_version if config is not None else None, + os.getenv(_ENV_AGENT_VERSION), + __version__, + ) + running_location = _first_non_empty( + config.running_location if config is not None else None, + os.getenv(_ENV_AGENT_RUNNING_LOCATION), + _DEFAULT_RUNNING_LOCATION, + ) + capabilities = _resolve_capabilities(config.capabilities if config is not None else None) + ttl = ( + config.ttl + if config is not None and config.ttl is not None + else _env_int(_ENV_AGENT_RUNTIME_TTL) + ) + external_task_ref = ( + _normalize_str(config.external_task_ref) + if config is not None and config.external_task_ref is not None + else None + ) + + return ResolvedOpenAIAgentRuntimeAuthConfig( + agent_harness_id=agent_harness_id, + agent_version=agent_version, + running_location=running_location, + capabilities=capabilities, + ttl=ttl, + external_task_ref=external_task_ref, + ) + + +class OpenAIAgentRuntimeAuthManager: + def __init__(self, config: ResolvedOpenAIAgentRuntimeAuthConfig) -> None: + self._config = config + self._lock = asyncio.Lock() + self._private_key: Any | None = None + self._agent_public_key: str | None = None + self._agent_runtime_id: str | None = None + self._task_id: str | None = None + + async def authorization_header(self, client: _OpenAIClientWithPost) -> str: + await self._ensure_registered(client) + if self._private_key is None or self._agent_runtime_id is None or self._task_id is None: + raise UserError("OpenAI agent runtime auth registration did not complete.") + + timestamp = _utc_timestamp() + signature = _sign_agent_assertion( + private_key=self._private_key, + agent_runtime_id=self._agent_runtime_id, + task_id=self._task_id, + timestamp=timestamp, + ) + assertion = _serialize_agent_assertion( + agent_runtime_id=self._agent_runtime_id, + task_id=self._task_id, + timestamp=timestamp, + signature=signature, + ) + return f"{_AGENT_ASSERTION_SCHEME} {assertion}" + + async def _ensure_registered(self, client: _OpenAIClientWithPost) -> None: + if self._agent_runtime_id is not None and self._task_id is not None: + return + + async with self._lock: + if self._agent_runtime_id is not None and self._task_id is not None: + return + + private_key, agent_public_key = _generate_agent_keypair() + register_agent_body: dict[str, object] = { + "abom": { + "agent_version": self._config.agent_version, + "agent_harness_id": self._config.agent_harness_id, + "running_location": self._config.running_location, + }, + "agent_public_key": agent_public_key, + "capabilities": list(self._config.capabilities), + } + if self._config.ttl is not None: + register_agent_body["ttl"] = self._config.ttl + + agent_response = await client.post( + "/agent/register", + cast_to=_RegisterAgentResponse, + body=register_agent_body, + ) + agent_runtime_id = agent_response.agent_runtime_id + + timestamp = _utc_timestamp() + task_body: dict[str, str] = { + "timestamp": timestamp, + "signature": _sign_task_registration( + private_key=private_key, + agent_runtime_id=agent_runtime_id, + timestamp=timestamp, + ), + } + if self._config.external_task_ref is not None: + task_body["external_task_ref"] = self._config.external_task_ref + + task_response = await client.post( + f"/agent/{agent_runtime_id}/task/register", + cast_to=_RegisterTaskResponse, + body=task_body, + ) + task_id = _decrypt_task_id( + encrypted_task_id=task_response.encrypted_task_id, + private_key=private_key, + ) + + self._private_key = private_key + self._agent_public_key = agent_public_key + self._agent_runtime_id = agent_runtime_id + self._task_id = task_id + + +def add_agent_assertion_header( + model_settings: ModelSettings, + *, + authorization_header: str, +) -> ModelSettings: + extra_headers = dict(model_settings.extra_headers or {}) + for header_name in extra_headers: + if header_name.lower() == "authorization": + raise UserError( + "Sandbox agent runtime auth cannot be combined with an explicit Authorization " + "header in ModelSettings.extra_headers." + ) + extra_headers["Authorization"] = authorization_header + return replace(model_settings, extra_headers=extra_headers) + + +def _first_non_empty(*values: str | None) -> str: + for value in values: + normalized = _normalize_str(value) + if normalized is not None: + return normalized + raise UserError("Expected at least one non-empty value.") + + +def _normalize_str(value: str | None) -> str | None: + if value is None: + return None + stripped = value.strip() + return stripped or None + + +def _env_bool(name: str) -> bool: + value = os.getenv(name) + return value is not None and value.strip().lower() in {"1", "true", "t", "yes", "y", "on"} + + +def _env_int(name: str) -> int | None: + value = os.getenv(name) + if value is None or not value.strip(): + return None + try: + return int(value) + except ValueError as exc: + raise UserError(f"{name} must be an integer.") from exc + + +def _resolve_capabilities(configured: Sequence[str] | None) -> tuple[str, ...]: + values: Sequence[str] + if configured is not None: + values = configured + else: + env_value = os.getenv(_ENV_AGENT_CAPABILITIES) + values = env_value.split(",") if env_value else _DEFAULT_CAPABILITIES + + capabilities = tuple(normalized for value in values if (normalized := _normalize_str(value))) + return capabilities or _DEFAULT_CAPABILITIES + + +def _utc_timestamp() -> str: + return dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z") + + +def _urlsafe_b64encode(data: bytes) -> str: + return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=") + + +def _serialize_agent_assertion( + *, + agent_runtime_id: str, + task_id: str, + timestamp: str, + signature: str, +) -> str: + payload = json.dumps( + { + "agent_runtime_id": agent_runtime_id, + "task_id": task_id, + "timestamp": timestamp, + "signature": signature, + }, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + return _urlsafe_b64encode(payload) + + +def _generate_agent_keypair() -> tuple[Any, str]: + serialization, Ed25519PrivateKey = _cryptography() + private_key = Ed25519PrivateKey.generate() + agent_public_key = private_key.public_key().public_bytes( + encoding=serialization.Encoding.OpenSSH, + format=serialization.PublicFormat.OpenSSH, + ) + return private_key, agent_public_key.decode("utf-8") + + +def _sign_payload(*, private_key: Any, payload: bytes) -> str: + return base64.b64encode(private_key.sign(payload)).decode("ascii") + + +def _sign_task_registration(*, private_key: Any, agent_runtime_id: str, timestamp: str) -> str: + return _sign_payload( + private_key=private_key, + payload=f"{agent_runtime_id}:{timestamp}".encode(), + ) + + +def _sign_agent_assertion( + *, + private_key: Any, + agent_runtime_id: str, + task_id: str, + timestamp: str, +) -> str: + return _sign_payload( + private_key=private_key, + payload=f"{agent_runtime_id}:{task_id}:{timestamp}".encode(), + ) + + +def _decrypt_task_id(*, encrypted_task_id: str, private_key: Any) -> str: + serialization, _ = _cryptography() + crypto_sign_ed25519_sk_to_curve25519, PrivateKey, SealedBox = _pynacl() + private_seed = private_key.private_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PrivateFormat.Raw, + encryption_algorithm=serialization.NoEncryption(), + ) + public_key = private_key.public_key().public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw, + ) + curve25519_private_key = PrivateKey( + crypto_sign_ed25519_sk_to_curve25519(private_seed + public_key) + ) + plaintext: bytes = SealedBox(curve25519_private_key).decrypt( + base64.b64decode(encrypted_task_id) + ) + return plaintext.decode("utf-8") + + +def _cryptography() -> tuple[Any, Any]: + try: + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey + except ImportError as exc: + raise UserError( + "Sandbox agent runtime auth requires cryptography. Install the encrypt extra with " + "`pip install 'openai-agents[encrypt]'`." + ) from exc + return serialization, Ed25519PrivateKey + + +def _pynacl() -> tuple[Any, Any, Any]: + try: + from nacl.bindings import crypto_sign_ed25519_sk_to_curve25519 + from nacl.public import PrivateKey, SealedBox + except ImportError as exc: + raise UserError( + "Sandbox agent runtime auth requires PyNaCl. Install the encrypt extra with " + "`pip install 'openai-agents[encrypt]'`." + ) from exc + return crypto_sign_ed25519_sk_to_curve25519, PrivateKey, SealedBox diff --git a/src/agents/run.py b/src/agents/run.py index 014271a5ea..da6c54be80 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -1213,6 +1213,7 @@ def _finalize_result(result: RunResult) -> RunResult: reasoning_item_id_policy=resolved_reasoning_item_id_policy, prompt_cache_key_resolver=prompt_cache_key_resolver, error_handlers=error_handlers, + sandbox_runtime=sandbox_runtime, ) ) @@ -1269,6 +1270,7 @@ def _finalize_result(result: RunResult) -> RunResult: reasoning_item_id_policy=resolved_reasoning_item_id_policy, prompt_cache_key_resolver=prompt_cache_key_resolver, error_handlers=error_handlers, + sandbox_runtime=sandbox_runtime, ) finally: attach_usage_to_span( diff --git a/src/agents/run_config.py b/src/agents/run_config.py index fcc9b01315..3674ee5b07 100644 --- a/src/agents/run_config.py +++ b/src/agents/run_config.py @@ -22,6 +22,7 @@ if TYPE_CHECKING: from .agent import Agent + from .models.openai_agent_runtime_auth import OpenAIAgentRuntimeAuthConfig from .run_context import RunContextWrapper from .sandbox.manifest import Manifest from .sandbox.session.base_sandbox_session import BaseSandboxSession @@ -199,6 +200,9 @@ class SandboxRunConfig: Use `SandboxArchiveLimits()` to enable SDK defaults. """ + agent_runtime_auth: OpenAIAgentRuntimeAuthConfig | None = None + """Opt-in verified runtime/task auth for sandbox-backed OpenAI Responses calls.""" + @dataclass class RunConfig: diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index 45f09c0fa0..3b01d75dae 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -1036,6 +1036,7 @@ async def _save_stream_items_without_count( reasoning_item_id_policy=resolved_reasoning_item_id_policy, prompt_cache_key_resolver=prompt_cache_key_resolver, error_handlers=error_handlers, + sandbox_runtime=sandbox_runtime, ) finally: attach_usage_to_span( @@ -1255,6 +1256,7 @@ async def run_single_turn_streamed( reasoning_item_id_policy: ReasoningItemIdPolicy | None = None, prompt_cache_key_resolver: PromptCacheKeyResolver | None = None, error_handlers: RunErrorHandlers[TContext] | None = None, + sandbox_runtime: SandboxRuntime[TContext] | None = None, ) -> SingleStepResult: """Run a single streamed turn and emit events as results arrive.""" public_agent = bindings.public_agent @@ -1448,6 +1450,12 @@ def _tool_search_fingerprint(raw_item: Any) -> str: else None ) model_settings = model_settings_with_prompt_cache_key(model_settings, prompt_cache_key) + if sandbox_runtime is not None: + model_settings = await sandbox_runtime.prepare_model_settings( + public_agent=public_agent, + model=model, + model_settings=model_settings, + ) async def rewind_model_request() -> None: items_to_rewind = session_items_to_rewind if session_items_to_rewind is not None else [] @@ -1722,6 +1730,7 @@ async def run_single_turn( reasoning_item_id_policy: ReasoningItemIdPolicy | None = None, prompt_cache_key_resolver: PromptCacheKeyResolver | None = None, error_handlers: RunErrorHandlers[TContext] | None = None, + sandbox_runtime: SandboxRuntime[TContext] | None = None, ) -> SingleStepResult: """Run a single non-streaming turn of the agent loop.""" public_agent = bindings.public_agent @@ -1776,6 +1785,7 @@ async def run_single_turn( session=session, session_items_to_rewind=session_items_to_rewind, prompt_cache_key_resolver=prompt_cache_key_resolver, + sandbox_runtime=sandbox_runtime, ) return await get_single_step_result_from_response( @@ -1811,6 +1821,7 @@ async def get_new_response( session: Session | None = None, session_items_to_rewind: list[TResponseInputItem] | None = None, prompt_cache_key_resolver: PromptCacheKeyResolver | None = None, + sandbox_runtime: SandboxRuntime[TContext] | None = None, ) -> ModelResponse: """Call the model and return the raw response, handling retries and hooks.""" public_agent = bindings.public_agent @@ -1872,6 +1883,12 @@ async def get_new_response( else None ) model_settings = model_settings_with_prompt_cache_key(model_settings, prompt_cache_key) + if sandbox_runtime is not None: + model_settings = await sandbox_runtime.prepare_model_settings( + public_agent=public_agent, + model=model, + model_settings=model_settings, + ) async def rewind_model_request() -> None: items_to_rewind = session_items_to_rewind if session_items_to_rewind is not None else [] diff --git a/src/agents/sandbox/runtime.py b/src/agents/sandbox/runtime.py index d273a54411..fd4d2d35ee 100644 --- a/src/agents/sandbox/runtime.py +++ b/src/agents/sandbox/runtime.py @@ -9,6 +9,12 @@ from ..agent import Agent from ..exceptions import UserError from ..items import TResponseInputItem +from ..model_settings import ModelSettings +from ..models.openai_agent_runtime_auth import ( + OpenAIAgentRuntimeAuthManager, + add_agent_assertion_header, + resolve_openai_agent_runtime_auth_config, +) from ..result import RunResult, RunResultStreaming from ..run_config import RunConfig from ..run_context import RunContextWrapper, TContext @@ -73,6 +79,16 @@ def __init__( ) -> None: self._sandbox_config = run_config.sandbox if run_config is not None else None self._run_config_model = run_config.model if run_config is not None else None + self._agent_runtime_auth_manager: OpenAIAgentRuntimeAuthManager | None = None + if run_config is not None and self._sandbox_config is not None: + runtime_auth_config = resolve_openai_agent_runtime_auth_config( + self._sandbox_config.agent_runtime_auth, + model_provider=run_config.model_provider, + ) + if runtime_auth_config is not None: + self._agent_runtime_auth_manager = OpenAIAgentRuntimeAuthManager( + runtime_auth_config + ) # The runner resolves this before constructing the runtime. It can be None only when # sandbox is disabled or tests instantiate the runtime directly. self._rollout_id = rollout_id @@ -120,6 +136,34 @@ def assert_agent_supported(self, agent: Agent[TContext]) -> None: if isinstance(agent, SandboxAgent) and self._sandbox_config is None: raise UserError("SandboxAgent execution requires `RunConfig(sandbox=...)`") + async def prepare_model_settings( + self, + *, + public_agent: Agent[TContext], + model: object, + model_settings: ModelSettings, + ) -> ModelSettings: + if self._agent_runtime_auth_manager is None or not isinstance(public_agent, SandboxAgent): + return model_settings + + # AgentAssertion is only supported for OpenAI Responses requests. + from ..models.openai_responses import OpenAIResponsesModel + + if not isinstance(model, OpenAIResponsesModel): + return model_settings + + get_client = getattr(model, "_get_client", None) + if not callable(get_client): + return model_settings + + authorization_header = await self._agent_runtime_auth_manager.authorization_header( + get_client() + ) + return add_agent_assertion_header( + model_settings, + authorization_header=authorization_header, + ) + async def enqueue_memory_result( self, result: RunResult | RunResultStreaming, diff --git a/tests/models/test_openai_agent_runtime_auth.py b/tests/models/test_openai_agent_runtime_auth.py new file mode 100644 index 0000000000..42c0d89cbb --- /dev/null +++ b/tests/models/test_openai_agent_runtime_auth.py @@ -0,0 +1,220 @@ +from __future__ import annotations + +import base64 +import json +from typing import Any + +import httpx +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey +from nacl.bindings import crypto_sign_ed25519_pk_to_curve25519 +from nacl.public import PublicKey, SealedBox + +from agents import ( + ModelSettings, + OpenAIAgentRuntimeAuthConfig, + OpenAIResponsesModel, + OpenAIResponsesWSModel, + RunConfig, +) +from agents.exceptions import UserError +from agents.models.openai_agent_runtime_auth import ( + OpenAIAgentRuntimeAuthManager, + add_agent_assertion_header, + resolve_openai_agent_runtime_auth_config, +) +from agents.sandbox import SandboxAgent, SandboxRunConfig +from agents.sandbox.runtime import SandboxRuntime + + +class _FakeOpenAIClient: + def __init__(self, *, task_id: str = "task_456") -> None: + self.task_id = task_id + self.calls: list[tuple[str, dict[str, Any]]] = [] + self.agent_public_key: str | None = None + self.base_url = httpx.URL("https://api.openai.com/v1/") + self.websocket_base_url = None + self.default_headers = {"Authorization": "Bearer sk-test", "X-Client": "1"} + self.default_query: dict[str, str] = {} + + async def post( + self, + path: str, + *, + cast_to: type[Any], + body: object | None = None, + **kwargs: Any, + ) -> Any: + _ = kwargs + assert isinstance(body, dict) + self.calls.append((path, body)) + if path == "/agent/register": + self.agent_public_key = str(body["agent_public_key"]) + return cast_to.model_validate({"agent_runtime_id": "agent_123"}) + if path == "/agent/agent_123/task/register": + assert self.agent_public_key is not None + return cast_to.model_validate( + {"encrypted_task_id": _encrypt_task_id(self.task_id, self.agent_public_key)} + ) + raise AssertionError(f"Unexpected path: {path}") + + +def test_runtime_auth_config_is_env_gated(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("OPENAI_AGENT_RUNTIME_AUTH_ENABLED", raising=False) + + assert resolve_openai_agent_runtime_auth_config(None, model_provider=object()) is None + + monkeypatch.setenv("OPENAI_AGENT_RUNTIME_AUTH_ENABLED", "true") + resolved = resolve_openai_agent_runtime_auth_config(None, model_provider=object()) + + assert resolved is not None + assert resolved.agent_harness_id == "agents-sdk-python" + assert resolved.running_location == "client" + assert resolved.capabilities == ("responsesapi",) + + +@pytest.mark.asyncio +async def test_runtime_auth_manager_registers_runtime_task_and_builds_assertion() -> None: + resolved = resolve_openai_agent_runtime_auth_config( + OpenAIAgentRuntimeAuthConfig( + agent_harness_id="test-harness", + agent_version="1.2.3", + running_location="test-location", + capabilities=("responsesapi", "sandbox"), + ttl=3600, + external_task_ref="external-run", + ), + model_provider=object(), + ) + assert resolved is not None + client = _FakeOpenAIClient() + manager = OpenAIAgentRuntimeAuthManager(resolved) + + header = await manager.authorization_header(client) + + assert len(client.calls) == 2 + register_path, register_body = client.calls[0] + assert register_path == "/agent/register" + assert register_body["abom"] == { + "agent_version": "1.2.3", + "agent_harness_id": "test-harness", + "running_location": "test-location", + } + assert register_body["capabilities"] == ["responsesapi", "sandbox"] + assert register_body["ttl"] == 3600 + task_path, task_body = client.calls[1] + assert task_path == "/agent/agent_123/task/register" + assert task_body["external_task_ref"] == "external-run" + + payload = _decode_assertion(header) + assert payload["agent_runtime_id"] == "agent_123" + assert payload["task_id"] == "task_456" + assert client.agent_public_key is not None + _verify_agent_assertion(payload, client.agent_public_key) + + +@pytest.mark.asyncio +async def test_sandbox_runtime_adds_agent_assertion_header_for_responses_model() -> None: + client = _FakeOpenAIClient() + model = OpenAIResponsesModel(model="gpt-4.1", openai_client=client) # type: ignore[arg-type] + sandbox_agent = SandboxAgent(name="sandbox") + runtime = SandboxRuntime( + starting_agent=sandbox_agent, + run_config=RunConfig( + sandbox=SandboxRunConfig( + agent_runtime_auth=OpenAIAgentRuntimeAuthConfig( + agent_harness_id="test-harness", + running_location="test-location", + ) + ) + ), + run_state=None, + ) + + updated = await runtime.prepare_model_settings( + public_agent=sandbox_agent, + model=model, + model_settings=ModelSettings(extra_headers={"X-Test": "1"}), + ) + + assert updated.extra_headers is not None + assert updated.extra_headers["X-Test"] == "1" + assert str(updated.extra_headers["Authorization"]).startswith("AgentAssertion ") + assert len(client.calls) == 2 + + +@pytest.mark.asyncio +async def test_sandbox_runtime_adds_agent_assertion_header_for_responses_ws_model() -> None: + client = _FakeOpenAIClient() + model = OpenAIResponsesWSModel(model="gpt-4.1", openai_client=client) # type: ignore[arg-type] + sandbox_agent = SandboxAgent(name="sandbox") + runtime = SandboxRuntime( + starting_agent=sandbox_agent, + run_config=RunConfig( + sandbox=SandboxRunConfig( + agent_runtime_auth=OpenAIAgentRuntimeAuthConfig( + agent_harness_id="test-harness", + running_location="test-location", + ) + ) + ), + run_state=None, + ) + updated = await runtime.prepare_model_settings( + public_agent=sandbox_agent, + model=model, + model_settings=ModelSettings(extra_headers={"X-Test": "1"}), + ) + _, _, handshake_headers = await model._prepare_websocket_request( + { + "model": "gpt-4.1", + "input": [], + "extra_headers": updated.extra_headers, + } + ) + + assert updated.extra_headers is not None + assert updated.extra_headers["X-Test"] == "1" + assert str(updated.extra_headers["Authorization"]).startswith("AgentAssertion ") + assert handshake_headers["Authorization"].startswith("AgentAssertion ") + assert handshake_headers["X-Client"] == "1" + assert len(client.calls) == 2 + + +def test_add_agent_assertion_header_rejects_explicit_authorization() -> None: + with pytest.raises(UserError, match="explicit Authorization"): + add_agent_assertion_header( + ModelSettings(extra_headers={"authorization": "Bearer sk-test"}), + authorization_header="AgentAssertion token", + ) + + +def _encrypt_task_id(task_id: str, agent_public_key: str) -> str: + loaded_key = serialization.load_ssh_public_key(agent_public_key.encode("utf-8")) + assert isinstance(loaded_key, Ed25519PublicKey) + raw_public_key = loaded_key.public_bytes( + encoding=serialization.Encoding.Raw, + format=serialization.PublicFormat.Raw, + ) + curve25519_public_key = PublicKey(crypto_sign_ed25519_pk_to_curve25519(raw_public_key)) + encrypted_task_id = SealedBox(curve25519_public_key).encrypt(task_id.encode("utf-8")) + return base64.b64encode(encrypted_task_id).decode("ascii") + + +def _decode_assertion(header: str) -> dict[str, str]: + scheme, token = header.split(" ", 1) + assert scheme == "AgentAssertion" + padding = "=" * (-len(token) % 4) + payload = json.loads(base64.urlsafe_b64decode(f"{token}{padding}").decode("utf-8")) + assert isinstance(payload, dict) + return {str(key): str(value) for key, value in payload.items()} + + +def _verify_agent_assertion(payload: dict[str, str], agent_public_key: str) -> None: + loaded_key = serialization.load_ssh_public_key(agent_public_key.encode("utf-8")) + assert isinstance(loaded_key, Ed25519PublicKey) + loaded_key.verify( + base64.b64decode(payload["signature"]), + (f"{payload['agent_runtime_id']}:{payload['task_id']}:{payload['timestamp']}").encode(), + ) diff --git a/tests/sandbox/test_compatibility_guards.py b/tests/sandbox/test_compatibility_guards.py index 5a11e5bf77..aac9af190f 100644 --- a/tests/sandbox/test_compatibility_guards.py +++ b/tests/sandbox/test_compatibility_guards.py @@ -365,6 +365,7 @@ def test_sandbox_dataclass_constructor_field_order_is_stable() -> None: "snapshot", "concurrency_limits", "archive_limits", + "agent_runtime_auth", ) diff --git a/uv.lock b/uv.lock index 3e5cb31b70..7a31665407 100644 --- a/uv.lock +++ b/uv.lock @@ -2471,6 +2471,7 @@ e2b = [ ] encrypt = [ { name = "cryptography" }, + { name = "pynacl" }, ] litellm = [ { name = "litellm" }, @@ -2532,6 +2533,7 @@ dev = [ { name = "mypy" }, { name = "playwright" }, { name = "pymongo" }, + { name = "pynacl" }, { name = "pynput" }, { name = "pyright" }, { name = "pytest" }, @@ -2571,6 +2573,7 @@ requires-dist = [ { name = "openai", specifier = ">=2.26.0,<3" }, { name = "pydantic", specifier = ">=2.12.2,<3" }, { name = "pymongo", marker = "extra == 'mongodb'", specifier = ">=4.14" }, + { name = "pynacl", marker = "extra == 'encrypt'", specifier = ">=1.5,<1.6" }, { name = "redis", marker = "extra == 'redis'", specifier = ">=7" }, { name = "requests", specifier = ">=2.0,<3" }, { name = "runloop-api-client", marker = "extra == 'runloop'", specifier = ">=1.16.0,<2.0.0" }, @@ -2606,6 +2609,7 @@ dev = [ { name = "mypy" }, { name = "playwright", specifier = "==1.50.0" }, { name = "pymongo", specifier = ">=4.14" }, + { name = "pynacl", specifier = ">=1.5,<1.6" }, { name = "pynput" }, { name = "pyright", specifier = "==1.1.408" }, { name = "pytest" }, @@ -3192,6 +3196,26 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/32/cd/ddc794cdc8500f6f28c119c624252fb6dfb19481c6d7ed150f13cf468a6d/pymongo-4.16.0-cp314-cp314t-win_arm64.whl", hash = "sha256:6b2a20edb5452ac8daa395890eeb076c570790dfce6b7a44d788af74c2f8cf96", size = 1047725, upload-time = "2026-01-07T18:05:28.47Z" }, ] +[[package]] +name = "pynacl" +version = "1.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cffi" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a7/22/27582568be639dfe22ddb3902225f91f2f17ceff88ce80e4db396c8986da/PyNaCl-1.5.0.tar.gz", hash = "sha256:8ac7448f09ab85811607bdd21ec2464495ac8b7c66d146bf545b0f08fb9220ba", size = 3392854, upload-time = "2022-01-07T22:05:41.134Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/75/0b8ede18506041c0bf23ac4d8e2971b4161cd6ce630b177d0a08eb0d8857/PyNaCl-1.5.0-cp36-abi3-macosx_10_10_universal2.whl", hash = "sha256:401002a4aaa07c9414132aaed7f6836ff98f59277a234704ff66878c2ee4a0d1", size = 349920, upload-time = "2022-01-07T22:05:49.156Z" }, + { url = "https://files.pythonhosted.org/packages/59/bb/fddf10acd09637327a97ef89d2a9d621328850a72f1fdc8c08bdf72e385f/PyNaCl-1.5.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.manylinux_2_24_aarch64.whl", hash = "sha256:52cb72a79269189d4e0dc537556f4740f7f0a9ec41c1322598799b0bdad4ef92", size = 601722, upload-time = "2022-01-07T22:05:50.989Z" }, + { url = "https://files.pythonhosted.org/packages/5d/70/87a065c37cca41a75f2ce113a5a2c2aa7533be648b184ade58971b5f7ccc/PyNaCl-1.5.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a36d4a9dda1f19ce6e03c9a784a2921a4b726b02e1c736600ca9c22029474394", size = 680087, upload-time = "2022-01-07T22:05:52.539Z" }, + { url = "https://files.pythonhosted.org/packages/ee/87/f1bb6a595f14a327e8285b9eb54d41fef76c585a0edef0a45f6fc95de125/PyNaCl-1.5.0-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_24_x86_64.whl", hash = "sha256:0c84947a22519e013607c9be43706dd42513f9e6ae5d39d3613ca1e142fba44d", size = 856678, upload-time = "2022-01-07T22:05:54.251Z" }, + { url = "https://files.pythonhosted.org/packages/66/28/ca86676b69bf9f90e710571b67450508484388bfce09acf8a46f0b8c785f/PyNaCl-1.5.0-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:06b8f6fa7f5de8d5d2f7573fe8c863c051225a27b61e6860fd047b1775807858", size = 1133660, upload-time = "2022-01-07T22:05:56.056Z" }, + { url = "https://files.pythonhosted.org/packages/3d/85/c262db650e86812585e2bc59e497a8f59948a005325a11bbbc9ecd3fe26b/PyNaCl-1.5.0-cp36-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:a422368fc821589c228f4c49438a368831cb5bbc0eab5ebe1d7fac9dded6567b", size = 663824, upload-time = "2022-01-07T22:05:57.434Z" }, + { url = "https://files.pythonhosted.org/packages/fd/1a/cc308a884bd299b651f1633acb978e8596c71c33ca85e9dc9fa33a5399b9/PyNaCl-1.5.0-cp36-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:61f642bf2378713e2c2e1de73444a3778e5f0a38be6fee0fe532fe30060282ff", size = 1117912, upload-time = "2022-01-07T22:05:58.665Z" }, + { url = "https://files.pythonhosted.org/packages/25/2d/b7df6ddb0c2a33afdb358f8af6ea3b8c4d1196ca45497dd37a56f0c122be/PyNaCl-1.5.0-cp36-abi3-win32.whl", hash = "sha256:e46dae94e34b085175f8abb3b0aaa7da40767865ac82c928eeb9e57e1ea8a543", size = 204624, upload-time = "2022-01-07T22:06:00.085Z" }, + { url = "https://files.pythonhosted.org/packages/5e/22/d3db169895faaf3e2eda892f005f433a62db2decbcfbc2f61e6517adfa87/PyNaCl-1.5.0-cp36-abi3-win_amd64.whl", hash = "sha256:20f42270d27e1b6a29f54032090b972d97f0a1b0948cc52392041ef7831fee93", size = 212141, upload-time = "2022-01-07T22:06:01.861Z" }, +] + [[package]] name = "pynput" version = "1.8.1"