diff --git a/README.md b/README.md index e9f2091..4114663 100644 --- a/README.md +++ b/README.md @@ -20,8 +20,11 @@ Go 言語（ジェネリクス）で書かれた、軽量・高性能・型安 ### 核心特性 - **强类型状态管理：** 利用 Go 泛型定义全局状态 `Graph[S]`，彻底告别松散的 `map[string]any` 和繁琐的类型断言。 - **支持循环结构：** 突破了传统 DAG（有向无环图）工作流的限制，原生支持节点间的双向循环和条件路由。 -- **并发分支与合并：** 原生支持多个 Agent 节点的并发执行（基于 `sync.WaitGroup` 调度），并提供线程安全的“分流值拷贝”与“汇合状态合并（Merger）”机制。 +- **并发分支与短路取消：** 原生支持多个 Agent 节点的并发执行，基于 `context.WithCancel` 实现短路机制——任意分支出错立即取消其余兄弟 goroutine，避免算力空转。 +- **并发安全深拷贝：** 通过 `Engine.WithStateCloner` 注入自定义深拷贝函数，彻底消除含引用类型的状态在并发分支中的数据竞争。 - **人机协同中断 (HITL)：** 支持在特定节点前自动挂起，返回执行线程快照，等待人工修改状态或审批通过后一键恢复 (`Resume`)。 +- **死循环硬性熔断：** 通过 `Engine.WithMaxSteps` 设置步数上限，图中的任何意外环路都会被明确报错而非无限阻塞。 +- **轻量级生命周期 Hooks：** `Engine.WithPreNodeHook` / `WithPostNodeHook` 提供节点执行前后的钩子，无需修改任何 `NodeFn`，即可接入日志、链路追踪（OpenTelemetry）或流式输出。 - **开箱即用持久化：** 提供 `Checkpointer[S]` 接口与内置的 `FileCheckpointer[S]`，支持将进度保存为 JSON 文件，实现停机或重启后进度无损恢复。 - **纯粹的高性能：** 基于内存流转，零外部依赖，极低的上下文切换开销。 @@ -33,36 +36,34 @@ type MyState struct { Query string Response string Value int - Log []string + Messages []string // 含引用类型，需要配合 StateCloner 使用 } ``` -#### 2. 并发分支与合并示例 +#### 2. 直接使用 CompiledGraph（原始 API，向后兼容） ```go package main import ( "context" - "fmt" - "GopherGraph" + GopherGraph "github.com/unclesam-LY/GopherGraph" ) func main() { g := GopherGraph.NewGraph[MyState]() - // 注册节点 g.AddNode("start", startNode) g.AddNode("task1", task1Node) g.AddNode("task2", task2Node) g.AddNode("end", endNode) - // 定义合并函数：将并发分支状态合并到主状态中 - merger := func(ctx context.Context, parent MyState, branches []MyState) (MyState, error) { - for _, b := range branches { - parent.Value += b.Value - } - return parent, nil - } + // 定义合并函数：将并发分支状态合并到主状态中 + merger := func(ctx context.Context, parent MyState, branches []MyState) (MyState, error) { + for _, b := range branches { + parent.Value += b.Value + } + return parent, nil + } // 建立并发边：从 start 并发分流到 task1 和 task2，执行完后通过 merger 合并状态并去往 end 节点 g.AddParallelEdges("start", []string{"task1", "task2"}, "end", merger) @@ -72,7 +73,56 @@ func main() { } ``` -#### 3. 文件持久化 (Checkpointer) 示例 +#### 3. 使用 Engine 包装器（推荐：生产级增强） +```go +cg, _ := g.Compile() + +engine := GopherGraph.NewEngine(cg). + // 深拷贝：消除并发分支对 Messages 切片的数据竞争 + WithStateCloner(func(s MyState) MyState { + clone := s + clone.Messages = append([]string{}, s.Messages...) + return clone + }). + // 步数熔断：防御图中意外形成的死循环 + WithMaxSteps(100). + // Pre Hook：节点执行前触发（日志、链路追踪、流式推送等） + WithPreNodeHook(func(ctx context.Context, name string, s MyState) { + log.Printf("[PRE] node=%-12s query=%q", name, s.Query) + }). + // Post Hook：节点执行后触发（指标上报、状态快照等） + WithPostNodeHook(func(ctx context.Context, name string, s MyState) { + log.Printf("[POST] node=%-12s value=%d", name, s.Value) + }) + +thread, err := engine.Start(ctx, "start", MyState{Query: "你好"}) +``` + +#### 4. 流式输出（通过 Context 注入 Channel） +不需要修改任何 `NodeFn` 签名，流式输出是节点的副作用，通过 Context 传递 channel 即可： +```go +type streamKey struct{} + +streamCh := make(chan string, 10) +ctx := context.WithValue(context.Background(), streamKey{}, streamCh) + +// 在 Post Hook 里从 Context 取出 channel，写入流式事件 +engine := GopherGraph.NewEngine(cg). + WithPostNodeHook(func(ctx context.Context, name string, s MyState) { + if ch, ok := ctx.Value(streamKey{}).(chan string); ok { + ch <- fmt.Sprintf("节点 [%s] 完成", name) + } + }) + +go func() { + for msg := range streamCh { fmt.Println(msg) } // 前端消费 +}() + +engine.Start(ctx, "start", MyState{}) +close(streamCh) +``` + +#### 5. 文件持久化 (Checkpointer) 示例 ```go // 创建一个文件存储器，指定存放状态文件的目录 fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints") @@ -87,14 +137,18 @@ thread, _ = cg.Resume(context.Background(), loadedThread, modifiedState) ``` #### 运行内置示例 -项目包含了一个完整的“AI翻译 -> 质量检测 -> 人工审核 -> 发布”的交互式 Demo： ```bash +# 完整的"AI翻译 -> 质量检测 -> 人工审核 -> 发布"交互式 Demo go run examples/translation/main.go + +# Engine 增强特性演示（Hooks、StateCloner、流式输出） +go run examples/hooks/main.go ``` #### 运行单元测试 ```bash -go test -v ./... +# 含竞态检测器，确保并发分支无数据竞争 +go test -v -race ./... ``` --- @@ -107,8 +161,11 @@ go test -v ./... ### Key Features - **Strictly Typed State:** Bind your custom struct to `Graph[S]` via Go generics. Say goodbye to unsafe `map[string]any` and runtime type assertions. - **Support for Cycles/Loops:** Unlike traditional DAG (Directed Acyclic Graph) engines, GopherGraph natively supports loops and dynamic routing based on agent evaluation. -- **Parallel Branching & Merging:** Run multiple agent nodes concurrently using `sync.WaitGroup`. Thread-safety is achieved through value-copying on fan-out and a custom `ParallelMergeFn` on fan-in. +- **Parallel Branching with Short-Circuit Cancellation:** Run multiple agent nodes concurrently. Built on `context.WithCancel`, a failure in any branch immediately cancels all sibling goroutines, preventing wasted compute on LLM API calls. +- **Deep-Clone for Concurrency Safety:** Inject a custom `StateCloner` via `Engine.WithStateCloner` to safely deep-copy reference types (slices, maps, pointers) before fan-out, eliminating data races at the root. - **Human-in-the-Loop (HITL):** Pause workflow execution *before* a designated node, capture a snapshot (`Thread`), modify the state, and `Resume` seamlessly. +- **Infinite Loop Circuit Breaker:** Set a hard step limit with `Engine.WithMaxSteps`. Any accidental cycle in the graph returns a clear error instead of blocking forever. +- **Lightweight Lifecycle Hooks:** `Engine.WithPreNodeHook` / `WithPostNodeHook` fire before and after each node without touching any `NodeFn` signature — perfect for logging, OpenTelemetry tracing, or streaming output. - **State Persistence (Checkpointer):** Generic `Checkpointer[S]` interface with built-in `FileCheckpointer[S]` for saving and loading execution snapshots to/from JSON files. - **Zero-Dependency & High-Performance:** Written in pure Go with zero external dependencies, leveraging in-memory queues for lightning-fast orchestration. @@ -120,18 +177,17 @@ type MyState struct { Query string Response string Value int - Log []string + Messages []string // reference type — pair with StateCloner for concurrency safety } ``` -#### 2. Parallel execution example +#### 2. Use CompiledGraph directly (original API, fully backwards-compatible) ```go package main import ( "context" - "fmt" - "GopherGraph" + GopherGraph "github.com/unclesam-LY/GopherGraph" ) func main() { @@ -158,7 +214,55 @@ func main() { } ``` -#### 3. State Persistence (Checkpointer) Example +#### 3. Use the Engine wrapper (recommended for production) +```go +cg, _ := g.Compile() + +engine := GopherGraph.NewEngine(cg). + // Deep-copy state before fan-out to eliminate data races on reference types + WithStateCloner(func(s MyState) MyState { + clone := s + clone.Messages = append([]string{}, s.Messages...) + return clone + }). + // Hard step limit — prevents accidental infinite loops + WithMaxSteps(100). + // Pre-hook: fires before each node (logging, tracing, streaming notifications) + WithPreNodeHook(func(ctx context.Context, name string, s MyState) { + log.Printf("[PRE] node=%-12s query=%q", name, s.Query) + }). + // Post-hook: fires after each node (metrics, snapshots) + WithPostNodeHook(func(ctx context.Context, name string, s MyState) { + log.Printf("[POST] node=%-12s value=%d", name, s.Value) + }) + +thread, err := engine.Start(ctx, "start", MyState{Query: "hello"}) +``` + +#### 4. Streaming output via Context-injected channel +No need to change any `NodeFn` signature. Streaming is a side effect — pass a channel through the Context: +```go +type streamKey struct{} + +streamCh := make(chan string, 10) +ctx := context.WithValue(context.Background(), streamKey{}, streamCh) + +engine := GopherGraph.NewEngine(cg). + WithPostNodeHook(func(ctx context.Context, name string, s MyState) { + if ch, ok := ctx.Value(streamKey{}).(chan string); ok { + ch <- fmt.Sprintf("node [%s] done", name) + } + }) + +go func() { + for msg := range streamCh { fmt.Println(msg) } // consume on the frontend +}() + +engine.Start(ctx, "start", MyState{}) +close(streamCh) +``` + +#### 5. State Persistence (Checkpointer) Example ```go // Initialize a local directory storage fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints") @@ -173,14 +277,18 @@ thread, _ = cg.Resume(context.Background(), loadedThread, modifiedState) ``` #### Run the Interactive Demo -An interactive "Translation -> Review -> Human Approval -> Publish" demo is included: ```bash +# Interactive "Translation -> Review -> Human Approval -> Publish" demo go run examples/translation/main.go + +# Engine features demo (Hooks, StateCloner, Streaming) +go run examples/hooks/main.go ``` #### Run Unit Tests ```bash -go test -v ./... +# The -race flag verifies zero data races in parallel branches +go test -v -race ./... ``` --- @@ -193,8 +301,11 @@ go test -v ./... ### 主な特徴 - **型安全な状態管理:** ジェネリクスを用いて `Graph[S]` に構造体をバインドします。冗長な `map[string]any` やランタイム時の型アサーションから解放されます。 - **ループと条件付き分岐のサポート:** 従来の DAG（有向非巡回グラフ）の制限を超え、エージェントの判定に基づく条件付きルートや双方向のループをサポート。 -- **並行処理とマージ（分流と合流）:** `sync.WaitGroup` に基づく複数エージェントノードの並行実行に対応。分流時の値渡しコピーによるスレッドセーフの確保と、合流時の `ParallelMergeFn` による状態マージを実現。 +- **並行処理とショートサーキットキャンセル:** `context.WithCancel` に基づく並行ブランチ実行を実装。いずれかのブランチでエラーが発生した瞬間に他の全 goroutine へキャンセルシグナルを伝播し、LLM API の無駄な呼び出しを防ぎます。 +- **データ競合を防ぐ深いコピー（DeepClone）:** `Engine.WithStateCloner` でカスタムの深コピー関数を注入。参照型（スライス・マップ・ポインタ）を持つ状態でも、並行ブランチでのデータ競合を根本から解消します。 - **Human-in-the-Loop (HITL) の一時停止と再開:** 特定のノードの実行前に処理を自動停止し、スナップショット（`Thread`）を返します。承認や状態修正の後にシームレスに処理を再開（`Resume`）できます。 +- **無限ループのサーキットブレーカー:** `Engine.WithMaxSteps` でステップ数の上限を設定。グラフ内に意図しないサイクルが発生しても、永久ブロックせず明確なエラーを返します。 +- **軽量なライフサイクルフック:** `Engine.WithPreNodeHook` / `WithPostNodeHook` により、`NodeFn` のシグネチャを一切変更せずにロギング・トレーシング・ストリーミング出力を実装できます。 - **状態の永続化 (Checkpointer):** 抽象的な `Checkpointer[S]` インタフェースと、JSONファイルへの書き出し・読み込みを行う組み込みの `FileCheckpointer[S]` をサポート。サーバー再起動後の進捗復旧が可能。 - **ピュア Go & 高性能:** 外部依存関係ゼロ。メモリ内チャネルを用いた高速なコンテキスト切り替え。 @@ -206,18 +317,17 @@ type MyState struct { Query string Response string Value int - Log []string + Messages []string // 参照型 — 並行処理では StateCloner と組み合わせて使用 } ``` -#### 2. 並行処理（Parallel）のコード例 +#### 2. CompiledGraph を直接使う（元の API、完全後方互換） ```go package main import ( "context" - "fmt" - "GopherGraph" + GopherGraph "github.com/unclesam-LY/GopherGraph" ) func main() { @@ -244,7 +354,55 @@ func main() { } ``` -#### 3. 状態の永続化（Checkpointer）のコード例 +#### 3. Engine ラッパーを使う（本番環境向け・推奨） +```go +cg, _ := g.Compile() + +engine := GopherGraph.NewEngine(cg). + // ファンアウト前に参照型を深くコピーしてデータ競合を排除 + WithStateCloner(func(s MyState) MyState { + clone := s + clone.Messages = append([]string{}, s.Messages...) + return clone + }). + // ステップ上限 — 意図しない無限ループを防止 + WithMaxSteps(100). + // Pre フック: 各ノード実行前に発火（ロギング・トレーシング・ストリーミング通知） + WithPreNodeHook(func(ctx context.Context, name string, s MyState) { + log.Printf("[PRE] node=%-12s query=%q", name, s.Query) + }). + // Post フック: 各ノード実行後に発火（メトリクス・スナップショット） + WithPostNodeHook(func(ctx context.Context, name string, s MyState) { + log.Printf("[POST] node=%-12s value=%d", name, s.Value) + }) + +thread, err := engine.Start(ctx, "start", MyState{Query: "こんにちは"}) +``` + +#### 4. Context 経由の Channel によるストリーミング出力 +`NodeFn` のシグネチャを変更する必要はありません。ストリーミングは副作用として Context 経由で channel を渡すだけです： +```go +type streamKey struct{} + +streamCh := make(chan string, 10) +ctx := context.WithValue(context.Background(), streamKey{}, streamCh) + +engine := GopherGraph.NewEngine(cg). + WithPostNodeHook(func(ctx context.Context, name string, s MyState) { + if ch, ok := ctx.Value(streamKey{}).(chan string); ok { + ch <- fmt.Sprintf("ノード [%s] 完了", name) + } + }) + +go func() { + for msg := range streamCh { fmt.Println(msg) } // フロントエンド側で消費 +}() + +engine.Start(ctx, "start", MyState{}) +close(streamCh) +``` + +#### 5. 状態の永続化（Checkpointer）のコード例 ```go // 保存先ディレクトリを指定してローカルファイルチェッカーを初期化 fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints") @@ -259,12 +417,16 @@ thread, _ = cg.Resume(context.Background(), loadedThread, modifiedState) ``` #### インタラクティブデモの実行 -「AI翻訳 -> 監査 -> 人間による確認 -> 公開」の一連のループを体験できるデモを用意しています： ```bash +# 「AI翻訳 -> 監査 -> 人間による確認 -> 公開」のインタラクティブデモ go run examples/translation/main.go + +# Engine 機能デモ（Hooks・StateCloner・ストリーミング） +go run examples/hooks/main.go ``` #### ユニットテストの実行 ```bash -go test -v ./... +# -race フラグで並行ブランチのデータ競合がゼロであることを検証 +go test -v -race ./... ``` diff --git a/engine.go b/engine.go index 64e9f74..5f4d64b 100644 --- a/engine.go +++ b/engine.go @@ -96,7 +96,7 @@ func (cg *CompiledGraph[S]) Start(ctx context.Context, startNode string, initial NextNode: startNode, } - return cg.run(ctx, thread) + return cg.run(ctx, thread, runOptions[S]{}) } // Resume 恢复执行一个被暂停的线程，并允许注入外部修改后的状态数据（例如人工审批修改后的结果） @@ -112,15 +112,38 @@ func (cg *CompiledGraph[S]) Resume(ctx context.Context, thread *Thread[S], modif thread.State = modifiedState thread.IsPaused = false - return cg.run(ctx, thread) + return cg.run(ctx, thread, runOptions[S]{}) +} + +// runOptions 封装引擎运行时的可选参数，由 Engine 包装器传入 +type runOptions[S any] struct { + maxSteps int + stateCloner func(S) S + preNodeHook HookFn[S] + postNodeHook HookFn[S] } // run 是引擎内部循环调度器，负责驱动节点向前流转 -func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S]) (*Thread[S], error) { +func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S], opts runOptions[S]) (*Thread[S], error) { + stepCount := 0 + for { - // 检查 Context，支持外部超时控制或取消 - if err := ctx.Err(); err != nil { - return thread, err + // 【熔断 1】检查步数上限，防御死循环（infinite loop protection） + if opts.maxSteps > 0 { + if stepCount >= opts.maxSteps { + return thread, fmt.Errorf( + "engine halt: max steps (%d) exceeded, possible infinite loop in graph", + opts.maxSteps, + ) + } + stepCount++ + } + + // 【熔断 2】检查 Context，支持外部超时控制或取消 + select { + case <-ctx.Done(): + return thread, ctx.Err() + default: } currentNodeName := thread.NextNode @@ -130,51 +153,38 @@ func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S]) (*Thread return thread, nil } - // 查找并执行当前节点 + // 查找当前节点 nodeFn, exists := cg.nodes[currentNodeName] if !exists { return thread, fmt.Errorf("runtime error: node %q not found", currentNodeName) } + // 触发 PreNodeHook + if opts.preNodeHook != nil { + opts.preNodeHook(ctx, currentNodeName, thread.State) + } + + // 执行当前节点 newState, err := nodeFn(ctx, thread.State) if err != nil { return thread, fmt.Errorf("node %q execution error: %w", currentNodeName, err) } thread.State = newState + // 触发 PostNodeHook + if opts.postNodeHook != nil { + opts.postNodeHook(ctx, currentNodeName, thread.State) + } + // 计算下一个该执行的节点名称 var nextNode string if step, isParallel := cg.parallels[currentNodeName]; isParallel { - // 【处理并发分流】 - var wg sync.WaitGroup - branches := make([]S, len(step.targets)) - errs := make([]error, len(step.targets)) - for i, target := range step.targets { - wg.Add(1) - // 启动并发协程运行分支节点。传入 thread.State 的副本（Go 默认是值拷贝） - go func(idx int, targetNode string, stateCopy S) { - defer wg.Done() - if ctx.Err() != nil { - errs[idx] = ctx.Err() - return - } - nodeFn := cg.nodes[targetNode] - resState, err := nodeFn(ctx, stateCopy) - if err != nil { - errs[idx] = err - return - } - branches[idx] = resState - }(i, target, thread.State) - } - // 等待所有分支协程执行完毕 - wg.Wait() - // 检查是否有任何分支报错 - for _, err := range errs { - if err != nil { - return thread, fmt.Errorf("parallel branch execution error: %w", err) - } + // 【处理并发分流】使用短路取消机制 + branches, err := cg.runParallelBranches(ctx, step.targets, thread.State, opts.stateCloner) + if err != nil { + return thread, fmt.Errorf("parallel branch execution error: %w", err) } + // 【状态合并】调用用户自定义的合并函数 mergedState, err := step.merger(ctx, thread.State, branches) if err != nil { @@ -194,7 +204,7 @@ func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S]) (*Thread nextNode = cg.edges[currentNodeName] } - // 更新快照中的“下一个节点” + // 更新快照中的"下一个节点" thread.NextNode = nextNode // 如果下一站是终点，直接进入下一次循环触发结束逻辑 @@ -209,3 +219,82 @@ func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S]) (*Thread } } } + +// runParallelBranches 并发执行多个目标节点，并实现首个错误的短路取消。 +// +// - 使用 context.WithCancel 派生子 Context，任意分支报错立即调用 cancel()， +// 通知其他仍在运行的分支通过 ctx.Done() 优雅退出，避免 AI 算力空转。 +// - 使用缓冲 channel（容量 = 分支数）收集结果，防止 goroutine 在写入时永久阻塞（goroutine 泄漏）。 +// - stateCloner 为可选的深拷贝函数；若为 nil，则传入值拷贝（对含引用类型的状态存在数据竞争风险，需用户知悉）。 +func (cg *CompiledGraph[S]) runParallelBranches( + ctx context.Context, + targets []string, + state S, + stateCloner func(S) S, +) ([]S, error) { + // 派生可取消的子 Context：任意分支出错时，立即广播取消信号 + cancelCtx, cancel := context.WithCancel(ctx) + defer cancel() // 确保函数退出时释放资源，不论成功或失败 + + type result struct { + idx int + state S + err error + } + + // 缓冲通道容量 = 分支数，确保所有 goroutine 均可不阻塞地写入结果 + resultCh := make(chan result, len(targets)) + + var wg sync.WaitGroup + for i, target := range targets { + wg.Add(1) + + // 按需深拷贝状态：有 stateCloner 则深拷贝，否则值拷贝（Go 默认语义） + var stateCopy S + if stateCloner != nil { + stateCopy = stateCloner(state) + } else { + stateCopy = state + } + + go func(idx int, nodeName string, s S) { + defer wg.Done() + + // 在执行真正的工作之前，先检查是否已被其他分支取消 + select { + case <-cancelCtx.Done(): + resultCh <- result{idx: idx, err: cancelCtx.Err()} + return + default: + } + + nodeFn := cg.nodes[nodeName] + resState, err := nodeFn(cancelCtx, s) + if err != nil { + cancel() // 关键：一处报错，立即取消其余所有兄弟 goroutine + resultCh <- result{idx: idx, err: err} + return + } + resultCh <- result{idx: idx, state: resState} + }(i, target, stateCopy) + } + + // 等待所有 goroutine 完成后关闭结果通道 + wg.Wait() + close(resultCh) + + // 按索引顺序收集结果，保证合并函数收到的 branches 顺序与 targets 一致 + branches := make([]S, len(targets)) + for r := range resultCh { + if r.err != nil { + // 过滤掉因取消产生的次生错误，只返回第一个真实错误 + if !errors.Is(r.err, context.Canceled) { + return nil, r.err + } + } else { + branches[r.idx] = r.state + } + } + + return branches, nil +} diff --git a/engine_options.go b/engine_options.go new file mode 100644 index 0000000..810db0e --- /dev/null +++ b/engine_options.go @@ -0,0 +1,140 @@ +package GopherGraph + +import "context" + +// Engine 是 CompiledGraph 的增强包装器，在完全兼容原有 API 的基础上， +// 提供以下四项生产级能力： +// +// 1. StateCloner — 自定义深拷贝逻辑，消除并发分支的数据竞争隐患 +// 2. MaxSteps — 步数硬性熔断，防止图中的环路导致无限循环 +// 3. PreNodeHook — 节点执行前的生命周期钩子（日志、链路追踪等） +// 4. PostNodeHook — 节点执行后的生命周期钩子（指标上报、流式推送等） +// +// 所有选项均为可选，通过链式调用注册。未注册的选项将退回到 CompiledGraph 的默认行为。 +// +// 用法示例： +// +// engine := GopherGraph.NewEngine(compiledGraph). +// WithStateCloner(func(s MyState) MyState { +// clone := s +// clone.Messages = append([]string{}, s.Messages...) +// return clone +// }). +// WithMaxSteps(100). +// WithPreNodeHook(func(ctx context.Context, name string, s MyState) { +// log.Printf("[PRE] node=%s", name) +// }). +// WithPostNodeHook(func(ctx context.Context, name string, s MyState) { +// log.Printf("[POST] node=%s value=%+v", name, s) +// }) +// +// thread, err := engine.Start(ctx, "start", initialState) +type Engine[S any] struct { + graph *CompiledGraph[S] + stateCloner func(S) S + maxSteps int + preNodeHook HookFn[S] + postNodeHook HookFn[S] +} + +// NewEngine 用一个已编译的图创建 Engine 实例。 +func NewEngine[S any](graph *CompiledGraph[S]) *Engine[S] { + return &Engine[S]{graph: graph} +} + +// WithStateCloner 注册自定义深拷贝函数。 +// +// 当图中存在并发分支（AddParallelEdges）时，引擎会在启动每个并发 goroutine 之前 +// 调用此函数对共享状态进行深拷贝，从根本上消除引用类型（切片、map、指针）的数据竞争。 +// +// 若不注册，引擎退化为 Go 默认的值拷贝语义（浅拷贝）。 +// 对于状态中含有引用类型的场景，强烈建议注册此函数。 +func (e *Engine[S]) WithStateCloner(fn func(S) S) *Engine[S] { + e.stateCloner = fn + return e +} + +// WithMaxSteps 设置引擎单次运行（Start/Resume）允许的最大节点执行步数。 +// +// 这是防御图中存在意外环路（infinite loop）的最后一道硬性熔断。 +// 超过上限后，引擎将返回明确的错误，而不是永久阻塞。 +// +// 建议根据工作流的预期最大深度合理设置（如 100~1000）。 +// 传入 0 或负数表示不限制（与不调用此方法等效）。 +func (e *Engine[S]) WithMaxSteps(n int) *Engine[S] { + e.maxSteps = n + return e +} + +// WithPreNodeHook 注册在每个节点执行【之前】触发的钩子函数。 +// +// 典型用途： +// - 打印结构化日志（node 名称、输入状态） +// - 记录 OpenTelemetry Span 的开始时间 +// - 向前端推送"正在执行 xxx 节点"的流式事件 +// +// 钩子函数接收到的 state 是节点执行【前】的状态快照，请勿修改（只读语义）。 +// 钩子中的 panic 会向上传播，建议在内部做好 recover。 +func (e *Engine[S]) WithPreNodeHook(fn HookFn[S]) *Engine[S] { + e.preNodeHook = fn + return e +} + +// WithPostNodeHook 注册在每个节点执行【之后】触发的钩子函数。 +// +// 典型用途： +// - 打印节点的输出状态（用于 Debug） +// - 上报节点耗时指标（Prometheus/DataDog） +// - 向前端推送节点执行完毕的流式事件 +// +// 钩子函数接收到的 state 是节点执行【后】更新过的状态，同为只读语义。 +func (e *Engine[S]) WithPostNodeHook(fn HookFn[S]) *Engine[S] { + e.postNodeHook = fn + return e +} + +// opts 将 Engine 的配置打包成 runOptions，传入底层 run() 调度器。 +func (e *Engine[S]) opts() runOptions[S] { + return runOptions[S]{ + maxSteps: e.maxSteps, + stateCloner: e.stateCloner, + preNodeHook: e.preNodeHook, + postNodeHook: e.postNodeHook, + } +} + +// Start 从指定节点启动工作流，行为与 CompiledGraph.Start 完全一致， +// 但会附加所有已注册的 Engine 增强选项（Hooks、MaxSteps、StateCloner）。 +func (e *Engine[S]) Start(ctx context.Context, startNode string, initialState S) (*Thread[S], error) { + thread := &Thread[S]{ + State: initialState, + NextNode: startNode, + } + return e.graph.run(ctx, thread, e.opts()) +} + +// Resume 恢复执行一个被暂停的线程，行为与 CompiledGraph.Resume 完全一致， +// 但会附加所有已注册的 Engine 增强选项（Hooks、MaxSteps、StateCloner）。 +func (e *Engine[S]) Resume(ctx context.Context, thread *Thread[S], modifiedState S) (*Thread[S], error) { + if !thread.IsPaused { + return nil, errNotPaused + } + if thread.IsFinished { + return nil, errAlreadyFinished + } + + thread.State = modifiedState + thread.IsPaused = false + + return e.graph.run(ctx, thread, e.opts()) +} + +// 复用与 CompiledGraph.Resume 中相同的哨兵错误，避免魔法字符串重复 +var ( + errNotPaused = errStr("cannot resume: thread is not paused") + errAlreadyFinished = errStr("cannot resume: thread is already finished") +) + +type errStr string + +func (e errStr) Error() string { return string(e) } diff --git a/examples/hooks/main.go b/examples/hooks/main.go new file mode 100644 index 0000000..598b110 --- /dev/null +++ b/examples/hooks/main.go @@ -0,0 +1,159 @@ +package main + +import ( + "context" + "fmt" + "log" + "time" + + GopherGraph "github.com/unclesam-LY/GopherGraph" +) + +// AgentState 是这个示例工作流的状态 +type AgentState struct { + Input string + Draft string + Reviewed string + Messages []string // 含引用类型，演示 StateCloner 的必要性 +} + +// deepClone 演示如何为含切片的状态做安全的深拷贝 +func deepClone(s AgentState) AgentState { + clone := s + clone.Messages = append([]string{}, s.Messages...) // 切片深拷贝 + return clone +} + +func main() { + g := GopherGraph.NewGraph[AgentState]() + + // 节点 1：起草 + g.AddNode("drafter", func(ctx context.Context, s AgentState) (AgentState, error) { + time.Sleep(80 * time.Millisecond) // 模拟耗时 AI 调用 + s.Draft = fmt.Sprintf("【草稿】%s 的初版内容", s.Input) + s.Messages = append(s.Messages, "drafter: 完成起草") + return s, nil + }) + + // 节点 2：审阅 + g.AddNode("reviewer", func(ctx context.Context, s AgentState) (AgentState, error) { + time.Sleep(60 * time.Millisecond) + s.Reviewed = s.Draft + "（已审阅 ✓）" + s.Messages = append(s.Messages, "reviewer: 审阅完成") + return s, nil + }) + + // 并发节点：同时翻译成英文和日文（演示 StateCloner） + g.AddNode("translate_en", func(ctx context.Context, s AgentState) (AgentState, error) { + time.Sleep(100 * time.Millisecond) + s.Messages = append(s.Messages, fmt.Sprintf("translate_en: [EN] %s", s.Reviewed)) + return s, nil + }) + g.AddNode("translate_ja", func(ctx context.Context, s AgentState) (AgentState, error) { + time.Sleep(120 * time.Millisecond) + s.Messages = append(s.Messages, fmt.Sprintf("translate_ja: [JA] %s", s.Reviewed)) + return s, nil + }) + + // 节点：发布 + g.AddNode("publisher", func(ctx context.Context, s AgentState) (AgentState, error) { + s.Messages = append(s.Messages, "publisher: 🎉 发布成功") + return s, nil + }) + + // 连线 + g.AddEdge("drafter", "reviewer") + g.AddParallelEdges("reviewer", []string{"translate_en", "translate_ja"}, "publisher", + func(ctx context.Context, parent AgentState, branches []AgentState) (AgentState, error) { + // 合并所有分支的 Messages + for _, b := range branches { + parent.Messages = append(parent.Messages, b.Messages...) + } + parent.Messages = append(parent.Messages, "merger: 翻译分支已合并") + return parent, nil + }, + ) + + cg, err := g.Compile() + if err != nil { + log.Fatalf("编译失败: %v", err) + } + + // ============================================================ + // 使用 Engine 包装器，挂载所有增强能力 + // ============================================================ + var nodeTimings = make(map[string]time.Time) + + engine := GopherGraph.NewEngine(cg). + // 1. 注册深拷贝函数，消除并发分支对 Messages 切片的数据竞争 + WithStateCloner(deepClone). + // 2. 设置最大步数熔断（防御意外死循环） + WithMaxSteps(50). + // 3. Pre Hook：记录节点开始时间 + 向控制台推送"流式进度" + WithPreNodeHook(func(ctx context.Context, name string, s AgentState) { + nodeTimings[name] = time.Now() + fmt.Printf(" ⏳ [PRE ] %-15s | 消息数: %d\n", name, len(s.Messages)) + }). + // 4. Post Hook：计算并打印节点耗时 + WithPostNodeHook(func(ctx context.Context, name string, s AgentState) { + elapsed := time.Since(nodeTimings[name]) + fmt.Printf(" ✅ [POST] %-15s | 耗时: %v | 消息数: %d\n", name, elapsed.Round(time.Millisecond), len(s.Messages)) + }) + + fmt.Println("🚀 启动工作流...") + fmt.Println("─────────────────────────────────────────────") + + totalStart := time.Now() + thread, err := engine.Start(context.Background(), "drafter", AgentState{ + Input: "GopherGraph 技术文档", + }) + if err != nil { + log.Fatalf("运行失败: %v", err) + } + + fmt.Println("─────────────────────────────────────────────") + fmt.Printf("🏁 工作流结束，总耗时: %v\n\n", time.Since(totalStart).Round(time.Millisecond)) + fmt.Println("📋 最终消息日志：") + for i, msg := range thread.State.Messages { + fmt.Printf(" %d. %s\n", i+1, msg) + } + + // ============================================================ + // 演示流式输出：通过 Context 注入 channel（不需要修改任何 NodeFn） + // ============================================================ + fmt.Println("\n─────────────────────────────────────────────") + fmt.Println("📡 演示：通过 Context 注入 channel 实现流式输出") + fmt.Println("─────────────────────────────────────────────") + + type streamCtxKey struct{} + streamCh := make(chan string, 10) + + // 在 Context 里注入 channel + streamCtx := context.WithValue(context.Background(), streamCtxKey{}, streamCh) + + // 启动一个 goroutine 消费流式事件（模拟前端 SSE/WebSocket） + go func() { + for msg := range streamCh { + fmt.Printf(" 📨 [STREAM] %s\n", msg) + } + }() + + // Post Hook 里从 Context 取出 channel，写入流式事件 + streamEngine := GopherGraph.NewEngine(cg). + WithStateCloner(deepClone). + WithPostNodeHook(func(ctx context.Context, name string, s AgentState) { + ch, ok := ctx.Value(streamCtxKey{}).(chan string) + if ok { + ch <- fmt.Sprintf("节点 [%s] 执行完毕，当前消息数: %d", name, len(s.Messages)) + } + }) + + _, err = streamEngine.Start(streamCtx, "drafter", AgentState{Input: "流式演示"}) + if err != nil { + log.Fatalf("流式演示运行失败: %v", err) + } + close(streamCh) + time.Sleep(50 * time.Millisecond) // 等待消费者 goroutine 打印完毕 + + fmt.Println("\n✨ 示例结束") +} diff --git a/examples/translation/main.go b/examples/translation/main.go index e2e6400..e08e03a 100644 --- a/examples/translation/main.go +++ b/examples/translation/main.go @@ -1,12 +1,13 @@ package main import ( - "GopherGraph" "bufio" "context" "fmt" "os" "strings" + + GopherGraph "github.com/unclesam-LY/GopherGraph" ) // TranslationState 代表我们工作流的共享状态（State） diff --git a/graph.go b/graph.go index 24ce836..98d67ad 100644 --- a/graph.go +++ b/graph.go @@ -73,3 +73,8 @@ func (g *Graph[S]) AddParallelEdges(from string, targets []string, next string, merger: merger, } } + +// HookFn 是节点执行前后的生命周期钩子函数签名。 +// nodeName 为当前执行的节点名称，state 为该节点执行时的状态快照（只读语义，请勿修改）。 +// 可通过此钩子接入日志、链路追踪（OpenTelemetry）、监控指标等能力，而无需修改任何 NodeFn。 +type HookFn[S any] func(ctx context.Context, nodeName string, state S) diff --git a/graph_test.go b/graph_test.go index 00b7e72..cf9bae1 100644 --- a/graph_test.go +++ b/graph_test.go @@ -317,3 +317,138 @@ func TestFileCheckpointer(t *testing.T) { t.Errorf("执行结果错误，最终值: %d, 是否结束: %t", thread.State.Value, thread.IsFinished) } } + +// TestMaxStepsCircuitBreaker 验证 Engine.WithMaxSteps 能在步数超限时熔断，防止死循环 +func TestMaxStepsCircuitBreaker(t *testing.T) { + g := NewGraph[TestState]() + + // 构造一个永远在 A -> B -> A 之间循环的图 + g.AddNode("A", func(ctx context.Context, s TestState) (TestState, error) { + s.Value++ + return s, nil + }) + g.AddNode("B", func(ctx context.Context, s TestState) (TestState, error) { + s.Value++ + return s, nil + }) + + // A -> B -> A：形成闭环 + g.AddEdge("A", "B") + g.AddEdge("B", "A") + + cg, err := g.Compile() + if err != nil { + t.Fatalf("编译图失败: %v", err) + } + + // 设置最大步数为 5，图会在第 5 步触发熔断 + engine := NewEngine(cg).WithMaxSteps(5) + _, err = engine.Start(context.Background(), "A", TestState{}) + if err == nil { + t.Fatal("期望返回步数超限错误，实际无错误返回") + } + if !strings.Contains(err.Error(), "max steps") { + t.Errorf("期望错误信息包含 'max steps'，实际: %v", err) + } +} + +// TestConcurrentCancellation 验证并发分支中一个节点报错时，其他兄弟节点能被及时取消 +func TestConcurrentCancellation(t *testing.T) { + g := NewGraph[TestState]() + + g.AddNode("start", func(ctx context.Context, s TestState) (TestState, error) { + return s, nil + }) + // task_fail：立即返回错误 + g.AddNode("task_fail", func(ctx context.Context, s TestState) (TestState, error) { + return s, errors.New("task_fail intentional error") + }) + // task_long：模拟耗时操作，应感知取消并提前退出 + g.AddNode("task_long", func(ctx context.Context, s TestState) (TestState, error) { + select { + case <-time.After(5 * time.Second): // 5 秒，正常不应该跑完 + s.Value = 999 + return s, nil + case <-ctx.Done(): + return s, ctx.Err() + } + }) + + g.AddParallelEdges("start", []string{"task_fail", "task_long"}, "", func(_ context.Context, parent TestState, _ []TestState) (TestState, error) { + return parent, nil + }) + + cg, err := g.Compile() + if err != nil { + t.Fatalf("编译图失败: %v", err) + } + + start := time.Now() + _, err = cg.Start(context.Background(), "start", TestState{}) + elapsed := time.Since(start) + + // 验证确实返回了错误 + if err == nil { + t.Fatal("期望返回并发分支错误，实际无错误返回") + } + if !strings.Contains(err.Error(), "task_fail intentional error") { + t.Errorf("期望错误来自 task_fail，实际: %v", err) + } + // 验证 task_long 被快速取消，整个并发组的耗时远小于 5 秒 + if elapsed > 2*time.Second { + t.Errorf("期望并发组在取消后快速结束（< 2s），实际耗时: %v", elapsed) + } +} + +// TestLifecycleHooks 验证 Engine 的 Pre/Post 钩子在每个节点执行前后均被正确触发 +func TestLifecycleHooks(t *testing.T) { + g := NewGraph[TestState]() + + g.AddNode("node1", func(ctx context.Context, s TestState) (TestState, error) { + s.Value += 1 + return s, nil + }) + g.AddNode("node2", func(ctx context.Context, s TestState) (TestState, error) { + s.Value += 10 + return s, nil + }) + g.AddEdge("node1", "node2") + + cg, err := g.Compile() + if err != nil { + t.Fatalf("编译图失败: %v", err) + } + + var hookLog []string + + engine := NewEngine(cg). + WithPreNodeHook(func(ctx context.Context, name string, s TestState) { + hookLog = append(hookLog, "pre:"+name) + }). + WithPostNodeHook(func(ctx context.Context, name string, s TestState) { + hookLog = append(hookLog, "post:"+name) + }) + + thread, err := engine.Start(context.Background(), "node1", TestState{}) + if err != nil { + t.Fatalf("启动图失败: %v", err) + } + if !thread.IsFinished { + t.Error("期望工作流运行结束") + } + // 最终值：node1(+1) -> node2(+10) = 11 + if thread.State.Value != 11 { + t.Errorf("期望 Value 为 11，实际为 %d", thread.State.Value) + } + + // 验证 Hook 触发顺序：pre:node1 -> post:node1 -> pre:node2 -> post:node2 + expected := []string{"pre:node1", "post:node1", "pre:node2", "post:node2"} + if len(hookLog) != len(expected) { + t.Fatalf("期望 %d 次 hook 触发，实际 %d 次: %v", len(expected), len(hookLog), hookLog) + } + for i, entry := range expected { + if hookLog[i] != entry { + t.Errorf("第 %d 次 hook 期望 %q，实际 %q", i, entry, hookLog[i]) + } + } +}