diff --git a/headroom/copilot_auth.py b/headroom/copilot_auth.py index 2603577b5..fd9a5504c 100644 --- a/headroom/copilot_auth.py +++ b/headroom/copilot_auth.py @@ -25,6 +25,7 @@ from headroom import paths from headroom._subprocess import run from headroom.copilot_linux_secret import read_copilot_oauth_token as read_linux_secret_token from headroom.copilot_macos_keychain import read_copilot_oauth_token as read_macos_keychain_token +from headroom.proxy import ssl_context as proxy_ssl_context logger = logging.getLogger(__name__) @@ -76,6 +77,15 @@ _OAUTH_TOKEN_KEYS = ( _EXPIRY_KEYS = ("expires_at", "expiresAt", "expiry", "expires") +def _urlopen(request: urllib_request.Request, *, timeout: float) -> Any: + """Open a GitHub request with Headroom's configured corporate trust roots.""" + + context = proxy_ssl_context.build_urlopen_context() + if context is not None: + return urllib_request.urlopen(request, timeout=timeout, context=context) + return urllib_request.urlopen(request, timeout=timeout) + + @dataclass(frozen=True) class CopilotAPIToken: """Short-lived API token exchanged from a GitHub OAuth token.""" @@ -662,7 +672,7 @@ def start_copilot_device_authorization( }, method="POST", ) - with urllib_request.urlopen(request, timeout=timeout) as response: + with _urlopen(request, timeout=timeout) as response: payload = json.loads(response.read().decode("utf-8", errors="replace")) if not isinstance(payload, dict): raise RuntimeError("GitHub device authorization returned an invalid response.") @@ -700,7 +710,7 @@ def poll_copilot_device_authorization( }, method="POST", ) - with urllib_request.urlopen(request, timeout=timeout) as response: + with _urlopen(request, timeout=timeout) as response: payload = json.loads(response.read().decode("utf-8", errors="replace")) if not isinstance(payload, dict): raise RuntimeError("GitHub device authorization returned an invalid response.") @@ -1341,7 +1351,7 @@ def _fetch_copilot_user_info(token: str) -> dict[str, Any] | None: headers = _copilot_token_exchange_headers(token) request = urllib_request.Request(_user_info_url(), headers=headers, method="GET") try: - with urllib_request.urlopen(request, timeout=10.0) as response: + with _urlopen(request, timeout=10.0) as response: payload = json.loads(response.read().decode("utf-8")) except Exception as exc: logger.debug("Unable to resolve Copilot API URL from user info: %s", exc) @@ -1457,7 +1467,7 @@ class CopilotTokenProvider: def _exchange_token_sync(headers: dict[str, str]) -> dict[str, Any]: request = urllib_request.Request(_token_exchange_url(), headers=headers, method="GET") try: - with urllib_request.urlopen(request, timeout=10.0) as response: + with _urlopen(request, timeout=10.0) as response: payload = json.loads(response.read().decode("utf-8")) if not isinstance(payload, dict): return {} diff --git a/headroom/proxy/ssl_context.py b/headroom/proxy/ssl_context.py index 7ecba7c31..058c70d1f 100644 --- a/headroom/proxy/ssl_context.py +++ b/headroom/proxy/ssl_context.py @@ -188,6 +188,19 @@ def build_httpx_verify() -> ssl.SSLContext | bool: return True +def build_urlopen_context() -> ssl.SSLContext | None: + """Return Headroom's configured TLS context for ``urllib.request.urlopen``. + + ``urlopen`` already handles Python's default trust configuration when no + explicit context is passed. Return only a custom context here so callers + retain that default while sharing Headroom's corporate CA and strict-mode + handling when it is configured. + """ + + verify = build_httpx_verify() + return verify if isinstance(verify, ssl.SSLContext) else None + + def apply_global_tls_relaxation() -> bool: """Strip ``VERIFY_X509_STRICT`` from urllib3's context builder when opted in. diff --git a/tests/test_copilot_auth.py b/tests/test_copilot_auth.py index dd377c987..797b70568 100644 --- a/tests/test_copilot_auth.py +++ b/tests/test_copilot_auth.py @@ -10,6 +10,7 @@ from urllib import error as urllib_error import pytest from headroom import copilot_auth +from headroom.proxy import ssl_context def test_device_authorization_uses_form_encoded_request(monkeypatch: pytest.MonkeyPatch) -> None: @@ -1615,3 +1616,36 @@ def test_exchange_token_sync_returns_payload_on_success(monkeypatch: pytest.Monk ) assert result == payload + + +def test_exchange_token_sync_uses_configured_corporate_tls_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Copilot refresh must use the same corporate trust config as upstream I/O.""" + payload = {"token": "copilot-api", "expires_at": int(time.time()) + 3600} + tls_context = object() + captured: dict[str, object] = {} + + class FakeResponse: + def read(self) -> bytes: + return json.dumps(payload).encode() + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def fake_urlopen(*args, **kwargs): # noqa: ANN002, ANN003, ANN202 + captured.update(kwargs) + return FakeResponse() + + monkeypatch.setattr(ssl_context, "build_urlopen_context", lambda: tls_context) + monkeypatch.setattr(copilot_auth.urllib_request, "urlopen", fake_urlopen) + + result = copilot_auth.CopilotTokenProvider._exchange_token_sync( + {"Authorization": "Bearer gho_test"} # noqa: S105 + ) + + assert result == payload + assert captured["context"] is tls_context diff --git a/tests/test_integrations/agno/test_model.py b/tests/test_integrations/agno/test_model.py index 5a0bf990e..a3fe45c40 100644 --- a/tests/test_integrations/agno/test_model.py +++ b/tests/test_integrations/agno/test_model.py @@ -31,6 +31,24 @@ from headroom import HeadroomConfig, HeadroomMode pytestmark = pytest.mark.skipif(not AGNO_AVAILABLE, reason="Agno not installed") +def _response_usage(input_tokens: int, output_tokens: int, total_tokens: int): + """Build response usage across Agno 2.x and 3.x module layouts.""" + + try: + from agno.metrics import MessageMetrics + + metrics_type = MessageMetrics + except ImportError: + from agno.models.metrics import Metrics + + metrics_type = Metrics + return metrics_type( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + ) + + @pytest.fixture def mock_agno_model(): """Create a mock Agno model (OpenAIChat-like).""" @@ -59,11 +77,7 @@ def mock_agno_model(): return ModelResponse( role="assistant", content="Hello! I'm a mock response.", - response_usage=Metrics( - input_tokens=10, - output_tokens=5, - total_tokens=15, - ), + response_usage=_response_usage(10, 5, 15), ) mock.invoke = MagicMock(side_effect=mock_invoke) @@ -79,11 +93,7 @@ def mock_agno_model(): yield ModelResponse( role="assistant", content="Streaming...", - response_usage=Metrics( - input_tokens=10, - output_tokens=5, - total_tokens=15, - ), + response_usage=_response_usage(10, 5, 15), ) mock.invoke_stream = MagicMock(side_effect=mock_invoke_stream)