From fb2b3b796c57246306cd7384b60fb671ce371143 Mon Sep 17 00:00:00 2001 From: weichen Date: Sat, 3 Oct 2026 17:27:04 +0000 Subject: [PATCH] Record tool intent before execution and add file checkpoints Why: a crash after a tool was about to run left the call unpaired, and workspace writes could not be rewound. Evidence: GOWORK=off go test ./runtime/ ./runtime/llm/ ./runtime/checkpoint/ ./runtime/telemetry/ ./runtime/tools/ ./runtime/session/ ./adapters/native/ Impact: pending calls rerun only when both the stored and current replay declarations are safe; otherwise the result is an error. Empty fallback lists and a nil checkpoint log keep the previous behavior. Refs: feature/durable-rewind Co-authored-by: Cursor --- adapters/native/backend.go | 22 ++- runtime/checkpoint/checkpoint.go | 242 ++++++++++++++++++++++++++ runtime/checkpoint/checkpoint_test.go | 80 +++++++++ runtime/durable_test.go | 210 ++++++++++++++++++++++ runtime/llm/client.go | 1 + runtime/llm/fallback.go | 79 +++++++++ runtime/llm/fallback_test.go | 74 ++++++++ runtime/llm/types.go | 2 + runtime/loop.go | 104 ++++++++++- runtime/session/intent.go | 57 ++++++ runtime/session/tree.go | 1 + runtime/telemetry/telemetry.go | 112 ++++++++++++ runtime/telemetry/telemetry_test.go | 42 +++++ runtime/tools/builtin.go | 30 +++- runtime/tools/checkpoint_hook_test.go | 42 +++++ runtime/tools/glob.go | 3 +- runtime/tools/registry.go | 58 +++++- 17 files changed, 1138 insertions(+), 21 deletions(-) create mode 100644 runtime/checkpoint/checkpoint.go create mode 100644 runtime/checkpoint/checkpoint_test.go create mode 100644 runtime/durable_test.go create mode 100644 runtime/llm/fallback.go create mode 100644 runtime/llm/fallback_test.go create mode 100644 runtime/session/intent.go create mode 100644 runtime/telemetry/telemetry.go create mode 100644 runtime/telemetry/telemetry_test.go create mode 100644 runtime/tools/checkpoint_hook_test.go diff --git a/adapters/native/backend.go b/adapters/native/backend.go index ff7edf7..d5aeab0 100644 --- a/adapters/native/backend.go +++ b/adapters/native/backend.go @@ -14,6 +14,7 @@ import ( "agent-native/exec" "agent-native/runtime" + "agent-native/runtime/checkpoint" "agent-native/runtime/filestate" "agent-native/runtime/hooks" "agent-native/runtime/llm" @@ -39,13 +40,15 @@ type ModelOptions struct { // ToolOptions 是可选工具装配。nil 字段不注册对应工具。 type ToolOptions struct { - Search tools.SearchProvider - Exec *tools.ExecConfig - Fetch *tools.FetchConfig - MCP *mcp.Manager - SkillsRoots []string - Subagent subagent.Factory - HostTools []tools.Tool + Search tools.SearchProvider + Exec *tools.ExecConfig + Fetch *tools.FetchConfig + MCP *mcp.Manager + SkillsRoots []string + Subagent subagent.Factory + HostTools []tools.Tool + Checkpoints *checkpoint.Log + CheckpointAnchor func() string } // SessionOptions 是单次执行的会话预算与压缩。 @@ -58,6 +61,8 @@ type SessionOptions struct { type ObserveOptions struct { Log *logsys.Logger ExecutionID string + // Trace 交给运行循环。nil 不采集。 + Trace func(name string, attrs map[string]string) func() } // GateOptions 把直连已经使用的停止、提醒和工具钩子交给小队成员。零值不改变行为。 @@ -184,7 +189,7 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, var reg *tools.Registry if profile == exec.ToolProfileWorkspace { - reg = tools.WorkspaceToolsTracked(workDir, filestate.New()) + reg = tools.WorkspaceToolsCheckpointed(workDir, filestate.New(), b.opts.Checkpoints, b.opts.CheckpointAnchor) skills.RegisterLoader(reg, skillStore) if b.opts.Search != nil { tools.RegisterSearch(reg, b.opts.Search) @@ -348,6 +353,7 @@ func (b *Backend) run(ctx context.Context, model string, budget *runtime.Budget, Reminders: b.opts.Reminders, RepeatGuard: b.opts.RepeatGuard, Hooks: gateBus(b.opts.GateOptions), + Trace: b.opts.Trace, OnEvent: func(ev runtime.Event) { logs.WriteString(ev.Type) if ev.Tool != "" { diff --git a/runtime/checkpoint/checkpoint.go b/runtime/checkpoint/checkpoint.go new file mode 100644 index 0000000..2800c98 --- /dev/null +++ b/runtime/checkpoint/checkpoint.go @@ -0,0 +1,242 @@ +// Package checkpoint 记录写文件之前的内容，并按用户消息把文件退回去。 +// 只追踪显式交给它的写入。shell 与 python 的改动不经过这里。 +package checkpoint + +import ( + "bytes" + "encoding/json" + "os" + "path/filepath" + "sync" +) + +// Record 是一次被追踪的写入。Before 是写入前的字节，After 是写入后的字节。 +type Record struct { + Seq int `json:"seq"` + Anchor string `json:"anchor,omitempty"` + Root string `json:"root"` + Path string `json:"path"` + Existed bool `json:"existed"` + Before []byte `json:"before,omitempty"` + After []byte `json:"after,omitempty"` +} + +// Result 是一次回滚的分类。Skipped 是当前内容既不是写入后也不是写入前的文件。 +type Result struct { + Restored []string + Skipped []string +} + +// Log 是追加写的检查点文件。 +type Log struct { + path string + mu sync.Mutex +} + +func Open(path string) (*Log, error) { + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return nil, err + } + f, err := os.OpenFile(path, os.O_CREATE, 0o644) + if err != nil { + return nil, err + } + _ = f.Close() + return &Log{path: path}, nil +} + +func (l *Log) Path() string { return l.path } + +// Record 追加一条检查点。rel 是相对 root 的路径。 +func (l *Log) Record(anchor, root, rel string, before []byte, existed bool, after []byte) error { + if l == nil { + return nil + } + l.mu.Lock() + defer l.mu.Unlock() + all, err := l.read() + if err != nil { + return err + } + rec := Record{ + Seq: len(all) + 1, Anchor: anchor, Root: root, Path: filepath.ToSlash(rel), + Existed: existed, Before: append([]byte(nil), before...), After: append([]byte(nil), after...), + } + raw, err := json.Marshal(rec) + if err != nil { + return err + } + f, err := os.OpenFile(l.path, os.O_APPEND|os.O_WRONLY, 0o644) + if err != nil { + return err + } + defer f.Close() + _, err = f.Write(append(raw, '\n')) + return err +} + +// Plan 列出从 anchor 起（含）会被回滚的文件，不改磁盘。 +// ordered 是用户消息 id 的时间顺序。 +func (l *Log) Plan(anchor string, ordered []string) (Result, error) { + if l == nil { + return Result{}, nil + } + l.mu.Lock() + defer l.mu.Unlock() + all, err := l.read() + if err != nil { + return Result{}, err + } + return classify(all, anchor, ordered), nil +} + +// Rewind 把 anchor 及其后的被追踪写入退回写入前。 +// 当前内容与最后一次记录的 After 不一致时跳过并列入 Skipped。 +func (l *Log) Rewind(anchor string, ordered []string) (Result, error) { + if l == nil { + return Result{}, nil + } + l.mu.Lock() + defer l.mu.Unlock() + all, err := l.read() + if err != nil { + return Result{}, err + } + plan := classify(all, anchor, ordered) + for _, rel := range plan.Restored { + group := groupFor(selected(all, anchor, ordered), rel) + if err := restore(group); err != nil { + return Result{}, err + } + } + return plan, nil +} + +func (l *Log) read() ([]Record, error) { + b, err := os.ReadFile(l.path) + if err != nil { + return nil, err + } + var out []Record + for _, line := range bytes.Split(b, []byte("\n")) { + line = bytes.TrimSpace(line) + if len(line) == 0 { + continue + } + var rec Record + if err := json.Unmarshal(line, &rec); err != nil { + return nil, err + } + out = append(out, rec) + } + return out, nil +} + +func selected(all []Record, anchor string, ordered []string) []Record { + keep := map[string]bool{} + on := false + for _, id := range ordered { + if id == anchor { + on = true + } + if on { + keep[id] = true + } + } + if anchor != "" && len(ordered) == 0 { + keep[anchor] = true + } + var out []Record + for _, rec := range all { + if keep[rec.Anchor] { + out = append(out, rec) + } + } + return out +} + +func classify(all []Record, anchor string, ordered []string) Result { + sel := selected(all, anchor, ordered) + seen := map[string]bool{} + var restored, skipped []string + for _, rec := range sel { + if seen[rec.Path] { + continue + } + seen[rec.Path] = true + group := groupFor(sel, rec.Path) + switch fileState(group) { + case stateRestore: + restored = append(restored, rec.Path) + case stateSkip: + skipped = append(skipped, rec.Path) + } + } + return Result{Restored: restored, Skipped: skipped} +} + +func groupFor(sel []Record, rel string) []Record { + var group []Record + for _, rec := range sel { + if rec.Path == rel { + group = append(group, rec) + } + } + return group +} + +const ( + stateAlready = iota + stateRestore + stateSkip +) + +func fileState(group []Record) int { + if len(group) == 0 { + return stateAlready + } + first, last := group[0], group[len(group)-1] + abs := filepath.Join(last.Root, filepath.FromSlash(last.Path)) + current, err := os.ReadFile(abs) + missing := os.IsNotExist(err) + if err != nil && !missing { + return stateSkip + } + if missing { + if !last.Existed && len(last.After) == 0 { + return stateAlready + } + if !first.Existed && bytes.Equal(nil, first.Before) && len(last.After) > 0 { + return stateSkip + } + return stateSkip + } + if bytes.Equal(current, last.After) { + if !first.Existed { + return stateRestore + } + if bytes.Equal(current, first.Before) { + return stateAlready + } + return stateRestore + } + if first.Existed && bytes.Equal(current, first.Before) { + return stateAlready + } + return stateSkip +} + +func restore(group []Record) error { + if len(group) == 0 { + return nil + } + first := group[0] + abs := filepath.Join(first.Root, filepath.FromSlash(first.Path)) + if !first.Existed { + return os.Remove(abs) + } + if err := os.MkdirAll(filepath.Dir(abs), 0o755); err != nil { + return err + } + return os.WriteFile(abs, first.Before, 0o644) +} diff --git a/runtime/checkpoint/checkpoint_test.go b/runtime/checkpoint/checkpoint_test.go new file mode 100644 index 0000000..32c2383 --- /dev/null +++ b/runtime/checkpoint/checkpoint_test.go @@ -0,0 +1,80 @@ +package checkpoint + +import ( + "os" + "path/filepath" + "testing" +) + +func TestRewindRestoresTrackedWriteAndSkipsExternalChange(t *testing.T) { + root := t.TempDir() + log, err := Open(filepath.Join(t.TempDir(), "cp.jsonl")) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(root, "keep.txt"), []byte("old"), 0o644); err != nil { + t.Fatal(err) + } + if err := log.Record("m1", root, "keep.txt", []byte("old"), true, []byte("new")); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(root, "keep.txt"), []byte("new"), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(root, "side.txt"), []byte("from-shell"), 0o644); err != nil { + t.Fatal(err) + } + if err := log.Record("m1", root, "moved.txt", []byte("a"), true, []byte("b")); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(root, "moved.txt"), []byte("someone-else"), 0o644); err != nil { + t.Fatal(err) + } + + got, err := log.Rewind("m1", []string{"m1"}) + if err != nil { + t.Fatal(err) + } + if len(got.Restored) != 1 || got.Restored[0] != "keep.txt" { + t.Fatalf("restored %#v", got.Restored) + } + if len(got.Skipped) != 1 || got.Skipped[0] != "moved.txt" { + t.Fatalf("skipped %#v", got.Skipped) + } + body, err := os.ReadFile(filepath.Join(root, "keep.txt")) + if err != nil || string(body) != "old" { + t.Fatalf("keep.txt %q %v", body, err) + } + side, err := os.ReadFile(filepath.Join(root, "side.txt")) + if err != nil || string(side) != "from-shell" { + t.Fatalf("untracked file changed: %q %v", side, err) + } + moved, _ := os.ReadFile(filepath.Join(root, "moved.txt")) + if string(moved) != "someone-else" { + t.Fatalf("external edit was reverted: %q", moved) + } +} + +func TestRewindRemovesCreatedFile(t *testing.T) { + root := t.TempDir() + log, err := Open(filepath.Join(t.TempDir(), "cp.jsonl")) + if err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(root, "new.txt"), []byte("hi"), 0o644); err != nil { + t.Fatal(err) + } + if err := log.Record("m2", root, "new.txt", nil, false, []byte("hi")); err != nil { + t.Fatal(err) + } + got, err := log.Rewind("m2", []string{"m1", "m2"}) + if err != nil { + t.Fatal(err) + } + if len(got.Restored) != 1 { + t.Fatalf("restored %#v", got.Restored) + } + if _, err := os.Stat(filepath.Join(root, "new.txt")); !os.IsNotExist(err) { + t.Fatalf("created file still present: %v", err) + } +} diff --git a/runtime/durable_test.go b/runtime/durable_test.go new file mode 100644 index 0000000..435d893 --- /dev/null +++ b/runtime/durable_test.go @@ -0,0 +1,210 @@ +package runtime + +import ( + "context" + "path/filepath" + "strings" + "testing" + + "agent-native/runtime/llm" + "agent-native/runtime/session" + "agent-native/runtime/tools" +) + +type replayTool struct { + countingTool + replay string +} + +func (r *replayTool) Replay() string { + if r.replay == "" { + return tools.ReplayNever + } + return r.replay +} + +func TestPendingUnsafeToolIsNotRerun(t *testing.T) { + dir := t.TempDir() + st, err := session.Open(filepath.Join(dir, "s.jsonl")) + if err != nil { + t.Fatal(err) + } + if err := st.Append(llm.Message{ + Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, + ToolCalls: []llm.ToolCall{{ID: "c1", Name: "write_file", Arguments: `{}`}}, + }); err != nil { + t.Fatal(err) + } + if err := st.AppendToolIntent(session.ToolIntent{CallID: "c1", Name: "write_file", Arguments: `{}`, Replay: tools.ReplayNever}); err != nil { + t.Fatal(err) + } + tool := &countingTool{name: "write_file"} + reg := tools.NewRegistry(dir) + reg.Register(tool) + _, err = Run(context.Background(), Config{ + Model: "fake", Session: st, Tools: reg, MaxTurns: 2, + Stream: script(llm.Message{Role: llm.RoleAssistant, Content: "later", StopReason: llm.StopEnd}), + }, llm.Message{Role: llm.RoleUser, Content: "next"}) + if err != nil { + t.Fatal(err) + } + if tool.calls != 0 { + t.Fatalf("unsafe tool reran: %d", tool.calls) + } + msgs, err := st.Load() + if err != nil { + t.Fatal(err) + } + var paired bool + for _, m := range msgs { + if m.Role == llm.RoleTool && m.ToolCallID == "c1" && m.Content == unknownSideEffect { + paired = true + } + } + if !paired { + t.Fatalf("pair missing: %#v", msgs) + } +} + +func TestPendingSafeToolRerunsOnce(t *testing.T) { + dir := t.TempDir() + st, err := session.Open(filepath.Join(dir, "s.jsonl")) + if err != nil { + t.Fatal(err) + } + if err := st.Append(llm.Message{ + Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, + ToolCalls: []llm.ToolCall{{ID: "c1", Name: "read_file", Arguments: `{}`}}, + }); err != nil { + t.Fatal(err) + } + if err := st.AppendToolIntent(session.ToolIntent{CallID: "c1", Name: "read_file", Arguments: `{}`, Replay: tools.ReplaySafe}); err != nil { + t.Fatal(err) + } + tool := &replayTool{countingTool: countingTool{name: "read_file"}, replay: tools.ReplaySafe} + reg := tools.NewRegistry(dir) + reg.Register(tool) + _, err = Run(context.Background(), Config{ + Model: "fake", Session: st, Tools: reg, MaxTurns: 2, + Stream: script(llm.Message{Role: llm.RoleAssistant, Content: "later", StopReason: llm.StopEnd}), + }, llm.Message{Role: llm.RoleUser, Content: "next"}) + if err != nil { + t.Fatal(err) + } + if tool.calls != 1 { + t.Fatalf("calls %d", tool.calls) + } + msgs, err := st.Load() + if err != nil { + t.Fatal(err) + } + for _, m := range msgs { + if m.Role == llm.RoleTool && m.ToolCallID == "c1" { + if m.Content != "ok" || strings.Contains(m.Content, unknownSideEffect) { + t.Fatalf("result %#v", m) + } + return + } + } + t.Fatal("missing tool result") +} + +func TestStoredSafeCurrentNeverDoesNotRerun(t *testing.T) { + dir := t.TempDir() + st, err := session.Open(filepath.Join(dir, "s.jsonl")) + if err != nil { + t.Fatal(err) + } + if err := st.AppendToolIntent(session.ToolIntent{CallID: "c1", Name: "read_file", Arguments: `{}`, Replay: tools.ReplaySafe}); err != nil { + t.Fatal(err) + } + tool := &replayTool{countingTool: countingTool{name: "read_file"}, replay: tools.ReplayNever} + reg := tools.NewRegistry(dir) + reg.Register(tool) + if err := recoverPending(context.Background(), Config{Session: st, Tools: reg}); err != nil { + t.Fatal(err) + } + if tool.calls != 0 { + t.Fatalf("calls %d", tool.calls) + } +} + +func TestFollowUpIsPersistedBeforeNextRequest(t *testing.T) { + dir := t.TempDir() + st, err := session.Open(filepath.Join(dir, "s.jsonl")) + if err != nil { + t.Fatal(err) + } + sent := false + var saw bool + stream := func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 1) + go func() { + defer close(ch) + for _, m := range messages { + if m.Content == "follow-please" { + loaded, err := st.Load() + if err != nil { + t.Errorf("load %v", err) + } + for _, got := range loaded { + if got.Content == "follow-please" { + saw = true + } + } + } + } + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}} + }() + return ch + } + _, err = Run(context.Background(), Config{ + Model: "fake", Session: st, MaxTurns: 3, Stream: stream, + GetFollowUp: func() []llm.Message { + if sent { + return nil + } + sent = true + return []llm.Message{{Role: llm.RoleUser, Content: "follow-please"}} + }, + }, llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + if !saw { + t.Fatal("follow-up entered the request before it was stored") + } +} + +func TestSpansUseAgreedNamesWithoutPrompt(t *testing.T) { + var names []string + var leaked bool + reg := tools.NewRegistry(t.TempDir()) + reg.Register(&countingTool{name: "read_file"}) + _, err := Run(context.Background(), Config{ + Model: "fake", Tools: reg, MaxTurns: 2, + Trace: func(name string, attrs map[string]string) func() { + names = append(names, name) + if _, ok := attrs["prompt"]; ok || attrs["content"] != "" { + leaked = true + } + return func() {} + }, + Stream: script( + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ID: "1", Name: "read_file", Arguments: `{}`}}}, + llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}, + ), + }, llm.Message{Role: llm.RoleUser, Content: "用户的秘密提示词"}) + if err != nil { + t.Fatal(err) + } + joined := strings.Join(names, ",") + for _, want := range []string{"invoke_agent", "chat", "execute_tool"} { + if !strings.Contains(joined, want) { + t.Fatalf("names %s", joined) + } + } + if leaked { + t.Fatal("prompt leaked into span attrs") + } +} diff --git a/runtime/llm/client.go b/runtime/llm/client.go index 08e1469..acda4df 100644 --- a/runtime/llm/client.go +++ b/runtime/llm/client.go @@ -92,6 +92,7 @@ func (c *Client) run(ctx context.Context, model string, messages []Message, tool msg := fmt.Sprintf("provider status %d: %s", resp.StatusCode, strings.TrimSpace(string(slurp))) ev := Event{ Kind: KindError, + HTTPStatus: resp.StatusCode, Retryable: StatusRetryable(resp.StatusCode) && !ContextOverflow(msg), ErrorMessage: msg, Err: &StreamError{Message: msg, Retryable: StatusRetryable(resp.StatusCode) && !ContextOverflow(msg)}, diff --git a/runtime/llm/fallback.go b/runtime/llm/fallback.go new file mode 100644 index 0000000..3a799e2 --- /dev/null +++ b/runtime/llm/fallback.go @@ -0,0 +1,79 @@ +package llm + +import "context" + +// WithFallback 在尚未转发任何增量时，遇到 5xx 就改用链上下一个模型。 +// fallbacks 为空时返回原来的 StreamFunc，不改变行为。 +func WithFallback(base StreamFunc, fallbacks []string, onSwitch func(from, to string)) StreamFunc { + chain := nonEmpty(fallbacks) + if base == nil || len(chain) == 0 { + return base + } + return func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event { + out := make(chan Event) + go func() { + defer close(out) + models := append([]string{model}, chain...) + for i, current := range models { + forwarded := false + switchNext := false + for ev := range base(ctx, current, messages, tools, opts) { + if isDelta(ev) { + forwarded = true + } + if ev.Kind == KindError && !forwarded && serverError(ev) && i+1 < len(models) { + switchNext = true + if onSwitch != nil { + onSwitch(current, models[i+1]) + } + break + } + select { + case out <- ev: + case <-ctx.Done(): + return + } + if ev.Kind == KindMessageEnd || ev.Kind == KindError { + return + } + } + if !switchNext { + return + } + } + }() + return out + } +} + +func nonEmpty(in []string) []string { + var out []string + seen := map[string]bool{} + for _, s := range in { + s = trimModel(s) + if s == "" || seen[s] { + continue + } + seen[s] = true + out = append(out, s) + } + return out +} + +func trimModel(s string) string { + for len(s) > 0 && (s[0] == ' ' || s[0] == '\t') { + s = s[1:] + } + for len(s) > 0 && (s[len(s)-1] == ' ' || s[len(s)-1] == '\t') { + s = s[:len(s)-1] + } + return s +} + +func isDelta(ev Event) bool { + return ev.Kind == KindTextDelta || ev.Kind == KindReasoningDelta || ev.Kind == KindToolCallDelta +} + +func serverError(ev Event) bool { + return ev.HTTPStatus >= 500 && ev.HTTPStatus <= 599 +} diff --git a/runtime/llm/fallback_test.go b/runtime/llm/fallback_test.go new file mode 100644 index 0000000..97c5be8 --- /dev/null +++ b/runtime/llm/fallback_test.go @@ -0,0 +1,74 @@ +package llm + +import ( + "context" + "testing" +) + +func TestFallbackSwitchesOn5xxBeforeOutput(t *testing.T) { + var switched string + base := func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event { + ch := make(chan Event, 2) + go func() { + defer close(ch) + if model == "primary" { + ch <- Event{Kind: KindError, HTTPStatus: 503, ErrorMessage: "provider status 503", Message: Message{Role: RoleAssistant, StopReason: StopError}} + return + } + ch <- Event{Kind: KindTextDelta, Text: "ok"} + ch <- Event{Kind: KindMessageEnd, Message: Message{Role: RoleAssistant, Content: "ok", StopReason: StopEnd}} + }() + return ch + } + stream := WithFallback(base, []string{"backup"}, func(from, to string) { + switched = from + "->" + to + }) + var text string + var sawError bool + for ev := range stream(context.Background(), "primary", nil, nil, StreamOptions{}) { + if ev.Kind == KindTextDelta { + text += ev.Text + } + if ev.Kind == KindError { + sawError = true + } + } + if switched != "primary->backup" || text != "ok" || sawError { + t.Fatalf("switched %q text %q error %v", switched, text, sawError) + } +} + +func TestFallbackStaysAfterFirstDelta(t *testing.T) { + calls := 0 + base := func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event { + calls++ + ch := make(chan Event, 2) + go func() { + defer close(ch) + ch <- Event{Kind: KindTextDelta, Text: "partial"} + ch <- Event{Kind: KindError, HTTPStatus: 500, Message: Message{Role: RoleAssistant, StopReason: StopError}} + }() + return ch + } + stream := WithFallback(base, []string{"backup"}, func(from, to string) { + t.Fatal("switched after output") + }) + n := 0 + for range stream(context.Background(), "primary", nil, nil, StreamOptions{}) { + n++ + } + if calls != 1 || n != 2 { + t.Fatalf("calls %d events %d", calls, n) + } +} + +func TestEmptyFallbackIsIdentity(t *testing.T) { + base := func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event { + ch := make(chan Event) + close(ch) + return ch + } + if WithFallback(base, nil, nil) == nil { + t.Fatal("nil stream") + } +} diff --git a/runtime/llm/types.go b/runtime/llm/types.go index b000cac..b886af2 100644 --- a/runtime/llm/types.go +++ b/runtime/llm/types.go @@ -103,6 +103,8 @@ type Event struct { Err error Retryable bool ErrorMessage string + // HTTPStatus 是提供方的 HTTP 状态。0 表示这条错误不是 HTTP 响应。 + HTTPStatus int // RetryAfter is the provider's requested wait. HasRetryAfter distinguishes // a real zero from a missing or unparseable header. RetryAfter time.Duration diff --git a/runtime/loop.go b/runtime/loop.go index 93caa78..e6f8882 100644 --- a/runtime/loop.go +++ b/runtime/loop.go @@ -47,11 +47,14 @@ type Event struct { // Config drives one agent session. type Config struct { - Model string - SystemPrompt string - Tools *tools.Registry - Stream llm.StreamFunc - Session *session.Store + Model string + SystemPrompt string + Tools *tools.Registry + Stream llm.StreamFunc + Session *session.Store + // Trace 开始一个 span，返回的函数结束它。nil 表示不采集。 + // 属性里不要放提示词正文。 + Trace func(name string, attrs map[string]string) func() MaxTurns int ToolExecution string GetSteering func() []llm.Message @@ -154,6 +157,8 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { cfg.ToolExecution = ExecSequential } serializeHostCallbacks(&cfg) + endRun := startSpan(cfg, "invoke_agent", map[string]string{"model": cfg.Model}) + defer endRun() var messages []llm.Message if cfg.Session != nil { loaded, err := cfg.Session.Load() @@ -162,6 +167,16 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { } messages = loaded } + if err := recoverPending(ctx, cfg); err != nil { + return Result{}, err + } + if cfg.Session != nil { + loaded, err := cfg.Session.Load() + if err != nil { + return Result{}, err + } + messages = loaded + } if len(messages) == 0 && len(cfg.History) > 0 { messages = append([]llm.Message{}, cfg.History...) } @@ -349,6 +364,8 @@ func systemText(cfg Config) string { } func consume(ctx context.Context, cfg Config, messages []llm.Message) (llm.Event, error) { + endSpan := startSpan(cfg, "chat", map[string]string{"model": cfg.Model}) + defer endSpan() var decls []llm.ToolDecl if cfg.Tools != nil { for _, d := range cfg.Tools.Decls() { @@ -576,6 +593,13 @@ func runOne(ctx context.Context, cfg Config, call llm.ToolCall) llm.Message { } } emit(cfg, Event{Type: "tool_execution_update", Tool: call.Name, Detail: call.Arguments}) + if err := commitToolIntent(cfg, call); err != nil { + msg := llm.Message{Role: llm.RoleTool, Name: call.Name, ToolCallID: call.ID, Content: err.Error(), IsError: true} + emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: msg.Content}) + return msg + } + endSpan := startSpan(cfg, "execute_tool", map[string]string{"tool": call.Name}) + defer endSpan() res := tools.Result{Content: "no tools configured", IsError: true} known := cfg.Tools != nil && cfg.Tools.Has(call.Name) switch { @@ -609,10 +633,78 @@ func settle(cfg Config, last llm.Message, messages []llm.Message, usage llm.Usag } func prepareRequest(cfg Config, messages []llm.Message, turn int, repeats *repeatState, lengthNudge *bool) []llm.Message { - request := clearOldResults(cfg, messages) + request := clearOldResults(cfg, persistedProjection(cfg, messages)) return reminders.Attach(request, turnReminderTexts(cfg, turn, repeats, lengthNudge)) } +// persistedProjection 在有会话存储时只采用已经落盘的消息。 +// steer 与 follow-up 必须先写入存储，下一轮请求才会看见它们。 +func persistedProjection(cfg Config, messages []llm.Message) []llm.Message { + if cfg.Session == nil { + return messages + } + loaded, err := cfg.Session.Load() + if err != nil { + return messages + } + return loaded +} + +const unknownSideEffect = "副作用未知，未重试" + +func recoverPending(ctx context.Context, cfg Config) error { + if cfg.Session == nil { + return nil + } + pending, err := cfg.Session.PendingIntents() + if err != nil { + return err + } + for _, in := range pending { + current := tools.ReplayNever + if cfg.Tools != nil { + current = cfg.Tools.Replay(in.Name) + } + msg := llm.Message{ + Role: llm.RoleTool, Name: in.Name, ToolCallID: in.CallID, + Content: unknownSideEffect, IsError: true, + } + if in.Replay == tools.ReplaySafe && current == tools.ReplaySafe && cfg.Tools != nil { + res := cfg.Tools.Execute(ctx, in.Name, json.RawMessage(in.Arguments)) + msg.Content = res.Content + msg.IsError = res.IsError + } + if err := persist(cfg, msg); err != nil { + return err + } + } + return nil +} + +func commitToolIntent(cfg Config, call llm.ToolCall) error { + if cfg.Session == nil { + return nil + } + replay := tools.ReplayNever + if cfg.Tools != nil { + replay = cfg.Tools.Replay(call.Name) + } + return cfg.Session.AppendToolIntent(session.ToolIntent{ + CallID: call.ID, Name: call.Name, Arguments: call.Arguments, Replay: replay, + }) +} + +func startSpan(cfg Config, name string, attrs map[string]string) func() { + if cfg.Trace == nil { + return func() {} + } + end := cfg.Trace(name, attrs) + if end == nil { + return func() {} + } + return end +} + func clearOldResults(cfg Config, messages []llm.Message) []llm.Message { window := cfg.ContextWindow if window <= 0 { diff --git a/runtime/session/intent.go b/runtime/session/intent.go new file mode 100644 index 0000000..750e707 --- /dev/null +++ b/runtime/session/intent.go @@ -0,0 +1,57 @@ +package session + +import ( + "encoding/json" + + "agent-native/runtime/llm" +) + +// ToolIntent 是一次尚未配对结果的工具调用声明。 +type ToolIntent struct { + CallID string `json:"callId"` + Name string `json:"name"` + Arguments string `json:"arguments,omitempty"` + Replay string `json:"replay"` +} + +// AppendToolIntent 先落下「即将执行」。结果另写一条 tool 消息。 +func (s *Store) AppendToolIntent(in ToolIntent) error { + raw, err := json.Marshal(in) + if err != nil { + return err + } + _, err = s.AppendEntry(Entry{Type: TypeToolIntent, Data: raw}) + return err +} + +// PendingIntents 返回活动分支上还没有对应 tool 结果的意图。 +func (s *Store) PendingIntents() ([]ToolIntent, error) { + s.mu.Lock() + defer s.mu.Unlock() + entries, err := s.readEntriesLocked() + if err != nil { + return nil, err + } + answered := map[string]bool{} + var intents []ToolIntent + for _, e := range pathTo(entries, s.leafLocked(entries)) { + if e.Type == TypeMessage && e.Message != nil && e.Message.Role == llm.RoleTool && e.Message.ToolCallID != "" { + answered[e.Message.ToolCallID] = true + } + if e.Type != TypeToolIntent { + continue + } + var in ToolIntent + if err := json.Unmarshal(e.Data, &in); err != nil { + return nil, err + } + intents = append(intents, in) + } + var pending []ToolIntent + for _, in := range intents { + if !answered[in.CallID] { + pending = append(pending, in) + } + } + return pending, nil +} diff --git a/runtime/session/tree.go b/runtime/session/tree.go index 3c349a4..2afbf7e 100644 --- a/runtime/session/tree.go +++ b/runtime/session/tree.go @@ -26,6 +26,7 @@ const ( TypeLabel EntryType = "label" TypeCustom EntryType = "custom" TypeSessionInfo EntryType = "session_info" + TypeToolIntent EntryType = "tool_intent" ) // Entry is one node in the session tree. Files written before the tree diff --git a/runtime/telemetry/telemetry.go b/runtime/telemetry/telemetry.go new file mode 100644 index 0000000..42d6f92 --- /dev/null +++ b/runtime/telemetry/telemetry.go @@ -0,0 +1,112 @@ +// Package telemetry 记录 GenAI span。默认不采集提示词正文。 +package telemetry + +import ( + "bytes" + "crypto/rand" + "encoding/hex" + "encoding/json" + "net/http" + "os" + "strings" + "time" + + "logging-system/logsys" +) + +// Span 是一次 invoke_agent、chat 或 execute_tool。 +type Span struct { + Name string + TraceID string + SpanID string + Start time.Time + End time.Time + Attrs map[string]string +} + +// StartFunc 返回循环使用的追踪函数。设置了 OTEL_EXPORTER_OTLP_ENDPOINT 时， +// 结束 span 后以 OTLP/HTTP JSON 上报，不引入新的依赖。 +func StartFunc(log *logsys.Logger) func(name string, attrs map[string]string) func() { + endpoint := strings.TrimRight(os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT"), "/") + return func(name string, attrs map[string]string) func() { + sp := Span{ + Name: name, TraceID: newID(16), SpanID: newID(8), + Start: time.Now(), Attrs: publicAttrs(attrs), + } + return func() { + sp.End = time.Now() + if log != nil { + kvs := []any{"span", sp.Name, "trace_id", sp.TraceID, "span_id", sp.SpanID} + for k, v := range sp.Attrs { + kvs = append(kvs, k, v) + } + log.Info("genai span", kvs...) + } + if endpoint != "" { + _ = ExportOTLP(endpoint, sp) + } + } + } +} + +// publicAttrs 丢掉提示词正文。调用方即使传入也不会进入 span。 +func publicAttrs(in map[string]string) map[string]string { + out := map[string]string{} + for k, v := range in { + if k == "prompt" || k == "messages" || k == "content" { + continue + } + out[k] = v + } + return out +} + +// ExportOTLP 把一个 span POST 到端点的 /v1/traces。 +func ExportOTLP(endpoint string, sp Span) error { + endpoint = strings.TrimRight(endpoint, "/") + if !strings.HasSuffix(endpoint, "/v1/traces") { + endpoint += "/v1/traces" + } + body, err := json.Marshal(otlpBody(sp)) + if err != nil { + return err + } + req, err := http.NewRequest(http.MethodPost, endpoint, bytes.NewReader(body)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + return nil +} + +func otlpBody(sp Span) map[string]any { + var attrs []map[string]any + for k, v := range publicAttrs(sp.Attrs) { + attrs = append(attrs, map[string]any{"key": k, "value": map[string]any{"stringValue": v}}) + } + return map[string]any{ + "resourceSpans": []any{map[string]any{ + "scopeSpans": []any{map[string]any{ + "spans": []any{map[string]any{ + "traceId": sp.TraceID, + "spanId": sp.SpanID, + "name": sp.Name, + "startTimeUnixNano": sp.Start.UnixNano(), + "endTimeUnixNano": sp.End.UnixNano(), + "attributes": attrs, + }}, + }}, + }}, + } +} + +func newID(n int) string { + b := make([]byte, n) + _, _ = rand.Read(b) + return hex.EncodeToString(b) +} diff --git a/runtime/telemetry/telemetry_test.go b/runtime/telemetry/telemetry_test.go new file mode 100644 index 0000000..d398c96 --- /dev/null +++ b/runtime/telemetry/telemetry_test.go @@ -0,0 +1,42 @@ +package telemetry + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +func TestSpanOmitsPromptAndUsesAgreedName(t *testing.T) { + var body string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/traces" { + t.Errorf("path %s", r.URL.Path) + } + b, _ := io.ReadAll(r.Body) + body = string(b) + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + sp := Span{ + Name: "chat", TraceID: "abc", SpanID: "def", + Start: time.Unix(0, 0), End: time.Unix(1, 0), + Attrs: map[string]string{"model": "m", "prompt": "用户的秘密提示词"}, + } + if err := ExportOTLP(srv.URL, sp); err != nil { + t.Fatal(err) + } + if !strings.Contains(body, `"name":"chat"`) { + t.Fatalf("span name missing: %s", body) + } + if strings.Contains(body, "用户的秘密提示词") || strings.Contains(body, "prompt") { + t.Fatalf("prompt leaked: %s", body) + } + got := publicAttrs(map[string]string{"tool": "read_file", "content": "全文"}) + if _, ok := got["content"]; ok || got["tool"] != "read_file" { + t.Fatalf("%#v", got) + } +} diff --git a/runtime/tools/builtin.go b/runtime/tools/builtin.go index 590278d..4cb0627 100644 --- a/runtime/tools/builtin.go +++ b/runtime/tools/builtin.go @@ -17,7 +17,8 @@ type readFile struct { files *filestate.Tracker } -func (t *readFile) Name() string { return "read_file" } +func (t *readFile) Name() string { return "read_file" } +func (t *readFile) Replay() string { return ReplaySafe } func (t *readFile) Description() string { return "Read a UTF-8 text file inside the workspace. Without offset/limit returns the whole file (truncated at 200 KiB); with offset/limit returns numbered lines from a 1-based offset." } @@ -79,6 +80,7 @@ func (t *readFile) Execute(_ context.Context, args json.RawMessage) (Result, err type writeFile struct { root string files *filestate.Tracker + mark func(rel string, before []byte, existed bool, after []byte) } func (t *writeFile) Name() string { return "write_file" } @@ -106,15 +108,20 @@ func (t *writeFile) Execute(_ context.Context, args json.RawMessage) (Result, er if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { return Result{Content: err.Error(), IsError: true}, nil } + before, existed := fileBefore(path) if err := atomicWriteFile(path, []byte(obj.Content), 0o644); err != nil { return Result{Content: err.Error(), IsError: true}, nil } + if t.mark != nil { + t.mark(obj.Path, before, existed, []byte(obj.Content)) + } return Result{Content: "wrote " + obj.Path}, nil } type editFile struct { root string files *filestate.Tracker + mark func(rel string, before []byte, existed bool, after []byte) } func (t *editFile) Name() string { return "edit_file" } @@ -156,12 +163,16 @@ func (t *editFile) Execute(_ context.Context, args json.RawMessage) (Result, err if err := atomicWriteFile(path, []byte(updated), 0o644); err != nil { return Result{Content: err.Error(), IsError: true}, nil } + if t.mark != nil { + t.mark(obj.Path, b, true, []byte(updated)) + } return Result{Content: "edited " + obj.Path}, nil } type listDir struct{ root string } -func (t *listDir) Name() string { return "list_dir" } +func (t *listDir) Name() string { return "list_dir" } +func (t *listDir) Replay() string { return ReplaySafe } func (t *listDir) Description() string { return "List names in a workspace directory." } @@ -194,7 +205,8 @@ func (t *listDir) Execute(_ context.Context, args json.RawMessage) (Result, erro type grepTool struct{ root string } -func (t *grepTool) Name() string { return "grep" } +func (t *grepTool) Name() string { return "grep" } +func (t *grepTool) Replay() string { return ReplaySafe } func (t *grepTool) Description() string { return "Search workspace text files. Literal substring by default; regex=true switches to Go regexp, include filters file names (e.g. *.go), ignore_case folds case. Caps at 50 hits." } @@ -284,6 +296,7 @@ func (t *grepTool) Execute(_ context.Context, args json.RawMessage) (Result, err type multiEditTool struct { root string files *filestate.Tracker + mark func(rel string, before []byte, existed bool, after []byte) } func (t *multiEditTool) Name() string { return "multi_edit" } @@ -335,5 +348,16 @@ func (t *multiEditTool) Execute(_ context.Context, args json.RawMessage) (Result if err := atomicWriteFile(path, []byte(text), 0o644); err != nil { return Result{Content: err.Error(), IsError: true}, nil } + if t.mark != nil { + t.mark(obj.Path, b, true, []byte(text)) + } return Result{Content: fmt.Sprintf("applied %d edits to %s", len(obj.Edits), obj.Path)}, nil } + +func fileBefore(path string) ([]byte, bool) { + b, err := os.ReadFile(path) + if err != nil { + return nil, false + } + return b, true +} diff --git a/runtime/tools/checkpoint_hook_test.go b/runtime/tools/checkpoint_hook_test.go new file mode 100644 index 0000000..0d768b8 --- /dev/null +++ b/runtime/tools/checkpoint_hook_test.go @@ -0,0 +1,42 @@ +package tools + +import ( + "context" + "os" + "path/filepath" + "testing" + + "agent-native/runtime/checkpoint" +) + +func TestWriteAndEditRecordCheckpoints(t *testing.T) { + root := t.TempDir() + log, err := checkpoint.Open(filepath.Join(t.TempDir(), "cp.jsonl")) + if err != nil { + t.Fatal(err) + } + reg := WorkspaceToolsCheckpointed(root, nil, log, func() string { return "m1" }) + if reg.Execute(context.Background(), "write_file", []byte(`{"path":"a.txt","content":"one"}`)).IsError { + t.Fatal("write") + } + if reg.Execute(context.Background(), "edit_file", []byte(`{"path":"a.txt","old":"one","new":"two"}`)).IsError { + t.Fatal("edit") + } + if err := os.WriteFile(filepath.Join(root, "shell.txt"), []byte("nope"), 0o644); err != nil { + t.Fatal(err) + } + got, err := log.Rewind("m1", []string{"m1"}) + if err != nil { + t.Fatal(err) + } + if len(got.Restored) != 1 || got.Restored[0] != "a.txt" { + t.Fatalf("restored %#v skipped %#v", got.Restored, got.Skipped) + } + if _, err := os.Stat(filepath.Join(root, "a.txt")); !os.IsNotExist(err) { + t.Fatal("created file was kept") + } + side, err := os.ReadFile(filepath.Join(root, "shell.txt")) + if err != nil || string(side) != "nope" { + t.Fatalf("shell file %q %v", side, err) + } +} diff --git a/runtime/tools/glob.go b/runtime/tools/glob.go index 8c3ddcc..8812c6e 100644 --- a/runtime/tools/glob.go +++ b/runtime/tools/glob.go @@ -15,7 +15,8 @@ const maxGlobResults = 500 type globTool struct{ root string } -func (t *globTool) Name() string { return "glob" } +func (t *globTool) Name() string { return "glob" } +func (t *globTool) Replay() string { return ReplaySafe } func (t *globTool) Description() string { return "Find workspace files matching a glob pattern (** crosses directories, * within one segment). Returns up to 500 relative paths." diff --git a/runtime/tools/registry.go b/runtime/tools/registry.go index 791c1bf..d90e72b 100644 --- a/runtime/tools/registry.go +++ b/runtime/tools/registry.go @@ -9,6 +9,7 @@ import ( "sync" "time" + "agent-native/runtime/checkpoint" "agent-native/runtime/filestate" ) @@ -26,6 +27,26 @@ type Tool interface { Execute(ctx context.Context, args json.RawMessage) (Result, error) } +const ( + // ReplaySafe 表示恢复时可以再执行一次。ReplayNever 是缺省，不重跑。 + ReplaySafe = "safe" + ReplayNever = "never" +) + +// Replayer 是可选声明。未实现或返回其他值都按 never。 +type Replayer interface { + Replay() string +} + +// ReplayOf 读取工具的重放声明。 +func ReplayOf(t Tool) string { + r, ok := t.(Replayer) + if !ok || r.Replay() != ReplaySafe { + return ReplayNever + } + return ReplaySafe +} + // BeforeHook can block a call after schema validation. type BeforeHook func(ctx context.Context, name string, args json.RawMessage) (block bool, reason string) @@ -55,6 +76,18 @@ func (r *Registry) Has(name string) bool { return ok } +// Replay 返回已注册工具的重放声明。未知工具是 never。 +func (r *Registry) Replay(name string) string { + if r == nil { + return ReplayNever + } + t, ok := r.tools[name] + if !ok { + return ReplayNever + } + return ReplayOf(t) +} + func (r *Registry) Register(t Tool) { if _, ok := r.tools[t.Name()]; !ok { r.order = append(r.order, t.Name()) @@ -276,17 +309,36 @@ func WorkspaceTools(root string) *Registry { // WorkspaceToolsTracked 与 WorkspaceTools 相同，但改写已存在的文件前要先在本轮读过。 func WorkspaceToolsTracked(root string, files *filestate.Tracker) *Registry { + return WorkspaceToolsCheckpointed(root, files, nil, nil) +} + +// WorkspaceToolsCheckpointed 在写、改文件前把改前内容交给检查点。log 为空则不记录。 +func WorkspaceToolsCheckpointed(root string, files *filestate.Tracker, log *checkpoint.Log, anchor func() string) *Registry { + mark := writeMark(log, root, anchor) r := NewRegistry(root) r.Register(&readFile{root: root, files: files}) - r.Register(&writeFile{root: root, files: files}) - r.Register(&editFile{root: root, files: files}) - r.Register(&multiEditTool{root: root, files: files}) + r.Register(&writeFile{root: root, files: files, mark: mark}) + r.Register(&editFile{root: root, files: files, mark: mark}) + r.Register(&multiEditTool{root: root, files: files, mark: mark}) r.Register(&listDir{root: root}) r.Register(&grepTool{root: root}) r.Register(&globTool{root: root}) return r } +func writeMark(log *checkpoint.Log, root string, anchor func() string) func(rel string, before []byte, existed bool, after []byte) { + if log == nil { + return nil + } + return func(rel string, before []byte, existed bool, after []byte) { + id := "" + if anchor != nil { + id = anchor() + } + _ = log.Record(id, root, rel, before, existed, after) + } +} + func rejectStaleWrite(files *filestate.Tracker, path string) string { if files == nil { return "" -- 2.54.0