diff --git a/README.md b/README.md index 1b7e9d9..eea049d 100644 --- a/README.md +++ b/README.md @@ -46,7 +46,7 @@ package main import ( "context" - GopherGraph "github.com/uncesam-ly/GopherGraph" + GopherGraph "github.com/unclesam-ly/GopherGraph" ) func main() { @@ -187,7 +187,7 @@ package main import ( "context" - GopherGraph "github.com/uncesam-ly/GopherGraph" + GopherGraph "github.com/unclesam-ly/GopherGraph" ) func main() { @@ -327,7 +327,7 @@ package main import ( "context" - GopherGraph "github.com/uncesam-ly/GopherGraph" + GopherGraph "github.com/unclesam-ly/GopherGraph" ) func main() { diff --git a/checkpoint.go b/checkpoint.go index c758c1e..febb316 100644 --- a/checkpoint.go +++ b/checkpoint.go @@ -6,6 +6,7 @@ import ( "fmt" "os" "path/filepath" + "strings" ) // Checkpointer 定义了持久化状态的接口,S 是工作流的强类型状态。 @@ -31,8 +32,13 @@ func NewFileCheckpointer[S any](dir string) (*FileCheckpointer[S], error) { return &FileCheckpointer[S]{dir: dir}, nil } -// Save 将 Thread 序列化为 JSON 并写入本地文件 +// Save 将 Thread 序列化为 JSON 并写入本地文件。 +// 采用“先写临时文件再原子重命名”的方式,防止写入中途程序崩溃导致文件损坏。 func (fc *FileCheckpointer[S]) Save(ctx context.Context, threadID string, thread *Thread[S]) error { + if err := validateThreadID(threadID); err != nil { + return err + } + path := filepath.Join(fc.dir, threadID+".json") // 使用 MarshalIndent 方便人工阅读生成的 JSON 文件 @@ -41,11 +47,23 @@ func (fc *FileCheckpointer[S]) Save(ctx context.Context, threadID string, thread return fmt.Errorf("failed to marshal thread: %w", err) } - return os.WriteFile(path, data, 0644) + // 原子写入:先写临时文件,再 Rename 替换,确保断电/崩溃安全 + tmpPath := path + ".tmp" + 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 { + return fmt.Errorf("failed to rename temp file: %w", err) + } + return nil } // Load 从本地文件中读取 JSON 并反序列化为 Thread。 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") data, err := os.ReadFile(path) @@ -60,3 +78,15 @@ func (fc *FileCheckpointer[S]) Load(ctx context.Context, threadID string) (*Thre return &thread, 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 +} diff --git a/doc.go b/doc.go new file mode 100644 index 0000000..51b5351 --- /dev/null +++ b/doc.go @@ -0,0 +1,19 @@ +// 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] 持久化接口:内置 [FileCheckpointer] 实现基于 JSON 文件的状态存储。 +// +// # 零外部依赖 +// +// GopherGraph 仅使用 Go 标准库,无任何第三方依赖。 +package GopherGraph diff --git a/engine.go b/engine.go index 5f4d64b..f92deef 100644 --- a/engine.go +++ b/engine.go @@ -37,7 +37,7 @@ func (g *Graph[S]) Compile() (*CompiledGraph[S], error) { return nil, fmt.Errorf("compile error: edge origin %q does not exist", from) } 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 +79,60 @@ 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]{ - nodes: g.nodes, - edges: g.edges, - conditional: g.conditional, - interrupts: g.interrupts, - parallels: g.parallels, + nodes: nodesCopy, + edges: edgesCopy, + conditional: conditionalCopy, + interrupts: interruptsCopy, + parallels: parallelsCopy, }, 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") + // Start 从指定的起始节点开始运行图,并持续流转,直到遇到中断或运行结束 func (cg *CompiledGraph[S]) Start(ctx context.Context, startNode string, initialState S) (*Thread[S], error) { thread := &Thread[S]{ @@ -102,10 +146,10 @@ func (cg *CompiledGraph[S]) Start(ctx context.Context, startNode string, initial // Resume 恢复执行一个被暂停的线程,并允许注入外部修改后的状态数据(例如人工审批修改后的结果) func (cg *CompiledGraph[S]) Resume(ctx context.Context, thread *Thread[S], modifiedState S) (*Thread[S], error) { if !thread.IsPaused { - return nil, errors.New("cannot resume: thread is not paused") + return nil, ErrNotPaused } if thread.IsFinished { - return nil, errors.New("cannot resume: thread is already finished") + return nil, ErrAlreadyFinished } // 注入人工修改后的状态,并解除暂停标记 @@ -207,11 +251,6 @@ func (cg *CompiledGraph[S]) run(ctx context.Context, thread *Thread[S], opts run // 更新快照中的"下一个节点" thread.NextNode = nextNode - // 如果下一站是终点,直接进入下一次循环触发结束逻辑 - if nextNode == "" { - continue - } - // 核心中断机制:如果即将进入的下一个节点被标记为了中断节点,则在此挂起 if cg.interrupts[nextNode] { thread.IsPaused = true diff --git a/engine_options.go b/engine_options.go index 810db0e..329d822 100644 --- a/engine_options.go +++ b/engine_options.go @@ -76,7 +76,15 @@ func (e *Engine[S]) WithMaxSteps(n int) *Engine[S] { // 钩子函数接收到的 state 是节点执行【前】的状态快照,请勿修改(只读语义)。 // 钩子中的 panic 会向上传播,建议在内部做好 recover。 func (e *Engine[S]) WithPreNodeHook(fn HookFn[S]) *Engine[S] { - e.preNodeHook = fn + 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 } @@ -89,7 +97,15 @@ func (e *Engine[S]) WithPreNodeHook(fn HookFn[S]) *Engine[S] { // // 钩子函数接收到的 state 是节点执行【后】更新过的状态,同为只读语义。 func (e *Engine[S]) WithPostNodeHook(fn HookFn[S]) *Engine[S] { - e.postNodeHook = fn + 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 } @@ -117,10 +133,10 @@ func (e *Engine[S]) Start(ctx context.Context, startNode string, initialState S) // 但会附加所有已注册的 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 + return nil, ErrNotPaused } if thread.IsFinished { - return nil, errAlreadyFinished + return nil, ErrAlreadyFinished } thread.State = modifiedState @@ -129,12 +145,4 @@ func (e *Engine[S]) Resume(ctx context.Context, thread *Thread[S], modifiedState 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 index 598b110..9bfacdf 100644 --- a/examples/hooks/main.go +++ b/examples/hooks/main.go @@ -6,7 +6,7 @@ import ( "log" "time" - GopherGraph "github.com/unclesam-LY/GopherGraph" + GopherGraph "github.com/unclesam-ly/GopherGraph" ) // AgentState 是这个示例工作流的状态 diff --git a/examples/translation/main.go b/examples/translation/main.go index e08e03a..6ef91fa 100644 --- a/examples/translation/main.go +++ b/examples/translation/main.go @@ -7,7 +7,7 @@ import ( "os" "strings" - GopherGraph "github.com/unclesam-LY/GopherGraph" + GopherGraph "github.com/unclesam-ly/GopherGraph" ) // TranslationState 代表我们工作流的共享状态(State) diff --git a/graph_test.go b/graph_test.go index cf9bae1..d82492e 100644 --- a/graph_test.go +++ b/graph_test.go @@ -452,3 +452,226 @@ func TestLifecycleHooks(t *testing.T) { } } } + +// 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) + } +} + +// 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) + } +} +