feat：接入sse

app/api/agent.py

@@ -2,18 +2,21 @@
 
 from fastapi import APIRouter, Request
 
-from app.common.response import ApiResponse, success_response
+from app.common.sse import AgentEventResponse
-from app.schemas.agent import AgentResult, ResumeRequest
+from app.schemas.agent import ResumeRequest
 from app.schemas.request import DesignRequest
 
 router = APIRouter(prefix="/agent", tags=["Agent"])
 
 
-@router.post("/run", response_model=ApiResponse[AgentResult])
+@router.post("/run", response_class=AgentEventResponse, status_code=200)
-async def run_agent(body: DesignRequest, request: Request):
+async def stream_agent(body: DesignRequest, request: Request):
-    return success_response(await request.app.state.agent_service.run(body))
+    service = request.app.state.agent_service
+    return AgentEventResponse(service, service.stream_run(body), request)
 
 
-@router.post("/resume", response_model=ApiResponse[AgentResult])
+@router.post("/resume", response_class=AgentEventResponse, status_code=200)
-async def resume_agent(body: ResumeRequest, request: Request):
+async def stream_resume_agent(body: ResumeRequest, request: Request):
-    return success_response(await request.app.state.agent_service.resume(body))
+    service = request.app.state.agent_service
+    execution = await service.prepare_resume(body)
+    return AgentEventResponse(service, execution, request)

app/common/sse.py

@@ -0,0 +1,65 @@
+"""SSE 编码和连接生命周期，不保存全局任务队列"""
+
+import json
+from contextlib import aclosing
+
+import anyio
+from starlette.responses import StreamingResponse
+
+from app.common.response import error_content, success_response
+from app.core.exceptions import BizException, handle_unexpected_exception
+
+
+def pack_event(event, content):
+    return f"event: {event}\ndata: {json.dumps(content, ensure_ascii=False)}\n\n"
+
+
+async def encode_events(service, execution, request):
+    try:
+        async with aclosing(service.stream(*execution)) as events:
+            async for event, data in events:
+                yield pack_event(event, success_response(data).model_dump(mode="json"))
+    except BizException as exc:
+        yield pack_event("error", error_content(exc.code, exc.msg))
+    except Exception as exc:  # noqa: BLE001 流中异常统一封装后终止
+        # 流开始后不能修改 HTTP 状态，复用统一异常映射并发送终止事件
+        exc._agent_thread_id = execution[0]
+        response = await handle_unexpected_exception(request, exc)
+        yield pack_event("error", json.loads(response.body))
+
+
+class AgentEventResponse(StreamingResponse):
+    media_type = "text/event-stream"
+
+    def __init__(self, service, execution, request):
+        self.service = service
+        self.execution = execution
+        super().__init__(
+            encode_events(service, execution, request),
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+        )
+
+    async def stream_response(self, send):
+        try:
+            await super().stream_response(send)
+        finally:
+            # 在消费流的同一协程关闭生成器，避免追踪上下文跨协程退出
+            with anyio.CancelScope(shield=True):
+                await self.body_iterator.aclose()
+
+    async def __call__(self, scope, receive, send):
+        try:
+            # 同时监听断开连接，模型尚未产生下一个事件时也能及时取消
+            async with anyio.create_task_group() as group:
+
+                async def stream():
+                    await self.stream_response(send)
+                    group.cancel_scope.cancel()
+
+                group.start_soon(stream)
+                await self.listen_for_disconnect(receive)
+                group.cancel_scope.cancel()
+        finally:
+            # 即使连接在生成器开始前断开，也要释放预先占用的任务
+            self.service.release(self.execution[0])

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(
+            yield (
-                payload,
+                {**config, "callbacks": [handler], "run_name": "run-screen-graph"},
-                config={
+                root,
-                    **config,
-                    "callbacks": [handler],
-                    "run_name": "run-screen-graph",
-                },
             )
-        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]:

app/services/agent.py

@@ -1,12 +1,13 @@
 """异步执行与恢复大屏任务，当前使用单进程内存检查点"""
 
+from contextlib import aclosing
 from uuid import uuid4
 
 from langgraph.types import Command
 from loguru import logger
 
 from app.core.exceptions import BizException
