diff --git a/src/backend/base/langflow/services/background_execution/runner.py b/src/backend/base/langflow/services/background_execution/runner.py index 21cb5c1c34..fb28bc46b3 100644 --- a/src/backend/base/langflow/services/background_execution/runner.py +++ b/src/backend/base/langflow/services/background_execution/runner.py @@ -37,7 +37,14 @@ if TYPE_CHECKING: # A frame source yields (sse_frame_bytes, protocol_event_type) pairs. FrameSource = Callable[..., AsyncIterator[tuple[bytes, str]]] -__all__ = ["HUMAN_INPUT_REQUIRED_EVENT", "FrameSource", "JobRunner"] +# Receives the hydrated checkpoint (or None) + the human decision on resume. +ResumeHook = Callable[[Any, Any], Any] + +__all__ = ["HUMAN_INPUT_REQUIRED_EVENT", "FrameSource", "JobRunner", "ResumeHook"] + + +async def _passthrough_resume_hook(checkpoint: Any, decision: Any) -> None: # noqa: ARG001 + return None class JobRunner: @@ -51,11 +58,16 @@ class JobRunner: job_timeout: float | None = None, owner: str | None = None, heartbeat_interval_s: float = 15.0, + resume_hook: ResumeHook | None = None, + checkpoint_store: Any = None, ) -> None: self._jobs = job_service self._bus = live_bus self._adapter = adapter self._frame_source = frame_source + # Producer (LE-1447/1449) supplies the real decision fold-back; default no-op. + self._resume_hook = resume_hook or _passthrough_resume_hook + self._checkpoint_store = checkpoint_store # When set, the whole run is bounded by this wall-clock budget; an overrun # surfaces as asyncio.TimeoutError, which execute_with_status maps to # TIMED_OUT. None means unbounded (the prior behaviour). @@ -190,6 +202,7 @@ class JobRunner: async def _drive(self, *, job_id: UUID, source_kwargs: dict[str, Any]) -> None: """The wrapped coroutine: stream frames, persist, publish, finalize result/error.""" + await self._maybe_resume(job_id) last_durable_seq = 0 errored_payload: dict[str, Any] | None = None async for frame_bytes, event_type in self._frame_source(**source_kwargs): @@ -231,14 +244,34 @@ class JobRunner: raise RuntimeError(msg) await self._jobs.set_result(job_id, {"status": "completed"}) + async def _maybe_resume(self, job_id: UUID) -> None: + data = await self._pending_resume(job_id) + if data is None: + return + checkpoint = await self._load_checkpoint(job_id) + await self._resume_hook(checkpoint, data.get("decision")) + await self._jobs.consume_signals(job_id, SignalType.RESUME) + + async def _pending_resume(self, job_id: UUID) -> dict[str, Any] | None: + for signal in await self._jobs.unconsumed_signals(job_id): + if signal.signal_type == SignalType.RESUME: + return signal.data or {} + return None + + async def _load_checkpoint(self, job_id: UUID) -> Any: + store = self._checkpoint_store + if store is None: + from lfx.services.deps import get_checkpoint_service + + store = get_checkpoint_service() + return await store.load_by_run_id(str(job_id)) + async def _stop_requested(self, job_id: UUID) -> bool: signals = await self._jobs.unconsumed_signals(job_id) return any(s.signal_type == SignalType.STOP for s in signals) @staticmethod def _user_cancelled() -> asyncio.CancelledError: - # Tagging the cancel as user-initiated makes execute_with_status write - # CANCELLED (not FAILED). Same convention the queue service uses. exc = asyncio.CancelledError() exc.args = ("LANGFLOW_USER_CANCELLED",) return exc diff --git a/src/backend/tests/unit/background_execution/test_resume.py b/src/backend/tests/unit/background_execution/test_resume.py index d7a4c68103..8c422d3e00 100644 --- a/src/backend/tests/unit/background_execution/test_resume.py +++ b/src/backend/tests/unit/background_execution/test_resume.py @@ -137,3 +137,187 @@ async def test_stale_request_id_is_rejected(real_services_job_service) -> None: assert not [s for s in signals if s.signal_type == SignalType.RESUME] # no signal written for a stale id job = await job_service.get_job_by_job_id(job_id) assert job.status == JobStatus.SUSPENDED # untouched + + +# --------------------------------------------------------------------------- # +# Slice 3 — runner RESUME consume + checkpoint load + injected resume hook. +# --------------------------------------------------------------------------- # + + +def _durable_store(job_service): + from langflow.services.checkpoint.store import JobScopedCheckpointStore + + return JobScopedCheckpointStore(job_service) + + +def _graph_checkpoint(job_id): + from lfx.graph.checkpoint.schema import GraphCheckpoint + + return GraphCheckpoint( + run_id=str(job_id), + job_id=str(job_id), + flow_id="flow-1", + session_id="sess-1", + flow_payload={"nodes": [], "edges": []}, + vertices_to_run={"a"}, + ) + + +def _post_resume_source(): + async def _source(**_kwargs): + yield _frame("add_message", {"text": "resumed"}) + yield _frame("end", {}) + + return _source + + +def _resume_runner(job_service, job_id, *, resume_hook, store, source=None): + adapter = get_stream_adapter("langflow", StreamAdapterContext(run_id=str(job_id), thread_id="t")) + return JobRunner( + job_service=job_service, + live_bus=InMemoryLiveBus(), + adapter=adapter, + frame_source=source or _post_resume_source(), + resume_hook=resume_hook, + checkpoint_store=store, + ) + + +@pytest.mark.real_services +@pytest.mark.no_blockbuster +async def test_runner_consumes_resume_and_hands_checkpoint_and_decision_to_hook(real_services_job_service) -> None: + from langflow.services.database.models.jobs.model import SignalType + + job_service = real_services_job_service + job_id, flow_id = uuid4(), uuid4() + await job_service.create_job(job_id=job_id, flow_id=flow_id, user_id=uuid4()) + store = _durable_store(job_service) + await store.save(_graph_checkpoint(job_id)) + decision = {"choice": "approve"} + await job_service.write_signal(job_id, SignalType.RESUME, {"decision": decision, "request_id": "req-h"}) + + received: list[tuple] = [] + + async def hook(checkpoint, decision_arg): + received.append((checkpoint, decision_arg)) + + await _resume_runner(job_service, job_id, resume_hook=hook, store=store).run(job_id=job_id, source_kwargs={}) + + assert len(received) == 1 + checkpoint, decision_arg = received[0] + assert checkpoint is not None + assert str(checkpoint.job_id) == str(job_id) # loaded by job_id (run_id carries it) + assert decision_arg == decision + + +@pytest.mark.real_services +@pytest.mark.no_blockbuster +async def test_resume_signal_consumed_exactly_once(real_services_job_service) -> None: + from langflow.services.database.models.jobs.model import SignalType + + job_service = real_services_job_service + job_id, flow_id = uuid4(), uuid4() + await job_service.create_job(job_id=job_id, flow_id=flow_id, user_id=uuid4()) + store = _durable_store(job_service) + await store.save(_graph_checkpoint(job_id)) + await job_service.write_signal(job_id, SignalType.RESUME, {"decision": {"x": 1}, "request_id": "req-once"}) + + async def hook(_checkpoint, _decision): + pass + + await _resume_runner(job_service, job_id, resume_hook=hook, store=store).run(job_id=job_id, source_kwargs={}) + + unconsumed = await job_service.unconsumed_signals(job_id) + assert not [s for s in unconsumed if s.signal_type == SignalType.RESUME] + + +@pytest.mark.real_services +@pytest.mark.no_blockbuster +async def test_fresh_run_does_not_call_resume_hook(real_services_job_service) -> None: + job_service = real_services_job_service + job_id, flow_id = uuid4(), uuid4() + await job_service.create_job(job_id=job_id, flow_id=flow_id, user_id=uuid4()) + store = _durable_store(job_service) + + called = [] + + async def hook(_checkpoint, _decision): + called.append(1) + + await _resume_runner(job_service, job_id, resume_hook=hook, store=store).run(job_id=job_id, source_kwargs={}) + + assert called == [] + + +# --------------------------------------------------------------------------- # +# Slice 4 — same-stream seq continuation + no orphan-sweep double-run. +# --------------------------------------------------------------------------- # + + +def _resumed_add_message_factory(**_kwargs): + async def _source(**_kw): + yield _frame("add_message", {"text": "post-resume"}) + yield _frame("end", {}) + + return _source + + +async def _facade_with(factory): + from langflow.services.background_execution.service import BackgroundExecutionService + from langflow.services.deps import get_settings_service + + svc = BackgroundExecutionService(settings_service=get_settings_service(), frame_source_factory=factory) + await svc.start() + return svc + + +async def _wait_for_seq_gt(job_service, job_id, threshold, *, attempts=100): + import asyncio + + for _ in range(attempts): + events = await job_service.read_events(job_id, after_seq=0) + if events and max(e.seq for e in events) > threshold: + return events + await asyncio.sleep(0.05) + return await job_service.read_events(job_id, after_seq=0) + + +@pytest.mark.real_services +@pytest.mark.no_blockbuster +async def test_resume_continues_same_event_stream(real_services_job_service) -> None: + """Post-resume frames extend the same (job_id, seq) sequence: strictly greater, no reset.""" + job_service = real_services_job_service + user_id, job_id = await _suspend_a_job(job_service, request_id="req-seq") + pre = await job_service.read_events(job_id, after_seq=0) + pre_max = max(e.seq for e in pre) + + svc = await _facade_with(_resumed_add_message_factory) + try: + await svc.resume_job(job_id, _StubUser(user_id), request_id="req-seq", decision={"choice": "go"}) + events = await _wait_for_seq_gt(job_service, job_id, pre_max) + finally: + await svc.stop() + + seqs = sorted(e.seq for e in events) + assert seqs == list(range(1, len(seqs) + 1)) # gap-free, no reset + assert max(seqs) > pre_max # post-resume frames strictly extend the stream + post = [e for e in events if e.seq > pre_max and e.event_type == "add_message"] + assert post + assert post[0].payload["data"]["text"] == "post-resume" + + +@pytest.mark.real_services +@pytest.mark.no_blockbuster +async def test_resume_does_not_leave_a_queued_row_for_the_sweep(real_services_job_service) -> None: + """Resume flips SUSPENDED->IN_PROGRESS, so the startup QUEUED re-enqueue can't double-run it.""" + job_service = real_services_job_service + user_id, job_id = await _suspend_a_job(job_service, request_id="req-nosweep") + + svc = await _facade_with(_noop_factory) + try: + await svc.resume_job(job_id, _StubUser(user_id), request_id="req-nosweep", decision={"choice": "go"}) + # The sweep only re-enqueues QUEUED workflow rows; a resumed row must not appear there. + queued_ids = await job_service.queued_workflow_job_ids() + assert job_id not in queued_ids + finally: + await svc.stop()