fix: mask attributes when starting span (#892)

python/openinference-instrumentation/src/openinference/instrumentation/config.py

@@ -1,13 +1,16 @@
 import os
 from contextlib import contextmanager
 from dataclasses import dataclass, field, fields
+from functools import partial
+from inspect import signature
 from typing import (
     Any,
     Callable,
     Dict,
     Iterator,
     Optional,
     Union,
+    cast,
     get_args,
 )
 
@@ -26,6 +29,7 @@
     detach,
     set_value,
 )
+from opentelemetry.trace import Tracer
 from opentelemetry.util.types import AttributeValue
 
 from .logging import logger
@@ -324,21 +328,35 @@ def set_attribute(
             span.set_attribute(key, value)
 
 
+class _TracerSignatures:
+    # remove `self` via partial
+    start_as_current_span = signature(partial(Tracer.start_as_current_span, None))
+    start_span = signature(partial(Tracer.start_span, None))
+
+
 class OITracer(wrapt.ObjectProxy):  # type: ignore[misc]
     def __init__(self, wrapped: trace_api.Tracer, config: TraceConfig) -> None:
         super().__init__(wrapped)
         self._self_config = config
 
     @contextmanager
-    def start_as_current_span(self, *args, **kwargs) -> Iterator[trace_api.Span]:
-        with self.__wrapped__.start_as_current_span(*args, **kwargs) as span:
-            yield _MaskedSpan(span, self._self_config)
-
-    def start_span(self, *args, **kwargs) -> trace_api.Span:
-        return _MaskedSpan(
-            self.__wrapped__.start_span(*args, **kwargs),
-            config=self._self_config,
-        )
+    def start_as_current_span(self, *args: Any, **kwargs: Any) -> Iterator[trace_api.Span]:
+        kwargs = _TracerSignatures.start_as_current_span.bind(*args, **kwargs).arguments
+        attributes = cast(Optional[Dict[str, AttributeValue]], kwargs.pop("attributes", None))
+        with self.__wrapped__.start_as_current_span(**kwargs) as span:
+            span = _MaskedSpan(span, self._self_config)
+            if attributes:
+                span.set_attributes(attributes)
+            yield span
+
+    def start_span(self, *args: Any, **kwargs: Any) -> trace_api.Span:
+        kwargs = _TracerSignatures.start_span.bind(*args, **kwargs).arguments
+        attributes = cast(Optional[Dict[str, AttributeValue]], kwargs.pop("attributes", None))
+        span = self.__wrapped__.start_span(**kwargs)
+        span = _MaskedSpan(span, config=self._self_config)
+        if attributes:
+            span.set_attributes(attributes)
+        return span
 
 
 def is_base64_url(url: str) -> bool:

python/openinference-instrumentation/tests/conftest.py

@@ -0,0 +1,17 @@
+import pytest
+from opentelemetry.sdk import trace as trace_sdk
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import TracerProvider
+
+
+@pytest.fixture(scope="module")
+def in_memory_span_exporter() -> InMemorySpanExporter:
+    return InMemorySpanExporter()
+
+
+@pytest.fixture(scope="module")
+def tracer_provider(in_memory_span_exporter: InMemorySpanExporter) -> TracerProvider:
+    tracer_provider = trace_sdk.TracerProvider()
+    tracer_provider.add_span_processor(SimpleSpanProcessor(in_memory_span_exporter))
+    return tracer_provider

python/openinference-instrumentation/tests/test_config.py

@@ -20,7 +20,11 @@
     OPENINFERENCE_HIDE_OUTPUT_MESSAGES,
     OPENINFERENCE_HIDE_OUTPUT_TEXT,
     OPENINFERENCE_HIDE_OUTPUTS,
+    REDACTED_VALUE,
+    OITracer,
 )
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import TracerProvider
 
 
 def test_default_settings() -> None:
@@ -35,6 +39,17 @@ def test_default_settings() -> None:
     assert config.base64_image_max_length == DEFAULT_BASE64_IMAGE_MAX_LENGTH
 
 
+def test_oi_tracer(
+    tracer_provider: TracerProvider,
+    in_memory_span_exporter: InMemorySpanExporter,
+) -> None:
+    tracer = OITracer(tracer_provider.get_tracer(__name__), TraceConfig(hide_inputs=True))
+    with tracer.start_as_current_span("a", attributes={"input.value": "c"}):
+        tracer.start_span("b", attributes={"input.value": "d"}).end()
+    for span in in_memory_span_exporter.get_finished_spans():
+        assert (span.attributes or {}).get("input.value") == REDACTED_VALUE
+
+
 @pytest.mark.parametrize("hide_inputs", [False, True])
 @pytest.mark.parametrize("hide_outputs", [False, True])
 @pytest.mark.parametrize("hide_input_messages", [False, True])

python/tox.ini

@@ -42,7 +42,7 @@ changedir =
   litellm: instrumentation/openinference-instrumentation-litellm/
   instructor: instrumentation/openinference-instrumentation-instructor/
 commands_pre =
-  instrumentation: uv pip install --reinstall {toxinidir}/openinference-instrumentation
+  instrumentation: uv pip install --reinstall {toxinidir}/openinference-instrumentation[test]
   semconv: uv pip install --reinstall {toxinidir}/openinference-semantic-conventions
   bedrock: uv pip install --reinstall {toxinidir}/instrumentation/openinference-instrumentation-bedrock[test]
   bedrock-latest: uv pip install -U boto3