feat：接入sse

app/observability/tracing.py

@@ -67,10 +67,43 @@ async def ainvoke_screen(
     session_id 由调用方在首次与恢复时一致传入（通常为 conversation_id）。
     未传时始终使用 thread_id，不从仅首次请求才有的 payload 推断。
     """
+    with _screen_trace(payload, thread_id, session_id) as (config, root):
+        result = await graph.ainvoke(payload, config=config)
+        _finish_trace(root, result)
+        return result
+
+
+async def astream_screen(graph, payload, *, thread_id, session_id=None):
+    """同一轮流保持同一个追踪上下文，关闭流时结束追踪"""
+    from contextlib import aclosing
+
+    result = {}
+    with _screen_trace(payload, thread_id, session_id) as (config, root):
+        async with aclosing(
+            graph.astream(payload, config=config, stream_mode=["custom", "values"])
+        ) as events:
+            async for mode, data in events:
+                if mode == "values":
+                    result = data
+                yield mode, data
+        _finish_trace(root, result)
+
+
+def _finish_trace(root, result):
+    if root is not None:
+        output = _run_output(result)
+        root.update(
+            output=output, level="ERROR" if output["outcome"] == "failed" else "DEFAULT"
+        )
+
+
+@contextmanager
+def _screen_trace(payload, thread_id, session_id):
     config = {"configurable": {"thread_id": thread_id}}
     client = _ensure_client() if tracing_enabled() else None
     if client is None:
-        return await graph.ainvoke(payload, config=config)
+        yield config, None
+        return
 
     from langfuse import propagate_attributes
     from langfuse.langchain import CallbackHandler
@@ -103,15 +136,11 @@ async def ainvoke_screen(
         ),
     ):
         try:
-            result = await graph.ainvoke(
-                payload,
-                config={
-                    **config,
-                    "callbacks": [handler],
-                    "run_name": "run-screen-graph",
-                },
+            yield (
+                {**config, "callbacks": [handler], "run_name": "run-screen-graph"},
+                root,
             )
-        except CancelledError:
+        except (CancelledError, GeneratorExit):
             root.update(output={"outcome": "cancelled"})
             raise
         except Exception as error:
@@ -121,12 +150,6 @@ async def ainvoke_screen(
                 status_message=str(error),
             )
             raise
-        output = _run_output(result)
-        root.update(
-            output=output,
-            level="ERROR" if output["outcome"] == "failed" else "DEFAULT",
-        )
-        return result
 
 
 def _run_output(state: dict[str, Any]) -> dict[str, Any]: