fix(langchain): handle early LangGraph stream closure

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

@@ -1053,6 +1053,11 @@ def _handle_error(
             return
 
         span = self._get_span(run_id)
+        if isinstance(error, GeneratorExit):
+            _set_span_attribute(span, SpanAttributes.GEN_AI_TASK_STATUS, "success")
+            self._end_span(span, run_id)
+            return
+
         # Set task status to failure
         _set_span_attribute(span, SpanAttributes.GEN_AI_TASK_STATUS, "failure")
         span.set_attribute(ERROR_TYPE, type(error).__name__)

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

@@ -28,6 +28,16 @@ def _get_current_node_contextvar():
 LANGGRAPH_FIRST_CHILD_PENDING_KEY = "langgraph_first_child_pending"
 
 
+def _safe_detach_context(token) -> None:
+    """Detach tokens that may be finalized in a copied async context."""
+    from opentelemetry.context import _RUNTIME_CONTEXT
+
+    try:
+        _RUNTIME_CONTEXT.detach(token)
+    except (RuntimeError, ValueError):
+        pass
+
+
 def _set_graph_span_attributes(
     graph_span: Span,
     instance: Any,
@@ -193,14 +203,14 @@ async def async_wrapper(wrapped, instance, args, kwargs):
         try:
             async for item in wrapped(*args, **kwargs):
                 yield item
-        except BaseException as e:
+        except Exception as e:
             graph_span.set_status(Status(StatusCode.ERROR, str(e)))
             graph_span.record_exception(e)
             raise
         finally:
             graph_span.end()
-            context_api.detach(graph_span_ctx)
-            context_api.detach(langgraph_ctx)
+            _safe_detach_context(graph_span_ctx)
+            _safe_detach_context(langgraph_ctx)
 
     return async_wrapper if is_async else wrapper
 

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

@@ -10,7 +10,7 @@
 )
 from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GenAiOperationNameValues
 from opentelemetry.semconv_ai import GenAICustomOperationName, SpanAttributes
-from opentelemetry.trace import INVALID_SPAN
+from opentelemetry.trace import INVALID_SPAN, StatusCode
 
 
 @pytest.mark.vcr
@@ -119,6 +119,91 @@ def calculate(state: State):
     assert graph_span.attributes[GenAIAttributes.GEN_AI_PROVIDER_NAME] == "langgraph"
 
 
+@pytest.mark.asyncio
+async def test_langgraph_astream_early_return_is_not_error(
+    instrument_legacy, span_exporter, caplog
+):
+    import asyncio
+    import gc
+
+    class State(TypedDict):
+        value: int
+
+    def first(state: State):
+        return {"value": state["value"] + 1}
+
+    def second(state: State):
+        return {"value": state["value"] + 1}
+
+    workflow = StateGraph(State)
+    workflow.add_node("first", first)
+    workflow.add_node("second", second)
+    workflow.set_entry_point("first")
+    workflow.add_edge("first", "second")
+    graph = workflow.compile()
+
+    async def consume_first_result():
+        async for event in graph.astream({"value": 0}):
+            return event
+        raise AssertionError("stream produced no events")
+
+    baseline_tasks = asyncio.all_tasks()
+    assert await consume_first_result() == {"first": {"value": 1}}
+
+    gc.collect()
+
+    async def finish_cleanup_tasks():
+        while not any(
+            span.name == "LangGraph.workflow"
+            for span in span_exporter.get_finished_spans()
+        ):
+            await asyncio.sleep(0)
+            cleanup_tasks = (
+                asyncio.all_tasks() - baseline_tasks - {asyncio.current_task()}
+            )
+            if cleanup_tasks:
+                await asyncio.gather(*cleanup_tasks)
+
+    await asyncio.wait_for(finish_cleanup_tasks(), timeout=1)
+
+    spans = span_exporter.get_finished_spans()
+    workflow_span = next(span for span in spans if span.name == "LangGraph.workflow")
+    graph_span = next(span for span in spans if span.name == "invoke_agent LangGraph")
+
+    assert workflow_span.status.status_code is StatusCode.UNSET
+    assert graph_span.status.status_code is StatusCode.UNSET
+    assert all(event.name != "exception" for event in workflow_span.events)
+    assert all(event.name != "exception" for event in graph_span.events)
+    assert "Failed to detach context" not in caplog.messages
+
+
+@pytest.mark.asyncio
+async def test_langgraph_astream_error_is_recorded(instrument_legacy, span_exporter):
+    class State(TypedDict):
+        value: int
+
+    def fail(_state: State):
+        raise ValueError("expected failure")
+
+    workflow = StateGraph(State)
+    workflow.add_node("fail", fail)
+    workflow.set_entry_point("fail")
+    graph = workflow.compile()
+
+    with pytest.raises(ValueError, match="expected failure"):
+        async for _ in graph.astream({"value": 0}):
+            pass
+
+    spans = span_exporter.get_finished_spans()
+    workflow_span = next(span for span in spans if span.name == "LangGraph.workflow")
+    graph_span = next(span for span in spans if span.name == "invoke_agent LangGraph")
+
+    assert workflow_span.status.status_code is StatusCode.ERROR
+    assert graph_span.status.status_code is StatusCode.ERROR
+    assert any(event.name == "exception" for event in workflow_span.events)
+    assert any(event.name == "exception" for event in graph_span.events)
+
+
 @pytest.mark.vcr
 def test_langgraph_double_invoke(instrument_legacy, span_exporter):
     class DummyGraphState(TypedDict):