feat(langgraph): optional on_final_ai_message callback for usage capture

Adds an additive on_final_ai_message=None parameter to
convert_langgraph_to_agentex_events so callers can capture AIMessage
usage_metadata without re-traversing the stream. No behavior change when
omitted. Also adds a DeprecationWarning to create_langgraph_tracing_handler
and its module docstring, pointing to the unified harness surface, and
updates the sync module docstring with the preferred unified path.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: declan-scale

src/agentex/lib/adk/_modules/_langgraph_sync.py

@@ -3,10 +3,28 @@
 Converts LangGraph graph.astream() events into Agentex TaskMessageUpdate
 events that are yielded back over the HTTP response. For use with sync ACP
 agents that stream via HTTP yields rather than Redis.
+
+Unified sync path
+-----------------
+Prefer using ``LangGraphTurn`` with ``UnifiedEmitter.yield_turn`` for new
+agents, which adds usage capture and optional tracing via the shared harness
+surface::
+
+    from agentex.lib.core.harness.emitter import UnifiedEmitter
+    from agentex.lib.adk._modules._langgraph_turn import LangGraphTurn
+
+    turn = LangGraphTurn(stream)
+    emitter = UnifiedEmitter(task_id=task_id, trace_id=trace_id, parent_span_id=span_id)
+    async for event in emitter.yield_turn(turn):
+        yield event
+
+``convert_langgraph_to_agentex_events`` remains available as a lower-level
+primitive (e.g. for callers that need the raw event stream without the
+harness envelope).
 """
 
 
-async def convert_langgraph_to_agentex_events(stream):
+async def convert_langgraph_to_agentex_events(stream, on_final_ai_message=None):
     """Convert LangGraph streaming events to Agentex TaskMessageUpdate events.
 
     Expects the stream from graph.astream() called with
@@ -22,8 +40,17 @@ async def convert_langgraph_to_agentex_events(stream):
     Supports both regular models (chunk.content is a str) and reasoning models
     like gpt-5/o1/o3 (chunk.content is a list of typed content blocks).
 
+    AGX1-377 note: LangGraph emits tool requests as ``StreamTaskMessageFull`` (from
+    "updates" events), NOT Start+Delta+Done like pydantic-ai. No coalesce_tool_requests
+    option is needed for LangGraph.
+
     Args:
         stream: Async iterator from graph.astream(..., stream_mode=["messages", "updates"])
+        on_final_ai_message: Optional callback ``(msg: AIMessage) -> None`` called for
+            each ``AIMessage`` in an "agent" node update. Use this to capture
+            ``usage_metadata`` for token accounting without re-traversing the stream.
+            The callback fires *after* all events for that message are yielded.
+            No-op when ``None`` (default).
 
     Yields:
         TaskMessageUpdate events (Start, Delta, Done, Full)
@@ -205,6 +232,13 @@ async def convert_langgraph_to_agentex_events(stream):
                                 )
                                 message_index += 1
 
+                        # Notify caller of the final AIMessage (e.g. for usage capture)
+                        if on_final_ai_message is not None:
+                            from langchain_core.messages import AIMessage as _AIMessage
+
+                            if isinstance(msg, _AIMessage):
+                                on_final_ai_message(msg)
+
                 elif node_name == "tools":
                     messages = state_update.get("messages", [])
                     for msg in messages:

src/agentex/lib/adk/_modules/_langgraph_tracing.py

@@ -1,9 +1,20 @@
-"""LangChain callback handler that creates Agentex spans for LLM calls and tool executions."""
+"""LangChain callback handler that creates Agentex spans for LLM calls and tool executions.
+
+.. deprecated::
+    ``AgentexLangGraphTracingHandler`` and ``create_langgraph_tracing_handler`` are
+    superseded by the unified harness surface (``LangGraphTurn`` +
+    ``UnifiedEmitter``), which derives spans automatically from the canonical
+    event stream without requiring a LangChain callback handler.
+
+    They remain importable and functional for backward compatibility, but new
+    agents should use the unified path instead.
+"""
 # ruff: noqa: ARG002
 # Callback methods must accept all arguments defined by LangChain's AsyncCallbackHandler interface.
 
 from __future__ import annotations
 
+import warnings
 from uuid import UUID
 from typing import Any, override
 
@@ -31,6 +42,11 @@ class AgentexLangGraphTracingHandler(AsyncCallbackHandler):
           ├── llm:<model>       (LLM call)
           ├── tool:<tool_name>  (tool execution)
           └── llm:<model>       (LLM call)
