LangGraph 源码深潜：StateGraph 与 Workflow Agent 运行机制解剖

1
def __init__(
2
    self,
3
    state_schema: type[StateT],
4
    context_schema: type[ContextT] | None = None,
5
    *,
6
    input_schema: type[InputT] | None = None,
7
    output_schema: type[OutputT] | None = None,
8
    **kwargs: Unpack[DeprecatedKwargs],
9
) -> None:
10
    ...

1
用户代码 StateGraph(TravelState)
2
  ↓
3
StateGraph.__init__
4
  ↓
5
StateGraph._add_schema(state_schema)
6
  ↓
7
_get_channels(schema)

1
def __init__(state_schema, context_schema=None,
2
             input_schema=None, output_schema=None, **kwargs):
3
    # 1. 处理 deprecated 参数
4
    if "config_schema" in kwargs and context_schema is None:
5
        warn_deprecated("config_schema")
6
        context_schema = kwargs["config_schema"]
7

8
    if "input" in kwargs and input_schema is None:
9
        warn_deprecated("input")
10
        input_schema = kwargs["input"]
11

12
    if "output" in kwargs and output_schema is None:
13
        warn_deprecated("output")
14
        output_schema = kwargs["output"]
15

16
    # 2. 初始化 builder 容器
17
    self.nodes = {}
18
    self.edges = set()
19
    self.branches = defaultdict(dict)
20
    self.schemas = {}
21
    self.channels = {}
22
    self.managed = {}
23
    self.waiting_edges = set()
24
    self.compiled = False
25

26
    # 3. 保存 schema 边界
27
    self.state_schema = state_schema
28
    self.input_schema = input_schema or state_schema
29
    self.output_schema = output_schema or state_schema
30
    self.context_schema = context_schema
31

32
    # 4. 初始化节点默认策略
33
    self._node_defaults = _NodeDefaults()
34

35
    # 5. 解析 state / input / output schema
36
    self._add_schema(self.state_schema)
37
    self._add_schema(self.input_schema, allow_managed=False)
38
    self._add_schema(self.output_schema, allow_managed=False)

1
def _add_schema(
2
    self,
3
    schema: type[Any],
4
    /,
5
    allow_managed: bool = True,
6
) -> None:
7
    ...

1
StateGraph.__init__
2
  ↓
3
StateGraph._add_schema
4
  ↓
5
_get_channels(schema)
6
  ↓
7
_get_channel / _is_field_binop / _is_field_channel

1
def _add_schema(schema, allow_managed=True):
2
    # 1. 避免重复解析同一个 schema
3
    if schema in self.schemas:
4
        return
5

6
    # 2. 对非法 schema 给出警告
7
    warn_if_schema_not_type_or_annotated(schema)
8

9
    # 3. 从 schema 提取 channels 与 managed values
10
    channels, managed, type_hints = _get_channels(schema)
11

12
    # 4. input/output schema 不允许 managed values
13
    if managed and not allow_managed:
14
        raise ValueError("Managed channels are not permitted")
15

16
    # 5. 记录 schema 对应的所有字段定义
17
    self.schemas[schema] = {**channels, **managed}
18

19
    # 6. 合并普通 channels
20
    for key, channel in channels.items():
21
        if key in self.channels:
22
            if self.channels[key] != channel:
23
                if isinstance(channel, LastValue):
24
                    # 默认 LastValue 可以兼容已有定义
25
                    pass
26
                else:
27
                    raise ValueError("Channel already exists with different type")
28
        else:
29
            self.channels[key] = channel
30

31
    # 7. 合并 managed values
32
    for key, managed_value in managed.items():
33
        if key in self.managed and self.managed[key] != managed_value:
34
            raise ValueError("Managed value already exists with different type")
35
        self.managed[key] = managed_value

1
TravelState TypedDict
2
  ↓
3
_get_channels
4
  ↓
5
user_request → LastValue(str)
6
destination → LastValue(str | None)
7
days → LastValue(int | None)
8
plan → LastValue(str | None)
9
  ↓
10
self.channels 保存字段通道

1
def add_node(
2
    self,
3
    node: str | StateNode[NodeInputT, ContextT],
4
    action: StateNode[NodeInputT, ContextT] | None = None,
5
    *,
6
    defer: bool = False,
7
    metadata: dict[str, Any] | None = None,
8
    input_schema: type[NodeInputT] | None = None,
9
    retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
10
    cache_policy: CachePolicy | None = None,
11
    error_handler: StateNode[Any, ContextT] | None = None,
12
    destinations: dict[str, str] | tuple[str, ...] | None = None,
13
    timeout: float | timedelta | TimeoutPolicy | None = None,
14
    **kwargs: Unpack[DeprecatedKwargs],
15
) -> Self:
16
    ...

