fix(api): close the migration checks' source when the client goes away

[ogabrielluiz]

Stacked on the copy destination PR.

POST /api/v1/migration/checks streams the source checks from a child process. A client that goes away is meant to end the run: the child is stopped and the record says cancelled. That happened only when the disconnect found the run waiting for its child. When it found an event waiting to be sent, the response ended and left the run going.

What that did. The record kept saying running and the child ran on to the end of its checks. For as long as the child lived, the next POST /checks answered 409 already_running. Later the garbage collector closed the abandoned stream, and the stream then wrote its own run into the record as cancelled, over whatever run the record held by that time. A check that had passed in between read as cancelled again, with no report.

The fix. The backup download already returns a small response class that closes its source when the response ends, however it ends. This PR names that class for both uses (_ClosingStream) and returns it from POST /checks as well, so the run ends inside its own request in both cases. Nothing else changes.

Why it was rare. Starlette cancels a streamed response when its client disconnects. A cancellation that arrives while the stream waits inside its source is raised there, and the source cleans up. One that arrives while the stream waits to send never reaches the source. A check run spends nearly all of its time waiting for the child, so most disconnects land well. Each middleware layer of the app takes an event over before it passes it on, and for those moments the route is waiting to send. test_a_second_run_waits_for_the_live_one hangs up just as the first event arrives, and it failed once in 33 runs of its file with the 409 above.

The copy-run events stream. GET /steps/{step_id}/runs/{run_id}/events is a plain streamed response too, and I left it as it is. A disconnect leaves its source open in the same way, but that source has nothing to clean up: it opens the run's log for each read and holds nothing in between, and a page that goes away is meant to stop nothing. Closing it inside the request would also take more than this class, because the route wraps the follower in a second generator and closing that one leaves the follower open.

How I tested it

One new test in test_migration.py. It takes the response that the route returns, with the real migration-preflight child behind it, and plays the client by hand: the first event is never taken and the client disconnects. When the response returns, the record on disk has to say cancelled with an end time and the child has to be gone, with nothing awaited in between. Before the fix it failed 5 runs of 5 with the record saying running.

The test drives the response and leaves the app's middleware out. Through the whole app the same disconnect lands in either place as timing has it, and I found no way to hold it in the bad one without a sleep. That is also why counting runs of the old test cannot show the fix: on the old code it passed 30 runs of 30 alone.

Beyond the test:

• Before the fix, a disconnect while the run still waits for its child ended the run, 3 runs of 3. That is the half that already worked.
• Before the fix, a disconnect with an event waiting to be sent left the record running and the child alive, and the next POST /checks answered 409 already_running. After the child had ended by itself a second run went through and read done. Then I dropped the first response and ran the collector, and the record read cancelled with no report. With the fix the record says cancelled at the disconnect, the next run starts at once, and it still reads done after the collector has run.
• On the events stream, a disconnect left the follower open, with no file of the run held open and the run still going. Closing the response's source by hand left the follower open until the event loop closed it a few turns later.

test_migration.py ran 5 times in a row each way, with the drivers, a live PostgreSQL and an S3 server, and with no driver importable, as in the CI unit job. On this branch the five migration test files give 240 passed and 1 skipped against a local PostgreSQL and S3 server. The skip is the test that needs an S3 server that checks keys. ruff check and ruff format --check are clean on the two files.

Summary by CodeRabbit

• Bug Fixes
  • Migration checks now stop cleanly if the connection is interrupted.
  • An interrupted older check no longer overwrites the status of a newer check.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

• Configuration used: Repository: langflow-ai/langflow/.coderabbit.yaml
• Review profile: CHILL
• Plan: Advanced
• Run ID: 11c6b82f-a9f8-4780-aff6-8de5615decc5

---

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: langflow-ai/langflow/.coderabbit.yaml
• Review profile: CHILL
• Plan: Advanced
• Run ID: 02bb15e4-b536-40d8-b541-a4138371e1ff