+
+    .. deprecated::
+        Use ``LangGraphTurn`` with ``UnifiedEmitter`` instead. The unified
+        harness derives equivalent spans from the canonical event stream,
+        removing the need for a LangChain callback handler entirely.
     """
 
     def __init__(
@@ -237,7 +253,28 @@ def create_langgraph_tracing_handler(
 
     Returns:
         An ``AgentexLangGraphTracingHandler`` instance ready to use as a LangChain callback.
+
+    .. deprecated::
+        Use ``LangGraphTurn`` with ``UnifiedEmitter`` instead. The unified harness
+        derives equivalent spans from the canonical event stream automatically, with
+        no LangChain callback required::
+
+            from agentex.lib.core.harness.emitter import UnifiedEmitter
+            from agentex.lib.adk._modules._langgraph_turn import LangGraphTurn
+
+            turn = LangGraphTurn(stream)
+            emitter = UnifiedEmitter(task_id=task_id, trace_id=trace_id, parent_span_id=span_id)
+            result = await emitter.auto_send_turn(turn)
+
+        This function remains available for backward compatibility.
     """
+    warnings.warn(
+        "create_langgraph_tracing_handler is deprecated. Use LangGraphTurn with "
+        "UnifiedEmitter instead — the unified harness derives equivalent spans from "
+        "the canonical event stream without a LangChain callback handler.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     return AgentexLangGraphTracingHandler(
         trace_id=trace_id,
         parent_span_id=parent_span_id,

tests/lib/adk/test_langgraph_sync.py

@@ -0,0 +1,209 @@
+"""Tests for the sync LangGraph -> Agentex stream event converter.
+
+Covers:
+- Basic text, tool call, and tool response emission
+- on_final_ai_message callback for usage capture
+- Deprecation warning emitted by create_langgraph_tracing_handler
+
+NOTE: langchain_core imports must be deferred to test-function scope because
+conftest.py stubs out ``langchain_core.messages`` with MagicMock for ADK
+package-level tests. The real classes are imported lazily inside each test.
+"""
+
+from __future__ import annotations
+
+import sys
+import warnings
+from typing import Any, AsyncIterator
+
+import pytest
+
+from agentex.types.task_message_update import (
+    StreamTaskMessageFull,
+)
+from agentex.types.tool_request_content import ToolRequestContent
+from agentex.types.tool_response_content import ToolResponseContent
+from agentex.lib.adk._modules._langgraph_sync import convert_langgraph_to_agentex_events
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+async def _collect(stream: AsyncIterator[Any]) -> list[Any]:
+    return [e async for e in stream]
+
+
+def _make_stream(events: list[tuple[str, Any]]) -> AsyncIterator[tuple[str, Any]]:
+    async def _gen():
+        for e in events:
+            yield e
+
+    return _gen()
+
+
+# ---------------------------------------------------------------------------
+# Remove the conftest stubs for langchain_core so real classes are used
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(autouse=True)
+def _real_langchain_core():
+    """Remove conftest MagicMock stubs so real langchain_core types are used."""
+    stub_keys = [k for k in sys.modules if k.startswith("langchain_core") or k.startswith("langgraph")]
+    saved = {k: sys.modules.pop(k) for k in stub_keys}
+    # Re-import the real modules
+    import importlib
+
+    importlib.import_module("langchain_core.messages")
+    yield
+    # Restore stubs after the test
+    sys.modules.update(saved)
+
+
+class TestTextStreaming:
+    async def test_plain_text_emits_start_delta_done(self):
+        from langchain_core.messages import AIMessage, AIMessageChunk
+
+        chunk = AIMessageChunk(content="Hello, world!")
+        events = [
+            ("messages", (chunk, {})),
+            ("updates", {"agent": {"messages": [AIMessage(content="Hello, world!")]}}),
+        ]
+        out = await _collect(convert_langgraph_to_agentex_events(_make_stream(events)))
+        types = [type(e).__name__ for e in out]
+        assert "StreamTaskMessageStart" in types
+        assert "StreamTaskMessageDelta" in types
+        assert "StreamTaskMessageDone" in types
+
+    async def test_empty_chunk_content_is_skipped(self):
+        from langchain_core.messages import AIMessageChunk
+
+        chunk = AIMessageChunk(content="")
+        events = [("messages", (chunk, {}))]
+        out = await _collect(convert_langgraph_to_agentex_events(_make_stream(events)))
+        assert out == []
+
+
+class TestToolCallEmission:
+    async def test_tool_call_emits_full_message(self):
+        from langchain_core.messages import AIMessage
+
+        tc = {"id": "call_1", "name": "get_weather", "args": {"city": "Paris"}}
+        ai_msg = AIMessage(content="", tool_calls=[tc])
+        events = [("updates", {"agent": {"messages": [ai_msg]}})]
+        out = await _collect(convert_langgraph_to_agentex_events(_make_stream(events)))
+        assert len(out) == 1
+        assert isinstance(out[0], StreamTaskMessageFull)
+        content = out[0].content
+        assert isinstance(content, ToolRequestContent)
+        assert content.tool_call_id == "call_1"
+        assert content.name == "get_weather"
+        assert content.arguments == {"city": "Paris"}
+        assert content.author == "agent"
+
+    async def test_tool_response_emits_full_message(self):
+        from langchain_core.messages import ToolMessage
+
+        tool_msg = ToolMessage(content="Sunny, 72F", tool_call_id="call_1", name="get_weather")
+        events = [("updates", {"tools": {"messages": [tool_msg]}})]
+        out = await _collect(convert_langgraph_to_agentex_events(_make_stream(events)))
+        assert len(out) == 1
+        assert isinstance(out[0], StreamTaskMessageFull)
+        content = out[0].content
+        assert isinstance(content, ToolResponseContent)
+        assert content.tool_call_id == "call_1"
+        assert content.name == "get_weather"
+        assert content.content == "Sunny, 72F"
+        assert content.author == "agent"
+
+
+class TestOnFinalAiMessageCallback:
+    async def test_callback_called_for_ai_message_in_agent_node(self):
+        from langchain_core.messages import AIMessage
+
+        captured: list[Any] = []
+        ai_msg = AIMessage(content="Hello!")
+
+        events = [("updates", {"agent": {"messages": [ai_msg]}})]
+        await _collect(convert_langgraph_to_agentex_events(_make_stream(events), on_final_ai_message=captured.append))
+        assert len(captured) == 1
+        assert captured[0] is ai_msg
+
+    async def test_callback_not_called_for_tool_messages(self):
+        from langchain_core.messages import ToolMessage
+
+        captured: list[Any] = []
+        tool_msg = ToolMessage(content="result", tool_call_id="c1", name="t")
+
+        events = [("updates", {"tools": {"messages": [tool_msg]}})]
+        await _collect(convert_langgraph_to_agentex_events(_make_stream(events), on_final_ai_message=captured.append))
+        assert captured == []
+
+    async def test_callback_receives_usage_metadata(self):
+        from langchain_core.messages import AIMessage
+
+        captured: list[Any] = []
+        usage = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
+        ai_msg = AIMessage(content="Answer.", usage_metadata=usage)
+
+        events = [("updates", {"agent": {"messages": [ai_msg]}})]
+        await _collect(convert_langgraph_to_agentex_events(_make_stream(events), on_final_ai_message=captured.append))
+        assert len(captured) == 1
+        assert captured[0].usage_metadata == usage
+
+    async def test_no_callback_is_noop(self):
+        from langchain_core.messages import AIMessage
+
+        ai_msg = AIMessage(content="Hello!")
+        events = [("updates", {"agent": {"messages": [ai_msg]}})]
+        out = await _collect(convert_langgraph_to_agentex_events(_make_stream(events)))
+        assert isinstance(out, list)
+
+    async def test_callback_called_multiple_times_for_multi_step(self):
+        from langchain_core.messages import AIMessage
+
+        captured: list[Any] = []
+        ai_msg_1 = AIMessage(content="Step 1")
+        ai_msg_2 = AIMessage(content="Step 2")
+
+        events = [
+            ("updates", {"agent": {"messages": [ai_msg_1]}}),
+            ("updates", {"agent": {"messages": [ai_msg_2]}}),
+        ]
+        await _collect(convert_langgraph_to_agentex_events(_make_stream(events), on_final_ai_message=captured.append))
+        assert len(captured) == 2
+        assert captured[0] is ai_msg_1
+        assert captured[1] is ai_msg_2
+
+    async def test_callback_called_after_tool_call_events_yielded(self):
+        """The callback fires after all events for that AIMessage are yielded."""
+        from langchain_core.messages import AIMessage
+
+        yield_order: list[str] = []
+
+        async def _gen():
+            tc = {"id": "c1", "name": "t", "args": {}}
+            ai_msg = AIMessage(content="", tool_calls=[tc])
+            yield ("updates", {"agent": {"messages": [ai_msg]}})
+
+        def _cb(msg):
+            yield_order.append("callback")
+
+        async for _ in convert_langgraph_to_agentex_events(_gen(), on_final_ai_message=_cb):
+            yield_order.append("event")
+
+        # The tool call Full event is emitted before the callback fires
+        assert yield_order.index("event") < yield_order.index("callback")
+
+
+class TestDeprecationWarning:
+    def test_create_langgraph_tracing_handler_emits_deprecation_warning(self):
+        from agentex.lib.adk._modules._langgraph_tracing import create_langgraph_tracing_handler
+
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            create_langgraph_tracing_handler(trace_id="t1")
+        assert any(issubclass(warning.category, DeprecationWarning) for warning in w), (
+            "create_langgraph_tracing_handler must emit a DeprecationWarning"
+        )