This commit is contained in:
Quanzheng Long
2026-03-19 16:28:15 -07:00
parent 04ed16fa20
commit fa5030dfee
4 changed files with 126 additions and 1 deletions
@@ -222,6 +222,13 @@ class Context:
return resumed
raise WaitRequested(_target_to_suspend_payload(target))
def is_resume(self) -> bool:
return self._run._is_resume_execution()
# Go-style alias.
def IsResume(self) -> bool:
return self.is_resume()
def publish_to_channel(self, channel: str, value: Any) -> None:
self._run.publish_nowait(channel, value)
@@ -418,6 +425,7 @@ class _GraphEngineRun:
) -> dict[str, Any]:
node_input, resume_event = _unwrap_resume_input(node_input)
self._set_resume_event(resume_event)
self._set_is_resume(resume_event is not None)
if node_name not in self._nodes:
raise ValueError(f"Unknown node `{node_name}`")
node = self._nodes[node_name]
@@ -431,6 +439,7 @@ class _GraphEngineRun:
return {"suspend": suspend.payload}
finally:
self._set_resume_event(None)
self._set_is_resume(False)
if isinstance(result, Command):
update = result.update
@@ -450,6 +459,12 @@ class _GraphEngineRun:
def _set_resume_event(self, event: dict[str, Any] | None) -> None:
self._local.resume_event = event
def _set_is_resume(self, is_resume: bool) -> None:
self._local.is_resume = is_resume
def _is_resume_execution(self) -> bool:
return bool(getattr(self._local, "is_resume", False))
def _consume_resume_event(
self, target: WaitCondition | AnyOfCondition | AllOfCondition
) -> WaitForResult | None:
@@ -216,3 +216,35 @@ async def test_all_of_channel_and_timer_marks_both_conditions() -> None:
assert result["logs"] == ["all_of_channel_timer"]
assert result["done"] == "ok"
async def test_is_resume_avoids_duplicate_side_effects() -> None:
graph = AdvancedStateGraph(PrimitiveState)
db_writes: list[str] = []
async def start_node(state: PrimitiveState) -> Command:
return Command(
update={"counter": 0, "logs": [], "done": None},
goto=Send("wait_node", None),
)
async def wait_node(ctx: Context, _input: None, state: PrimitiveState) -> Command:
if not ctx.IsResume():
# Simulate one-time side effect (e.g. database write).
db_writes.append("write")
await ctx.wait_for(timer_condition(seconds=0.02))
state["logs"].append(f"resume={ctx.IsResume()}")
return Command(update=state, goto=Send("finish_node", None))
async def finish_node(_input: None, state: PrimitiveState) -> dict[str, object]:
return {"counter": state["counter"], "logs": state["logs"], "done": "ok"}
graph.add_entry_node(start_node)
graph.add_node(wait_node)
graph.add_finish_node(finish_node)
result = await graph.compile().ainvoke({"counter": 0, "logs": [], "done": None})
assert db_writes == ["write"]
assert result["counter"] == 0
assert result["logs"] == ["resume=True"]
assert result["done"] == "ok"