more

langgraph-go/advancedgraph/graph.go

@@ -31,6 +31,10 @@ func NewAdvancedStateGraph[StateT any]() *AdvancedStateGraph[StateT] {
 
 func (g *AdvancedStateGraph[StateT]) AddNode(fn any) string {
 	name := NodeName(fn)
+	return g.AddNodeAs(name, fn)
+}
+
+func (g *AdvancedStateGraph[StateT]) AddNodeAs(name string, fn any) string {
 	if _, exists := g.nodes[name]; exists {
 		panic(fmt.Sprintf("node `%s` already exists", name))
 	}
@@ -47,13 +51,23 @@ func (g *AdvancedStateGraph[StateT]) AddAsyncChannel(name string) {
 }
 
 func (g *AdvancedStateGraph[StateT]) AddEntryNode(fn any) string {
-	name := g.AddNode(fn)
+	name := NodeName(fn)
+	return g.AddEntryNodeAs(name, fn)
+}
+
+func (g *AdvancedStateGraph[StateT]) AddEntryNodeAs(name string, fn any) string {
+	name = g.AddNodeAs(name, fn)
 	g.entryPoint = name
 	return name
 }
 
 func (g *AdvancedStateGraph[StateT]) AddFinishNode(fn any) string {
-	name := g.AddNode(fn)
+	name := NodeName(fn)
+	return g.AddFinishNodeAs(name, fn)
+}
+
+func (g *AdvancedStateGraph[StateT]) AddFinishNodeAs(name string, fn any) string {
+	name = g.AddNodeAs(name, fn)
 	g.finishPoint = name
 	return name
 }
@@ -331,4 +345,3 @@ func unwrapResumeInput(input any) (any, *WaitEvent) {
 	}
 	return rawArg, &event
 }
-

langgraph-go/advancedgraph/rust_engine.go

@@ -12,6 +12,7 @@ import "C"
 import (
 	"encoding/json"
 	"fmt"
+	"reflect"
 	"runtime/cgo"
 	"unsafe"
 )
@@ -59,8 +60,12 @@ func goNodeCallback(userData C.ulong, node *C.char, argJSON *C.char, stateJSON *
 
 	sends := make([]map[string]any, 0, len(cmd.Goto))
 	for _, send := range cmd.Goto {
+		targetNode, err := resolveSendTarget(send.Node)
+		if err != nil {
+			return cCallbackEnvelopeError(err.Error())
+		}
 		sends = append(sends, map[string]any{
-			"node": NodeName(send.Node),
+			"node": targetNode,
 			"arg":  send.NodeInput,
 		})
 	}
@@ -215,6 +220,20 @@ func cCallbackEnvelopeSuspend(cond AnyOfCondition) *C.char {
 	return C.CString(string(raw))
 }
 
+func resolveSendTarget(target any) (string, error) {
+	if name, ok := target.(string); ok {
+		if name == "" {
+			return "", fmt.Errorf("send target cannot be empty string")
+		}
+		return name, nil
+	}
+	rv := reflect.ValueOf(target)
+	if rv.IsValid() && rv.Kind() == reflect.Func {
+		return NodeName(target), nil
+	}
+	return "", fmt.Errorf("unsupported send target type %T", target)
+}
+
 func coerceJSONValue(v any) any {
 	switch t := v.(type) {
 	case map[string]any:

langgraph-go/advancedgraph/state_graph.go

@@ -0,0 +1,245 @@
+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
+}

langgraph-go/tests/state_graph_test.go

@@ -0,0 +1,164 @@
+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
+}