diff --git a/langgraph-go/advancedgraph/graph.go b/langgraph-go/advancedgraph/graph.go index 19d92820a..5744a1825 100644 --- a/langgraph-go/advancedgraph/graph.go +++ b/langgraph-go/advancedgraph/graph.go @@ -103,6 +103,7 @@ type CompiledGraph[StateT any] struct { type Context struct { engine *RustEngine resumeEvent *WaitEvent + isResume bool } func (c *Context) WaitFor(target WaitTarget) (WaitForResult, error) { @@ -117,6 +118,10 @@ func (c *Context) WaitFor(target WaitTarget) (WaitForResult, error) { return WaitForResult{}, ErrWaitRequested{Target: target} } +func (c *Context) IsResume() bool { + return c.isResume +} + func (c *Context) PublishToChannel(channel string, value any) error { return c.engine.Publish(channel, value) } @@ -215,7 +220,15 @@ func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT, str return Command{}, fmt.Errorf("node `%s` expected map state argument", node) } resolvedInput, resumeEvent := unwrapResumeInput(nodeInput) - return fn(&Context{engine: engine, resumeEvent: resumeEvent}, resolvedInput, fallbackState) + return fn( + &Context{ + engine: engine, + resumeEvent: resumeEvent, + isResume: resumeEvent != nil, + }, + resolvedInput, + fallbackState, + ) }, ) if err != nil { diff --git a/langgraph-go/tests/test_primitives_test.go b/langgraph-go/tests/test_primitives_test.go index f872f4bd0..b827c3a43 100644 --- a/langgraph-go/tests/test_primitives_test.go +++ b/langgraph-go/tests/test_primitives_test.go @@ -8,6 +8,7 @@ import ( ) type primitiveWorkflow struct { + dbWriteCount int } type primitiveState struct { @@ -427,3 +428,67 @@ func TestAllOfChannelAndTimerMarksBothConditions(t *testing.T) { t.Fatalf("unexpected done: %v", result.Done) } } + +func (w *primitiveWorkflow) startResumeFlagNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) { + state.Logs = []string{} + state.Count = 0 + return ag.Command{ + Update: state, + Goto: []ag.Send{ + {Node: w.resumeFlagNode, NodeInput: nil}, + }, + }, nil +} + +func (w *primitiveWorkflow) resumeFlagNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) { + if !ctx.IsResume() { + // Simulate one-time side effect (e.g. DB write). + w.dbWriteCount += 1 + } + if _, err := ctx.WaitFor(ag.AnyOf(ag.TimerCondition{Seconds: 0.02})); err != nil { + return ag.Command{}, err + } + state.Logs = append(state.Logs, fmt.Sprintf("resume=%v", ctx.IsResume())) + return ag.Command{ + Update: state, + Goto: []ag.Send{ + {Node: w.finishResumeFlagNode, NodeInput: nil}, + }, + }, nil +} + +func (w *primitiveWorkflow) finishResumeFlagNode(_ *ag.Context, _ any, state primitiveState) (ag.Command, error) { + state.Done = "ok" + return ag.Command{Update: state}, nil +} + +func TestIsResumeAvoidsDuplicateSideEffects(t *testing.T) { + workflow := &primitiveWorkflow{} + graph := ag.NewAdvancedStateGraph[primitiveState]() + graph.AddEntryNode(workflow.startResumeFlagNode) + graph.AddNode(workflow.resumeFlagNode) + graph.AddFinishNode(workflow.finishResumeFlagNode) + + handler, err := graph.Compile().Start(nil, primitiveState{ + Count: 0, + Logs: []string{}, + Done: "", + }) + if err != nil { + t.Fatalf("start failed: %v", err) + } + + result, err := handler.WaitForResult() + if err != nil { + t.Fatalf("result failed: %v", err) + } + if workflow.dbWriteCount != 1 { + t.Fatalf("db write should happen once, got=%v", workflow.dbWriteCount) + } + if len(result.Logs) != 1 || result.Logs[0] != "resume=true" { + t.Fatalf("unexpected logs: %#v", result.Logs) + } + if result.Done != "ok" { + t.Fatalf("unexpected done: %v", result.Done) + } +} diff --git a/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py b/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py index 948f74629..b591c7f87 100644 --- a/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py +++ b/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py @@ -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: diff --git a/saf-python-sdk/tests/advanced-graph/test_primitives.py b/saf-python-sdk/tests/advanced-graph/test_primitives.py index ff11b70ed..2d967a7c3 100644 --- a/saf-python-sdk/tests/advanced-graph/test_primitives.py +++ b/saf-python-sdk/tests/advanced-graph/test_primitives.py @@ -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" +