diff --git a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py index 15ed817e15..46741df994 100644 --- a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py +++ b/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__) diff --git a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/patch.py b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/patch.py index 1adfed4f39..abdcb2a133 100644 --- a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/patch.py +++ b/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 diff --git a/packages/opentelemetry-instrumentation-langchain/tests/test_langgraph.py b/packages/opentelemetry-instrumentation-langchain/tests/test_langgraph.py index 6e330e0e45..0ce2c85a10 100644 --- a/packages/opentelemetry-instrumentation-langchain/tests/test_langgraph.py +++ b/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):