feat(langchain): emit AG-UI text and tool events

packages/client/tests/test_agui.py

@@ -149,6 +149,16 @@ def test_fail_open_masks_a_call_that_never_finished(self) -> None:
         assert events[-1]["content"] == ""
         assert events[-1]["metadata"]["launchdarkly"]["toolFailed"] is True
 
+    def test_tool_end_extends_a_streamed_argument_prefix(self) -> None:
+        emitter = TurnEmitter(capture_content=True)
+        events = emitter.tool_start("a", "search")
+        events.extend(emitter.tool_args("a", '{"q":'))
+        events.extend(emitter.tool_end("a", arguments={"q": "x"}))
+        deltas = [
+            event["delta"] for event in events if event["type"] == "TOOL_CALL_ARGS"
+        ]
+        assert "".join(deltas) == '{"q":"x"}'
+
     def test_a_second_start_for_the_same_id_is_dropped(self) -> None:
         emitter = TurnEmitter(capture_content=True)
         assert emitter.tool_start("a", "search")

packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py

@@ -31,6 +31,7 @@
     set_input_content_attributes,
     set_output_content_attributes,
 )
+from launchdarkly_ai_server.agui import TurnEmitter
 from launchdarkly_ai_server.lifecycle import effective_capture_content
 
 from ._version import PACKAGE_NAME, __version__
@@ -474,6 +475,7 @@ async def _stream_gen(
     # stream; a CancelledError means something cancelled the run, usually a timeout, and the
     # consumer chose nothing. The blocking path already tells these apart.
     cancelled = False
+    emitter = TurnEmitter(capture_content=capture_content)
     try:
         # Inside the guard, because serialising the prompt raises on anything that is not
         # JSON-serialisable. A raise out here would leave the root open with the `finally` never
@@ -519,8 +521,36 @@ async def _stream_gen(
                     if getattr(msg, "type", None) == "ai":
                         text = lang_chain_content_text(msg.content)
                         if text:
-                            yield {"type": "chunk", "text": text}
+                            for rich in emitter.text(text):
+                                yield rich
                             full_output = text
+                        for tc in getattr(msg, "tool_calls", None) or []:
+                            if not isinstance(tc, dict):
+                                continue
+                            call_id = str(tc.get("id") or tc.get("name") or "")
+                            name = str(tc.get("name") or "")
+                            for rich in emitter.tool_start(call_id, name):
+                                yield rich
+                            for rich in emitter.tool_end(
+                                call_id, arguments=tc.get("args")
+                            ):
+                                yield rich
+                    elif getattr(msg, "type", None) == "tool":
+                        call_id = str(getattr(msg, "tool_call_id", "") or "")
+                        if call_id:
+                            for rich in emitter.tool_result(
+                                call_id,
+                                getattr(msg, "content", None),
+                                failed=getattr(msg, "status", None) == "error",
+                            ):
+                                yield rich
+            for rich in emitter.end_text():
+                yield rich
+
+        # A step can start a call the graph never answers with a tool message.
+        # Close it before the run finishes, or the client keeps it open.
+        for rich in emitter.fail_open():
+            yield rich
 
         # The same reconciliation the blocking path does, for the same reason. This walk reads
         # usage_metadata off the astream payloads, and sees nothing at all unless the graph is
@@ -562,6 +592,8 @@ async def _stream_gen(
         if span_callbacks.run_usage.reported:
             finish_root_span(span, config, span_callbacks.run_usage.total)
         fail_span(span, exc, ended)
+        for rich in emitter.fail_open():
+            yield rich
         raise
     finally:
         # A no-op on the success and failure paths, because both already ended their spans through

packages/langchain-agents/tests/test_handler.py

@@ -1427,7 +1427,11 @@ async def _mock_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]:
                     llm=MagicMock()
                 ).stream(_make_config(), "q")
             ]
-        assert {"type": "chunk", "text": "visible answer"} in events
+        assert any(
+            event.get("type") == "TEXT_MESSAGE_CONTENT"
+            and event.get("delta") == "visible answer"
+            for event in events
+        )
         assert events[-1]["output"] == "visible answer"
 
     @pytest.mark.asyncio
@@ -1447,6 +1451,111 @@ async def _bad_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]:
                     async for _ in await h.stream(_make_config(), "hi"):
                         pass
 
