bugfix Alex 2026-06-30

tests/unit/test_telemetry_langfuse_compact.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+import pytest
+
+from agent_framework.observability.context import clear_observability_context
+from agent_framework.observability.telemetry import Telemetry
+
+
+class Settings:
+    ENABLE_LANGFUSE = False
+    ENABLE_OTEL = False
+    LANGFUSE_TRACE_MODE = "compact"
+
+
+class FakeObservation:
+    _next_id = 1
+
+    def __init__(self, kwargs):
+        self.kwargs = kwargs
+        self.id = f"obs-{FakeObservation._next_id}"
+        self.trace_id = "trace-123"
+        FakeObservation._next_id += 1
+        self.updates = []
+        self.trace_updates = []
+        self.trace_io_updates = []
+
+    def update(self, **kwargs):
+        self.updates.append(kwargs)
+
+    def update_trace(self, **kwargs):
+        self.trace_updates.append(kwargs)
+
+    def set_trace_io(self, **kwargs):
+        self.trace_io_updates.append(kwargs)
+
+
+class FakeContextManager:
+    def __init__(self, observation):
+        self.observation = observation
+
+    def __enter__(self):
+        return self.observation
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+
+class FakePropagationContext:
+    def __init__(self, owner, kwargs):
+        self.owner = owner
+        self.kwargs = kwargs
+
+    def __enter__(self):
+        self.owner.propagations.append(self.kwargs)
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+
+class FakeLangfuse:
+    def __init__(self, *, legacy_api: bool = False):
+        self.observations = []
+        self.propagations = []
+        self.flush_count = 0
+        self.api = FakeApi() if legacy_api else None
+
+    def start_as_current_observation(self, **kwargs):
+        observation = FakeObservation(kwargs)
+        self.observations.append(observation)
+        return FakeContextManager(observation)
+
+    def propagate_attributes(self, **kwargs):
+        return FakePropagationContext(self, kwargs)
+
+    def flush(self):
+        self.flush_count += 1
+
+
+class FakeIngestionResponse:
+    errors = []
+    successes = []
+
+
+class FakeIngestion:
+    def __init__(self):
+        self.batches = []
+
+    def batch(self, *, batch, metadata=None):
+        self.batches.append({"batch": batch, "metadata": metadata})
+        return FakeIngestionResponse()
+
+
+class FakeApi:
+    def __init__(self):
+        self.ingestion = FakeIngestion()
+
+
+def telemetry_with_fake_langfuse(*, legacy_api: bool = False):
+    FakeObservation._next_id = 1
+    telemetry = Telemetry(Settings())
+    telemetry.enabled = True
+    telemetry.langfuse = FakeLangfuse(legacy_api=legacy_api)
+    return telemetry
+
+
+@pytest.mark.asyncio
+async def test_compact_keeps_root_output_and_shows_only_aga_noc_events():
+    clear_observability_context()
+    telemetry = telemetry_with_fake_langfuse()
+
+    async with telemetry.span("agent.gateway_message", session_id="s1", input={"request": "cms"}, _root_span=True) as span:
+        await telemetry.event("IC.INTERNAL", {"step": "hidden"}, kind="ic")
+        await telemetry.event("NOC.001", {"step": "visible"}, kind="noc")
+        await telemetry.event("AGA.010", {"step": "visible"}, kind="ic")
+        span.set_output({"answer": "ok"})
+
+    names = [obs.kwargs["name"] for obs in telemetry.langfuse.observations]
+    assert names == ["agent.gateway_message", "NOC.001", "AGA.010"]
+
+    root = telemetry.langfuse.observations[0]
+    assert root.updates[-1]["input"] == {"request": "cms"}
+    assert root.updates[-1]["output"] == {"answer": "ok"}
+    assert root.trace_io_updates[-1] == {"input": {"request": "cms"}, "output": {"answer": "ok"}}
+    assert telemetry.langfuse.observations[1].kwargs.get("trace_context") is None
+    assert telemetry.langfuse.observations[2].kwargs.get("trace_context") is None
+    assert telemetry.langfuse.propagations[-1]["trace_name"] == "agent.gateway_message"
+
+    aggregated = root.updates[-1]["metadata"]["aggregated_events"]
+    assert [event["name"] for event in aggregated] == ["IC.INTERNAL", "NOC.001", "AGA.010"]
+
+
+@pytest.mark.asyncio
+async def test_compact_generation_records_io_model_parameters_and_usage_details():
+    clear_observability_context()
+    telemetry = telemetry_with_fake_langfuse()
+
+    async with telemetry.span("agent.gateway_message", session_id="s1", input={"request": "cms"}, _root_span=True):
+        async with telemetry.generation_span(
+            name="llm.test",
+            model="test-model",
+            input=[{"role": "user", "content": "ping"}],
+            metadata={"profile_name": "test"},
+            model_parameters={"temperature": 0.2, "max_tokens": 100},
+        ) as generation:
+            generation.set_output("pong")
+            generation.set_usage({"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2, "cost_usd": 0.01})
+
+    generation = telemetry.langfuse.observations[1]
+    assert generation.kwargs["name"] == "llm.test"
+    assert generation.kwargs["as_type"] == "generation"
+    assert generation.kwargs["input"] == [{"role": "user", "content": "ping"}]
+    assert generation.kwargs["model"] == "test-model"
+    assert generation.kwargs["model_parameters"] == {"temperature": 0.2, "max_tokens": 100}
+    assert "usage" not in generation.kwargs
+    assert "usage_details" not in generation.kwargs
+    assert generation.updates[-1]["input"] == [{"role": "user", "content": "ping"}]
+    assert generation.updates[-1]["output"] == "pong"
+    assert generation.updates[-1]["usage_details"] == {"input": 1, "output": 1}
+    assert generation.updates[-1]["cost_details"] == {"total": 0.01}
+    assert generation.kwargs.get("trace_context") is None
+
+
+@pytest.mark.asyncio
+async def test_legacy_io_fallback_updates_same_root_and_generation_observations():
+    clear_observability_context()
+    telemetry = telemetry_with_fake_langfuse(legacy_api=True)
+
+    async with telemetry.span("agent.gateway_message", session_id="s1", input={"request": "cms"}, _root_span=True) as root:
+        async with telemetry.generation_span(
+            name="llm.test",
+            model="test-model",
+            input=[{"role": "user", "content": "ping"}],
+            model_parameters={"temperature": 0.2},
+        ) as generation:
+            generation.set_output("pong")
+            generation.set_usage({"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5})
+        root.set_output({"answer": "ok"})
+
+    batches = telemetry.langfuse.api.ingestion.batches
+    assert [event.type for item in batches for event in item["batch"]] == ["generation-update", "span-update"]
+
+    generation_event = batches[0]["batch"][0]
+    assert generation_event.body.id == "obs-2"
+    assert generation_event.body.trace_id == "trace-123"
+    assert generation_event.body.input == [{"role": "user", "content": "ping"}]
+    assert generation_event.body.output == "pong"
+    assert generation_event.body.usage_details == {"input": 2, "output": 3}
+
+    root_event = batches[1]["batch"][0]
+    assert root_event.body.id == "obs-1"
+    assert root_event.body.trace_id == "trace-123"
+    assert root_event.body.input == {"request": "cms"}
+    assert root_event.body.output == {"answer": "ok"}
+    assert len(telemetry.langfuse.observations) == 2