From 3987f9cd6335638c13769f66bd2470110fdf0b1c Mon Sep 17 00:00:00 2001 From: Quanzheng Long Date: Sat, 14 Mar 2026 17:47:33 -0700 Subject: [PATCH] rewrite-state-graph --- langgraph-go/advancedgraph/state_graph.go | 245 --------------- langgraph-go/stategraph/state_graph.go | 320 ++++++++++++++++++++ langgraph-go/stategraph/state_graph_test.go | 218 +++++++++++++ langgraph-go/tests/state_graph_test.go | 164 ---------- 4 files changed, 538 insertions(+), 409 deletions(-) delete mode 100644 langgraph-go/advancedgraph/state_graph.go create mode 100644 langgraph-go/stategraph/state_graph.go create mode 100644 langgraph-go/stategraph/state_graph_test.go delete mode 100644 langgraph-go/tests/state_graph_test.go diff --git a/langgraph-go/advancedgraph/state_graph.go b/langgraph-go/advancedgraph/state_graph.go deleted file mode 100644 index d117c1d38..000000000 --- a/langgraph-go/advancedgraph/state_graph.go +++ /dev/null @@ -1,245 +0,0 @@ -package advancedgraph - -import ( - "fmt" - "slices" - "sync" -) - -type StateNodeFunc[StateT any] func(ctx *Context, state StateT) (StateT, error) - -type BasicStateGraph[StateT any] struct { - nodes map[string]StateNodeFunc[StateT] - edges map[string][]string - interruptChannel string - interruptSteps map[int]struct{} -} - -func NewBasicStateGraph[StateT any]() *BasicStateGraph[StateT] { - return &BasicStateGraph[StateT]{ - nodes: make(map[string]StateNodeFunc[StateT]), - edges: make(map[string][]string), - interruptSteps: make(map[int]struct{}), - } -} - -func (g *BasicStateGraph[StateT]) AddNode(name string, fn StateNodeFunc[StateT]) { - if name == "" { - panic("node name cannot be empty") - } - if _, exists := g.nodes[name]; exists { - panic(fmt.Sprintf("node `%s` already exists", name)) - } - g.nodes[name] = fn -} - -func (g *BasicStateGraph[StateT]) AddEdge(from string, to string) { - if _, ok := g.nodes[from]; !ok { - panic(fmt.Sprintf("source node `%s` does not exist", from)) - } - if _, ok := g.nodes[to]; !ok { - panic(fmt.Sprintf("target node `%s` does not exist", to)) - } - g.edges[from] = append(g.edges[from], to) -} - -// EnableInterruptOnSuperstep enables pause/resume before dispatching a superstep. -// superstep=1 means "after first superstep has completed, before second starts". -func (g *BasicStateGraph[StateT]) EnableInterruptOnSuperstep(superstep int, channel string) { - if superstep <= 0 { - panic("interrupt superstep must be >= 1") - } - if channel == "" { - panic("interrupt channel cannot be empty") - } - if g.interruptChannel != "" && g.interruptChannel != channel { - panic("all interrupts must use the same channel") - } - g.interruptChannel = channel - g.interruptSteps[superstep] = struct{}{} -} - -type CompiledBasicStateGraph[StateT any] struct { - inner *CompiledGraph[StateT] -} - -func (g *BasicStateGraph[StateT]) Compile() *CompiledBasicStateGraph[StateT] { - if len(g.nodes) == 0 { - panic("graph has no nodes") - } - levels, err := g.computeSupersteps() - if err != nil { - panic(err) - } - adv := NewAdvancedStateGraph[StateT]() - - const finalNodeName = "__stategraph_finish" - finalNode := func(_ *Context, _ any, state StateT) (Command, error) { - return Command{Update: state}, nil - } - adv.AddFinishNodeAs(finalNodeName, finalNode) - - if g.interruptChannel != "" { - adv.AddAsyncChannel(g.interruptChannel) - } - - for stepIdx, stepNodes := range levels { - for _, nodeName := range stepNodes { - userFn := g.nodes[nodeName] - nextBarrier := fmt.Sprintf("__stategraph_barrier_%d", stepIdx+1) - wrapper := func(ctx *Context, _ any, state StateT) (Command, error) { - updated, err := userFn(ctx, state) - if err != nil { - return Command{}, err - } - return Command{ - Update: updated, - Goto: []Send{{Node: nextBarrier}}, - }, nil - } - adv.AddNodeAs(fmt.Sprintf("__stategraph_node_%s", nodeName), wrapper) - } - } - - type barrierCounter struct { - mu sync.Mutex - counts map[*RustEngine]int - } - counters := make(map[int]*barrierCounter) - for barrierStep := 1; barrierStep <= len(levels); barrierStep++ { - counters[barrierStep] = &barrierCounter{ - counts: make(map[*RustEngine]int), - } - } - - lastBarrier := len(levels) - for barrierStep := 0; barrierStep <= lastBarrier; barrierStep++ { - barrierName := fmt.Sprintf("__stategraph_barrier_%d", barrierStep) - nextStep := barrierStep - barrier := func(ctx *Context, _ any, state StateT) (Command, error) { - if nextStep > 0 { - counter := counters[nextStep] - counter.mu.Lock() - counter.counts[ctx.engine]++ - current := counter.counts[ctx.engine] - needed := len(levels[nextStep-1]) - if current < needed { - counter.mu.Unlock() - return Command{Update: state}, nil - } - delete(counter.counts, ctx.engine) - counter.mu.Unlock() - } - if _, needsInterrupt := g.interruptSteps[nextStep]; needsInterrupt { - cond := AnyOf(ChannelCondition{Channel: g.interruptChannel, N: 1}) - if _, err := ctx.WaitFor(cond); err != nil { - return Command{}, err - } - } - if nextStep >= len(levels) { - return Command{ - Update: state, - Goto: []Send{{Node: finalNodeName}}, - }, nil - } - - sends := make([]Send, 0, len(levels[nextStep])) - for _, nodeName := range levels[nextStep] { - sends = append(sends, Send{ - Node: fmt.Sprintf("__stategraph_node_%s", nodeName), - }) - } - return Command{ - Update: state, - Goto: sends, - }, nil - } - if barrierStep == 0 { - adv.AddEntryNodeAs(barrierName, barrier) - } else { - adv.AddNodeAs(barrierName, barrier) - } - } - - return &CompiledBasicStateGraph[StateT]{ - inner: adv.Compile(), - } -} - -func (g *CompiledBasicStateGraph[StateT]) Start(initialState StateT) (*Handler[StateT], error) { - return g.inner.Start(nil, initialState) -} - -func (g *CompiledBasicStateGraph[StateT]) Invoke(initialState StateT) (StateT, error) { - handler, err := g.Start(initialState) - if err != nil { - var zero StateT - return zero, err - } - return handler.WaitForResult() -} - -func (g *BasicStateGraph[StateT]) computeSupersteps() ([][]string, error) { - indegree := make(map[string]int, len(g.nodes)) - for name := range g.nodes { - indegree[name] = 0 - } - for from, tos := range g.edges { - if _, ok := g.nodes[from]; !ok { - return nil, fmt.Errorf("edge source `%s` does not exist", from) - } - for _, to := range tos { - if _, ok := g.nodes[to]; !ok { - return nil, fmt.Errorf("edge target `%s` does not exist", to) - } - indegree[to]++ - } - } - - queue := make([]string, 0, len(g.nodes)) - level := make(map[string]int, len(g.nodes)) - for name, deg := range indegree { - if deg == 0 { - queue = append(queue, name) - } - } - if len(queue) == 0 { - return nil, fmt.Errorf("graph has no entry nodes (cycle suspected)") - } - - processed := 0 - for len(queue) > 0 { - curr := queue[0] - queue = queue[1:] - processed++ - currLevel := level[curr] - for _, to := range g.edges[curr] { - if level[to] < currLevel+1 { - level[to] = currLevel + 1 - } - indegree[to]-- - if indegree[to] == 0 { - queue = append(queue, to) - } - } - } - if processed != len(g.nodes) { - return nil, fmt.Errorf("graph contains a cycle") - } - - maxLevel := 0 - for _, lv := range level { - if lv > maxLevel { - maxLevel = lv - } - } - levels := make([][]string, maxLevel+1) - for nodeName := range g.nodes { - lv := level[nodeName] - levels[lv] = append(levels[lv], nodeName) - } - for i := range levels { - slices.Sort(levels[i]) - } - return levels, nil -} diff --git a/langgraph-go/stategraph/state_graph.go b/langgraph-go/stategraph/state_graph.go new file mode 100644 index 000000000..33e238856 --- /dev/null +++ b/langgraph-go/stategraph/state_graph.go @@ -0,0 +1,320 @@ +package stategraph + +import ( + "encoding/json" + "errors" + "fmt" + "slices" + + ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph" +) + +type StateNodeFunc[StateT any] func(ctx *Context, state StateT) (StateT, error) + +const ( + internalBarrierChannel = "__stategraph_barrier" + internalInterruptChannel = "__stategraph_interrupt" +) + +type Context struct { + inner *ag.Context +} + +func (c *Context) Interrupt(name string) (any, error) { + if name == "" { + return nil, fmt.Errorf("interrupt name cannot be empty") + } + event, err := c.inner.WaitFor(ag.AnyOf(ag.ChannelCondition{ + Channel: internalInterruptChannel, + N: 1, + })) + if err != nil { + if waitReq, ok := ag.AsErrWaitRequested(err); ok { + return nil, errInterruptRequested{ + Name: name, + Condition: waitReq.Condition, + } + } + return nil, err + } + if len(event.Value) == 0 { + return nil, nil + } + + var payload interruptPayload + if err := json.Unmarshal(event.Value, &payload); err != nil { + var value any + if err := json.Unmarshal(event.Value, &value); err != nil { + return nil, fmt.Errorf("decode interrupt `%s` value: %w", name, err) + } + return value, nil + } + if payload.Name != "" && payload.Name != name { + return nil, fmt.Errorf("interrupt name mismatch: expected `%s`, got `%s`", name, payload.Name) + } + if len(payload.Value) == 0 { + return nil, nil + } + var value any + if err := json.Unmarshal(payload.Value, &value); err != nil { + return nil, fmt.Errorf("decode interrupt `%s` payload: %w", name, err) + } + return value, nil +} + +type errInterruptRequested struct { + Name string + Condition ag.AnyOfCondition +} + +func (e errInterruptRequested) Error() string { + if e.Name == "" { + return "interrupt requested" + } + return fmt.Sprintf("interrupt requested: %s", e.Name) +} + +func asErrInterruptRequested(err error) (errInterruptRequested, bool) { + var target errInterruptRequested + if !errors.As(err, &target) { + return target, false + } + return target, true +} + +type BasicStateGraph[StateT any] struct { + nodes map[string]StateNodeFunc[StateT] + edges map[string][]string +} + +type interruptPayload struct { + Name string `json:"name"` + Value json.RawMessage `json:"value"` +} + +func NewBasicStateGraph[StateT any]() *BasicStateGraph[StateT] { + return &BasicStateGraph[StateT]{ + nodes: make(map[string]StateNodeFunc[StateT]), + edges: make(map[string][]string), + } +} + +func (g *BasicStateGraph[StateT]) AddNode(fn StateNodeFunc[StateT]) string { + name := ag.NodeName(fn) + if _, exists := g.nodes[name]; exists { + panic(fmt.Sprintf("node `%s` already exists", name)) + } + g.nodes[name] = fn + return name +} + +func (g *BasicStateGraph[StateT]) AddEdge(from StateNodeFunc[StateT], to StateNodeFunc[StateT]) { + fromName := ag.NodeName(from) + toName := ag.NodeName(to) + if _, ok := g.nodes[fromName]; !ok { + panic(fmt.Sprintf("source node `%s` does not exist", fromName)) + } + if _, ok := g.nodes[toName]; !ok { + panic(fmt.Sprintf("target node `%s` does not exist", toName)) + } + g.edges[fromName] = append(g.edges[fromName], toName) +} + +type CompiledBasicStateGraph[StateT any] struct { + inner *ag.CompiledGraph[StateT] +} + +type Handler[StateT any] struct { + inner *ag.Handler[StateT] +} + +func (h *Handler[StateT]) WaitForResult() (StateT, error) { + return h.inner.WaitForResult() +} + +func (h *Handler[StateT]) Interrupt(name string, value any) error { + if name == "" { + return fmt.Errorf("interrupt name cannot be empty") + } + return h.inner.PublishToChannel(internalInterruptChannel, map[string]any{ + "name": name, + "value": value, + }) +} + +func (g *BasicStateGraph[StateT]) Compile() *CompiledBasicStateGraph[StateT] { + if len(g.nodes) == 0 { + panic("graph has no nodes") + } + levels, err := g.computeSupersteps() + if err != nil { + panic(err) + } + adv := ag.NewAdvancedStateGraph[StateT]() + adv.AddAsyncChannel(internalBarrierChannel) + adv.AddAsyncChannel(internalInterruptChannel) + + const finalNodeName = "__stategraph_finish" + finalNode := func(_ *ag.Context, _ any, state StateT) (ag.Command, error) { + return ag.Command{Update: state}, nil + } + adv.AddFinishNodeAs(finalNodeName, finalNode) + + for stepIdx, stepNodes := range levels { + for _, nodeName := range stepNodes { + userFn := g.nodes[nodeName] + nextBarrier := fmt.Sprintf("__stategraph_barrier_%d", stepIdx+1) + wrapper := func(ctx *ag.Context, _ any, state StateT) (ag.Command, error) { + updated, err := userFn(&Context{inner: ctx}, state) + if err != nil { + if interruptReq, ok := asErrInterruptRequested(err); ok { + cond := interruptReq.Condition + if len(cond.Conditions) == 0 { + cond = ag.AnyOf(ag.ChannelCondition{ + Channel: internalInterruptChannel, + N: 1, + }) + } + return ag.Command{}, ag.ErrWaitRequested{Condition: cond} + } + return ag.Command{}, err + } + if err := ctx.PublishToChannel(internalBarrierChannel, map[string]any{ + "step": stepIdx, + }); err != nil { + return ag.Command{}, err + } + return ag.Command{ + Update: updated, + Goto: []ag.Send{{Node: nextBarrier}}, + }, nil + } + adv.AddNodeAs(fmt.Sprintf("__stategraph_node_%s", nodeName), wrapper) + } + } + + lastBarrier := len(levels) + for barrierStep := 0; barrierStep <= lastBarrier; barrierStep++ { + barrierName := fmt.Sprintf("__stategraph_barrier_%d", barrierStep) + nextStep := barrierStep + barrier := func(ctx *ag.Context, _ any, state StateT) (ag.Command, error) { + if nextStep > 0 { + needed := len(levels[nextStep-1]) + _, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{ + Channel: internalBarrierChannel, + N: needed, + })) + if err != nil { + return ag.Command{}, err + } + } + if nextStep >= len(levels) { + return ag.Command{ + Update: state, + Goto: []ag.Send{{Node: finalNodeName}}, + }, nil + } + sends := make([]ag.Send, 0, len(levels[nextStep])) + for _, nodeName := range levels[nextStep] { + sends = append(sends, ag.Send{ + Node: fmt.Sprintf("__stategraph_node_%s", nodeName), + }) + } + return ag.Command{ + Update: state, + Goto: sends, + }, nil + } + if barrierStep == 0 { + adv.AddEntryNodeAs(barrierName, barrier) + } else { + adv.AddNodeAs(barrierName, barrier) + } + } + + return &CompiledBasicStateGraph[StateT]{ + inner: adv.Compile(), + } +} + +func (g *CompiledBasicStateGraph[StateT]) Start(initialState StateT) (*Handler[StateT], error) { + raw, err := g.inner.Start(nil, initialState) + if err != nil { + return nil, err + } + return &Handler[StateT]{inner: raw}, nil +} + +func (g *CompiledBasicStateGraph[StateT]) Invoke(initialState StateT) (StateT, error) { + handler, err := g.Start(initialState) + if err != nil { + var zero StateT + return zero, err + } + return handler.WaitForResult() +} + +func (g *BasicStateGraph[StateT]) computeSupersteps() ([][]string, error) { + indegree := make(map[string]int, len(g.nodes)) + for name := range g.nodes { + indegree[name] = 0 + } + for from, tos := range g.edges { + if _, ok := g.nodes[from]; !ok { + return nil, fmt.Errorf("edge source `%s` does not exist", from) + } + for _, to := range tos { + if _, ok := g.nodes[to]; !ok { + return nil, fmt.Errorf("edge target `%s` does not exist", to) + } + indegree[to]++ + } + } + + queue := make([]string, 0, len(g.nodes)) + level := make(map[string]int, len(g.nodes)) + for name, deg := range indegree { + if deg == 0 { + queue = append(queue, name) + } + } + if len(queue) == 0 { + return nil, fmt.Errorf("graph has no entry nodes (cycle suspected)") + } + + processed := 0 + for len(queue) > 0 { + curr := queue[0] + queue = queue[1:] + processed++ + currLevel := level[curr] + for _, to := range g.edges[curr] { + if level[to] < currLevel+1 { + level[to] = currLevel + 1 + } + indegree[to]-- + if indegree[to] == 0 { + queue = append(queue, to) + } + } + } + if processed != len(g.nodes) { + return nil, fmt.Errorf("graph contains a cycle") + } + + maxLevel := 0 + for _, lv := range level { + if lv > maxLevel { + maxLevel = lv + } + } + levels := make([][]string, maxLevel+1) + for nodeName := range g.nodes { + lv := level[nodeName] + levels[lv] = append(levels[lv], nodeName) + } + for i := range levels { + slices.Sort(levels[i]) + } + return levels, nil +} diff --git a/langgraph-go/stategraph/state_graph_test.go b/langgraph-go/stategraph/state_graph_test.go new file mode 100644 index 000000000..7e03f8e41 --- /dev/null +++ b/langgraph-go/stategraph/state_graph_test.go @@ -0,0 +1,218 @@ +package stategraph_test + +import ( + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + sg "github.com/langchain-ai/langgraph/langgraph-go/stategraph" +) + +type stateGraphState struct { + Noop bool `json:"noop"` +} + +type orderRecorder struct { + mu sync.Mutex + orders map[string]int32 + seq int32 +} + +func newOrderRecorder() *orderRecorder { + return &orderRecorder{orders: make(map[string]int32)} +} + +func (f *orderRecorder) record(name string) { + idx := atomic.AddInt32(&f.seq, 1) + f.mu.Lock() + f.orders[name] = idx + f.mu.Unlock() +} + +type orderGraph struct { + recorder *orderRecorder +} + +func newOrderGraph() *orderGraph { + return &orderGraph{ + recorder: newOrderRecorder(), + } +} + +func (g *orderGraph) A(_ *sg.Context, state stateGraphState) (stateGraphState, error) { + g.recorder.record("A") + return state, nil +} + +func (g *orderGraph) B1(_ *sg.Context, state stateGraphState) (stateGraphState, error) { + g.recorder.record("B1") + return state, nil +} + +func (g *orderGraph) B2(_ *sg.Context, state stateGraphState) (stateGraphState, error) { + g.recorder.record("B2") + return state, nil +} + +func (g *orderGraph) C1(_ *sg.Context, state stateGraphState) (stateGraphState, error) { + g.recorder.record("C1") + return state, nil +} + +func (g *orderGraph) C2(_ *sg.Context, state stateGraphState) (stateGraphState, error) { + g.recorder.record("C2") + return state, nil +} + +func (g *orderGraph) C3(_ *sg.Context, state stateGraphState) (stateGraphState, error) { + g.recorder.record("C3") + return state, nil +} + +func (g *orderGraph) D(_ *sg.Context, state stateGraphState) (stateGraphState, error) { + g.recorder.record("D") + return state, nil +} + +func TestBasicStateGraphWithoutInterrupt(t *testing.T) { + fixture := newOrderGraph() + graph := sg.NewBasicStateGraph[stateGraphState]() + graph.AddNode(fixture.A) + graph.AddNode(fixture.B1) + graph.AddNode(fixture.B2) + graph.AddNode(fixture.C1) + graph.AddNode(fixture.C2) + graph.AddNode(fixture.C3) + graph.AddNode(fixture.D) + + graph.AddEdge(fixture.A, fixture.B1) + graph.AddEdge(fixture.A, fixture.B2) + graph.AddEdge(fixture.B1, fixture.C1) + graph.AddEdge(fixture.B1, fixture.C2) + graph.AddEdge(fixture.B2, fixture.C3) + graph.AddEdge(fixture.C1, fixture.D) + graph.AddEdge(fixture.C2, fixture.D) + graph.AddEdge(fixture.C3, fixture.D) + + _, err := graph.Compile().Invoke(stateGraphState{}) + if err != nil { + t.Fatalf("invoke failed: %v", err) + } + + fixture.recorder.mu.Lock() + orders := make(map[string]int32, len(fixture.recorder.orders)) + for k, v := range fixture.recorder.orders { + orders[k] = v + } + fixture.recorder.mu.Unlock() + + for _, name := range []string{"A", "B1", "B2", "C1", "C2", "C3", "D"} { + if _, ok := orders[name]; !ok { + t.Fatalf("node %s did not execute; orders=%v", name, orders) + } + } + maxB := maxInt32(orders["B1"], orders["B2"]) + minC := minInt32(orders["C1"], minInt32(orders["C2"], orders["C3"])) + maxC := maxInt32(orders["C1"], maxInt32(orders["C2"], orders["C3"])) + if !(orders["A"] < orders["B1"] && orders["A"] < orders["B2"]) { + t.Fatalf("A should run before B-step, orders=%v", orders) + } + if !(maxB < minC) { + t.Fatalf("B-step should finish before C-step, orders=%v", orders) + } + if !(maxC < orders["D"]) { + t.Fatalf("C-step should finish before D, orders=%v", orders) + } +} + +type interruptState struct { + A bool `json:"a"` + B bool `json:"b"` +} + +type interruptFixture struct{} + +func (f *interruptFixture) A(_ *sg.Context, state interruptState) (interruptState, error) { + state.A = true + return state, nil +} + +func (f *interruptFixture) B(ctx *sg.Context, state interruptState) (interruptState, error) { + if !state.A { + return state, fmt.Errorf("B should observe A=true") + } + value, err := ctx.Interrupt("resume_channel") + if err != nil { + return state, err + } + s, ok := value.(string) + if !ok || s != "go" { + return state, fmt.Errorf("unexpected interrupt payload: %#v", value) + } + state.B = true + return state, nil +} + +func TestBasicStateGraphWithInterrupt(t *testing.T) { + fixture := &interruptFixture{} + graph := sg.NewBasicStateGraph[interruptState]() + graph.AddNode(fixture.A) + graph.AddNode(fixture.B) + graph.AddEdge(fixture.A, fixture.B) + + handler, err := graph.Compile().Start(interruptState{}) + if err != nil { + t.Fatalf("start failed: %v", err) + } + + doneCh := make(chan interruptState, 1) + errCh := make(chan error, 1) + go func() { + result, runErr := handler.WaitForResult() + if runErr != nil { + errCh <- runErr + return + } + doneCh <- result + }() + + select { + case <-doneCh: + t.Fatalf("run should pause for interrupt, but completed early") + case err := <-errCh: + t.Fatalf("run should pause for interrupt, but failed early: %v", err) + case <-time.After(120 * time.Millisecond): + // expected: paused + } + + if err := handler.Interrupt("resume_channel", "go"); err != nil { + t.Fatalf("resume interrupt failed: %v", err) + } + + select { + case err := <-errCh: + t.Fatalf("run failed after interrupt: %v", err) + case result := <-doneCh: + if !(result.A && result.B) { + t.Fatalf("unexpected final state: %#v", result) + } + case <-time.After(2 * time.Second): + t.Fatalf("timeout waiting for resumed run completion") + } +} + +func minInt32(a int32, b int32) int32 { + if a < b { + return a + } + return b +} + +func maxInt32(a int32, b int32) int32 { + if a > b { + return a + } + return b +} diff --git a/langgraph-go/tests/state_graph_test.go b/langgraph-go/tests/state_graph_test.go deleted file mode 100644 index 002b6c99e..000000000 --- a/langgraph-go/tests/state_graph_test.go +++ /dev/null @@ -1,164 +0,0 @@ -package tests - -import ( - "fmt" - "sync" - "sync/atomic" - "testing" - "time" - - ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph" -) - -type stateGraphState struct { - Noop bool `json:"noop"` -} - -func TestBasicStateGraphWithoutInterrupt(t *testing.T) { - var ( - orderMu sync.Mutex - orders = make(map[string]int32) - seq int32 - ) - record := func(name string) { - idx := atomic.AddInt32(&seq, 1) - orderMu.Lock() - orders[name] = idx - orderMu.Unlock() - } - graph := ag.NewBasicStateGraph[stateGraphState]() - graph.AddNode("A", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) { - record("A") - return state, nil - }) - graph.AddNode("B1", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) { - record("B1") - return state, nil - }) - graph.AddNode("B2", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) { - record("B2") - return state, nil - }) - graph.AddNode("C1", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) { - record("C1") - return state, nil - }) - graph.AddNode("C2", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) { - record("C2") - return state, nil - }) - graph.AddNode("C3", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) { - record("C3") - return state, nil - }) - graph.AddNode("D", func(_ *ag.Context, state stateGraphState) (stateGraphState, error) { - record("D") - return state, nil - }) - - graph.AddEdge("A", "B1") - graph.AddEdge("A", "B2") - graph.AddEdge("B1", "C1") - graph.AddEdge("B1", "C2") - graph.AddEdge("B2", "C3") - graph.AddEdge("C1", "D") - graph.AddEdge("C2", "D") - graph.AddEdge("C3", "D") - - _, err := graph.Compile().Invoke(stateGraphState{}) - if err != nil { - t.Fatalf("invoke failed: %v", err) - } - for _, name := range []string{"A", "B1", "B2", "C1", "C2", "C3", "D"} { - if _, ok := orders[name]; !ok { - t.Fatalf("node %s did not execute; orders=%v", name, orders) - } - } - maxB := maxInt32(orders["B1"], orders["B2"]) - minC := minInt32(orders["C1"], minInt32(orders["C2"], orders["C3"])) - maxC := maxInt32(orders["C1"], maxInt32(orders["C2"], orders["C3"])) - if !(orders["A"] < orders["B1"] && orders["A"] < orders["B2"]) { - t.Fatalf("A should run before B-step, orders=%v", orders) - } - if !(maxB < minC) { - t.Fatalf("B-step should finish before C-step, orders=%v", orders) - } - if !(maxC < orders["D"]) { - t.Fatalf("C-step should finish before D, orders=%v", orders) - } -} - -func TestBasicStateGraphWithInterrupt(t *testing.T) { - type interruptState struct { - A bool `json:"a"` - B bool `json:"b"` - } - graph := ag.NewBasicStateGraph[interruptState]() - graph.AddNode("A", func(_ *ag.Context, state interruptState) (interruptState, error) { - state.A = true - return state, nil - }) - graph.AddNode("B", func(_ *ag.Context, state interruptState) (interruptState, error) { - if !state.A { - return state, fmt.Errorf("B should observe A=true") - } - state.B = true - return state, nil - }) - graph.AddEdge("A", "B") - graph.EnableInterruptOnSuperstep(1, "resume_channel") - - handler, err := graph.Compile().Start(interruptState{}) - if err != nil { - t.Fatalf("start failed: %v", err) - } - - doneCh := make(chan interruptState, 1) - errCh := make(chan error, 1) - go func() { - result, runErr := handler.WaitForResult() - if runErr != nil { - errCh <- runErr - return - } - doneCh <- result - }() - - select { - case <-doneCh: - t.Fatalf("run should pause for interrupt, but completed early") - case err := <-errCh: - t.Fatalf("run should pause for interrupt, but failed early: %v", err) - case <-time.After(120 * time.Millisecond): - // expected: paused - } - - if err := handler.PublishToChannel("resume_channel", "go"); err != nil { - t.Fatalf("publish interrupt failed: %v", err) - } - - select { - case err := <-errCh: - t.Fatalf("run failed after interrupt: %v", err) - case result := <-doneCh: - if !(result.A && result.B) { - t.Fatalf("unexpected final state: %#v", result) - } - case <-time.After(2 * time.Second): - t.Fatalf("timeout waiting for resumed run completion") - } -} - -func minInt32(a int32, b int32) int32 { - if a < b { - return a - } - return b -} - -func maxInt32(a int32, b int32) int32 { - if a > b { - return a - } - return b -}