diff --git a/langgraph-go/advancedgraph/graph.go b/langgraph-go/advancedgraph/graph.go index debea5cf9..19d92820a 100644 --- a/langgraph-go/advancedgraph/graph.go +++ b/langgraph-go/advancedgraph/graph.go @@ -105,13 +105,16 @@ type Context struct { resumeEvent *WaitEvent } -func (c *Context) WaitFor(cond AnyOfCondition) (WaitForResult, 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 waitForResultFromRaw(cond, event), nil + return waitForResultFromRaw(target, event), nil } - return WaitForResult{}, ErrWaitRequested{Condition: cond} + return WaitForResult{}, ErrWaitRequested{Target: target} } func (c *Context) PublishToChannel(channel string, value any) error { @@ -403,16 +406,26 @@ func unwrapResumeInput(input any) (any, *WaitEvent) { return rawArg, &event } -func waitForResultFromRaw(cond AnyOfCondition, event WaitEvent) WaitForResult { +func waitForResultFromRaw(target WaitTarget, event WaitEvent) WaitForResult { + conditions := target.waitConditions() result := WaitForResult{ - Conditions: make([]ConditionResult, len(cond.Conditions)), + 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, raw := range cond.Conditions { - if kind, _ := raw["kind"].(string); kind == "timer" { + for i, cond := range conditions { + if isTimerCondition(cond) { result.Conditions[i] = ConditionResult{Met: true} - break + if target.waitKind() == "any_of" { + break + } } } return result @@ -422,19 +435,18 @@ func waitForResultFromRaw(cond AnyOfCondition, event WaitEvent) WaitForResult { return result } - if event.Channel == "__any_of__" { + 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, raw := range cond.Conditions { - kind, _ := raw["kind"].(string) - if kind != "channel" || cursor >= len(matched) { + for i, cond := range conditions { + channelName, ok := channelNameOfCondition(cond) + if !ok || cursor >= len(matched) { continue } - channelName, _ := raw["channel"].(string) if channelName == matched[cursor].Channel { result.Conditions[i] = ConditionResult{ Met: true, @@ -449,12 +461,11 @@ func waitForResultFromRaw(cond AnyOfCondition, event WaitEvent) WaitForResult { var value any _ = json.Unmarshal(event.Value, &value) - for i, raw := range cond.Conditions { - kind, _ := raw["kind"].(string) - if kind != "channel" { + for i, cond := range conditions { + channelName, ok := channelNameOfCondition(cond) + if !ok { continue } - channelName, _ := raw["channel"].(string) if channelName == event.Channel { result.Conditions[i] = ConditionResult{ Met: true, @@ -467,6 +478,29 @@ func waitForResultFromRaw(cond AnyOfCondition, event WaitEvent) WaitForResult { 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{} diff --git a/langgraph-go/advancedgraph/rust_engine.go b/langgraph-go/advancedgraph/rust_engine.go index adbff2177..086300117 100644 --- a/langgraph-go/advancedgraph/rust_engine.go +++ b/langgraph-go/advancedgraph/rust_engine.go @@ -80,7 +80,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()) } @@ -234,6 +234,35 @@ 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, @@ -309,13 +338,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)) } diff --git a/langgraph-go/advancedgraph/types.go b/langgraph-go/advancedgraph/types.go index 8c56ec297..059fcace4 100644 --- a/langgraph-go/advancedgraph/types.go +++ b/langgraph-go/advancedgraph/types.go @@ -3,10 +3,12 @@ package advancedgraph import ( "encoding/json" "errors" + "fmt" + "reflect" ) type WaitCondition interface { - toAny() map[string]any + json.Marshaler } type ChannelCondition struct { @@ -15,40 +17,128 @@ type ChannelCondition struct { Max int } -func (c ChannelCondition) toAny() map[string]any { +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 map[string]any{ - "kind": "channel", - "channel": c.Channel, - "min": min, - "max": c.Max, - } + 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 { @@ -79,7 +169,7 @@ type Command struct { } type ErrWaitRequested struct { - Condition AnyOfCondition + Target WaitTarget } func (e ErrWaitRequested) Error() string { diff --git a/langgraph-go/stategraph/state_graph.go b/langgraph-go/stategraph/state_graph.go index b553ae09e..af417f3e7 100644 --- a/langgraph-go/stategraph/state_graph.go +++ b/langgraph-go/stategraph/state_graph.go @@ -32,7 +32,7 @@ func (c *Context) Interrupt(name string) (any, error) { if waitReq, ok := ag.AsErrWaitRequested(err); ok { return nil, errInterruptRequested{ Name: name, - Condition: waitReq.Condition, + Condition: waitReq.Target, } } return nil, err @@ -73,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 { @@ -178,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, Min: 1, }) } - return ag.Command{}, ag.ErrWaitRequested{Condition: cond} + return ag.Command{}, ag.ErrWaitRequested{Target: cond} } return ag.Command{}, err } diff --git a/langgraph-go/tests/test_primitives_test.go b/langgraph-go/tests/test_primitives_test.go index 5653be19c..f872f4bd0 100644 --- a/langgraph-go/tests/test_primitives_test.go +++ b/langgraph-go/tests/test_primitives_test.go @@ -291,3 +291,139 @@ func TestAnyOfConsumesAllReadyChannels(t *testing.T) { 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) + } +} diff --git a/rust-core/include/langgraph_rust_core.h b/rust-core/include/langgraph_rust_core.h index f01c3dbe7..f8d4daeec 100644 --- a/rust-core/include/langgraph_rust_core.h +++ b/rust-core/include/langgraph_rust_core.h @@ -20,6 +20,7 @@ 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, const char* stream_name); char* rc_send_custom_stream_event(Engine* ptr, const char* stream_name, const char* value_json); diff --git a/rust-core/src/engine.rs b/rust-core/src/engine.rs index cefb1fef9..d4a56ae8f 100644 --- a/rust-core/src/engine.rs +++ b/rust-core/src/engine.rs @@ -32,6 +32,11 @@ pub struct AnyOfCondition { pub conditions: Vec, } +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct AllOfCondition { + pub conditions: Vec, +} + #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(tag = "condition")] pub enum WaitEvent { @@ -51,6 +56,8 @@ pub enum WaitRequest { Condition { condition: WaitCondition }, #[serde(rename = "any_of")] AnyOf { any_of: AnyOfCondition }, + #[serde(rename = "all_of")] + AllOf { all_of: AllOfCondition }, } pub struct SendPayload { @@ -422,6 +429,7 @@ impl Engine { self.wait_for_any_of_async(&any_of).await } WaitRequest::AnyOf { any_of } => self.wait_for_any_of_async(any_of).await, + WaitRequest::AllOf { all_of } => self.wait_for_all_of_async(all_of).await, } } @@ -438,15 +446,8 @@ impl Engine { } let started = Instant::now(); - let mut min_timer: Option = None; - for cond in &any_of.conditions { - if let WaitCondition::Timer { seconds } = cond { - if *seconds <= 0.0 { - return Err("timer condition must be > 0".to_string()); - } - min_timer = Some(min_timer.map_or(*seconds, |x| x.min(*seconds))); - } - } + let timer_bounds = validate_wait_conditions(&any_of.conditions)?; + let min_timer = timer_bounds.min_seconds; loop { if let Some(event) = self.try_take_any_of_channel_events(any_of)? { @@ -472,75 +473,205 @@ impl Engine { } } + pub async fn wait_for_all_of_async( + &self, + all_of: &AllOfCondition, + ) -> Result { + debug_log(&format!( + "Engine::wait_for_all_of_async(conditions={})", + all_of.conditions.len() + )); + if all_of.conditions.is_empty() { + return Err("all_of requires at least one condition".to_string()); + } + + let started = Instant::now(); + let timer_bounds = validate_wait_conditions(&all_of.conditions)?; + let max_timer = timer_bounds.max_seconds; + let has_channel_condition = all_of + .conditions + .iter() + .any(|c| matches!(c, WaitCondition::Channel { .. })); + + let mut consumed_channel_event: Option = None; + let mut channels_met = !has_channel_condition; + + loop { + if !channels_met { + if let Some(event) = self.try_take_all_of_channel_events(all_of)? { + consumed_channel_event = Some(event); + channels_met = true; + } + } + + let timers_met = if let Some(seconds) = max_timer { + started.elapsed() >= Duration::from_secs_f64(seconds) + } else { + true + }; + + if channels_met && timers_met { + if let Some(event) = consumed_channel_event { + return Ok(event); + } + return Ok(WaitEvent::Timer { + seconds: max_timer.unwrap_or(0.0), + }); + } + + if !channels_met && !timers_met { + let timeout = Duration::from_secs_f64(max_timer.unwrap_or(0.0)); + let elapsed = started.elapsed(); + let remaining = timeout.saturating_sub(elapsed); + tokio::select! { + _ = self.channel_notify.notified() => {} + _ = tokio::time::sleep(remaining) => {} + } + } else if !channels_met { + self.channel_notify.notified().await; + } else { + let timeout = Duration::from_secs_f64(max_timer.unwrap_or(0.0)); + let elapsed = started.elapsed(); + let remaining = timeout.saturating_sub(elapsed); + tokio::time::sleep(remaining).await; + } + } + } + fn try_take_any_of_channel_events( &self, any_of: &AnyOfCondition, ) -> Result, String> { let mut channels = self.channels.lock().expect("channels mutex poisoned"); - let mut consumed_per_channel: HashMap = HashMap::new(); - let mut plans: Vec<(String, usize)> = Vec::new(); - - for cond in &any_of.conditions { - let WaitCondition::Channel { channel, min, max } = cond else { - continue; - }; - if *min < 1 { - return Err("channel condition min must be >= 1".to_string()); - } - if *max != 0 && *max < *min { - return Err("channel condition max must be 0 or >= min".to_string()); - } - - let queue = channels - .get(channel) - .ok_or_else(|| format!("Unknown channel `{channel}`"))?; - let already_planned = consumed_per_channel.get(channel).copied().unwrap_or(0); - let available = queue.len().saturating_sub(already_planned); - if available < *min { - continue; - } - - let take_count = if *max == 0 { *min } else { available.min(*max) }; - plans.push((channel.clone(), take_count)); - consumed_per_channel - .entry(channel.clone()) - .and_modify(|v| *v += take_count) - .or_insert(take_count); - } - + let plans = collect_channel_plans(&channels, &any_of.conditions, false)?; if plans.is_empty() { return Ok(None); } - if plans.len() == 1 { - let (channel, take_count) = &plans[0]; - let queue = channels - .get_mut(channel) - .ok_or_else(|| format!("Unknown channel `{channel}`"))?; - let value = pop_queue_value(queue, *take_count); - return Ok(Some(WaitEvent::Channel { - channel: channel.clone(), - value, - })); - } - - let mut matched = Vec::with_capacity(plans.len()); - for (channel, take_count) in plans { - let queue = channels - .get_mut(&channel) - .ok_or_else(|| format!("Unknown channel `{channel}`"))?; - let value = pop_queue_value(queue, take_count); - matched.push(json!({ - "channel": channel, - "value": value, - })); - } - - Ok(Some(WaitEvent::Channel { - channel: "__any_of__".to_string(), - value: Value::Array(matched), - })) + build_channel_wait_event(&mut channels, plans, "__any_of__").map(Some) } + + fn try_take_all_of_channel_events( + &self, + all_of: &AllOfCondition, + ) -> Result, String> { + let mut channels = self.channels.lock().expect("channels mutex poisoned"); + let plans = collect_channel_plans(&channels, &all_of.conditions, true)?; + if plans.is_empty() { + return Ok(None); + } + build_channel_wait_event(&mut channels, plans, "__all_of__").map(Some) + } +} + +struct TimerBounds { + min_seconds: Option, + max_seconds: Option, +} + +fn validate_wait_conditions(conditions: &[WaitCondition]) -> Result { + let mut min_timer: Option = None; + let mut max_timer: Option = None; + for cond in conditions { + match cond { + WaitCondition::Timer { seconds } => { + if *seconds <= 0.0 { + return Err("timer condition must be > 0".to_string()); + } + min_timer = Some(min_timer.map_or(*seconds, |x| x.min(*seconds))); + max_timer = Some(max_timer.map_or(*seconds, |x| x.max(*seconds))); + } + WaitCondition::Channel { min, max, .. } => { + if *min < 1 { + return Err("channel condition min must be >= 1".to_string()); + } + if *max != 0 && *max < *min { + return Err("channel condition max must be 0 or >= min".to_string()); + } + } + } + } + Ok(TimerBounds { + min_seconds: min_timer, + max_seconds: max_timer, + }) +} + +fn collect_channel_plans( + channels: &HashMap>, + conditions: &[WaitCondition], + require_all_channels: bool, +) -> Result, String> { + let mut consumed_per_channel: HashMap = HashMap::new(); + let mut plans: Vec<(String, usize)> = Vec::new(); + let mut channel_condition_count = 0usize; + + for cond in conditions { + let WaitCondition::Channel { channel, min, max } = cond else { + continue; + }; + channel_condition_count += 1; + let queue = channels + .get(channel) + .ok_or_else(|| format!("Unknown channel `{channel}`"))?; + let already_planned = consumed_per_channel.get(channel).copied().unwrap_or(0); + let available = queue.len().saturating_sub(already_planned); + if available < *min { + if require_all_channels { + return Ok(Vec::new()); + } + continue; + } + + let take_count = if *max == 0 { *min } else { available.min(*max) }; + plans.push((channel.clone(), take_count)); + consumed_per_channel + .entry(channel.clone()) + .and_modify(|v| *v += take_count) + .or_insert(take_count); + } + + if require_all_channels && channel_condition_count > 0 && plans.len() != channel_condition_count + { + return Ok(Vec::new()); + } + + Ok(plans) +} + +fn build_channel_wait_event( + channels: &mut HashMap>, + plans: Vec<(String, usize)>, + aggregate_channel_name: &str, +) -> Result { + if plans.len() == 1 { + let (channel, take_count) = &plans[0]; + let queue = channels + .get_mut(channel) + .ok_or_else(|| format!("Unknown channel `{channel}`"))?; + let value = pop_queue_value(queue, *take_count); + return Ok(WaitEvent::Channel { + channel: channel.clone(), + value, + }); + } + + let mut matched = Vec::with_capacity(plans.len()); + for (channel, take_count) in plans { + let queue = channels + .get_mut(&channel) + .ok_or_else(|| format!("Unknown channel `{channel}`"))?; + let value = pop_queue_value(queue, take_count); + matched.push(json!({ + "channel": channel, + "value": value, + })); + } + + Ok(WaitEvent::Channel { + channel: aggregate_channel_name.to_string(), + value: Value::Array(matched), + }) } fn pop_queue_value(queue: &mut VecDeque, take_count: usize) -> Value { diff --git a/rust-core/src/lib_c.rs b/rust-core/src/lib_c.rs index 20e3dbaf9..b0ff25c9c 100644 --- a/rust-core/src/lib_c.rs +++ b/rust-core/src/lib_c.rs @@ -1,6 +1,6 @@ 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::ffi::{CStr, CString}; @@ -144,6 +144,39 @@ 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`. diff --git a/rust-core/src/lib_py.rs b/rust-core/src/lib_py.rs index d9e56fe6f..6c4301111 100644 --- a/rust-core/src/lib_py.rs +++ b/rust-core/src/lib_py.rs @@ -1,7 +1,7 @@ #[cfg(feature = "python-bindings")] use crate::engine::{ - run_graph_with_callback, run_loop_block_on, AnyOfCondition, Engine, NodeExecResult, - NodeOutcome, SendPayload, WaitCondition, WaitEvent, WaitRequest, + 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; @@ -92,6 +92,12 @@ impl PyRustEngine { json_string_to_py_obj(py, &event_json) } + fn wait_all_of_obj(&self, py: Python<'_>, all_of_payload: Py) -> PyResult> { + 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 { let any_of: AnyOfCondition = serde_json::from_str(any_of_json) .map_err(|e| PyValueError::new_err(format!("Invalid any_of JSON: {e}")))?; @@ -101,6 +107,15 @@ impl PyRustEngine { .map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}"))) } + fn _wait_all_of_json(&self, all_of_json: &str) -> PyResult { + 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 { let cond: WaitCondition = serde_json::from_str(cond_json) .map_err(|e| PyValueError::new_err(format!("Invalid condition JSON: {e}")))?; diff --git a/saf-python-sdk/Makefile b/saf-python-sdk/Makefile index a95abb0e5..a09541b14 100644 --- a/saf-python-sdk/Makefile +++ b/saf-python-sdk/Makefile @@ -2,7 +2,7 @@ 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: 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 227f891ae..2e62ea7cb 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 @@ -1,5 +1,6 @@ from .state import ( AdvancedStateGraph, + AllOfCondition, AnyOfCondition, ChannelCondition, ConditionResult, @@ -8,6 +9,7 @@ from .state import ( GraphRunHandler, TimerCondition, WaitForResult, + all_of, any_of, channel_condition, timer_condition, @@ -21,10 +23,12 @@ __all__ = [ "ChannelCondition", "TimerCondition", "AnyOfCondition", + "AllOfCondition", "ConditionResult", "WaitForResult", "channel_condition", "timer_condition", "any_of", + "all_of", ] 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 7d2aec017..948f74629 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 @@ -40,6 +40,11 @@ class AnyOfCondition: conditions: tuple[WaitCondition, ...] +@dataclass(frozen=True) +class AllOfCondition: + conditions: tuple[WaitCondition, ...] + + @dataclass(frozen=True) class ConditionResult: met: bool @@ -209,7 +214,9 @@ class Context: def __init__(self, run: _GraphEngineRun) -> None: self._run = run - async def wait_for(self, target: WaitCondition | AnyOfCondition) -> WaitForResult: + async def wait_for( + self, target: WaitCondition | AnyOfCondition | AllOfCondition + ) -> WaitForResult: resumed = self._run._consume_resume_event(target) if resumed is not None: return resumed @@ -316,7 +323,9 @@ class _GraphEngineRun: def publish_nowait(self, channel: str, value: Any) -> None: self._publish_sync(channel, value) - async def wait_for(self, target: WaitCondition | AnyOfCondition) -> WaitForResult: + async def wait_for( + self, target: WaitCondition | AnyOfCondition | AllOfCondition + ) -> WaitForResult: if isinstance(target, ChannelCondition): value = await self._wait_for_channel_values( target.channel, min=target.min, max=target.max @@ -341,6 +350,9 @@ class _GraphEngineRun: if isinstance(target, AnyOfCondition): 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( @@ -375,6 +387,19 @@ 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) @@ -426,7 +451,7 @@ class _GraphEngineRun: self._local.resume_event = event def _consume_resume_event( - self, target: WaitCondition | AnyOfCondition + self, target: WaitCondition | AnyOfCondition | AllOfCondition ) -> WaitForResult | None: event = cast(dict[str, Any] | None, getattr(self._local, "resume_event", None)) if event is None: @@ -527,6 +552,12 @@ 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 @@ -534,7 +565,7 @@ def _normalize_channel_values(value: Any) -> list[Any]: def _wait_for_result_from_resume_event( - target: WaitCondition | AnyOfCondition, event: dict[str, Any] + target: WaitCondition | AnyOfCondition | AllOfCondition, event: dict[str, Any] ) -> WaitForResult: if isinstance(target, ChannelCondition): return WaitForResult( @@ -548,6 +579,8 @@ def _wait_for_result_from_resume_event( ) 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) @@ -603,6 +636,52 @@ def _wait_for_result_from_any_of_event( 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 { @@ -616,7 +695,9 @@ def _condition_to_rust(condition: WaitCondition) -> dict[str, Any]: 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", @@ -624,6 +705,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)} diff --git a/saf-python-sdk/tests/advanced-graph/test_primitives.py b/saf-python-sdk/tests/advanced-graph/test_primitives.py index e6b131c40..ff11b70ed 100644 --- a/saf-python-sdk/tests/advanced-graph/test_primitives.py +++ b/saf-python-sdk/tests/advanced-graph/test_primitives.py @@ -4,8 +4,10 @@ from typing_extensions import TypedDict 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 @@ -146,3 +148,71 @@ async def test_any_of_consumes_all_ready_channels() -> None: ] 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" +