diff --git a/langgraph-go/advancedgraph/graph.go b/langgraph-go/advancedgraph/graph.go index 5744a1825..315d07ba5 100644 --- a/langgraph-go/advancedgraph/graph.go +++ b/langgraph-go/advancedgraph/graph.go @@ -10,8 +10,28 @@ 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 @@ -25,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)) } @@ -47,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 } @@ -58,24 +80,24 @@ func (g *AdvancedStateGraph[StateT]) AddCustomOutputStream(name string) { g.customStreams = append(g.customStreams, name) } -func (g *AdvancedStateGraph[StateT]) AddEntryNode(fn any) string { +func (g *AdvancedStateGraph[StateT]) AddEntryNode(fn any, nodeOption ...NodeOption) string { name := NodeName(fn) - return g.AddEntryNodeAs(name, fn) + return g.AddEntryNodeAs(name, fn, nodeOption...) } -func (g *AdvancedStateGraph[StateT]) AddEntryNodeAs(name string, fn any) string { - name = g.AddNodeAs(name, fn) +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 } @@ -83,6 +105,7 @@ 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, @@ -93,6 +116,7 @@ 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 @@ -211,6 +235,7 @@ func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT, str streamModeForRun, initialState, initialInput, + g.nodeLockedFields(), func(node string, nodeInput any, fallbackState map[string]any) (Command, error) { fn, ok := g.nodes[node] if !ok { @@ -243,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 { diff --git a/langgraph-go/advancedgraph/rust_engine.go b/langgraph-go/advancedgraph/rust_engine.go index 086300117..cf7a9a3ba 100644 --- a/langgraph-go/advancedgraph/rust_engine.go +++ b/langgraph-go/advancedgraph/rust_engine.go @@ -6,6 +6,28 @@ package advancedgraph #include "langgraph_rust_core.h" #include 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" @@ -269,6 +291,7 @@ func (e *RustEngine) RunGraph( 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) @@ -279,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) @@ -291,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)) } @@ -298,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) diff --git a/langgraph-go/tests/test_primitives_test.go b/langgraph-go/tests/test_primitives_test.go index b827c3a43..1d46d4ec6 100644 --- a/langgraph-go/tests/test_primitives_test.go +++ b/langgraph-go/tests/test_primitives_test.go @@ -2,13 +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 { @@ -492,3 +496,94 @@ func TestIsResumeAvoidsDuplicateSideEffects(t *testing.T) { 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]) + } +} diff --git a/rust-core/include/langgraph_rust_core.h b/rust-core/include/langgraph_rust_core.h index f8d4daeec..1f863b3aa 100644 --- a/rust-core/include/langgraph_rust_core.h +++ b/rust-core/include/langgraph_rust_core.h @@ -32,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 ); diff --git a/rust-core/src/engine.rs b/rust-core/src/engine.rs index d4a56ae8f..906dd005b 100644 --- a/rust-core/src/engine.rs +++ b/rust-core/src/engine.rs @@ -1,7 +1,7 @@ use serde::{Deserialize, Serialize}; use serde_json::json; use serde_json::Value; -use std::collections::{HashMap, VecDeque}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::env; use std::future::Future; use std::sync::mpsc; @@ -723,12 +723,19 @@ enum SchedulerEvent { } struct NodeExecution { + execution_id: u64, node: String, arg: A, outcome: NodeOutcome, } +struct PendingNodeExecution { + node: String, + arg: A, +} + fn spawn_node_task( + execution_id: u64, node: String, arg: A, state_snapshot: State, @@ -745,6 +752,7 @@ where let node_for_result = node.clone(); let arg_for_result = arg.clone(); let result = callback(node, arg, state_snapshot).map(|outcome| NodeExecution { + execution_id, node: node_for_result, arg: arg_for_result, outcome, @@ -753,6 +761,113 @@ where }) } +fn try_spawn_or_enqueue( + node: String, + arg: A, + pending: &mut VecDeque>, + locked_fields: &mut HashSet, + execution_locks: &mut HashMap>, + next_execution_id: &mut u64, + active: &mut usize, + state_for_spawn: &Arc>, + tx_for_spawn: &tokio_mpsc::UnboundedSender>, + callback: &Arc, + locked_fields_for_node: &Arc, +) -> Result<(), String> +where + State: Clone + Send + 'static, + U: Send + 'static, + A: Clone + Send + 'static, + F: Fn(String, A, State) -> Result, String> + Send + Sync + 'static, + FLock: Fn(&str) -> Vec + Send + Sync + 'static, +{ + let requested = locked_fields_for_node(node.as_str()) + .into_iter() + .filter(|field| !field.is_empty()) + .collect::>(); + let conflict = requested.iter().any(|field| locked_fields.contains(field)); + if conflict { + debug_log(&format!( + "enqueue node={} due to lock conflict (requested={:?})", + node, requested + )); + pending.push_back(PendingNodeExecution { node, arg }); + return Ok(()); + } + for field in &requested { + locked_fields.insert(field.clone()); + } + let execution_id = *next_execution_id; + *next_execution_id = next_execution_id.saturating_add(1); + execution_locks.insert(execution_id, requested); + *active = active.saturating_add(1); + let snapshot = state_for_spawn + .lock() + .expect("state mutex poisoned") + .clone(); + spawn_node_task( + execution_id, + node, + arg, + snapshot, + tx_for_spawn.clone(), + Arc::clone(callback), + ) +} + +fn drain_pending( + pending: &mut VecDeque>, + locked_fields: &mut HashSet, + execution_locks: &mut HashMap>, + next_execution_id: &mut u64, + active: &mut usize, + state_for_spawn: &Arc>, + tx_for_spawn: &tokio_mpsc::UnboundedSender>, + callback: &Arc, + locked_fields_for_node: &Arc, +) -> Result<(), String> +where + State: Clone + Send + 'static, + U: Send + 'static, + A: Clone + Send + 'static, + F: Fn(String, A, State) -> Result, String> + Send + Sync + 'static, + FLock: Fn(&str) -> Vec + Send + Sync + 'static, +{ + if pending.is_empty() { + return Ok(()); + } + loop { + let mut progressed = false; + let round = pending.len(); + for _ in 0..round { + let Some(item) = pending.pop_front() else { + continue; + }; + let active_before = *active; + try_spawn_or_enqueue( + item.node, + item.arg, + pending, + locked_fields, + execution_locks, + next_execution_id, + active, + state_for_spawn, + tx_for_spawn, + callback, + locked_fields_for_node, + )?; + if *active > active_before { + progressed = true; + } + } + if !progressed { + break; + } + } + Ok(()) +} + pub async fn run_graph_with_callback( entry_point: String, finish_point: String, @@ -762,6 +877,7 @@ pub async fn run_graph_with_callback( callback: FCallback, merge_update: FMerge, wrap_resume_arg: FWrap, + locked_fields_for_node: impl Fn(&str) -> Vec + Send + Sync + 'static, ) -> Result where State: Clone + Send + 'static, @@ -774,23 +890,31 @@ where let callback = Arc::new(callback); let merge_update = Arc::new(merge_update); let wrap_resume_arg = Arc::new(wrap_resume_arg); + let locked_fields_for_node = Arc::new(locked_fields_for_node); let (tx, mut rx) = tokio_mpsc::unbounded_channel::>(); let state = Arc::new(StdMutex::new(initial_state)); let tx_for_spawn = tx.clone(); let state_for_spawn = Arc::clone(&state); - let mut active: usize = 1; + let mut active: usize = 0; let mut waiting: usize = 0; + let mut next_execution_id: u64 = 1; + let mut pending: VecDeque> = VecDeque::new(); + let mut locked_fields: HashSet = HashSet::new(); + let mut execution_locks: HashMap> = HashMap::new(); - spawn_node_task( + try_spawn_or_enqueue( entry_point, initial_input, - state_for_spawn - .lock() - .expect("state mutex poisoned") - .clone(), - tx_for_spawn.clone(), - Arc::clone(&callback), + &mut pending, + &mut locked_fields, + &mut execution_locks, + &mut next_execution_id, + &mut active, + &state_for_spawn, + &tx_for_spawn, + &callback, + &locked_fields_for_node, )?; while active > 0 || waiting > 0 { @@ -802,6 +926,11 @@ where SchedulerEvent::Node(result) => { active = active.saturating_sub(1); let exec = result?; + if let Some(fields) = execution_locks.remove(&exec.execution_id) { + for field in fields { + locked_fields.remove(&field); + } + } match exec.outcome { NodeOutcome::Completed(node_result) => { let mut guard = state.lock().expect("state mutex poisoned"); @@ -813,17 +942,18 @@ where } for send in node_result.sends { - active += 1; - let snapshot = state_for_spawn - .lock() - .expect("state mutex poisoned") - .clone(); - spawn_node_task( + try_spawn_or_enqueue( send.node, send.arg, - snapshot, - tx_for_spawn.clone(), - Arc::clone(&callback), + &mut pending, + &mut locked_fields, + &mut execution_locks, + &mut next_execution_id, + &mut active, + &state_for_spawn, + &tx_for_spawn, + &callback, + &locked_fields_for_node, )?; } } @@ -849,22 +979,34 @@ where } SchedulerEvent::Resume { node, arg, event } => { waiting = waiting.saturating_sub(1); - active += 1; - let snapshot = state_for_spawn - .lock() - .expect("state mutex poisoned") - .clone(); let resume_arg = wrap_resume_arg(arg, event)?; - spawn_node_task( + try_spawn_or_enqueue( node, resume_arg, - snapshot, - tx_for_spawn.clone(), - Arc::clone(&callback), + &mut pending, + &mut locked_fields, + &mut execution_locks, + &mut next_execution_id, + &mut active, + &state_for_spawn, + &tx_for_spawn, + &callback, + &locked_fields_for_node, )?; } SchedulerEvent::WaitError(e) => return Err(e), } + drain_pending( + &mut pending, + &mut locked_fields, + &mut execution_locks, + &mut next_execution_id, + &mut active, + &state_for_spawn, + &tx_for_spawn, + &callback, + &locked_fields_for_node, + )?; } let final_state = state.lock().expect("state mutex poisoned").clone(); @@ -878,6 +1020,7 @@ pub async fn run_graph_json_with_callback( initial_input: Value, engine: Engine, callback: F, + locked_fields_by_node: HashMap>, ) -> Result where F: Fn(String, Value, Value) -> Result, String> @@ -902,6 +1045,7 @@ where "__lg_resume_event__": event, })) }, + move |node: &str| locked_fields_by_node.get(node).cloned().unwrap_or_default(), ) .await } diff --git a/rust-core/src/lib_c.rs b/rust-core/src/lib_c.rs index b0ff25c9c..823e578c1 100644 --- a/rust-core/src/lib_c.rs +++ b/rust-core/src/lib_c.rs @@ -3,6 +3,7 @@ use crate::engine::{ 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; @@ -304,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, ) -> *mut c_char { @@ -353,6 +355,22 @@ 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> = 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 Some(mode) = stream_mode.as_deref() { if let Err(e) = (*ptr).start_stream(Some(mode)) { @@ -404,6 +422,7 @@ pub unsafe extern "C" fn rc_run_graph_json( initial_input, run_engine.clone(), callback_wrapper, + node_locked_fields, ) .await; run_engine.close_all_streams(); diff --git a/rust-core/src/lib_py.rs b/rust-core/src/lib_py.rs index 6c4301111..86e3a8e27 100644 --- a/rust-core/src/lib_py.rs +++ b/rust-core/src/lib_py.rs @@ -14,6 +14,8 @@ 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")] @@ -165,7 +167,7 @@ impl PyRustEngine { self.inner.close_all_streams(); } - #[pyo3(signature = (entry_point, finish_point, initial_state, callback, stream_mode=None))] + #[pyo3(signature = (entry_point, finish_point, initial_state, callback, stream_mode=None, node_locked_fields=None))] fn run_graph_py( &self, py: Python<'_>, @@ -174,6 +176,7 @@ impl PyRustEngine { initial_state: Py, callback: Py, stream_mode: Option<&str>, + node_locked_fields: Option>, ) -> PyResult> { if let Some(mode) = stream_mode { self.inner @@ -187,6 +190,15 @@ impl PyRustEngine { 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> = + 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 || { @@ -229,6 +241,7 @@ impl PyRustEngine { |arg: Arc>, event: WaitEvent| -> Result>, String> { wrap_resume_arg(arg.as_ref(), &event).map(Arc::new) }, + move |node: &str| locked_fields_by_node.get(node).cloned().unwrap_or_default(), )) }); diff --git a/saf-python-sdk/python/saf_python_sdk/advanced_graph/__init__.py b/saf-python-sdk/python/saf_python_sdk/advanced_graph/__init__.py index 2e62ea7cb..628408086 100644 --- a/saf-python-sdk/python/saf_python_sdk/advanced_graph/__init__.py +++ b/saf-python-sdk/python/saf_python_sdk/advanced_graph/__init__.py @@ -7,6 +7,7 @@ from .state import ( CompiledGraphEngine, Context, GraphRunHandler, + NodeStateOption, TimerCondition, WaitForResult, all_of, @@ -20,6 +21,7 @@ __all__ = [ "CompiledGraphEngine", "Context", "GraphRunHandler", + "NodeStateOption", "ChannelCondition", "TimerCondition", "AnyOfCondition", diff --git a/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py b/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py index b591c7f87..6ebd58bc6 100644 --- a/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py +++ b/saf-python-sdk/python/saf_python_sdk/advanced_graph/state.py @@ -23,6 +23,11 @@ class _ChannelSpec: typ: Any +@dataclass(frozen=True) +class NodeStateOption: + locked_fields: tuple[str, ...] = () + + @dataclass(frozen=True) class ChannelCondition: channel: str @@ -97,6 +102,7 @@ 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 @@ -106,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): @@ -121,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: @@ -137,13 +146,23 @@ class AdvancedStateGraph(Generic[StateT]): 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]) -> str: - node_name = self.add_node(node) + 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 @@ -164,6 +183,7 @@ 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, @@ -178,12 +198,14 @@ 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 @@ -200,6 +222,7 @@ class CompiledGraphEngine(Generic[StateT]): 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, @@ -273,12 +296,14 @@ 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() @@ -312,6 +337,7 @@ class _GraphEngineRun: initial_state, self._execute_node_for_rust, None, + _node_locked_fields_payload(self._node_options), ) finally: self._stream_ready.set() @@ -524,6 +550,43 @@ def _normalize_goto(goto: Any, *, default_input: Any) -> list[Send]: return [] +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") diff --git a/saf-python-sdk/tests/advanced-graph/test_primitives.py b/saf-python-sdk/tests/advanced-graph/test_primitives.py index 2d967a7c3..10e1b205e 100644 --- a/saf-python-sdk/tests/advanced-graph/test_primitives.py +++ b/saf-python-sdk/tests/advanced-graph/test_primitives.py @@ -1,3 +1,5 @@ +import asyncio +import time import pytest from typing_extensions import TypedDict @@ -248,3 +250,56 @@ async def test_is_resume_avoids_duplicate_side_effects() -> None: 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']}" +