1
用户代码 builder.add_node("parse_request", parse_request)
2
  ↓
3
StateGraph.add_node
4
  ↓
5
coerce_to_runnable(action)
6
  ↓
7
StateNodeSpec(...)
8
  ↓
9
self.nodes[name] = spec

1
def add_node(node, action=None, *, input_schema=None,
2
             retry_policy=None, cache_policy=None,
3
             error_handler=None, defer=False,
4
             metadata=None, destinations=None, timeout=None, **kwargs):
5
    # 1. 处理 deprecated 参数
6
    if "retry" in kwargs and retry_policy is None:
7
        warn_deprecated("retry")
8
        retry_policy = kwargs["retry"]
9

10
    if "input" in kwargs and input_schema is None:
11
        warn_deprecated("input")
12
        input_schema = kwargs["input"]
13

14
    # 2. 规范化 timeout policy
15
    timeout = coerce_timeout_policy(timeout)
16

17
    # 3. 如果第一个参数不是字符串，说明用户写的是 add_node(func)
18
    if not isinstance(node, str):
19
        action = node
20
        if isinstance(action, Runnable):
21
            node_name = action.get_name()
22
        else:
23
            node_name = getattr(action, "__name__", action.__class__.__name__)
24
    else:
25
        node_name = node
26

27
    # 4. 基本校验
28
    if node_name is None:
29
        raise ValueError("Node name must be provided")
30

31
    if self.compiled:
32
        logger.warning("Adding node after compile will not affect compiled graph")
33

34
    if node_name in {START, END}:
35
        raise ValueError("START/END cannot be used as normal node name")
36

37
    if node_name in self.nodes:
38
        raise ValueError("Node already exists")
39

40
    if node_name in self.channels:
41
        raise ValueError("Node name conflicts with state channel")
42

43
    # 5. 推断或注册节点输入 schema
44
    resolved_input_schema = input_schema or self.state_schema
45
    self._add_schema(resolved_input_schema)
46

47
    # 6. 将函数 / Runnable 归一化为 Runnable
48
    runnable = coerce_to_runnable(action, name=node_name)
49

50
    # 7. 保存节点规格
51
    self.nodes[node_name] = StateNodeSpec(
52
        runnable=runnable,
53
        metadata=metadata,
54
        input_schema=resolved_input_schema,
55
        retry_policy=retry_policy,
56
        cache_policy=cache_policy,
57
        error_handler=error_handler,
58
        defer=defer,
59
        destinations=destinations,
60
        timeout=timeout,
61
    )
62

63
    return self

1
builder.add_node(parse_request)

1
parse_request

1
builder.add_node("parse_request", parse_request)

1
add_node("parse_request", parse_request)
2
  ↓
3
node_name = "parse_request"
4
  ↓
5
input_schema = TravelState
6
  ↓
7
coerce_to_runnable(parse_request)
8
  ↓
9
StateNodeSpec(...)
10
  ↓
11
self.nodes["parse_request"] = spec

1
def add_edge(
2
    self,
3
    start_key: str | list[str],
4
    end_key: str,
5
) -> Self:
6
    ...

1
用户代码 builder.add_edge(START, "parse_request")
2
  ↓
3
StateGraph.add_edge
4
  ↓
5
self.edges.add((START, "parse_request"))

1
def add_edge(start_key, end_key):
2
    # 1. 已编译后添加边只影响 builder，不影响旧 compiled graph
3
    if self.compiled:
4
        logger.warning("Adding edge after compile will not affect compiled graph")
5

6
    # 2. 校验特殊节点方向
7
    if start_key == END:
8
        raise ValueError("END cannot be a start node")
9

10
    if end_key == START:
11
        raise ValueError("START cannot be an end node")
12

13
    # 3. 多起点边：等待多个节点都完成后再触发 end
14
    if isinstance(start_key, list):
15
        for key in start_key:
16
            if key == END:
17
                raise ValueError("END cannot be a start node")
18
        self.waiting_edges.add((tuple(start_key), end_key))
19
        return self
20

21
    # 4. 普通固定边
22
    self.edges.add((start_key, end_key))
23
    return self

1
START 只能作为入口起点
2
END 只能作为终点

1
builder.add_edge(END, "x")

1
add_edge(START, "parse_request")
2
add_edge("parse_request", "generate_plan")
3
add_edge("generate_plan", END)
4
  ↓
