locking

.gitignore

@@ -100,3 +100,10 @@ dmypy.json
 .turbo
 .editorconfig
 .scratch
+
+# macOS debug symbol bundles generated during local Rust builds
+saf-python-sdk/python/saf_python_sdk/*.dSYM/
+
+# Local PyO3 extension artifacts for saf-python-sdk
+saf-python-sdk/python/saf_python_sdk/langgraph_rust_core*.so
+saf-python-sdk/python/saf_python_sdk/langgraph_rust_core*.pyd

langgraph-go/advancedgraph/graph.go

@@ -10,9 +10,30 @@ import (
 
 type nodeExecutor func(ctx *Context, input any, state map[string]any) (Command, error)
 
+type nodeConfig struct {
+	lockedFields []string
+}
+
+type NodeOption interface {
+	applyToNodeConfig(*nodeConfig)
+}
+
+type NodeStateOption struct {
+	LockedFields []string
+}
+
+func (o NodeStateOption) applyToNodeConfig(cfg *nodeConfig) {
+	if cfg == nil {
+		return
+	}
+	cfg.lockedFields = append(cfg.lockedFields[:0], o.LockedFields...)
+}
+
 type AdvancedStateGraph[StateT any] struct {
 	nodes         map[string]nodeExecutor
+	nodeOptions   map[string]nodeConfig
 	asyncChannels []string
+	customStreams []string
 	entryPoint    string
 	finishPoint   string
 	stateType     reflect.Type
@@ -24,20 +45,21 @@ func NewAdvancedStateGraph[StateT any]() *AdvancedStateGraph[StateT] {
 		panic(fmt.Sprintf("StateT must be a struct, got %s", stateType.String()))
 	}
 	return &AdvancedStateGraph[StateT]{
-		nodes:     make(map[string]nodeExecutor),
-		stateType: stateType,
+		nodes:       make(map[string]nodeExecutor),
+		nodeOptions: make(map[string]nodeConfig),
+		stateType:   stateType,
 	}
 }
 
 // AddNode keeps `fn` as `any` because advanced graph nodes can have different
 // input argument types per node, while only `StateT` is globally constrained.
 // We validate and adapt node signatures at runtime in compileNodeExecutor.
-func (g *AdvancedStateGraph[StateT]) AddNode(fn any) string {
+func (g *AdvancedStateGraph[StateT]) AddNode(fn any, nodeOption ...NodeOption) string {
 	name := NodeName(fn)
-	return g.AddNodeAs(name, fn)
+	return g.AddNodeAs(name, fn, nodeOption...)
 }
 
-func (g *AdvancedStateGraph[StateT]) AddNodeAs(name string, fn any) string {
+func (g *AdvancedStateGraph[StateT]) AddNodeAs(name string, fn any, nodeOption ...NodeOption) string {
 	if _, exists := g.nodes[name]; exists {
 		panic(fmt.Sprintf("node `%s` already exists", name))
 	}
@@ -46,6 +68,7 @@ func (g *AdvancedStateGraph[StateT]) AddNodeAs(name string, fn any) string {
 		panic(err)
 	}
 	g.nodes[name] = exec
+	g.nodeOptions[name] = resolveNodeConfig(nodeOption...)
 	return name
 }
 
@@ -53,24 +76,28 @@ func (g *AdvancedStateGraph[StateT]) AddAsyncChannel(name string) {
 	g.asyncChannels = append(g.asyncChannels, name)
 }
 
-func (g *AdvancedStateGraph[StateT]) AddEntryNode(fn any) string {
-	name := NodeName(fn)
-	return g.AddEntryNodeAs(name, fn)
+func (g *AdvancedStateGraph[StateT]) AddCustomOutputStream(name string) {
+	g.customStreams = append(g.customStreams, name)
 }
 
-func (g *AdvancedStateGraph[StateT]) AddEntryNodeAs(name string, fn any) string {
-	name = g.AddNodeAs(name, fn)
+func (g *AdvancedStateGraph[StateT]) AddEntryNode(fn any, nodeOption ...NodeOption) string {
+	name := NodeName(fn)
+	return g.AddEntryNodeAs(name, fn, nodeOption...)
+}
+
+func (g *AdvancedStateGraph[StateT]) AddEntryNodeAs(name string, fn any, nodeOption ...NodeOption) string {
+	name = g.AddNodeAs(name, fn, nodeOption...)
 	g.entryPoint = name
 	return name
 }
 
-func (g *AdvancedStateGraph[StateT]) AddFinishNode(fn any) string {
+func (g *AdvancedStateGraph[StateT]) AddFinishNode(fn any, nodeOption ...NodeOption) string {
 	name := NodeName(fn)
-	return g.AddFinishNodeAs(name, fn)
+	return g.AddFinishNodeAs(name, fn, nodeOption...)
 }
 
-func (g *AdvancedStateGraph[StateT]) AddFinishNodeAs(name string, fn any) string {
-	name = g.AddNodeAs(name, fn)
+func (g *AdvancedStateGraph[StateT]) AddFinishNodeAs(name string, fn any, nodeOption ...NodeOption) string {
+	name = g.AddNodeAs(name, fn, nodeOption...)
 	g.finishPoint = name
 	return name
 }
@@ -78,7 +105,9 @@ func (g *AdvancedStateGraph[StateT]) AddFinishNodeAs(name string, fn any) string
 func (g *AdvancedStateGraph[StateT]) Compile() *CompiledGraph[StateT] {
 	return &CompiledGraph[StateT]{
 		nodes:         g.nodes,
+		nodeOptions:   g.nodeOptions,
 		asyncChannels: g.asyncChannels,
+		customStreams: g.customStreams,
 		entryPoint:    g.entryPoint,
 		finishPoint:   g.finishPoint,
 		stateType:     g.stateType,
@@ -87,7 +116,9 @@ func (g *AdvancedStateGraph[StateT]) Compile() *CompiledGraph[StateT] {
 
 type CompiledGraph[StateT any] struct {
 	nodes         map[string]nodeExecutor
+	nodeOptions   map[string]nodeConfig
 	asyncChannels []string
+	customStreams []string
 	entryPoint    string
 	finishPoint   string
 	stateType     reflect.Type
@@ -96,28 +127,37 @@ type CompiledGraph[StateT any] struct {
 type Context struct {
 	engine      *RustEngine
 	resumeEvent *WaitEvent
+	isResume    bool
 }
 
-func (c *Context) WaitFor(cond AnyOfCondition) (WaitEvent, error) {
+func (c *Context) WaitFor(target WaitTarget) (WaitForResult, error) {
+	if target == nil {
+		return WaitForResult{}, fmt.Errorf("wait target cannot be nil")
+	}
 	if c.resumeEvent != nil {
 		event := *c.resumeEvent
 		c.resumeEvent = nil
-		return event, nil
+		return waitForResultFromRaw(target, event), nil
 	}
-	return WaitEvent{}, ErrWaitRequested{Condition: cond}
+	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)
 }
 
-func (c *Context) SendCustomStreamEvent(value any) error {
-	return c.engine.SendCustomStreamEvent(value)
+func (c *Context) SendCustomStreamEvent(streamName string, value any) error {
+	return c.engine.SendCustomStreamEvent(streamName, value)
 }
 
 type Handler[StateT any] struct {
-	engine *RustEngine
-	done   chan resultOrErr[StateT]
+	engine       *RustEngine
+	done         chan resultOrErr[StateT]
+	streamReadyC chan struct{}
 }
 
 type resultOrErr[StateT any] struct {
@@ -134,8 +174,9 @@ func (h *Handler[StateT]) WaitForResult() (StateT, error) {
 	return res.state, res.err
 }
 
-func (h *Handler[StateT]) ReceiveStream() (any, error) {
-	event, hasEvent, err := h.engine.ReceiveStream()
+func (h *Handler[StateT]) ReceiveStream(streamName string) (any, error) {
+	<-h.streamReadyC
+	event, hasEvent, err := h.engine.ReceiveStream(streamName)
 	if err != nil {
 		return nil, err
 	}
@@ -145,8 +186,9 @@ func (h *Handler[StateT]) ReceiveStream() (any, error) {
 	return event, nil
 }
 
-func (h *Handler[StateT]) CloseStream() error {
-	return h.engine.CloseStream()
+func (h *Handler[StateT]) CloseAllStreams() error {
+	<-h.streamReadyC
+	return h.engine.CloseAllStreams()
 }
 
 func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT, streamMode ...string) (*Handler[StateT], error) {
@@ -163,19 +205,37 @@ func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT, str
 			return nil, err
 		}
 	}
+	for _, streamName := range g.customStreams {
+		if err := engine.AddCustomOutputStream(streamName); err != nil {
+			return nil, err
+		}
+	}
 
 	handler := &Handler[StateT]{
-		engine: engine,
-		done:   make(chan resultOrErr[StateT], 1),
+		engine:       engine,
+		done:         make(chan resultOrErr[StateT], 1),
+		streamReadyC: make(chan struct{}),
 	}
 	go func() {
 		defer engine.Close()
+		streamModeForRun := resolvedStreamMode
+		if resolvedStreamMode != "" {
+			if err := engine.StartStream(resolvedStreamMode); err != nil {
+				close(handler.streamReadyC)
+				handler.done <- resultOrErr[StateT]{err: err}
+				close(handler.done)
+				return
+			}
+			streamModeForRun = ""
+		}
+		close(handler.streamReadyC)
 		rawState, err := engine.RunGraph(
 			g.entryPoint,
 			g.finishPoint,
-			resolvedStreamMode,
+			streamModeForRun,
 			initialState,
 			initialInput,
+			g.nodeLockedFields(),
 			func(node string, nodeInput any, fallbackState map[string]any) (Command, error) {
 				fn, ok := g.nodes[node]
 				if !ok {
@@ -185,7 +245,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 {
@@ -200,6 +268,37 @@ func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT, str
 	return handler, nil
 }
 
+func (g *CompiledGraph[StateT]) nodeLockedFields() map[string][]string {
+	result := make(map[string][]string, len(g.nodeOptions))
+	for nodeName, option := range g.nodeOptions {
+		if len(option.lockedFields) == 0 {
+			continue
+		}
+		fields := make([]string, 0, len(option.lockedFields))
+		for _, field := range option.lockedFields {
+			if field == "" {
+				continue
+			}
+			fields = append(fields, field)
+		}
+		if len(fields) > 0 {
+			result[nodeName] = fields
+		}
+	}
+	return result
+}
+
+func resolveNodeConfig(nodeOption ...NodeOption) nodeConfig {
+	cfg := nodeConfig{}
+	for _, option := range nodeOption {
+		if option == nil {
+			panic("node option cannot be nil")
+		}
+		option.applyToNodeConfig(&cfg)
+	}
+	return cfg
+}
+
 func NodeName(fn any) string {
 	rv := reflect.ValueOf(fn)
 	if !rv.IsValid() || rv.Kind() != reflect.Func {
@@ -375,3 +474,108 @@ func unwrapResumeInput(input any) (any, *WaitEvent) {
 	}
 	return rawArg, &event
 }
+
+func waitForResultFromRaw(target WaitTarget, event WaitEvent) WaitForResult {
+	conditions := target.waitConditions()
+	result := WaitForResult{
+		Conditions: make([]ConditionResult, len(conditions)),
+	}
+	if target.waitKind() == "all_of" {
+		for i, cond := range conditions {
+			if isTimerCondition(cond) {
+				result.Conditions[i] = ConditionResult{Met: true}
+			}
+		}
+	}
+
+	if event.Condition == "timer" {
+		for i, cond := range conditions {
+			if isTimerCondition(cond) {
+				result.Conditions[i] = ConditionResult{Met: true}
+				if target.waitKind() == "any_of" {
+					break
+				}
+			}
+		}
+		return result
+	}
+
+	if event.Condition != "channel" {
+		return result
+	}
+
+	if event.Channel == "__any_of__" || event.Channel == "__all_of__" {
+		var matched []struct {
+			Channel string `json:"channel"`
+			Value   any    `json:"value"`
+		}
+		_ = json.Unmarshal(event.Value, &matched)
+		cursor := 0
+		for i, cond := range conditions {
+			channelName, ok := channelNameOfCondition(cond)
+			if !ok || cursor >= len(matched) {
+				continue
+			}
+			if channelName == matched[cursor].Channel {
+				result.Conditions[i] = ConditionResult{
+					Met:         true,
+					ChannelName: channelName,
+					Values:      toValues(matched[cursor].Value),
+				}
+				cursor++
+			}
+		}
+		return result
+	}
+
+	var value any
+	_ = json.Unmarshal(event.Value, &value)
+	for i, cond := range conditions {
+		channelName, ok := channelNameOfCondition(cond)
+		if !ok {
+			continue
+		}
+		if channelName == event.Channel {
+			result.Conditions[i] = ConditionResult{
+				Met:         true,
+				ChannelName: channelName,
+				Values:      toValues(value),
+			}
+			break
+		}
+	}
+	return result
+}
+
+func isTimerCondition(cond WaitCondition) bool {
+	switch cond.(type) {
+	case TimerCondition, *TimerCondition:
+		return true
+	default:
+		return false
+	}
+}
+
+func channelNameOfCondition(cond WaitCondition) (string, bool) {
+	switch c := cond.(type) {
+	case ChannelCondition:
+		return c.Channel, true
+	case *ChannelCondition:
+		if c == nil {
+			return "", false
+		}
+		return c.Channel, true
+	default:
+		return "", false
+	}
+}
+
+func toValues(value any) []any {
+	if value == nil {
+		return []any{}
+	}
+	if vals, ok := value.([]any); ok {
+		return vals
+	}
+	return []any{value}
+}

langgraph-go/advancedgraph/rust_engine.go

@@ -6,6 +6,28 @@ package advancedgraph
 #include "langgraph_rust_core.h"
 #include <stdlib.h>
 extern char* goNodeCallback(unsigned long user_data, char* node, char* arg_json, char* state_json);
+static inline char* rc_run_graph_json_with_go_callback(
+    Engine* ptr,
+    const char* entry_point,
+    const char* finish_point,
+    const char* initial_state_json,
+    const char* initial_input_json,
+    const char* stream_mode,
+    const char* node_locked_fields_json,
+    unsigned long user_data
+) {
+    return rc_run_graph_json(
+        ptr,
+        entry_point,
+        finish_point,
+        initial_state_json,
+        initial_input_json,
+        stream_mode,
+        node_locked_fields_json,
+        user_data,
+        goNodeCallback
+    );
+}
 */
 import "C"
 
