diff --git a/agentops/sdk/exporters.py b/agentops/sdk/exporters.py index 69e872135..d2969ec3a 100644 --- a/agentops/sdk/exporters.py +++ b/agentops/sdk/exporters.py @@ -66,7 +66,10 @@ def __init__( if compression is not None: parent_kwargs["compression"] = compression - super().__init__(endpoint=endpoint, **parent_kwargs) + # OpenTelemetry >=1.45 dropped the exporter's `_headers` and `_session` attributes, so keep our own session + # and hand it to the parent through its public `session` argument. + self._auth_session = requests.Session() + super().__init__(endpoint=endpoint, session=self._auth_session, **parent_kwargs) def _get_current_jwt(self) -> Optional[str]: """Get the current JWT token from the provider or stored JWT.""" @@ -104,7 +107,7 @@ def _filter_user_headers(self, headers: Optional[Dict[str, str]]) -> Optional[Di def _prepare_headers(self, headers: Optional[Dict[str, str]] = None) -> Dict[str, str]: """Prepare headers with current JWT token.""" # Start with base headers - prepared_headers = dict(self._headers) + prepared_headers = dict(self._auth_session.headers) # Add any additional headers, but only allow non-critical headers filtered_headers = self._filter_user_headers(headers) @@ -137,8 +140,8 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: current_headers = self._prepare_headers() # Temporarily update the session headers for this request - original_headers = dict(self._session.headers) - self._session.headers.update(current_headers) + original_headers = dict(self._auth_session.headers) + self._auth_session.headers.update(current_headers) try: # Call parent export method @@ -153,8 +156,8 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: finally: # Restore original headers - self._session.headers.clear() - self._session.headers.update(original_headers) + self._auth_session.headers.clear() + self._auth_session.headers.update(original_headers) except requests.exceptions.HTTPError as e: if e.response and e.response.status_code in (401, 403): diff --git a/tests/unit/sdk/test_exporters.py b/tests/unit/sdk/test_exporters.py index 65f5a4e3c..a7016609e 100644 --- a/tests/unit/sdk/test_exporters.py +++ b/tests/unit/sdk/test_exporters.py @@ -6,8 +6,9 @@ from unittest.mock import Mock, patch import requests -from opentelemetry.sdk.trace import ReadableSpan -from opentelemetry.sdk.trace.export import SpanExportResult +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.exporter.otlp.proto.http import Compression from agentops.sdk.exporters import AuthenticatedOTLPExporter @@ -233,6 +234,38 @@ def test_full_export_cycle(self): # Test clear method exporter.clear() # Should not raise any exception + def test_export_sends_jwt_and_custom_headers(self): + """The real parent export path must put the JWT and allowed custom headers on the request.""" + memory_exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(memory_exporter)) + with provider.get_tracer("test").start_as_current_span("span"): + pass + + sent = [] + + def fake_send(session, request, **kwargs): + sent.append(request) + response = requests.Response() + response.status_code = 200 + response._content = b"" + response.request = request + return response + + exporter = AuthenticatedOTLPExporter( + endpoint=self.endpoint, + jwt=self.jwt, + headers={"X-Custom-Header": "test-value", "Authorization": "Malicious-Auth malicious-token"}, + ) + + with patch.object(requests.Session, "send", autospec=True, side_effect=fake_send): + result = exporter.export(memory_exporter.get_finished_spans()) + + self.assertEqual(result, SpanExportResult.SUCCESS) + self.assertEqual(len(sent), 1) + self.assertEqual(sent[0].headers["Authorization"], f"Bearer {self.jwt}") + self.assertEqual(sent[0].headers["X-Custom-Header"], "test-value") + def test_export_with_different_compression_types(self): """Test exporter with different compression types.""" compression_types = [Compression.Gzip, Compression.Deflate, None]