anyof

libs/langgraph/langgraph/advanced_graph/__init__.py

@@ -1,13 +1,25 @@
 from langgraph.advanced_graph.state import (
     AdvancedStateGraph,
+    AnyOfCondition,
+    ChannelCondition,
     CompiledGraphEngine,
+    TimerCondition,
+    any_of,
+    channel_condition,
     publish_to_channel,
+    timer_condition,
     wait_for,
 )
 
 __all__ = (
     "AdvancedStateGraph",
+    "AnyOfCondition",
+    "ChannelCondition",
     "CompiledGraphEngine",
+    "TimerCondition",
+    "any_of",
+    "channel_condition",
     "publish_to_channel",
+    "timer_condition",
     "wait_for",
 )

libs/langgraph/langgraph/advanced_graph/state.py

@@ -5,6 +5,7 @@ import contextvars
 import inspect
 from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass
+from datetime import timedelta
 from typing import Any, Generic, TypeVar, cast
 
 from langgraph.types import Command, Send
@@ -22,6 +23,25 @@ class _ChannelSpec:
     maxsize: int
 
 
+@dataclass(frozen=True)
+class ChannelCondition:
+    channel: str
+    n: int = 1
+
+
+@dataclass(frozen=True)
+class TimerCondition:
+    seconds: float
+
+
+@dataclass(frozen=True)
+class AnyOfCondition:
+    conditions: tuple[WaitCondition, ...]
+
+
+WaitCondition = ChannelCondition | TimerCondition | AnyOfCondition
+
+
 class AdvancedStateGraph(Generic[StateT]):
     """Experimental in-memory graph engine with async channels."""
 
@@ -199,7 +219,24 @@ class _GraphEngineRun:
         queue = self._get_channel(channel)
         queue.put_nowait(value)
 
-    async def wait_for(self, channel: str, n: int = 1) -> Any:
+    async def wait_for(self, target: str | WaitCondition, n: int = 1) -> Any:
+        if isinstance(target, str):
+            return await self._wait_for_channel_values(target, n=n)
+        if isinstance(target, ChannelCondition):
+            value = await self._wait_for_channel_values(target.channel, n=target.n)
+            return {
+                "condition": "channel",
+                "channel": target.channel,
+                "value": value,
+            }
+        if isinstance(target, TimerCondition):
+            await asyncio.sleep(target.seconds)
+            return {"condition": "timer", "seconds": target.seconds}
+        if isinstance(target, AnyOfCondition):
+            return await self._wait_for_any_of(target)
+        raise ValueError(f"Unsupported wait condition type: {type(target)!r}")
+
+    async def _wait_for_channel_values(self, channel: str, n: int) -> Any:
         if n < 1:
             raise ValueError("wait_for count `n` must be >= 1")
         queue = self._get_channel(channel)
@@ -210,6 +247,21 @@ class _GraphEngineRun:
             values.append(await queue.get())
         return values
 
+    async def _wait_for_any_of(self, condition: AnyOfCondition) -> Any:
+        if not condition.conditions:
+            raise ValueError("any_of() requires at least one condition")
+
+        tasks = [
+            asyncio.create_task(self.wait_for(inner_condition, n=1))
+            for inner_condition in condition.conditions
+        ]
+        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+        for task in pending:
+            task.cancel()
+        await asyncio.gather(*pending, return_exceptions=True)
+        first = done.pop()
+        return first.result()
+
     def _get_channel(self, channel: str) -> asyncio.Queue[Any]:
         if channel not in self._channels:
             raise ValueError(f"Unknown channel `{channel}`")
@@ -321,11 +373,11 @@ def _normalize_goto(goto: Any, *, default_arg: Any) -> list[Send]:
     return []
 
 
-async def wait_for(channel: str, n: int = 1) -> Any:
+async def wait_for(target: str | WaitCondition, n: int = 1) -> Any:
     run = _CURRENT_RUN.get()
     if run is None:
         raise RuntimeError("wait_for() can only be used inside advanced_graph nodes")
-    return await run.wait_for(channel, n=n)
+    return await run.wait_for(target, n=n)
 
 
 def publish_to_channel(channel: str, value: Any) -> None:
@@ -337,6 +389,45 @@ def publish_to_channel(channel: str, value: Any) -> None:
     run.publish_nowait(channel, value)
 
 
