From c0bd2dbc07f8d73959e86abed4a8aab592ae07cd Mon Sep 17 00:00:00 2001 From: Quanzheng Long Date: Wed, 18 Mar 2026 11:41:05 -0700 Subject: [PATCH] anyofall --- langgraph-go/advancedgraph/graph.go | 80 +++++++++++- langgraph-go/advancedgraph/types.go | 10 ++ langgraph-go/stategraph/state_graph.go | 15 ++- langgraph-go/tests/sub_agents_test.go | 27 ++-- langgraph-go/tests/test_primitives_test.go | 110 +++++++++++++++- rust-core/src/engine.rs | 98 ++++++++++++-- .../saf_python_sdk/advanced_graph/__init__.py | 4 + .../saf_python_sdk/advanced_graph/state.py | 120 ++++++++++++++++-- .../tests/advanced-graph/test_primitives.py | 62 ++++++++- .../tests/advanced-graph/test_sub_agents.py | 40 ++++-- 10 files changed, 505 insertions(+), 61 deletions(-) diff --git a/langgraph-go/advancedgraph/graph.go b/langgraph-go/advancedgraph/graph.go index bf4eabf4d..904b97d87 100644 --- a/langgraph-go/advancedgraph/graph.go +++ b/langgraph-go/advancedgraph/graph.go @@ -98,13 +98,13 @@ type Context struct { resumeEvent *WaitEvent } -func (c *Context) WaitFor(cond AnyOfCondition) (WaitEvent, error) { +func (c *Context) WaitFor(cond AnyOfCondition) (WaitForResult, error) { if c.resumeEvent != nil { event := *c.resumeEvent c.resumeEvent = nil - return event, nil + return waitForResultFromRaw(cond, event), nil } - return WaitEvent{}, ErrWaitRequested{Condition: cond} + return WaitForResult{}, ErrWaitRequested{Condition: cond} } func (c *Context) PublishToChannel(channel string, value any) error { @@ -390,3 +390,77 @@ func unwrapResumeInput(input any) (any, *WaitEvent) { } return rawArg, &event } + +func waitForResultFromRaw(cond AnyOfCondition, event WaitEvent) WaitForResult { + result := WaitForResult{ + Conditions: make([]ConditionResult, len(cond.Conditions)), + } + + if event.Condition == "timer" { + for i, raw := range cond.Conditions { + if kind, _ := raw["kind"].(string); kind == "timer" { + result.Conditions[i] = ConditionResult{Met: true} + break + } + } + return result + } + + if event.Condition != "channel" { + return result + } + + if event.Channel == "__any_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) { + continue + } + channelName, _ := raw["channel"].(string) + 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, raw := range cond.Conditions { + kind, _ := raw["kind"].(string) + if kind != "channel" { + continue + } + channelName, _ := raw["channel"].(string) + if channelName == event.Channel { + result.Conditions[i] = ConditionResult{ + Met: true, + ChannelName: channelName, + Values: toValues(value), + } + break + } + } + return result +} + +func toValues(value any) []any { + if value == nil { + return []any{} + } + if vals, ok := value.([]any); ok { + return vals + } + return []any{value} +} diff --git a/langgraph-go/advancedgraph/types.go b/langgraph-go/advancedgraph/types.go index 6c2d9064b..8c56ec297 100644 --- a/langgraph-go/advancedgraph/types.go +++ b/langgraph-go/advancedgraph/types.go @@ -58,6 +58,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 diff --git a/langgraph-go/stategraph/state_graph.go b/langgraph-go/stategraph/state_graph.go index 7d56c3f80..b553ae09e 100644 --- a/langgraph-go/stategraph/state_graph.go +++ b/langgraph-go/stategraph/state_graph.go @@ -37,14 +37,23 @@ func (c *Context) Interrupt(name string) (any, error) { } 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 diff --git a/langgraph-go/tests/sub_agents_test.go b/langgraph-go/tests/sub_agents_test.go index 3516be3e5..73a505d2f 100644 --- a/langgraph-go/tests/sub_agents_test.go +++ b/langgraph-go/tests/sub_agents_test.go @@ -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 } diff --git a/langgraph-go/tests/test_primitives_test.go b/langgraph-go/tests/test_primitives_test.go index b6319e468..5653be19c 100644 --- a/langgraph-go/tests/test_primitives_test.go +++ b/langgraph-go/tests/test_primitives_test.go @@ -1,7 +1,6 @@ package tests import ( - "encoding/json" "fmt" "testing" @@ -143,7 +142,7 @@ func (w *primitiveWorkflow) startWaitBatchNode(ctx *ag.Context, _ any, state pri } func (w *primitiveWorkflow) waitBatchNode(ctx *ag.Context, _ any, state primitiveState) (ag.Command, error) { - event, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{ + result, err := ctx.WaitFor(ag.AnyOf(ag.ChannelCondition{ Channel: "events", Min: 2, Max: 4, @@ -151,9 +150,13 @@ func (w *primitiveWorkflow) waitBatchNode(ctx *ag.Context, _ any, state primitiv if err != nil { return ag.Command{}, err } - var values []string - if err := json.Unmarshal(event.Value, &values); 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 @@ -191,3 +194,100 @@ func TestChannelWaitRespectsMaxM(t *testing.T) { 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) + } +} diff --git a/rust-core/src/engine.rs b/rust-core/src/engine.rs index 9caf57235..fe0a54a18 100644 --- a/rust-core/src/engine.rs +++ b/rust-core/src/engine.rs @@ -1,4 +1,5 @@ use serde::{Deserialize, Serialize}; +use serde_json::json; use serde_json::Value; use std::collections::{HashMap, VecDeque}; use std::env; @@ -432,18 +433,8 @@ impl Engine { } loop { - for cond in &any_of.conditions { - if let WaitCondition::Channel { channel, min, max } = cond { - 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()); - } - if let Some(event) = self.try_take_channel_event(channel, *min, *max)? { - return Ok(event); - } - } + if let Some(event) = self.try_take_any_of_channel_events(any_of)? { + return Ok(event); } if let Some(seconds) = min_timer { @@ -507,6 +498,89 @@ impl Engine { value: serde_json::Value::Array(values), })) } + + 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); + } + + 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), + })) + } +} + +fn pop_queue_value(queue: &mut VecDeque, take_count: usize) -> Value { + if take_count == 1 { + return queue.pop_front().unwrap_or(Value::Null); + } + let mut values = Vec::with_capacity(take_count); + for _ in 0..take_count { + if let Some(v) = queue.pop_front() { + values.push(v); + } + } + Value::Array(values) } #[derive(Debug, Deserialize)] 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 d0cd77021..227f891ae 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 @@ -2,10 +2,12 @@ from .state import ( AdvancedStateGraph, AnyOfCondition, ChannelCondition, + ConditionResult, CompiledGraphEngine, Context, GraphRunHandler, TimerCondition, + WaitForResult, any_of, channel_condition, timer_condition, @@ -19,6 +21,8 @@ __all__ = [ "ChannelCondition", "TimerCondition", "AnyOfCondition", + "ConditionResult", + "WaitForResult", "channel_condition", "timer_condition", "any_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 4071551d7..fe8c83019 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,18 @@ class AnyOfCondition: 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() @@ -183,7 +195,7 @@ 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) -> WaitForResult: resumed = self._run._consume_resume_event(target) if resumed is not None: return resumed @@ -287,25 +299,31 @@ 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) -> WaitForResult: if isinstance(target, ChannelCondition): value = await self._wait_for_channel_values( target.channel, min=target.min, max=target.max ) - return { - "condition": "channel", - "channel": target.channel, - "value": value, - } + 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) raise ValueError(f"Unsupported wait condition type: {type(target)!r}") async def _wait_for_channel_values( @@ -327,7 +345,7 @@ class _GraphEngineRun: ) 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 = { @@ -390,12 +408,14 @@ 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 _consume_resume_event( + self, target: WaitCondition | AnyOfCondition + ) -> 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 @@ -490,6 +510,82 @@ def any_of(*conditions: WaitCondition) -> AnyOfCondition: return AnyOfCondition(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, 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)]) + 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 _condition_to_rust(condition: WaitCondition) -> dict[str, Any]: if isinstance(condition, ChannelCondition): return { diff --git a/saf-python-sdk/tests/advanced-graph/test_primitives.py b/saf-python-sdk/tests/advanced-graph/test_primitives.py index 9d0d61ee4..e6b131c40 100644 --- a/saf-python-sdk/tests/advanced-graph/test_primitives.py +++ b/saf-python-sdk/tests/advanced-graph/test_primitives.py @@ -1,7 +1,12 @@ import pytest from typing_extensions import TypedDict -from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context, channel_condition +from saf_python_sdk.advanced_graph import ( + AdvancedStateGraph, + Context, + any_of, + channel_condition, +) from saf_python_sdk.types import Command, Send pytestmark = pytest.mark.anyio @@ -77,8 +82,8 @@ async def test_channel_wait_respects_max_m() -> None: return Command(update=state, goto=Send("wait_node", None)) async def wait_node(ctx: Context, _input: None, state: PrimitiveState) -> dict[str, object]: - event = await ctx.wait_for(channel_condition("events", min=2, max=4)) - values = event["value"] + 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"} @@ -90,3 +95,54 @@ async def test_channel_wait_respects_max_m() -> None: 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" + diff --git a/saf-python-sdk/tests/advanced-graph/test_sub_agents.py b/saf-python-sdk/tests/advanced-graph/test_sub_agents.py index 73cc0ea57..51ef63713 100644 --- a/saf-python-sdk/tests/advanced-graph/test_sub_agents.py +++ b/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)