From b0cf060966b1ce8fa3e76fcde080bc2c9a6be038 Mon Sep 17 00:00:00 2001 From: cristhianzl Date: Tue, 23 Jun 2026 09:27:43 -0300 Subject: [PATCH] gh suggestions fixes --- .../a1f4c9d27b30_add_job_checkpoints_table.py | 3 +- src/backend/base/langflow/api/build.py | 18 ++-- .../langflow/services/checkpoint/store.py | 10 +++ .../services/database/models/jobs/model.py | 2 +- .../api/test_resume_rerun_predecessors.py | 83 +++++++++++++------ .../background_execution/test_service.py | 3 + .../controllers/API/agui/run-flow-bridge.ts | 3 +- src/lfx/src/lfx/graph/checkpoint/resume.py | 5 ++ src/lfx/src/lfx/graph/graph/base.py | 4 + .../unit/graph/checkpoint/test_resume.py | 13 +++ 10 files changed, 107 insertions(+), 37 deletions(-) diff --git a/src/backend/base/langflow/alembic/versions/a1f4c9d27b30_add_job_checkpoints_table.py b/src/backend/base/langflow/alembic/versions/a1f4c9d27b30_add_job_checkpoints_table.py index ea98d88171..d200780b9a 100644 --- a/src/backend/base/langflow/alembic/versions/a1f4c9d27b30_add_job_checkpoints_table.py +++ b/src/backend/base/langflow/alembic/versions/a1f4c9d27b30_add_job_checkpoints_table.py @@ -39,7 +39,8 @@ def upgrade() -> None: sa.UniqueConstraint("job_id", "kind", name="uq_job_checkpoints_job_id_kind"), ) with op.batch_alter_table("job_checkpoints", schema=None) as batch_op: - batch_op.create_index(batch_op.f("ix_job_checkpoints_id"), ["id"], unique=False) + # No index on ``id``: the PRIMARY KEY already provides a unique index, so a separate + # ix_job_checkpoints_id would be dead duplicate DDL. batch_op.create_index(batch_op.f("ix_job_checkpoints_job_id"), ["job_id"], unique=False) diff --git a/src/backend/base/langflow/api/build.py b/src/backend/base/langflow/api/build.py index 67ef70f54e..75aaca6305 100644 --- a/src/backend/base/langflow/api/build.py +++ b/src/backend/base/langflow/api/build.py @@ -84,13 +84,16 @@ def _output_meta_for_vertex(graph: Graph, vertex_id: str) -> dict: def _rerun_non_input_predecessors(graph: Graph, vertex_id: str) -> None: - """Un-build the paused vertex's non-input predecessors so they re-run on resume. + """Un-build only the upstream producers whose live output was dropped from the checkpoint. - A checkpoint cannot serialize non-JSON outputs (Tools, models), so a producer like - an Agent would receive ``None`` tools after restore. Re-running the upstream - definitions regenerates valid inputs; input vertices (e.g. Chat Input) keep their - restored value and are not re-run. + A checkpoint cannot serialize non-JSON outputs (Tools, model clients), so a producer of one of + those must re-run on resume to regenerate it. Producers whose output round-tripped (e.g. an + Agent's ``Message``) keep their restored result — re-running them would re-bill the LLM call and + re-fire side-effecting tools. The dropped set is computed at restore; if it's empty (no live + objects were dropped) nothing re-runs. Input vertices (e.g. Chat Input) are never re-run. + The whole predecessor chain is still walked so a dropped producer behind a restored one is found. """ + dropped: set[str] = getattr(graph, "checkpoint_opaque_dropped_ids", set()) visited: set[str] = set() stack = list(graph.predecessor_map.get(vertex_id, [])) while stack: @@ -102,9 +105,8 @@ def _rerun_non_input_predecessors(graph: Graph, vertex_id: str) -> None: pred = graph.get_vertex(pred_id) except ValueError: continue - if pred.is_input: - continue - pred.built = False + if not pred.is_input and pred_id in dropped: + pred.built = False stack.extend(graph.predecessor_map.get(pred_id, [])) diff --git a/src/backend/base/langflow/services/checkpoint/store.py b/src/backend/base/langflow/services/checkpoint/store.py index b7894494a2..6136e8af07 100644 --- a/src/backend/base/langflow/services/checkpoint/store.py +++ b/src/backend/base/langflow/services/checkpoint/store.py @@ -103,6 +103,16 @@ class JobScopedCheckpointStore(CheckpointStore, Service): return [cp for _job_id, cp in await self._all() if cp.session_id == session_id and not _expired(cp)] async def _all(self) -> list[tuple[UUID, GraphCheckpoint]]: + """Decode every stored ``graph`` checkpoint across ALL jobs/users. + + This scan is NOT tenant-scoped (the row has no user_id; the store has no + request user). The callers that use it — ``load(checkpoint_id)``, + ``delete(checkpoint_id)``, ``list_by_session`` — must therefore never be + exposed to a cross-user request: the durable multi-tenant resume path uses + only ``load_by_run_id`` (job-scoped) behind ``resume_job``'s ownership + check. Adding a user filter here would require threading the user through + the CheckpointStore ABC. + """ decoded = [] for job_id, blob in await self._jobs.all_checkpoints(_KIND): checkpoint = _decode(blob) diff --git a/src/backend/base/langflow/services/database/models/jobs/model.py b/src/backend/base/langflow/services/database/models/jobs/model.py index 8c0403cd3b..1e3216311f 100644 --- a/src/backend/base/langflow/services/database/models/jobs/model.py +++ b/src/backend/base/langflow/services/database/models/jobs/model.py @@ -174,7 +174,7 @@ class JobCheckpoint(SQLModel, table=True): # type: ignore[call-arg] __tablename__ = "job_checkpoints" __table_args__ = (UniqueConstraint("job_id", "kind", name="uq_job_checkpoints_job_id_kind"),) - id: UUID = Field(default_factory=uuid4, primary_key=True, index=True) + id: UUID = Field(default_factory=uuid4, primary_key=True) job_id: UUID = Field(index=True, nullable=False) kind: str = Field(sa_column=Column(String, nullable=False)) blob: str = Field(sa_column=Column(Text, nullable=False)) diff --git a/src/backend/tests/unit/api/test_resume_rerun_predecessors.py b/src/backend/tests/unit/api/test_resume_rerun_predecessors.py index dc6c444860..32b6513149 100644 --- a/src/backend/tests/unit/api/test_resume_rerun_predecessors.py +++ b/src/backend/tests/unit/api/test_resume_rerun_predecessors.py @@ -1,8 +1,9 @@ -"""``_rerun_non_input_predecessors`` (HITL resume): regenerate dropped upstream outputs. +"""``_rerun_non_input_predecessors`` (HITL resume): regenerate only the dropped upstream outputs. -A graph checkpoint cannot serialize non-JSON outputs (Tools, models), so on resume the -paused producer's non-input predecessors must re-run to rebuild valid inputs. Input -vertices (Chat Input) keep their restored value and are not re-run. +A graph checkpoint cannot serialize non-JSON outputs (Tools, model clients), so on resume the +producers of those — and only those — must re-run to rebuild valid inputs. Producers whose output +round-tripped (e.g. an Agent's Message) keep their restored result, so resume does not re-bill an +LLM call or re-fire side-effecting tools. Input vertices (Chat Input) are never re-run. """ from __future__ import annotations @@ -12,7 +13,7 @@ from types import SimpleNamespace from langflow.api.build import _rerun_non_input_predecessors -def _graph(vertices, predecessor_map): +def _graph(vertices, predecessor_map, dropped=None): by_id = {v.id: v for v in vertices} def get_vertex(vid): @@ -20,46 +21,76 @@ def _graph(vertices, predecessor_map): raise ValueError(vid) return by_id[vid] - return SimpleNamespace(predecessor_map=predecessor_map, get_vertex=get_vertex) + return SimpleNamespace( + predecessor_map=predecessor_map, + get_vertex=get_vertex, + checkpoint_opaque_dropped_ids=set(dropped or set()), + ) def _vertex(vid, *, is_input=False, built=True): return SimpleNamespace(id=vid, is_input=is_input, built=built) -def test_reruns_non_input_predecessors_and_keeps_inputs_built(): +def test_reruns_only_dropped_predecessor_and_keeps_inputs_built(): chat = _vertex("chat", is_input=True) tool = _vertex("tool") agent = _vertex("agent") - graph = _graph([chat, tool, agent], {"agent": ["chat", "tool"], "tool": [], "chat": []}) - - _rerun_non_input_predecessors(graph, "agent") - - assert tool.built is False # non-input predecessor (a tool) re-runs - assert chat.built is True # input vertex keeps its restored value - assert agent.built is True # the paused vertex is un-built by the caller, not here - - -def test_walks_predecessors_transitively(): - chat = _vertex("chat", is_input=True) - mid = _vertex("mid") - tool = _vertex("tool") - agent = _vertex("agent") + # The tool's live output (a Tool object) was opaque-dropped; the agent's Message round-tripped. graph = _graph( - [chat, mid, tool, agent], - {"agent": ["tool"], "tool": ["mid"], "mid": ["chat"], "chat": []}, + [chat, tool, agent], + {"agent": ["chat", "tool"], "tool": [], "chat": []}, + dropped={"tool"}, ) _rerun_non_input_predecessors(graph, "agent") - assert tool.built is False - assert mid.built is False + assert tool.built is False # dropped live output → re-run to regenerate it + assert chat.built is True # input vertex keeps its restored value + assert agent.built is True # the paused vertex is un-built by the caller, not here + + +def test_does_not_rerun_producer_whose_output_round_tripped(): + # Agent -> HumanInput: the Agent's output is a serializable Message (not dropped), so resume must + # NOT re-run it — that would re-bill the LLM call and re-fire any side-effecting tools. + chat = _vertex("chat", is_input=True) + agent = _vertex("agent") + human = _vertex("human") + graph = _graph( + [chat, agent, human], + {"human": ["agent"], "agent": ["chat"], "chat": []}, + dropped=set(), + ) + + _rerun_non_input_predecessors(graph, "human") + + assert agent.built is True # round-tripped producer is reused, not re-executed + assert chat.built is True + + +def test_walks_chain_and_reruns_only_dropped_links(): + # chat(input) -> tool(dropped) -> mid(round-tripped) -> agent(paused). The walk must reach the + # dropped tool even though the intermediate mid stays built. + chat = _vertex("chat", is_input=True) + tool = _vertex("tool") + mid = _vertex("mid") + agent = _vertex("agent") + graph = _graph( + [chat, mid, tool, agent], + {"agent": ["mid"], "mid": ["tool"], "tool": ["chat"], "chat": []}, + dropped={"tool"}, + ) + + _rerun_non_input_predecessors(graph, "agent") + + assert tool.built is False # dropped → re-run + assert mid.built is True # round-tripped → reused, but the walk passed through it assert chat.built is True def test_no_predecessors_is_a_noop(): agent = _vertex("agent") - graph = _graph([agent], {"agent": []}) + graph = _graph([agent], {"agent": []}, dropped={"agent"}) _rerun_non_input_predecessors(graph, "agent") diff --git a/src/backend/tests/unit/services/background_execution/test_service.py b/src/backend/tests/unit/services/background_execution/test_service.py index 56a0f2afe9..5932a9fde0 100644 --- a/src/backend/tests/unit/services/background_execution/test_service.py +++ b/src/backend/tests/unit/services/background_execution/test_service.py @@ -27,6 +27,9 @@ def _frame(event_type: str, data: dict) -> tuple[bytes, str]: async def _scripted_source(**_kwargs) -> AsyncIterator[tuple[bytes, str]]: + # Why: these tests isolate the service/runner plumbing (submit, status, dedupe, cross-user + # ownership, terminal handling) from graph execution. The REAL frame source + # (generate_flow_events over a live graph) is exercised end-to-end in test_build_pause_seam. yield _frame("build_start", {}) yield _frame("end_vertex", {"id": "n1"}) yield _frame("end", {}) diff --git a/src/frontend/src/controllers/API/agui/run-flow-bridge.ts b/src/frontend/src/controllers/API/agui/run-flow-bridge.ts index 5465ae0791..e6d4f2ea27 100644 --- a/src/frontend/src/controllers/API/agui/run-flow-bridge.ts +++ b/src/frontend/src/controllers/API/agui/run-flow-bridge.ts @@ -26,6 +26,7 @@ import type { VertexBuildTypeAPI, VertexDataTypeAPI, } from "@/types/api"; +import { getFetchCredentials } from "@/customization/utils/get-fetch-credentials"; import { api } from "../api"; import { injectHumanInputCard } from "./human-input-card"; import { @@ -576,7 +577,7 @@ export async function consumeBackgroundEvents( const response = await fetch(url, { method: "GET", headers, - credentials: "same-origin", + credentials: getFetchCredentials(), signal: opts.signal, }); if (!response.ok || !response.body) { diff --git a/src/lfx/src/lfx/graph/checkpoint/resume.py b/src/lfx/src/lfx/graph/checkpoint/resume.py index e521322a0d..2a7897f5c3 100644 --- a/src/lfx/src/lfx/graph/checkpoint/resume.py +++ b/src/lfx/src/lfx/graph/checkpoint/resume.py @@ -99,6 +99,11 @@ def restore_graph_from_checkpoint(checkpoint: GraphCheckpoint, *, store: Checkpo # Vertices already built at checkpoint time must not have their async generators re-consumed on # resume (the original run exhausted them); the output-collection loop reads this set to skip them. graph.checkpoint_restored_built_ids = {vid for vid, vd in checkpoint.vertex_results.items() if vd.built} + # Built vertices whose live output was opaque-dropped to None re-run on resume; others reuse the + # restored result instead of re-executing (avoids re-billing an Agent whose Message round-tripped). + graph.checkpoint_opaque_dropped_ids = { + vid for vid, vd in checkpoint.vertex_results.items() if vd.built and vd.built_object is None + } _restore_run_manager(graph, checkpoint) _restore_vertices(graph, checkpoint) graph._run_queue.clear() # noqa: SLF001 diff --git a/src/lfx/src/lfx/graph/graph/base.py b/src/lfx/src/lfx/graph/graph/base.py index abc3e1a8a7..11891a7672 100644 --- a/src/lfx/src/lfx/graph/graph/base.py +++ b/src/lfx/src/lfx/graph/graph/base.py @@ -153,6 +153,10 @@ class Graph: # Vertices already built at checkpoint time: on resume their async generators are exhausted, # so the output-collection loop must NOT re-consume them. Empty for fresh (non-resume) runs. self.checkpoint_restored_built_ids: set[str] = set() + # Built vertices whose live output (Tool/model client) was opaque-dropped to None in the + # checkpoint: only these re-run on resume to regenerate the object, so producers whose output + # round-tripped (e.g. an Agent's Message) are not needlessly re-executed. Empty otherwise. + self.checkpoint_opaque_dropped_ids: set[str] = set() self.pause_probe: Callable[[str], Any] | None = None if context and not isinstance(context, dict): diff --git a/src/lfx/tests/unit/graph/checkpoint/test_resume.py b/src/lfx/tests/unit/graph/checkpoint/test_resume.py index a83cff3297..b83ac4c215 100644 --- a/src/lfx/tests/unit/graph/checkpoint/test_resume.py +++ b/src/lfx/tests/unit/graph/checkpoint/test_resume.py @@ -101,6 +101,19 @@ async def test_resume_does_not_reconsume_restored_built_output_vertex(): await resumed.arun([{}], fallback_to_env_vars=False) +async def test_resume_flags_only_opaque_dropped_producers(): + _, checkpoint = await _paused_checkpoint() + # A built producer whose live output (Tool/model) was opaque-dropped is stored with built_object + # None; one whose output round-tripped keeps a value. Only the former must be flagged for re-run. + checkpoint.vertex_results["chat_input"].built = True + checkpoint.vertex_results["chat_input"].built_object = None + + resumed = Graph.resume_from_checkpoint(checkpoint) + + assert "chat_input" in resumed.checkpoint_opaque_dropped_ids + assert "chat_output" not in resumed.checkpoint_opaque_dropped_ids + + async def test_resume_restores_cycle_vertices_from_graph(): _, checkpoint = await _paused_checkpoint() resumed = Graph.resume_from_checkpoint(checkpoint)