Compare commits

...
Author SHA1 Message Date
Quanzheng Long 388d6b3593 done 2026-03-12 16:05:12 -07:00
Quanzheng Long 667b679694 more 2026-03-12 15:58:42 -07:00
Quanzheng Long ccb9f4c41a more 2026-03-12 15:58:21 -07:00
Quanzheng Long 789be99634 ref 2026-03-12 15:44:35 -07:00
Quanzheng Long e33842ff54 mre 2026-03-12 15:39:46 -07:00
Quanzheng Long 89d37a0f9b more 2026-03-12 15:38:22 -07:00
Quanzheng Long 8f717d3874 more 2026-03-12 15:27:25 -07:00
Quanzheng Long 508193272a anyof 2026-03-12 15:15:49 -07:00
Quanzheng Long 788ae6cb72 rm 2026-03-12 15:03:58 -07:00
Quanzheng Long b2dca399e8 1 2026-03-12 13:31:54 -07:00
Quanzheng Long 4b2167dd25 1stpass 2026-03-12 13:13:08 -07:00
Quanzheng Long be46e91180 doc 2026-03-12 12:37:34 -07:00
Quanzheng Long bbb308259b poc 2026-03-12 12:31:46 -07:00
4 changed files with 819 additions and 0 deletions
@@ -0,0 +1,25 @@
from langgraph.advanced_graph.state import (
AdvancedStateGraph,
AnyOfCondition,
ChannelCondition,
CompiledGraphEngine,
Context,
GraphRunHandler,
TimerCondition,
any_of,
channel_condition,
timer_condition,
)
__all__ = (
"AdvancedStateGraph",
"AnyOfCondition",
"ChannelCondition",
"Context",
"CompiledGraphEngine",
"GraphRunHandler",
"TimerCondition",
"any_of",
"channel_condition",
"timer_condition",
)
@@ -0,0 +1,130 @@
# Evolve/Extend LangGraph with next level of orchestration
## LangGraph Today: A Strong Foundation with Creative Innovation
LangGraph is already an exceptional orchestration framework. It has introduced a number of creative features that no other workflow engine on the market has even attempted.
**First-class streaming.** No workflow engine has ever integrated streaming as seamlessly as LangGraph. Streaming is not an afterthought bolted on top — it is woven into the core execution model, allowing every node, every tool call, and every LLM interaction to emit incremental output naturally.
**Flexible durability modes.** LangGraph defaults to asynchronous execution and supports sync and "exit" modes as well. This is a significant departure from traditional workflow engines, which typically only offer synchronous execution. The ability to choose a durability mode gives developers fine-grained control over the trade-off between persistence guarantees and execution speed.
**Reusable checkpoints.** The checkpoint system allows state to be captured at any point during graph execution and freely replayed, forked, or resumed later. This enables powerful patterns like time-travel debugging, human-in-the-loop approval flows, and long-running conversations that can be picked up exactly where they left off.
**Double texting.** LangGraph natively handles the real-world scenario where a user sends a new message while a previous one is still being processed — a problem most orchestration frameworks simply ignore.
Beyond these innovative features, LangGraph provides solid support for the foundational workflow execution patterns that developers rely on daily. Sequential execution, or loops and conditional branching. Basic parallelism is also well supported: when multiple LLM calls or tool invocations are independent of each other, they can run concurrently to avoid the latency cost of sequential execution, and their results are merged back into the shared state for downstream processing.
LangGraph also offers a simple and intuitive mechanism for human-in-the-loop interactions, allowing a graph to pause execution and wait for user input before continuing.
Combined with the broader LangChain ecosystem, these have made LangGraph a significant success in the market.
## Emerging Gaps: What LangGraph Struggles to Support
As adoption has grown and use cases have become more sophisticated, we have discovered an increasing number of scenarios and design patterns that LangGraph cannot support well today.
**Complex sub-agent coordination.** A main agent often needs to manage multiple sub-agents, but the coordination involved is far more nuanced than simply launching a batch of sub-agents, waiting for all of them to finish, and then moving on. In practice, a main agent may launch a sub-agent, continue doing other work, spawn additional sub-agents later, wait selectively for certain results, retry with a different strategy if one sub-agent fails, or dynamically decide what to do next based on partial results that arrive at unpredictable times.
LangGraph today lacks the coordination primitives to express this. The current parallelism model groups multiple nodes into a single superstep — all of them execute concurrently, but _all_ must complete before the graph can advance to the next step. There is no way for one node to proceed independently while others are still running, and no built-in mechanism for selective waiting, partial result handling, or dynamic task spawning mid-execution.
Sub-agents also cannot simply be modeled as subgraphs, because subgraphs today execute within the same run. They cannot be scaled up independently — if a sub-agent is resource-intensive, there is no straightforward way to run it on a separate machine. Ideally, launching a sub-agent should be(or opt in) as simple as dispatching it for distributed execution across multiple machines.
**Concurrent input and output (e.g. audio agents).** Audio agents also present a particularly clear example of a pattern LangGraph cannot express today. In a voice interaction, speech input and speech output may happen simultaneously — the agent should be able to process a previous utterance, continue receiving new audio input, and produce output all at the same time. These three activities should not be mutually exclusive.
The closest workaround today is double texting, but it has a fundamental flaw: when a new audio input arrives, the previous one is interrupted and canceled rather than being allowed to gracefully complete. The workflow code itself should have the control to decide whether to stop running.
LangGraph is, at its core, a general-purpose workflow engine. Although we focus primarily on agent development, none of the primitives it offers are exclusive to agents or dedicated solely to agentic use cases. Conversely, there is nothing that a general-purpose workflow engine provides that we can safely assume agent development will _never_ need.
The difference is probably only priority. For example, durable timer where a step can sleep for hours, days or months before resuming. Traditional workflow engines — those built for general microservice orchestration(which doesn't need streaming) -- they may need durable timers. In the agent development world today, most agents are still relatively simple. There are not yet many scenarios that require a step to wait for hours or days before proceeding.
## Deriving What's Needed from First Principles
Before jumping to solutions, it is worth stepping back and asking a fundamental question: what is an orchestration engine, and what do users expect it to provide?
At its most fundamental level, a workflow engine's value proposition is making a long-running process execute reliably. If a machine crashes, execution should smoothly fail over to another machine and resume from the last point where it was interrupted — not start over from the beginning. So we can reason about what is needed by asking: what would a developer do if they had to build a long-running process _without_ a workflow engine?
Starting from the simple. A developer could write a simple `main` function — a single-threaded program, just like everyone writes when they first learn to code. It would have `if/else` branches, `for` loops, and maybe it would wait for command-line input. Many early agent use cases look exactly like this: execute a sequence of steps, make decisions along the way, loop when necessary.
But if that machine crashes, you probably do not want the process to start over from scratch. You want it to resume from the last step that completed successfully. And if a step fails, you might want it to retry automatically before giving up.
LangGraph handles this case very well.
There is an important constraint worth calling out explicitly: LangGraph requires the developer to organize their code into **nodes**, which serve as the boundaries at which checkpoints can be taken. This is a constraint shared by every workflow engine — it is simply not feasible to persist a checkpoint after every single line of arbitrary code.
### From Single-Threaded to Concurrent: Where the Model Breaks Down
But as product requirements grow more complex, a single-threaded program is no longer sufficient. The process becomes multi-threaded or multi-process. And in a multi-threaded program, each thread executes independently — when one thread finishes a step and moves on to its next step, it does not need to wait for another thread to finish _its_ current step first.
This is precisely why LangGraph's superstep restriction feels awkward in practice. In the superstep model, all concurrently executing nodes must complete before any of them can advance. But that is not how independent threads work. Each thread should be able to progress at its own pace, checkpoint its own state, and move to its next step without being blocked by unrelated work happening in parallel.
Multiple threads and processes do, however, need to coordinate with each other. In concurrent programming, channels are an essential primitive precisely because they provide a safe, structured way for threads to communicate and synchronize without relying on shared mutable memory — avoiding data races and deadlocks. In some cases, threads may use locking for coordination, but the preferred approach is message passing through channels.
NOTE: "channel" is overloaded term here as it's also an internal term within current LangGraph pregel algorithm.
LangGraph already has a mechanism that is closely related: `interrupt`. A run can be interrupted, and then another run can resume it. If we look at this through the lens of channels, `interrupt` is essentially a **channel with size 0** — a synchronous rendezvous point where one side blocks until the other side is ready.
The natural extension:
1. **Variable-size channels.** The channel buffer size should be configurable — size 0 for synchronous handoff (like `interrupt` today), size N for buffered communication where the sender can proceed without waiting, and unbounded for fully asynchronous fire-and-forget messaging.
2. **Channels across boundaries.** Channels should not be limited to communication between separate runs. Nodes within the same graph should also be able to send and receive through channels — mirroring the way both multi-process communication (between runs) and multi-thread communication (between nodes within a run) work in ordinary concurrent programs.
3. **Node-level blocking, not run-level pausing.** When a node waits on a channel (i.e. `interrupt`), only that node should block — the rest of the graph should continue executing. Today, `interrupt` pauses the entire run. In a concurrent program, when one thread blocks on a channel read, the other threads keep running. The same should be true: an interrupt should suspend the individual node, not halt the whole run.
## Summmary of all extension opportunity
### P1: urgently needed
#### Remove the Superstep Restriction
Today, when multiple nodes execute in parallel, they are grouped into a superstep. All nodes in a superstep must complete before any downstream node can begin. This means that even if `b1` finishes quickly and its successor `b11` is ready to run, it must wait for `b2` to finish first.
With the superstep restriction removed, each parallel branch progresses independently. As soon as a node completes, its downstream successor can begin immediately — regardless of what is happening in other branches.
**Current behavior (superstep model):**
```
Step 1: a
Step 2: b1, b2 ← both must finish before step 3
Step 3: b11, b22 ← both start together
```
Even if `b1` finishes in 1 second and `b2` takes 30 seconds, `b11` cannot start until `b2` is done.
**Proposed behavior (independent branches):**
```
Branch 1: a → b1 → b11 → ...
Branch 2: a → b2 → b22 → ...
```
Each branch advances at its own pace. `b1` finishing triggers `b11` immediately, without waiting for `b2`.
No API change is needed from the user's perspective — the graph definition stays the same. The change is in the execution semantics: the engine no longer forces all parallel nodes to synchronize at each step boundary. Each branch is checkpointed independently, so if `b1 → b11` completes while `b2` is still running, `b11`'s result is already persisted.
This is necessary for the next one -- Light-weight Interrupt: Only Block the Current Node. Because we want to let other nodes continue to run while a node is waiting on something.
#### Light-weight Interrupt -- wait_for API: Only Block the Current Node Until Channel Has Enough Messages
Today, `interrupt` pauses the entire run. Every node stops, and nothing can proceed until the interrupt is resolved externally. This is the right behavior for a simple single-threaded workflow, but it breaks down when multiple branches are executing concurrently — one branch needing input should not freeze all the others.
The proposed change has three parts:
1. **Named channels.** A graph can declare named channels as coordination points. These are distinct from the graph's state — they are message-passing primitives, not shared memory.
2. **`wait_for` blocks only the current node.** When a node calls `wait_for`, it suspends itself and waits for messages on the specified channel. All other nodes in the graph continue executing normally.
3. A channel can be published from both external and internal
The `wait_for` call takes a channel name and optionally a count `N`, meaning "wait until N messages have arrived on this channel before resuming."
**Prototype:**
See [test_sub_agents.py](../libs/langgraph/tests/advanced-graph/test_sub_agents.py)
### P2: likely needed
#### subGraph redesign
#### durable timers
#### more flexiable waiting conditions on interrupts
#### locking on state fields
### P3: future needed or nice to have
#### RPC
@@ -0,0 +1,445 @@
from __future__ import annotations
import asyncio
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
StateT = TypeVar("StateT")
@dataclass(frozen=True)
class _ChannelSpec:
typ: Any
@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
class AdvancedStateGraph(Generic[StateT]):
"""Experimental in-memory graph engine with async channels."""
def __init__(self, state_schema: type[StateT]) -> None:
self.state_schema = state_schema
self._nodes: dict[str, Callable[..., Any]] = {}
self._async_channels: dict[str, _ChannelSpec] = {}
self._entry_point: str | None = None
self._finish_point: str | None = None
def add_node(
self,
name_or_node: str | Callable[..., Any],
node: Callable[..., Any] | None = None,
) -> str:
if node is None:
if not callable(name_or_node):
raise TypeError("add_node() expects a callable when name is omitted")
node_name = _infer_node_name(name_or_node)
node_fn = name_or_node
else:
if not isinstance(name_or_node, str):
raise TypeError("add_node() expects a string node name")
node_name = name_or_node
node_fn = node
if node_name in self._nodes:
raise ValueError(f"Node `{node_name}` already exists")
self._nodes[node_name] = node_fn
return node_name
def add_async_channel(self, name: str, typ: Any) -> None:
if name in self._async_channels:
raise ValueError(f"Channel `{name}` already exists")
self._async_channels[name] = _ChannelSpec(typ=typ)
def set_entry_point(self, name_or_node: str | Callable[..., Any]) -> None:
self._entry_point = self._resolve_node_name(name_or_node)
def set_finish_point(self, name_or_node: str | Callable[..., Any]) -> None:
self._finish_point = self._resolve_node_name(name_or_node)
def add_entry_node(self, node: Callable[..., Any]) -> str:
node_name = self.add_node(node)
self.set_entry_point(node_name)
return node_name
def add_finish_node(self, node: Callable[..., Any]) -> str:
node_name = self.add_node(node)
self.set_finish_point(node_name)
return node_name
def _resolve_node_name(self, name_or_node: str | Callable[..., Any]) -> str:
if isinstance(name_or_node, str):
return name_or_node
node_name = _infer_node_name(name_or_node)
if node_name not in self._nodes:
self._nodes[node_name] = name_or_node
return node_name
def compile(self) -> CompiledGraphEngine[StateT]:
if self._entry_point is None:
raise ValueError("Entry point is not set")
if self._finish_point is None:
raise ValueError("Finish point is not set")
if self._entry_point not in self._nodes:
raise ValueError(f"Entry point node `{self._entry_point}` does not exist")
if self._finish_point not in self._nodes:
raise ValueError(f"Finish point node `{self._finish_point}` does not exist")
return CompiledGraphEngine(
nodes=dict(self._nodes),
async_channels=dict(self._async_channels),
entry_point=self._entry_point,
finish_point=self._finish_point,
)
class CompiledGraphEngine(Generic[StateT]):
"""Executable runtime for `AdvancedStateGraph`."""
def __init__(
self,
*,
nodes: dict[str, Callable[..., Any]],
async_channels: dict[str, _ChannelSpec],
entry_point: str,
finish_point: str,
) -> None:
self._nodes = nodes
self._async_channels = async_channels
self._entry_point = entry_point
self._finish_point = finish_point
async def ainvoke(self, initial_state: StateT) -> StateT:
handler = await self.astart(initial_state)
return await handler
async def astart(self, initial_state: StateT) -> GraphRunHandler[StateT]:
run = _GraphEngineRun(
nodes=self._nodes,
async_channel_specs=self._async_channels,
entry_point=self._entry_point,
finish_point=self._finish_point,
)
task = asyncio.create_task(run.run(initial_state))
return GraphRunHandler(run=run, task=task)
class Context:
"""Per-run context injected into advanced graph nodes."""
def __init__(self, run: _GraphEngineRun) -> None:
self._run = run
async def wait_for(self, target: WaitCondition | AnyOfCondition) -> Any:
return await self._run.wait_for(target)
def publish_to_channel(self, channel: str, value: Any) -> None:
self._run.publish_nowait(channel, value)
async def apublish_to_channel(self, channel: str, value: Any) -> None:
await self._run.publish(channel, value)
class GraphRunHandler(Generic[StateT]):
"""Handle for an active in-memory run."""
def __init__(self, *, run: _GraphEngineRun, task: asyncio.Task[StateT]) -> None:
self._run = run
self._task = task
async def apublish_to_channel(self, channel: str, value: Any) -> None:
if self._task.done():
raise RuntimeError("Run has already completed")
await self._run.publish(channel, value)
async def aresult(self) -> StateT:
return await self._task
def __await__(self) -> Any:
return self._task.__await__()
class _GraphEngineRun:
def __init__(
self,
*,
nodes: dict[str, Callable[..., Any]],
async_channel_specs: dict[str, _ChannelSpec],
entry_point: str,
finish_point: str,
) -> None:
self._nodes = nodes
self._entry_point = entry_point
self._finish_point = finish_point
self._async_channels: dict[str, asyncio.Queue[Any]] = {
name: asyncio.Queue() for name, _spec in async_channel_specs.items()
}
self._tasks: set[asyncio.Task[list[Send]]] = set()
self._finished = False
self._state: Any = None
self.context = Context(self)
async def run(self, initial_state: StateT) -> StateT:
self._state = initial_state
self._schedule(Send(self._entry_point, initial_state))
try:
while self._tasks and not self._finished:
done, _ = await asyncio.wait(
self._tasks, return_when=asyncio.FIRST_COMPLETED
)
for task in done:
self._tasks.remove(task)
exc = task.exception()
if exc is not None:
await self._cancel_all_tasks()
raise exc
sends = task.result()
for send in sends:
self._schedule(send)
if self._finished:
await self._cancel_all_tasks()
return cast(StateT, self._state)
finally:
await self._cancel_all_tasks()
async def publish(self, channel: str, value: Any) -> None:
queue = self._get_async_channel(channel)
await queue.put(value)
def publish_nowait(self, channel: str, value: Any) -> None:
queue = self._get_async_channel(channel)
queue.put_nowait(value)
async def wait_for(self, target: WaitCondition | AnyOfCondition) -> Any:
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_async_channel(channel)
if n == 1:
return await queue.get()
values: list[Any] = []
for _ in range(n):
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))
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_async_channel(self, channel: str) -> asyncio.Queue[Any]:
if channel not in self._async_channels:
raise ValueError(f"Unknown channel `{channel}`")
return self._async_channels[channel]
def _schedule(self, send: Send) -> None:
if self._finished:
return
task: asyncio.Task[list[Send]] = asyncio.create_task(self._execute_send(send))
self._tasks.add(task)
async def _cancel_all_tasks(self) -> None:
if not self._tasks:
return
to_cancel = list(self._tasks)
for task in to_cancel:
task.cancel()
await asyncio.gather(*to_cancel, return_exceptions=True)
self._tasks.clear()
async def _execute_send(self, send: Send) -> list[Send]:
node_name = _resolve_target_name(send.node)
if node_name not in self._nodes:
raise ValueError(f"Unknown node `{node_name}`")
node = self._nodes[node_name]
result = _invoke_node(node, self.context, send.arg)
if inspect.isawaitable(result):
result = await result
if isinstance(result, Command):
self._apply_update(result.update)
next_sends = _normalize_goto(result.goto, default_arg=self._state)
else:
self._apply_update(result)
next_sends = _normalize_result_to_sends(result, default_arg=self._state)
if node_name == self._finish_point:
self._finished = True
return []
return next_sends
def _apply_update(self, update: Any) -> None:
if update is None:
return
if isinstance(update, Mapping):
if isinstance(self._state, Mapping):
# Keep semantics simple: in-place update for mapping-like state.
cast(dict[str, Any], self._state).update(update)
return
if isinstance(update, Sequence) and not isinstance(update, (str, bytes)):
pairs = list(update)
if all(
isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str)
for item in pairs
):
if isinstance(self._state, Mapping):
cast(dict[str, Any], self._state).update(
cast(dict[str, Any], pairs)
)
return
def _normalize_result_to_sends(result: Any, *, default_arg: Any) -> list[Send]:
if result is None:
return []
if isinstance(result, Send):
return [result]
if callable(result):
return [Send(_infer_node_name(result), default_arg)]
if isinstance(result, str):
return [Send(result, default_arg)]
if isinstance(result, Sequence) and not isinstance(result, (str, bytes)):
sends: list[Send] = []
for item in result:
if isinstance(item, Send):
sends.append(item)
elif callable(item):
sends.append(Send(_infer_node_name(item), default_arg))
elif isinstance(item, str):
sends.append(Send(item, default_arg))
return sends
return []
def _normalize_goto(goto: Any, *, default_arg: Any) -> list[Send]:
if not goto:
return []
if isinstance(goto, Send):
return [goto]
if callable(goto):
return [Send(_infer_node_name(goto), default_arg)]
if isinstance(goto, str):
return [Send(goto, default_arg)]
if isinstance(goto, Sequence):
sends: list[Send] = []
for item in goto:
if isinstance(item, Send):
sends.append(item)
elif callable(item):
sends.append(Send(_infer_node_name(item), default_arg))
elif isinstance(item, str):
sends.append(Send(item, default_arg))
return sends
return []
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>":
raise ValueError("Cannot infer node name from anonymous callable")
return node_name
def _resolve_target_name(target: Any) -> str:
if isinstance(target, str):
return target
if callable(target):
return _infer_node_name(target)
raise ValueError(f"Unsupported node target type: {type(target)!r}")
def _invoke_node(node: Callable[..., Any], ctx: Context, state: Any) -> Any:
try:
params = list(inspect.signature(node).parameters.values())
except (TypeError, ValueError):
params = []
if len(params) >= 2:
return node(ctx, state)
if len(params) == 1:
return node(state)
return node()
@@ -0,0 +1,219 @@
import asyncio
from dataclasses import dataclass
from typing import Any, Literal
import pytest
from typing_extensions import TypedDict
from langgraph.advanced_graph import (
AdvancedStateGraph,
Context,
any_of,
channel_condition,
timer_condition,
)
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from langgraph.types import Command, Send
pytestmark = pytest.mark.anyio
class MainAgentState(TypedDict):
input: str
output: list[str]
done: str | None
class SubAgentState(TypedDict):
input: str
output: str
@dataclass(frozen=True)
class Decision:
type: Literal["end", "sub_agent", "tool"]
sub_agent: str | None = None
tool: str | None = None
complete: str | None = None
class MockPlanner:
def __init__(self) -> None:
self.responses: list[list[Decision]] = []
self._idx = 0
async def ainvoke(self, _: MainAgentState) -> list[Decision]:
if self._idx >= len(self.responses):
return []
response = self.responses[self._idx]
self._idx += 1
return response
def build_sub_agent() -> Any:
# Sub-agent uses the regular/simple StateGraph API.
sub_agent = StateGraph(SubAgentState)
async def research_node(state: SubAgentState) -> dict[str, str]:
# Intentionally slower than timer_condition(seconds=1) to validate timer path.
await asyncio.sleep(5)
return {"output": f"research sub agent completed for: {state['input']}"}
sub_agent.add_node("research_node", research_node)
sub_agent.add_edge(START, "research_node")
sub_agent.add_edge("research_node", END)
return sub_agent.compile()
def build_main_agent(planner: MockPlanner, sub_agent: Any) -> Any:
async def llm_node(state: MainAgentState) -> Command:
# Planner decides whether to call a tool, spawn a sub-agent, or finish.
decisions = await planner.ainvoke(state)
sends: list[Send] = []
for decision in decisions:
if decision.type == "end":
# NOTE: this can be simplified further in the future with a dedicated
# complete primitive, instead of routing to a finish node manually.
return Command(
goto=Send(
order_food_node,
{
"state": state,
"complete": decision.complete or "order flow completed",
},
)
)
if decision.type == "sub_agent" and decision.sub_agent:
sends.append(Send("sub_agent_node", decision.sub_agent))
if decision.type == "tool" and decision.tool:
sends.append(Send("tool_node", decision.tool))
# Keep the main loop responsive: wait for one inbound message and continue.
sends.append(Send("wait_node", state))
return Command(goto=sends)
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(
any_of(
channel_condition("tool_completion_channel"),
channel_condition("subagent_completion_channel"),
channel_condition("user_input_channel"),
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}")
# State changed -> ask planner what to do next.
return Command(goto=Send("llm_node", state))
else:
state["output"].append("timer: no updates yet")
# No meaningful state change -> keep waiting without calling planner.
return Command(goto=Send("wait_node", state))
async def tool_node(ctx: Context, tool_input: str) -> None:
await asyncio.sleep(0.1)
# Fire-and-forget style completion: publish result to inbox and exit.
# (i.e., just complete without explicitly going to a next node)
ctx.publish_to_channel(
"tool_completion_channel",
f"tool completed for: {tool_input}",
)
async def sub_agent_node(ctx: Context, sub_agent_input: str) -> None:
# Sub-agent remains a regular StateGraph, compiled independently.
sub_agent_output = await sub_agent.ainvoke(
{"input": sub_agent_input, "output": ""}
)
# Same pattern as tool node: publish result and complete current node.
ctx.publish_to_channel(
"subagent_completion_channel",
sub_agent_output["output"],
)
async def order_food_node(payload: dict[str, Any]) -> dict[str, Any]:
state = payload["state"]
complete_message = payload["complete"]
return {
"done": complete_message,
"output": [*state["output"], f"order_food: {complete_message}"],
}
advanced_flow = AdvancedStateGraph(MainAgentState)
# Default behavior is an unbounded async channel like Rust channel
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)
# nodes are the same as in the regular StateGraph API
advanced_flow.add_entry_node(llm_node)
advanced_flow.add_node(wait_node)
advanced_flow.add_node(tool_node)
advanced_flow.add_node(sub_agent_node)
advanced_flow.add_finish_node(order_food_node)
return advanced_flow.compile()
async def test_async_sub_graph() -> None:
planner = MockPlanner()
sub_agent = build_sub_agent()
main_agent = build_main_agent(planner, sub_agent)
planner.responses = [
[
# First planner pass triggers one slow sub-agent.
Decision(type="sub_agent", sub_agent="research lunch options"),
Decision(type="tool", tool="slack_tool"),
],
# After user input.
[],
# After tool completion.
[],
# After first sub-agent completion, planner decides to run second research.
[Decision(type="sub_agent", sub_agent="find vegetarian fallback")],
# After second sub-agent completion, planner decides to end.
[Decision(type="end", complete="order submitted")],
]
handler = await main_agent.astart(
{"input": "help me get something for lunch", "output": [], "done": None}
)
# External input can be injected while graph execution is in progress.
await asyncio.sleep(0.01)
await handler.apublish_to_channel("user_input_channel", "No spicy food please")
result = await handler.aresult()
assert result["input"] == "help me get something for lunch"
assert result["done"] == "order submitted"
output = result["output"]
assert output.count("timer: no updates yet") >= 3
assert "user_input: No spicy food please" in output
assert "tool: tool completed for: slack_tool" in output
assert (
"sub_agent: research sub agent completed for: research lunch options" in output
)
assert (
"sub_agent: research sub agent completed for: find vegetarian fallback" in output
)
assert output[-1] == "order_food: order submitted"
first_sub_idx = output.index(
"sub_agent: research sub agent completed for: research lunch options"
)
second_sub_idx = output.index(
"sub_agent: research sub agent completed for: find vegetarian fallback"
)
order_food_idx = output.index("order_food: order submitted")
assert first_sub_idx < second_sub_idx < order_food_idx
assert planner._idx == len(planner.responses)
import json
print(json.dumps(result, ensure_ascii=False, indent=2))