+    @pytest.mark.asyncio
+    async def test_a_failed_run_closes_started_tool_calls(self) -> None:
+        mocks = _make_langchain_mock()
+        msg = _make_ai_msg("looking")
+        msg.tool_calls = [
+            {"name": "search", "id": "c1", "args": {"q": "x"}},
+            {"name": "other", "id": "d1", "args": {"n": 1}},
+        ]
+
+        async def _astream(*_a: Any, **_kw: Any) -> AsyncIterator[Any]:
+            yield {"agent": {"messages": [msg]}}
+            raise RuntimeError("tool failed")
+
+        mocks["_agent"].astream = _astream
+        seen: list[dict[str, Any]] = []
+        with (
+            _patch_lc(mocks),
+            patch.object(spans_mod, "_HAS_OTEL", False),
+            pytest.raises(RuntimeError, match="tool failed"),
+        ):
+            handler = create_langchain_agents_handler(llm=MagicMock())
+            async for event in await handler.stream(_make_config(), "q"):
+                seen.append(event)
+        results = {
+            event["toolCallId"] for event in seen if event["type"] == "TOOL_CALL_RESULT"
+        }
+        assert results == {"c1", "d1"}
+
+    @pytest.mark.asyncio
+    async def test_tool_messages_close_parallel_calls(self) -> None:
+        mocks = _make_langchain_mock()
+        ai = _make_ai_msg("looking")
+        ai.tool_calls = [
+            {"name": "search", "id": "c1", "args": {"q": "x"}},
+            {"name": "other", "id": "d1", "args": {"n": 1}},
+        ]
+
+        def tool(call_id: str, content: str) -> Any:
+            message = MagicMock()
+            message.type = "tool"
+            message.tool_call_id = call_id
+            message.content = content
+            message.status = "success"
+            message.usage_metadata = None
+            return message
+
+        async def _astream(*_a: Any, **_kw: Any) -> AsyncIterator[Any]:
+            yield {"agent": {"messages": [ai]}}
+            yield {"tools": {"messages": [tool("c1", "found"), tool("d1", "also")]}}
+
+        mocks["_agent"].astream = _astream
+        with _patch_lc(mocks), patch.object(spans_mod, "_HAS_OTEL", False):
+            events = [
+                event
+                async for event in await create_langchain_agents_handler(
+                    llm=MagicMock(), capture_content=True
+                ).stream(_make_config(), "q")
+            ]
+        by_id: dict[str, list[str]] = {}
+        for event in events:
+            if not str(event.get("type", "")).startswith("TOOL_CALL"):
+                continue
+            by_id.setdefault(event["toolCallId"], []).append(event["type"])
+        assert by_id["c1"] == [
+            "TOOL_CALL_START",
+            "TOOL_CALL_ARGS",
+            "TOOL_CALL_END",
+            "TOOL_CALL_RESULT",
+        ]
+        assert by_id["d1"] == [
+            "TOOL_CALL_START",
+            "TOOL_CALL_ARGS",
+            "TOOL_CALL_END",
+            "TOOL_CALL_RESULT",
+        ]
+        contents = {
+            event["toolCallId"]: event["content"]
+            for event in events
+            if event["type"] == "TOOL_CALL_RESULT"
+        }
+        assert contents == {"c1": "found", "d1": "also"}
+
+    @pytest.mark.asyncio
+    async def test_a_successful_run_closes_a_call_with_no_tool_message(self) -> None:
+        mocks = _make_langchain_mock()
+        ai = _make_ai_msg("looking")
+        ai.tool_calls = [{"name": "search", "id": "c1", "args": {"q": "x"}}]
+
+        async def _astream(*_a: Any, **_kw: Any) -> AsyncIterator[Any]:
+            yield {"agent": {"messages": [ai]}}
+
+        mocks["_agent"].astream = _astream
+        with _patch_lc(mocks), patch.object(spans_mod, "_HAS_OTEL", False):
+            events = [
+                event
+                async for event in await create_langchain_agents_handler(
+                    llm=MagicMock()
+                ).stream(_make_config(), "q")
+            ]
+        result = next(event for event in events if event["type"] == "TOOL_CALL_RESULT")
+        assert result["toolCallId"] == "c1"
+        assert result["metadata"]["launchdarkly"]["toolFailed"] is True
+        assert events[-1]["type"] == "done"
+        assert events.index(result) < len(events) - 1
+
 
 # ---------------------------------------------------------------------------
 # §1.5 Streaming telemetry (Appendix A.5 — do not patch _HAS_OTEL=False)
@@ -2586,5 +2695,6 @@ def factory(config: Any) -> Any:
             ]
         assert seen[0]["model"]["parameters"] == {"temperature": 0.2}
         assert any(
-            e.get("type") == "chunk" and e.get("text") == "streamed" for e in events
+            e.get("type") == "TEXT_MESSAGE_CONTENT" and e.get("delta") == "streamed"
+            for e in events
         )