Fix Langfuse observation tracing

gaia_agent/observability.py

@@ -45,8 +45,10 @@ def traced_step(trace: Any, name: str, fn: Callable[[], dict[str, Any]]) -> dict
     span = trace.start_observation(name=name, as_type="span")
     try:
         output = fn()
-        span.end(output=output)
+        span.update(output=output)
+        span.end()
         return output
     except Exception as exc:
-        span.end(level="ERROR", status_message=str(exc))
+        span.update(level="ERROR", status_message=str(exc))
+        span.end()
         raise

tests/test_observability.py

@@ -0,0 +1,47 @@
+import pytest
+
+from gaia_agent.observability import traced_step
+
+
+class FakeSpan:
+    def __init__(self):
+        self.updates = []
+        self.ended = False
+
+    def update(self, **kwargs):
+        self.updates.append(kwargs)
+
+    def end(self):
+        self.ended = True
+
+
+class FakeTrace:
+    def __init__(self):
+        self.span = FakeSpan()
+
+    def start_observation(self, name, as_type):
+        self.name = name
+        self.as_type = as_type
+        return self.span
+
+
+def test_traced_step_records_output_then_ends_span():
+    trace = FakeTrace()
+
+    output = traced_step(trace, "step", lambda: {"answer": "Paris"})
+
+    assert output == {"answer": "Paris"}
+    assert trace.name == "step"
+    assert trace.as_type == "span"
+    assert trace.span.updates == [{"output": {"answer": "Paris"}}]
+    assert trace.span.ended is True
+
+
+def test_traced_step_records_error_then_reraises():
+    trace = FakeTrace()
+
+    with pytest.raises(ValueError):
+        traced_step(trace, "step", lambda: (_ for _ in ()).throw(ValueError("bad")))
+
+    assert trace.span.updates == [{"level": "ERROR", "status_message": "bad"}]
+    assert trace.span.ended is True