diff --git a/src/cwsandbox/_auth.py b/src/cwsandbox/_auth.py index a856f55..4cd8095 100644 --- a/src/cwsandbox/_auth.py +++ b/src/cwsandbox/_auth.py @@ -20,6 +20,10 @@ from collections.abc import Callable from dataclasses import dataclass +from cwsandbox._client_metadata import ( + HEADER_CWSANDBOX_CLIENT_VERSION, + client_metadata_headers, +) from cwsandbox.exceptions import CWSandboxAuthenticationError logger = logging.getLogger(__name__) @@ -107,11 +111,18 @@ def resolve_auth() -> AuthHeaders: def resolve_auth_metadata() -> tuple[tuple[str, str], ...]: """Resolve authentication credentials and return as gRPC metadata tuples. - Convenience wrapper around resolve_auth() that returns metadata in the - format expected by gRPC calls (lowercase key-value tuples). + Convenience wrapper around resolve_auth() that returns auth headers plus + cwsandbox-managed client metadata in the format expected by gRPC calls + (lowercase key-value tuples). Returns: Tuple of (key, value) pairs suitable for gRPC metadata parameter """ auth = resolve_auth() - return tuple((k.lower(), v) for k, v in auth.headers.items()) + metadata = {k.lower(): v for k, v in auth.headers.items()} + for key, value in client_metadata_headers(): + if key == HEADER_CWSANDBOX_CLIENT_VERSION: + metadata[key] = value + else: + metadata.setdefault(key, value) + return tuple(metadata.items()) diff --git a/src/cwsandbox/_client_metadata.py b/src/cwsandbox/_client_metadata.py new file mode 100644 index 0000000..c67e1f4 --- /dev/null +++ b/src/cwsandbox/_client_metadata.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: 2025 CoreWeave, Inc. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-PackageName: cwsandbox-client + +"""Request-scoped client metadata sent as gRPC metadata.""" + +from __future__ import annotations + +from importlib import metadata as importlib_metadata + +HEADER_CWSANDBOX_CLIENT_VERSION = "x-cwsandbox-client-version" +HEADER_SANDBOX_INTEGRATION = "x-sandbox-integration" + +_INTEGRATION_METADATA_VALUE = "" + + +def _cwsandbox_version() -> str: + try: + return importlib_metadata.version("cwsandbox") + except importlib_metadata.PackageNotFoundError: + from cwsandbox import __version__ + + return __version__ + + +def set_integration_metadata(integration: str) -> None: + """Set the integration name attached to future sandbox API requests. + + The value is process-global and optional. Passing an empty string clears the + integration metadata. + """ + global _INTEGRATION_METADATA_VALUE + _INTEGRATION_METADATA_VALUE = integration + + +def _reset_client_metadata_for_testing() -> None: + """Reset process-global client metadata for unit-test isolation.""" + set_integration_metadata("") + + +def client_metadata_headers() -> tuple[tuple[str, str], ...]: + """Return cwsandbox-managed client metadata headers.""" + headers = [(HEADER_CWSANDBOX_CLIENT_VERSION, _cwsandbox_version())] + if _INTEGRATION_METADATA_VALUE: + headers.append((HEADER_SANDBOX_INTEGRATION, _INTEGRATION_METADATA_VALUE)) + return tuple(headers) diff --git a/tests/unit/cwsandbox/conftest.py b/tests/unit/cwsandbox/conftest.py index 665936a..7990f08 100644 --- a/tests/unit/cwsandbox/conftest.py +++ b/tests/unit/cwsandbox/conftest.py @@ -8,12 +8,14 @@ import asyncio import concurrent.futures +from collections.abc import Generator from typing import TypeVar from unittest.mock import MagicMock import pytest from cwsandbox import OperationRef +from cwsandbox._client_metadata import _reset_client_metadata_for_testing from cwsandbox._types import Process, ProcessResult, StreamReader T = TypeVar("T") @@ -33,7 +35,7 @@ @pytest.fixture(autouse=True) -def clean_auth_env(monkeypatch: pytest.MonkeyPatch) -> None: +def clean_auth_env(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: """Clear all auth-related env vars before each test. This runs automatically for every test (autouse=True) and ensures: @@ -43,6 +45,9 @@ def clean_auth_env(monkeypatch: pytest.MonkeyPatch) -> None: """ for var in AUTH_ENV_VARS: monkeypatch.delenv(var, raising=False) + _reset_client_metadata_for_testing() + yield + _reset_client_metadata_for_testing() @pytest.fixture(autouse=True) diff --git a/tests/unit/cwsandbox/test_auth.py b/tests/unit/cwsandbox/test_auth.py index 6605b8c..67b6025 100644 --- a/tests/unit/cwsandbox/test_auth.py +++ b/tests/unit/cwsandbox/test_auth.py @@ -6,6 +6,7 @@ import pytest +from cwsandbox import __version__ from cwsandbox._auth import ( AuthHeaders, _reset_auth_mode_for_testing, @@ -13,6 +14,11 @@ resolve_auth_metadata, set_auth_mode, ) +from cwsandbox._client_metadata import ( + HEADER_CWSANDBOX_CLIENT_VERSION, + HEADER_SANDBOX_INTEGRATION, + set_integration_metadata, +) from cwsandbox.exceptions import CWSandboxAuthenticationError @@ -206,7 +212,10 @@ def test_returns_lowercased_metadata_tuples(self, monkeypatch: pytest.MonkeyPatc result = resolve_auth_metadata() - assert result == (("authorization", "Bearer test-key"),) + assert result == ( + ("authorization", "Bearer test-key"), + (HEADER_CWSANDBOX_CLIENT_VERSION, __version__), + ) def test_returns_registered_auth_mode_metadata( self, @@ -224,12 +233,50 @@ def test_returns_registered_auth_mode_metadata( finally: _reset_auth_mode_for_testing() - assert result == (("x-api-key", "mode-key"),) + assert result == ( + ("x-api-key", "mode-key"), + (HEADER_CWSANDBOX_CLIENT_VERSION, __version__), + ) - def test_returns_empty_tuple_when_no_auth(self, monkeypatch: pytest.MonkeyPatch) -> None: - """Test returns empty tuple when no credentials found.""" + def test_returns_client_metadata_when_no_auth(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Test cwsandbox client metadata is sent even without auth credentials.""" monkeypatch.delenv("CWSANDBOX_API_KEY", raising=False) result = resolve_auth_metadata() - assert result == () + assert result == ((HEADER_CWSANDBOX_CLIENT_VERSION, __version__),) + + def test_returns_integration_metadata_when_set(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Test integration metadata is included only when explicitly set.""" + monkeypatch.setenv("CWSANDBOX_API_KEY", "test-key") + set_integration_metadata("harbor") + + result = resolve_auth_metadata() + + assert result == ( + ("authorization", "Bearer test-key"), + (HEADER_CWSANDBOX_CLIENT_VERSION, __version__), + (HEADER_SANDBOX_INTEGRATION, "harbor"), + ) + + def test_auth_mode_metadata_can_override_integration_metadata( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Test auth modes may provide their own integration metadata header.""" + monkeypatch.delenv("CWSANDBOX_API_KEY", raising=False) + set_integration_metadata("client-default") + set_auth_mode( + "auth-mode-test", + lambda: AuthHeaders( + headers={HEADER_SANDBOX_INTEGRATION: "auth-integration"}, + strategy="auth_mode", + ), + ) + + try: + result = dict(resolve_auth_metadata()) + finally: + _reset_auth_mode_for_testing() + + assert result[HEADER_SANDBOX_INTEGRATION] == "auth-integration" + assert result[HEADER_CWSANDBOX_CLIENT_VERSION] == __version__