---

📥 Commits

Reviewing files that changed from the base of the PR and between 0a0888a and d77998e.

---

📒 Files selected for processing (2)

• src/backend/base/langflow/api/v1/migration.py
• src/backend/tests/unit/api/v1/test_migration.py

---

Included review availability: This review used your included allowance. Your plan provides up to 10 included reviews per hour; 6 remain after this review.

---

---

Walkthrough

Migration source-check and database-backup responses now use _ClosingStream. Source-check state updates only save when the record still matches the run’s start time. Tests cover client disconnects, child-process termination, and newer-run record preservation.

Changes

Migration stream handling

Layer / File(s)	Summary
Shared closing stream for migration responses src/backend/base/langflow/api/v1/migration.py	The source-check and database-backup endpoints return _ClosingStream. This replaces the _Download response class.
Disconnect cancellation and state protection src/backend/base/langflow/api/v1/migration.py, src/backend/tests/unit/api/v1/test_migration.py	The check-state callback skips a save if the record has a different started_at. Tests verify cancellation state, completion time, child-process termination, and preservation of a newer run’s record.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Merge Risk: ⚪ Minimal · up to d7799

No actionable merge-blocking issue was identified in the changed migration streams.

🚥 Pre-merge checks | ✅ 8 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 14 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (8 passed)

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main change: closing the migration checks' source stream when the client disconnects.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Test Coverage For New Implementations		The PR includes the matching backend test file, src/backend/tests/unit/api/v1/test_migration.py. It adds regression coverage for disconnects while an event waits to be sent, including cancelled, `…
Test Quality And Coverage		Tests cover the changed behavior. The existing migration tests verify successful source-check streaming and validation errors, and backup tests verify successful and rejected responses. The new async …
Test File Naming And Structure		The pull request keeps the backend test at src/backend/tests/unit/api/v1/test_migration.py, which matches the required test_*.py pattern and pytest unit-test location. The added tests use `async d…
Excessive Mock Usage Warning		The changed tests do not add mock objects or mock libraries. They call the real run_checks endpoint, start the real child process, inspect the real migration record, and verify the child process wit…

---

---

✨ Finishing Touches 💡 2

📝 Generate docstrings 💡

• Commit to this branch
• Create a new PR

---

🛠️ Fix failing CI checks 💡

• Commit to this branch
• Create a new PR

---

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

---

---

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

---

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

✅ Test Coverage Advisor

No source changes detected without accompanying tests. Thanks for keeping coverage up! 🎉

Advisory check only — never blocks merge.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 68.96%. Comparing base (567cf2e) to head (ef34471).

Additional details and impacted files

@@                         Coverage Diff                         @@
##           feat/migration-copy-destination   #15618      +/-   ##
===================================================================
+ Coverage                            68.80%   68.96%   +0.16%     
===================================================================
  Files                                 2770     2775       +5     
  Lines                               298870   299572     +702     
  Branches                             41158    40301     -857     
===================================================================
+ Hits                                205637   206604     +967     
+ Misses                               90739    90474     -265     
  Partials                              2494     2494              

Flag	Coverage Δ
backend	78.67% <100.00%> (-0.12%)	⬇️
frontend	65.45% <ø> (+0.34%)	⬆️
lfx	67.49% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/backend/base/langflow/api/v1/migration.py	97.18% <100.00%> (+0.01%)	⬆️

... and 250 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[github-actions]

Frontend Unit Test Coverage Report

Coverage Summary

Lines	Statements	Branches	Functions
	59.49% (97538/163942)	73.98% (15125/20443)	54.61% (2413/4418)

Unit Test Results

Tests	Skipped	Failures	Errors	Time
7402	0 💤	0 ❌	0 🔥	21m 53s ⏱️

[Cristhianzl]

⚠️ Important (preferably this PR)

I1 — The late record write still clobbers a newer run, so the Jira's second symptom survives

_stream_checks' finally assigns the step unconditionally — record["steps"]["check_source"] = step at src/backend/base/langflow/api/v1/migration.py:762 — with no check that the record still belongs to this run. Before this PR the window was wide open: the abandoned stream was closed by the garbage collector at an arbitrary later time, which is the symptom LE-2965 describes as "when the abandoned stream was collected later it wrote cancelled over whatever run the record held by then". This PR moves that close inside the request, which removes the garbage-collector window, but it does not remove the unguarded write, and one path still reaches it late.

The surviving path needs a client that stops reading without closing the socket. The generator parks at yield _event(event) for the first event while await send(...) blocks on backpressure; migration-preflight writes the rest of its (small) output into the 64 KB stdout pipe and exits; _is_live at migration.py:217 then reads status == "running" with a dead pid and returns False, so the next POST /checks is admitted and writes its own check_source. When the stalled client finally goes away, this request's cleanup stamps the old run's status: "cancelled", old pid and report: None over the live one — the page shows the new run as cancelled with no report, which is the exact failure the description says the fix prevents.

The repo already has the idiom for this one function away: _settle_copies guards its write with if saved["steps"][step_id]["run_id"] == step["run_id"] at migration.py:652-654. check_source has no run_id, but started_at (or pid) identifies the run just as well, so the same compare-and-set closes it.

Code reference — src/backend/base/langflow/api/v1/migration.py:758-763

        if step["status"] == "running":
            step["status"] = "cancelled"
        step["finished_at"] = _now()
        record = _read_record()
        record["steps"]["check_source"] = step
        _write_record(record)

Note on scope: the line itself comes from the base PR and this diff does not touch it. I am raising it anyway because the PR description and the Jira both present the overwrite as fixed by this change, and after reading the code it is narrowed rather than fixed. If you would rather keep this PR to the three lines it has, saying so in the description — and carrying the guard into 15621 — is a fine disposition; what I would not leave standing is the claim that the overwrite is gone.

[erichare]

@ogabrielluiz Approving, with one small fix pushed for @Cristhianzl's I1. Thanks for fixing this by reusing the backup download's close-on-end response, now _ClosingStream, instead of adding a second mechanism, and for a test that holds the disconnect in the "event waiting in send" window without a sleep.

What I verified at 2ece9bcf24:

• POST /checks now returns _ClosingStream. In the finally of _stream_checks, the drain cancel, the kill, the cancelled stamp and the synchronous _save all run before the first await. So a close from a cancelled scope cannot lose the record or leave the child running. No stale _Download references remain, and the backup download behaves the same.
• The new test fails on the old code: with run_checks reverted to a plain StreamingResponse, it fails with assert 'running' == 'cancelled'.

Fixes I pushed

• 85d19fc fix(api): keep a late check cancel off a newer run's record. The late write still reached a newer run (I1), and the description says that overwrite is gone. A client that stops reading without closing the socket keeps its request open after its child exits, so _is_live lets the next POST /checks in. When the stalled client finally went away, its cleanup stamped the old run's cancelled over the new run:

  RUN2 record after it finished: done 25078 ...15:10:31 report=True
  FINAL record: cancelled 23128 ...15:09:59 report=None
  

  keep() now saves only while the record's check_source still has this run's started_at, using the same compare-and-set that _settle_copies uses with run_id. test_a_run_that_ends_after_a_newer_run_took_the_record_leaves_it_alone drives the response by hand, like your disconnect test. While the first event waits to be sent, another run takes the record and finishes, and then the client goes away. The test asserts that the newer step is untouched and the first child is gone. Without the guard, the record ends cancelled instead of the newer done.

The PRs stacked above (#15621, #15622, #15623, #15624) will need a restack onto the new head.

Verification

• test_migration.py plus test_migration_copy_runs.py: 165 passed, 23 skipped (the skips need PostgreSQL or S3 servers).
• Mutation check: run against the code before the guard, the new test fails as described above.
• ruff check and ruff format --check are clean on both changed files.