by-framework-langgraph 0.0.3.dev0__tar.gz → 0.0.3.dev2__tar.gz

{by_framework_langgraph-0.0.3.dev0 → by_framework_langgraph-0.0.3.dev2}/.gitignore

@@ -68,6 +68,6 @@ Thumbs.db
 .claude/settings.local.json
 .claude/settings.json
 
-core
-
 .coderfleet-uploads/
+
+.claude/skills

{by_framework_langgraph-0.0.3.dev0 → by_framework_langgraph-0.0.3.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: by-framework-langgraph
-Version: 0.0.3.dev0
+Version: 0.0.3.dev2
 Summary: LangGraph integration for by-framework
 Requires-Python: >=3.12
 Requires-Dist: by-framework>=0.2.1

{by_framework_langgraph-0.0.3.dev0 → by_framework_langgraph-0.0.3.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "by-framework-langgraph"
-version = "0.0.3.dev0"
+version = "0.0.3.dev2"
 description = "LangGraph integration for by-framework"
 readme = "README.md"
 requires-python = ">=3.12"

{by_framework_langgraph-0.0.3.dev0 → by_framework_langgraph-0.0.3.dev2}/src/by_framework_langgraph/adapter.py

@@ -16,7 +16,8 @@ from by_framework.common.logger import logger
 from by_framework.core.protocol.agent_state import AgentState
 from by_framework.core.protocol.commands import ResumeCommand
 from by_framework.core.protocol.events import StreamChunkEvent
-from by_framework.observability.span_recorder import (str_to_uint64, str_to_uint128)
+from by_framework.trace.span_recorder import str_to_uint64, str_to_uint128
+from langchain_core.callbacks import BaseCallbackHandler
 from langchain_core.messages import HumanMessage
 from langgraph.types import Command
 
@@ -31,6 +32,172 @@ if TYPE_CHECKING:
 LANGFUSE_OBSERVATION_ATTR = "_langfuse_observation"
 
 
+class _TokenAccumulatingCallbackHandler(BaseCallbackHandler):
+    """LangChain callback handler that accumulates LLM token usage into AgentContext.
+
+    Extracts token usage from whichever location the provider populates:
+      - ``llm_output["token_usage"]``  (OpenAI-style via LangChain)
+      - ``llm_output["usage"]``         (Anthropic-style / raw provider mapping)
+      - ``generation.message.usage_metadata``  (LangChain >= 0.2 standard)
+      - ``generation.generation_info``  (some community integrations)
+
+    ``run_id`` deduplication prevents double-counting when both ``on_llm_end``
+    and ``on_chat_model_end`` fire for the same call (LangChain >= 0.2).
+    """
+
+    def __init__(self, context: Any) -> None:
+        super().__init__()
+        self._context = context
+        self._seen_run_ids: set = set()
+
+    # ------------------------------------------------------------------
+    # LangChain callback entry point
+    # ------------------------------------------------------------------
+
+    def on_llm_end(self, response: Any, *, run_id: Any = None, **_kwargs: Any) -> None:
+        # Guard: only mark run as seen when we actually extracted tokens so that
+        # the on_chat_model_end event path in _stream_invoke can still fire as a
+        # fallback when the callback found nothing (e.g. stream_options not set).
+        if run_id is not None and run_id in self._seen_run_ids:
+            return
+        self._handle_llm_result(response, run_id=run_id)
+
+    # ------------------------------------------------------------------
+    # Internal extraction logic
+    # ------------------------------------------------------------------
+
+    def _handle_llm_result(self, response: Any, *, run_id: Any = None) -> None:
+        context = self._context
+        if context is None:
+            return
+
+        prompt, completion = self._extract_tokens(response)
+
+        if not (prompt or completion):
+            # Log at WARNING (always visible) to help diagnose providers whose
+            # token format is not yet handled, or where stream_options is missing.
+            gens = getattr(response, "generations", []) or []
+            first = (gens[0] or [None])[0] if gens else None
+            msg_type = type(getattr(first, "message", None)).__name__
+            logger.warning(
+                "[TokenAccumulator] on_llm_end fired but extracted 0 tokens. "
+                "For OpenAI-compatible streaming APIs add "
+                "stream_options={'include_usage': True} to your ChatModel. "
+                "llm_output=%r  first_gen_message_type=%s",
+                getattr(response, "llm_output", None),
+                msg_type,
+            )
+            # Do NOT mark run_id as seen — let the on_chat_model_end event path
+            # in _stream_invoke attempt extraction from the merged message.
+            return
+
+        if run_id is not None:
+            self._seen_run_ids.add(run_id)
+        try:
+            context.record_token_usage(
+                prompt_tokens=prompt,
+                completion_tokens=completion,
+            )
+        except Exception:  # pylint: disable=broad-exception-caught
+            pass
+
+    @staticmethod
+    def _extract_tokens(response: Any) -> tuple[int, int]:
+        """Return (prompt_tokens, completion_tokens) from an LLMResult.
+
+        Checks every known location across providers:
+          1. llm_output["token_usage"]        — OpenAI via LangChain
+          2. llm_output["usage"]              — Anthropic / raw provider mapping
+          3. message.usage_metadata           — LangChain >= 0.2 standard
+          4. message.response_metadata        — some community integrations
+          5. generation.generation_info       — older / custom integrations
+        """
+        prompt, completion = 0, 0
+
+        llm_output = getattr(response, "llm_output", None) or {}
+        if isinstance(llm_output, dict):
+            for key in ("token_usage", "usage"):
+                usage = llm_output.get(key) or {}
+                if usage:
+                    prompt = int(
+                        usage.get("prompt_tokens") or usage.get("input_tokens") or 0
+                    )
+                    completion = int(
+                        usage.get("completion_tokens")
+                        or usage.get("output_tokens")
+                        or 0
+                    )
+                    break
+
+        if prompt or completion:
+            return prompt, completion
+
+        # Iterate all generations
+        for gen_list in getattr(response, "generations", []) or []:
+            for gen in (gen_list if isinstance(gen_list, list) else [gen_list]):
+                msg = getattr(gen, "message", None)
+
+                # LangChain >= 0.2: message.usage_metadata
+                meta = getattr(msg, "usage_metadata", None)
+                if meta:
+                    prompt += int(
+                        meta.get("input_tokens") or meta.get("prompt_tokens") or 0
+                    )
+                    completion += int(
+                        meta.get("output_tokens") or meta.get("completion_tokens") or 0
+                    )
+                    continue
+
+                # response_metadata (e.g. MiniMax, Qwen, some Chinese providers)
+                resp_meta = getattr(msg, "response_metadata", None) or {}
+                if isinstance(resp_meta, dict):
+                    for key in ("token_usage", "usage"):
+                        usage = resp_meta.get(key) or {}
+                        if usage:
+                            prompt += int(
+                                usage.get("prompt_tokens")
+                                or usage.get("input_tokens")
+                                or 0
+                            )
+                            completion += int(
+                                usage.get("completion_tokens")
+                                or usage.get("output_tokens")
+                                or 0
+                            )
+                            break
+                    # Flat keys at root of response_metadata
+                    if not (prompt or completion):
+                        prompt += int(
+                            resp_meta.get("prompt_tokens")
+                            or resp_meta.get("input_tokens")
+                            or 0
+                        )
+                        completion += int(
+                            resp_meta.get("completion_tokens")
+                            or resp_meta.get("output_tokens")
+                            or 0
+                        )
+                    if prompt or completion:
+                        continue
+
+                # generation_info fallback
+                info = getattr(gen, "generation_info", None) or {}
+                for key in ("token_usage", "usage"):
+                    sub = info.get(key) or {}
+                    if sub:
+                        prompt += int(
+                            sub.get("prompt_tokens") or sub.get("input_tokens") or 0
+                        )
+                        completion += int(
+                            sub.get("completion_tokens")
+                            or sub.get("output_tokens")
+                            or 0
+                        )
+                        break
+
+        return prompt, completion
+
+
 @dataclass(frozen=True)
 class _AdapterTracingConfig:
     """Tracing-related adapter config kept separate from core graph handles."""
@@ -167,6 +334,47 @@ class LangGraphAdapter:
                         await self._context.emit_chunk(
                             chunk.content, content_type="text"
                         )
+                elif kind == "on_chat_model_end":
+                    # Fallback: capture token usage from the event's merged output
+                    # message when the on_llm_end callback found nothing (e.g. the
+                    # provider requires stream_options but it was not set).
+                    # _TokenAccumulatingCallbackHandler marks run_id as seen only
+                    # after a successful extraction, so this path fires only when
+                    # the callback got 0 tokens.
+                    run_id = event.get("run_id")
+                    token_handler = next(
+                        (
+                            cb
+                            for cb in scoped_callbacks
+                            if isinstance(cb, _TokenAccumulatingCallbackHandler)
+                        ),
+                        None,
+                    )
+                    if token_handler is not None and run_id not in (
+                        token_handler._seen_run_ids  # pylint: disable=protected-access
+                    ):
+                        output = event["data"].get("output")
+                        meta = getattr(output, "usage_metadata", None)
+                        if meta:
+                            prompt = int(
+                                meta.get("input_tokens")
+                                or meta.get("prompt_tokens")
+                                or 0
+                            )
+                            completion = int(
+                                meta.get("output_tokens")
+                                or meta.get("completion_tokens")
+                                or 0
+                            )
+                            if prompt or completion:
+                                token_handler._seen_run_ids.add(run_id)  # pylint: disable=protected-access
+                                try:
+                                    self._context.record_token_usage(
+                                        prompt_tokens=prompt,
+                                        completion_tokens=completion,
+                                    )
+                                except Exception:  # pylint: disable=broad-exception-caught
+                                    pass
                 elif kind == "on_tool_start":
                     tool_name = event["name"]
                     tool_input = event["data"].get("input")
@@ -184,7 +392,7 @@ class LangGraphAdapter:
                         ),
                     )
 