5
self.edges = {
6
  (START, "parse_request"),
7
  ("parse_request", "generate_plan"),
8
  ("generate_plan", END)
9
}

1
def compile(
2
    self,
3
    checkpointer: Checkpointer = None,
4
    *,
5
    cache: BaseCache | None = None,
6
    store: BaseStore | None = None,
7
    interrupt_before: All | list[str] | None = None,
8
    interrupt_after: All | list[str] | None = None,
9
    debug: bool = False,
10
    name: str | None = None,
11
) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]:
12
    ...

1
用户代码 graph = builder.compile()
2
  ↓
3
StateGraph.compile
4
  ↓
5
validate graph
6
  ↓
7
CompiledStateGraph(...)
8
  ↓
9
compiled.attach_node / attach_edge / attach_branch
10
  ↓
11
compiled.validate()

1
def compile(checkpointer=None, *, cache=None, store=None,
2
            interrupt_before=None, interrupt_after=None,
3
            debug=False, name=None):
4
    # 1. 规范化 interrupt 配置
5
    interrupt_before = interrupt_before or []
6
    interrupt_after = interrupt_after or []
7

8
    # 2. 校验图结构
9
    self.validate(
10
        interrupt_before=interrupt_before,
11
        interrupt_after=interrupt_after,
12
    )
13

14
    # 3. 标记 builder 已编译
15
    self.compiled = True
16

17
    # 4. 创建 CompiledStateGraph / Pregel 应用
18
    compiled = CompiledStateGraph(
19
        builder=self,
20
        channels={
21
            **self.channels,
22
            START: EphemeralValue(self.input_schema),
23
        },
24
        input_channels=START,
25
        output_channels=output_keys_from_schema(self.output_schema),
26
        stream_channels=stream_keys_from_state_channels(self.channels),
27
        checkpointer=checkpointer,
28
        cache=cache,
29
        store=store,
30
        interrupt_before=interrupt_before,
31
        interrupt_after=interrupt_after,
32
        debug=debug,
33
        name=name,
34
    )
35

36
    # 5. 附加 START 虚拟节点
37
    compiled.attach_node(START, None)
38

39
    # 6. 附加用户节点
40
    for node_name, node_spec in self.nodes.items():
41
        compiled.attach_node(node_name, node_spec)
42

43
    # 7. 附加固定边
44
    for start, end in self.edges:
45
        compiled.attach_edge(start, end)
46

47
    # 8. 附加等待边
48
    for starts, end in self.waiting_edges:
49
        compiled.attach_edge(starts, end)
50

51
    # 9. 附加条件分支，本篇不展开
52
    for start, branches in self.branches.items():
53
        for name, branch in branches.items():
54
            compiled.attach_branch(start, name, branch)
55

56
    # 10. 校验 compiled graph 并返回
57
    return compiled.validate()

1
builder.compile()
2
  ↓
3
validate builder
4
  ↓
5
CompiledStateGraph(...)
6
  ↓
7
attach START
8
  ↓
9
attach parse_request
10
  ↓
11
attach generate_plan
12
  ↓
13
attach START → parse_request
14
  ↓
15
attach parse_request → generate_plan
16
  ↓
17
attach generate_plan → END
18
  ↓
19
return executable graph

1
class CompiledStateGraph(Pregel[...]):
2
    builder: StateGraph
3
    ...

1
StateGraph.compile
2
  ↓
3
CompiledStateGraph(...)
4
  ↓
5
CompiledStateGraph.attach_node / attach_edge
6
  ↓
7
Pregel.invoke / stream

1
class CompiledStateGraph(Pregel):
2
    def attach_node(self, key, node_spec):
3
        # 1. START 是特殊入口节点
4
        if key == START:
5
            self.nodes[START] = PregelNode(
6
                triggers=[START],
7
                channels=[START],
8
                writers=[...],
9
            )
10
            return
11

12
        # 2. 普通节点读取 state channels
13
        input_schema = node_spec.input_schema
14
        read_channels = keys_from_schema(input_schema)
15

16
        # 3. 包装用户 runnable
17
        runnable = node_spec.runnable
18

19
        # 4. 生成写入器：只允许写 state 中合法字段
20
        writers = [
21
            ChannelWrite(
22
                entries=[ChannelWriteTupleEntry(...)]
23
            )
24
        ]
25

26
        # 5. 保存 PregelNode
27
        self.nodes[key] = PregelNode(
28
            triggers=[key],
29
            channels=read_channels,
30
            mapper=state_mapper,
31
            writers=writers,
32
        ) | runnable | writers

1
CompiledStateGraph.attach_node("parse_request", node_spec)
2
  ↓
