allof

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:

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",
 ]
 

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)}
 
 

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"
+