14 Commits
Author SHA1 Message Date
unclesam b20d6a3750 fix(sqlite_checkpoint): harden concurrent safety and close lifecycle
- validateTableName: guard against empty string input to prevent malformed SQL
- Save/PruneVersions/Delete: use mu.Lock to serialize SELECT MAX(version)+INSERT
  preventing duplicate version numbers under concurrent goroutines
- Load/LoadVersion/ListVersions: retain mu.RLock to allow concurrent reads
- All methods: add nil db check after acquiring lock to fail fast after Close()
- NewSQLiteCheckpointer: set SetMaxOpenConns(1)+SetMaxIdleConns(1) to ensure
  PRAGMA WAL/busy_timeout takes effect on the single shared connection and
  eliminate SQLITE_BUSY errors under concurrent writes
- ListVersions: initialize result as make([]CheckpointVersion, 0) instead of
  nil slice for consistent caller semantics
- PruneVersions: add comment explaining the double threadID parameter
2026-08-13 20:18:48 +08:00
unclesam 85a1294fa3 Merge remote-tracking branch 'origin/master'
# Conflicts:
#	sqlite_checkpoint.go
2026-08-12 13:14:52 +08:00
unclesam 3a4c1e9c84 feat: implement SQLite-based append-only multi-version state persistence for thread checkpoints 2026-08-12 13:14:18 +08:00
unclesam e247b504a4 feat: implement SQLite-based append-only multi-version state persistence for thread checkpoints 2026-08-12 00:09:10 +08:00
unclesam 7028571a4f docs: update README to include Checkpointer.Delete example 2026-07-04 20:47:59 +08:00
unclesam 5782504351 fix(engine): resolve parallel execution races and improve max steps semantics
- Fix data race in runParallelBranches where first error could be masked
- Improve MaxSteps to allow exactly N steps
- Consolidate Resume validation logic and fix silent overwrite in Graph AddNode/AddEdge
- Add sync.RWMutex to Engine for thread-safe configuration
- Fix P2 issues: Add immediate check for router output, add mutex to hooks example, add context timeout to translation example, and update docs for Checkpointer.Delete
2026-07-04 20:46:36 +08:00
unclesam d2d8eb5a3a docs: 更新README以反映新增特性与安全校验 2026-06-17 11:02:46 +08:00
unclesam de457f670a feat: production-grade hardening v1.1.0
- fix: Compile() defensive map copy, prevents post-compile Graph mutation
- fix: correct 'origin'→'destination' in edge validation error message
- feat: export ErrNotPaused/ErrAlreadyFinished for errors.Is() support
- feat: WithPreNodeHook/WithPostNodeHook now support chaining (multiple hooks)
- feat: Compile detects conflicting outgoing edge types on same node
- fix: FileCheckpointer atomic write (tmp+rename, crash-safe)
- fix: FileCheckpointer threadID path traversal validation
- refactor: remove redundant continue in run() loop
- docs: add doc.go package-level documentation
- test: +7 test cases covering error paths, edge conflicts, hook chaining
2026-06-17 10:44:59 +08:00
unclesam 235f66c9ba Fix(README.md) GopherGraph import path casing
修复 README.md 示例错误导入路径
2026-06-13 21:16:32 +08:00
unclesam 62b23e6bfa chore: 将模块路径中的组织名更改为小写 2026-06-11 20:53:51 +08:00
unclesam 25150551a0 Merge remote-tracking branch 'origin/master' 2026-06-11 20:31:21 +08:00
unclesam 3a15b53b84 feat(engine): 增加运行熔断、并发短路取消及生命周期钩子 2026-06-11 20:31:00 +08:00
unclesam b9218f2a6a feat(engine): 增加运行熔断、并发短路取消及生命周期钩子 2026-06-11 20:23:52 +08:00
unclesam 693c107473 build: 更新模块路径为标准格式 2026-06-11 03:29:55 +08:00
15 changed files with 2197 additions and 101 deletions

No files matched your search

