resume-by-re-enqueue facade + job_id threading (LE-1446, partial)

This commit is contained in:
cristhianzl
2026-06-15 10:29:11 -03:00
parent 49d9fd4366
commit 484aac4e13
5 changed files with 200 additions and 1 deletions

View File

@ -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

View File

@ -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

View File

@ -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):

View File

@ -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:

View 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