+def channel_condition(channel: str, n: int = 1) -> ChannelCondition:
+    if n < 1:
+        raise ValueError("channel_condition `n` must be >= 1")
+    return ChannelCondition(channel=channel, n=n)
+
+
+def timer_condition(
+    timeout: float | timedelta | None = None,
+    *,
+    seconds: float | None = None,
+    minutes: float | None = None,
+) -> TimerCondition:
+    if timeout is not None and (seconds is not None or minutes is not None):
+        raise ValueError(
+            "Provide either `timeout` or named `seconds`/`minutes`, not both"
+        )
+
+    if isinstance(timeout, timedelta):
+        resolved_seconds = timeout.total_seconds()
+    elif isinstance(timeout, (int, float)):
+        resolved_seconds = float(timeout)
+    else:
+        resolved_seconds = 0.0
+        if seconds is not None:
+            resolved_seconds += float(seconds)
+        if minutes is not None:
+            resolved_seconds += float(minutes) * 60.0
+
+    if resolved_seconds <= 0:
+        raise ValueError("timer_condition must be greater than 0 seconds")
+    return TimerCondition(seconds=resolved_seconds)
+
+
+def any_of(*conditions: WaitCondition) -> AnyOfCondition:
+    if not conditions:
+        raise ValueError("any_of() requires at least one condition")
+    return AnyOfCondition(conditions=tuple(conditions))
+
+
 def _infer_node_name(node: Callable[..., Any]) -> str:
     node_name = getattr(node, "__name__", "")
     if not node_name or node_name == "<lambda>":

libs/langgraph/tests/advanced-graph/test_sub_agents.py

@@ -5,7 +5,14 @@ from typing import Any, Literal
 import pytest
 from typing_extensions import TypedDict
 
-from langgraph.advanced_graph import AdvancedStateGraph, publish_to_channel, wait_for
+from langgraph.advanced_graph import (
+    AdvancedStateGraph,
+    any_of,
+    channel_condition,
+    publish_to_channel,
+    timer_condition,
+    wait_for,
+)
 from langgraph.constants import END, START
 from langgraph.graph import StateGraph
 from langgraph.types import Command, Send
@@ -90,9 +97,26 @@ async def test_async_sub_graph() -> None:
         return Command(goto=sends)
 
     async def wait_node(state: MainAgentState) -> Command:
-        # Lightweight interrupt: only this node blocks on inbox.
-        msg = await wait_for("inbox")
-        state["output"].append(f"{msg['type']}: {msg['payload']}")
+        # Lightweight interrupt: only this node blocks for the next relevant signal.
+        event = await wait_for(
+            any_of(
+                channel_condition("tool_completion_channel"),
+                channel_condition("subagent_completion_channel"),
+                channel_condition("user_input_channel"),
+                timer_condition(minutes=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}")
+        else:
+            state["output"].append("timer: no updates yet")
         # Loop back to planner with updated output.
         return Command(goto=Send("llm_node", state))
 
@@ -101,8 +125,8 @@ async def test_async_sub_graph() -> None:
         # Fire-and-forget style completion: publish result to inbox and exit.
         # (i.e., just complete without explicitly going to a next node)
         publish_to_channel(
-            "inbox",
-            {"type": "tool", "payload": f"tool completed for: {tool_input}"},
+            "tool_completion_channel",
+            f"tool completed for: {tool_input}",
         )
 
     async def sub_agent_node(sub_agent_input: str) -> None:
@@ -112,8 +136,8 @@ async def test_async_sub_graph() -> None:
         )
         # Same pattern as tool node: publish result and complete current node.
         publish_to_channel(
-            "inbox",
-            {"type": "sub_agent", "payload": sub_agent_output["output"]},
+            "subagent_completion_channel",
+            sub_agent_output["output"],
         )
 
     async def order_food_node(complete_message: str) -> dict[str, str]:
@@ -121,7 +145,9 @@ async def test_async_sub_graph() -> None:
 
     advanced_flow = AdvancedStateGraph(MainAgentState)
     # Default behavior is an unbounded async channel (maxsize=None).
-    advanced_flow.add_async_channel("inbox", dict)
+    advanced_flow.add_async_channel("tool_completion_channel", str)
+    advanced_flow.add_async_channel("subagent_completion_channel", str)
+    advanced_flow.add_async_channel("user_input_channel", str)
     advanced_flow.add_entry_node(llm_node)
     advanced_flow.node(wait_node)
     advanced_flow.node(tool_node)
@@ -151,9 +177,7 @@ async def test_async_sub_graph() -> None:
 
     # External input can be injected while graph execution is in progress.
     await asyncio.sleep(0.01)
-    await main_agent.apublish_to_channel(
-        "inbox", {"type": "user_input", "payload": "No spicy food please"}
-    )
+    await main_agent.apublish_to_channel("user_input_channel", "No spicy food please")
     result = await started
 
     assert result == {