@@ -80,7 +102,7 @@ func goNodeCallback(userData C.ulong, node *C.char, argJSON *C.char, stateJSON *
 	cmd, err := ctx.exec(nodeName, nodeInput, state)
 	if err != nil {
 		if waitReq, ok := AsErrWaitRequested(err); ok {
-			return cCallbackEnvelopeSuspend(waitReq.Condition)
+			return cCallbackEnvelopeSuspend(waitReq.Target)
 		}
 		return cCallbackEnvelopeError(err.Error())
 	}
@@ -128,6 +150,13 @@ func (e *RustEngine) AddAsyncChannel(channel string) error {
 	return parseRustStatus(resp)
 }
 
+func (e *RustEngine) AddCustomOutputStream(streamName string) error {
+	cname := C.CString(streamName)
+	defer C.free(unsafe.Pointer(cname))
+	resp := C.rc_add_custom_output_stream(e.ptr, cname)
+	return parseRustStatus(resp)
+}
+
 func (e *RustEngine) StartStream(streamMode string) error {
 	var cmode *C.char
 	if streamMode != "" {
@@ -138,8 +167,10 @@ func (e *RustEngine) StartStream(streamMode string) error {
 	return parseRustStatus(resp)
 }
 
-func (e *RustEngine) ReceiveStream() (any, bool, error) {
-	resp := C.rc_receive_stream_json(e.ptr)
+func (e *RustEngine) ReceiveStream(streamName string) (any, bool, error) {
+	cname := C.CString(streamName)
+	defer C.free(unsafe.Pointer(cname))
+	resp := C.rc_receive_stream_json(e.ptr, cname)
 	defer C.rc_string_free(resp)
 
 	raw := C.GoString(resp)
@@ -165,19 +196,21 @@ func (e *RustEngine) ReceiveStream() (any, bool, error) {
 	return coerceJSONValue(event), true, nil
 }
 
-func (e *RustEngine) SendCustomStreamEvent(value any) error {
+func (e *RustEngine) SendCustomStreamEvent(streamName string, value any) error {
 	payload, err := json.Marshal(value)
 	if err != nil {
 		return fmt.Errorf("marshal stream event: %w", err)
 	}
+	cname := C.CString(streamName)
 	cval := C.CString(string(payload))
+	defer C.free(unsafe.Pointer(cname))
 	defer C.free(unsafe.Pointer(cval))
-	resp := C.rc_send_custom_stream_event(e.ptr, cval)
+	resp := C.rc_send_custom_stream_event(e.ptr, cname, cval)
 	return parseRustStatus(resp)
 }
 
-func (e *RustEngine) CloseStream() error {
-	resp := C.rc_close_stream(e.ptr)
+func (e *RustEngine) CloseAllStreams() error {
+	resp := C.rc_close_all_streams(e.ptr)
 	return parseRustStatus(resp)
 }
 
@@ -223,12 +256,42 @@ func (e *RustEngine) WaitAnyOf(cond AnyOfCondition) (WaitEvent, error) {
 	return event, nil
 }
 
+func (e *RustEngine) WaitAllOf(cond AllOfCondition) (WaitEvent, error) {
+	payload, err := json.Marshal(cond)
+	if err != nil {
+		return WaitEvent{}, fmt.Errorf("marshal all_of: %w", err)
+	}
+	cpayload := C.CString(string(payload))
+	defer C.free(unsafe.Pointer(cpayload))
+	resp := C.rc_wait_all_of_json(e.ptr, cpayload)
+	defer C.rc_string_free(resp)
+
+	raw := C.GoString(resp)
+	var status struct {
+		OK    bool            `json:"ok"`
+		Error string          `json:"error"`
+		Event json.RawMessage `json:"event"`
+	}
+	if err := json.Unmarshal([]byte(raw), &status); err != nil {
+		return WaitEvent{}, fmt.Errorf("decode rust wait response: %w", err)
+	}
+	if !status.OK {
+		return WaitEvent{}, fmt.Errorf("rust wait failed: %s", status.Error)
+	}
+	var event WaitEvent
+	if err := json.Unmarshal(status.Event, &event); err != nil {
+		return WaitEvent{}, fmt.Errorf("decode wait event: %w", err)
+	}
+	return event, nil
+}
+
 func (e *RustEngine) RunGraph(
 	entryPoint string,
 	finishPoint string,
 	streamMode string,
 	initialState any,
 	initialInput any,
+	nodeLockedFields map[string][]string,
 	exec func(node string, nodeInput any, state map[string]any) (Command, error),
 ) (map[string]any, error) {
 	initialJSON, err := json.Marshal(initialState)
@@ -239,10 +302,15 @@ func (e *RustEngine) RunGraph(
 	if err != nil {
 		return nil, fmt.Errorf("marshal initial input: %w", err)
 	}
+	lockedFieldsJSON, err := json.Marshal(nodeLockedFields)
+	if err != nil {
+		return nil, fmt.Errorf("marshal node locked fields: %w", err)
+	}
 	centry := C.CString(entryPoint)
 	cfinish := C.CString(finishPoint)
 	cinitial := C.CString(string(initialJSON))
 	cinitialInput := C.CString(string(initialInputJSON))
+	clockedFields := C.CString(string(lockedFieldsJSON))
 	var cstreamMode *C.char
 	if streamMode != "" {
 		cstreamMode = C.CString(streamMode)
@@ -251,6 +319,7 @@ func (e *RustEngine) RunGraph(
 	defer C.free(unsafe.Pointer(cfinish))
 	defer C.free(unsafe.Pointer(cinitial))
 	defer C.free(unsafe.Pointer(cinitialInput))
+	defer C.free(unsafe.Pointer(clockedFields))
 	if cstreamMode != nil {
 		defer C.free(unsafe.Pointer(cstreamMode))
 	}
@@ -258,15 +327,15 @@ func (e *RustEngine) RunGraph(
 	callbackID := registerRunGraphCallbackCtx(&runGraphCallbackCtx{exec: exec})
 	defer unregisterRunGraphCallbackCtx(callbackID)
 
-	resp := C.rc_run_graph_json(
+	resp := C.rc_run_graph_json_with_go_callback(
 		e.ptr,
 		centry,
 		cfinish,
 		cinitial,
 		cinitialInput,
 		cstreamMode,
+		clockedFields,
 		C.ulong(callbackID),
-		(C.rc_node_callback_t)(C.goNodeCallback),
 	)
 	defer C.rc_string_free(resp)
 
@@ -298,13 +367,29 @@ func cCallbackEnvelopeError(message string) *C.char {
 	return C.CString(string(raw))
 }
 
-func cCallbackEnvelopeSuspend(cond AnyOfCondition) *C.char {
-	raw, _ := json.Marshal(map[string]any{
-		"ok": true,
-		"suspend": map[string]any{
+func cCallbackEnvelopeSuspend(target WaitTarget) *C.char {
+	if target == nil {
+		return cCallbackEnvelopeError("wait requested with nil target")
+	}
+	kind := target.waitKind()
+	if kind != "any_of" && kind != "all_of" {
+		return cCallbackEnvelopeError(fmt.Sprintf("unsupported wait target kind `%s`", kind))
+	}
+	var payload map[string]any
+	if kind == "any_of" {
+		payload = map[string]any{
 			"kind":   "any_of",
-			"any_of": cond,
-		},
+			"any_of": AnyOfCondition{Conditions: target.waitConditions()},
+		}
+	} else {
+		payload = map[string]any{
+			"kind":   "all_of",
+			"all_of": AllOfCondition{Conditions: target.waitConditions()},
+		}
+	}
+	raw, _ := json.Marshal(map[string]any{
+		"ok":      true,
+		"suspend": payload,
 	})
 	return C.CString(string(raw))
 }

langgraph-go/advancedgraph/types.go

@@ -3,50 +3,142 @@ package advancedgraph
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
+	"reflect"
 )
 
 type WaitCondition interface {
-	toAny() map[string]any
+	json.Marshaler
 }
 
 type ChannelCondition struct {
 	Channel string
-	N       int
+	Min     int
+	Max     int
 }
 
-func (c ChannelCondition) toAny() map[string]any {
-	n := c.N
-	if n <= 0 {
-		n = 1
-	}
-	return map[string]any{
-		"kind":    "channel",
-		"channel": c.Channel,
-		"n":       n,
+type channelConditionJSON struct {
+	Kind    string `json:"kind"`
+	Channel string `json:"channel"`
+	Min     int    `json:"min"`
+	Max     int    `json:"max"`
+}
+
+func (c ChannelCondition) MarshalJSON() ([]byte, error) {
+	min := c.Min
+	if min <= 0 {
+		min = 1
 	}
+	return json.Marshal(channelConditionJSON{
+		Kind:    "channel",
+		Channel: c.Channel,
+		Min:     min,
+		Max:     c.Max,
+	})
 }
 
 type TimerCondition struct {
 	Seconds float64
 }
 
-func (t TimerCondition) toAny() map[string]any {
-	return map[string]any{
-		"kind":    "timer",
-		"seconds": t.Seconds,
-	}
+type timerConditionJSON struct {
+	Kind    string  `json:"kind"`
+	Seconds float64 `json:"seconds"`
+}
+
+func (t TimerCondition) MarshalJSON() ([]byte, error) {
+	return json.Marshal(timerConditionJSON{
+		Kind:    "timer",
+		Seconds: t.Seconds,
+	})
 }
 
 type AnyOfCondition struct {
-	Conditions []map[string]any `json:"conditions"`
+	Conditions []WaitCondition
+}
+
+type AllOfCondition struct {
+	Conditions []WaitCondition
+}
+
+type WaitTarget interface {
+	waitKind() string
+	waitConditions() []WaitCondition
+}
+
+func (a AnyOfCondition) waitKind() string {
+	return "any_of"
+}
+
+func (a AnyOfCondition) waitConditions() []WaitCondition {
+	return a.Conditions
+}
+
+func (a AllOfCondition) waitKind() string {
+	return "all_of"
+}
+
+func (a AllOfCondition) waitConditions() []WaitCondition {
+	return a.Conditions
 }
 
 func AnyOf(conditions ...WaitCondition) AnyOfCondition {
-	result := AnyOfCondition{Conditions: make([]map[string]any, 0, len(conditions))}
-	for _, cond := range conditions {
-		result.Conditions = append(result.Conditions, cond.toAny())
+	return AnyOfCondition{Conditions: append([]WaitCondition{}, conditions...)}
+}
+
+func AllOf(conditions ...WaitCondition) AllOfCondition {
+	return AllOfCondition{Conditions: append([]WaitCondition{}, conditions...)}
+}
+
+func (a AnyOfCondition) MarshalJSON() ([]byte, error) {
+	result := struct {
+		Conditions []json.RawMessage `json:"conditions"`
+	}{
+		Conditions: make([]json.RawMessage, 0, len(a.Conditions)),
+	}
+	for i, cond := range a.Conditions {
+		if isNilWaitCondition(cond) {
+			return nil, fmt.Errorf("any_of condition[%d] is nil", i)
+		}
+		raw, err := cond.MarshalJSON()
+		if err != nil {
+			return nil, fmt.Errorf("encode any_of condition[%d]: %w", i, err)
+		}
+		result.Conditions = append(result.Conditions, json.RawMessage(raw))
+	}
+	return json.Marshal(result)
+}
+
+func (a AllOfCondition) MarshalJSON() ([]byte, error) {
+	result := struct {
+		Conditions []json.RawMessage `json:"conditions"`
+	}{
+		Conditions: make([]json.RawMessage, 0, len(a.Conditions)),
+	}
+	for i, cond := range a.Conditions {
+		if isNilWaitCondition(cond) {
+			return nil, fmt.Errorf("all_of condition[%d] is nil", i)
+		}
+		raw, err := cond.MarshalJSON()
+		if err != nil {
+			return nil, fmt.Errorf("encode all_of condition[%d]: %w", i, err)
+		}
+		result.Conditions = append(result.Conditions, json.RawMessage(raw))
+	}
+	return json.Marshal(result)
+}
+
+func isNilWaitCondition(cond WaitCondition) bool {
+	if cond == nil {
+		return true
+	}
+	v := reflect.ValueOf(cond)
+	switch v.Kind() {
+	case reflect.Ptr, reflect.Interface, reflect.Slice, reflect.Map, reflect.Func:
+		return v.IsNil()
+	default:
+		return false
 	}
-	return result
 }
 
 type WaitEvent struct {
@@ -56,6 +148,16 @@ type WaitEvent struct {
 	Seconds   float64         `json:"seconds,omitempty"`
 }
 
+type ConditionResult struct {
+	Met         bool   `json:"met"`
+	ChannelName string `json:"channel_name,omitempty"`
+	Values      []any  `json:"values,omitempty"`
+}
+
+type WaitForResult struct {
+	Conditions []ConditionResult `json:"conditions"`
+}
+
 type Send struct {
 	Node      any
 	NodeInput any
@@ -67,7 +169,7 @@ type Command struct {
 }
 
 type ErrWaitRequested struct {
-	Condition AnyOfCondition
+	Target WaitTarget
 }
 
 func (e ErrWaitRequested) Error() string {

langgraph-go/stategraph/state_graph.go

@@ -26,25 +26,34 @@ func (c *Context) Interrupt(name string) (any, error) {
 	}
 	event, err := c.inner.WaitFor(ag.AnyOf(ag.ChannelCondition{
 		Channel: internalInterruptChannel,
-		N:       1,
+		Min:     1,
 	}))
 	if err != nil {
 		if waitReq, ok := ag.AsErrWaitRequested(err); ok {
 			return nil, errInterruptRequested{
 				Name:      name,
-				Condition: waitReq.Condition,
+				Condition: waitReq.Target,
 			}
 		}
 		return nil, err
 	}
-	if len(event.Value) == 0 {
+	if len(event.Conditions) == 0 || !event.Conditions[0].Met {
 		return nil, nil
 	}
+	values := event.Conditions[0].Values
+	if len(values) == 0 {
+		return nil, nil
+	}
+	rawValue := values[0]
+	valueBytes, err := json.Marshal(rawValue)
+	if err != nil {
+		return nil, fmt.Errorf("encode interrupt `%s` value: %w", name, err)
+	}
 
 	var payload interruptPayload
-	if err := json.Unmarshal(event.Value, &payload); err != nil {
+	if err := json.Unmarshal(valueBytes, &payload); err != nil {
 		var value any
-		if err := json.Unmarshal(event.Value, &value); err != nil {
+		if err := json.Unmarshal(valueBytes, &value); err != nil {
 			return nil, fmt.Errorf("decode interrupt `%s` value: %w", name, err)
 		}
 		return value, nil
@@ -64,7 +73,7 @@ func (c *Context) Interrupt(name string) (any, error) {
 
 type errInterruptRequested struct {
 	Name      string
-	Condition ag.AnyOfCondition
+	Condition ag.WaitTarget
 }
 
 func (e errInterruptRequested) Error() string {
@@ -169,13 +178,13 @@ func (g *BasicStateGraph[StateT]) Compile() *CompiledBasicStateGraph[StateT] {
 				if err != nil {
 					if interruptReq, ok := asErrInterruptRequested(err); ok {
 						cond := interruptReq.Condition
-						if len(cond.Conditions) == 0 {
+						if cond == nil {
 							cond = ag.AnyOf(ag.ChannelCondition{
 								Channel: internalInterruptChannel,
-								N:       1,
+								Min:     1,
 							})
 						}
-						return ag.Command{}, ag.ErrWaitRequested{Condition: cond}
+						return ag.Command{}, ag.ErrWaitRequested{Target: cond}
 					}
 					return ag.Command{}, err
 				}
@@ -202,7 +211,7 @@ func (g *BasicStateGraph[StateT]) Compile() *CompiledBasicStateGraph[StateT] {
 				needed := len(levels[nextStep-1])
 				_, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{
 					Channel: internalBarrierChannel,
-					N:       needed,
+					Min:     needed,
 				}))
 				if err != nil {
 					return ag.Command{}, err

langgraph-go/tests/sub_agents_test.go

@@ -64,9 +64,9 @@ func (w *lunchWorkflow) llmNode(ctx *ag.Context, _ any, _ lunchState) (ag.Comman
 func (w *lunchWorkflow) waitNode(ctx *ag.Context, _ any, state lunchState) (ag.Command, error) {
 	event, err := ctx.WaitFor(
 		ag.AnyOf(
-			ag.ChannelCondition{Channel: "tool_completion_channel", N: 1},
-			ag.ChannelCondition{Channel: "subagent_completion_channel", N: 1},
-			ag.ChannelCondition{Channel: "user_input_channel", N: 1},
+			ag.ChannelCondition{Channel: "tool_completion_channel"},
+			ag.ChannelCondition{Channel: "subagent_completion_channel"},
+			ag.ChannelCondition{Channel: "user_input_channel"},
 			ag.TimerCondition{Seconds: 1},
 		),
 	)
@@ -75,16 +75,25 @@ func (w *lunchWorkflow) waitNode(ctx *ag.Context, _ any, state lunchState) (ag.C
 	}
 
 	output := append([]string(nil), state.Output...)
-	if event.Condition == "channel" {
-		payload := ag.DecodeString(event.Value)
-		switch event.Channel {
-		case "tool_completion_channel":
-			output = append(output, "tool: "+payload)
-		case "subagent_completion_channel":
-			output = append(output, "sub_agent: "+payload)
-		case "user_input_channel":
-			output = append(output, "user_input: "+payload)
+	hadChannel := false
+	for _, cond := range event.Conditions {
+		if !cond.Met || cond.ChannelName == "" {
+			continue
 		}
+		for _, raw := range cond.Values {
+			payload, _ := raw.(string)
+			switch cond.ChannelName {
+			case "tool_completion_channel":
+				output = append(output, "tool: "+payload)
+			case "subagent_completion_channel":
+				output = append(output, "sub_agent: "+payload)
+			case "user_input_channel":
+				output = append(output, "user_input: "+payload)
+			}
+		}
+		hadChannel = true
+	}
+	if hadChannel {
 		state.Output = output
 		return ag.Command{Goto: []ag.Send{{Node: w.llmNode}}, Update: state}, nil
 	}

langgraph-go/tests/test_primitives_test.go

@@ -2,12 +2,17 @@ package tests
 
 import (
 	"fmt"
+	"sync"
 	"testing"
+	"time"
 
 	ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph"
 )
 
 type primitiveWorkflow struct {
+	dbWriteCount int
+	intervalMu   sync.Mutex
+	intervals    map[string][2]time.Time
 }
 
 type primitiveState struct {
@@ -122,3 +127,463 @@ func TestRunEndsWithoutFinishNode(t *testing.T) {
 		t.Fatalf("unexpected logs: %#v", result.Logs)
 	}
 }
+
+func (w *primitiveWorkflow) startWaitBatchNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	if err := ctx.PublishToChannel("events", "a"); err != nil {
+		return ag.Command{}, err
+	}
+	if err := ctx.PublishToChannel("events", "b"); err != nil {
+		return ag.Command{}, err
+	}
+	if err := ctx.PublishToChannel("events", "c"); err != nil {
+		return ag.Command{}, err
+	}
+	return ag.Command{
+		Update: state,
+		Goto: []ag.Send{
+			{Node: w.waitBatchNode, NodeInput: nil},
+		},
+	}, nil
+}
+
+func (w *primitiveWorkflow) waitBatchNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	result, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{
+		Channel: "events",
+		Min:     2,
+		Max:     4,
+	}))
+	if err != nil {
+		return ag.Command{}, err
+	}
+	if len(result.Conditions) != 1 || !result.Conditions[0].Met {
+		return ag.Command{}, fmt.Errorf("expected one met condition")
+	}
+	values := make([]string, 0, len(result.Conditions[0].Values))
+	for _, v := range result.Conditions[0].Values {
+		s, _ := v.(string)
+		values = append(values, s)
+	}
+	state.Count = len(values)
+	state.Logs = values
+	state.Done = "ok"
+	return ag.Command{Update: state}, nil
+}
+
+func TestChannelWaitRespectsMaxM(t *testing.T) {
+	workflow := &primitiveWorkflow{}
+	graph := ag.NewAdvancedStateGraph[primitiveState]()
+	graph.AddAsyncChannel("events")
+	graph.AddEntryNode(workflow.startWaitBatchNode)
+	graph.AddFinishNode(workflow.waitBatchNode)
+
+	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 result.Count != 3 {
+		t.Fatalf("unexpected count: %v", result.Count)
+	}
+	if len(result.Logs) != 3 || result.Logs[0] != "a" || result.Logs[1] != "b" || result.Logs[2] != "c" {
+		t.Fatalf("unexpected logs: %#v", result.Logs)
+	}
+	if result.Done != "ok" {
+		t.Fatalf("unexpected done: %v", result.Done)
+	}
+}
+
+func (w *primitiveWorkflow) startAnyOfTwoChannelsNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	if err := ctx.PublishToChannel("alpha", "a1"); err != nil {
+		return ag.Command{}, err
+	}
+	if err := ctx.PublishToChannel("beta", "b1"); err != nil {
+		return ag.Command{}, err
+	}
+	return ag.Command{
+		Update: state,
+		Goto: []ag.Send{
+			{Node: w.waitAnyOfTwoChannelsNode, NodeInput: nil},
+		},
+	}, nil
+}
+
+func (w *primitiveWorkflow) waitAnyOfTwoChannelsNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	first, err := ctx.WaitFor(ag.AnyOf(
+		ag.ChannelCondition{Channel: "alpha"},
+		ag.ChannelCondition{Channel: "beta"},
+	))
+	if err != nil {
+		return ag.Command{}, err
+	}
+	if len(first.Conditions) != 2 {
+		return ag.Command{}, fmt.Errorf("expected 2 condition results, got %d", len(first.Conditions))
+	}
+	if !first.Conditions[0].Met || first.Conditions[0].ChannelName != "alpha" || len(first.Conditions[0].Values) != 1 || first.Conditions[0].Values[0] != "a1" {
+		return ag.Command{}, fmt.Errorf("unexpected first condition result: %#v", first.Conditions[0])
+	}
+	if !first.Conditions[1].Met || first.Conditions[1].ChannelName != "beta" || len(first.Conditions[1].Values) != 1 || first.Conditions[1].Values[0] != "b1" {
+		return ag.Command{}, fmt.Errorf("unexpected second condition result: %#v", first.Conditions[1])
+	}
+	if err := ctx.PublishToChannel("beta", "b2"); err != nil {
+		return ag.Command{}, err
+	}
+	state.Count = 1
+	state.Logs = []string{
+		"matched=2",
+	}
+	return ag.Command{
+		Update: state,
+		Goto: []ag.Send{
+			{Node: w.verifyBetaAfterAnyOfNode, NodeInput: nil},
+		},
+	}, nil
+}
+
+func (w *primitiveWorkflow) verifyBetaAfterAnyOfNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	second, err := ctx.WaitFor(ag.AnyOf(
+		ag.ChannelCondition{Channel: "beta"},
+	))
+	if err != nil {
+		return ag.Command{}, err
+	}
+	if len(second.Conditions) != 1 || !second.Conditions[0].Met || second.Conditions[0].ChannelName != "beta" || len(second.Conditions[0].Values) != 1 {
+		return ag.Command{}, fmt.Errorf("unexpected beta condition result: %#v", second.Conditions)
+	}
+	payload, _ := second.Conditions[0].Values[0].(string)
+	state.Count = 2
+	state.Logs = append(state.Logs, fmt.Sprintf("beta=%s", payload))
+	state.Done = "ok"
+	return ag.Command{Update: state}, nil
+}
+
+func TestAnyOfConsumesAllReadyChannels(t *testing.T) {
+	workflow := &primitiveWorkflow{}
+	graph := ag.NewAdvancedStateGraph[primitiveState]()
+	graph.AddAsyncChannel("alpha")
+	graph.AddAsyncChannel("beta")
+	graph.AddEntryNode(workflow.startAnyOfTwoChannelsNode)
+	graph.AddNode(workflow.waitAnyOfTwoChannelsNode)
+	graph.AddFinishNode(workflow.verifyBetaAfterAnyOfNode)
+
+	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 result.Count != 2 {
+		t.Fatalf("unexpected count: %v", result.Count)
+	}
+	if len(result.Logs) != 2 || result.Logs[0] != "matched=2" || result.Logs[1] != "beta=b2" {
+		t.Fatalf("unexpected logs: %#v", result.Logs)
+	}
+	if result.Done != "ok" {
+		t.Fatalf("unexpected done: %v", result.Done)
+	}
+}
+
+func (w *primitiveWorkflow) startAllOfTwoChannelsNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	if err := ctx.PublishToChannel("alpha", "a1"); err != nil {
+		return ag.Command{}, err
+	}
+	return ag.Command{
+		Update: state,
+		Goto: []ag.Send{
+			{Node: w.waitAllOfTwoChannelsNode, NodeInput: nil},
+		},
+	}, nil
+}
+
+func (w *primitiveWorkflow) waitAllOfTwoChannelsNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	result, err := ctx.WaitFor(ag.AllOf(
+		ag.ChannelCondition{Channel: "alpha"},
+		ag.ChannelCondition{Channel: "beta"},
+	))
+	if err != nil {
+		return ag.Command{}, err
+	}
+	if len(result.Conditions) != 2 {
+		return ag.Command{}, fmt.Errorf("expected 2 condition results, got %d", len(result.Conditions))
+	}
+	if !result.Conditions[0].Met || result.Conditions[0].ChannelName != "alpha" || len(result.Conditions[0].Values) != 1 || result.Conditions[0].Values[0] != "a1" {
+		return ag.Command{}, fmt.Errorf("unexpected alpha condition result: %#v", result.Conditions[0])
+	}
+	if !result.Conditions[1].Met || result.Conditions[1].ChannelName != "beta" || len(result.Conditions[1].Values) != 1 || result.Conditions[1].Values[0] != "b1" {
+		return ag.Command{}, fmt.Errorf("unexpected beta condition result: %#v", result.Conditions[1])
+	}
+	state.Count = 2
+	state.Logs = []string{"all_of_channels_ok"}
+	state.Done = "ok"
+	return ag.Command{Update: state}, nil
+}
+
+func TestAllOfWaitsUntilAllChannelsAreReady(t *testing.T) {
+	workflow := &primitiveWorkflow{}
+	graph := ag.NewAdvancedStateGraph[primitiveState]()
+	graph.AddAsyncChannel("alpha")
+	graph.AddAsyncChannel("beta")
+	graph.AddEntryNode(workflow.startAllOfTwoChannelsNode)
+	graph.AddFinishNode(workflow.waitAllOfTwoChannelsNode)
+
+	handler, err := graph.Compile().Start(nil, primitiveState{
+		Count: 0,
+		Logs:  []string{},
+		Done:  "",
+	})
+	if err != nil {
+		t.Fatalf("start failed: %v", err)
+	}
+	if err := handler.PublishToChannel("beta", "b1"); err != nil {
+		t.Fatalf("publish beta failed: %v", err)
+	}
+
+	result, err := handler.WaitForResult()
+	if err != nil {
+		t.Fatalf("result failed: %v", err)
+	}
+	if result.Count != 2 {
+		t.Fatalf("unexpected count: %v", result.Count)
+	}
+	if len(result.Logs) != 1 || result.Logs[0] != "all_of_channels_ok" {
+		t.Fatalf("unexpected logs: %#v", result.Logs)
+	}
+	if result.Done != "ok" {
+		t.Fatalf("unexpected done: %v", result.Done)
+	}
+}
+
+func (w *primitiveWorkflow) startAllOfChannelTimerNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	if err := ctx.PublishToChannel("alpha", "a1"); err != nil {
+		return ag.Command{}, err
+	}
+	return ag.Command{
+		Update: state,
+		Goto: []ag.Send{
+			{Node: w.waitAllOfChannelTimerNode, NodeInput: nil},
+		},
+	}, nil
+}
+
+func (w *primitiveWorkflow) waitAllOfChannelTimerNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	result, err := ctx.WaitFor(ag.AllOf(
+		ag.ChannelCondition{Channel: "alpha"},
+		ag.TimerCondition{Seconds: 0.05},
+	))
+	if err != nil {
+		return ag.Command{}, err
+	}
+	if len(result.Conditions) != 2 {
+		return ag.Command{}, fmt.Errorf("expected 2 condition results, got %d", len(result.Conditions))
+	}
+	if !result.Conditions[0].Met || result.Conditions[0].ChannelName != "alpha" || len(result.Conditions[0].Values) != 1 || result.Conditions[0].Values[0] != "a1" {
+		return ag.Command{}, fmt.Errorf("unexpected channel condition result: %#v", result.Conditions[0])
+	}
+	if !result.Conditions[1].Met {
+		return ag.Command{}, fmt.Errorf("timer condition should be met: %#v", result.Conditions[1])
+	}
+	state.Count = 1
+	state.Logs = []string{"all_of_channel_timer_ok"}
+	state.Done = "ok"
+	return ag.Command{Update: state}, nil
+}
+
+func TestAllOfChannelAndTimerMarksBothConditions(t *testing.T) {
+	workflow := &primitiveWorkflow{}
+	graph := ag.NewAdvancedStateGraph[primitiveState]()
+	graph.AddAsyncChannel("alpha")
+	graph.AddEntryNode(workflow.startAllOfChannelTimerNode)
+	graph.AddFinishNode(workflow.waitAllOfChannelTimerNode)
+
+	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 result.Count != 1 {
+		t.Fatalf("unexpected count: %v", result.Count)
+	}
+	if len(result.Logs) != 1 || result.Logs[0] != "all_of_channel_timer_ok" {
+		t.Fatalf("unexpected logs: %#v", result.Logs)
+	}
+	if result.Done != "ok" {
+		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)
+	}
+}
+
+func (w *primitiveWorkflow) startLockedWorkersNode(_ *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	w.intervalMu.Lock()
+	w.intervals = map[string][2]time.Time{}
+	w.intervalMu.Unlock()
+	return ag.Command{
+		Update: state,
+		Goto: []ag.Send{
+			{Node: w.lockedWorkerANode, NodeInput: nil},
+			{Node: w.lockedWorkerBNode, NodeInput: nil},
+			{Node: w.waitLockedWorkersNode, NodeInput: nil},
+		},
+	}, nil
+}
+
+func (w *primitiveWorkflow) lockedWorkerANode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	start := time.Now()
+	time.Sleep(40 * time.Millisecond)
+	end := time.Now()
+	w.intervalMu.Lock()
+	w.intervals["a"] = [2]time.Time{start, end}
+	w.intervalMu.Unlock()
+	if err := ctx.PublishToChannel("done", "a"); err != nil {
+		return ag.Command{}, err
+	}
+	return ag.Command{Update: state}, nil
+}
+
+func (w *primitiveWorkflow) lockedWorkerBNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	start := time.Now()
+	time.Sleep(40 * time.Millisecond)
+	end := time.Now()
+	w.intervalMu.Lock()
+	w.intervals["b"] = [2]time.Time{start, end}
+	w.intervalMu.Unlock()
+	if err := ctx.PublishToChannel("done", "b"); err != nil {
+		return ag.Command{}, err
+	}
+	return ag.Command{Update: state}, nil
+}
+
+func (w *primitiveWorkflow) waitLockedWorkersNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) {
+	if _, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{Channel: "done", Min: 2})); err != nil {
+		return ag.Command{}, err
+	}
+	return ag.Command{
+		Update: state,
+		Goto:   []ag.Send{{Node: w.finishResumeFlagNode, NodeInput: nil}},
+	}, nil
+}
+
+func TestStateFieldLockingSerializesConflictingNodes(t *testing.T) {
+	workflow := &primitiveWorkflow{}
+	graph := ag.NewAdvancedStateGraph[primitiveState]()
+	graph.AddAsyncChannel("done")
+	graph.AddEntryNode(workflow.startLockedWorkersNode)
+	graph.AddNode(
+		workflow.lockedWorkerANode,
+		ag.NodeStateOption{LockedFields: []string{"counter"}},
+	)
+	graph.AddNode(
+		workflow.lockedWorkerBNode,
+		ag.NodeStateOption{LockedFields: []string{"counter"}},
+	)
+	graph.AddNode(workflow.waitLockedWorkersNode)
+	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)
+	}
+	if _, err := handler.WaitForResult(); err != nil {
+		t.Fatalf("result failed: %v", err)
+	}
+
+	workflow.intervalMu.Lock()
+	ia, okA := workflow.intervals["a"]
+	ib, okB := workflow.intervals["b"]
+	workflow.intervalMu.Unlock()
+	if !okA || !okB {
+		t.Fatalf("missing worker intervals: %#v", workflow.intervals)
+	}
+	serialized := !ia[1].After(ib[0]) || !ib[1].After(ia[0])
+	if !serialized {
+		t.Fatalf("expected serialized execution, got overlap: a=%v..%v b=%v..%v", ia[0], ia[1], ib[0], ib[1])
+	}
+}

langgraph-go/tests/test_streaming_test.go

@@ -15,11 +15,11 @@ type streamState struct {
 type streamWorkflow struct{}
 
 func (w *streamWorkflow) startNode(ctx *ag.Context, _ any, state streamState) (ag.Command, error) {
-	if err := ctx.SendCustomStreamEvent(map[string]any{"step": "start", "value": 1}); err != nil {
+	if err := ctx.SendCustomStreamEvent("high", map[string]any{"step": "start", "value": 1}); err != nil {
 		return ag.Command{}, err
 	}
 	time.Sleep(80 * time.Millisecond)
-	if err := ctx.SendCustomStreamEvent(map[string]any{"step": "start", "value": 2}); err != nil {
+	if err := ctx.SendCustomStreamEvent("regular", map[string]any{"step": "start", "value": 2}); err != nil {
 		return ag.Command{}, err
 	}
 	return ag.Command{
@@ -38,6 +38,8 @@ func (w *streamWorkflow) finishNode(ctx *ag.Context, _ any, state streamState) (
 func TestCustomStreamReceiveAndClose(t *testing.T) {
 	workflow := &streamWorkflow{}
 	graph := ag.NewAdvancedStateGraph[streamState]()
+	graph.AddCustomOutputStream("high")
+	graph.AddCustomOutputStream("regular")
 	graph.AddEntryNode(workflow.startNode)
 	graph.AddFinishNode(workflow.finishNode)
 
@@ -46,7 +48,7 @@ func TestCustomStreamReceiveAndClose(t *testing.T) {
 		t.Fatalf("start failed: %v", err)
 	}
 
-	event, err := handler.ReceiveStream()
+	event, err := handler.ReceiveStream("high")
 	if err != nil {
 		t.Fatalf("receive stream failed: %v", err)
 	}
@@ -61,11 +63,23 @@ func TestCustomStreamReceiveAndClose(t *testing.T) {
 		t.Fatalf("unexpected stream event payload: %#v", eventMap)
 	}
 
-	if err := handler.CloseStream(); err != nil {
-		t.Fatalf("close stream failed: %v", err)
+	eventRegular, err := handler.ReceiveStream("regular")
+	if err != nil {
+		t.Fatalf("receive regular stream failed: %v", err)
+	}
+	eventRegularMap, ok := eventRegular.(map[string]any)
+	if !ok {
+		t.Fatalf("unexpected regular event type: %T", eventRegular)
+	}
+	if eventRegularMap["value"] != float64(2) {
+		t.Fatalf("unexpected regular stream payload: %#v", eventRegularMap)
 	}
 
-	closedEvent, err := handler.ReceiveStream()
+	if err := handler.CloseAllStreams(); err != nil {
+		t.Fatalf("close all streams failed: %v", err)
+	}
+
+	closedEvent, err := handler.ReceiveStream("high")
 	if err != nil {
 		t.Fatalf("receive stream after close failed: %v", err)
 	}
@@ -85,6 +99,7 @@ func TestCustomStreamReceiveAndClose(t *testing.T) {
 func TestOnlyCustomStreamModeSupported(t *testing.T) {
 	workflow := &streamWorkflow{}
 	graph := ag.NewAdvancedStateGraph[streamState]()
+	graph.AddCustomOutputStream("regular")
 	graph.AddEntryNode(workflow.startNode)
 	graph.AddFinishNode(workflow.finishNode)
 

rust-core/Cargo.lock

@@ -2,6 +2,12 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "autocfg"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
+
 [[package]]
 name = "bitflags"
 version = "2.11.0"
@@ -14,6 +20,21 @@ version = "1.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
 
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "indoc"
+version = "2.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706"
+dependencies = [
+ "rustversion",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.17"
@@ -26,6 +47,7 @@ version = "0.1.0"
 dependencies = [
  "libc",
  "parking_lot",
+ "pyo3",
  "serde",
  "serde_json",
  "tokio",
@@ -52,6 +74,21 @@ version = "2.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
 
+[[package]]
+name = "memoffset"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "once_cell"
+version = "1.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
+
 [[package]]
 name = "parking_lot"
 version = "0.12.5"
@@ -81,6 +118,12 @@ version = "0.2.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
 
+[[package]]
+name = "portable-atomic"
+version = "1.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.106"
@@ -90,6 +133,69 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "pyo3"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872"
+dependencies = [
+ "cfg-if",
+ "indoc",
+ "libc",
+ "memoffset",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+ "unindent",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb"
+dependencies = [
+ "once_cell",
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "pyo3-build-config",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "quote"
 version = "1.0.45"
@@ -108,6 +214,12 @@ dependencies = [
  "bitflags",
 ]
 
+[[package]]
+name = "rustversion"
+version = "1.0.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
+
 [[package]]
 name = "scopeguard"
 version = "1.2.0"
@@ -174,6 +286,12 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "target-lexicon"
+version = "0.12.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
+
 [[package]]
 name = "tokio"
 version = "1.50.0"
@@ -201,6 +319,12 @@ version = "1.0.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
 
+[[package]]
+name = "unindent"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
+
 [[package]]
 name = "windows-link"
 version = "0.2.1"

rust-core/Cargo.toml

@@ -7,7 +7,12 @@ edition = "2021"
 name = "langgraph_rust_core"
 crate-type = ["cdylib", "rlib"]
 
+[features]
+default = ["python-bindings"]
+python-bindings = ["dep:pyo3"]
+
 [dependencies]
+pyo3 = { version = "0.23.5", features = ["extension-module"], optional = true }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 parking_lot = "0.12"

rust-core/include/langgraph_rust_core.h

@@ -17,12 +17,14 @@ Engine* rc_engine_new(void);
 void rc_engine_free(Engine* ptr);
 
 char* rc_add_async_channel(Engine* ptr, const char* channel);
+char* rc_add_custom_output_stream(Engine* ptr, const char* stream_name);
 char* rc_publish_json(Engine* ptr, const char* channel, const char* value_json);
 char* rc_wait_any_of_json(Engine* ptr, const char* any_of_json);
+char* rc_wait_all_of_json(Engine* ptr, const char* all_of_json);
 char* rc_start_stream(Engine* ptr, const char* stream_mode);
-char* rc_receive_stream_json(Engine* ptr);
-char* rc_send_custom_stream_event(Engine* ptr, const char* value_json);
-char* rc_close_stream(Engine* ptr);
+char* rc_receive_stream_json(Engine* ptr, const char* stream_name);
+char* rc_send_custom_stream_event(Engine* ptr, const char* stream_name, const char* value_json);
+char* rc_close_all_streams(Engine* ptr);
 char* rc_run_graph_json(
     Engine* ptr,
     const char* entry_point,
@@ -30,6 +32,7 @@ char* rc_run_graph_json(
     const char* initial_state_json,
     const char* initial_input_json,
     const char* stream_mode,
+    const char* node_locked_fields_json,
     unsigned long user_data,
     rc_node_callback_t callback
 );

rust-core/src/lib.rs

@@ -1,2 +1,4 @@
 mod engine;
 mod lib_c;
+#[cfg(feature = "python-bindings")]
+mod lib_py;

rust-core/src/lib_c.rs

@@ -1,8 +1,9 @@
 use crate::engine::{
     parse_callback_envelope_json, run_graph_json_with_callback, run_loop_block_on, run_loop_spawn,
-    AnyOfCondition, Engine, NodeOutcome,
+    AllOfCondition, AnyOfCondition, Engine, NodeOutcome,
 };
 use serde_json::Value;
+use std::collections::HashMap;
 use std::ffi::{CStr, CString};
 use std::os::raw::c_char;
 use std::sync::mpsc;
@@ -144,6 +145,60 @@ pub unsafe extern "C" fn rc_wait_any_of_json(
     }
 }
 
+#[no_mangle]
+/// # Safety
+/// `ptr` must be a valid engine pointer from `rc_engine_new`.
+/// `all_of_json` must be a valid null-terminated UTF-8 string pointer.
+pub unsafe extern "C" fn rc_wait_all_of_json(
+    ptr: *mut Engine,
+    all_of_json: *const c_char,
+) -> *mut c_char {
+    if ptr.is_null() {
+        return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
+    }
+    let all_of_json = match cstr_to_str(all_of_json) {
+        Ok(v) => v,
+        Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
+    };
+    let all_of: AllOfCondition = match serde_json::from_str(all_of_json) {
+        Ok(v) => v,
+        Err(e) => {
+            return into_c_ptr(format!(
+                "{{\"ok\":false,\"error\":\"invalid all_of JSON: {e}\"}}"
+            ))
+        }
+    };
+    let result = run_loop_block_on((*ptr).wait_for_all_of_async(&all_of));
+    match result {
+        Ok(event) => match serde_json::to_string(&event) {
+            Ok(s) => into_c_ptr(format!("{{\"ok\":true,\"event\":{s}}}")),
+            Err(e) => into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
+        },
+        Err(e) => into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
+    }
+}
+
+#[no_mangle]
+/// # Safety
+/// `ptr` must be a valid engine pointer from `rc_engine_new`.
+/// `stream_name` must be a valid null-terminated UTF-8 string pointer.
+pub unsafe extern "C" fn rc_add_custom_output_stream(
+    ptr: *mut Engine,
+    stream_name: *const c_char,
+) -> *mut c_char {
+    if ptr.is_null() {
+        return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
+    }
+    let stream_name = match cstr_to_str(stream_name) {
+        Ok(v) => v,
+        Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
+    };
+    match (*ptr).add_custom_output_stream(stream_name) {
+        Ok(()) => into_c_ptr("{\"ok\":true}".to_string()),
+        Err(e) => into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
+    }
+}
+
 #[no_mangle]
 /// # Safety
 /// `ptr` must be a valid engine pointer from `rc_engine_new`.
@@ -172,11 +227,19 @@ pub unsafe extern "C" fn rc_start_stream(
 #[no_mangle]
 /// # Safety
 /// `ptr` must be a valid engine pointer from `rc_engine_new`.
-pub unsafe extern "C" fn rc_receive_stream_json(ptr: *mut Engine) -> *mut c_char {
+/// `stream_name` must be a valid null-terminated UTF-8 string pointer.
+pub unsafe extern "C" fn rc_receive_stream_json(
+    ptr: *mut Engine,
+    stream_name: *const c_char,
+) -> *mut c_char {
     if ptr.is_null() {
         return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
     }
-    let event = run_loop_block_on((*ptr).receive_stream_async());
+    let stream_name = match cstr_to_str(stream_name) {
+        Ok(v) => v,
+        Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
+    };
+    let event = run_loop_block_on((*ptr).receive_stream_async(stream_name));
     match event {
         Some(value) => match serde_json::to_string(&value) {
             Ok(s) => into_c_ptr(format!("{{\"ok\":true,\"has_event\":true,\"event\":{s}}}")),
@@ -189,14 +252,20 @@ pub unsafe extern "C" fn rc_receive_stream_json(ptr: *mut Engine) -> *mut c_char
 #[no_mangle]
 /// # Safety
 /// `ptr` must be a valid engine pointer from `rc_engine_new`.
+/// `stream_name` must be a valid null-terminated UTF-8 string pointer.
 /// `value_json` must be a valid null-terminated UTF-8 string pointer.
 pub unsafe extern "C" fn rc_send_custom_stream_event(
     ptr: *mut Engine,
+    stream_name: *const c_char,
     value_json: *const c_char,
 ) -> *mut c_char {
     if ptr.is_null() {
         return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
     }
+    let stream_name = match cstr_to_str(stream_name) {
+        Ok(v) => v,
+        Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
+    };
     let value_json = match cstr_to_str(value_json) {
         Ok(v) => v,
         Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
@@ -209,18 +278,18 @@ pub unsafe extern "C" fn rc_send_custom_stream_event(
             ))
         }
     };
-    (*ptr).send_custom_stream_event(value);
+    (*ptr).send_custom_stream_event(stream_name, value);
     into_c_ptr("{\"ok\":true}".to_string())
 }
 
 #[no_mangle]
 /// # Safety
 /// `ptr` must be a valid engine pointer from `rc_engine_new`.
-pub unsafe extern "C" fn rc_close_stream(ptr: *mut Engine) -> *mut c_char {
+pub unsafe extern "C" fn rc_close_all_streams(ptr: *mut Engine) -> *mut c_char {
     if ptr.is_null() {
         return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
     }
-    (*ptr).close_stream();
+    (*ptr).close_all_streams();
     into_c_ptr("{\"ok\":true}".to_string())
 }
 
@@ -236,6 +305,7 @@ pub unsafe extern "C" fn rc_run_graph_json(
     initial_state_json: *const c_char,
     initial_input_json: *const c_char,
     stream_mode: *const c_char,
+    node_locked_fields_json: *const c_char,
     user_data: libc::c_ulong,
     callback: Option<CNodeCallback>,
 ) -> *mut c_char {
@@ -285,9 +355,27 @@ pub unsafe extern "C" fn rc_run_graph_json(
             Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
         }
     };
+    let node_locked_fields: HashMap<String, Vec<String>> = if node_locked_fields_json.is_null() {
+        HashMap::new()
+    } else {
+        let payload = match cstr_to_str(node_locked_fields_json) {
+            Ok(v) => v,
+            Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
+        };
+        match serde_json::from_str(payload) {
+            Ok(v) => v,
+            Err(e) => {
+                return into_c_ptr(format!(
+                    "{{\"ok\":false,\"error\":\"invalid node_locked_fields JSON: {e}\"}}"
+                ))
+            }
+        }
+    };
 
-    if let Err(e) = (*ptr).start_stream(stream_mode.as_deref()) {
-        return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}"));
+    if let Some(mode) = stream_mode.as_deref() {
+        if let Err(e) = (*ptr).start_stream(Some(mode)) {
+            return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}"));
+        }
     }
 
     let (tx, rx) = mpsc::channel::<Result<Value, String>>();
@@ -334,9 +422,10 @@ pub unsafe extern "C" fn rc_run_graph_json(
             initial_input,
             run_engine.clone(),
             callback_wrapper,
+            node_locked_fields,
         )
         .await;
-        run_engine.close_stream();
+        run_engine.close_all_streams();
         let _ = tx.send(out);
     });
     if let Err(e) = submit {

rust-core/src/lib_py.rs

@@ -0,0 +1,394 @@
+#[cfg(feature = "python-bindings")]
+use crate::engine::{
+    run_graph_with_callback, run_loop_block_on, AllOfCondition, AnyOfCondition, Engine,
+    NodeExecResult, NodeOutcome, SendPayload, WaitCondition, WaitEvent, WaitRequest,
+};
+#[cfg(feature = "python-bindings")]
+use pyo3::exceptions::PyValueError;
+#[cfg(feature = "python-bindings")]
+use pyo3::prelude::*;
+#[cfg(feature = "python-bindings")]
+use pyo3::types::PyAny;
+#[cfg(feature = "python-bindings")]
+use pyo3::types::{PyDict, PyList, PyTuple};
+#[cfg(feature = "python-bindings")]
+use serde_json::Value;
+#[cfg(feature = "python-bindings")]
+use std::collections::HashMap;
+#[cfg(feature = "python-bindings")]
+use std::sync::Arc;
+
+#[cfg(feature = "python-bindings")]
+#[pyclass]
+struct PyRustEngine {
+    inner: Engine,
+}
+
+#[cfg(feature = "python-bindings")]
+#[pymethods]
+impl PyRustEngine {
+    #[new]
+    fn new() -> Self {
+        Self {
+            inner: Engine::new(),
+        }
+    }
+
+    fn add_async_channel(&self, name: &str) {
+        self.inner.add_async_channel(name);
+    }
+
+    fn add_custom_output_stream(&self, stream_name: &str) -> PyResult<()> {
+        self.inner
+            .add_custom_output_stream(stream_name)
+            .map_err(PyValueError::new_err)
+    }
+
+    fn publish_obj(&self, py: Python<'_>, channel: &str, value: Py<PyAny>) -> PyResult<()> {
+        let value_json = py_obj_to_json_string(py, &value.bind(py))?;
+        let parsed: Value = serde_json::from_str(&value_json)
+            .map_err(|e| PyValueError::new_err(format!("Invalid Python JSON value: {e}")))?;
+        self.inner
+            .publish_json(channel, parsed)
+            .map_err(PyValueError::new_err)
+    }
+
+    #[pyo3(signature = (channel, min, max=None))]
+    fn wait_channel(
+        &self,
+        py: Python<'_>,
+        channel: &str,
+        min: usize,
+        max: Option<usize>,
+    ) -> PyResult<Py<PyAny>> {
+        let cond = WaitCondition::Channel {
+            channel: channel.to_string(),
+            min,
+            max: max.unwrap_or(0),
+        };
+        let any_of = AnyOfCondition {
+            conditions: vec![cond],
+        };
+        let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of))
+            .map_err(PyValueError::new_err)?;
+        let event_json = serde_json::to_string(&event)
+            .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?;
+        json_string_to_py_obj(py, &event_json)
+    }
+
+    fn wait_timer(&self, py: Python<'_>, seconds: f64) -> PyResult<Py<PyAny>> {
+        let cond = WaitCondition::Timer { seconds };
+        let any_of = AnyOfCondition {
+            conditions: vec![cond],
+        };
+        let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of))
+            .map_err(PyValueError::new_err)?;
+        let event_json = serde_json::to_string(&event)
+            .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?;
+        json_string_to_py_obj(py, &event_json)
+    }
+
+    fn wait_any_of_obj(&self, py: Python<'_>, any_of_payload: Py<PyAny>) -> PyResult<Py<PyAny>> {
+        let payload_json = py_obj_to_json_string(py, &any_of_payload.bind(py))?;
+        let event_json = self._wait_any_of_json(&payload_json)?;
+        json_string_to_py_obj(py, &event_json)
+    }
+
+    fn wait_all_of_obj(&self, py: Python<'_>, all_of_payload: Py<PyAny>) -> PyResult<Py<PyAny>> {
+        let payload_json = py_obj_to_json_string(py, &all_of_payload.bind(py))?;
+        let event_json = self._wait_all_of_json(&payload_json)?;
+        json_string_to_py_obj(py, &event_json)
+    }
+
+    fn _wait_any_of_json(&self, any_of_json: &str) -> PyResult<String> {
+        let any_of: AnyOfCondition = serde_json::from_str(any_of_json)
+            .map_err(|e| PyValueError::new_err(format!("Invalid any_of JSON: {e}")))?;
+        let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of))
+            .map_err(PyValueError::new_err)?;
+        serde_json::to_string(&event)
+            .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))
+    }
+
+    fn _wait_all_of_json(&self, all_of_json: &str) -> PyResult<String> {
+        let all_of: AllOfCondition = serde_json::from_str(all_of_json)
+            .map_err(|e| PyValueError::new_err(format!("Invalid all_of JSON: {e}")))?;
+        let event = run_loop_block_on(self.inner.wait_for_all_of_async(&all_of))
+            .map_err(PyValueError::new_err)?;
+        serde_json::to_string(&event)
+            .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))
+    }
+
+    fn wait_condition_json(&self, cond_json: &str) -> PyResult<String> {
+        let cond: WaitCondition = serde_json::from_str(cond_json)
+            .map_err(|e| PyValueError::new_err(format!("Invalid condition JSON: {e}")))?;
+        let any_of = AnyOfCondition {
+            conditions: vec![cond],
+        };
+        let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of))
+            .map_err(PyValueError::new_err)?;
+        serde_json::to_string(&event)
+            .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))
+    }
+
+    #[pyo3(signature = (stream_mode=None))]
+    fn start_stream(&self, stream_mode: Option<&str>) -> PyResult<()> {
+        self.inner
+            .start_stream(stream_mode)
+            .map_err(PyValueError::new_err)
+    }
+
+    fn receive_stream_obj(&self, py: Python<'_>, stream_name: &str) -> PyResult<Py<PyAny>> {
+        let event =
+            py.allow_threads(|| run_loop_block_on(self.inner.receive_stream_async(stream_name)));
+        match event {
+            Some(value) => {
+                let event_json = serde_json::to_string(&value)
+                    .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?;
+                json_string_to_py_obj(py, &event_json)
+            }
+            None => Ok(py.None()),
+        }
+    }
+
+    fn send_custom_stream_event_obj(
+        &self,
+        py: Python<'_>,
+        stream_name: &str,
+        value: Py<PyAny>,
+    ) -> PyResult<()> {
+        let value_json = py_obj_to_json_string(py, &value.bind(py))?;
+        let parsed: Value = serde_json::from_str(&value_json)
+            .map_err(|e| PyValueError::new_err(format!("Invalid Python JSON value: {e}")))?;
+        self.inner.send_custom_stream_event(stream_name, parsed);
+        Ok(())
+    }
+
+    fn close_all_streams(&self) {
+        self.inner.close_all_streams();
+    }
+
+    #[pyo3(signature = (entry_point, finish_point, initial_state, callback, stream_mode=None, node_locked_fields=None))]
+    fn run_graph_py(
+        &self,
+        py: Python<'_>,
+        entry_point: &str,
+        finish_point: &str,
+        initial_state: Py<PyAny>,
+        callback: Py<PyAny>,
+        stream_mode: Option<&str>,
+        node_locked_fields: Option<Py<PyAny>>,
+    ) -> PyResult<Py<PyAny>> {
+        if let Some(mode) = stream_mode {
+            self.inner
+                .start_stream(Some(mode))
+                .map_err(PyValueError::new_err)?;
+        }
+
+        let engine = self.inner.clone();
+        let callback = Arc::new(callback);
+        let entry_point = entry_point.to_string();
+        let finish_point = finish_point.to_string();
+        let initial_state = Arc::new(initial_state);
+        let initial_input = Arc::clone(&initial_state);
+        let locked_fields_by_node: HashMap<String, Vec<String>> =
+            if let Some(locked_fields_obj) = node_locked_fields {
+                let json_str = py_obj_to_json_string(py, &locked_fields_obj.bind(py))?;
+                serde_json::from_str(&json_str).map_err(|e| {
+                    PyValueError::new_err(format!("invalid node_locked_fields payload: {e}"))
+                })?
+            } else {
+                HashMap::new()
+            };
+
+        let run_result =
+            py.allow_threads(move || {
+                run_loop_block_on(run_graph_with_callback(
+                    entry_point,
+                    finish_point,
+                    initial_state,
+                    initial_input,
+                    engine.clone(),
+                    {
+                        let callback = Arc::clone(&callback);
+                    move |node: String,
+                          arg: Arc<Py<PyAny>>,
+                          state_snapshot: Arc<Py<PyAny>>|
+                          -> Result<NodeOutcome<Arc<Py<PyAny>>, Arc<Py<PyAny>>>, String> {
+                        Python::with_gil(|py| -> Result<NodeOutcome<Arc<Py<PyAny>>, Arc<Py<PyAny>>>, String> {
+                            let callback_bound = callback.as_ref().bind(py);
+                            let payload_obj = callback_bound
+                                .call1((
+                                    node.as_str(),
+                                    arg.as_ref().clone_ref(py),
+                                    state_snapshot.as_ref().clone_ref(py),
+                                ))
+                                .map_err(|e| format!("callback failed for node `{node}`: {e}"))?;
+                            parse_node_outcome_arc(py, &payload_obj).map_err(|e| {
+                                format!("invalid callback payload for `{node}`: {e}")
+                            })
+                        })
+                    }
+                    },
+                |state: &mut Arc<Py<PyAny>>, update: Option<Arc<Py<PyAny>>>| -> Result<(), String> {
+                        Python::with_gil(|py| -> Result<(), String> {
+                            if let Some(update) = update {
+                            apply_update_to_state(py, state.as_ref(), update.as_ref())
+                                    .map_err(|e| format!("state merge failed: {e}"))?;
+                            }
+                            Ok(())
+                        })
+                    },
+                |arg: Arc<Py<PyAny>>, event: WaitEvent| -> Result<Arc<Py<PyAny>>, String> {
+                    wrap_resume_arg(arg.as_ref(), &event).map(Arc::new)
+                    },
+                move |node: &str| locked_fields_by_node.get(node).cloned().unwrap_or_default(),
+                ))
+            });
+
+        self.inner.close_all_streams();
+        let out = run_result.map_err(PyValueError::new_err)?;
+        Ok(out.as_ref().clone_ref(py))
+    }
+}
+
+#[cfg(feature = "python-bindings")]
+fn parse_node_outcome_arc(
+    py: Python<'_>,
+    payload_obj: &Bound<'_, PyAny>,
+) -> Result<NodeOutcome<Arc<Py<PyAny>>, Arc<Py<PyAny>>>, String> {
+    let payload_dict = payload_obj
+        .downcast::<PyDict>()
+        .map_err(|_| "payload must be a dict".to_string())?;
+    let suspended_item = payload_dict
+        .get_item("suspend")
+        .map_err(|e| format!("failed to read suspend: {e}"))?;
+    if let Some(wait_obj) = suspended_item {
+        let wait_json = py_obj_to_json_string(py, &wait_obj)
+            .map_err(|e| format!("failed to encode suspend payload: {e}"))?;
+        let wait: WaitRequest = serde_json::from_str(&wait_json)
+            .map_err(|e| format!("invalid suspend payload: {e}"))?;
+        return Ok(NodeOutcome::Suspended { wait });
+    }
+
+    let update_item = payload_dict
+        .get_item("update")
+        .map_err(|e| format!("failed to read update: {e}"))?;
+    let update = match update_item {
+        Some(v) if !v.is_none() => Some(Arc::new(v.unbind())),
+        _ => None,
+    };
+
+    let sends_obj = payload_dict
+        .get_item("sends")
+        .map_err(|e| format!("failed to read sends: {e}"))?
+        .ok_or_else(|| "missing sends".to_string())?;
+    let sends_list = sends_obj
+        .downcast::<PyList>()
+        .map_err(|_| "sends must be a list".to_string())?;
+
+    let mut sends = Vec::with_capacity(sends_list.len());
+    for item in sends_list.iter() {
+        let send_dict = item
+            .downcast::<PyDict>()
+            .map_err(|_| "send item must be a dict".to_string())?;
+        let node_obj = send_dict
+            .get_item("node")
+            .map_err(|e| format!("failed to read send.node: {e}"))?
+            .ok_or_else(|| "send.node is required".to_string())?;
+        let node = node_obj
+            .extract::<String>()
+            .map_err(|e| format!("send.node must be string: {e}"))?;
+        let arg: Arc<Py<PyAny>> = match send_dict.get_item("arg") {
+            Ok(Some(v)) => Arc::new(v.unbind()),
+            Ok(None) => Arc::new(Python::with_gil(|py| py.None())),
+            Err(e) => return Err(format!("failed to read send.arg: {e}")),
+        };
+        sends.push(SendPayload { node, arg });
+    }
+
+    Ok(NodeOutcome::Completed(NodeExecResult { update, sends }))
+}
+
+#[cfg(feature = "python-bindings")]
+fn wrap_resume_arg(arg: &Py<PyAny>, event: &WaitEvent) -> Result<Py<PyAny>, String> {
+    Python::with_gil(|py| -> Result<Py<PyAny>, String> {
+        let wrapper = PyDict::new(py);
+        wrapper
+            .set_item("__lg_resume_arg__", arg.clone_ref(py))
+            .map_err(|e| format!("failed to set resume arg: {e}"))?;
+        let event_json = serde_json::to_string(event)
+            .map_err(|e| format!("failed to encode wait event: {e}"))?;
+        let event_obj = json_string_to_py_obj(py, &event_json)
+            .map_err(|e| format!("failed to parse event: {e}"))?;
+        wrapper
+            .set_item("__lg_resume_event__", event_obj.bind(py))
+            .map_err(|e| format!("failed to set resume event: {e}"))?;
+        Ok(wrapper.unbind().into_any())
+    })
+}
+
+#[cfg(feature = "python-bindings")]
+fn apply_update_to_state(py: Python<'_>, state: &Py<PyAny>, update: &Py<PyAny>) -> PyResult<()> {
+    let state_obj = state.bind(py);
+    let update_obj = update.bind(py);
+
+    if update_obj.is_none() {
+        return Ok(());
+    }
+    if state_obj.is_instance_of::<PyDict>() && update_obj.is_instance_of::<PyDict>() {
+        let state_dict = state_obj.downcast::<PyDict>()?;
+        let update_dict = update_obj.downcast::<PyDict>()?;
+        state_dict.call_method1("update", (update_dict,))?;
+        return Ok(());
+    }
+    if let Ok(tuple_like) = update_obj.downcast::<PyList>() {
+        apply_pair_updates(state_obj, tuple_like)?;
+        return Ok(());
+    }
+    if let Ok(tuple_like) = update_obj.downcast::<PyTuple>() {
+        let list = PyList::new(py, tuple_like)?;
+        apply_pair_updates(state_obj, &list)?;
+    }
+    Ok(())
+}
+
+#[cfg(feature = "python-bindings")]
+fn apply_pair_updates(state_obj: &Bound<'_, PyAny>, entries: &Bound<'_, PyList>) -> PyResult<()> {
+    if !state_obj.is_instance_of::<PyDict>() {
+        return Ok(());
+    }
+    let state_dict = state_obj.downcast::<PyDict>()?;
+    for entry in entries.iter() {
+        if let Ok(pair) = entry.downcast::<PyTuple>() {
+            if pair.len() == 2 {
+                let key_obj = pair.get_item(0)?;
+                if let Ok(key) = key_obj.extract::<String>() {
+                    let value_obj = pair.get_item(1)?;
+                    state_dict.set_item(key, value_obj)?;
+                }
+            }
+        }
+    }
+    Ok(())
+}
+
+#[cfg(feature = "python-bindings")]
+fn py_obj_to_json_string(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult<String> {
+    let json_mod = py.import("json")?;
+    let dumped = json_mod.call_method1("dumps", (obj,))?;
+    dumped.extract::<String>()
+}
+
+#[cfg(feature = "python-bindings")]
+fn json_string_to_py_obj(py: Python<'_>, value: &str) -> PyResult<Py<PyAny>> {
+    let json_mod = py.import("json")?;
+    let loaded = json_mod.call_method1("loads", (value,))?;
+    Ok(loaded.unbind())
+}
+
+#[cfg(feature = "python-bindings")]
+#[pymodule]
+fn langgraph_rust_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<PyRustEngine>()?;
+    Ok(())
+}

saf-python-sdk/Cargo.lock

@@ -2,6 +2,12 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "autocfg"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
+
 [[package]]
 name = "bitflags"
 version = "2.11.0"
@@ -14,6 +20,21 @@ version = "1.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
 
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "indoc"
+version = "2.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706"
+dependencies = [
+ "rustversion",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.17"
@@ -26,6 +47,7 @@ version = "0.1.0"
 dependencies = [
  "libc",
  "parking_lot",
+ "pyo3",
  "serde",
  "serde_json",
  "tokio",
@@ -52,6 +74,21 @@ version = "2.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
 
+[[package]]
+name = "memoffset"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "once_cell"
+version = "1.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
+
 [[package]]
 name = "parking_lot"
 version = "0.12.5"
@@ -81,6 +118,12 @@ version = "0.2.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
 
+[[package]]
+name = "portable-atomic"
+version = "1.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.106"
@@ -90,6 +133,69 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "pyo3"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872"
+dependencies = [
+ "cfg-if",
+ "indoc",
+ "libc",
+ "memoffset",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+ "unindent",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb"
+dependencies = [
+ "once_cell",
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "pyo3-build-config",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "quote"
 version = "1.0.45"
@@ -108,6 +214,12 @@ dependencies = [
  "bitflags",
 ]
 
+[[package]]
+name = "rustversion"
+version = "1.0.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
+
 [[package]]
 name = "scopeguard"
 version = "1.2.0"
@@ -174,6 +286,12 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "target-lexicon"
+version = "0.12.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
+
 [[package]]
 name = "tokio"
 version = "1.50.0"
@@ -201,6 +319,12 @@ version = "1.0.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
 
+[[package]]
+name = "unindent"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
+
 [[package]]
 name = "windows-link"
 version = "0.2.1"

saf-python-sdk/Cargo.toml

@@ -8,7 +8,12 @@ name = "langgraph_rust_core"
 path = "../rust-core/src/lib.rs"
 crate-type = ["cdylib", "rlib"]
 
+[features]
+default = ["python-bindings"]
+python-bindings = ["dep:pyo3"]
+
 [dependencies]
+pyo3 = { version = "0.23.5", features = ["extension-module"], optional = true }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 parking_lot = "0.12"

saf-python-sdk/Makefile

@@ -2,14 +2,15 @@
 PYPI_REPOSITORY ?= pypi
 PYPI_TOKEN ?=
 PYTHON_VERSION ?= 3.13
-TEST_BASE = uv run --python $(PYTHON_VERSION) --with pytest --with anyio --with typing_extensions --with pydantic --with langgraph pytest -q -s
+TEST_BASE = uv run --python $(PYTHON_VERSION) --reinstall-package saf-python-sdk --with pytest --with anyio --with typing_extensions --with pydantic --with langgraph pytest -q -s
 
 .PHONY: publish-to-pypi-saf-python-sdk all-tests test_primitives test_sub_agents test_update_elision test_run_pool_size test-benchmark
 publish-to-pypi-saf-python-sdk:
 	@test -n "$(PYPI_TOKEN)" || (echo "PYPI_TOKEN is required"; exit 1)
 	cd "$(CURDIR)" && \
+	PATH="$(HOME)/.cargo/bin:$$PATH" \
 	MATURIN_PYPI_TOKEN="$(PYPI_TOKEN)" \
-	uvx maturin publish --repository $(PYPI_REPOSITORY) --non-interactive --skip-existing --no-sdist
+	uv run --with maturin maturin publish --repository $(PYPI_REPOSITORY) --non-interactive --skip-existing --no-sdist
 
 all-tests:
 	$(TEST_BASE) tests/advanced-graph/test_*.py

saf-python-sdk/README.md

@@ -5,7 +5,7 @@ Standalone Python SDK for the `advanced_graph` runtime backed by the Rust engine
 This package intentionally contains only:
 
 - `saf_python_sdk.advanced_graph` (Python API)
-- Rust core engine via C bindings (`ctypes`)
+- `langgraph_rust_core` (Rust execution engine via PyO3)
 
 It does not package the original `langgraph` `stategraph` stack.
 

saf-python-sdk/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "saf-python-sdk"
-version = "0.1.2"
+version = "0.1.4"
 description = "Standalone advanced graph runtime powered by Rust engine"
 readme = "README.md"
 requires-python = ">=3.10"

saf-python-sdk/python/saf_python_sdk/advanced_graph/__init__.py

@@ -1,11 +1,16 @@
 from .state import (
     AdvancedStateGraph,
+    AllOfCondition,
     AnyOfCondition,
     ChannelCondition,
+    ConditionResult,
     CompiledGraphEngine,
     Context,
     GraphRunHandler,
+    NodeStateOption,
     TimerCondition,
+    WaitForResult,
+    all_of,
     any_of,
     channel_condition,
     timer_condition,
@@ -16,11 +21,16 @@ __all__ = [
     "CompiledGraphEngine",
     "Context",
     "GraphRunHandler",
+    "NodeStateOption",
     "ChannelCondition",
     "TimerCondition",
     "AnyOfCondition",
+    "AllOfCondition",
+    "ConditionResult",
+    "WaitForResult",
     "channel_condition",
     "timer_condition",
     "any_of",
+    "all_of",
 ]
 

saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py

@@ -11,7 +11,7 @@ from dataclasses import dataclass
 from datetime import timedelta
 from typing import Any, Generic, TypeVar, cast
 
-from saf_python_sdk.rust_core_cffi import PyRustEngine
+from saf_python_sdk.langgraph_rust_core import PyRustEngine  # type: ignore[import-untyped]
 
 from saf_python_sdk.types import Command, Send
 
@@ -23,10 +23,16 @@ class _ChannelSpec:
     typ: Any
 
 
+@dataclass(frozen=True)
+class NodeStateOption:
+    locked_fields: tuple[str, ...] = ()
+
+
 @dataclass(frozen=True)
 class ChannelCondition:
     channel: str
-    n: int = 1
+    min: int = 1
+    max: int = 0
 
 
 @dataclass(frozen=True)
@@ -39,6 +45,23 @@ class AnyOfCondition:
     conditions: tuple[WaitCondition, ...]
 
 
+@dataclass(frozen=True)
+class AllOfCondition:
+    conditions: tuple[WaitCondition, ...]
+
+
+@dataclass(frozen=True)
+class ConditionResult:
+    met: bool
+    channel_name: str | None = None
+    values: list[Any] | None = None
+
+
+@dataclass(frozen=True)
+class WaitForResult:
+    conditions: list[ConditionResult]
+
+
 WaitCondition = ChannelCondition | TimerCondition
 
 _EXECUTOR_LOCK = threading.Lock()
@@ -79,7 +102,9 @@ class AdvancedStateGraph(Generic[StateT]):
     def __init__(self, state_schema: type[StateT]) -> None:
         self.state_schema = state_schema
         self._nodes: dict[str, Callable[..., Any]] = {}
+        self._node_options: dict[str, NodeStateOption] = {}
         self._async_channels: dict[str, _ChannelSpec] = {}
+        self._custom_output_streams: dict[str, _ChannelSpec] = {}
         self._entry_point: str | None = None
         self._finish_point: str | None = None
 
@@ -87,6 +112,8 @@ class AdvancedStateGraph(Generic[StateT]):
         self,
         name_or_node: str | Callable[..., Any],
         node: Callable[..., Any] | None = None,
+        *,
+        state_option: dict[str, Any] | NodeStateOption | None = None,
     ) -> str:
         if node is None:
             if not callable(name_or_node):
@@ -102,6 +129,7 @@ class AdvancedStateGraph(Generic[StateT]):
         if node_name in self._nodes:
             raise ValueError(f"Node `{node_name}` already exists")
         self._nodes[node_name] = node_fn
+        self._node_options[node_name] = _normalize_node_state_option(state_option)
         return node_name
 
     def add_async_channel(self, name: str, typ: Any) -> None:
@@ -109,13 +137,32 @@ class AdvancedStateGraph(Generic[StateT]):
             raise ValueError(f"Channel `{name}` already exists")
         self._async_channels[name] = _ChannelSpec(typ=typ)
 
-    def add_entry_node(self, node: Callable[..., Any]) -> str:
-        node_name = self.add_node(node)
+    def add_custom_outout_stream(self, name: str, typ: Any) -> None:
+        if name in self._custom_output_streams:
+            raise ValueError(f"Custom output stream `{name}` already exists")
+        self._custom_output_streams[name] = _ChannelSpec(typ=typ)
+
+    # Alias with corrected spelling.
+    def add_custom_output_stream(self, name: str, typ: Any) -> None:
+        self.add_custom_outout_stream(name, typ)
+
+    def add_entry_node(
+        self,
+        node: Callable[..., Any],
+        *,
+        state_option: dict[str, Any] | NodeStateOption | None = None,
+    ) -> str:
+        node_name = self.add_node(node, state_option=state_option)
         self._entry_point = self._resolve_node_name(node_name)
         return node_name
 
-    def add_finish_node(self, node: Callable[..., Any]) -> str:
-        node_name = self.add_node(node)
+    def add_finish_node(
+        self,
+        node: Callable[..., Any],
+        *,
+        state_option: dict[str, Any] | NodeStateOption | None = None,
+    ) -> str:
+        node_name = self.add_node(node, state_option=state_option)
         self._finish_point = self._resolve_node_name(node_name)
         return node_name
 
@@ -136,7 +183,9 @@ class AdvancedStateGraph(Generic[StateT]):
             raise ValueError(f"Finish point node `{self._finish_point}` does not exist")
         return CompiledGraphEngine(
             nodes=dict(self._nodes),
+            node_options=dict(self._node_options),
             async_channels=dict(self._async_channels),
+            custom_output_streams=dict(self._custom_output_streams),
             entry_point=self._entry_point,
             finish_point=self._finish_point,
         )
@@ -149,12 +198,16 @@ class CompiledGraphEngine(Generic[StateT]):
         self,
         *,
         nodes: dict[str, Callable[..., Any]],
+        node_options: dict[str, NodeStateOption],
         async_channels: dict[str, _ChannelSpec],
+        custom_output_streams: dict[str, _ChannelSpec],
         entry_point: str,
         finish_point: str | None,
     ) -> None:
         self._nodes = nodes
+        self._node_options = node_options
         self._async_channels = async_channels
+        self._custom_output_streams = custom_output_streams
         self._entry_point = entry_point
         self._finish_point = finish_point
 
@@ -168,6 +221,8 @@ class CompiledGraphEngine(Generic[StateT]):
         run = _GraphEngineRun(
             nodes=self._nodes,
             async_channel_specs=self._async_channels,
+            custom_output_stream_specs=self._custom_output_streams,
+            node_options=self._node_options,
             entry_point=self._entry_point,
             finish_point=self._finish_point,
             stream_mode=stream_mode,
@@ -182,20 +237,29 @@ class Context:
     def __init__(self, run: _GraphEngineRun) -> None:
         self._run = run
 
-    async def wait_for(self, target: WaitCondition | AnyOfCondition) -> Any:
+    async def wait_for(
+        self, target: WaitCondition | AnyOfCondition | AllOfCondition
+    ) -> WaitForResult:
         resumed = self._run._consume_resume_event(target)
         if resumed is not None:
             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)
 
     async def apublish_to_channel(self, channel: str, value: Any) -> None:
         await self._run.publish(channel, value)
 
-    def send_custom_stream_event(self, value: Any) -> None:
-        self._run.send_custom_stream_event(value)
+    def send_custom_stream_event(self, stream_name: str, value: Any) -> None:
+        self._run.send_custom_stream_event(stream_name, value)
 
 
 class GraphRunHandler(Generic[StateT]):
@@ -210,19 +274,13 @@ class GraphRunHandler(Generic[StateT]):
             raise RuntimeError("Run has already completed")
         await self._run.publish(channel, value)
 
-    async def receive_stream(self) -> Any | None:
-        while True:
-            loop = asyncio.get_running_loop()
-            event = await loop.run_in_executor(
-                _advanced_graph_executor(),
-                self._run.receive_stream_sync,
-            )
-            if event is not None or self._task.done():
-                return event
-            await asyncio.sleep(0.005)
+    async def receive_stream(self, stream_name: str) -> Any | None:
+        # Use a separate thread pool from graph execution to avoid deadlock
+        # when LANGGRAPH_ADVANCED_GRAPH_PY_THREADS is configured to 1.
+        return await asyncio.to_thread(self._run.receive_stream_sync, stream_name)
 
-    def close_stream(self) -> None:
-        self._run.close_stream_sync()
+    def close_all_streams(self) -> None:
+        self._run.close_all_streams_sync()
 
     async def aresult(self) -> StateT:
         return await self._task
@@ -237,17 +295,25 @@ class _GraphEngineRun:
         *,
         nodes: dict[str, Callable[..., Any]],
         async_channel_specs: dict[str, _ChannelSpec],
+        custom_output_stream_specs: dict[str, _ChannelSpec],
+        node_options: dict[str, NodeStateOption],
         entry_point: str,
         finish_point: str | None,
         stream_mode: str | None,
     ) -> None:
         self._nodes = nodes
         self._entry_point = entry_point
+        self._node_options = node_options
         self._finish_point = finish_point
         self._stream_mode = stream_mode
         self._rust_engine = PyRustEngine()
         for name in async_channel_specs:
             self._rust_engine.add_async_channel(name)
+        for stream_name in custom_output_stream_specs:
+            self._rust_engine.add_custom_output_stream(stream_name)
+        self._stream_ready = threading.Event()
+        if self._stream_mode is None:
+            self._stream_ready.set()
         self._tasks: set[asyncio.Task[list[Send]]] = set()
         self._finished = False
         self._state: Any = None
@@ -256,16 +322,25 @@ class _GraphEngineRun:
 
     async def run(self, initial_state: StateT) -> StateT:
         finish_point = self._finish_point or ""
+        if self._stream_mode is not None:
+            try:
+                self._rust_engine.start_stream(self._stream_mode)
+            finally:
+                self._stream_ready.set()
         loop = asyncio.get_running_loop()
-        result_obj = await loop.run_in_executor(
-            _advanced_graph_executor(),
-            self._rust_engine.run_graph_py,
-            self._entry_point,
-            finish_point,
-            initial_state,
-            self._execute_node_for_rust,
-            self._stream_mode,
-        )
+        try:
+            result_obj = await loop.run_in_executor(
+                _advanced_graph_executor(),
+                self._rust_engine.run_graph_py,
+                self._entry_point,
+                finish_point,
+                initial_state,
+                self._execute_node_for_rust,
+                None,
+                _node_locked_fields_payload(self._node_options),
+            )
+        finally:
+            self._stream_ready.set()
         self._state = result_obj
         return cast(StateT, self._state)
 
@@ -281,38 +356,58 @@ class _GraphEngineRun:
     def publish_nowait(self, channel: str, value: Any) -> None:
         self._publish_sync(channel, value)
 
-    async def wait_for(self, target: WaitCondition | AnyOfCondition) -> Any:
+    async def wait_for(
+        self, target: WaitCondition | AnyOfCondition | AllOfCondition
+    ) -> WaitForResult:
         if isinstance(target, ChannelCondition):
-            value = await self._wait_for_channel_values(target.channel, n=target.n)
-            return {
-                "condition": "channel",
-                "channel": target.channel,
-                "value": value,
-            }
+            value = await self._wait_for_channel_values(
+                target.channel, min=target.min, max=target.max
+            )
+            return WaitForResult(
+                conditions=[
+                    ConditionResult(
+                        met=True,
+                        channel_name=target.channel,
+                        values=_normalize_channel_values(value),
+                    )
+                ]
+            )
         if isinstance(target, TimerCondition):
             loop = asyncio.get_running_loop()
-            return await loop.run_in_executor(
+            await loop.run_in_executor(
                 _advanced_graph_executor(),
                 self._rust_engine.wait_timer,
                 target.seconds,
             )
+            return WaitForResult(conditions=[ConditionResult(met=True)])
         if isinstance(target, AnyOfCondition):
-            return await self._wait_for_any_of(target)
+            raw_event = await self._wait_for_any_of(target)
+            return _wait_for_result_from_any_of_event(target, raw_event)
+        if isinstance(target, AllOfCondition):
+            raw_event = await self._wait_for_all_of(target)
+            return _wait_for_result_from_all_of_event(target, raw_event)
         raise ValueError(f"Unsupported wait condition type: {type(target)!r}")
 
-    async def _wait_for_channel_values(self, channel: str, n: int) -> Any:
-        if n < 1:
-            raise ValueError("wait_for count `n` must be >= 1")
+    async def _wait_for_channel_values(
+        self, channel: str, min: int, max: int = 0
+    ) -> Any:
+        if min < 1:
+            raise ValueError("wait_for count `min` must be >= 1")
+        if max < 0:
+            raise ValueError("wait_for max count `max` must be >= 0")
+        if max != 0 and max < min:
+            raise ValueError("wait_for max count `max` must be 0 or >= min")
         loop = asyncio.get_running_loop()
         event = await loop.run_in_executor(
             _advanced_graph_executor(),
             self._rust_engine.wait_channel,
             channel,
-            n,
+            min,
+            max,
         )
         return event["value"]
 
-    async def _wait_for_any_of(self, condition: AnyOfCondition) -> Any:
+    async def _wait_for_any_of(self, condition: AnyOfCondition) -> dict[str, Any]:
         if not condition.conditions:
             raise ValueError("any_of() requires at least one condition")
         payload = {
@@ -325,23 +420,38 @@ class _GraphEngineRun:
             payload,
         )
 
+    async def _wait_for_all_of(self, condition: AllOfCondition) -> dict[str, Any]:
+        if not condition.conditions:
+            raise ValueError("all_of() requires at least one condition")
+        payload = {
+            "conditions": [_condition_to_rust(cond) for cond in condition.conditions]
+        }
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(
+            _advanced_graph_executor(),
+            self._rust_engine.wait_all_of_obj,
+            payload,
+        )
+
     def _publish_sync(self, channel: str, value: Any) -> None:
         self._rust_engine.publish_obj(channel, value)
 
-    def send_custom_stream_event(self, value: Any) -> None:
-        self._rust_engine.send_custom_stream_event_obj(value)
+    def send_custom_stream_event(self, stream_name: str, value: Any) -> None:
+        self._rust_engine.send_custom_stream_event_obj(stream_name, value)
 
-    def receive_stream_sync(self) -> Any | None:
-        return self._rust_engine.receive_stream_obj()
+    def receive_stream_sync(self, stream_name: str) -> Any | None:
+        self._stream_ready.wait()
+        return self._rust_engine.receive_stream_obj(stream_name)
 
-    def close_stream_sync(self) -> None:
-        self._rust_engine.close_stream()
+    def close_all_streams_sync(self) -> None:
+        self._rust_engine.close_all_streams()
 
     def _execute_node_for_rust(
         self, node_name: str, node_input: Any, state: Any
     ) -> 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]
@@ -355,6 +465,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
@@ -374,12 +485,20 @@ class _GraphEngineRun:
     def _set_resume_event(self, event: dict[str, Any] | None) -> None:
         self._local.resume_event = event
 
-    def _consume_resume_event(self, target: WaitCondition | AnyOfCondition) -> Any | None:
+    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:
         event = cast(dict[str, Any] | None, getattr(self._local, "resume_event", None))
         if event is None:
             return None
         self._local.resume_event = None
-        return event
+        return _wait_for_result_from_resume_event(target, event)
 
     def _run_awaitable_in_worker(self, awaitable: Coroutine[Any, Any, Any]) -> Any:
         # Create and close a dedicated loop per execution to avoid
@@ -431,10 +550,51 @@ def _normalize_goto(goto: Any, *, default_input: Any) -> list[Send]:
     return []
 
 
-def channel_condition(channel: str, n: int = 1) -> ChannelCondition:
-    if n < 1:
-        raise ValueError("channel_condition `n` must be >= 1")
-    return ChannelCondition(channel=channel, n=n)
+def _normalize_node_state_option(
+    state_option: dict[str, Any] | NodeStateOption | None,
+) -> NodeStateOption:
+    if state_option is None:
+        return NodeStateOption()
+    if isinstance(state_option, NodeStateOption):
+        return state_option
+    if not isinstance(state_option, dict):
+        raise TypeError("state_option must be a dict or NodeStateOption")
+    locked_fields_raw = state_option.get("locked_fields", ())
+    if locked_fields_raw is None:
+        return NodeStateOption()
+    if not isinstance(locked_fields_raw, Sequence) or isinstance(
+        locked_fields_raw, (str, bytes)
+    ):
+        raise TypeError("state_option['locked_fields'] must be a sequence of strings")
+    locked_fields: list[str] = []
+    for field in locked_fields_raw:
+        if not isinstance(field, str):
+            raise TypeError("locked field names must be strings")
+        if not field:
+            continue
+        locked_fields.append(field)
+    return NodeStateOption(locked_fields=tuple(locked_fields))
+
+
+def _node_locked_fields_payload(
+    node_options: dict[str, NodeStateOption],
+) -> dict[str, list[str]]:
+    payload: dict[str, list[str]] = {}
+    for node_name, option in node_options.items():
+        if not option.locked_fields:
+            continue
+        payload[node_name] = list(option.locked_fields)
+    return payload
+
+
+def channel_condition(channel: str, min: int = 1, max: int = 0) -> ChannelCondition:
+    if min < 1:
+        raise ValueError("channel_condition `min` must be >= 1")
+    if max < 0:
+        raise ValueError("channel_condition `max` must be >= 0")
+    if max != 0 and max < min:
+        raise ValueError("channel_condition `max` must be 0 or >= min")
+    return ChannelCondition(channel=channel, min=min, max=max)
 
 
 def timer_condition(
@@ -470,15 +630,152 @@ def any_of(*conditions: WaitCondition) -> AnyOfCondition:
     return AnyOfCondition(conditions=tuple(conditions))
 
 
+def all_of(*conditions: WaitCondition) -> AllOfCondition:
+    if not conditions:
+        raise ValueError("all_of() requires at least one condition")
+    return AllOfCondition(conditions=tuple(conditions))
+
+
+def _normalize_channel_values(value: Any) -> list[Any]:
+    if isinstance(value, list):
+        return value
+    return [value]
+
+
+def _wait_for_result_from_resume_event(
+    target: WaitCondition | AnyOfCondition | AllOfCondition, event: dict[str, Any]
+) -> WaitForResult:
+    if isinstance(target, ChannelCondition):
+        return WaitForResult(
+            conditions=[
+                ConditionResult(
+                    met=True,
+                    channel_name=target.channel,
+                    values=_normalize_channel_values(event.get("value")),
+                )
+            ]
+        )
+    if isinstance(target, TimerCondition):
+        return WaitForResult(conditions=[ConditionResult(met=True)])
+    if isinstance(target, AllOfCondition):
+        return _wait_for_result_from_all_of_event(target, event)
+    return _wait_for_result_from_any_of_event(target, event)
+
+
+def _wait_for_result_from_any_of_event(
+    target: AnyOfCondition, event: dict[str, Any]
+) -> WaitForResult:
+    results = [ConditionResult(met=False) for _ in target.conditions]
+    condition = event.get("condition")
+
+    if condition == "timer":
+        for idx, cond in enumerate(target.conditions):
+            if isinstance(cond, TimerCondition):
+                results[idx] = ConditionResult(met=True)
+                break
+        return WaitForResult(conditions=results)
+
+    if condition != "channel":
+        return WaitForResult(conditions=results)
+
+    channel = cast(str | None, event.get("channel"))
+    value = event.get("value")
+
+    if channel == "__any_of__" and isinstance(value, list):
+        matched = list(value)
+        cursor = 0
+        for idx, cond in enumerate(target.conditions):
+            if not isinstance(cond, ChannelCondition):
+                continue
+            if cursor >= len(matched):
+                continue
+            item = matched[cursor]
+            if (
+                isinstance(item, dict)
+                and item.get("channel") == cond.channel
+                and "value" in item
+            ):
+                results[idx] = ConditionResult(
+                    met=True,
+                    channel_name=cond.channel,
+                    values=_normalize_channel_values(item.get("value")),
+                )
+                cursor += 1
+        return WaitForResult(conditions=results)
+
+    for idx, cond in enumerate(target.conditions):
+        if isinstance(cond, ChannelCondition) and cond.channel == channel:
+            results[idx] = ConditionResult(
+                met=True,
+                channel_name=cond.channel,
+                values=_normalize_channel_values(value),
+            )
+            break
+    return WaitForResult(conditions=results)
+
+
+def _wait_for_result_from_all_of_event(
+    target: AllOfCondition, event: dict[str, Any]
+) -> WaitForResult:
+    results = [ConditionResult(met=False) for _ in target.conditions]
+    condition = event.get("condition")
+
+    # all_of completion implies all timer conditions are satisfied.
+    for idx, cond in enumerate(target.conditions):
+        if isinstance(cond, TimerCondition):
+            results[idx] = ConditionResult(met=True)
+
+    if condition != "channel":
+        return WaitForResult(conditions=results)
+
+    channel = cast(str | None, event.get("channel"))
+    value = event.get("value")
+
+    if channel == "__all_of__" and isinstance(value, list):
+        matched_by_channel: dict[str, Any] = {}
+        for item in value:
+            if isinstance(item, dict) and isinstance(item.get("channel"), str):
+                matched_by_channel[cast(str, item["channel"])] = item.get("value")
+
+        for idx, cond in enumerate(target.conditions):
+            if not isinstance(cond, ChannelCondition):
+                continue
+            if cond.channel not in matched_by_channel:
+                continue
+            results[idx] = ConditionResult(
+                met=True,
+                channel_name=cond.channel,
+                values=_normalize_channel_values(matched_by_channel[cond.channel]),
+            )
+        return WaitForResult(conditions=results)
+
+    for idx, cond in enumerate(target.conditions):
+        if isinstance(cond, ChannelCondition) and cond.channel == channel:
+            results[idx] = ConditionResult(
+                met=True,
+                channel_name=cond.channel,
+                values=_normalize_channel_values(value),
+            )
+            break
+    return WaitForResult(conditions=results)
+
+
 def _condition_to_rust(condition: WaitCondition) -> dict[str, Any]:
     if isinstance(condition, ChannelCondition):
-        return {"kind": "channel", "channel": condition.channel, "n": condition.n}
+        return {
+            "kind": "channel",
+            "channel": condition.channel,
+            "min": condition.min,
+            "max": condition.max,
+        }
     if isinstance(condition, TimerCondition):
         return {"kind": "timer", "seconds": condition.seconds}
     raise TypeError(f"Unsupported condition type: {type(condition)!r}")
 
 
-def _target_to_suspend_payload(target: WaitCondition | AnyOfCondition) -> dict[str, Any]:
+def _target_to_suspend_payload(
+    target: WaitCondition | AnyOfCondition | AllOfCondition,
+) -> dict[str, Any]:
     if isinstance(target, AnyOfCondition):
         return {
             "kind": "any_of",
@@ -486,6 +783,13 @@ def _target_to_suspend_payload(target: WaitCondition | AnyOfCondition) -> dict[s
                 "conditions": [_condition_to_rust(cond) for cond in target.conditions]
             },
         }
+    if isinstance(target, AllOfCondition):
+        return {
+            "kind": "all_of",
+            "all_of": {
+                "conditions": [_condition_to_rust(cond) for cond in target.conditions]
+            },
+        }
     return {"kind": "condition", "condition": _condition_to_rust(target)}
 
 

saf-python-sdk/python/saf_python_sdk/rust_core_cffi.py

@@ -1,347 +0,0 @@
-from __future__ import annotations
-
-import ctypes
-import dataclasses
-import json
-import subprocess
-from copy import deepcopy
-from pathlib import Path
-import threading
-from typing import Any, Callable, get_args, get_origin
-
-
-class PyRustEngine:
-    def __init__(self) -> None:
-        self._lib = _load_rust_lib()
-        self._engine = self._lib.rc_engine_new()
-        if not self._engine:
-            raise RuntimeError("failed to create rust engine")
-
-    def __del__(self) -> None:
-        engine = getattr(self, "_engine", None)
-        if engine:
-            self._lib.rc_engine_free(engine)
-            self._engine = None
-
-    def add_async_channel(self, name: str) -> None:
-        self._call_status(self._lib.rc_add_async_channel, name.encode())
-
-    def publish_obj(self, channel: str, value: Any) -> None:
-        payload = json.dumps(value, ensure_ascii=False).encode()
-        self._call_status(self._lib.rc_publish_json, channel.encode(), payload)
-
-    def wait_any_of_obj(self, any_of_payload: Any) -> Any:
-        payload = json.dumps(any_of_payload, ensure_ascii=False).encode()
-        raw = self._consume_json_ptr(self._lib.rc_wait_any_of_json(self._engine, payload))
-        if not raw.get("ok"):
-            raise ValueError(raw.get("error", "rust wait_any_of failed"))
-        return raw["event"]
-
-    def wait_channel(self, channel: str, n: int) -> Any:
-        return self.wait_any_of_obj({"conditions": [{"kind": "channel", "channel": channel, "n": n}]})
-
-    def wait_timer(self, seconds: float) -> Any:
-        return self.wait_any_of_obj({"conditions": [{"kind": "timer", "seconds": seconds}]})
-
-    def start_stream(self, stream_mode: str | None) -> None:
-        encoded = stream_mode.encode() if stream_mode is not None else None
-        self._call_status(self._lib.rc_start_stream, encoded)
-
-    def receive_stream_obj(self) -> Any | None:
-        raw = self._consume_json_ptr(self._lib.rc_receive_stream_json(self._engine))
-        if not raw.get("ok"):
-            raise ValueError(raw.get("error", "rust receive_stream failed"))
-        if not raw.get("has_event", False):
-            return None
-        return raw.get("event")
-
-    def send_custom_stream_event_obj(self, value: Any) -> None:
-        payload = json.dumps(value, ensure_ascii=False).encode()
-        self._call_status(self._lib.rc_send_custom_stream_event, payload)
-
-    def close_stream(self) -> None:
-        self._call_status(self._lib.rc_close_stream)
-
-    def run_graph_py(
-        self,
-        entry_point: str,
-        finish_point: str,
-        initial_state: Any,
-        callback: Callable[[str, Any, Any], dict[str, Any]],
-        stream_mode: str | None = None,
-    ) -> Any:
-        shared_state = initial_state
-        shared_state_lock = threading.Lock()
-        state_type = type(initial_state)
-        use_shared_state = dataclasses.is_dataclass(initial_state)
-        initial_state_json = json.dumps(_to_jsonable(shared_state), ensure_ascii=False).encode()
-        callback_c = _make_node_callback(
-            callback,
-            state_type,
-            use_shared_state,
-            shared_state,
-            shared_state_lock,
-        )
-        stream_mode_encoded = stream_mode.encode() if stream_mode is not None else None
-        out = self._consume_json_ptr(
-            self._lib.rc_run_graph_json(
-                self._engine,
-                entry_point.encode(),
-                finish_point.encode(),
-                initial_state_json,
-                initial_state_json,
-                stream_mode_encoded,
-                ctypes.c_ulong(0),
-                callback_c,
-            )
-        )
-        if not out.get("ok"):
-            raise ValueError(out.get("error", "rust run_graph failed"))
-        if use_shared_state:
-            return shared_state
-        return _coerce_for_type(out["state"], state_type)
-
-    def _call_status(self, func: Any, *args: Any) -> None:
-        raw = self._consume_json_ptr(func(self._engine, *args))
-        if not raw.get("ok"):
-            raise ValueError(raw.get("error", "rust call failed"))
-
-    def _consume_json_ptr(self, ptr: ctypes.c_void_p) -> dict[str, Any]:
-        if not ptr:
-            raise RuntimeError("rust returned null string pointer")
-        try:
-            text = ctypes.cast(ptr, ctypes.c_char_p).value
-            if text is None:
-                raise RuntimeError("rust returned empty string pointer")
-            return json.loads(text.decode())
-        finally:
-            self._lib.rc_string_free(ptr)
-
-
-def _make_node_callback(
-    callback: Callable[[str, Any, Any], dict[str, Any]],
-    state_type: type[Any],
-    use_shared_state: bool,
-    shared_state: Any,
-    shared_state_lock: threading.Lock,
-) -> ctypes.CFUNCTYPE:  # type: ignore[type-arg]
-    cb_type = ctypes.CFUNCTYPE(
-        ctypes.c_void_p,
-        ctypes.c_ulong,
-        ctypes.c_char_p,
-        ctypes.c_char_p,
-        ctypes.c_char_p,
-    )
-    libc = ctypes.CDLL(None)
-    libc.malloc.argtypes = [ctypes.c_size_t]
-    libc.malloc.restype = ctypes.c_void_p
-
-    @cb_type
-    def _callback(
-        _user_data: int,
-        node_ptr: bytes,
-        arg_ptr: bytes,
-        state_ptr: bytes,
-    ) -> ctypes.c_void_p:
-        try:
-            node = node_ptr.decode()
-            arg = json.loads(arg_ptr.decode())
-            if use_shared_state:
-                with shared_state_lock:
-                    before = deepcopy(_to_jsonable(shared_state))
-                    result = callback(node, arg, shared_state)
-                    state_after_call = shared_state
-            else:
-                state_raw = json.loads(state_ptr.decode())
-                state_snapshot = _coerce_for_type(state_raw, state_type)
-                before = deepcopy(_to_jsonable(state_snapshot))
-                result = callback(node, arg, state_snapshot)
-                state_after_call = state_snapshot
-            if "suspend" in result:
-                envelope = {"ok": True, "suspend": result["suspend"]}
-            else:
-                update = result.get("update")
-                if update is not None:
-                    update_json = _to_jsonable(update)
-                    if update_json == before:
-                        update = None
-                        update_json = None
-                    else:
-                        if use_shared_state:
-                            _apply_update_to_state(shared_state, update)
-                after = _to_jsonable(state_after_call)
-                if update is None and after != before:
-                    update_json = after
-                elif update is None:
-                    update_json = None
-                else:
-                    update_json = _to_jsonable(update)
-                sends = []
-                for item in result.get("sends", []):
-                    if not isinstance(item, dict):
-                        continue
-                    sends.append(
-                        {
-                            "node": item.get("node"),
-                            "arg": _to_jsonable(item.get("arg")),
-                        }
-                    )
-                envelope = {
-                    "ok": True,
-                    "payload": {
-                        "update": update_json,
-                        "sends": sends,
-                    },
-                }
-        except Exception as exc:  # noqa: BLE001
-            envelope = {"ok": False, "error": f"python callback failed: {exc}"}
-        return _malloc_c_string(json.dumps(envelope, ensure_ascii=False).encode(), libc)
-
-    return _callback
-
-
-def _malloc_c_string(payload: bytes, libc: Any) -> ctypes.c_void_p:
-    size = len(payload) + 1
-    ptr = libc.malloc(size)
-    if not ptr:
-        return ctypes.c_void_p(0)
-    ctypes.memmove(ptr, payload, len(payload))
-    ctypes.memset(ctypes.c_void_p(ptr + len(payload)), 0, 1)
-    return ptr
-
-
-def _load_rust_lib() -> ctypes.CDLL:
-    env = Path.cwd()
-    root = _find_repo_root(env)
-    rust_core = root / "rust-core"
-    lib_path = _resolve_lib_path(rust_core)
-    if not lib_path.exists():
-        subprocess.run(["cargo", "build"], cwd=rust_core, check=True)
-    lib = ctypes.CDLL(str(lib_path))
-    _configure_signatures(lib)
-    return lib
-
-
-def _find_repo_root(start: Path) -> Path:
-    current = start.resolve()
-    for candidate in [current, *current.parents]:
-        if (candidate / "rust-core").exists() and (candidate / "saf-python-sdk").exists():
-            return candidate
-    here = Path(__file__).resolve()
-    return here.parents[3]
-
-
-def _resolve_lib_path(rust_core: Path) -> Path:
-    if (rust_core / "target" / "debug" / "liblanggraph_rust_core.dylib").exists():
-        return rust_core / "target" / "debug" / "liblanggraph_rust_core.dylib"
-    if (rust_core / "target" / "debug" / "liblanggraph_rust_core.so").exists():
-        return rust_core / "target" / "debug" / "liblanggraph_rust_core.so"
-    if (rust_core / "target" / "debug" / "langgraph_rust_core.dll").exists():
-        return rust_core / "target" / "debug" / "langgraph_rust_core.dll"
-    return rust_core / "target" / "debug" / "liblanggraph_rust_core.dylib"
-
-
-def _configure_signatures(lib: ctypes.CDLL) -> None:
-    cb_type = ctypes.CFUNCTYPE(
-        ctypes.c_void_p,
-        ctypes.c_ulong,
-        ctypes.c_char_p,
-        ctypes.c_char_p,
-        ctypes.c_char_p,
-    )
-    lib.rc_engine_new.argtypes = []
-    lib.rc_engine_new.restype = ctypes.c_void_p
-    lib.rc_engine_free.argtypes = [ctypes.c_void_p]
-    lib.rc_engine_free.restype = None
-    lib.rc_string_free.argtypes = [ctypes.c_void_p]
-    lib.rc_string_free.restype = None
-    lib.rc_add_async_channel.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
-    lib.rc_add_async_channel.restype = ctypes.c_void_p
-    lib.rc_publish_json.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p]
-    lib.rc_publish_json.restype = ctypes.c_void_p
-    lib.rc_wait_any_of_json.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
-    lib.rc_wait_any_of_json.restype = ctypes.c_void_p
-    lib.rc_start_stream.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
-    lib.rc_start_stream.restype = ctypes.c_void_p
-    lib.rc_receive_stream_json.argtypes = [ctypes.c_void_p]
-    lib.rc_receive_stream_json.restype = ctypes.c_void_p
-    lib.rc_send_custom_stream_event.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
-    lib.rc_send_custom_stream_event.restype = ctypes.c_void_p
-    lib.rc_close_stream.argtypes = [ctypes.c_void_p]
-    lib.rc_close_stream.restype = ctypes.c_void_p
-    lib.rc_run_graph_json.argtypes = [
-        ctypes.c_void_p,
-        ctypes.c_char_p,
-        ctypes.c_char_p,
-        ctypes.c_char_p,
-        ctypes.c_char_p,
-        ctypes.c_char_p,
-        ctypes.c_ulong,
-        cb_type,
-    ]
-    lib.rc_run_graph_json.restype = ctypes.c_void_p
-
-
-def _to_jsonable(value: Any) -> Any:
-    if value is None or isinstance(value, (str, int, float, bool)):
-        return value
-    if dataclasses.is_dataclass(value):
-        return {field.name: _to_jsonable(getattr(value, field.name)) for field in dataclasses.fields(value)}
-    model_dump = getattr(value, "model_dump", None)
-    if callable(model_dump):
-        return _to_jsonable(model_dump())
-    if isinstance(value, dict):
-        return {str(k): _to_jsonable(v) for k, v in value.items()}
-    if isinstance(value, (list, tuple, set)):
-        return [_to_jsonable(v) for v in value]
-    return value
-
-
-def _coerce_for_type(value: Any, typ: Any) -> Any:
-    if value is None:
-        return None
-    origin = get_origin(typ)
-    args = get_args(typ)
-    if origin is not None:
-        if origin in (list, tuple, set):
-            item_type = args[0] if args else Any
-            items = [_coerce_for_type(v, item_type) for v in value]
-            if origin is tuple:
-                return tuple(items)
-            if origin is set:
-                return set(items)
-            return items
-        if origin is dict:
-            value_type = args[1] if len(args) == 2 else Any
-            return {k: _coerce_for_type(v, value_type) for k, v in value.items()}
-    if isinstance(typ, type):
-        if dataclasses.is_dataclass(typ):
-            kwargs = {}
-            for field in dataclasses.fields(typ):
-                kwargs[field.name] = _coerce_for_type(value.get(field.name), field.type)
-            return typ(**kwargs)
-        model_validate = getattr(typ, "model_validate", None)
-        if callable(model_validate):
-            return model_validate(value)
-    return value
-
-
-def _apply_update_to_state(state: Any, update: Any) -> None:
-    if update is None:
-        return
-    if isinstance(state, dict):
-        if isinstance(update, dict):
-            state.update(update)
-            return
-        if dataclasses.is_dataclass(update):
-            state.update(_to_jsonable(update))
-            return
-        return
-    if dataclasses.is_dataclass(state):
-        if dataclasses.is_dataclass(update):
-            for field in dataclasses.fields(state):
-                setattr(state, field.name, getattr(update, field.name))
-            return
-        if isinstance(update, dict):
-            for key, val in update.items():
-                setattr(state, key, val)

saf-python-sdk/tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py

@@ -40,7 +40,7 @@ def build_advanced_parallel() -> Any:
         return Command(goto=sends)
 
     async def end_node(ctx: Any, state: dict[str, Any]) -> dict[str, Any]:
-        await ctx.wait_for(channel_condition(done_channel, n=MIDDLE_COUNT))
+        await ctx.wait_for(channel_condition(done_channel, min=MIDDLE_COUNT))
         out = dict(state)
         out["done"] = True
         return out
@@ -161,7 +161,7 @@ def build_advanced_parallel_blocking() -> Any:
         return Command(goto=sends)
 
     async def end_node_blocking(ctx: Any, state: dict[str, Any]) -> dict[str, Any]:
-        await ctx.wait_for(channel_condition(done_channel, n=MIDDLE_COUNT))
+        await ctx.wait_for(channel_condition(done_channel, min=MIDDLE_COUNT))
         out = dict(state)
         out["done"] = True
         return out

saf-python-sdk/tests/advanced-graph/test_primitives.py

@@ -1,7 +1,16 @@
+import asyncio
+import time
 import pytest
 from typing_extensions import TypedDict
 
-from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context
+from saf_python_sdk.advanced_graph import (
+    AdvancedStateGraph,
+    Context,
+    all_of,
+    any_of,
+    channel_condition,
+    timer_condition,
+)
 from saf_python_sdk.types import Command, Send
 
 pytestmark = pytest.mark.anyio
@@ -65,3 +74,232 @@ async def test_run_ends_without_finish_node() -> None:
     assert result["done"] == "stopped"
     assert result["logs"] == ["start", "middle:from_start"]
 
+
+async def test_channel_wait_respects_max_m() -> None:
+    graph = AdvancedStateGraph(PrimitiveState)
+    graph.add_async_channel("events", list[str])
+
+    async def start_node(ctx: Context, state: PrimitiveState) -> Command:
+        ctx.publish_to_channel("events", "a")
+        ctx.publish_to_channel("events", "b")
+        ctx.publish_to_channel("events", "c")
+        return Command(update=state, goto=Send("wait_node", None))
+
+    async def wait_node(ctx: Context, _input: None, state: PrimitiveState) -> dict[str, object]:
+        result = await ctx.wait_for(channel_condition("events", min=2, max=4))
+        values = result.conditions[0].values or []
+        assert isinstance(values, list)
+        return {"counter": len(values), "logs": values, "done": "ok"}
+
+    graph.add_entry_node(start_node)
+    graph.add_finish_node(wait_node)
+
+    result = await graph.compile().ainvoke({"counter": 0, "logs": [], "done": None})
+    assert result["counter"] == 3
+    assert result["logs"] == ["a", "b", "c"]
+    assert result["done"] == "ok"
+
+
+async def test_any_of_consumes_all_ready_channels() -> None:
+    graph = AdvancedStateGraph(PrimitiveState)
+    graph.add_async_channel("alpha", str)
+    graph.add_async_channel("beta", str)
+
+    async def start_node(ctx: Context, state: PrimitiveState) -> Command:
+        ctx.publish_to_channel("alpha", "a1")
+        ctx.publish_to_channel("beta", "b1")
+        return Command(update=state, goto=Send("wait_node", None))
+
+    async def wait_node(ctx: Context, _input: None, state: PrimitiveState) -> Command:
+        first = await ctx.wait_for(
+            any_of(channel_condition("alpha"), channel_condition("beta"))
+        )
+        assert len(first.conditions) == 2
+        assert first.conditions[0].met is True
+        assert first.conditions[0].channel_name == "alpha"
+        assert first.conditions[0].values == ["a1"]
+        assert first.conditions[1].met is True
+        assert first.conditions[1].channel_name == "beta"
+        assert first.conditions[1].values == ["b1"]
+        ctx.publish_to_channel("beta", "b2")
+        return Command(
+            update={"counter": 1, "logs": ["matched=2"], "done": None},
+            goto=Send("verify_node", None),
+        )
+
+    async def verify_node(
+        ctx: Context, _input: None, state: PrimitiveState
+    ) -> dict[str, object]:
+        second = await ctx.wait_for(channel_condition("beta"))
+        values = second.conditions[0].values or []
+        return {
+            "counter": 2,
+            "logs": [*state["logs"], f"beta={values[0]}"],
+            "done": "ok",
+        }
+
+    graph.add_entry_node(start_node)
+    graph.add_node(wait_node)
+    graph.add_finish_node(verify_node)
+
+    result = await graph.compile().ainvoke({"counter": 0, "logs": [], "done": None})
+    assert result["counter"] == 2
+    assert result["logs"] == [
+        "matched=2",
+        "beta=b2",
+    ]
+    assert result["done"] == "ok"
+
+
+async def test_all_of_waits_until_all_channels_are_ready() -> None:
+    graph = AdvancedStateGraph(PrimitiveState)
+    graph.add_async_channel("alpha", str)
+    graph.add_async_channel("beta", str)
+
+    async def start_node(ctx: Context, state: PrimitiveState) -> Command:
+        ctx.publish_to_channel("alpha", "a1")
+        return Command(
+            update=state,
+            goto=[Send("wait_node", None), Send("publish_beta_node", None)],
+        )
+
+    async def publish_beta_node(ctx: Context, _input: None, state: PrimitiveState) -> Command:
+        await ctx.wait_for(timer_condition(seconds=0.02))
+        ctx.publish_to_channel("beta", "b1")
+        return Command(update=state)
+
+    async def wait_node(ctx: Context, _input: None, state: PrimitiveState) -> dict[str, object]:
+        waited = await ctx.wait_for(
+            all_of(channel_condition("alpha"), channel_condition("beta"))
+        )
+        assert len(waited.conditions) == 2
+        assert waited.conditions[0].met is True
+        assert waited.conditions[0].channel_name == "alpha"
+        assert waited.conditions[0].values == ["a1"]
+        assert waited.conditions[1].met is True
+        assert waited.conditions[1].channel_name == "beta"
+        assert waited.conditions[1].values == ["b1"]
+        return {"counter": 1, "logs": ["all_of_channels"], "done": "ok"}
+
+    graph.add_entry_node(start_node)
+    graph.add_node(publish_beta_node)
+    graph.add_finish_node(wait_node)
+
+    result = await graph.compile().ainvoke({"counter": 0, "logs": [], "done": None})
+    assert result["counter"] == 1
+    assert result["logs"] == ["all_of_channels"]
+    assert result["done"] == "ok"
+
+
+async def test_all_of_channel_and_timer_marks_both_conditions() -> None:
+    graph = AdvancedStateGraph(PrimitiveState)
+    graph.add_async_channel("alpha", str)
+
+    async def start_node(ctx: Context, state: PrimitiveState) -> Command:
+        ctx.publish_to_channel("alpha", "a1")
+        return Command(update=state, goto=Send("wait_node", None))
+
+    async def wait_node(ctx: Context, _input: None, state: PrimitiveState) -> dict[str, object]:
+        waited = await ctx.wait_for(
+            all_of(channel_condition("alpha"), timer_condition(seconds=0.02))
+        )
+        assert len(waited.conditions) == 2
+        assert waited.conditions[0].met is True
+        assert waited.conditions[0].channel_name == "alpha"
+        assert waited.conditions[0].values == ["a1"]
+        assert waited.conditions[1].met is True
+        return {"counter": 1, "logs": ["all_of_channel_timer"], "done": "ok"}
+
+    graph.add_entry_node(start_node)
+    graph.add_finish_node(wait_node)
+
+    result = await graph.compile().ainvoke({"counter": 0, "logs": [], "done": None})
+    assert result["counter"] == 1
+    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"
+
+
+async def test_state_field_locking_serializes_conflicting_nodes() -> None:
+    graph = AdvancedStateGraph(PrimitiveState)
+    graph.add_async_channel("done", str)
+    intervals: dict[str, tuple[float, float]] = {}
+
+    async def start_node(state: PrimitiveState) -> Command:
+        return Command(
+            update=state,
+            goto=[
+                Send("worker_a", None),
+                Send("worker_b", None),
+                Send("wait_node", None),
+            ],
+        )
+
+    async def worker_a(ctx: Context, _input: None, state: PrimitiveState) -> Command:
+        started = time.perf_counter()
+        await asyncio.sleep(0.04)
+        ended = time.perf_counter()
+        intervals["a"] = (started, ended)
+        ctx.publish_to_channel("done", "a")
+        return Command(update=state)
+
+    async def worker_b(ctx: Context, _input: None, state: PrimitiveState) -> Command:
+        started = time.perf_counter()
+        await asyncio.sleep(0.04)
+        ended = time.perf_counter()
+        intervals["b"] = (started, ended)
+        ctx.publish_to_channel("done", "b")
+        return Command(update=state)
+
+    async def wait_node(ctx: Context, _input: None, state: PrimitiveState) -> Command:
+        await ctx.wait_for(channel_condition("done", min=2))
+        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(worker_a, state_option={"locked_fields": ["counter"]})
+    graph.add_node(worker_b, state_option={"locked_fields": ["counter"]})
+    graph.add_node(wait_node)
+    graph.add_finish_node(finish_node)
+
+    result = await graph.compile().ainvoke({"counter": 0, "logs": [], "done": None})
+    assert result["done"] == "ok"
+    assert "a" in intervals and "b" in intervals
+    a_start, a_end = intervals["a"]
+    b_start, b_end = intervals["b"]
+    serialized = (a_end <= b_start) or (b_end <= a_start)
+    assert serialized, f"expected serialized execution, got overlap: a={intervals['a']} b={intervals['b']}"
+

saf-python-sdk/tests/advanced-graph/test_streaming.py

@@ -15,11 +15,13 @@ class StreamState(TypedDict):
 
 async def test_custom_stream_receive_and_close() -> None:
     graph: AdvancedStateGraph[StreamState] = AdvancedStateGraph(StreamState)
+    graph.add_custom_outout_stream("high", dict[str, int | str])
+    graph.add_custom_outout_stream("regular", dict[str, int | str])
 
     async def start_node(ctx: Context, state: StreamState) -> Command:
-        ctx.send_custom_stream_event({"step": "start", "value": 1})
+        ctx.send_custom_stream_event("high", {"step": "start", "value": 1})
         await asyncio.sleep(0.08)
-        ctx.send_custom_stream_event({"step": "start", "value": 2})
+        ctx.send_custom_stream_event("regular", {"step": "start", "value": 2})
         return Command(update=state, goto=Send("finish_node", None))
 
     async def finish_node(state: StreamState) -> dict[str, bool]:
@@ -30,13 +32,18 @@ async def test_custom_stream_receive_and_close() -> None:
 
     handler = await graph.compile().astart({"done": False}, stream_mode="custom")
 
-    event = await handler.receive_stream()
+    event = await handler.receive_stream("high")
     assert isinstance(event, dict)
     assert event["step"] == "start"
     assert event["value"] == 1
 
-    handler.close_stream()
-    assert await handler.receive_stream() is None
+    event_regular = await handler.receive_stream("regular")
+    assert isinstance(event_regular, dict)
+    assert event_regular["value"] == 2
+
+    handler.close_all_streams()
+    assert await handler.receive_stream("high") is None
+    assert await handler.receive_stream("regular") is None
 
     result = await handler.aresult()
     assert result["done"] is True
@@ -44,9 +51,10 @@ async def test_custom_stream_receive_and_close() -> None:
 
 async def test_only_custom_stream_mode_supported() -> None:
     graph: AdvancedStateGraph[StreamState] = AdvancedStateGraph(StreamState)
+    graph.add_custom_outout_stream("regular", dict[str, str])
 
     async def start_node(ctx: Context, state: StreamState) -> Command:
-        ctx.send_custom_stream_event({"hello": "world"})
+        ctx.send_custom_stream_event("regular", {"hello": "world"})
         return Command(update=state)
 
     graph.add_entry_node(start_node)

saf-python-sdk/tests/advanced-graph/test_sub_agents.py

@@ -92,7 +92,7 @@ def build_main_agent(planner: MockLLM, sub_agent: Any) -> Any:
 
     async def wait_node(ctx: Context, state: MainAgentState) -> Command:
         # Lightweight interrupt: only this node blocks for the next relevant signal.
-        event = await ctx.wait_for(
+        result = await ctx.wait_for(
             any_of(
                 channel_condition("tool_completion_channel"),
                 channel_condition("subagent_completion_channel"),
@@ -100,21 +100,33 @@ def build_main_agent(planner: MockLLM, sub_agent: Any) -> Any:
                 timer_condition(seconds=1),
             )
         )
-        if event["condition"] == "channel":
-            channel = event["channel"]
-            payload = event["value"]
-            if channel == "tool_completion_channel":
-                state["output"].append(f"tool: {payload}")
-            elif channel == "subagent_completion_channel":
-                state["output"].append(f"sub_agent: {payload}")
-            elif channel == "user_input_channel":
-                state["output"].append(f"user_input: {payload}")
+        had_channel_update = False
+        for item in result.conditions:
+            if not item.met:
+                continue
+            if item.channel_name == "tool_completion_channel":
+                payloads = item.values or []
+                for payload in payloads:
+                    state["output"].append(f"tool: {payload}")
+                had_channel_update = True
+            elif item.channel_name == "subagent_completion_channel":
+                payloads = item.values or []
+                for payload in payloads:
+                    state["output"].append(f"sub_agent: {payload}")
+                had_channel_update = True
+            elif item.channel_name == "user_input_channel":
+                payloads = item.values or []
+                for payload in payloads:
+                    state["output"].append(f"user_input: {payload}")
+                had_channel_update = True
+
+        if had_channel_update:
             # State changed -> ask planner what to do next.
             return Command(update=state, goto=Send("llm_node", None))
-        else:
-            state["output"].append("timer: no updates yet")
-            # No meaningful state change -> keep waiting without calling planner.
-            return Command(update=state, goto=Send("wait_node", None))
+
+        state["output"].append("timer: no updates yet")
+        # No meaningful state change -> keep waiting without calling planner.
+        return Command(update=state, goto=Send("wait_node", None))
 
     async def tool_node(ctx: Context, tool_input: str) -> None:
         await asyncio.sleep(0.1)

saf-python-sdk/uv.lock

@@ -4,5 +4,5 @@ requires-python = ">=3.10"
 
 [[package]]
 name = "saf-python-sdk"
-version = "0.1.2"
+version = "0.1.4"
 source = { editable = "." }