fix(langchain): changed dictionary access from spans[run_id] to spans.get(run_id)

packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py

@@ -229,16 +229,18 @@ def __call__(
         run_manager = kwargs.get("run_manager")
         if run_manager:
             run_id = run_manager.run_id
-            span_holder = self._callback_manager.spans[run_id]
-
-            extra_headers = kwargs.get("extra_headers", {})
-
-            # Inject tracing context into the extra headers
-            ctx = set_span_in_context(span_holder.span)
-            TraceContextTextMapPropagator().inject(extra_headers, context=ctx)
-
-            # Update kwargs to include the modified headers
-            kwargs["extra_headers"] = extra_headers
+            span_holder = self._callback_manager.spans.get(run_id)
+
+            if span_holder:
+                extra_headers = kwargs.get("extra_headers", {})
+                ctx = set_span_in_context(span_holder.span)
+                TraceContextTextMapPropagator().inject(extra_headers, context=ctx)
+                kwargs["extra_headers"] = extra_headers
+            else:
+                logger.debug(
+                    "No span found for run_id %s, skipping header injection",
+                    run_id
+                )
 
         # In legacy chains like LLMChain, suppressing model instrumentations
         # within create_llm_span doesn't work, so this should helps as a fallback

packages/opentelemetry-instrumentation-langchain/tests/test_external_run_id.py

@@ -0,0 +1,138 @@
+"""Tests for handling external run_ids from systems like LangSmith."""
+
+import logging
+from unittest.mock import MagicMock, Mock, patch
+from uuid import uuid4
+
+import pytest
+from opentelemetry.instrumentation.langchain.callback_handler import (
+    TraceloopCallbackHandler,
+)
+
+
+@pytest.fixture
+def callback_handler(tracer_provider):
+    """Create a TraceloopCallbackHandler instance."""
+    return TraceloopCallbackHandler(tracer_provider=tracer_provider)
+
+
+@pytest.fixture
+def mock_run_manager():
+    """Create a mock run manager with an external run_id."""
+    manager = Mock()
+    manager.run_id = uuid4()
+    return manager
+
+
+def test_external_run_id_no_keyerror(
+    callback_handler, mock_run_manager, instrument_legacy, caplog
+):
+    """Test that external run_ids (e.g., from LangSmith) don't cause KeyError."""
+    from opentelemetry.instrumentation.langchain import _OpenAITracingWrapper
+
+    # Create the wrapper
+    wrapper = _OpenAITracingWrapper(callback_handler)
+
+    # Mock the wrapped function
+    mock_wrapped = Mock(return_value="test_result")
+
+    # Ensure the run_id is NOT in the spans dictionary
+    # (simulating an external system like LangSmith creating the run_id)
+    assert mock_run_manager.run_id not in callback_handler.spans
+
+    # Create kwargs with the external run_manager
+    kwargs = {
+        "run_manager": mock_run_manager,
+        "extra_headers": {},
+    }
+
+    # Capture debug logs
+    with caplog.at_level(logging.DEBUG):
+        # Call the wrapper - should NOT raise KeyError
+        result = wrapper(mock_wrapped, None, [], kwargs)
+
+    # Verify the function was called successfully
+    assert result == "test_result"
+    mock_wrapped.assert_called_once()
+
+    # Verify debug log was generated
+    assert any(
+        "No span found for run_id" in record.message
+        and "skipping header injection" in record.message
+        for record in caplog.records
+    )
+
+    # Verify extra_headers were not modified since there's no span
+    # (the original empty dict should still be there)
+    assert "traceparent" not in kwargs["extra_headers"]
+
+
+def test_internal_run_id_injects_headers(
+    callback_handler, mock_run_manager, instrument_legacy
+):
+    """Test that internal run_ids (created by OTEL) get headers injected."""
+    from opentelemetry.instrumentation.langchain import _OpenAITracingWrapper
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.trace import SpanKind
+
+    # Create a real span and add it to the callback handler's spans
+    tracer = TracerProvider().get_tracer(__name__)
+    span = tracer.start_span("test_span", kind=SpanKind.CLIENT)
+
+    # Create a span holder (mimicking what the callback handler does)
+    span_holder = Mock()
+    span_holder.span = span
+
+    # Add the span to the callback handler's spans dictionary
+    callback_handler.spans[mock_run_manager.run_id] = span_holder
+
+    # Create the wrapper
+    wrapper = _OpenAITracingWrapper(callback_handler)
+
+    # Mock the wrapped function
+    mock_wrapped = Mock(return_value="test_result")
+
+    # Create kwargs with the internal run_manager
+    kwargs = {
+        "run_manager": mock_run_manager,
+        "extra_headers": {},
+    }
+
+    # Call the wrapper
+    result = wrapper(mock_wrapped, None, [], kwargs)
+
+    # Verify the function was called successfully
+    assert result == "test_result"
+    mock_wrapped.assert_called_once()
+
+    # Verify headers were injected (traceparent should exist)
+    assert "traceparent" in kwargs["extra_headers"]
+    assert kwargs["extra_headers"]["traceparent"]  # Should have a value
+
+    # Clean up
+    span.end()
+
+
+def test_no_run_manager_continues_normally(callback_handler, instrument_legacy):
+    """Test that missing run_manager doesn't cause issues."""
+    from opentelemetry.instrumentation.langchain import _OpenAITracingWrapper
+
+    # Create the wrapper
+    wrapper = _OpenAITracingWrapper(callback_handler)
+
+    # Mock the wrapped function
+    mock_wrapped = Mock(return_value="test_result")
+
+    # Create kwargs without run_manager
+    kwargs = {"extra_headers": {}}
+
+    # Call the wrapper - should work fine
+    result = wrapper(mock_wrapped, None, [], kwargs)
+
+    # Verify the function was called successfully
+    assert result == "test_result"
+    mock_wrapped.assert_called_once()
+
+    # Verify no headers were injected
+    assert "traceparent" not in kwargs["extra_headers"]
+