+287 -32
View File
@@ -20,9 +20,13 @@ Go 言語(ジェネリクス)で書かれた、軽量・高性能・型安
### 核心特性 ### 核心特性
- **强类型状态管理:** 利用 Go 泛型定义全局状态 `Graph[S]`,彻底告别松散的 `map[string]any` 和繁琐的类型断言。 - **强类型状态管理:** 利用 Go 泛型定义全局状态 `Graph[S]`,彻底告别松散的 `map[string]any` 和繁琐的类型断言。
- **支持循环结构:** 突破了传统 DAG(有向无环图)工作流的限制,原生支持节点间的双向循环和条件路由。 - **支持循环结构:** 突破了传统 DAG(有向无环图)工作流的限制,原生支持节点间的双向循环和条件路由。
- **并发分支与合并:** 原生支持多个 Agent 节点的并发执行(基于 `sync.WaitGroup` 调度),并提供线程安全的“分流值拷贝”与“汇合状态合并(Merger)”机制。 - **并发分支与短路取消:** 原生支持多个 Agent 节点的并发执行,基于 `context.WithCancel` 实现短路机制——任意分支出错立即取消其余兄弟 goroutine,避免算力空转。
- **人机协同中断 (HITL):** 支持在特定节点前自动挂起,返回执行线程快照,等待人工修改状态或审批通过后一键恢复 (`Resume`)。 - **并发安全深拷贝:** 通过 `Engine.WithStateCloner` 注入自定义深拷贝函数,彻底消除含引用类型的状态在并发分支中的数据竞争。
- **开箱即用持久化:** 提供 `Checkpointer[S]` 接口与内置的 `FileCheckpointer[S]`,支持将进度保存为 JSON 文件,实现停机或重启后进度无损恢复。 - **人机协同中断 (HITL):** 支持在特定节点前自动挂起,返回执行线程快照,等待人工修改状态或审批通过后一键恢复 (`Resume`)。支持 `ErrNotPaused`、`ErrAlreadyFinished` 哨兵错误,便于精准判断异常。
- **死循环硬性熔断:** 通过 `Engine.WithMaxSteps` 设置步数上限,图中的任何意外环路都会被明确报错而非无限阻塞。
- **轻量级生命周期 Hooks(支持多重链式累加):** `Engine.WithPreNodeHook` / `WithPostNodeHook` 支持多次调用链式累加(不覆盖已有钩子),无需修改 `NodeFn`,极易解耦日志、链路追踪(OpenTelemetry)或流式输出等逻辑。
- **编译期结构校验:** 图在编译 (`Compile`) 时,不仅会验证所有节点和边是否存在,还会严格检测同一节点是否冲突地定义了多种出边类型(如同时定义静态边、条件边与并发边),规避多义性逻辑隐患。
- **开箱即用持久化与安全校验:** 提供 `Checkpointer[S]` 接口,内置基于 JSON 文件的 `FileCheckpointer[S]`(原子重命名防损坏)以及基于 SQLite 的 `SQLiteCheckpointer[S]`(支持多版本追加写入 Append-Only、历史版本回溯与剪枝),并对 `threadID` 进行路径穿越防御性校验,实现进程意外退出后进度安全无损地恢复。
- **纯粹的高性能:** 基于内存流转,零外部依赖,极低的上下文切换开销。 - **纯粹的高性能:** 基于内存流转,零外部依赖,极低的上下文切换开销。
### 快速开始 ### 快速开始
@@ -33,24 +37,22 @@ type MyState struct {
Query string Query string
Response string Response string
Value int Value int
Log []string Messages []string // 含引用类型,需要配合 StateCloner 使用
} }
``` ```
#### 2. 并发分支与合并示例 #### 2. 直接使用 CompiledGraph(原始 API,向后兼容)
```go ```go
package main package main
import ( import (
"context" "context"
"fmt" GopherGraph "github.com/unclesam-ly/GopherGraph"
"GopherGraph"
) )
func main() { func main() {
g := GopherGraph.NewGraph[MyState]() g := GopherGraph.NewGraph[MyState]()
// 注册节点
g.AddNode("start", startNode) g.AddNode("start", startNode)
g.AddNode("task1", task1Node) g.AddNode("task1", task1Node)
g.AddNode("task2", task2Node) g.AddNode("task2", task2Node)
@@ -72,9 +74,60 @@ func main() {
} }
``` ```
#### 3. 文件持久化 (Checkpointer) 示例 #### 3. 使用 Engine 包装器(推荐:生产级增强)
```go ```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) 示例
##### A. 基于 JSON 本地文件的 Checkpointer
```go
// 创建文件存储器,指定存放状态文件的目录
fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints") fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints")
// 在遇到中断挂起时,将 thread 保存到磁盘 // 在遇到中断挂起时,将 thread 保存到磁盘
@@ -84,17 +137,49 @@ fc.Save(context.Background(), sessionID, thread)
// 重启程序后,可以从磁盘重新加载进度并恢复执行 // 重启程序后,可以从磁盘重新加载进度并恢复执行
loadedThread, _ := fc.Load(context.Background(), sessionID) loadedThread, _ := fc.Load(context.Background(), sessionID)
thread, _ = cg.Resume(context.Background(), loadedThread, modifiedState) thread, _ = cg.Resume(context.Background(), loadedThread, modifiedState)
// 任务彻底跑完后,清理磁盘上的存档文件
if thread.IsFinished {
fc.Delete(context.Background(), sessionID)
}
```
##### B. 基于 SQLite 的多版本追加写入 (Append-Only) Checkpointer
```go
// 创建 SQLite 存储器,自动开启 WAL 高并发模式
sc, _ := GopherGraph.NewSQLiteCheckpointer[MyState]("./checkpoints.db")
defer sc.Close()
// 每次 Save 都会追加写入一个新的版本号 (v1, v2, v3...)
sessionID := "user-session-123"
sc.Save(ctx, sessionID, thread)
// Load 默认获取最新版本的线程快照
latestThread, _ := sc.Load(ctx, sessionID)
// LoadVersion 支持精确回溯并恢复指定历史版本的快照(版本重放/回滚)
v1Thread, _ := sc.LoadVersion(ctx, sessionID, 1)
// ListVersions 查询该 session 的所有历史版本元数据
versions, _ := sc.ListVersions(ctx, sessionID)
// PruneVersions 保留最近 N 个版本,自动清理过期历史
sc.PruneVersions(ctx, sessionID, 5)
``` ```
#### 运行内置示例 #### 运行内置示例
项目包含了一个完整的“AI翻译 -> 质量检测 -> 人工审核 -> 发布”的交互式 Demo:
```bash ```bash
# 完整的"AI翻译 -> 质量检测 -> 人工审核 -> 发布"交互式 Demo
go run examples/translation/main.go go run examples/translation/main.go
# Engine 增强特性演示(Hooks、StateCloner、流式输出)
go run examples/hooks/main.go
``` ```
#### 运行单元测试 #### 运行单元测试
```bash ```bash
go test -v ./... # 含竞态检测器,确保并发分支无数据竞争
go test -v -race ./...
``` ```
--- ---
@@ -107,9 +192,13 @@ go test -v ./...
### Key Features ### 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. - **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. - **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.
- **Human-in-the-Loop (HITL):** Pause workflow execution *before* a designated node, capture a snapshot (`Thread`), modify the state, and `Resume` seamlessly. - **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.
- **State Persistence (Checkpointer):** Generic `Checkpointer[S]` interface with built-in `FileCheckpointer[S]` for saving and loading execution snapshots to/from JSON files. - **Human-in-the-Loop (HITL):** Pause workflow execution *before* a designated node, capture a snapshot (`Thread`), modify the state, and `Resume` seamlessly. Returns exported sentinel errors `ErrNotPaused` or `ErrAlreadyFinished` when resume pre-conditions are not met.
- **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 (Additive Chaining):** `Engine.WithPreNodeHook` / `WithPostNodeHook` fire before and after each node without touching any `NodeFn` signature. Supports multiple chained registrations (additive) instead of overwriting, making it easy to decouple concerns like logging, tracing (OpenTelemetry), and streaming.
- **Compile-Time Structural Validation:** During `Compile()`, the engine validates node/edge existence and strictly checks for outgoing edge conflicts on each node (e.g., preventing a node from simultaneously defining static, conditional, and parallel edges) to ensure unambiguous flow execution.
- **State Persistence & Path Safety:** Generic `Checkpointer[S]` interface with built-in `FileCheckpointer[S]` (atomic file rename) and `SQLiteCheckpointer[S]` (append-only multi-versioning, time travel history & pruning), featuring path traversal validation on `threadID`.
- **Zero-Dependency & High-Performance:** Written in pure Go with zero external dependencies, leveraging in-memory queues for lightning-fast orchestration. - **Zero-Dependency & High-Performance:** Written in pure Go with zero external dependencies, leveraging in-memory queues for lightning-fast orchestration.
### Quick Start ### Quick Start
@@ -120,18 +209,17 @@ type MyState struct {
Query string Query string
Response string Response string
Value int 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 ```go
package main package main
import ( import (
"context" "context"
"fmt" GopherGraph "github.com/unclesam-ly/GopherGraph"
"GopherGraph"
) )
func main() { func main() {
@@ -158,7 +246,57 @@ 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
##### A. JSON File Checkpointer
```go ```go
// Initialize a local directory storage // Initialize a local directory storage
fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints") fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints")
@@ -170,17 +308,49 @@ fc.Save(context.Background(), sessionID, thread)
// Reload the snapshot (e.g. after a process restart) and resume // Reload the snapshot (e.g. after a process restart) and resume
loadedThread, _ := fc.Load(context.Background(), sessionID) loadedThread, _ := fc.Load(context.Background(), sessionID)
thread, _ = cg.Resume(context.Background(), loadedThread, modifiedState) thread, _ = cg.Resume(context.Background(), loadedThread, modifiedState)
// Clean up the storage file once the workflow is completely finished
if thread.IsFinished {
fc.Delete(context.Background(), sessionID)
}
```
##### B. SQLite Append-Only Multi-Version Checkpointer
```go
// Initialize SQLite checkpointer (WAL mode enabled automatically)
sc, _ := GopherGraph.NewSQLiteCheckpointer[MyState]("./checkpoints.db")
defer sc.Close()
// Append a new version snapshot on every Save (v1, v2, v3...)
sessionID := "user-session-123"
sc.Save(ctx, sessionID, thread)
// Load gets the latest version by default
latestThread, _ := sc.Load(ctx, sessionID)
// LoadVersion fetches a specific historical version for replay/rollback
v1Thread, _ := sc.LoadVersion(ctx, sessionID, 1)
// ListVersions retrieves version metadata for the thread
versions, _ := sc.ListVersions(ctx, sessionID)
// PruneVersions retains the N most recent versions and purges older records
sc.PruneVersions(ctx, sessionID, 5)
``` ```
#### Run the Interactive Demo #### Run the Interactive Demo
An interactive "Translation -> Review -> Human Approval -> Publish" demo is included:
```bash ```bash
# Interactive "Translation -> Review -> Human Approval -> Publish" demo
go run examples/translation/main.go go run examples/translation/main.go
# Engine features demo (Hooks, StateCloner, Streaming)
go run examples/hooks/main.go
``` ```
#### Run Unit Tests #### Run Unit Tests
```bash ```bash
go test -v ./... # The -race flag verifies zero data races in parallel branches
go test -v -race ./...
``` ```
--- ---
@@ -193,9 +363,13 @@ go test -v ./...
### 主な特徴 ### 主な特徴
- **型安全な状態管理:** ジェネリクスを用いて `Graph[S]` に構造体をバインドします。冗長な `map[string]any` やランタイム時の型アサーションから解放されます。 - **型安全な状態管理:** ジェネリクスを用いて `Graph[S]` に構造体をバインドします。冗長な `map[string]any` やランタイム時の型アサーションから解放されます。
- **ループと条件付き分岐のサポート:** 従来の DAG(有向非巡回グラフ)の制限を超え、エージェントの判定に基づく条件付きルートや双方向のループをサポート。 - **ループと条件付き分岐のサポート:** 従来の DAG(有向非巡回グラフ)の制限を超え、エージェントの判定に基づく条件付きルートや双方向のループをサポート。
- **並行処理とマージ(分流と合流):** `sync.WaitGroup` に基づく複数エージェントノードの並行実行に対応。分流時の値渡しコピーによるスレッドセーフの確保と、合流時の `ParallelMergeFn` による状態マージを実現。 - **並行処理とショートサーキットキャンセル:** `context.WithCancel` に基づく並行ブランチ実行を実装。いずれかのブランチでエラーが発生した瞬間に他の全 goroutine へキャンセルシグナルを伝播し、LLM API の無駄な呼び出しを防ぎます。
- **Human-in-the-Loop (HITL) の一時停止と再開:** 特定のノードの実行前に処理を自動停止し、スナップショット(`Thread`)を返します。承認や状態修正の後にシームレスに処理を再開(`Resume`)できます。 - **データ競合を防ぐ深いコピー(DeepClone):** `Engine.WithStateCloner` でカスタムの深コピー関数を注入。参照型(スライス・マップ・ポインタ)を持つ状態でも、並行ブランチでのデータ競合を根本から解消します。
- **状態の永続化 (Checkpointer):** 抽象的な `Checkpointer[S]` インタフェースと、JSONファイルへの書き出し・読み込みを行う組み込みの `FileCheckpointer[S]` をサポート。サーバー再起動後の進捗復旧が可能。 - **Human-in-the-Loop (HITL) の一時停止と再開:** 特定のノードの実行前に処理を自动停止し、スナップショット(`Thread`)を返します。承認や状態修正の後にシームレスに処理を再開(`Resume`)できます。再開判定用に `ErrNotPaused` や `ErrAlreadyFinished` といった明確なセンチネルエラーを提供しています。
- **無限ループのサーキットブレーカー:** `Engine.WithMaxSteps` でステップ数の上限を設定。グラフ内に意図しないサイクルが発生しても、永久ブロックせず明確なエラーを返します。
- **軽量なライフサイクルフック(複数チェーン登録対応):** `Engine.WithPreNodeHook` / `WithPostNodeHook` により、`NodeFn` のシグネチャを変更することなくロギング、トレーシング(OpenTelemetry)、ストリーミング出力を実装できます。**複数回呼び出すことで複数のフックを順番に実行(チェーン蓄積)**可能です。
- **コンパイル時の構造検証:** `Compile()` 実行時に、すべてのノードや辺の存在をチェックするだけでなく、同一ノードに対する出辺タイプ(静的辺、条件付き辺、並行辺)の重複競合も厳密に検出し、意図しない曖昧な挙動を防ぎます。
- **状態の永続化とパス安全性:** 抽象的な `Checkpointer[S]` インタフェースと、アトミック書き込み対応の `FileCheckpointer[S]`、およびマルチバージョン追記(Append-Only)・履歴復元に対応した `SQLiteCheckpointer[S]` を内蔵。`threadID` のパス・トラバーサル検証によりセキュアに状態の復旧が可能です。
- **ピュア Go & 高性能:** 外部依存関係ゼロ。メモリ内チャネルを用いた高速なコンテキスト切り替え。 - **ピュア Go & 高性能:** 外部依存関係ゼロ。メモリ内チャネルを用いた高速なコンテキスト切り替え。
### クイックスタート ### クイックスタート
@@ -206,18 +380,17 @@ type MyState struct {
Query string Query string
Response string Response string
Value int Value int
Log []string Messages []string // 参照型 — 並行処理では StateCloner と組み合わせて使用
} }
``` ```
#### 2. 並行処理(Parallel)のコード例 #### 2. CompiledGraph を直接使う(元の API、完全後方互換)
```go ```go
package main package main
import ( import (
"context" "context"
"fmt" GopherGraph "github.com/unclesam-ly/GopherGraph"
"GopherGraph"
) )
func main() { func main() {
@@ -244,7 +417,57 @@ 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)のコード例
##### A. JSON ファイル Checkpointer
```go ```go
// 保存先ディレクトリを指定してローカルファイルチェッカーを初期化 // 保存先ディレクトリを指定してローカルファイルチェッカーを初期化
fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints") fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints")
@@ -256,15 +479,47 @@ fc.Save(context.Background(), sessionID, thread)
// スナップショットを読み込み、プロセス再起動後に処理を再開 // スナップショットを読み込み、プロセス再起動後に処理を再開
loadedThread, _ := fc.Load(context.Background(), sessionID) loadedThread, _ := fc.Load(context.Background(), sessionID)
thread, _ = cg.Resume(context.Background(), loadedThread, modifiedState) thread, _ = cg.Resume(context.Background(), loadedThread, modifiedState)
// 処理が完全に終了した後、ストレージ上のファイルをクリーンアップする
if thread.IsFinished {
fc.Delete(context.Background(), sessionID)
}
```
##### B. SQLite マルチバージョン追記 (Append-Only) Checkpointer
```go
// SQLite チェッカーを初期化(WAL 高並行モードが自動有効化)
sc, _ := GopherGraph.NewSQLiteCheckpointer[MyState]("./checkpoints.db")
defer sc.Close()
// Save のたびに新しいバージョン(v1, v2, v3...)を追記保存
sessionID := "user-session-123"
sc.Save(ctx, sessionID, thread)
// Load はデフォルトで最新バージョンのスナップショットを取得
latestThread, _ := sc.Load(ctx, sessionID)
// LoadVersion は特定バージョンを正確に復元(タイムトラベル / 状態ロールバック)
v1Thread, _ := sc.LoadVersion(ctx, sessionID, 1)
// ListVersions で全履歴のメタデータ一覧を取得
versions, _ := sc.ListVersions(ctx, sessionID)
// PruneVersions で最新 N 件のみを保持し、古い履歴を削除
sc.PruneVersions(ctx, sessionID, 5)
``` ```
#### インタラクティブデモの実行 #### インタラクティブデモの実行
「AI翻訳 -> 監査 -> 人間による確認 -> 公開」の一連のループを体験できるデモを用意しています:
```bash ```bash
# 「AI翻訳 -> 監査 -> 人間による確認 -> 公開」のインタラクティブデモ
go run examples/translation/main.go go run examples/translation/main.go
# Engine 機能デモ(Hooks・StateCloner・ストリーミング)
go run examples/hooks/main.go
``` ```
#### ユニットテストの実行 #### ユニットテストの実行
```bash ```bash
go test -v ./... # -race フラグで並行ブランチのデータ競合がゼロであることを検証
go test -v -race ./...
``` ```
+95 -7
View File
@@ -6,6 +6,7 @@ import (
"fmt" "fmt"
"os" "os"
"path/filepath" "path/filepath"
"strings"
) )
// Checkpointer 定义了持久化状态的接口,S 是工作流的强类型状态。 // Checkpointer 定义了持久化状态的接口,S 是工作流的强类型状态。
@@ -14,38 +15,93 @@ type Checkpointer[S any] interface {
Save(ctx context.Context, threadID string, thread *Thread[S]) error Save(ctx context.Context, threadID string, thread *Thread[S]) error
// Load 读取之前保存的线程快照,用于恢复执行。 // Load 读取之前保存的线程快照,用于恢复执行。
Load(ctx context.Context, threadID string) (*Thread[S], error) Load(ctx context.Context, threadID string) (*Thread[S], error)
// Delete 删除指定线程的快照文件,用于工作流完结后的清理。
Delete(ctx context.Context, threadID string) error
} }
// FileCheckpointer 实现了基于本地 JSON 文件的状态持久化。 // FileCheckpointer 实现了基于本地 JSON 文件的状态持久化。
type FileCheckpointer[S any] struct { type FileCheckpointer[S any] struct {
dir string // 存放状态文件的文件夹路径 dir string // 存放状态文件的文件夹路径
pretty bool // 是否使用缩进格式(方便人工阅读,但体积更大)
}
// FileCheckpointerOption 是 FileCheckpointer 的可选配置项。
type FileCheckpointerOption func(*fileCheckpointerConfig)
type fileCheckpointerConfig struct {
pretty bool
}
// WithPrettyJSON 开启 JSON 缩进格式(方便人工阅读,适合开发调试阶段)。
// 默认关闭,使用紧凑格式以减少磁盘 I/O 和文件体积。
func WithPrettyJSON() FileCheckpointerOption {
return func(c *fileCheckpointerConfig) {
c.pretty = true
}
} }
// NewFileCheckpointer 创建一个文件检查点管理器。 // NewFileCheckpointer 创建一个文件检查点管理器。
func NewFileCheckpointer[S any](dir string) (*FileCheckpointer[S], error) { // 支持通过 opts 传入可选配置(如 WithPrettyJSON())。
func NewFileCheckpointer[S any](dir string, opts ...FileCheckpointerOption) (*FileCheckpointer[S], error) {
// 确保文件夹存在,不存在则创建 // 确保文件夹存在,不存在则创建
if err := os.MkdirAll(dir, 0755); err != nil { if err := os.MkdirAll(dir, 0755); err != nil {
return nil, fmt.Errorf("failed to create checkpoint dir: %w", err) return nil, fmt.Errorf("failed to create checkpoint dir: %w", err)
} }
return &FileCheckpointer[S]{dir: dir}, nil cfg := &fileCheckpointerConfig{}
for _, opt := range opts {
opt(cfg)
}
return &FileCheckpointer[S]{dir: dir, pretty: cfg.pretty}, nil
} }
// Save 将 Thread 序列化为 JSON 并写入本地文件 // Save 将 Thread 序列化为 JSON 并写入本地文件。
// 采用"先写临时文件再原子重命名"的方式,防止写入中途程序崩溃导致文件损坏。
// 若存在上次崩溃残留的临时文件,会在写入前自动清理。
func (fc *FileCheckpointer[S]) Save(ctx context.Context, threadID string, thread *Thread[S]) error { func (fc *FileCheckpointer[S]) Save(ctx context.Context, threadID string, thread *Thread[S]) error {
path := filepath.Join(fc.dir, threadID+".json") if err := validateThreadID(threadID); err != nil {
return err
}
// 使用 MarshalIndent 方便人工阅读生成的 JSON 文件 path := filepath.Join(fc.dir, threadID+".json")
data, err := json.MarshalIndent(thread, "", " ") tmpPath := path + ".tmp"
// 序列化:生产环境默认使用紧凑格式以减少磁盘 I/O;开发阶段可用 WithPrettyJSON() 开启缩进
var data []byte
var err error
if fc.pretty {
data, err = json.MarshalIndent(thread, "", " ")
} else {
data, err = json.Marshal(thread)
}
if err != nil { if err != nil {
return fmt.Errorf("failed to marshal thread: %w", err) return fmt.Errorf("failed to marshal thread: %w", err)
} }
return os.WriteFile(path, data, 0644) // 清理上次崩溃可能残留的临时文件(忽略不存在的错误)
_ = os.Remove(tmpPath)
// 原子写入:先写临时文件,再 Rename 替换,确保断电/崩溃安全
if err := os.WriteFile(tmpPath, data, 0644); err != nil {
return fmt.Errorf("failed to write temp file: %w", err)
}
if err := os.Rename(tmpPath, path); err != nil {
// Rename 失败时,尝试清理临时文件,避免残留
_ = os.Remove(tmpPath)
return fmt.Errorf("failed to rename temp file: %w", err)
}
return nil
} }
// Load 从本地文件中读取 JSON 并反序列化为 Thread。 // Load 从本地文件中读取 JSON 并反序列化为 Thread。
// 若加载到一个已完结(IsFinished == true)的线程,会包装 ErrAlreadyFinished 作为额外信息返回,
// 方便调用方做幂等处理(注意:此时 thread 依然有效,err 可通过 errors.Is 判断)。
func (fc *FileCheckpointer[S]) Load(ctx context.Context, threadID string) (*Thread[S], error) { func (fc *FileCheckpointer[S]) Load(ctx context.Context, threadID string) (*Thread[S], error) {
if err := validateThreadID(threadID); err != nil {
return nil, err
}
path := filepath.Join(fc.dir, threadID+".json") path := filepath.Join(fc.dir, threadID+".json")
data, err := os.ReadFile(path) data, err := os.ReadFile(path)
@@ -58,5 +114,37 @@ func (fc *FileCheckpointer[S]) Load(ctx context.Context, threadID string) (*Thre
return nil, fmt.Errorf("failed to unmarshal thread: %w", err) return nil, fmt.Errorf("failed to unmarshal thread: %w", err)
} }
// 如果加载到已完结的线程,附带哨兵错误提示,方便调用方感知
if thread.IsFinished {
return &thread, fmt.Errorf("thread %q has already finished: %w", threadID, ErrAlreadyFinished)
}
return &thread, nil return &thread, nil
} }
// Delete 删除指定 threadID 对应的快照文件。
// 通常在工作流完结后调用,防止快照文件无限积累占用磁盘空间。
// 若文件不存在,静默返回 nil(幂等操作)。
func (fc *FileCheckpointer[S]) Delete(ctx context.Context, threadID string) error {
if err := validateThreadID(threadID); err != nil {
return err
}
path := filepath.Join(fc.dir, threadID+".json")
if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("failed to delete checkpoint file: %w", err)
}
return nil
}
// validateThreadID 校验 threadID 的合法性,防止路径穿越(path traversal)攻击。
// threadID 中不允许出现目录分隔符或 ".." 序列。
func validateThreadID(id string) error {
if id == "" {
return fmt.Errorf("threadID must not be empty")
}
if strings.ContainsAny(id, "/\\") || strings.Contains(id, "..") {
return fmt.Errorf("threadID %q contains invalid characters", id)
}
return nil
}
+15
View File
@@ -0,0 +1,15 @@
// Package GopherGraph 是一个用 Go 语言编写的代码即图(Code-as-Graph)智能体编排引擎,
// 设计灵感来源于 Python 生态的 LangGraph。
//
// 它利用 Go 1.18+ 的泛型机制,让工作流的上下文状态在编译期就具备强类型约束,
// 并依靠 Go 原生的 Goroutines 和 Channels 实现极高并发的本地数据流转。
//
// # 核心概念
//
// - [Graph] 图构建器:通过 [NewGraph] 创建,使用 AddNode/AddEdge 系列方法声明节点和连线。
// - [CompiledGraph] 编译后的只读图:通过 [Graph.Compile] 生成,可直接调用 Start/Resume 运行。
// - [Engine] 增强包装器:通过 [NewEngine] 创建,在 CompiledGraph 基础上叠加
// StateCloner(深拷贝)、MaxSteps(死循环熔断)、PreNodeHook/PostNodeHook(生命周期钩子)。
// - [Thread] 执行快照:承载当前状态、下一节点、暂停/结束标记,可序列化用于断点恢复。
// - [Checkpointer] 持久化接口:提供 Save/Load/Delete 用于工作流进度存储和清理,内置 [FileCheckpointer] 基于 JSON 文件存储,以及 [SQLiteCheckpointer] 支持多版本追加写入(Append-Only)与版本回溯。
package GopherGraph
+208 -52
View File
@@ -27,6 +27,11 @@ type CompiledGraph[S any] struct {
// Compile 校验图结构的合法性,并生成可运行的 CompiledGraph。 // Compile 校验图结构的合法性,并生成可运行的 CompiledGraph。
func (g *Graph[S]) Compile() (*CompiledGraph[S], error) { func (g *Graph[S]) Compile() (*CompiledGraph[S], error) {
// 首先展示所有注册阶段积累的错误(如重复注册节点或边)
if err := errorsJoin(g.errs); err != nil {
return nil, fmt.Errorf("graph registration errors:\n%w", err)
}
if len(g.nodes) == 0 { if len(g.nodes) == 0 {
return nil, errors.New("cannot compile graph: graph contains no nodes") return nil, errors.New("cannot compile graph: graph contains no nodes")
} }
@@ -37,7 +42,7 @@ func (g *Graph[S]) Compile() (*CompiledGraph[S], error) {
return nil, fmt.Errorf("compile error: edge origin %q does not exist", from) return nil, fmt.Errorf("compile error: edge origin %q does not exist", from)
} }
if _, exists := g.nodes[to]; !exists { if _, exists := g.nodes[to]; !exists {
return nil, fmt.Errorf("compile error: edge origin %q does not exist", to) return nil, fmt.Errorf("compile error: edge destination %q does not exist", to)
} }
} }
@@ -79,16 +84,72 @@ func (g *Graph[S]) Compile() (*CompiledGraph[S], error) {
} }
} }
// 返回编译好的只读图 // 校验同一节点的出边冲突:不能同时存在于 edges、conditional、parallels 中
for node := range g.nodes {
conflictCount := 0
if _, ok := g.edges[node]; ok {
conflictCount++
}
if _, ok := g.conditional[node]; ok {
conflictCount++
}
if _, ok := g.parallels[node]; ok {
conflictCount++
}
if conflictCount > 1 {
return nil, fmt.Errorf("compile error: node %q has conflicting outgoing edges (defined in multiple edge types simultaneously)", node)
}
}
// 返回编译好的只读图(防御性拷贝,阻止编译后修改 Graph 对 CompiledGraph 造成污染)
nodesCopy := make(map[string]NodeFn[S], len(g.nodes))
for k, v := range g.nodes {
nodesCopy[k] = v
}
edgesCopy := make(map[string]string, len(g.edges))
for k, v := range g.edges {
edgesCopy[k] = v
}
conditionalCopy := make(map[string]RouterFn[S], len(g.conditional))
for k, v := range g.conditional {
conditionalCopy[k] = v
}
interruptsCopy := make(map[string]bool, len(g.interrupts))
for k, v := range g.interrupts {
interruptsCopy[k] = v
}
parallelsCopy := make(map[string]parallelStep[S], len(g.parallels))
for k, v := range g.parallels {
parallelsCopy[k] = v
}
return &CompiledGraph[S]{ return &CompiledGraph[S]{
nodes: g.nodes, nodes: nodesCopy,
edges: g.edges, edges: edgesCopy,
conditional: g.conditional, conditional: conditionalCopy,
interrupts: g.interrupts, interrupts: interruptsCopy,
parallels: g.parallels, parallels: parallelsCopy,
}, nil }, nil
} }
// ErrNotPaused 在调用 Resume 时线程未处于暂停状态时返回。可用于 errors.Is 判断。
var ErrNotPaused = errors.New("cannot resume: thread is not paused")
// ErrAlreadyFinished 在调用 Resume 时线程已结束时返回。可用于 errors.Is 判断。
var ErrAlreadyFinished = errors.New("cannot resume: thread is already finished")
// validateResume 验证线程是否满足恢复执行的前提条件。
// CompiledGraph.Resume 和 Engine.Resume 共用此逻辑,避免重复实现。
func validateResume[S any](thread *Thread[S]) error {
if !thread.IsPaused {
return ErrNotPaused
}
if thread.IsFinished {
return ErrAlreadyFinished
}
return nil
}
// Start 从指定的起始节点开始运行图,并持续流转,直到遇到中断或运行结束 // Start 从指定的起始节点开始运行图,并持续流转,直到遇到中断或运行结束
func (cg *CompiledGraph[S]) Start(ctx context.Context, startNode string, initialState S) (*Thread[S], error) { func (cg *CompiledGraph[S]) Start(ctx context.Context, startNode string, initialState S) (*Thread[S], error) {
thread := &Thread[S]{ thread := &Thread[S]{
@@ -96,31 +157,40 @@ func (cg *CompiledGraph[S]) Start(ctx context.Context, startNode string, initial
NextNode: startNode, NextNode: startNode,
} }
return cg.run(ctx, thread) return cg.run(ctx, thread, runOptions[S]{})
} }
// Resume 恢复执行一个被暂停的线程,并允许注入外部修改后的状态数据(例如人工审批修改后的结果) // Resume 恢复执行一个被暂停的线程,并允许注入外部修改后的状态数据(例如人工审批修改后的结果)
func (cg *CompiledGraph[S]) Resume(ctx context.Context, thread *Thread[S], modifiedState S) (*Thread[S], error) { func (cg *CompiledGraph[S]) Resume(ctx context.Context, thread *Thread[S], modifiedState S) (*Thread[S], error) {
if !thread.IsPaused { if err := validateResume(thread); err != nil {
return nil, errors.New("cannot resume: thread is not paused") return nil, err
}
if thread.IsFinished {
return nil, errors.New("cannot resume: thread is already finished")
} }
// 注入人工修改后的状态,并解除暂停标记 // 注入人工修改后的状态,并解除暂停标记
thread.State = modifiedState thread.State = modifiedState
thread.IsPaused = false 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 是引擎内部循环调度器,负责驱动节点向前流转 // 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 { for {
// 检查 Context,支持外部超时控制或取消 // 【熔断 2】检查 Context,支持外部超时控制或取消
if err := ctx.Err(); err != nil { select {
return thread, err case <-ctx.Done():
return thread, ctx.Err()
default:
} }
currentNodeName := thread.NextNode currentNodeName := thread.NextNode
@@ -130,51 +200,49 @@ func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S]) (*Thread
return thread, nil return thread, nil
} }
// 查找并执行当前节点 // 【熔断 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++
}
// 查找当前节点
nodeFn, exists := cg.nodes[currentNodeName] nodeFn, exists := cg.nodes[currentNodeName]
if !exists { if !exists {
return thread, fmt.Errorf("runtime error: node %q not found", currentNodeName) 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) newState, err := nodeFn(ctx, thread.State)
if err != nil { if err != nil {
return thread, fmt.Errorf("node %q execution error: %w", currentNodeName, err) return thread, fmt.Errorf("node %q execution error: %w", currentNodeName, err)
} }
thread.State = newState thread.State = newState
// 触发 PostNodeHook
if opts.postNodeHook != nil {
opts.postNodeHook(ctx, currentNodeName, thread.State)
}
// 计算下一个该执行的节点名称 // 计算下一个该执行的节点名称
var nextNode string var nextNode string
if step, isParallel := cg.parallels[currentNodeName]; isParallel { if step, isParallel := cg.parallels[currentNodeName]; isParallel {
// 【处理并发分流】 // 【处理并发分流】使用短路取消机制
var wg sync.WaitGroup branches, err := cg.runParallelBranches(ctx, step.targets, thread.State, opts.stateCloner)
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 { if err != nil {
return thread, fmt.Errorf("parallel branch execution error: %w", err) return thread, fmt.Errorf("parallel branch execution error: %w", err)
} }
}
// 【状态合并】调用用户自定义的合并函数 // 【状态合并】调用用户自定义的合并函数
mergedState, err := step.merger(ctx, thread.State, branches) mergedState, err := step.merger(ctx, thread.State, branches)
if err != nil { if err != nil {
@@ -188,20 +256,21 @@ func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S]) (*Thread
if err != nil { if err != nil {
return thread, fmt.Errorf("router for node %q execution error: %w", currentNodeName, err) return thread, fmt.Errorf("router for node %q execution error: %w", currentNodeName, err)
} }
// 立即校验返回的节点是否存在,提供更清晰的报错
if next != "" {
if _, exists := cg.nodes[next]; !exists {
return thread, fmt.Errorf("router for node %q returned non-existent node %q (not found)", currentNodeName, next)
}
}
nextNode = next nextNode = next
} else { } else {
// 否则使用静态连线边 // 否则使用静态连线边
nextNode = cg.edges[currentNodeName] nextNode = cg.edges[currentNodeName]
} }
// 更新快照中的“下一个节点” // 更新快照中的"下一个节点"
thread.NextNode = nextNode thread.NextNode = nextNode
// 如果下一站是终点,直接进入下一次循环触发结束逻辑
if nextNode == "" {
continue
}
// 核心中断机制:如果即将进入的下一个节点被标记为了中断节点,则在此挂起 // 核心中断机制:如果即将进入的下一个节点被标记为了中断节点,则在此挂起
if cg.interrupts[nextNode] { if cg.interrupts[nextNode] {
thread.IsPaused = true thread.IsPaused = true
@@ -209,3 +278,90 @@ 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)
// 【P0 修复】先收集所有结果,再判断是否有真实错误。
// 原先在循环中遇到非 context.Canceled 错误就立即 return,但由于 channel 中的写入顺序是不确定的,
// 存在以下竞态:分支 A 先把 context.Canceled 写入 channel,紧接着分支 B 才把真实错误写入。
// 若此时早返回,真实错误永远不会被读取,调用方可能收到 nil error 而实际上有分支失败了。
// 解决方案:等全部结果到齐后,再过滤出第一个非取消的真实错误。
branches := make([]S, len(targets))
var firstErr error
for r := range resultCh {
if r.err != nil {
// 过滤掉因取消产生的次生错误,只记录第一个真实错误
if !errors.Is(r.err, context.Canceled) && firstErr == nil {
firstErr = r.err
}
} else {
branches[r.idx] = r.state
}
}
if firstErr != nil {
return nil, firstErr
}
return branches, nil
}
+167
View File
@@ -0,0 +1,167 @@
package GopherGraph
import (
"context"
"sync"
)
// Engine 是 CompiledGraph 的增强包装器,在完全兼容原有 API 的基础上,
// 提供以下四项生产级能力:
//
// 1. StateCloner — 自定义深拷贝逻辑,消除并发分支的数据竞争隐患
// 2. MaxSteps — 步数硬性熔断,防止图中的环路导致无限循环
// 3. PreNodeHook — 节点执行前的生命周期钩子(日志、链路追踪等)
// 4. PostNodeHook — 节点执行后的生命周期钩子(指标上报、流式推送等)
//
// 所有选项均为可选,通过链式调用注册。未注册的选项将退回到 CompiledGraph 的默认行为。
//
// # 线程安全说明
//
// Engine 使用 sync.RWMutex 保护配置字段,WithXxx 配置方法和 opts() 读取方法均加锁。
// 因此,允许在 Start/Resume 运行期间并发调用 WithXxx 修改配置(但不推荐这样做)。
// 常见的安全模式:在服务启动时一次性完成链式配置,之后多个 Goroutine 并发调用 Start/Resume。
//
// 用法示例:
//
// 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 {
mu sync.RWMutex
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.mu.Lock()
defer e.mu.Unlock()
e.stateCloner = fn
return e
}
// WithMaxSteps 设置引擎单次运行(Start/Resume)允许的最大节点执行步数。
//
// 这是防御图中存在意外环路(infinite loop)的最后一道硬性熔断。
// 超过上限后,引擎将返回明确的错误,而不是永久阻塞。
//
// 建议根据工作流的预期最大深度合理设置(如 100~1000)。
// 传入 0 或负数表示不限制(与不调用此方法等效)。
func (e *Engine[S]) WithMaxSteps(n int) *Engine[S] {
e.mu.Lock()
defer e.mu.Unlock()
e.maxSteps = n
return e
}
// WithPreNodeHook 注册在每个节点执行【之前】触发的钩子函数。
//
// 典型用途:
// - 打印结构化日志(node 名称、输入状态)
// - 记录 OpenTelemetry Span 的开始时间
// - 向前端推送"正在执行 xxx 节点"的流式事件
//
// 支持多次调用:每次调用将新钩子追加到执行链末尾,按注册顺序依次执行,不覆盖已有钩子。
// 钩子函数接收到的 state 是节点执行【前】的状态快照,请勿修改(只读语义)。
// 钩子中的 panic 会向上传播,建议在内部做好 recover。
func (e *Engine[S]) WithPreNodeHook(fn HookFn[S]) *Engine[S] {
e.mu.Lock()
defer e.mu.Unlock()
if e.preNodeHook == nil {
e.preNodeHook = fn
} else {
prev := e.preNodeHook
e.preNodeHook = func(ctx context.Context, name string, s S) {
prev(ctx, name, s)
fn(ctx, name, s)
}
}
return e
}
// WithPostNodeHook 注册在每个节点执行【之后】触发的钩子函数。
//
// 典型用途:
// - 打印节点的输出状态(用于 Debug)
// - 上报节点耗时指标(Prometheus/DataDog)
// - 向前端推送节点执行完毕的流式事件
//
// 支持多次调用:每次调用将新钩子追加到执行链末尾,按注册顺序依次执行,不覆盖已有钩子。
// 钩子函数接收到的 state 是节点执行【后】更新过的状态,同为只读语义。
func (e *Engine[S]) WithPostNodeHook(fn HookFn[S]) *Engine[S] {
e.mu.Lock()
defer e.mu.Unlock()
if e.postNodeHook == nil {
e.postNodeHook = fn
} else {
prev := e.postNodeHook
e.postNodeHook = func(ctx context.Context, name string, s S) {
prev(ctx, name, s)
fn(ctx, name, s)
}
}
return e
}
// opts 将 Engine 的配置快照打包成 runOptions,传入底层 run() 调度器。
// 使用 RLock 保证读取配置时的并发安全。
func (e *Engine[S]) opts() runOptions[S] {
e.mu.RLock()
defer e.mu.RUnlock()
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)。
// 复用 validateResume 进行前提条件校验,与 CompiledGraph.Resume 保持一致。
func (e *Engine[S]) Resume(ctx context.Context, thread *Thread[S], modifiedState S) (*Thread[S], error) {
if err := validateResume(thread); err != nil {
return nil, err
}
thread.State = modifiedState
thread.IsPaused = false
return e.graph.run(ctx, thread, e.opts())
}
+167
View File
@@ -0,0 +1,167 @@
package main
import (
"context"
"fmt"
"log"
"sync"
"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)
nodeTimingsMu sync.Mutex
)
engine := GopherGraph.NewEngine(cg).
// 1. 注册深拷贝函数,消除并发分支对 Messages 切片的数据竞争
WithStateCloner(deepClone).
// 2. 设置最大步数熔断(防御意外死循环)
WithMaxSteps(50).
// 3. Pre Hook:记录节点开始时间 + 向控制台推送"流式进度"
WithPreNodeHook(func(ctx context.Context, name string, s AgentState) {
nodeTimingsMu.Lock()
nodeTimings[name] = time.Now()
nodeTimingsMu.Unlock()
fmt.Printf(" ⏳ [PRE ] %-15s | 消息数: %d\n", name, len(s.Messages))
}).
// 4. Post Hook:计算并打印节点耗时
WithPostNodeHook(func(ctx context.Context, name string, s AgentState) {
nodeTimingsMu.Lock()
elapsed := time.Since(nodeTimings[name])
nodeTimingsMu.Unlock()
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✨ 示例结束")
}
+5 -2
View File
@@ -1,12 +1,14 @@
package main package main
import ( import (
"GopherGraph"
"bufio" "bufio"
"context" "context"
"fmt" "fmt"
"os" "os"
"strings" "strings"
"time"
GopherGraph "github.com/unclesam-ly/GopherGraph"
) )
// TranslationState 代表我们工作流的共享状态(State) // TranslationState 代表我们工作流的共享状态(State)
@@ -121,7 +123,8 @@ func main() {
// 包含敏感词“秘密”,会触发人工审核中断: // 包含敏感词“秘密”,会触发人工审核中断:
inputText := "这是一个秘密你好" inputText := "这是一个秘密你好"
fmt.Printf("🚀 启动工作流,输入原文: %q\n", inputText) fmt.Printf("🚀 启动工作流,输入原文: %q\n", inputText)
ctx := context.Background() ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
defer cancel()
thread, err := cg.Start(ctx, "translator", TranslationState{InputText: inputText}) thread, err := cg.Start(ctx, "translator", TranslationState{InputText: inputText})
if err != nil { if err != nil {
fmt.Printf("运行失败: %v\n", err) fmt.Printf("运行失败: %v\n", err)
+14 -1
View File
@@ -1,3 +1,16 @@
module GopherGraph module github.com/unclesam-ly/GopherGraph
go 1.26.2 go 1.26.2
require (
github.com/dustin/go-humanize v1.0.1 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/mattn/go-isatty v0.0.24 // indirect
github.com/ncruces/go-strftime v1.0.0 // indirect
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
golang.org/x/sys v0.47.0 // indirect
modernc.org/libc v1.74.4 // indirect
modernc.org/mathutil v1.7.1 // indirect
modernc.org/memory v1.11.0 // indirect
modernc.org/sqlite v1.56.0 // indirect
)
+20
View File
@@ -0,0 +1,20 @@
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI=
github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A=
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
modernc.org/libc v1.74.4 h1:fX1Omw4o2/1C2iRkkIsrQTasJQldLhRmuPreXLoWs9k=
modernc.org/libc v1.74.4/go.mod h1:eeQAS9W3sZeKYMFubydxJpII9ybHWshk+7or7bLG9co=
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
modernc.org/sqlite v1.56.0 h1:/D8e2RfFqoy/Zc6PuC76U28zFwmI/sYx1Kjm4yEn9e0=
modernc.org/sqlite v1.56.0/go.mod h1:yCJ2cmAaIkHQ25oXWrF8H4O1lIfPYPR26yCEDj2P3pQ=
+46 -6
View File
@@ -1,6 +1,10 @@
package GopherGraph package GopherGraph
import "context" import (
"context"
"errors"
"fmt"
)
// NodeFn 代表图中的一个执行节点(Agent)。 // NodeFn 代表图中的一个执行节点(Agent)。
// 它接收当前的上下文和状态 S,处理后返回更新后的状态 S,或者返回错误。 // 它接收当前的上下文和状态 S,处理后返回更新后的状态 S,或者返回错误。
@@ -27,6 +31,7 @@ type Graph[S any] struct {
conditional map[string]RouterFn[S] // 条件路由边:源节点 -> 路由函数 conditional map[string]RouterFn[S] // 条件路由边:源节点 -> 路由函数
interrupts map[string]bool // 标记需要中断的节点:执行完该节点后挂起 interrupts map[string]bool // 标记需要中断的节点:执行完该节点后挂起
parallels map[string]parallelStep[S] // 并发分支:源节点 -> 并发步骤配置 parallels map[string]parallelStep[S] // 并发分支:源节点 -> 并发步骤配置
errs []error // 注册阶段的错误积累(重复注册等),在 Compile 时统一报告
} }
// NewGraph 创建并初始化一个全新强类型的图构建器 // NewGraph 创建并初始化一个全新强类型的图构建器
@@ -36,27 +41,43 @@ func NewGraph[S any]() *Graph[S] {
edges: make(map[string]string), edges: make(map[string]string),
conditional: make(map[string]RouterFn[S]), conditional: make(map[string]RouterFn[S]),
interrupts: make(map[string]bool), interrupts: make(map[string]bool),
parallels: make(map[string]parallelStep[S]), // 初始化 parallels parallels: make(map[string]parallelStep[S]),
} }
} }
// AddNode 向图中注册一个节点(Agent) // AddNode 向图中注册一个节点(Agent)。
// 若同名节点已被注册,不会覆盖原有节点,错误将在 Compile() 时统一报告。
func (g *Graph[S]) AddNode(name string, fn NodeFn[S]) { func (g *Graph[S]) AddNode(name string, fn NodeFn[S]) {
if _, exists := g.nodes[name]; exists {
g.errs = append(g.errs, fmt.Errorf("AddNode: node %q is already registered", name))
return
}
g.nodes[name] = fn g.nodes[name] = fn
} }
// AddEdge 建立一条从 from 节点到 to 节点的静态连接线 // AddEdge 建立一条从 from 节点到 to 节点的静态连接线。
// 若 from 节点已有静态出边,不会覆盖,错误将在 Compile() 时统一报告。
func (g *Graph[S]) AddEdge(from, to string) { func (g *Graph[S]) AddEdge(from, to string) {
if _, exists := g.edges[from]; exists {
g.errs = append(g.errs, fmt.Errorf("AddEdge: node %q already has a static edge registered", from))
return
}
g.edges[from] = to g.edges[from] = to
} }
// AddConditionalEdges 建立一条从 from 节点出发的条件路由 // AddConditionalEdges 建立一条从 from 节点出发的条件路由
// 到底去哪,由 router 函数在运行时根据当前 State 动态决定 // 到底去哪,由 router 函数在运行时根据当前 State 动态决定。
// 若 from 节点已有条件出边,不会覆盖,错误将在 Compile() 时统一报告。
func (g *Graph[S]) AddConditionalEdges(from string, router RouterFn[S]) { func (g *Graph[S]) AddConditionalEdges(from string, router RouterFn[S]) {
if _, exists := g.conditional[from]; exists {
g.errs = append(g.errs, fmt.Errorf("AddConditionalEdges: node %q already has a conditional edge registered", from))
return
}
g.conditional[from] = router g.conditional[from] = router
} }
// AddInterrupt 标记在执行 nodeName 节点“之前”进行中断挂起,等待人工确认 // AddInterrupt 标记在执行 nodeName 节点"之前"进行中断挂起,等待人工确认。
// 重复标记同一节点为幂等操作(设置 true 再设置 true,无副作用)。
func (g *Graph[S]) AddInterrupt(nodeName string) { func (g *Graph[S]) AddInterrupt(nodeName string) {
g.interrupts[nodeName] = true g.interrupts[nodeName] = true
} }
@@ -66,10 +87,29 @@ func (g *Graph[S]) AddInterrupt(nodeName string) {
// targets: 需要并发执行的多个目标节点名称(如 "translate_en", "translate_fr") // targets: 需要并发执行的多个目标节点名称(如 "translate_en", "translate_fr")
// next: 所有并发节点运行完并合并后,下一步要去的节点(如 "publisher") // next: 所有并发节点运行完并合并后,下一步要去的节点(如 "publisher")
// merger: 用于合并各个并发分支状态的函数 // merger: 用于合并各个并发分支状态的函数
// 若 from 节点已有并发出边,不会覆盖,错误将在 Compile() 时统一报告。
func (g *Graph[S]) AddParallelEdges(from string, targets []string, next string, merger ParallelMergeFn[S]) { func (g *Graph[S]) AddParallelEdges(from string, targets []string, next string, merger ParallelMergeFn[S]) {
if _, exists := g.parallels[from]; exists {
g.errs = append(g.errs, fmt.Errorf("AddParallelEdges: node %q already has parallel edges registered", from))
return
}
g.parallels[from] = parallelStep[S]{ g.parallels[from] = parallelStep[S]{
targets: targets, targets: targets,
next: next, next: next,
merger: merger, merger: merger,
} }
} }
// HookFn 是节点执行前后的生命周期钩子函数签名。
// nodeName 为当前执行的节点名称,state 为该节点执行时的状态快照(只读语义,请勿修改)。
// 可通过此钩子接入日志、链路追踪(OpenTelemetry)、监控指标等能力,而无需修改任何 NodeFn。
type HookFn[S any] func(ctx context.Context, nodeName string, state S)
// errorsJoin 兼容多个错误合并为一个,避免 errors.Join 输出格式不够友好。
// 仅供内部 graph.errs 汇总使用。
func errorsJoin(errs []error) error {
if len(errs) == 0 {
return nil
}
return errors.Join(errs...)
}
+442 -1
View File
@@ -3,6 +3,7 @@ package GopherGraph
import ( import (
"context" "context"
"errors" "errors"
"os"
"strings" "strings"
"testing" "testing"
"time" "time"
@@ -256,7 +257,7 @@ func TestParallelExecution(t *testing.T) {
func TestFileCheckpointer(t *testing.T) { func TestFileCheckpointer(t *testing.T) {
// 创建测试临时文件夹(测试结束后 Go 会自动清理) // 创建测试临时文件夹(测试结束后 Go 会自动清理)
tmpDir := t.TempDir() tmpDir := t.TempDir()
fc, err := NewFileCheckpointer[TestState](tmpDir) fc, err := NewFileCheckpointer[TestState](tmpDir, WithPrettyJSON())
if err != nil { if err != nil {
t.Fatalf("创建文件检查点管理器失败: %v", err) t.Fatalf("创建文件检查点管理器失败: %v", err)
} }
@@ -316,4 +317,444 @@ func TestFileCheckpointer(t *testing.T) {
if !thread.IsFinished || thread.State.Value != 700 { if !thread.IsFinished || thread.State.Value != 700 {
t.Errorf("执行结果错误,最终值: %d, 是否结束: %t", thread.State.Value, thread.IsFinished) t.Errorf("执行结果错误,最终值: %d, 是否结束: %t", thread.State.Value, thread.IsFinished)
} }
// 5. 保存已完结的线程,Load 后应附带 ErrAlreadyFinished 提示
if err := fc.Save(ctx, sessionID, thread); err != nil {
t.Fatalf("保存已完结线程失败: %v", err)
}
_, loadErr := fc.Load(ctx, sessionID)
if !errors.Is(loadErr, ErrAlreadyFinished) {
t.Errorf("Load 已完结线程应附带 ErrAlreadyFinished,实际: %v", loadErr)
}
// 6. Delete 后文件应不存在
if err := fc.Delete(ctx, sessionID); err != nil {
t.Fatalf("删除检查点失败: %v", err)
}
_, err = fc.Load(ctx, sessionID)
if err == nil {
t.Error("删除后 Load 应报错(文件不存在)")
}
// 幂等:重复 Delete 不应报错
if err := fc.Delete(ctx, sessionID); err != nil {
t.Errorf("重复 Delete 应幂等不报错,实际: %v", err)
}
}
// 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])
}
}
}
// TestResumeErrors 验证 Resume 对非暂停/已结束线程的错误路径
func TestResumeErrors(t *testing.T) {
g := NewGraph[TestState]()
g.AddNode("A", func(ctx context.Context, s TestState) (TestState, error) {
return s, nil
})
cg, err := g.Compile()
if err != nil {
t.Fatalf("编译图失败: %v", err)
}
// 1. Resume 非暂停的线程
thread := &Thread[TestState]{NextNode: "A"}
_, err = cg.Resume(context.Background(), thread, TestState{})
if !errors.Is(err, ErrNotPaused) {
t.Errorf("期望 ErrNotPaused,实际: %v", err)
}
// 2. Resume 已结束的线程
thread2 := &Thread[TestState]{IsPaused: true, IsFinished: true}
_, err = cg.Resume(context.Background(), thread2, TestState{})
if !errors.Is(err, ErrAlreadyFinished) {
t.Errorf("期望 ErrAlreadyFinished,实际: %v", err)
}
// 3. Engine.Resume 也应返回相同的哨兵错误
engine := NewEngine(cg)
_, err = engine.Resume(context.Background(), &Thread[TestState]{NextNode: "A"}, TestState{})
if !errors.Is(err, ErrNotPaused) {
t.Errorf("Engine.Resume: 期望 ErrNotPaused,实际: %v", err)
}
}
// TestCompileValidationErrors 验证 Compile 对各种非法图结构的校验
func TestCompileValidationErrors(t *testing.T) {
// 1. 空图
g1 := NewGraph[TestState]()
_, err := g1.Compile()
if err == nil || !strings.Contains(err.Error(), "no nodes") {
t.Errorf("空图应报错: %v", err)
}
// 2. 边指向不存在的节点
g2 := NewGraph[TestState]()
g2.AddNode("A", func(ctx context.Context, s TestState) (TestState, error) { return s, nil })
g2.AddEdge("A", "nonexistent")
_, err = g2.Compile()
if err == nil || !strings.Contains(err.Error(), "destination") {
t.Errorf("边指向不存在节点应报 destination 错误: %v", err)
}
// 3. 中断节点不存在
g3 := NewGraph[TestState]()
g3.AddNode("A", func(ctx context.Context, s TestState) (TestState, error) { return s, nil })
g3.AddInterrupt("ghost")
_, err = g3.Compile()
if err == nil || !strings.Contains(err.Error(), "interrupt") {
t.Errorf("不存在的中断节点应报错: %v", err)
}
// 4. 重复注册节点和边
g4 := NewGraph[TestState]()
g4.AddNode("A", func(ctx context.Context, s TestState) (TestState, error) { return s, nil })
g4.AddNode("A", func(ctx context.Context, s TestState) (TestState, error) { return s, nil }) // 重复节点
g4.AddEdge("A", "B")
g4.AddEdge("A", "C") // 重复静态边
g4.AddConditionalEdges("A", func(ctx context.Context, s TestState) (string, error) { return "B", nil })
g4.AddConditionalEdges("A", func(ctx context.Context, s TestState) (string, error) { return "C", nil }) // 重复条件边
g4.AddParallelEdges("A", []string{"B"}, "", func(ctx context.Context, parent TestState, branches []TestState) (TestState, error) { return parent, nil })
g4.AddParallelEdges("A", []string{"C"}, "", func(ctx context.Context, parent TestState, branches []TestState) (TestState, error) { return parent, nil }) // 重复并发边
_, err = g4.Compile()
if err == nil {
t.Fatal("重复注册应报错")
}
errMsg := err.Error()
if !strings.Contains(errMsg, "already registered") {
t.Errorf("应包含节点重复错误: %v", err)
}
if !strings.Contains(errMsg, "static edge registered") {
t.Errorf("应包含静态边重复错误: %v", err)
}
if !strings.Contains(errMsg, "conditional edge registered") {
t.Errorf("应包含条件边重复错误: %v", err)
}
if !strings.Contains(errMsg, "parallel edges registered") {
t.Errorf("应包含并发边重复错误: %v", err)
}
}
// TestCompileEdgeConflict 验证同一节点不能同时存在于多种出边类型中
func TestCompileEdgeConflict(t *testing.T) {
g := NewGraph[TestState]()
g.AddNode("A", func(ctx context.Context, s TestState) (TestState, error) { return s, nil })
g.AddNode("B", func(ctx context.Context, s TestState) (TestState, error) { return s, nil })
// 同时为 A 设置静态边和条件路由
g.AddEdge("A", "B")
g.AddConditionalEdges("A", func(ctx context.Context, s TestState) (string, error) {
return "B", nil
})
_, err := g.Compile()
if err == nil || !strings.Contains(err.Error(), "conflicting") {
t.Errorf("出边冲突应报错: %v", err)
}
}
// TestRouterReturnsNonExistentNode 验证路由函数返回不存在的节点名时的运行时错误
func TestRouterReturnsNonExistentNode(t *testing.T) {
g := NewGraph[TestState]()
g.AddNode("start", func(ctx context.Context, s TestState) (TestState, error) { return s, nil })
g.AddConditionalEdges("start", func(ctx context.Context, s TestState) (string, error) {
return "nonexistent_node", nil
})
cg, err := g.Compile()
if err != nil {
t.Fatalf("编译图失败: %v", err)
}
_, err = cg.Start(context.Background(), "start", TestState{})
if err == nil || !strings.Contains(err.Error(), "not found") {
t.Errorf("路由到不存在节点应返回 runtime error: %v", err)
}
}
// TestParallelWithStateCloner 验证 StateCloner 在并发分支中正确隔离引用类型
func TestParallelWithStateCloner(t *testing.T) {
g := NewGraph[TestState]()
g.AddNode("start", func(ctx context.Context, s TestState) (TestState, error) {
s.Log = append(s.Log, "start")
return s, nil
})
// 两个分支都向 Log 切片追加元素
g.AddNode("branch1", func(ctx context.Context, s TestState) (TestState, error) {
s.Log = append(s.Log, "branch1")
s.Value = 10
return s, nil
})
g.AddNode("branch2", func(ctx context.Context, s TestState) (TestState, error) {
s.Log = append(s.Log, "branch2")
s.Value = 20
return s, nil
})
g.AddParallelEdges("start", []string{"branch1", "branch2"}, "",
func(ctx context.Context, parent TestState, branches []TestState) (TestState, error) {
for _, b := range branches {
parent.Value += b.Value
}
return parent, nil
},
)
cg, err := g.Compile()
if err != nil {
t.Fatalf("编译图失败: %v", err)
}
// 使用 Engine + StateCloner 运行
engine := NewEngine(cg).WithStateCloner(func(s TestState) TestState {
clone := s
clone.Log = append([]string{}, s.Log...)
return clone
})
thread, err := engine.Start(context.Background(), "start", TestState{})
if err != nil {
t.Fatalf("启动图失败: %v", err)
}
if thread.State.Value != 30 {
t.Errorf("期望 Value 为 30,实际为 %d", thread.State.Value)
}
}
// TestHookChaining 验证多次调用 WithPreNodeHook/WithPostNodeHook 时钩子按顺序叠加
func TestHookChaining(t *testing.T) {
g := NewGraph[TestState]()
g.AddNode("A", func(ctx context.Context, s TestState) (TestState, error) {
return s, nil
})
cg, _ := g.Compile()
var log []string
engine := NewEngine(cg).
WithPreNodeHook(func(ctx context.Context, name string, s TestState) {
log = append(log, "pre1:"+name)
}).
WithPreNodeHook(func(ctx context.Context, name string, s TestState) {
log = append(log, "pre2:"+name)
}).
WithPostNodeHook(func(ctx context.Context, name string, s TestState) {
log = append(log, "post1:"+name)
}).
WithPostNodeHook(func(ctx context.Context, name string, s TestState) {
log = append(log, "post2:"+name)
})
engine.Start(context.Background(), "A", TestState{})
expected := []string{"pre1:A", "pre2:A", "post1:A", "post2:A"}
if len(log) != len(expected) {
t.Fatalf("期望 %d 次 hook 触发,实际 %d 次: %v", len(expected), len(log), log)
}
for i, entry := range expected {
if log[i] != entry {
t.Errorf("第 %d 次 hook 期望 %q,实际 %q", i, entry, log[i])
}
}
}
// TestCheckpointerThreadIDValidation 验证 FileCheckpointer 对恶意 threadID 的防御
func TestCheckpointerThreadIDValidation(t *testing.T) {
tmpDir := t.TempDir()
fc, err := NewFileCheckpointer[TestState](tmpDir)
if err != nil {
t.Fatalf("创建 FileCheckpointer 失败: %v", err)
}
thread := &Thread[TestState]{State: TestState{Value: 1}}
ctx := context.Background()
// 空 ID
if err := fc.Save(ctx, "", thread); err == nil {
t.Error("空 threadID 应报错")
}
// 路径穿越
if err := fc.Save(ctx, "../escape", thread); err == nil {
t.Error("包含 .. 的 threadID 应报错")
}
// 目录分隔符
if err := fc.Save(ctx, "sub/dir", thread); err == nil {
t.Error("包含 / 的 threadID 应报错")
}
// Load 同样校验
_, err = fc.Load(ctx, "../escape")
if err == nil {
t.Error("Load 包含 .. 的 threadID 应报错")
}
// 合法 ID 应成功
if err := fc.Save(ctx, "valid-id_123", thread); err != nil {
t.Errorf("合法 threadID 不应报错: %v", err)
}
// Delete 校验
if err := fc.Delete(ctx, "../escape"); err == nil {
t.Error("Delete 包含 .. 的 threadID 应报错")
}
}
// TestFileCheckpointerCompactJSON 验证默认格式为紧凑 JSON(不带缩进)
func TestFileCheckpointerCompactJSON(t *testing.T) {
tmpDir := t.TempDir()
// 默认紧凑格式
fc, _ := NewFileCheckpointer[TestState](tmpDir)
thread := &Thread[TestState]{State: TestState{Value: 42}, NextNode: "A"}
_ = fc.Save(context.Background(), "compact", thread)
data, _ := os.ReadFile(tmpDir + "/compact.json")
if len(data) == 0 {
t.Fatal("文件为空")
}
// 紧凑格式不应包含换行符
if strings.Contains(string(data), "\n") {
t.Error("默认模式应为紧凑 JSON,不应包含换行符")
}
// WithPrettyJSON 模式应包含换行符
fcPretty, _ := NewFileCheckpointer[TestState](tmpDir, WithPrettyJSON())
_ = fcPretty.Save(context.Background(), "pretty", thread)
dataPretty, _ := os.ReadFile(tmpDir + "/pretty.json")
if !strings.Contains(string(dataPretty), "\n") {
t.Error("WithPrettyJSON 模式应包含换行符")
}
} }
+5
View File
@@ -0,0 +1,5 @@
module testmod
go 1.26.2
require github.com/unclesam-ly/GopherGraph v1.1.1 // indirect
+2
View File
@@ -0,0 +1,2 @@
github.com/unclesam-ly/GopherGraph v1.1.1 h1:1jou9QFtsom/eK69Evx+J3HLkZTjKjnuEuRvkir8qi4=
github.com/unclesam-ly/GopherGraph v1.1.1/go.mod h1:+99WeqLyTIYfOs2OSob/SQ3Pnw+R7pxW8+o6D7S4m6g=
+394
View File
@@ -0,0 +1,394 @@
package GopherGraph
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"regexp"
"sync"
"time"
_ "modernc.org/sqlite"
)
// CheckpointVersion 保存多版本检查点的元数据信息。
type CheckpointVersion struct {
ID int64 `json:"id"`
ThreadID string `json:"thread_id"`
Version int64 `json:"version"`
NextNode string `json:"next_node"`
IsPaused bool `json:"is_paused"`
IsFinished bool `json:"is_finished"`
CreatedAt time.Time `json:"created_at"`
}
// SQLiteCheckpointer 实现了基于 SQLite 数据库的多版本追加写入(Append-Only Multi-Version)状态持久化。
type SQLiteCheckpointer[S any] struct {
db *sql.DB
tableName string
pretty bool
ownsDB bool
mu sync.RWMutex
}
// SQLiteCheckpointerOption 是 SQLiteCheckpointer 的可选配置项。
type SQLiteCheckpointerOption func(*sqliteCheckpointerConfig)
type sqliteCheckpointerConfig struct {
tableName string
pretty bool
}
// WithSQLiteTableName 指定状态持久化存储的表名(默认 "checkpoints")。
func WithSQLiteTableName(tableName string) SQLiteCheckpointerOption {
return func(c *sqliteCheckpointerConfig) {
if tableName != "" {
c.tableName = tableName
}
}
}
// WithSQLitePrettyJSON 开启 JSON 缩进格式存储(适合开发调试阶段)。
func WithSQLitePrettyJSON() SQLiteCheckpointerOption {
return func(c *sqliteCheckpointerConfig) {
c.pretty = true
}
}
var validTableNameRegex = regexp.MustCompile(`^[a-zA-Z0-9_]+$`)
func validateTableName(name string) error {
if name == "" {
return fmt.Errorf("table name must not be empty")
}
if !validTableNameRegex.MatchString(name) {
return fmt.Errorf("invalid table name %q: must contain only alphanumeric characters and underscores", name)
}
return nil
}
// NewSQLiteCheckpointer 通过 SQLite 数据库文件路径创建一个多版本状态检查点管理器。
// 会自动打开 SQLite 连接并开启 WAL 模式以保障高并发读写性能。
func NewSQLiteCheckpointer[S any](dbPath string, opts ...SQLiteCheckpointerOption) (*SQLiteCheckpointer[S], error) {
if dbPath == "" {
return nil, fmt.Errorf("dbPath must not be empty")
}
// 连接 SQLite,配置 WAL 模式与 busy_timeout
dsn := fmt.Sprintf("%s?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)", dbPath)
db, err := sql.Open("sqlite", dsn)
if err != nil {
return nil, fmt.Errorf("failed to open sqlite db: %w", err)
}
// 限制为单一共享连接:
// 1. 确保 PRAGMA WAL / busy_timeout 对所有请求生效(不会因多连接而失效);
// 2. SQLite 单写者模型下,单连接消除 SQLITE_BUSY 并发写冲突。
// 读操作在 WAL 模式下仍可与写操作并行,性能不受影响。
db.SetMaxOpenConns(1)
db.SetMaxIdleConns(1)
sc, err := NewSQLiteCheckpointerWithDB[S](db, opts...)
if err != nil {
_ = db.Close()
return nil, err
}
sc.ownsDB = true
return sc, nil
}
// NewSQLiteCheckpointerWithDB 使用已有的 *sql.DB 实例创建一个 SQLite 检查点管理器。
func NewSQLiteCheckpointerWithDB[S any](db *sql.DB, opts ...SQLiteCheckpointerOption) (*SQLiteCheckpointer[S], error) {
if db == nil {
return nil, fmt.Errorf("db must not be nil")
}
cfg := &sqliteCheckpointerConfig{
tableName: "checkpoints",
}
for _, opt := range opts {
opt(cfg)
}
if err := validateTableName(cfg.tableName); err != nil {
return nil, err
}
sc := &SQLiteCheckpointer[S]{
db: db,
tableName: cfg.tableName,
pretty: cfg.pretty,
ownsDB: false,
}
if err := sc.initTable(); err != nil {
return nil, err
}
return sc, nil
}
func (sc *SQLiteCheckpointer[S]) initTable() error {
query := fmt.Sprintf(`
CREATE TABLE IF NOT EXISTS %s (
id INTEGER PRIMARY KEY AUTOINCREMENT,
thread_id TEXT NOT NULL,
version INTEGER NOT NULL,
next_node TEXT NOT NULL,
is_paused BOOLEAN NOT NULL DEFAULT 0,
is_finished BOOLEAN NOT NULL DEFAULT 0,
data TEXT NOT NULL,
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS idx_%s_thread_ver ON %s(thread_id, version DESC);
`, sc.tableName, sc.tableName, sc.tableName)
_, err := sc.db.Exec(query)
if err != nil {
return fmt.Errorf("failed to initialize checkpoints table %q: %w", sc.tableName, err)
}
return nil
}
// Save 将当前 Thread 追加保存为一个新的版本号记录(Append-Only)。
func (sc *SQLiteCheckpointer[S]) Save(ctx context.Context, threadID string, thread *Thread[S]) error {
if err := validateThreadID(threadID); err != nil {
return err
}
if thread == nil {
return fmt.Errorf("thread must not be nil")
}
// Save 使用写锁(Lock)而非读锁:
// SELECT MAX(version) + INSERT 是两步操作,必须在应用层互斥,
// 防止并发 goroutine 读到相同的 maxVersion 并写入重复的版本号。
sc.mu.Lock()
defer sc.mu.Unlock()
if sc.db == nil {
return fmt.Errorf("SQLiteCheckpointer is already closed")
}
var data []byte
var err error
if sc.pretty {
data, err = json.MarshalIndent(thread, "", " ")
} else {
data, err = json.Marshal(thread)
}
if err != nil {
return fmt.Errorf("failed to marshal thread: %w", err)
}
tx, err := sc.db.BeginTx(ctx, nil)
if err != nil {
return fmt.Errorf("failed to begin tx: %w", err)
}
defer func() {
_ = tx.Rollback()
}()
// 获取当前 threadID 的最大版本号
var maxVersion sql.NullInt64
queryMax := fmt.Sprintf("SELECT MAX(version) FROM %s WHERE thread_id = ?", sc.tableName)
if err := tx.QueryRowContext(ctx, queryMax, threadID).Scan(&maxVersion); err != nil {
return fmt.Errorf("failed to query max version: %w", err)
}
newVersion := int64(1)
if maxVersion.Valid {
newVersion = maxVersion.Int64 + 1
}
// 追加插入新版本记录
insertQuery := fmt.Sprintf(`
INSERT INTO %s (thread_id, version, next_node, is_paused, is_finished, data, created_at)
VALUES (?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
`, sc.tableName)
_, err = tx.ExecContext(ctx, insertQuery, threadID, newVersion, thread.NextNode, thread.IsPaused, thread.IsFinished, string(data))
if err != nil {
return fmt.Errorf("failed to insert checkpoint version %d: %w", newVersion, err)
}
if err := tx.Commit(); err != nil {
return fmt.Errorf("failed to commit checkpoint transaction: %w", err)
}
return nil
}
// Load 读取指定 threadID 的最新版本快照。
// 若最新线程状态为已完结 (IsFinished == true),会包装 ErrAlreadyFinished 作为额外信息返回。
func (sc *SQLiteCheckpointer[S]) Load(ctx context.Context, threadID string) (*Thread[S], error) {
if err := validateThreadID(threadID); err != nil {
return nil, err
}
sc.mu.RLock()
defer sc.mu.RUnlock()
if sc.db == nil {
return nil, fmt.Errorf("SQLiteCheckpointer is already closed")
}
query := fmt.Sprintf("SELECT data FROM %s WHERE thread_id = ? ORDER BY version DESC LIMIT 1", sc.tableName)
var dataStr string
err := sc.db.QueryRowContext(ctx, query, threadID).Scan(&dataStr)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil, fmt.Errorf("checkpoint for thread %q not found", threadID)
}
return nil, fmt.Errorf("failed to load latest checkpoint: %w", err)
}
var thread Thread[S]
if err := json.Unmarshal([]byte(dataStr), &thread); err != nil {
return nil, fmt.Errorf("failed to unmarshal thread data: %w", err)
}
if thread.IsFinished {
return &thread, fmt.Errorf("thread %q has already finished: %w", threadID, ErrAlreadyFinished)
}
return &thread, nil
}
// LoadVersion 读取指定 threadID 的特定历史版本号 (version) 的快照。
func (sc *SQLiteCheckpointer[S]) LoadVersion(ctx context.Context, threadID string, version int64) (*Thread[S], error) {
if err := validateThreadID(threadID); err != nil {
return nil, err
}
if version <= 0 {
return nil, fmt.Errorf("version must be positive")
}
sc.mu.RLock()
defer sc.mu.RUnlock()
if sc.db == nil {
return nil, fmt.Errorf("SQLiteCheckpointer is already closed")
}
query := fmt.Sprintf("SELECT data FROM %s WHERE thread_id = ? AND version = ?", sc.tableName)
var dataStr string
err := sc.db.QueryRowContext(ctx, query, threadID, version).Scan(&dataStr)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil, fmt.Errorf("checkpoint for thread %q version %d not found", threadID, version)
}
return nil, fmt.Errorf("failed to load checkpoint version %d: %w", version, err)
}
var thread Thread[S]
if err := json.Unmarshal([]byte(dataStr), &thread); err != nil {
return nil, fmt.Errorf("failed to unmarshal thread data: %w", err)
}
if thread.IsFinished {
return &thread, fmt.Errorf("thread %q (version %d) has already finished: %w", threadID, version, ErrAlreadyFinished)
}
return &thread, nil
}
// ListVersions 获取指定 threadID 的所有历史版本元数据列表(按 version 升序排列)。
func (sc *SQLiteCheckpointer[S]) ListVersions(ctx context.Context, threadID string) ([]CheckpointVersion, error) {
if err := validateThreadID(threadID); err != nil {
return nil, err
}
sc.mu.RLock()
defer sc.mu.RUnlock()
if sc.db == nil {
return nil, fmt.Errorf("SQLiteCheckpointer is already closed")
}
query := fmt.Sprintf("SELECT id, thread_id, version, next_node, is_paused, is_finished, created_at FROM %s WHERE thread_id = ? ORDER BY version ASC", sc.tableName)
rows, err := sc.db.QueryContext(ctx, query, threadID)
if err != nil {
return nil, fmt.Errorf("failed to list versions: %w", err)
}
defer rows.Close()
// 初始化为空 slice 而非 nil,使调用方无需区分 nil 与空列表两种情况
versions := make([]CheckpointVersion, 0)
for rows.Next() {
var v CheckpointVersion
if err := rows.Scan(&v.ID, &v.ThreadID, &v.Version, &v.NextNode, &v.IsPaused, &v.IsFinished, &v.CreatedAt); err != nil {
return nil, fmt.Errorf("failed to scan checkpoint version: %w", err)
}
versions = append(versions, v)
}
if err := rows.Err(); err != nil {
return nil, err
}
return versions, nil
}
// PruneVersions 对指定 threadID 的历史检查点进行剪枝清理,仅保留最新的 keepN 个版本。
func (sc *SQLiteCheckpointer[S]) PruneVersions(ctx context.Context, threadID string, keepN int) error {
if err := validateThreadID(threadID); err != nil {
return err
}
if keepN <= 0 {
return fmt.Errorf("keepN must be greater than 0")
}
sc.mu.Lock()
defer sc.mu.Unlock()
if sc.db == nil {
return fmt.Errorf("SQLiteCheckpointer is already closed")
}
query := fmt.Sprintf(`
DELETE FROM %s
WHERE thread_id = ?
AND id NOT IN (
SELECT id FROM %s WHERE thread_id = ? ORDER BY version DESC LIMIT ?
)
`, sc.tableName, sc.tableName)
// threadID 传两次:第一个参数对应外层 WHERE,第二个对应子查询 WHERE
_, err := sc.db.ExecContext(ctx, query, threadID, threadID, keepN)
if err != nil {
return fmt.Errorf("failed to prune versions: %w", err)
}
return nil
}
// Delete 删除指定 threadID 对应的所有版本历史快照记录(幂等操作)。
func (sc *SQLiteCheckpointer[S]) Delete(ctx context.Context, threadID string) error {
if err := validateThreadID(threadID); err != nil {
return err
}
sc.mu.Lock()
defer sc.mu.Unlock()
if sc.db == nil {
return fmt.Errorf("SQLiteCheckpointer is already closed")
}
query := fmt.Sprintf("DELETE FROM %s WHERE thread_id = ?", sc.tableName)
_, err := sc.db.ExecContext(ctx, query, threadID)
if err != nil {
return fmt.Errorf("failed to delete checkpoints for thread %q: %w", threadID, err)
}
return nil
}
// Close 关闭由 NewSQLiteCheckpointer 创建的底层 SQLite 数据库连接。
// 若数据库连接是由外部创建并传给 NewSQLiteCheckpointerWithDB 的,则调用 Close 为无操作(由外部管理)。
//
// Close 使用写锁(Lock),会等待所有进行中的 Save/Load 等操作(持有 RLock)全部完成后
// 才将 sc.db 置 nil 并关闭连接,保证不会出现 use-after-close 的竞态问题。
func (sc *SQLiteCheckpointer[S]) Close() error {
sc.mu.Lock()
defer sc.mu.Unlock()
if sc.ownsDB && sc.db != nil {
err := sc.db.Close()
sc.db = nil
return err
}
return nil
}
+330
View File
@@ -0,0 +1,330 @@
package GopherGraph
import (
"context"
"database/sql"
"errors"
"path/filepath"
"sync"
"testing"
"time"
)
type SQLiteTestState struct {
Query string
Count int
}
func TestSQLiteCheckpointerBasic(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "test_checkpoints.db")
sc, err := NewSQLiteCheckpointer[SQLiteTestState](dbPath)
if err != nil {
t.Fatalf("Failed to create SQLiteCheckpointer: %v", err)
}
defer sc.Close()
ctx := context.Background()
threadID := "thread-1"
// 验证未存入时 Load 返回错误
_, err = sc.Load(ctx, threadID)
if err == nil {
t.Fatalf("Expected error when loading non-existent thread, got nil")
}
thread := &Thread[SQLiteTestState]{
State: SQLiteTestState{Query: "hello", Count: 1},
NextNode: "node1",
IsPaused: true,
IsFinished: false,
}
// 存入快照
if err := sc.Save(ctx, threadID, thread); err != nil {
t.Fatalf("Failed to save thread: %v", err)
}
// 读取快照
loaded, err := sc.Load(ctx, threadID)
if err != nil {
t.Fatalf("Failed to load thread: %v", err)
}
if loaded.State.Query != "hello" || loaded.State.Count != 1 {
t.Errorf("Loaded state mismatch, got %+v", loaded.State)
}
if loaded.NextNode != "node1" || !loaded.IsPaused || loaded.IsFinished {
t.Errorf("Loaded thread metadata mismatch, got %+v", loaded)
}
// 删除线程记录
if err := sc.Delete(ctx, threadID); err != nil {
t.Fatalf("Failed to delete thread: %v", err)
}
// 校验删除后 Load 失败
_, err = sc.Load(ctx, threadID)
if err == nil {
t.Fatalf("Expected error after delete, got nil")
}
}
func TestSQLiteCheckpointerMultiVersion(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "multiver.db")
sc, err := NewSQLiteCheckpointer[SQLiteTestState](dbPath, WithSQLitePrettyJSON())
if err != nil {
t.Fatalf("Failed to create SQLiteCheckpointer: %v", err)
}
defer sc.Close()
ctx := context.Background()
threadID := "thread-multi"
// 连续保存 3 个不同版本的线程状态
t1 := &Thread[SQLiteTestState]{State: SQLiteTestState{Query: "step1", Count: 10}, NextNode: "node1"}
t2 := &Thread[SQLiteTestState]{State: SQLiteTestState{Query: "step2", Count: 20}, NextNode: "node2"}
t3 := &Thread[SQLiteTestState]{State: SQLiteTestState{Query: "step3", Count: 30}, NextNode: "node3", IsFinished: true}
if err := sc.Save(ctx, threadID, t1); err != nil {
t.Fatalf("Save v1 failed: %v", err)
}
if err := sc.Save(ctx, threadID, t2); err != nil {
t.Fatalf("Save v2 failed: %v", err)
}
if err := sc.Save(ctx, threadID, t3); err != nil {
t.Fatalf("Save v3 failed: %v", err)
}
// 1. 验证 ListVersions
versions, err := sc.ListVersions(ctx, threadID)
if err != nil {
t.Fatalf("ListVersions failed: %v", err)
}
if len(versions) != 3 {
t.Fatalf("Expected 3 versions, got %d", len(versions))
}
for i, v := range versions {
expectedVer := int64(i + 1)
if v.Version != expectedVer {
t.Errorf("Version index %d expected version %d, got %d", i, expectedVer, v.Version)
}
}
// 2. 验证 Load 默认取最新的 v3 (并且因为 IsFinished == true,包装了 ErrAlreadyFinished)
latest, err := sc.Load(ctx, threadID)
if !errors.Is(err, ErrAlreadyFinished) {
t.Fatalf("Expected ErrAlreadyFinished when loading finished latest state, got: %v", err)
}
if latest.State.Query != "step3" || latest.State.Count != 30 {
t.Errorf("Latest state mismatch, got %+v", latest.State)
}
// 3. 验证 LoadVersion 取历史版本 v1 和 v2
v1State, err := sc.LoadVersion(ctx, threadID, 1)
if err != nil {
t.Fatalf("LoadVersion(1) failed: %v", err)
}
if v1State.State.Query != "step1" || v1State.State.Count != 10 || v1State.NextNode != "node1" {
t.Errorf("v1 state mismatch, got %+v", v1State)
}
v2State, err := sc.LoadVersion(ctx, threadID, 2)
if err != nil {
t.Fatalf("LoadVersion(2) failed: %v", err)
}
if v2State.State.Query != "step2" || v2State.State.Count != 20 || v2State.NextNode != "node2" {
t.Errorf("v2 state mismatch, got %+v", v2State)
}
// 4. 验证 LoadVersion(3) 也因为 IsFinished 报 ErrAlreadyFinished
v3State, err := sc.LoadVersion(ctx, threadID, 3)
if !errors.Is(err, ErrAlreadyFinished) {
t.Fatalf("Expected ErrAlreadyFinished for LoadVersion(3), got: %v", err)
}
if v3State.State.Query != "step3" {
t.Errorf("v3 state mismatch, got %+v", v3State)
}
}
func TestSQLiteCheckpointerPruneVersions(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "prune.db")
sc, err := NewSQLiteCheckpointer[SQLiteTestState](dbPath)
if err != nil {
t.Fatalf("Failed to create SQLiteCheckpointer: %v", err)
}
defer sc.Close()
ctx := context.Background()
threadID := "thread-prune"
// 存入 5 个版本
for i := 1; i <= 5; i++ {
tState := &Thread[SQLiteTestState]{
State: SQLiteTestState{Count: i},
NextNode: "node",
}
if err := sc.Save(ctx, threadID, tState); err != nil {
t.Fatalf("Save v%d failed: %v", i, err)
}
}
// 保留最新的 2 个版本 (v4, v5)
if err := sc.PruneVersions(ctx, threadID, 2); err != nil {
t.Fatalf("PruneVersions failed: %v", err)
}
versions, err := sc.ListVersions(ctx, threadID)
if err != nil {
t.Fatalf("ListVersions failed: %v", err)
}
if len(versions) != 2 {
t.Fatalf("Expected 2 versions after prune, got %d", len(versions))
}
if versions[0].Version != 4 || versions[1].Version != 5 {
t.Errorf("Expected remaining versions 4 and 5, got %d and %d", versions[0].Version, versions[1].Version)
}
// 验证旧版本 v1 已不可读取
_, err = sc.LoadVersion(ctx, threadID, 1)
if err == nil {
t.Fatalf("Expected error loading pruned version 1, got nil")
}
// 验证 v4 和 v5 依然能正常读取
v4, err := sc.LoadVersion(ctx, threadID, 4)
if err != nil || v4.State.Count != 4 {
t.Errorf("LoadVersion(4) failed or mismatch: %v, state: %+v", err, v4)
}
}
func TestSQLiteCheckpointerCustomTableAndExternalDB(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "custom.db")
rawDB, err := sql.Open("sqlite", dbPath)
if err != nil {
t.Fatalf("Failed to open sqlite raw DB: %v", err)
}
sc, err := NewSQLiteCheckpointerWithDB[SQLiteTestState](rawDB, WithSQLiteTableName("my_custom_checkpoints"))
if err != nil {
rawDB.Close()
t.Fatalf("Failed to create SQLiteCheckpointerWithDB: %v", err)
}
ctx := context.Background()
threadID := "thread-custom"
thread := &Thread[SQLiteTestState]{State: SQLiteTestState{Query: "custom table"}}
if err := sc.Save(ctx, threadID, thread); err != nil {
t.Fatalf("Save failed: %v", err)
}
loaded, err := sc.Load(ctx, threadID)
if err != nil {
t.Fatalf("Load failed: %v", err)
}
if loaded.State.Query != "custom table" {
t.Errorf("Loaded state mismatch, got %s", loaded.State.Query)
}
// 测试 Close 不会关闭外部传入的 rawDB (ownsDB == false)
if err := sc.Close(); err != nil {
t.Fatalf("Close failed: %v", err)
}
if err := rawDB.Ping(); err != nil {
t.Fatalf("External DB should remain open after sc.Close(): %v", err)
}
rawDB.Close()
}
func TestSQLiteCheckpointerThreadIDValidation(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "val.db")
sc, err := NewSQLiteCheckpointer[SQLiteTestState](dbPath)
if err != nil {
t.Fatalf("Failed to create SQLiteCheckpointer: %v", err)
}
defer sc.Close()
ctx := context.Background()
invalidIDs := []string{"", "../malicious", "foo/bar", "foo\\bar"}
dummyThread := &Thread[SQLiteTestState]{State: SQLiteTestState{Query: "test"}}
for _, id := range invalidIDs {
if err := sc.Save(ctx, id, dummyThread); err == nil {
t.Errorf("Expected Save error for invalid threadID %q, got nil", id)
}
if _, err := sc.Load(ctx, id); err == nil {
t.Errorf("Expected Load error for invalid threadID %q, got nil", id)
}
if err := sc.Delete(ctx, id); err == nil {
t.Errorf("Expected Delete error for invalid threadID %q, got nil", id)
}
if _, err := sc.ListVersions(ctx, id); err == nil {
t.Errorf("Expected ListVersions error for invalid threadID %q, got nil", id)
}
}
}
func TestSQLiteCheckpointerConcurrency(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "concurrent.db")
sc, err := NewSQLiteCheckpointer[SQLiteTestState](dbPath)
if err != nil {
t.Fatalf("Failed to create SQLiteCheckpointer: %v", err)
}
defer sc.Close()
ctx := context.Background()
threadID := "thread-concurrent"
var wg sync.WaitGroup
numRoutines := 10
savesPerRoutine := 5
for i := 0; i < numRoutines; i++ {
wg.Add(1)
go func(routineID int) {
defer wg.Done()
for j := 0; j < savesPerRoutine; j++ {
tState := &Thread[SQLiteTestState]{
State: SQLiteTestState{Count: routineID*100 + j},
NextNode: "step",
}
if err := sc.Save(ctx, threadID, tState); err != nil {
t.Errorf("Concurrent Save failed for routine %d step %d: %v", routineID, j, err)
}
time.Sleep(1 * time.Millisecond)
}
}(i)
}
wg.Wait()
versions, err := sc.ListVersions(ctx, threadID)
if err != nil {
t.Fatalf("ListVersions failed after concurrency test: %v", err)
}
expectedTotal := numRoutines * savesPerRoutine
if len(versions) != expectedTotal {
t.Fatalf("Expected %d total versions, got %d", expectedTotal, len(versions))
}
// 验证版本号单调递增且不重复 (1 ~ expectedTotal)
for i, v := range versions {
expectedVer := int64(i + 1)
if v.Version != expectedVer {
t.Errorf("Version mismatch at index %d, expected %d, got %d", i, expectedVer, v.Version)
}
}
}