fix: Enhance interrupt handling in StreamHandler to prevent re-interruption after resume #49
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -99,6 +99,11 @@ async def _check_interrupted( | |
| logger.info( | ||
| "Resuming from interrupted state at node '%s'", state.execution_meta.current_node | ||
| ) | ||
| # Save the interrupted node info before clearing so we don't re-interrupt | ||
| config["_skip_interrupt_at"] = { | ||
| "node": state.execution_meta.interrupted_node, | ||
| "status": state.execution_meta.status, | ||
| } | ||
| # This is a resume case - clear interrupt and merge input data | ||
| if input_data: | ||
| config["resume_data"] = input_data | ||
|
|
@@ -143,6 +148,24 @@ async def _check_and_handle_interrupt( | |
| self.interrupt_before if interrupt_type == "before" else self.interrupt_after | ||
| ) or [] | ||
|
|
||
| # Check if we just resumed from an interrupt at this node with this type | ||
| skip_info = config.get("_skip_interrupt_at", {}) | ||
|
|
||
| if skip_info.get("node") == current_node: | ||
| expected_status = ( | ||
| ExecutionStatus.INTERRUPTED_BEFORE | ||
| if interrupt_type == "before" | ||
| else ExecutionStatus.INTERRUPTED_AFTER | ||
| ) | ||
| if skip_info.get("status") == expected_status: | ||
| logger.debug( | ||
| "Skipping %s interrupt check for node '%s' - just resumed from it", | ||
| interrupt_type, | ||
| current_node, | ||
| ) | ||
| # Clear the flag after using it once | ||
| config.pop("_skip_interrupt_at", None) | ||
| return False | ||
|
Comment on lines
+151
to
+167
There was a problem hiding this comment. This change alters streaming interrupt/resume behavior, but there doesn’t appear to be an integration test that resumes a streaming run (via |
||
|
|
||
| if current_node in interrupt_nodes: | ||
| status = ( | ||
| ExecutionStatus.INTERRUPTED_BEFORE | ||
|
|
||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
On resume,
_skip_interrupt_atis set toexecution_meta.interrupted_nodeunconditionally. ForINTERRUPTED_AFTERthe state often advancescurrent_nodeto the next node before pausing (see_execute_graph), so this flag may never be consumed/cleared and can later suppress a legitimate interrupt if the graph cycles back to the old node. Consider only setting this skip flag wheninterrupted_node == current_node(i.e., when resuming will immediately re-check the same node), or otherwise clearing/avoiding the flag for cases wherecurrent_nodehas already moved on.