-from app.observability.tracing import ainvoke_screen
+from app.observability.tracing import astream_screen
 from app.schemas.agent import AgentResult, ResumeRequest
 from app.schemas.request import DesignRequest
 
@@ -16,7 +17,7 @@ class AgentService:
         self.graph = graph
         self._active: set[str] = set()
 
-    async def run(self, request: DesignRequest) -> AgentResult:
+    def _prepare_run(self, request: DesignRequest):
         # 新请求始终独立建任务，不用 conversation_id 充当检查点 ID
         thread_id = str(uuid4())
         # 暂无会话创建接口，未传会话 ID 时先分配标识，不创建数据库记录
@@ -24,7 +25,7 @@ class AgentService:
         payload = request.model_dump(mode="json")
         payload["context"]["conversation_id"] = conversation_id
 
-        return await self._invoke(
+        return (
             thread_id,
             {
                 "request": payload,
@@ -35,7 +36,13 @@ class AgentService:
             conversation_id,
         )
 
-    async def resume(self, request: ResumeRequest) -> AgentResult:
+    def stream_run(self, request: DesignRequest):
+        execution = self._prepare_run(request)
+        self._claim(execution[0])
+        return execution
+
+    async def prepare_resume(self, request: ResumeRequest):
+        # 在发送流响应头之前完成校验并占用任务
         thread_id = str(request.thread_id)
         # 检查和占用之间没有 await，同一事件循环内不会重复进入
         self._claim(thread_id)
@@ -71,7 +78,7 @@ class AgentService:
                 )
 
             # 每次有效补充或取消生成新的输入 ID，任务和会话保持不变
-            return await self._execute(
+            return (
                 thread_id,
                 Command(
                     resume=reply.model_dump(mode="json"),
@@ -79,7 +86,11 @@ class AgentService:
                 ),
                 conversation_id,
             )
-        finally:
+        except BaseException:
+            self.release(thread_id)
+            raise
+
+    def release(self, thread_id):
         self._active.discard(thread_id)
 
     def _claim(self, thread_id: str) -> None:
@@ -87,33 +98,8 @@ class AgentService:
             raise BizException(409, "任务正在执行，请勿重复提交")
         self._active.add(thread_id)
 
-    async def _invoke(self, thread_id, payload, session_id) -> AgentResult:
+    @staticmethod
-        self._claim(thread_id)
+    def _result(state, thread_id):
-        try:
-            return await self._execute(thread_id, payload, session_id)
-        finally:
-            # 异常和取消也必须释放占用，不阻塞后续请求
-            self._active.discard(thread_id)
-
-    async def _execute(self, thread_id, payload, session_id) -> AgentResult:
-        # 上下文只作用于当前异步任务，节点自动继承任务标识
-        with logger.contextualize(thread_id=thread_id):
-            logger.info("任务执行开始")
-            try:
-                state = await ainvoke_screen(
-                    self.graph, payload, thread_id=thread_id, session_id=session_id
-                )
-            except Exception as exc:
-                # API 异常处理发生在上下文退出后，保留任务标识用于定位
-                exc._agent_thread_id = thread_id
-                raise
-            logger.info(
-                "任务执行结束，状态 {}",
-                "waiting_for_input"
-                if state.get("__interrupt__")
-                else state.get("status"),
-            )
-
         waiting = bool(state.get("__interrupt__"))
 
         # 只交付校验通过的完整动作批次，失败草稿留在后端
@@ -129,3 +115,21 @@ class AgentService:
             clarification=state.get("clarification") if waiting else None,
             failure_reason=state.get("failure_reason"),
         )
+
+    async def stream(self, thread_id, payload, session_id):
+        """消费图事件，完整状态只用于整理最终结果，不直接暴露给客户端"""
+        state = None
+        with logger.contextualize(thread_id=thread_id):
+            async with aclosing(
+                astream_screen(
+                    self.graph, payload, thread_id=thread_id, session_id=session_id
+                )
+            ) as events:
+                async for mode, data in events:
+                    if mode == "values":
+                        state = data
+                    elif data.get("type") == "node_progress" and data.get("label"):
+                        yield "progress", data
+        if state is None:
+            raise RuntimeError("图未返回状态")
+        yield "final", self._result(state, thread_id).model_dump(mode="json")