From 52d241d754e0cd0858bb6cdf4b3bf9e1095954ac Mon Sep 17 00:00:00 2001 From: Charles Yao Date: Mon, 31 Aug 2026 21:36:31 -0700 Subject: [PATCH] feat(dify-builder): streaming create_fix/build/edit_session with subscribe-before-dispatch Co-Authored-By: Claude Opus 4.8 --- api/services/dify_builder/service.py | 130 +++++++++++++++--- .../services/dify_builder/test_service.py | 72 ++++++++++ 2 files changed, 181 insertions(+), 21 deletions(-) diff --git a/api/services/dify_builder/service.py b/api/services/dify_builder/service.py index ebb41b1beba..37516cd79fa 100644 --- a/api/services/dify_builder/service.py +++ b/api/services/dify_builder/service.py @@ -423,15 +423,17 @@ class DifyBuilderService: self._enqueue_fn = enqueue_fn self._subscribe_fn = subscribe_fn or (lambda _sid: None) - def create_fix_session( + def _prepare_fix_session( self, app_id: str, actor: Actor, failed_run_id: str | None = None, checklist_errors: list[ChecklistError] | None = None, model_config: dict | None = None, - ) -> SessionView: - """Port of Go ``CreateFixSession``, extended to record the failed run. + ) -> tuple[str, Action]: + """Shared setup for ``create_fix_session``/``create_fix_session_stream``: + everything through persisting the session (and the failed-run record) + up to -- but not including -- the ``request_fix`` dispatch. ``failed_run_id`` is the id of the **Dify workflow run** that failed (what the frontend has). In fix mode we record it as an immutable @@ -475,16 +477,53 @@ class DifyBuilderService: # advance's diagnose can resolve fc.failed_run_id. self._repo.save_run(s.id, failed_run) - self.dispatch(s.id, Action(kind="request_fix", base_version=1), actor) - return self.get_session_view(s.id, actor) + return s.id, Action(kind="request_fix", base_version=1) - def create_build_session( - self, app_id: str, actor: Actor, goal_text: str, model_config: dict | None = None + def create_fix_session( + self, + app_id: str, + actor: Actor, + failed_run_id: str | None = None, + checklist_errors: list[ChecklistError] | None = None, + model_config: dict | None = None, ) -> SessionView: - """Start a Build session at build.capability_check and dispatch the - initial ``send_goal`` (parallels ``create_fix_session``). The goal is - seeded as the user's opening bubble; the first advance's - ``handle_capability_check`` analyzes it into requirements.""" + """Port of Go ``CreateFixSession``, extended to record the failed run.""" + session_id, action = self._prepare_fix_session(app_id, actor, failed_run_id, checklist_errors, model_config) + if action is not None: + self.dispatch(session_id, action, actor) + return self.get_session_view(session_id, actor) + + def create_fix_session_stream( + self, + app_id: str, + actor: Actor, + failed_run_id: str | None = None, + checklist_errors: list[ChecklistError] | None = None, + model_config: dict | None = None, + ) -> Iterator[str]: + """Streaming counterpart of ``create_fix_session``: subscribes to the + session's progress channel BEFORE dispatching the initial + ``request_fix`` advance, so no progress frames are lost.""" + from services.dify_builder.wiring import stream_advance_frames + + session_id, action = self._prepare_fix_session(app_id, actor, failed_run_id, checklist_errors, model_config) + subscription = self._subscribe_fn(session_id) # BEFORE dispatch + try: + self.dispatch(session_id, action, actor) + except Exception: + if subscription is not None: + subscription.close() + raise + return stream_advance_frames( + asdict(self.get_session_view(session_id, actor)), subscription, expect_advance=True + ) + + def _prepare_build_session( + self, app_id: str, actor: Actor, goal_text: str, model_config: dict | None = None + ) -> tuple[str, Action]: + """Shared setup for ``create_build_session``/``create_build_session_stream``: + everything through persisting the session up to -- but not including -- + the ``send_goal`` dispatch.""" if model_config: validate_model_config(actor.tenant_id, model_config) policy = FeatureService.get_features(actor.tenant_id).skill_learning_policy @@ -498,16 +537,44 @@ class DifyBuilderService: ) items = [UserItem(text=goal_text).to_item(seq=0, at_version=0)] self._repo.create_session(s, fc, items) # assigns s.id, s.version = 1 - self.dispatch(s.id, Action(kind="send_goal", payload={"text": goal_text}, base_version=1), actor) - return self.get_session_view(s.id, actor) + return s.id, Action(kind="send_goal", payload={"text": goal_text}, base_version=1) - def create_edit_session(self, app_id: str, actor: Actor, model_config: dict | None = None) -> SessionView: - """Start an Edit session at edit.capability_check WITHOUT dispatching. - Mock 02-edit.txt:3 -- on open, show history + composer only; do not read - or lock the canvas. The graph read + impact analysis happen on the first - send_edit_goal action, not here. Unlike create_build_session, this takes - no goal and enqueues no advance -- the session simply rests, canvas - editable, until the user sends their change request.""" + def create_build_session( + self, app_id: str, actor: Actor, goal_text: str, model_config: dict | None = None + ) -> SessionView: + """Start a Build session at build.capability_check and dispatch the + initial ``send_goal`` (parallels ``create_fix_session``). The goal is + seeded as the user's opening bubble; the first advance's + ``handle_capability_check`` analyzes it into requirements.""" + session_id, action = self._prepare_build_session(app_id, actor, goal_text, model_config) + if action is not None: + self.dispatch(session_id, action, actor) + return self.get_session_view(session_id, actor) + + def create_build_session_stream( + self, app_id: str, actor: Actor, goal_text: str, model_config: dict | None = None + ) -> Iterator[str]: + """Streaming counterpart of ``create_build_session``: subscribes BEFORE + dispatching the initial ``send_goal`` advance.""" + from services.dify_builder.wiring import stream_advance_frames + + session_id, action = self._prepare_build_session(app_id, actor, goal_text, model_config) + subscription = self._subscribe_fn(session_id) # BEFORE dispatch + try: + self.dispatch(session_id, action, actor) + except Exception: + if subscription is not None: + subscription.close() + raise + return stream_advance_frames( + asdict(self.get_session_view(session_id, actor)), subscription, expect_advance=True + ) + + def _prepare_edit_session( + self, app_id: str, actor: Actor, model_config: dict | None = None + ) -> tuple[str, Action | None]: + """Shared setup for ``create_edit_session``/``create_edit_session_stream``: + creates the row (no dispatch), always returning a ``None`` action.""" if model_config: validate_model_config(actor.tenant_id, model_config) fc = DifyBuilderContext(model_config=model_config or {}) @@ -519,7 +586,28 @@ class DifyBuilderService: current_state=PcState.EDIT_CAPABILITY_CHECK, ) self._repo.create_session(s, fc, []) # empty conversation; assigns s.id, s.version = 1 - return self.get_session_view(s.id, actor) + return s.id, None + + def create_edit_session(self, app_id: str, actor: Actor, model_config: dict | None = None) -> SessionView: + """Start an Edit session at edit.capability_check WITHOUT dispatching. + Mock 02-edit.txt:3 -- on open, show history + composer only; do not read + or lock the canvas. The graph read + impact analysis happen on the first + send_edit_goal action, not here. Unlike create_build_session, this takes + no goal and enqueues no advance -- the session simply rests, canvas + editable, until the user sends their change request.""" + session_id, action = self._prepare_edit_session(app_id, actor, model_config) + if action is not None: + self.dispatch(session_id, action, actor) + return self.get_session_view(session_id, actor) + + def create_edit_session_stream(self, app_id: str, actor: Actor, model_config: dict | None = None) -> Iterator[str]: + """Streaming counterpart of ``create_edit_session``: settle-only -- + ``create_edit_session`` dispatches no advance, so this yields just the + snapshot frame and closes (no subscription).""" + from services.dify_builder.wiring import stream_advance_frames + + session_id, _action = self._prepare_edit_session(app_id, actor, model_config) # _action is None + return stream_advance_frames(asdict(self.get_session_view(session_id, actor)), None, expect_advance=False) def get_session_view(self, session_id: str, actor: Actor) -> SessionView: """Port of Go ``GetSessionView``.""" diff --git a/api/tests/unit_tests/services/dify_builder/test_service.py b/api/tests/unit_tests/services/dify_builder/test_service.py index d66c9adc9cf..97dec23d939 100644 --- a/api/tests/unit_tests/services/dify_builder/test_service.py +++ b/api/tests/unit_tests/services/dify_builder/test_service.py @@ -93,6 +93,35 @@ def service(repo: SqlDifyBuilderRepository, lock: FakeSessionLock, enqueued: lis return DifyBuilderService(repo=repo, session_lock=lock, enqueue_fn=enqueue_fn) +@pytest.fixture +def inmemory_service_factory_raw(repo: SqlDifyBuilderRepository, lock: FakeSessionLock): + """Build a ``DifyBuilderService`` (in-memory SQLite repo + fake lock) with + caller-supplied ``subscribe_fn``/``enqueue_fn``, for tests that need to + observe subscribe/dispatch ordering directly (rather than via the + capturing ``enqueued`` list the plain ``service`` fixture uses).""" + + def _factory(subscribe_fn=None, enqueue_fn=None) -> tuple[DifyBuilderService, Actor]: + svc = DifyBuilderService( + repo=repo, + session_lock=lock, + enqueue_fn=enqueue_fn or (lambda *_a, **_k: None), + subscribe_fn=subscribe_fn, + ) + return svc, _actor() + + return _factory + + +class _StateSub: + """Fake subscription: one terminal ``state`` frame, then closes.""" + + def receive(self, timeout=None): # noqa: ARG002 + return json.dumps({"kind": "state", "version": 2, "session_id": "sid"}).encode() + + def close(self): + pass + + def _actor(account_id: str = ACCOUNT_ID) -> Actor: return Actor(account_id=account_id, tenant_id=TENANT_ID) @@ -480,6 +509,49 @@ def test_submit_action_stream_raises_conflict_before_streaming(repo: SqlDifyBuil assert subscribed == [] # no subscription opened on a pre-flight error +def test_create_fix_session_stream_subscribes_before_dispatch(inmemory_service_factory_raw) -> None: + order: list[str] = [] + svc, actor = inmemory_service_factory_raw( + subscribe_fn=lambda _sid: (order.append("subscribe"), _StateSub())[1], + enqueue_fn=lambda *_a, **_k: order.append("enqueue"), + ) + + gen = svc.create_fix_session_stream( + app_id=APP_ID, actor=actor, failed_run_id="run-1", checklist_errors=None, model_config=None + ) + frames = list(gen) + + assert order == ["subscribe", "enqueue"] + assert frames[0].startswith("event: snapshot") + + +def test_create_build_session_stream_subscribes_before_dispatch(inmemory_service_factory_raw) -> None: + order: list[str] = [] + svc, actor = inmemory_service_factory_raw( + subscribe_fn=lambda _sid: (order.append("subscribe"), _StateSub())[1], + enqueue_fn=lambda *_a, **_k: order.append("enqueue"), + ) + + gen = svc.create_build_session_stream(app_id=APP_ID, actor=actor, goal_text="Build a report workflow") + frames = list(gen) + + assert order == ["subscribe", "enqueue"] + assert frames[0].startswith("event: snapshot") + + +def test_create_edit_session_stream_is_settle_only(inmemory_service_factory_raw) -> None: + subscribed: list[str] = [] + svc, actor = inmemory_service_factory_raw( + subscribe_fn=lambda sid: subscribed.append(sid), enqueue_fn=lambda *_a, **_k: None + ) + + frames = list(svc.create_edit_session_stream(app_id=APP_ID, actor=actor, model_config=None)) + + assert frames == [frames[0]] # snapshot only + assert frames[0].startswith("event: snapshot") + assert subscribed == [] # edit create dispatches no advance, so no subscription + + def test_dispatch_releases_lock_when_enqueue_fails(repo: SqlDifyBuilderRepository) -> None: lock = FakeSessionLock()