From fd0c652f33fe97d6d636be64233474651222fc97 Mon Sep 17 00:00:00 2001 From: XuQuanxin04 Date: Wed, 26 Aug 2026 17:03:01 +0800 Subject: [PATCH] fix(flow): reset cancellation flag at the start of each execution cancel() set self._cancelled = True, but that flag was only ever read in _execute() and never reset. Because run(), resume() and rerun_step() all share the same _execute() path on a LightFlow instance, a single cancel() call permanently poisoned the instance: every subsequent run skipped all of its steps with "cancelled before execution", even when cancellation was requested after a run had already finished. Reset the flag when a new execution begins. In-flight cancellation is unaffected because the flag is still re-checked before every step. Co-Authored-By: Claude --- LightAgent/flow.py | 6 +++++ tests/test_lightflow.py | 50 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/LightAgent/flow.py b/LightAgent/flow.py index cfefe6b..8346f56 100644 --- a/LightAgent/flow.py +++ b/LightAgent/flow.py @@ -401,6 +401,12 @@ def _execute( parent_trace_id: str | None = None, run_group_id: str | None = None, ) -> LightFlowResult | str | dict[str, Any]: + # A new execution must not inherit the cancellation state of a previous + # run on the same LightFlow instance. cancel() is still honored for the + # currently executing run because the flag is re-checked before every + # step below; resetting it here only prevents it from permanently + # poisoning subsequent run()/resume()/rerun_step() calls. + self._cancelled = False trace_id = uuid4().hex run_group = run_group_id or run_id recorder = TraceRecorder(enabled=trace, trace_id=trace_id, parent_trace_id=parent_trace_id, run_group_id=run_group) diff --git a/tests/test_lightflow.py b/tests/test_lightflow.py index 05e1f84..0210fa3 100644 --- a/tests/test_lightflow.py +++ b/tests/test_lightflow.py @@ -268,3 +268,53 @@ def rewrite_step(ctx): assert result.success is True assert agent.calls[0]["query"] == "rewritten by flow hook" assert any(event["type"] == "hook_decision" for event in result.trace) + + +class CancellingAgent(FakeAgent): + """An agent that calls flow.cancel() while the flow is executing its step.""" + + def __init__(self, name, flow_holder): + super().__init__(name, ["done"]) + self._flow_holder = flow_holder + + def run(self, query, **kwargs): + self.calls.append({"query": query, "kwargs": kwargs}) + self._flow_holder["flow"].cancel() + return RunResult(content="cancelled mid-run", trace=[]) + + +def test_lightflow_cancel_during_run_skips_remaining_steps_of_that_run(): + flow_holder = {} + first = CancellingAgent("first", flow_holder) + second = FakeAgent("second", ["done"]) + flow = LightFlow().step("first", agent=first).step("second", agent=second, depends_on=["first"]) + flow_holder["flow"] = flow + + result = flow.run("go") + + statuses = {step.name: step.status for step in result.steps} + assert statuses == {"first": "success", "second": "skipped"} + assert second.calls == [] + + +def test_lightflow_cancel_between_runs_does_not_poison_the_next_run(): + """Regression: cancel() set a sticky _cancelled flag that was never reset, + so a subsequent run()/resume()/rerun_step() on the same instance skipped + every step. The flag must be cleared at the start of each execution while + still being honored for the run in progress.""" + first = FakeAgent("first", ["one", "two"]) + second = FakeAgent("second", ["one", "two"]) + flow = LightFlow().step("first", agent=first).step("second", agent=second, depends_on=["first"]) + + first_run = flow.run("first") + assert first_run.success is True + assert len(first.calls) == 1 and len(second.calls) == 1 + + flow.cancel() # e.g. user cancels after the run has already finished + + second_run = flow.run("second") + assert second_run.success is True, [ + (step.name, step.status, step.error) for step in second_run.steps + ] + assert [step.status for step in second_run.steps] == ["success", "success"] + assert len(first.calls) == 2 and len(second.calls) == 2