3
生成 PregelNode
4
  ↓
5
PregelNode 订阅 parse_request trigger
6
  ↓
7
读取 TravelState channels
8
  ↓
9
执行 parse_request runnable
10
  ↓
11
将返回 dict 写入 destination / days channels

1
def invoke(
2
    self,
3
    input: InputT | Command | None,
4
    config: RunnableConfig | None = None,
5
    *,
6
    context: ContextT | None = None,
7
    stream_mode: StreamMode = "values",
8
    print_mode: StreamMode | Sequence[StreamMode] = (),
9
    output_keys: str | Sequence[str] | None = None,
10
    interrupt_before: All | Sequence[str] | None = None,
11
    interrupt_after: All | Sequence[str] | None = None,
12
    durability: Durability | None = None,
13
    **kwargs: Any,
14
) -> dict[str, Any] | Any:
15
    ...

1
用户代码 graph.invoke(initial_state)
2
  ↓
3
CompiledStateGraph.invoke
4
  ↓
5
Pregel.invoke
6
  ↓
7
Pregel.stream
8
  ↓
9
运行时循环

1
def invoke(input, config=None, *, context=None,
2
           stream_mode="values", output_keys=None,
3
           interrupt_before=None, interrupt_after=None, **kwargs):
4
    # 1. 同步 invoke 复用 stream 主链
5
    chunks = []
6

7
    for chunk in self.stream(
8
        input,
9
        config=config,
10
        context=context,
11
        stream_mode=stream_mode,
12
        output_keys=output_keys,
13
        interrupt_before=interrupt_before,
14
        interrupt_after=interrupt_after,
15
        **kwargs,
16
    ):
17
        # 2. 收集流式输出
18
        chunks.append(chunk)
19

20
    # 3. 根据 stream_mode 整理最终输出
21
    if stream_mode == "values":
22
        return last_value_chunk(chunks)
23

24
    if stream_mode == "updates":
25
        return merge_or_return_updates(chunks)
26

27
    return normalize_stream_result(chunks)

1
graph.invoke(initial_state)
2
  ↓
3
for chunk in graph.stream(...)
4
  ↓
5
记录最后一个 values chunk
6
  ↓
7
返回最终 state

1
Pregel.invoke
2
  ↓
3
Pregel.stream
4
  ↓
5
create loop / runtime context
6
  ↓
7
while loop.tick(): execute active tasks
8
  ↓
9
apply writes to channels
10
  ↓
11
yield stream chunks

1
def stream(input, config=None, *, stream_mode="values", output_keys=None,
2
           interrupt_before=None, interrupt_after=None, context=None, **kwargs):
3
    # 1. 归一化 config / context / stream 参数
4
    config = ensure_config(config)
5
    output_keys = resolve_output_keys(output_keys)
6
    stream_modes = normalize_stream_mode(stream_mode)
7

8
    # 2. 创建运行时 loop，包含 channels、checkpoint、tasks、step 计数
9
    with create_pregel_loop(
10
        graph=self,
11
        input=input,
12
        config=config,
13
        context=context,
14
        output_keys=output_keys,
15
        interrupt_before=interrupt_before,
16
        interrupt_after=interrupt_after,
17
    ) as loop:
18
        # 3. 将初始 input 写入 input channel
19
        loop.initialize_input(input)
20

21
        # 4. Pregel 主循环：每次 tick 对应一个 super-step
22
        while loop.tick():
23
            # 4.1 找到当前 super-step 中 active tasks
24
            tasks = loop.prepare_next_tasks()
25

26
            # 4.2 执行 active nodes，可并行
27
            results = runner.tick(tasks)
28

29
            # 4.3 收集每个节点的 writes
30
            writes = collect_writes(results)
31

32
            # 4.4 将 writes 应用到 channels
33
            loop.apply_writes(writes)
34

35
            # 4.5 根据 stream_mode 产出 updates / values
36
            for chunk in loop.emit_stream_chunks(stream_modes):
37
                yield chunk
38

39
        # 5. loop 结束后产出最终状态或清理
40
        yield from loop.emit_final_if_needed()

1
step 0:
2
  START channel 收到 input
3
  ↓
4
  parse_request 被激活
5

6
step 1:
7
  parse_request 执行
8
  ↓
9
  写入 destination / days
10
  ↓
11
  generate_plan 被激活
12

13
step 2:
14
  generate_plan 执行
15
  ↓
16
  写入 plan
17
  ↓
18
  END 被触发
19

20
step 3:
21
  无 active tasks
22
  ↓
23
  输出最终 state