diff --git a/api/controllers/console/workflow_copilot.py b/api/controllers/console/workflow_copilot.py index 89d463bf4bc..8c0cb87f8bd 100644 --- a/api/controllers/console/workflow_copilot.py +++ b/api/controllers/console/workflow_copilot.py @@ -30,7 +30,7 @@ argument. import functools from collections.abc import Callable -from flask import request +from flask import Response, request from flask_restx import Resource from controllers.console.wraps import ( @@ -42,7 +42,13 @@ from controllers.console.wraps import ( from core.workflow_copilot.models import Action, Actor from libs.login import login_required from services.feature_service import FeatureService -from services.workflow_copilot.wiring import build_service, copilot_error_response, session_view_to_dict +from services.workflow_copilot import progress_bus +from services.workflow_copilot.wiring import ( + build_service, + copilot_error_response, + session_view_to_dict, + stream_frames, +) from . import console_ns @@ -159,3 +165,29 @@ class WorkflowCopilotMessagesApi(Resource): @workflow_copilot_required def post(self, current_tenant_id, current_user, session_id): return _message(session_id, request.get_json(silent=True), _actor(current_user, current_tenant_id)) + + +@console_ns.route("/workflow-copilot/sessions//stream") +class WorkflowCopilotStreamApi(Resource): + @setup_required + @login_required + @account_initialization_required + @with_current_user + @with_current_tenant_id + @workflow_copilot_required + def get(self, current_tenant_id, current_user, session_id): + actor = _actor(current_user, current_tenant_id) + # owner check + authoritative snapshot; non-owner/missing → NotFoundError → 404 (generic) + try: + view = build_service().get_session_view(session_id, actor) + except Exception as exc: # re-raised below if not a known copilot error + mapped = copilot_error_response(exc) + if mapped is None: + raise + return mapped + subscription = progress_bus.subscribe(session_id) + return Response( + stream_frames(session_view_to_dict(view), subscription), + mimetype="text/event-stream", + headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"}, + ) diff --git a/api/services/workflow_copilot/wiring.py b/api/services/workflow_copilot/wiring.py index f4cf42d83c5..9e1788eddec 100644 --- a/api/services/workflow_copilot/wiring.py +++ b/api/services/workflow_copilot/wiring.py @@ -5,18 +5,25 @@ controller module so it is unit-testable without the Flask request stack. """ import dataclasses +import json +import time +from collections.abc import Iterator from sqlalchemy.orm import sessionmaker from core.workflow_copilot.errors import BusyError, ConflictError, NotFoundError from core.workflow_copilot.models import Action, Actor from extensions.ext_database import db +from libs.broadcast_channel.exc import SubscriptionClosedError from services.workflow_copilot import session_lock from services.workflow_copilot.repository import SqlCopilotRepository from services.workflow_copilot.service import SessionView, WorkflowCopilotService from tasks.workflow_copilot_advance_task import advance_session -__all__ = ["build_service", "copilot_error_response", "session_view_to_dict"] +__all__ = ["build_service", "copilot_error_response", "session_view_to_dict", "stream_frames"] + +_MAX_STREAM_SECONDS = 180 +_HEARTBEAT_SECONDS = 15 def _enqueue(session_id: str, action: Action, actor: Actor, token: str) -> None: @@ -42,3 +49,30 @@ def copilot_error_response(exc: Exception) -> tuple[dict, int] | None: if isinstance(exc, BusyError): return {"code": "session_busy"}, 409 return None + + +def stream_frames(view_dict: dict, subscription) -> Iterator[str]: + """Yield the snapshot frame, then relay progress events as SSE frames with a + heartbeat, bounded by _MAX_STREAM_SECONDS. Always closes the subscription. + + Frame format matches the frontend parser + the retired Go transport: + `event: \ndata: \n\n`. + """ + yield f"event: snapshot\ndata: {json.dumps(view_dict)}\n\n" + deadline = time.monotonic() + _MAX_STREAM_SECONDS + try: + while time.monotonic() < deadline: + try: + raw = subscription.receive(timeout=_HEARTBEAT_SECONDS) + except SubscriptionClosedError: + return + if raw is None: # receive timed out → heartbeat comment frame + yield ": keep-alive\n\n" + continue + try: + kind = json.loads(raw).get("kind", "message") + except (ValueError, TypeError): + kind = "message" + yield f"event: {kind}\ndata: {raw.decode()}\n\n" + finally: + subscription.close() diff --git a/api/tests/unit_tests/services/workflow_copilot/test_stream_frames.py b/api/tests/unit_tests/services/workflow_copilot/test_stream_frames.py new file mode 100644 index 00000000000..6e67a761f29 --- /dev/null +++ b/api/tests/unit_tests/services/workflow_copilot/test_stream_frames.py @@ -0,0 +1,43 @@ +import json + +from libs.broadcast_channel.exc import SubscriptionClosedError +from services.workflow_copilot import wiring + + +class _FakeSub: + def __init__(self, items): + self._items = list(items) # bytes | None entries; None => a receive() timeout + self.closed = False + self.timeouts: list[float | None] = [] + + def receive(self, timeout=None): + self.timeouts.append(timeout) + if self._items: + return self._items.pop(0) + raise SubscriptionClosedError("closed") + + def close(self): + self.closed = True + + +def test_stream_frames_snapshot_heartbeat_events_then_close(): + view = {"session_id": "s1", "state": "fix.verify"} + node = json.dumps({"kind": "node", "node_id": "output", "status": "running"}).encode() + state = json.dumps({"kind": "state", "version": 5, "state": "success"}).encode() + sub = _FakeSub([node, None, state]) # event, timeout(heartbeat), event, then closed + frames = list(wiring.stream_frames(view, sub)) + assert frames[0] == f"event: snapshot\ndata: {json.dumps(view)}\n\n" + assert frames[1] == f"event: node\ndata: {node.decode()}\n\n" + assert frames[2] == ": keep-alive\n\n" + assert frames[3] == f"event: state\ndata: {state.decode()}\n\n" + assert sub.closed is True # generator closed the subscription in finally + + +def test_stream_frames_malformed_json_falls_back_to_message_kind(): + view = {"session_id": "s1", "state": "fix.verify"} + malformed = b"not-json" + sub = _FakeSub([malformed]) + frames = list(wiring.stream_frames(view, sub)) + assert frames[0] == f"event: snapshot\ndata: {json.dumps(view)}\n\n" + assert frames[1] == f"event: message\ndata: {malformed.decode()}\n\n" + assert sub.closed is True