fix(runtime): prevent session restart until cancelled execution fully terminates#7001
Conversation
PR Requirements WarningThis PR does not meet the contribution requirements. PR Author: @kowshik-04 To fix:
Exception: To bypass this requirement, you can:
Micro-fix requirements (must meet ALL):
Why is this required? See #472 for details. |
github-actions
Bot
added
the
pr-requirements-warning
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughIntroduces a tri-state cancellation result and "cancelling" lifecycle state; preserves execution bookkeeping until tasks fully terminate; registers executions and performs cancellation transitions under the stream lock; prevents new executions when blocking executions exist; updates HTTP handlers to return 409 on conflicting starts; adds async regression tests for forced-cancel overlaps. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Server
participant ExecStream as ExecutionStream
participant Task as AsyncTask
Client->>Server: POST /trigger (start Execution A)
Server->>ExecStream: execute()
ExecStream->>ExecStream: acquire _lock\n_get_blocking_execution_ids_locked()\n(register task under lock)
ExecStream->>Task: start Task A
Note over Task: Task A running...
Client->>Server: POST /stop (cancel Execution A)
Server->>ExecStream: cancel_execution()
ExecStream->>ExecStream: acquire _lock\nset context.status -> "cancelling"\ncall task.cancel()
ExecStream->>ExecStream: release _lock\nwait grace (5s) but keep bookkeeping
Client->>Server: POST /trigger (retry start)
Server->>ExecStream: execute()
ExecStream->>ExecStream: acquire _lock\n_get_blocking_execution_ids_locked() -> blocking_ids
ExecStream-->>Server: raise ExecutionAlreadyRunningError(blocking_ids)
Server-->>Client: HTTP 409 Conflict (active_execution_ids)
Task->>ExecStream: Task A exits
ExecStream->>ExecStream: final cleanup (remove bookkeeping)
Client->>Server: POST /trigger (after exit)
Server->>ExecStream: execute() -> success (new execution)
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
core/framework/server/tests/test_api.py (1)
571-587: Extend this conflict test to/resumeand/replay.The PR adds the same
ExecutionAlreadyRunningErrormapping in three handlers, but this only locks down/trigger. A small parametrized case over/trigger,/resume, and/replaywould protect the full 409 contract and catch drift between the copied branches.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/framework/server/tests/test_api.py` around lines 571 - 587, Extend the existing test_trigger_returns_409_when_execution_still_running into a parametrized test over the three endpoints "/api/sessions/test_agent/trigger", "/api/sessions/test_agent/resume", and "/api/sessions/test_agent/replay"; set session.graph_runtime.trigger, session.graph_runtime.resume, and session.graph_runtime.replay to AsyncMock(side_effect=ExecutionAlreadyRunningError("default", ["session-1"])) (or set the method corresponding to the current endpoint dynamically), then POST the same JSON payload to each endpoint and assert resp.status == 409 and that the JSON contains stream_id == "default" and active_execution_ids == ["session-1"] so the 409 contract is verified for all three handlers.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/framework/runtime/execution_stream.py`:
- Around line 319-329: stop() currently clears both _active_executions and
_execution_tasks after its 5s timeout which allows orphaned tasks to overlap
with a later start()/execute(); change stop() to mirror the "_keep bookkeeping
until task exit" behavior used by _get_blocking_execution_ids_locked(): when
timing out, do not remove entries from _execution_tasks (and only remove from
_active_executions), ensure you only cancel tasks via their Task objects but
retain their bookkeeping until task.done() is true, and add cleanup logic that
removes entries from _execution_tasks once the associated Task finishes to avoid
leaks; update references in stop(), any task-cancellation helper, and the
task-completion callback to implement this behavior.
- Around line 1224-1246: The current cancellation routine stores a cancel reason
and calls task.cancel() but always returns True even when the task remains
running, causing callers (handle_pause()/handle_stop()) to treat "cancellation
requested" as "execution stopped"; change the function that manipulates
self._execution_tasks, self._cancel_reasons and self._active_executions so it
distinguishes outcomes: after awaiting asyncio.wait({task}, timeout=5.0) return
a distinct value (e.g. "cancelled" when done, "cancelling" when still alive)
instead of a boolean True/False, or return an enum/string and update callers
accordingly so handle_pause() and handle_stop() only mark execution as
stopped/staging when the result indicates the task actually finished; ensure the
code still sets ctx.status = "cancelling" before task.cancel() and preserve the
warning log path when not done.
---
Nitpick comments:
In `@core/framework/server/tests/test_api.py`:
- Around line 571-587: Extend the existing
test_trigger_returns_409_when_execution_still_running into a parametrized test
over the three endpoints "/api/sessions/test_agent/trigger",
"/api/sessions/test_agent/resume", and "/api/sessions/test_agent/replay"; set
session.graph_runtime.trigger, session.graph_runtime.resume, and
session.graph_runtime.replay to
AsyncMock(side_effect=ExecutionAlreadyRunningError("default", ["session-1"]))
(or set the method corresponding to the current endpoint dynamically), then POST
the same JSON payload to each endpoint and assert resp.status == 409 and that
the JSON contains stream_id == "default" and active_execution_ids ==
["session-1"] so the 409 contract is verified for all three handlers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 1f489c3b-6832-45f7-a51f-e1f0970c117f
📒 Files selected for processing (4)
core/framework/runtime/execution_stream.pycore/framework/runtime/tests/test_forced_cancel_dual_execution.pycore/framework/server/routes_execution.pycore/framework/server/tests/test_api.py
9099193 to
4939f90
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/framework/runtime/execution_stream.py (1)
1213-1253:⚠️ Potential issue | 🟠 Major
_cancel_reasonsstill leaks on the normal cancellation path.The reason is only drained in the outer
CancelledErrorhandler. WhenGraphExecutor.execute()turns cancellation into a pausedExecutionResultinstead—which your comment in_run_execution()says is expected—the entry survives and can be reused by a later resume/cancel on the same session ID. Please consume or clear the stored reason on the paused-result path as well, or unconditionally remove it infinallyafter emission.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/framework/runtime/execution_stream.py` around lines 1213 - 1253, The _cancel_reasons entry for execution_id is left in self._cancel_reasons when GraphExecutor.execute() converts the cancellation into a paused ExecutionResult (per the comment in _run_execution()), causing the reason to leak and be reused later; update the cancellation flow to consume/clear that stored reason when a paused result is produced (or, simpler, unconditionally remove self._cancel_reasons[execution_id] in the finally block after emitting the pause/fail event). Focus changes around cancel_execution, _run_execution and the CancelledError handler so the stored reason is removed once handled (or after emission) to avoid reuse on subsequent resume/cancel operations.
♻️ Duplicate comments (1)
core/framework/runtime/execution_stream.py (1)
321-331:⚠️ Potential issue | 🔴 Critical
stop()can still bypass this blocking guard.This helper only works if timed-out tasks stay registered, but
stop()still clears_execution_tasksand_active_executionsafter its 5s wait. That lets a later restart miss an orphaned task and overlap a new execution on the same session anyway.stop()needs to keep timed-out tasks tracked until their normalfinallycleanup runs.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/framework/runtime/execution_stream.py` around lines 321 - 331, _stop() currently clears _execution_tasks and _active_executions after a 5s wait which can drop timed-out/orphaned tasks and allow a restart to overlap a new execution; instead, modify stop() to preserve entries for tasks that are still running so _get_blocking_execution_ids_locked can see them. Concretely: when implementing stop(), do not unconditionally clear _execution_tasks or _active_executions after the timeout; instead cancel outstanding tasks, await their completion (or track them) and only remove entries when the task is done and its normal finally/cleanup has run (or remove only those whose task.done() is True). Keep using the same dictionaries (_execution_tasks, _active_executions) so _get_blocking_execution_ids_locked continues to detect blocking execution IDs until cleanup finishes.
🧹 Nitpick comments (2)
core/framework/server/tests/test_api.py (2)
572-586: Add-> Noneto the new async test signatures.Both new test functions are missing explicit return annotations, so this file is out of compliance with the repo's Python typing rule.
As per coding guidelines, "Use type hints on all function signatures".
Also applies to: 839-853
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/framework/server/tests/test_api.py` around lines 572 - 586, The async test functions (e.g., async def test_trigger_returns_409_when_execution_still_running and the other test at lines 839-853) are missing return type annotations; update their signatures to include an explicit return hint "-> None" (e.g., async def test_trigger_returns_409_when_execution_still_running(...) -> None:) so they comply with the project's type-hinting rule for all function signatures.
571-586: Add the same conflict-path coverage for/resumeand/replay.This PR changes the blocking contract for all start-like routes, but this file only asserts the new
409 Conflictbehavior for/trigger. If/resumeor/replayregresses back to 200/500 onExecutionAlreadyRunningError, this suite will miss it.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/framework/server/tests/test_api.py` around lines 571 - 586, Add two tests mirroring test_trigger_returns_409_when_execution_still_running that exercise the /resume and /replay routes: mock session.graph_runtime.resume and session.graph_runtime.replay to raise ExecutionAlreadyRunningError("default", ["session-1"]), create the app with _make_app_with_session(session), POST to "/api/sessions/test_agent/resume" and "/api/sessions/test_agent/replay" with appropriate JSON payloads (matching expected input keys used by those endpoints), then assert resp.status == 409 and that the returned JSON contains "stream_id" == "default" and "active_execution_ids" == ["session-1"] so the conflict-path coverage matches trigger.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/framework/runtime/execution_stream.py`:
- Around line 588-596: The execute() registration is racy with shutdown because
_running is checked before acquiring _lock and stop() never acquires _lock;
modify stop() to acquire self._lock while flipping self._running and when
snapshotting/cancelling tracked tasks (use the same lock used in execute), and
update execute() to re-check self._running inside the locked section before
adding to self._active_executions, self._completion_events and creating the task
via asyncio.create_task(self._run_execution(ctx)); if self._running is False
after acquiring the lock, raise ExecutionAlreadyRunningError (or an appropriate
shutdown error) instead of registering the execution so no untracked tasks can
start during shutdown.
In `@core/framework/tools/queen_lifecycle_tools.py`:
- Around line 1037-1047: The stop helper now returns a status field that can be
"cancelling" or "stopped"; callers that call stop_graph() must not
unconditionally advance lifecycle phases after await stop_graph() — instead
inspect the returned JSON's "status" (or the cancelling/cancelled booleans) and
only transition to editing/staging when status == "stopped" (or cancelling ==
False and cancelled == True). Locate the callers that do await stop_graph() and
follow it with a phase change, and wrap those transitions in a conditional that
checks stop_graph()'s returned status (e.g., call stop_graph(), parse the
returned dict, and if status == "stopped" then proceed to change phase;
otherwise keep the session blocked).
---
Outside diff comments:
In `@core/framework/runtime/execution_stream.py`:
- Around line 1213-1253: The _cancel_reasons entry for execution_id is left in
self._cancel_reasons when GraphExecutor.execute() converts the cancellation into
a paused ExecutionResult (per the comment in _run_execution()), causing the
reason to leak and be reused later; update the cancellation flow to
consume/clear that stored reason when a paused result is produced (or, simpler,
unconditionally remove self._cancel_reasons[execution_id] in the finally block
after emitting the pause/fail event). Focus changes around cancel_execution,
_run_execution and the CancelledError handler so the stored reason is removed
once handled (or after emission) to avoid reuse on subsequent resume/cancel
operations.
---
Duplicate comments:
In `@core/framework/runtime/execution_stream.py`:
- Around line 321-331: _stop() currently clears _execution_tasks and
_active_executions after a 5s wait which can drop timed-out/orphaned tasks and
allow a restart to overlap a new execution; instead, modify stop() to preserve
entries for tasks that are still running so _get_blocking_execution_ids_locked
can see them. Concretely: when implementing stop(), do not unconditionally clear
_execution_tasks or _active_executions after the timeout; instead cancel
outstanding tasks, await their completion (or track them) and only remove
entries when the task is done and its normal finally/cleanup has run (or remove
only those whose task.done() is True). Keep using the same dictionaries
(_execution_tasks, _active_executions) so _get_blocking_execution_ids_locked
continues to detect blocking execution IDs until cleanup finishes.
---
Nitpick comments:
In `@core/framework/server/tests/test_api.py`:
- Around line 572-586: The async test functions (e.g., async def
test_trigger_returns_409_when_execution_still_running and the other test at
lines 839-853) are missing return type annotations; update their signatures to
include an explicit return hint "-> None" (e.g., async def
test_trigger_returns_409_when_execution_still_running(...) -> None:) so they
comply with the project's type-hinting rule for all function signatures.
- Around line 571-586: Add two tests mirroring
test_trigger_returns_409_when_execution_still_running that exercise the /resume
and /replay routes: mock session.graph_runtime.resume and
session.graph_runtime.replay to raise ExecutionAlreadyRunningError("default",
["session-1"]), create the app with _make_app_with_session(session), POST to
"/api/sessions/test_agent/resume" and "/api/sessions/test_agent/replay" with
appropriate JSON payloads (matching expected input keys used by those
endpoints), then assert resp.status == 409 and that the returned JSON contains
"stream_id" == "default" and "active_execution_ids" == ["session-1"] so the
conflict-path coverage matches trigger.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 08415fbb-db4b-4ec9-a7b7-460a23e64ace
📒 Files selected for processing (6)
core/framework/runtime/agent_runtime.pycore/framework/runtime/execution_stream.pycore/framework/runtime/tests/test_forced_cancel_dual_execution.pycore/framework/server/routes_execution.pycore/framework/server/tests/test_api.pycore/framework/tools/queen_lifecycle_tools.py
✅ Files skipped from review due to trivial changes (1)
- core/framework/runtime/tests/test_forced_cancel_dual_execution.py
🚧 Files skipped from review as they are similar to previous changes (1)
- core/framework/server/routes_execution.py
4939f90 to
ef71924
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
core/framework/server/tests/test_api.py (1)
571-586: Please add the same 409-path coverage forresumeandreplay.This new test exercises
/trigger, but the PR changes the conflict response for/resumeand/replaytoo. A regression in either route would still pass this suite.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/framework/server/tests/test_api.py` around lines 571 - 586, Add two tests mirroring test_trigger_returns_409_when_execution_still_running for the resume and replay routes: create a session via _make_session(), set session.graph_runtime.resume and session.graph_runtime.replay to AsyncMock(side_effect=ExecutionAlreadyRunningError("default", ["session-1"])), build the app with _make_app_with_session(session), POST to "/api/sessions/test_agent/resume" and "/api/sessions/test_agent/replay" with the same JSON structure used by the trigger test, and assert resp.status == 409 and that the returned JSON contains "stream_id" == "default" and "active_execution_ids" == ["session-1"] using TestClient(TestServer(app)) as in the trigger test.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/framework/runtime/execution_stream.py`:
- Around line 423-429: The code currently uses
self._cancel_reasons.setdefault(execution_id, "Execution cancelled") which
preserves any stale cancel reason for reused execution IDs (e.g.,
resume_session_id reusing the same execution id); change this to explicitly
overwrite or clear previous reasons: replace the setdefault() call in the loop
with an assignment self._cancel_reasons[execution_id] = "Execution cancelled"
(so marking cancelling always replaces old reasons), and also ensure
resume_session_id (the resume path that reuses an execution id) clears any prior
entry (e.g., pop/exlicitly delete from _cancel_reasons) when resuming so resumed
executions cannot leak stale messages.
---
Nitpick comments:
In `@core/framework/server/tests/test_api.py`:
- Around line 571-586: Add two tests mirroring
test_trigger_returns_409_when_execution_still_running for the resume and replay
routes: create a session via _make_session(), set session.graph_runtime.resume
and session.graph_runtime.replay to
AsyncMock(side_effect=ExecutionAlreadyRunningError("default", ["session-1"])),
build the app with _make_app_with_session(session), POST to
"/api/sessions/test_agent/resume" and "/api/sessions/test_agent/replay" with the
same JSON structure used by the trigger test, and assert resp.status == 409 and
that the returned JSON contains "stream_id" == "default" and
"active_execution_ids" == ["session-1"] using TestClient(TestServer(app)) as in
the trigger test.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 7b7b9b97-c293-434a-bde2-1e0db80e966d
📒 Files selected for processing (6)
core/framework/runtime/agent_runtime.pycore/framework/runtime/execution_stream.pycore/framework/runtime/tests/test_forced_cancel_dual_execution.pycore/framework/server/routes_execution.pycore/framework/server/tests/test_api.pycore/framework/tools/queen_lifecycle_tools.py
🚧 Files skipped from review as they are similar to previous changes (3)
- core/framework/runtime/agent_runtime.py
- core/framework/server/routes_execution.py
- core/framework/tools/queen_lifecycle_tools.py
|
Closing PR because the contribution requirements were not resolved within the 24-hour grace period. |
github-actions
Bot
removed
the
pr-requirements-warning
Hundao
left a comment
Hundao
left a comment
There was a problem hiding this comment.
Diagnosis is right and the fix is the right shape. Tri-state return + Literal, 409 in the routes, _get_blocking_execution_ids_locked catching the "task still alive but missing from active_executions" gap, create_task moved inside the lock, the stop() rewrite. All the right calls. The tests with a real BlockingExecutor mock are also way better than monkeypatching asyncio.wait against a SimpleNamespace.
One thing I'd want before merging: a comment on cancel_execution (and probably _get_blocking_execution_ids_locked too) explaining why the stream stays locked until the task actually exits. If the task is truly stuck it can still mutate session state, so silently letting a replacement start would guarantee corruption. The permanent lock is doing real work, it's not an oversight.
The reason I want this written down: the competing PR #6970 added a done_callback based "deferred cleanup" that looks like a safety net but actually re-introduces the race, because resume_session_id makes the new execution reuse the same execution_id, and the deferred pop ends up running after the new execution has registered its bookkeeping. Without a comment explaining the intent, the next person looking at this is going to be tempted to add the same thing.
Minor / non-blocking: it might be worth emitting a periodic warning when a stream sits in cancelling for >N seconds, so we can spot a leaked task before someone notices the stream is wedged. Follow-up is fine.
Also when you push the comment update, please run ruff check and ruff format over the changed files so CI passes cleanly.
Other than that, looks ready to land.
ef71924 to
38c1cc2
Compare
|
@Hundao I also ran |
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
core/framework/runtime/execution_stream.py (1)
427-433:⚠️ Potential issue | 🟠 MajorClear stale
_cancel_reasonswhen reusing an execution ID.
resume_session_idreuses the old execution ID, and cooperative cancellation can leave_cancel_reasons[execution_id]behind.stop()still usessetdefault(), so a later shutdown can surface the previous run's reason on the resumed run.Suggested fix
- self._cancel_reasons.setdefault(execution_id, "Execution cancelled") + self._cancel_reasons[execution_id] = "Execution cancelled"async with self._lock: if not self._running: raise RuntimeError(f"ExecutionStream '{self.stream_id}' is not running") blocking_ids = self._get_blocking_execution_ids_locked() if blocking_ids: raise ExecutionAlreadyRunningError(self.stream_id, blocking_ids) + self._cancel_reasons.pop(execution_id, None) self._active_executions[execution_id] = ctx self._completion_events[execution_id] = asyncio.Event() self._execution_tasks[execution_id] = asyncio.create_task(self._run_execution(ctx))Also applies to: 557-559, 596-606
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/framework/runtime/execution_stream.py` around lines 427 - 433, The stale _cancel_reasons entry must be cleared when an execution ID is reused: locate places referencing _cancel_reasons (the loop shown where you setdefault on execution_id, the stop() code paths around the other occurrences, and any resume_session_id logic) and remove or override previous entries instead of using setdefault; specifically, ensure you either delete _cancel_reasons[execution_id] when starting/resuming a run or use direct assignment (e.g., _cancel_reasons[execution_id] = "<new reason>") so a prior run's reason cannot leak into the resumed execution. Ensure this change is applied to the loop that marks ctx.status = "cancelling" and the other similar sites (the occurrences around lines 557-559 and 596-606 referenced in the review).
🧹 Nitpick comments (3)
core/framework/server/tests/test_api.py (1)
571-586: Add the same 409 regression coverage for/resumeand/replay.This PR changes all three start-like routes to reject while an older execution is still shutting down, but this file only exercises the
/triggerbranch. Leaving/resumeand/replayuntested makes those new conflict paths easy to regress later.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/framework/server/tests/test_api.py` around lines 571 - 586, Add two tests mirroring test_trigger_returns_409_when_execution_still_running: create session via _make_session(), set session.graph_runtime.resume and session.graph_runtime.replay to AsyncMock that raise ExecutionAlreadyRunningError("default", ["session-1"]), build app with _make_app_with_session(session), POST to "/api/sessions/test_agent/resume" and "/api/sessions/test_agent/replay" with the same JSON payload used in the trigger test, and assert resp.status == 409 and that the returned JSON contains stream_id == "default" and active_execution_ids == ["session-1"]; use the same pytest.mark.asyncio, AsyncMock, TestClient(TestServer(app)) pattern as test_trigger_returns_409_when_execution_still_running.core/framework/runtime/tests/test_forced_cancel_dual_execution.py (2)
103-103: Reuse the explicit session ID in the restart attempts.Both tests start with
"session-1"but retry withexecution_id. That only exercises the same-session conflict if those identifiers stay equivalent. Reusing asession_idlocal here would make the regression target the bug more directly and decouple it fromexecute()return semantics.Also applies to: 119-128, 160-172
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/framework/runtime/tests/test_forced_cancel_dual_execution.py` at line 103, The test currently starts execution with session_state={"resume_session_id": "session-1"} but uses execution_id for subsequent restarts; change the restart attempts to reuse the explicit session id (e.g., "session-1") instead of execution_id so the conflict/restart logic is exercised against the same resume_session_id; update all restart calls in test_forced_cancel_dual_execution.py (the places around the stream.execute call that set execution_id and the retry blocks at lines ~119-128 and ~160-172) to pass the original session id value into the restart/re-execute calls rather than using execution_id.
22-73: Add the missing annotations on the test helpers and fixtures.
_build_stream,_install_blocking_executor, both test functions, and theimmediate_timeoutstubs still have untyped parameters or returns, so this new file is not aligned with the repo's Python rules.As per coding guidelines, "Use type hints on all function signatures."
Also applies to: 76-101, 135-156
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/framework/runtime/tests/test_forced_cancel_dual_execution.py` around lines 22 - 73, Add missing type hints to all helper and fixture functions referenced in this test file: annotate _build_stream to return ExecutionStream and accept tmp_path: Path and event_bus: EventBus | None, annotate _install_blocking_executor to accept monkeypatch: pytest.MonkeyPatch and release: asyncio.Event and return None, annotate immediate_timeout stub with appropriate return type (e.g., contextmanager or AsyncContextManager) and any parameters, and add type hints to both test functions' parameters and return types (e.g., -> None). Update imports as needed for Path, ExecutionStream, EventBus, pytest.MonkeyPatch, and typing constructs so the signatures are fully typed and satisfy the repository rule "Use type hints on all function signatures."
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/framework/tools/queen_lifecycle_tools.py`:
- Around line 1163-1164: After calling await
phase_state.switch_to_planning(source="tool") in the stop_graph_and_plan path,
persist the updated phase metadata to meta.json the same way the sibling
stop-and-transition helpers do: invoke the existing metadata persistence routine
used elsewhere (the helper that updates meta.json for other transitions) so the
successful planning transition is written to disk and restored sessions open in
the correct phase; place the persistence call immediately after
phase_state.switch_to_planning(source="tool") so it runs only on successful
transition.
---
Duplicate comments:
In `@core/framework/runtime/execution_stream.py`:
- Around line 427-433: The stale _cancel_reasons entry must be cleared when an
execution ID is reused: locate places referencing _cancel_reasons (the loop
shown where you setdefault on execution_id, the stop() code paths around the
other occurrences, and any resume_session_id logic) and remove or override
previous entries instead of using setdefault; specifically, ensure you either
delete _cancel_reasons[execution_id] when starting/resuming a run or use direct
assignment (e.g., _cancel_reasons[execution_id] = "<new reason>") so a prior
run's reason cannot leak into the resumed execution. Ensure this change is
applied to the loop that marks ctx.status = "cancelling" and the other similar
sites (the occurrences around lines 557-559 and 596-606 referenced in the
review).
---
Nitpick comments:
In `@core/framework/runtime/tests/test_forced_cancel_dual_execution.py`:
- Line 103: The test currently starts execution with
session_state={"resume_session_id": "session-1"} but uses execution_id for
subsequent restarts; change the restart attempts to reuse the explicit session
id (e.g., "session-1") instead of execution_id so the conflict/restart logic is
exercised against the same resume_session_id; update all restart calls in
test_forced_cancel_dual_execution.py (the places around the stream.execute call
that set execution_id and the retry blocks at lines ~119-128 and ~160-172) to
pass the original session id value into the restart/re-execute calls rather than
using execution_id.
- Around line 22-73: Add missing type hints to all helper and fixture functions
referenced in this test file: annotate _build_stream to return ExecutionStream
and accept tmp_path: Path and event_bus: EventBus | None, annotate
_install_blocking_executor to accept monkeypatch: pytest.MonkeyPatch and
release: asyncio.Event and return None, annotate immediate_timeout stub with
appropriate return type (e.g., contextmanager or AsyncContextManager) and any
parameters, and add type hints to both test functions' parameters and return
types (e.g., -> None). Update imports as needed for Path, ExecutionStream,
EventBus, pytest.MonkeyPatch, and typing constructs so the signatures are fully
typed and satisfy the repository rule "Use type hints on all function
signatures."
In `@core/framework/server/tests/test_api.py`:
- Around line 571-586: Add two tests mirroring
test_trigger_returns_409_when_execution_still_running: create session via
_make_session(), set session.graph_runtime.resume and
session.graph_runtime.replay to AsyncMock that raise
ExecutionAlreadyRunningError("default", ["session-1"]), build app with
_make_app_with_session(session), POST to "/api/sessions/test_agent/resume" and
"/api/sessions/test_agent/replay" with the same JSON payload used in the trigger
test, and assert resp.status == 409 and that the returned JSON contains
stream_id == "default" and active_execution_ids == ["session-1"]; use the same
pytest.mark.asyncio, AsyncMock, TestClient(TestServer(app)) pattern as
test_trigger_returns_409_when_execution_still_running.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: c1a0a038-9c0a-4c9c-b266-7d16b8ff2f7c
📒 Files selected for processing (6)
core/framework/runtime/agent_runtime.pycore/framework/runtime/execution_stream.pycore/framework/runtime/tests/test_forced_cancel_dual_execution.pycore/framework/server/routes_execution.pycore/framework/server/tests/test_api.pycore/framework/tools/queen_lifecycle_tools.py
🚧 Files skipped from review as they are similar to previous changes (2)
- core/framework/runtime/agent_runtime.py
- core/framework/server/routes_execution.py
38c1cc2 to
c099108
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
core/framework/runtime/agent_runtime.py (1)
1600-1615: LGTM! Return type aligns with tri-state cancellation semantics.The change from
booltostrcorrectly propagates theCancelExecutionResultoutcomes ("cancelled","cancelling","not_found") from the underlyingExecutionStream.cancel_execution()method. The docstring update at line 1610 accurately describes the new semantics.Consider using the
CancelExecutionResulttype alias for stronger type safety:from framework.runtime.execution_stream import CancelExecutionResult async def cancel_execution(...) -> CancelExecutionResult:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/framework/runtime/agent_runtime.py` around lines 1600 - 1615, Update the cancel_execution signature to use the CancelExecutionResult type alias for stronger typing: import CancelExecutionResult from framework.runtime.execution_stream and change the return annotation of AgentRuntime.cancel_execution to -> CancelExecutionResult; ensure any callers or related docstring remain consistent with the tri-state outcomes returned by ExecutionStream.cancel_execution.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@core/framework/runtime/agent_runtime.py`:
- Around line 1600-1615: Update the cancel_execution signature to use the
CancelExecutionResult type alias for stronger typing: import
CancelExecutionResult from framework.runtime.execution_stream and change the
return annotation of AgentRuntime.cancel_execution to -> CancelExecutionResult;
ensure any callers or related docstring remain consistent with the tri-state
outcomes returned by ExecutionStream.cancel_execution.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 86aac994-7fbc-4eb9-88ee-b3cb00b8aa74
📒 Files selected for processing (6)
core/framework/runtime/agent_runtime.pycore/framework/runtime/execution_stream.pycore/framework/runtime/tests/test_forced_cancel_dual_execution.pycore/framework/server/routes_execution.pycore/framework/server/tests/test_api.pycore/framework/tools/queen_lifecycle_tools.py
🚧 Files skipped from review as they are similar to previous changes (3)
- core/framework/server/tests/test_api.py
- core/framework/tools/queen_lifecycle_tools.py
- core/framework/server/routes_execution.py
|
@Hundao
Could you please take another look when you have a moment? |
c099108 to
42df196
Compare
|
Rebased onto the latest upstream/main and resolved the merge conflicts. Preserved the execution lifecycle fixes, aligned the changes with recent upstream updates, and re-ran the focused forced-cancel regression tests locally. Please let me know if anything else needs adjustment. |
42df196 to
65b02d0
Compare
|
Rebased onto the latest upstream/main and resolved the remaining conflict in Preserved the tri-state cancellation semantics and aligned with the updated upstream structure. |
- keep bookkeeping until task termination - block restart while any execution task is still alive - make execution registration atomic under lock - avoid premature cleanup on cancel timeout - add regression tests for forced-cancel restart scenarios
fbcdd86 to
00cf878
Compare
2 participants
Summary
Fixes a race condition in ExecutionStream where a force-cancelled execution could time out during cancellation, have its bookkeeping removed early, and allow a new execution to start on the same session while the previous task was still running.
This change ensures the session remains blocked until the original task has fully terminated.
Type of Change
Bug fix (non-breaking change that fixes an issue)
Related Issues
Closes #6967
Changes Made
Keep execution bookkeeping on cancel timeout instead of removing it early
Mark executions as cancelling while shutdown is in progress
Treat any still-alive execution task as blocking for new executions
Make execution registration fully atomic under the stream lock
Return 409 Conflict from trigger/resume/replay routes when execution is still shutting down
Add regression tests covering forced-cancel restart scenarios
Testing
Focused regression tests pass:
.venv/bin/pytest core/framework/runtime/tests/test_forced_cancel_dual_execution.py -q
Modified files also compile cleanly.
The following scenarios are covered:
forced cancel timeout keeps the stream blocked until task exit
repeated forced restarts do not allow overlapping executions
Note: Full aiohttp server tests could not be executed in this environment due to socket binding restrictions, but the runtime behavior relevant to this fix was validated by focused regression tests.
Checklist
My code follows the project's style guidelines
I have performed a self-review of my code
I have added tests that prove my fix is effective
Relevant focused tests pass locally
Summary by CodeRabbit
New Features
Bug Fixes
Tests