mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-23 23:13:58 +08:00
resume-by-re-enqueue facade + job_id threading (LE-1446, partial)
This commit is contained in:
@ -364,6 +364,7 @@ async def generate_flow_events(
|
||||
current_user: CurrentActiveUser,
|
||||
flow_name: str | None = None,
|
||||
source_flow_id: uuid.UUID | None = None,
|
||||
job_id: uuid.UUID | str | None = None,
|
||||
) -> None:
|
||||
"""Generate events for flow building process.
|
||||
|
||||
@ -381,7 +382,9 @@ async def generate_flow_events(
|
||||
start_time = time.perf_counter()
|
||||
components_count = 0
|
||||
graph = None
|
||||
run_id = str(uuid.uuid4())
|
||||
# On the durable background path the run_id MUST equal job_id so the HITL
|
||||
# checkpoint (keyed by job_id) is found on resume; foreground runs keep a uuid4.
|
||||
run_id = str(job_id) if job_id is not None else str(uuid.uuid4())
|
||||
try:
|
||||
flow_id_str = str(flow_id)
|
||||
# Create a fresh session for database operations
|
||||
|
||||
@ -575,6 +575,7 @@ async def _stream_event_frames(
|
||||
parsed: ParsedWorkflowRun,
|
||||
current_user: UserRead,
|
||||
source_flow_id: UUID | None = None,
|
||||
job_id: UUID | None = None,
|
||||
) -> AsyncIterator[tuple[bytes, str]]:
|
||||
"""Run a flow via the v1 build-vertex loop, dispatch its events through ``adapter``.
|
||||
|
||||
@ -635,6 +636,7 @@ async def _stream_event_frames(
|
||||
current_user=current_user,
|
||||
flow_name=flow_name,
|
||||
source_flow_id=source_flow_id,
|
||||
job_id=job_id,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
@ -772,6 +774,7 @@ def _default_frame_source_factory(*, request, flow_id, user, adapter, **_extra):
|
||||
background_tasks=fresh_background_tasks,
|
||||
parsed=parsed,
|
||||
current_user=user,
|
||||
job_id=job_id,
|
||||
):
|
||||
if terminal_error_type is not None and event_type == terminal_error_type:
|
||||
errored = True
|
||||
|
||||
@ -266,6 +266,33 @@ class BackgroundExecutionService(Service):
|
||||
await job_service.write_signal(job_id, SignalType.STOP)
|
||||
await self._executor.cancel(str(job_id))
|
||||
|
||||
async def resume_job(self, job_id: UUID, user: UserRead, *, request_id: str, decision: Any) -> bool:
|
||||
"""Carry a human decision back into a SUSPENDED run and re-enqueue it.
|
||||
|
||||
Returns True when the run was accepted for resume, False on a conflict
|
||||
(not suspended, stale request_id, or lost the single-flight flip) — the
|
||||
route maps False to 409. Ownership is enforced by ``_validate``.
|
||||
"""
|
||||
job = await self._validate(job_id, user)
|
||||
if job.status != JobStatus.SUSPENDED:
|
||||
return False
|
||||
pending = (job.job_metadata or {}).get("pending_request_id")
|
||||
if pending is not None and request_id != pending:
|
||||
return False
|
||||
job_service = get_job_service()
|
||||
# Win the single-flight flip BEFORE writing the RESUME signal, so exactly one
|
||||
# RESUME row exists per suspend and a loser never strands a stray decision.
|
||||
if not await job_service.claim_suspended_for_resume(job_id, owner=self._owner):
|
||||
return False
|
||||
await job_service.write_signal(job_id, SignalType.RESUME, {"decision": decision, "request_id": request_id})
|
||||
await self._enqueue(
|
||||
job_id=job_id,
|
||||
flow_id=job.flow_id,
|
||||
request=self._reconstruct_request(job),
|
||||
user=self._user_stub(job.user_id),
|
||||
)
|
||||
return True
|
||||
|
||||
async def _cancel_suspended(self, job_id: UUID, job_service) -> None:
|
||||
await job_service.update_job_status(job_id, JobStatus.CANCELLED, finished_timestamp=True)
|
||||
with contextlib.suppress(Exception):
|
||||
|
||||
@ -563,6 +563,33 @@ class JobService(Service):
|
||||
await session.flush()
|
||||
return result.rowcount == 1
|
||||
|
||||
async def claim_suspended_for_resume(self, job_id: UUID, *, owner: str | None = None) -> bool:
|
||||
"""Atomically flip SUSPENDED->IN_PROGRESS for resume; True iff this caller won.
|
||||
|
||||
The single conditional ``UPDATE ... WHERE status = SUSPENDED`` is the single-use
|
||||
primitive: only one concurrent resume gets ``rowcount == 1``. Going straight to
|
||||
IN_PROGRESS (with a fresh heartbeat) means a racing startup sweep sees a non-QUEUED,
|
||||
live row and cannot double-enqueue it.
|
||||
"""
|
||||
from sqlmodel import update
|
||||
|
||||
now = datetime.now(timezone.utc).isoformat()
|
||||
async with session_scope() as session:
|
||||
job = await session.get(Job, job_id)
|
||||
if job is None:
|
||||
return False
|
||||
merged = {**(job.job_metadata or {})}
|
||||
if owner is not None:
|
||||
merged.update({"owner": owner, "heartbeat_at": now})
|
||||
stmt = (
|
||||
update(Job)
|
||||
.where(Job.job_id == job_id, Job.status == JobStatus.SUSPENDED)
|
||||
.values(status=JobStatus.IN_PROGRESS, job_metadata=merged)
|
||||
)
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
await session.flush()
|
||||
return result.rowcount == 1
|
||||
|
||||
async def queued_workflow_job_ids(self) -> list[UUID]:
|
||||
"""Return the ids of every QUEUED workflow job (for strand recovery)."""
|
||||
async with session_scope() as session:
|
||||
|
||||
139
src/backend/tests/unit/background_execution/test_resume.py
Normal file
139
src/backend/tests/unit/background_execution/test_resume.py
Normal file
@ -0,0 +1,139 @@
|
||||
"""Resume-by-re-enqueue (LE-1446) against a real JobService + facade.
|
||||
|
||||
Slice 1 — the service-layer ``resume_job``: validate SUSPENDED, atomic single-flight
|
||||
flip, write a RESUME signal carrying the human decision, re-enqueue. Staleness and
|
||||
non-SUSPENDED resumes are rejected before any signal/re-enqueue.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langflow.api.v2.adapters import StreamAdapterContext, get_stream_adapter
|
||||
from langflow.services.background_execution.live_bus import InMemoryLiveBus
|
||||
from langflow.services.background_execution.runner import HUMAN_INPUT_REQUIRED_EVENT, JobRunner
|
||||
from langflow.services.database.models.jobs.model import JobStatus, SignalType
|
||||
|
||||
|
||||
def _frame(event_type: str, data: dict) -> tuple[bytes, str]:
|
||||
return (json.dumps({"event": event_type, "data": data}).encode(), event_type)
|
||||
|
||||
|
||||
def _pause_source(payload: dict):
|
||||
async def _source(**_kwargs):
|
||||
yield _frame("add_message", {"text": "working…"})
|
||||
yield (json.dumps(payload).encode(), HUMAN_INPUT_REQUIRED_EVENT)
|
||||
|
||||
return _source
|
||||
|
||||
|
||||
def _noop_factory(**_kwargs):
|
||||
async def _source(**_kw):
|
||||
yield _frame("end", {})
|
||||
|
||||
return _source
|
||||
|
||||
|
||||
class _StubUser:
|
||||
def __init__(self, user_id):
|
||||
self.id = user_id
|
||||
|
||||
|
||||
def _runner(job_service, job_id, source):
|
||||
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)
|
||||
|
||||
|
||||
async def _suspend_a_job(job_service, *, request_id="req-1"):
|
||||
user_id, flow_id, job_id = uuid4(), uuid4(), uuid4()
|
||||
await job_service.create_job(job_id=job_id, flow_id=flow_id, user_id=user_id)
|
||||
await job_service.update_job_metadata(job_id, {"request": {"flow_id": str(flow_id), "stream_protocol": "langflow"}})
|
||||
payload = {"reason": "human_input_required", "request_id": request_id, "options": ["approve", "reject"]}
|
||||
await _runner(job_service, job_id, _pause_source(payload)).run(job_id=job_id, source_kwargs={})
|
||||
return user_id, job_id
|
||||
|
||||
|
||||
async def _facade():
|
||||
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=_noop_factory)
|
||||
await svc.start()
|
||||
return svc
|
||||
|
||||
|
||||
@pytest.mark.real_services
|
||||
@pytest.mark.no_blockbuster
|
||||
async def test_resume_job_writes_resume_signal_with_decision(real_services_job_service) -> None:
|
||||
job_service = real_services_job_service
|
||||
user_id, job_id = await _suspend_a_job(job_service, request_id="req-a")
|
||||
|
||||
svc = await _facade()
|
||||
try:
|
||||
decision = {"choice": "approve", "note": "looks good"}
|
||||
accepted = await svc.resume_job(job_id, _StubUser(user_id), request_id="req-a", decision=decision)
|
||||
assert accepted is True
|
||||
finally:
|
||||
await svc.stop()
|
||||
|
||||
signals = await job_service.unconsumed_signals(job_id)
|
||||
resume = [s for s in signals if s.signal_type == SignalType.RESUME]
|
||||
assert len(resume) == 1
|
||||
assert resume[0].data["decision"] == decision
|
||||
assert resume[0].data["request_id"] == "req-a"
|
||||
|
||||
|
||||
@pytest.mark.real_services
|
||||
@pytest.mark.no_blockbuster
|
||||
async def test_resume_flips_out_of_suspended(real_services_job_service) -> None:
|
||||
job_service = real_services_job_service
|
||||
user_id, job_id = await _suspend_a_job(job_service, request_id="req-b")
|
||||
|
||||
svc = await _facade()
|
||||
try:
|
||||
await svc.resume_job(job_id, _StubUser(user_id), request_id="req-b", decision={"choice": "x"})
|
||||
finally:
|
||||
await svc.stop()
|
||||
|
||||
job = await job_service.get_job_by_job_id(job_id)
|
||||
assert job.status != JobStatus.SUSPENDED
|
||||
|
||||
|
||||
@pytest.mark.real_services
|
||||
@pytest.mark.no_blockbuster
|
||||
async def test_second_resume_for_same_request_is_rejected(real_services_job_service) -> None:
|
||||
job_service = real_services_job_service
|
||||
user_id, job_id = await _suspend_a_job(job_service, request_id="req-c")
|
||||
|
||||
svc = await _facade()
|
||||
try:
|
||||
first = await svc.resume_job(job_id, _StubUser(user_id), request_id="req-c", decision={"choice": "a"})
|
||||
second = await svc.resume_job(job_id, _StubUser(user_id), request_id="req-c", decision={"choice": "b"})
|
||||
assert first is True
|
||||
assert second is False # no longer SUSPENDED → rejected
|
||||
finally:
|
||||
await svc.stop()
|
||||
|
||||
signals = await job_service.unconsumed_signals(job_id)
|
||||
assert len([s for s in signals if s.signal_type == SignalType.RESUME]) == 1 # only one RESUME row
|
||||
|
||||
|
||||
@pytest.mark.real_services
|
||||
@pytest.mark.no_blockbuster
|
||||
async def test_stale_request_id_is_rejected(real_services_job_service) -> None:
|
||||
job_service = real_services_job_service
|
||||
user_id, job_id = await _suspend_a_job(job_service, request_id="req-real")
|
||||
|
||||
svc = await _facade()
|
||||
try:
|
||||
accepted = await svc.resume_job(job_id, _StubUser(user_id), request_id="req-WRONG", decision={"choice": "a"})
|
||||
assert accepted is False
|
||||
finally:
|
||||
await svc.stop()
|
||||
|
||||
signals = await job_service.unconsumed_signals(job_id)
|
||||
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
|
||||
Reference in New Issue
Block a user