-                    # Use a stable logical ID from metadata if available, fallback to run_id
+                    # Use stable metadata when available, otherwise fall back.
                     stable_id = (
                         event.get("metadata", {}).get("tool_call_id")
                         or event.get("metadata", {}).get("checkpoint_ns")
@@ -379,69 +587,50 @@ class LangGraphAdapter:
     @contextmanager
     def _langfuse_callback_manager(self, callbacks: list[Any]) -> Iterator[None]:
         """Prepare Langfuse callback and observation for LangChain."""
-        # Prefer AgentContext's callback factory so trace and parent ids align.
-        # Filter out auto-generated MagicMock attributes when tests use a mock
-        # context — real callback objects always come from a non-test module.
-        langfuse_callback_value = getattr(self._context, "langfuse_callback", None)
-        is_real_callback = langfuse_callback_value is not None and type(
-            langfuse_callback_value
-        ).__module__ not in ("unittest.mock",)
-
-        if is_real_callback:
-            handler = (
-                langfuse_callback_value()
-                if callable(langfuse_callback_value)
-                else langfuse_callback_value
-            )
-            if handler is not None:
-                callbacks.append(handler)
-                yield
-                return
+        # Always inject token accumulator — works regardless of Langfuse config.
+        callbacks.append(_TokenAccumulatingCallbackHandler(self._context))
 
-        # Fallback to local import if context method is missing
-        # pylint: disable=import-outside-toplevel
         try:
-            langfuse_config = import_module(
-                "by_framework_trace_langfuse"
-            ).LangfuseConfig
-            if langfuse_config.from_env() is None:
-                raise ImportError("Langfuse not configured")
-
-            callback_handler = import_module("langfuse.langchain").CallbackHandler
-            get_client = import_module("langfuse").get_client
+            build_langchain_callback = getattr(
+                import_module("by_framework_trace_langfuse"),
+                "build_langchain_callback",
+            )
         except (ImportError, AttributeError):
             yield
             return
 
-        callbacks.append(callback_handler())
-
-        framework_observation = getattr(self._context, LANGFUSE_OBSERVATION_ATTR, None)
-        if framework_observation is None:
-            yield
-            return
+        get_parent_observation_id = getattr(
+            self._context,
+            "get_trace_parent_observation_id",
+            None,
+        )
+        parent_observation_id = (
+            str(get_parent_observation_id() or "")
+            if callable(get_parent_observation_id)
+            else ""
+        )
+        if not parent_observation_id:
+            framework_observation = getattr(
+                self._context, LANGFUSE_OBSERVATION_ATTR, None
+            )
+            parent_observation_id = getattr(framework_observation, "id", "") or ""
+        if not parent_observation_id:
+            execution_id = getattr(self._context, "execution_id", "")
+            message_id = getattr(self._context, "message_id", "")
+            raw_parent_id = (
+                f"{execution_id}:worker.execute"
+                if execution_id
+                else f"{message_id}:worker.execute"
+            )
+            parent_observation_id = f"{str_to_uint64(raw_parent_id):016x}"
 
-        langfuse = get_client()
-        with langfuse.start_as_current_observation(
-            as_type="span",
-            name=self._tracing.run_name,
-            trace_context={
-                "trace_id": getattr(self._context, "trace_id", ""),
-                "parent_span_id": framework_observation.id,
-            },
-            metadata=self._default_metadata(),
-        ):
-            # Prevent the generated OTel span from being promoted to a trace root.
-            # The native LangfusePlugin sets the same attribute on its own path
-            # (via _SdkLangfuseTracer); this covers the LangGraph fallback path.
-            try:
-                from opentelemetry import trace
-
-                current_span = trace.get_current_span()
-                if current_span and hasattr(current_span, "set_attribute"):
-                    current_span.set_attribute("langfuse.internal.as_root", False)
-            except Exception:  # pylint: disable=broad-exception-caught
-                pass
-            yield
+        handler = build_langchain_callback(
+            trace_id=getattr(self._context, "trace_id", ""),
+            parent_observation_id=parent_observation_id,
+        )
+        if handler is not None:
+            callbacks.append(handler)
+        yield
 
     @staticmethod
     def _default_input_mapper(content: str) -> dict:

{by_framework_langgraph-0.0.3.dev0 → by_framework_langgraph-0.0.3.dev2}/src/by_framework_langgraph/tools.py

@@ -7,7 +7,7 @@ mechanism.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Annotated
+from typing import TYPE_CHECKING, Annotated, Any
 
 from langchain_core.tools import BaseTool, InjectedToolCallId, tool
 from langgraph.types import interrupt
@@ -16,6 +16,31 @@ if TYPE_CHECKING:
     from by_framework.worker.context import AgentContext
 
 
+def _langfuse_observation_id_from_callbacks(callbacks: Any) -> str:
+    """Return the Langfuse observation id for the active LangChain tool run."""
+    run_id = getattr(callbacks, "run_id", None) or getattr(
+        callbacks,
+        "parent_run_id",
+        None,
+    )
+    if not run_id:
+        return ""
+
+    handlers = [
+        *list(getattr(callbacks, "handlers", []) or []),
+        *list(getattr(callbacks, "inheritable_handlers", []) or []),
+    ]
+    for handler in handlers:
+        runs = getattr(handler, "_runs", None)
+        if not isinstance(runs, dict):
+            continue
+        observation = runs.get(run_id)
+        observation_id = getattr(observation, "id", None)
+        if observation_id:
+            return str(observation_id)
+    return ""
+
+
 def make_remote_agent_tool(
     context: AgentContext,
     tool_name: str,
@@ -50,6 +75,7 @@ def make_remote_agent_tool(
     async def remote_agent_tool(
         topic: str,
         tool_call_id: Annotated[str, InjectedToolCallId],
+        callbacks: Any = None,
     ) -> str:
         # Idempotency guard: checkpoint restore replays tool execution,
         # but we must not re-dispatch the command.
@@ -57,9 +83,19 @@ def make_remote_agent_tool(
         is_dispatched = await context.redis.exists(redis_key)
 
         if not is_dispatched:
+            metadata = {}
+            langfuse_parent_observation_id = _langfuse_observation_id_from_callbacks(
+                callbacks
+            )
+            if langfuse_parent_observation_id:
+                metadata["langfuse_parent_observation_id"] = (
+                    langfuse_parent_observation_id
+                )
+
             await context.call_agent(
                 target_agent_type=target_agent_type,
                 content=topic,
+                metadata=metadata,
             )
             await context.redis.set(redis_key, "1", ex=idempotency_ttl)
 

{by_framework_langgraph-0.0.3.dev0 → by_framework_langgraph-0.0.3.dev2}/tests/test_adapter.py

@@ -1,15 +1,17 @@
 """Tests for adapter and worker modules."""
 
 import sys
-from contextlib import contextmanager
 from types import SimpleNamespace
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from by_framework.core.protocol.commands import AskAgentCommand, ResumeCommand
 from by_framework.core.protocol.message_header import MessageHeader
 
-from by_framework_langgraph.adapter import LangGraphAdapter
+from by_framework_langgraph.adapter import (
+    LangGraphAdapter,
+    _TokenAccumulatingCallbackHandler,
+)
 from by_framework_langgraph.worker import LangGraphWorker
 
 
@@ -37,7 +39,7 @@ def _make_mock_context(session_id: str = "test-session"):
             metadata={"source": "test"},
         )
     )
-    ctx._langfuse_observation = SimpleNamespace(id="obs-framework")  # pylint: disable=protected-access
+    ctx.get_trace_parent_observation_id.return_value = "obs-framework"
     return ctx
 
 
@@ -184,7 +186,7 @@ class TestAdapterRun:
     async def test_includes_langfuse_callbacks_and_parent_trace_context(
         self, monkeypatch
     ):
-        """Verify adapter wires Langfuse callback handler into LangGraph config."""
+        """Verify adapter gets Langfuse callback from the trace provider package."""
         ctx = _make_mock_context()
         graph = MagicMock()
         graph.ainvoke = AsyncMock(
@@ -194,45 +196,18 @@ class TestAdapterRun:
         snapshot.next = ()
         graph.get_state.return_value = snapshot
 
-        observation_calls: list[dict] = []
-        callback_instances: list[object] = []
+        callback_handler = object()
+        callback_calls: list[dict[str, str]] = []
 
-        @contextmanager
-        def fake_observation_scope(**kwargs):
-            observation_calls.append(kwargs)
-            yield SimpleNamespace(id="obs-langgraph")
+        def fake_build_langchain_callback(**kwargs):
+            callback_calls.append(kwargs)
+            return callback_handler
 
-        class FakeCallbackHandler:  # pylint: disable=too-few-public-methods
-            """Minimal callback handler stub for adapter config assertions."""
-
-            def __init__(self):
-                callback_instances.append(self)
-
-        fake_langfuse_client = SimpleNamespace(
-            start_as_current_observation=fake_observation_scope
-        )
-
-        monkeypatch.setitem(
-            sys.modules,
-            "langfuse",
-            SimpleNamespace(get_client=MagicMock(return_value=fake_langfuse_client)),
-        )
-        monkeypatch.setitem(
-            sys.modules,
-            "langfuse.langchain",
-            SimpleNamespace(CallbackHandler=FakeCallbackHandler),
-        )
         monkeypatch.setitem(
             sys.modules,
             "by_framework_trace_langfuse",
-            SimpleNamespace(LangfuseConfig=MagicMock()),
+            SimpleNamespace(build_langchain_callback=fake_build_langchain_callback),
         )
-        sys.modules[
-            "by_framework_trace_langfuse"
-        ].LangfuseConfig.from_env.return_value = object()
-        monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-test")
-        monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-test")
-        monkeypatch.setenv("LANGFUSE_BASE_URL", "http://localhost:3000")
 
         adapter = LangGraphAdapter(graph, ctx, stream=False)
 
@@ -251,21 +226,27 @@ class TestAdapterRun:
         assert config["metadata"]["langfuse_user_id"] == "user-1"
         assert config["metadata"]["by_framework_message_id"] == "msg-ctx"
         assert config["metadata"]["langgraph_thread_id"] == "test-session"
-        assert config["callbacks"] == callback_instances
-        assert len(observation_calls) == 1
-        assert observation_calls[0]["trace_context"] == {
-            "trace_id": "trace-ctx",
-            "parent_span_id": "obs-framework",
-        }
-        assert observation_calls[0]["name"] == "planner:langgraph"
+        # Callbacks list includes the Langfuse handler(s) plus the token accumulator.
+        assert callback_handler in config["callbacks"]
+        assert any(
+            type(cb).__name__ == "_TokenAccumulatingCallbackHandler"
+            for cb in config["callbacks"]
+        )
+        assert callback_calls == [
+            {
+                "trace_id": "trace-ctx",
+                "parent_observation_id": "obs-framework",
+            }
+        ]
 
     @pytest.mark.asyncio
-    async def test_uses_context_langfuse_callback_property(self):
-        """Verify AgentContext.langfuse_callback property value is used directly as a handler."""
-        handler = object()
+    async def test_skips_langfuse_callback_when_provider_package_missing(
+        self, monkeypatch
+    ):
+        """LangGraph stays provider-agnostic when trace-langfuse is not installed."""
 
         # pylint: disable=too-few-public-methods,missing-class-docstring,missing-function-docstring
-        class ContextWithCallbackProperty:
+        class ContextWithoutProvider:
             session_id = "test-session"
             trace_id = "trace-ctx"
             message_id = "msg-ctx"
@@ -283,11 +264,15 @@ class TestAdapterRun:
             def __init__(self):
                 self.emit_chunk = AsyncMock()
 
-            @property
-            def langfuse_callback(self):
-                return handler
+        def fake_import_module(name):
+            if name == "by_framework_trace_langfuse":
+                raise ImportError(name)
+            raise AssertionError(f"unexpected import: {name}")
 
-        ctx = ContextWithCallbackProperty()
+        monkeypatch.setattr(
+            "by_framework_langgraph.adapter.import_module", fake_import_module
+        )
+        ctx = ContextWithoutProvider()
         graph = MagicMock()
         graph.ainvoke = AsyncMock(
             return_value={"messages": [MagicMock(content="hello")]}
@@ -301,7 +286,10 @@ class TestAdapterRun:
 
         assert result == "hello"
         _, kwargs = graph.ainvoke.call_args
-        assert kwargs["config"]["callbacks"] == [handler]
+        callbacks = kwargs["config"]["callbacks"]
+        assert any(
+            type(cb).__name__ == "_TokenAccumulatingCallbackHandler" for cb in callbacks
+        )
 
     @pytest.mark.asyncio
     async def test_skips_langfuse_tracing_silently_when_not_configured(
@@ -323,11 +311,8 @@ class TestAdapterRun:
         monkeypatch.setitem(
             sys.modules,
             "by_framework_trace_langfuse",
-            SimpleNamespace(LangfuseConfig=MagicMock()),
+            SimpleNamespace(build_langchain_callback=MagicMock(return_value=None)),
         )
-        sys.modules[
-            "by_framework_trace_langfuse"
-        ].LangfuseConfig.from_env.return_value = None
 
         adapter = LangGraphAdapter(graph, ctx, stream=False)
 
@@ -340,7 +325,13 @@ class TestAdapterRun:
         assert result == "hello"
         _, kwargs = graph.ainvoke.call_args
         config = kwargs["config"]
-        assert "callbacks" not in config
+        # Token accumulator is always present; only Langfuse callbacks should be absent.
+        non_token_callbacks = [
+            cb
+            for cb in config.get("callbacks", [])
+            if type(cb).__name__ != "_TokenAccumulatingCallbackHandler"
+        ]
+        assert non_token_callbacks == []
         assert "Langfuse" not in caplog.text
 
     @pytest.mark.asyncio
@@ -364,11 +355,8 @@ class TestAdapterRun:
         monkeypatch.setitem(
             sys.modules,
             "by_framework_trace_langfuse",
-            SimpleNamespace(LangfuseConfig=MagicMock()),
+            SimpleNamespace(build_langchain_callback=MagicMock(return_value=None)),
         )
-        sys.modules[
-            "by_framework_trace_langfuse"
-        ].LangfuseConfig.from_env.return_value = None
 
         adapter = LangGraphAdapter(graph, ctx, stream=False)
 
@@ -381,7 +369,12 @@ class TestAdapterRun:
         assert result == "hello"
         _, kwargs = graph.ainvoke.call_args
         config = kwargs["config"]
-        assert "callbacks" not in config
+        non_token_callbacks = [
+            cb
+            for cb in config.get("callbacks", [])
+            if type(cb).__name__ != "_TokenAccumulatingCallbackHandler"
+        ]
+        assert non_token_callbacks == []
         assert "Langfuse" not in caplog.text
 
 
@@ -450,3 +443,75 @@ class TestLangGraphWorkerHooks:  # pylint: disable=too-few-public-methods
         assert captured["run_name"] == "custom-run"
         assert captured["metadata"] == {"team": "alpha"}
         assert captured["callbacks"] == ["cb-1"]
+
+
+class TestTokenAccumulatingCallbackHandler:
+
+    def _make_llm_result(self, prompt=10, completion=20, style="openai"):
+        """Build a mock LLMResult in either openai or usage_metadata style."""
+        result = MagicMock()
+        if style == "openai":
+            result.llm_output = {
+                "token_usage": {
+                    "prompt_tokens": prompt,
+                    "completion_tokens": completion,
+                }
+            }
+            result.generations = []
+        else:
+            result.llm_output = {}
+            gen = MagicMock()
+            gen.message.usage_metadata = {
+                "input_tokens": prompt,
+                "output_tokens": completion,
+            }
+            result.generations = [[gen]]
+        return result
+
+    def test_accumulates_openai_style(self):
+        ctx = MagicMock()
+        handler = _TokenAccumulatingCallbackHandler(ctx)
+        handler.on_llm_end(self._make_llm_result(10, 20, "openai"))
+        ctx.record_token_usage.assert_called_once_with(
+            prompt_tokens=10, completion_tokens=20
+        )
+
+    def test_accumulates_usage_metadata_style(self):
+        ctx = MagicMock()
+        handler = _TokenAccumulatingCallbackHandler(ctx)
+        handler.on_llm_end(self._make_llm_result(5, 15, "metadata"))
+        ctx.record_token_usage.assert_called_once_with(
+            prompt_tokens=5, completion_tokens=15
+        )
+
+    def test_no_call_on_zero_tokens(self):
+        ctx = MagicMock()
+        handler = _TokenAccumulatingCallbackHandler(ctx)
+        result = MagicMock()
+        result.llm_output = {}
+        result.generations = []
+        handler.on_llm_end(result)
+        ctx.record_token_usage.assert_not_called()
+
+    def test_none_context_does_not_raise(self):
+        handler = _TokenAccumulatingCallbackHandler(None)
+        handler.on_llm_end(self._make_llm_result())
+
+    def test_callback_injected_in_tracing_scope(self):
+        """_langfuse_callback_manager injects _TokenAccumulatingCallbackHandler."""
+        ctx = MagicMock()
+        graph = MagicMock()
+        adapter = LangGraphAdapter(graph=graph, context=ctx)
+
+        callbacks = []
+        with patch.object(
+            adapter,
+            "_langfuse_callback_manager",
+            wraps=adapter._langfuse_callback_manager,
+        ):
+            with adapter._langfuse_callback_manager(callbacks):
+                pass
+
+        assert any(
+            isinstance(cb, _TokenAccumulatingCallbackHandler) for cb in callbacks
+        )

{by_framework_langgraph-0.0.3.dev0 → by_framework_langgraph-0.0.3.dev2}/tests/test_tools.py

@@ -1,8 +1,15 @@
 """Tests for tools module."""
 
+from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock
 
-from by_framework_langgraph.tools import (make_ask_user_tool, make_remote_agent_tool)
+import pytest
+
+from by_framework_langgraph.tools import (
+    _langfuse_observation_id_from_callbacks,
+    make_ask_user_tool,
+    make_remote_agent_tool,
+)
 
 
 def _make_mock_context(session_id: str = "test-session"):
@@ -32,6 +39,49 @@ class TestMakeRemoteAgentTool:
         )
         assert "Invoke the poet agent" in tool.description
 
+    def test_resolves_langfuse_observation_id_from_callbacks(self):
+        """The active tool observation can be read from Langfuse callback state."""
+        run_id = object()
+        handler = SimpleNamespace(_runs={run_id: SimpleNamespace(id="obs-tool")})
+        callbacks = SimpleNamespace(parent_run_id=run_id, handlers=[handler])
+
+        assert _langfuse_observation_id_from_callbacks(callbacks) == "obs-tool"
+
+    @pytest.mark.asyncio
+    async def test_passes_tool_observation_id_to_call_agent(self, monkeypatch):
+        """Remote calls are parented to the current LangGraph tool observation."""
+        ctx = _make_mock_context()
+        ctx.redis.exists.return_value = False
+        tool = make_remote_agent_tool(
+            ctx,
+            "query_weather",
+            "weather-agent",
+            "Query weather",
+        )
+        monkeypatch.setattr(
+            "by_framework_langgraph.tools.interrupt",
+            lambda _message: "queued",
+        )
+
+        run_id = object()
+        handler = SimpleNamespace(_runs={run_id: SimpleNamespace(id="obs-tool")})
+        callbacks = SimpleNamespace(parent_run_id=run_id, handlers=[handler])
+        run_manager = SimpleNamespace(get_child=lambda: callbacks)
+
+        result = await tool._arun(
+            "Beijing weather",
+            tool_call_id="tool-call-1",
+            run_manager=run_manager,
+            config={},
+        )
+
+        assert result == "queued"
+        ctx.call_agent.assert_awaited_once_with(
+            target_agent_type="weather-agent",
+            content="Beijing weather",
+            metadata={"langfuse_parent_observation_id": "obs-tool"},
+        )
+
 
 class TestMakeAskUserTool:
     """Tests for make_ask_user_tool."""