feat: add safe agent writeback flow

This commit is contained in:
Yoilun
2026-05-25 21:06:32 +08:00
parent ba975112de
commit 4b5237fda2
19 changed files with 969 additions and 24 deletions

View File

@@ -6,6 +6,7 @@ type AgentDefinition struct {
ID string `json:"id"`
FilePath string `json:"filePath"`
FileName string `json:"fileName"`
Content string `json:"content"`
Name string `json:"name"`
Description string `json:"description"`
DeveloperInstructions string `json:"developerInstructions"`
@@ -15,3 +16,25 @@ type AgentDefinition struct {
ParseError string `json:"parseError,omitempty"`
DraftStatus string `json:"draftStatus"`
}
type FieldChange struct {
Field string `json:"field"`
Before string `json:"before"`
After string `json:"after"`
}
type DraftValidation struct {
Valid bool `json:"valid"`
Errors []string `json:"errors"`
Diff string `json:"diff"`
TargetPath string `json:"targetPath"`
CurrentHash string `json:"currentHash"`
FieldChanges []FieldChange `json:"fieldChanges"`
}
type WriteResult struct {
Status string `json:"status"`
TargetPath string `json:"targetPath"`
BackupPath string `json:"backupPath"`
CurrentHash string `json:"currentHash"`
}

View File

@@ -90,6 +90,7 @@ func (s Store) readOne(fileName string) AgentDefinition {
def.ParseError = err.Error()
return def
}
def.Content = string(data)
values, err := parseSimpleTOML(string(data))
if err != nil {
def.ParseStatus = "invalid"

View File

@@ -0,0 +1,228 @@
package agents
import (
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"os"
"path/filepath"
"regexp"
"sort"
"strings"
"time"
"codex-agent-manager/internal/codexhome"
)
var ErrWriteConflict = errors.New("目标文件已在校验后发生变化")
var safeAgentID = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_-]*$`)
func (s Store) ValidateDraft(id string, content string) (DraftValidation, error) {
targetPath, current, _, err := s.readWriteTarget(id)
if err != nil {
return DraftValidation{}, err
}
result := DraftValidation{
TargetPath: targetPath,
CurrentHash: hashBytes(current),
}
currentFields, currentErr := parseSimpleTOML(string(current))
draftFields, draftErr := parseSimpleTOML(content)
if draftErr != nil {
result.Valid = false
result.Errors = []string{draftErr.Error()}
result.Diff = simpleDiff(string(current), content)
return result, nil
}
result.Valid = true
result.Errors = []string{}
result.Diff = simpleDiff(string(current), content)
if currentErr == nil {
result.FieldChanges = changedFields(currentFields, draftFields)
}
return result, nil
}
func (s Store) WriteDraft(id string, content string, expectedHash string) (WriteResult, error) {
validation, err := s.ValidateDraft(id, content)
if err != nil {
return WriteResult{}, err
}
if !validation.Valid {
return WriteResult{}, errors.New(strings.Join(validation.Errors, ""))
}
if expectedHash == "" || validation.CurrentHash != expectedHash {
return WriteResult{}, ErrWriteConflict
}
targetPath, current, mode, err := s.readWriteTarget(id)
if err != nil {
return WriteResult{}, err
}
if hashBytes(current) != expectedHash {
return WriteResult{}, ErrWriteConflict
}
backupPath, err := s.createBackup(targetPath, current, mode)
if err != nil {
return WriteResult{}, err
}
if err := atomicWrite(targetPath, []byte(content), mode); err != nil {
return WriteResult{}, err
}
return WriteResult{
Status: "written",
TargetPath: targetPath,
BackupPath: backupPath,
CurrentHash: hashBytes([]byte(content)),
}, nil
}
func (s Store) readWriteTarget(id string) (string, []byte, os.FileMode, error) {
if !safeAgentID.MatchString(id) {
return "", nil, 0, codexhome.ErrForbiddenPath
}
agentsPath := filepath.Join(s.CodexHome, "agents")
if info, err := os.Lstat(agentsPath); err != nil {
return "", nil, 0, err
} else if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
return "", nil, 0, codexhome.ErrForbiddenPath
}
fileName := id + ".toml"
targetPath, err := codexhome.ResolveAgentTOML(s.CodexHome, fileName)
if err != nil {
return "", nil, 0, err
}
info, err := os.Lstat(targetPath)
if err != nil {
return "", nil, 0, err
}
if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
return "", nil, 0, codexhome.ErrForbiddenPath
}
data, err := os.ReadFile(targetPath)
if err != nil {
return "", nil, 0, err
}
return targetPath, data, info.Mode().Perm(), nil
}
func (s Store) createBackup(targetPath string, content []byte, mode os.FileMode) (string, error) {
backupPath := fmt.Sprintf("%s.bak-%s", targetPath, time.Now().UTC().Format("20060102T150405.000000000Z"))
file, err := os.OpenFile(backupPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, mode)
if err != nil {
return "", err
}
if _, err := file.Write(content); err != nil {
_ = file.Close()
_ = os.Remove(backupPath)
return "", err
}
if err := file.Close(); err != nil {
_ = os.Remove(backupPath)
return "", err
}
return backupPath, nil
}
func atomicWrite(targetPath string, content []byte, mode os.FileMode) error {
dir := filepath.Dir(targetPath)
base := filepath.Base(targetPath)
tmp, err := os.CreateTemp(dir, "."+base+".tmp-*")
if err != nil {
return err
}
tmpPath := tmp.Name()
defer func() {
_ = os.Remove(tmpPath)
}()
if err := tmp.Chmod(mode); err != nil {
_ = tmp.Close()
return err
}
if _, err := tmp.Write(content); err != nil {
_ = tmp.Close()
return err
}
if err := tmp.Close(); err != nil {
return err
}
return os.Rename(tmpPath, targetPath)
}
func hashBytes(data []byte) string {
sum := sha256.Sum256(data)
return hex.EncodeToString(sum[:])
}
func changedFields(before map[string]string, after map[string]string) []FieldChange {
keys := map[string]bool{}
for key := range before {
keys[key] = true
}
for key := range after {
keys[key] = true
}
ordered := make([]string, 0, len(keys))
for key := range keys {
ordered = append(ordered, key)
}
sort.Strings(ordered)
changes := make([]FieldChange, 0)
for _, key := range ordered {
if before[key] == after[key] {
continue
}
changes = append(changes, FieldChange{Field: key, Before: before[key], After: after[key]})
}
return changes
}
func simpleDiff(before string, after string) string {
if before == after {
return "无差异\n"
}
beforeLines := splitLines(before)
afterLines := splitLines(after)
prefix := 0
for prefix < len(beforeLines) && prefix < len(afterLines) && beforeLines[prefix] == afterLines[prefix] {
prefix++
}
suffix := 0
for suffix < len(beforeLines)-prefix &&
suffix < len(afterLines)-prefix &&
beforeLines[len(beforeLines)-1-suffix] == afterLines[len(afterLines)-1-suffix] {
suffix++
}
var b strings.Builder
b.WriteString("--- current\n+++ draft\n@@\n")
for _, line := range beforeLines[prefix : len(beforeLines)-suffix] {
b.WriteString("-")
b.WriteString(line)
b.WriteString("\n")
}
for _, line := range afterLines[prefix : len(afterLines)-suffix] {
b.WriteString("+")
b.WriteString(line)
b.WriteString("\n")
}
return b.String()
}
func splitLines(input string) []string {
trimmed := strings.TrimSuffix(input, "\n")
if trimmed == "" {
return []string{}
}
return strings.Split(trimmed, "\n")
}

View File

@@ -0,0 +1,161 @@
package agents
import (
"errors"
"os"
"path/filepath"
"strings"
"testing"
)
func TestValidateInvalidTOMLReturnsInvalidAndDoesNotWrite(t *testing.T) {
store, target := writebackFixture(t, `name = "旧名称"`+"\n")
got, err := store.ValidateDraft("backend", `name = "未闭合`+"\n")
if err != nil {
t.Fatalf("ValidateDraft returned error: %v", err)
}
if got.Valid {
t.Fatalf("validation should be invalid: %#v", got)
}
if len(got.Errors) == 0 {
t.Fatalf("expected validation errors: %#v", got)
}
assertFileContent(t, target, `name = "旧名称"`+"\n")
}
func TestValidateValidTOMLReturnsDiffAndCurrentHash(t *testing.T) {
store, _ := writebackFixture(t, `name = "旧名称"`+"\ndescription = \"旧描述\"\n")
got, err := store.ValidateDraft("backend", `name = "新名称"`+"\ndescription = \"旧描述\"\n")
if err != nil {
t.Fatalf("ValidateDraft returned error: %v", err)
}
if !got.Valid {
t.Fatalf("validation should be valid: %#v", got)
}
if got.CurrentHash == "" {
t.Fatalf("current hash must be returned: %#v", got)
}
if !strings.Contains(got.Diff, `-name = "旧名称"`) || !strings.Contains(got.Diff, `+name = "新名称"`) {
t.Fatalf("diff does not show changed line: %q", got.Diff)
}
if len(got.FieldChanges) != 1 || got.FieldChanges[0].Field != "name" {
t.Fatalf("unexpected field changes: %#v", got.FieldChanges)
}
}
func TestWriteInvalidTOMLRejectsAndLeavesOriginal(t *testing.T) {
store, target := writebackFixture(t, `name = "旧名称"`+"\n")
validation, err := store.ValidateDraft("backend", `name = "新名称"`+"\n")
if err != nil {
t.Fatal(err)
}
_, err = store.WriteDraft("backend", `name = "未闭合`+"\n", validation.CurrentHash)
if err == nil {
t.Fatal("expected invalid TOML write to be rejected")
}
assertFileContent(t, target, `name = "旧名称"`+"\n")
}
func TestWriteExpectedHashMismatchRejectsAndLeavesOriginal(t *testing.T) {
store, target := writebackFixture(t, `name = "旧名称"`+"\n")
validation, err := store.ValidateDraft("backend", `name = "新名称"`+"\n")
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(target, []byte(`name = "用户已改"`+"\n"), 0o644); err != nil {
t.Fatal(err)
}
_, err = store.WriteDraft("backend", `name = "新名称"`+"\n", validation.CurrentHash)
if !errors.Is(err, ErrWriteConflict) {
t.Fatalf("expected conflict error, got %v", err)
}
assertFileContent(t, target, `name = "用户已改"`+"\n")
}
func TestWriteSuccessCreatesBackupAndAtomicallyWrites(t *testing.T) {
store, target := writebackFixture(t, `name = "旧名称"`+"\n")
validation, err := store.ValidateDraft("backend", `name = "新名称"`+"\n")
if err != nil {
t.Fatal(err)
}
got, err := store.WriteDraft("backend", `name = "新名称"`+"\n", validation.CurrentHash)
if err != nil {
t.Fatalf("WriteDraft returned error: %v", err)
}
if got.Status != "written" || got.TargetPath != target || got.BackupPath == "" {
t.Fatalf("unexpected write result: %#v", got)
}
assertFileContent(t, target, `name = "新名称"`+"\n")
assertFileContent(t, got.BackupPath, `name = "旧名称"`+"\n")
}
func TestWriteRejectsUnsafeIDAndSymlinks(t *testing.T) {
store, target := writebackFixture(t, `name = "旧名称"`+"\n")
validation, err := store.ValidateDraft("backend", `name = "新名称"`+"\n")
if err != nil {
t.Fatal(err)
}
for _, id := range []string{"../backend", "nested/backend", "/tmp/backend", "backend.toml"} {
t.Run("unsafe id "+id, func(t *testing.T) {
if _, err := store.ValidateDraft(id, `name = "x"`+"\n"); err == nil {
t.Fatal("expected unsafe validate id to be rejected")
}
if _, err := store.WriteDraft(id, `name = "x"`+"\n", validation.CurrentHash); err == nil {
t.Fatal("expected unsafe write id to be rejected")
}
})
}
if err := os.Remove(target); err != nil {
t.Fatal(err)
}
if err := os.Symlink("../config.toml", target); err != nil {
t.Fatal(err)
}
if _, err := store.WriteDraft("backend", `name = "新名称"`+"\n", validation.CurrentHash); err == nil {
t.Fatal("expected symlink leaf to be rejected")
}
root := t.TempDir()
if err := os.WriteFile(filepath.Join(root, "config.toml"), []byte(`name = "secret"`), 0o600); err != nil {
t.Fatal(err)
}
if err := os.Symlink(".", filepath.Join(root, "agents")); err != nil {
t.Fatal(err)
}
symlinkedStore := Store{CodexHome: root}
if _, err := symlinkedStore.WriteDraft("backend", `name = "x"`+"\n", validation.CurrentHash); err == nil {
t.Fatal("expected symlinked agents directory to be rejected")
}
}
func writebackFixture(t *testing.T, content string) (Store, string) {
t.Helper()
root := t.TempDir()
agentsDir := filepath.Join(root, "agents")
if err := os.MkdirAll(agentsDir, 0o755); err != nil {
t.Fatal(err)
}
target := filepath.Join(agentsDir, "backend.toml")
if err := os.WriteFile(target, []byte(content), 0o644); err != nil {
t.Fatal(err)
}
return Store{CodexHome: root}, target
}
func assertFileContent(t *testing.T, path string, want string) {
t.Helper()
got, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
if string(got) != want {
t.Fatalf("%s content = %q, want %q", path, string(got), want)
}
}