Merge pull request 'Record tool intent before execution and add file checkpoints' (#14) from feature/durable-rewind into main
This commit was merged in pull request #14.
This commit is contained in:
@@ -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 != "" {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)},
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
+98
-6
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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."
|
||||
|
||||
@@ -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 ""
|
||||
|
||||
Reference in New Issue
Block a user