diff --git a/adapters/intentgate/intentgate.go b/adapters/intentgate/intentgate.go index 7bb54e0..c3a9cbb 100644 --- a/adapters/intentgate/intentgate.go +++ b/adapters/intentgate/intentgate.go @@ -35,10 +35,35 @@ func Wrap(next exec.Backend, rec intent.Recognizer, snap *intent.Snapshot, h Han return &gated{next: next, rec: rec, snap: snap, handler: h} } +// WrapRegistry 持有 Registry，每次执行读取当前快照。Reload 之后不用重建装饰器。 +func WrapRegistry(next exec.Backend, rec intent.Recognizer, reg *intent.Registry, h Handler) exec.Backend { + if next == nil || rec == nil || reg == nil || h == nil { + return next + } + return &gatedRegistry{next: next, rec: rec, reg: reg, handler: h} +} + +type gatedRegistry struct { + next exec.Backend + rec intent.Recognizer + reg *intent.Registry + handler Handler +} + +func (g *gatedRegistry) Name() string { return g.next.Name() } + +func (g *gatedRegistry) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) { + return executeGated(ctx, g.next, g.rec, g.reg.Snapshot(), g.handler, req) +} + func (g *gated) Name() string { return g.next.Name() } func (g *gated) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) { - res, err := g.rec.Recognize(ctx, req.Prompt, g.snap) + return executeGated(ctx, g.next, g.rec, g.snap, g.handler, req) +} + +func executeGated(ctx context.Context, next exec.Backend, rec intent.Recognizer, snap *intent.Snapshot, handler Handler, req exec.Request) (*exec.Result, error) { + res, err := rec.Recognize(ctx, req.Prompt, snap) if err != nil { // 识别系统错误 fail-closed：不静默放行可能应拦截的内容，也不吞错。 return nil, err @@ -46,11 +71,11 @@ func (g *gated) Execute(ctx context.Context, req exec.Request) (*exec.Result, er switch res.Branch { case intent.BranchGuardHit, intent.BranchMatched, intent.BranchExpectationHit: default: - return g.next.Execute(ctx, req) // 无裁决：透传原后端 + return next.Execute(ctx, req) // 无裁决：透传原后端 } - out, herr, handled := g.handler(ctx, req, res) + out, herr, handled := handler(ctx, req, res) if !handled { - return g.next.Execute(ctx, req) + return next.Execute(ctx, req) } return out, herr } diff --git a/adapters/intentgate/intentgate_test.go b/adapters/intentgate/intentgate_test.go index 7bdf43a..8445827 100644 --- a/adapters/intentgate/intentgate_test.go +++ b/adapters/intentgate/intentgate_test.go @@ -5,6 +5,7 @@ import ( "errors" "reflect" "testing" + "testing/fstest" "agent-intent/intent" "agent-native/exec" @@ -164,3 +165,33 @@ func TestRecognizeErrorFailsClosed(t *testing.T) { t.Fatal("backend must not run on recognize error") } } + +func TestR2RegistryReloadIsVisibleWithoutRewrap(t *testing.T) { + fsys := fstest.MapFS{"intents.json": &fstest.MapFile{Data: []byte(`{"intents":[{"id":"guard.block","domain":"guard","version":1,"matcher":{"includes_any":["旧词"]}}]}`)}} + reg, err := intent.NewRegistryBuilder(intent.Options{GuardDomains: []string{"guard"}}).BuildFS(context.Background(), fsys, ".") + if err != nil { + t.Fatal(err) + } + next := &stubBackend{name: "native"} + var hits int + be := WrapRegistry(next, newRec(), reg, func(ctx context.Context, req exec.Request, res intent.Result) (*exec.Result, error, bool) { + hits++ + return &exec.Result{Success: false, ErrorReason: "blocked"}, nil, true + }) + if _, err := be.Execute(context.Background(), exec.Request{Prompt: "请处理新词"}); err != nil { + t.Fatal(err) + } + if hits != 0 || next.calls != 1 { + t.Fatalf("before reload hits=%d calls=%d", hits, next.calls) + } + fsys["intents.json"].Data = []byte(`{"intents":[{"id":"guard.block","domain":"guard","version":2,"matcher":{"includes_any":["新词"]}}]}`) + if _, err := reg.Reload(context.Background()); err != nil { + t.Fatal(err) + } + if _, err := be.Execute(context.Background(), exec.Request{Prompt: "请处理新词"}); err != nil { + t.Fatal(err) + } + if hits != 1 || next.calls != 1 { + t.Fatalf("after reload hits=%d calls=%d", hits, next.calls) + } +} diff --git a/adapters/native/backend.go b/adapters/native/backend.go index 0c2e42a..a2fee46 100644 --- a/adapters/native/backend.go +++ b/adapters/native/backend.go @@ -276,9 +276,7 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, if runRes.Err != nil { reason := runRes.Err.Error() - b.mu.Lock() - b.lastFailure[workDir] = reason - b.mu.Unlock() + b.noteFailure(workDir, reason) assessment := exec.AssessStepFailure(exec.StepFailureSignals{ Err: runRes.Err, WorkDir: workDir, @@ -293,9 +291,7 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, } return res, nil } - b.mu.Lock() - delete(b.lastFailure, workDir) - b.mu.Unlock() + b.clearFailure(workDir) patches, matErr := exec.MaterializeTargets(workDir, req.Targets, nil) if matErr != nil { @@ -433,9 +429,7 @@ var resumeUnsafeMarkers = []string{ // rotateIfPoisoned 上次失败属 resume-unsafe 且存在旧会话时归档旧会话。 func (b *Backend) rotateIfPoisoned(workDir string) { - b.mu.Lock() - reason := strings.ToLower(b.lastFailure[workDir]) - b.mu.Unlock() + reason := strings.ToLower(b.failureReason(workDir)) if reason == "" { return } @@ -469,6 +463,36 @@ func (b *Backend) logger() *logsys.Logger { return log.With("execution_id", b.opts.ExecutionID) } +const lastFailureFile = ".native-last-failure" + +func (b *Backend) noteFailure(workDir, reason string) { + b.mu.Lock() + b.lastFailure[workDir] = reason + b.mu.Unlock() + _ = os.WriteFile(filepath.Join(workDir, lastFailureFile), []byte(reason), 0o644) +} + +func (b *Backend) clearFailure(workDir string) { + b.mu.Lock() + delete(b.lastFailure, workDir) + b.mu.Unlock() + _ = os.Remove(filepath.Join(workDir, lastFailureFile)) +} + +func (b *Backend) failureReason(workDir string) string { + b.mu.Lock() + reason := b.lastFailure[workDir] + b.mu.Unlock() + if reason != "" { + return reason + } + raw, err := os.ReadFile(filepath.Join(workDir, lastFailureFile)) + if err != nil { + return "" + } + return string(raw) +} + // systemPrompt 经 promptx 组装：角色指令 + 环境事实 + 技能索引 + 交付物约束。 func (b *Backend) systemPrompt(req exec.Request, model, profile, workDir string, skillIndex []skills.Skill) string { var index []promptx.IndexEntry diff --git a/adapters/native/backend_test.go b/adapters/native/backend_test.go index 1c9c43e..ba350ce 100644 --- a/adapters/native/backend_test.go +++ b/adapters/native/backend_test.go @@ -41,7 +41,7 @@ func TestBackendWorkspaceProfileWritesDeliverable(t *testing.T) { dir := t.TempDir() b := New(Options{Model: "fake", Stream: script( llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ - ID: "c1", Name: "write_file", Arguments: `{"path":"output.md","content":"# 交付\n\n工具写入的交付物内容。","reason":"产出交付物"}`, + ID: "c1", Name: "write_file", Arguments: `{"path":"output.md","content":"# 交付\n\n工具写入的交付物内容。"}`, }}}, llm.Message{Role: llm.RoleAssistant, Content: "已写入 output.md", StopReason: llm.StopEnd}, )}) diff --git a/adapters/native/r2_regression_test.go b/adapters/native/r2_regression_test.go new file mode 100644 index 0000000..699c7f1 --- /dev/null +++ b/adapters/native/r2_regression_test.go @@ -0,0 +1,40 @@ +package native + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestR2PoisonMarkerSurvivesNewProcess(t *testing.T) { + dir := t.TempDir() + if err := os.MkdirAll(filepath.Join(dir, ".agentx"), 0o755); err != nil { + t.Fatal(err) + } + transcript := filepath.Join(dir, ".agentx", "transcript.jsonl") + if err := os.WriteFile(transcript, []byte("{}\n"), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, lastFailureFile), []byte("context length exceeded"), 0o644); err != nil { + t.Fatal(err) + } + b := New(Options{}) + b.rotateIfPoisoned(dir) + if _, err := os.Stat(transcript); !os.IsNotExist(err) { + t.Fatal("poisoned transcript still at the live path") + } + entries, err := os.ReadDir(filepath.Join(dir, ".agentx")) + if err != nil { + t.Fatal(err) + } + var renamed bool + for _, e := range entries { + if strings.Contains(e.Name(), "transcript.poisoned.") { + renamed = true + } + } + if !renamed { + t.Fatal("expected archived transcript") + } +} diff --git a/runtime/llm/r2_regression_test.go b/runtime/llm/r2_regression_test.go new file mode 100644 index 0000000..97bfa86 --- /dev/null +++ b/runtime/llm/r2_regression_test.go @@ -0,0 +1,29 @@ +package llm + +import ( + "context" + "strings" + "testing" +) + +func TestR2SSEJoinsMultilineData(t *testing.T) { + body := strings.Join([]string{ + `data: {"choices":[{"delta":{"content":"hi"}`, + `data: }]}`, + ``, + `data: [DONE]`, + ``, + }, "\n") + var got string + for ev := range ParseSSE(context.Background(), strings.NewReader(body)) { + if ev.Kind == KindTextDelta { + got += ev.Text + } + if ev.Kind == KindError { + t.Fatal(ev.Err) + } + } + if got != "hi" { + t.Fatalf("joined content %q", got) + } +} diff --git a/runtime/llm/sse.go b/runtime/llm/sse.go index 823b253..d7e8230 100644 --- a/runtime/llm/sse.go +++ b/runtime/llm/sse.go @@ -58,19 +58,16 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { finished := "" sc := bufio.NewScanner(r) sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) - for sc.Scan() { + reader := &sseEventReader{sc: sc} + for { + data, ok := reader.next() + if !ok { + break + } if ctx.Err() != nil { emitAbort(out, content.String(), ctx.Err()) return } - line := strings.TrimSpace(sc.Text()) - if line == "" || strings.HasPrefix(line, ":") { - continue - } - if !strings.HasPrefix(line, "data:") { - continue - } - data := strings.TrimSpace(strings.TrimPrefix(line, "data:")) if data == "[DONE]" { break } @@ -168,6 +165,57 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { } } +// sseEventReader 按空行切事件，并把同一事件里的多行 data: 用换行拼回。 +// 已经是完整 JSON 的 data 行单独成事件，这样单行流不用空行也能解析。 +type sseEventReader struct { + sc *bufio.Scanner + carry string +} + +func (r *sseEventReader) next() (string, bool) { + var parts []string + nextPiece := func() (string, bool) { + if r.carry != "" { + piece := r.carry + r.carry = "" + return piece, true + } + if !r.sc.Scan() { + return "", false + } + return r.sc.Text(), true + } + for { + line, ok := nextPiece() + if !ok { + break + } + if strings.TrimSpace(line) == "" { + if len(parts) == 0 { + continue + } + return strings.Join(parts, "\n"), true + } + trimmed := strings.TrimSpace(line) + if strings.HasPrefix(trimmed, ":") { + continue + } + if !strings.HasPrefix(trimmed, "data:") { + continue + } + piece := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + if len(parts) > 0 && json.Valid([]byte(strings.Join(parts, "\n"))) { + r.carry = line + return strings.Join(parts, "\n"), true + } + parts = append(parts, piece) + } + if len(parts) > 0 { + return strings.Join(parts, "\n"), true + } + return "", false +} + func mapFinish(reason string, hasTools bool) StopReason { switch reason { case "tool_calls", "function_call": diff --git a/runtime/loop.go b/runtime/loop.go index b79b9bb..cd9c9e4 100644 --- a/runtime/loop.go +++ b/runtime/loop.go @@ -199,6 +199,14 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { emit(cfg, Event{Type: "message_end", Message: last, PromptTokens: end.Usage.PromptTokens, CompletionTokens: end.Usage.CompletionTokens, TotalTokens: end.Usage.TotalTokens}) if last.StopReason == llm.StopToolCalls && len(last.ToolCalls) > 0 { + if reason := budgetExceeded(cfg.Budget, usage, toolCalls, time.Now()); reason != "" { + var err error + messages, err = appendSkippedToolResults(cfg, last.ToolCalls, messages) + if err != nil { + return Result{}, err + } + return budgetStop(cfg, last, messages, usage, reason), nil + } if cfg.Budget != nil && cfg.Budget.MaxToolCalls > 0 && toolCalls+len(last.ToolCalls) > cfg.Budget.MaxToolCalls { reason := fmt.Sprintf("budget_exceeded: tool calls %d+%d > max %d", toolCalls, len(last.ToolCalls), cfg.Budget.MaxToolCalls) var err error @@ -261,6 +269,11 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { if decision.Replacement != "" && len(messages) > 0 { messages[len(messages)-1].Content = decision.Replacement last.Content = decision.Replacement + note := llm.Message{Role: llm.RoleUser, Content: decision.Replacement} + if err := persist(cfg, note); err != nil { + return Result{}, err + } + messages = append(messages, note) } if decision.Reason != "" { reasonMsg := llm.Message{Role: llm.RoleUser, Content: decision.Reason} @@ -418,6 +431,10 @@ func executeTools(ctx context.Context, cfg Config, assistant llm.Message) []llm. } out := make([]llm.Message, 0, len(assistant.ToolCalls)) for _, call := range assistant.ToolCalls { + if err := ctx.Err(); err != nil { + out = append(out, cancelledToolMessage(call, err)) + continue + } out = append(out, runOne(ctx, cfg, call)) } return out @@ -437,7 +454,20 @@ func executeParallel(ctx context.Context, cfg Config, assistant llm.Message) []l return out } +func cancelledToolMessage(call llm.ToolCall, err error) llm.Message { + return llm.Message{ + Role: llm.RoleTool, + Name: call.Name, + ToolCallID: call.ID, + Content: "not executed: " + err.Error(), + IsError: true, + } +} + func runOne(ctx context.Context, cfg Config, call llm.ToolCall) llm.Message { + if err := ctx.Err(); err != nil { + return cancelledToolMessage(call, err) + } emit(cfg, Event{Type: "tool_execution_start", Tool: call.Name, Detail: call.Arguments}) if cfg.BeforeTool != nil { if err := cfg.BeforeTool(call); err != nil { diff --git a/runtime/mcp/http.go b/runtime/mcp/http.go index 0bf5bcd..a12ccab 100644 --- a/runtime/mcp/http.go +++ b/runtime/mcp/http.go @@ -92,7 +92,31 @@ func (c *httpClient) CallTool(ctx context.Context, name string, args json.RawMes return res, err } -func (c *httpClient) Close() error { return nil } +func (c *httpClient) Close() error { + c.mu.Lock() + sid := c.sessionID + url := c.cfg.URL + client := c.client + c.mu.Unlock() + if sid == "" || url == "" || client == nil { + return nil + } + req, err := http.NewRequest(http.MethodDelete, url, nil) + if err != nil { + return err + } + req.Header.Set("Mcp-Session-Id", sid) + resp, err := client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 2048)) + if resp.StatusCode >= 400 { + return fmt.Errorf("mcp session terminate: http %d", resp.StatusCode) + } + return nil +} func (c *httpClient) roundTrip(ctx context.Context, method string, params any, out any) error { c.mu.Lock() diff --git a/runtime/mcp/manager.go b/runtime/mcp/manager.go index f8e0e07..c5eedb6 100644 --- a/runtime/mcp/manager.go +++ b/runtime/mcp/manager.go @@ -96,7 +96,7 @@ func (m *Manager) RegisterInto(ctx context.Context, reg *tools.Registry) error { } desc := ti.Description if len(desc) > maxDescriptionBytes { - desc = desc[:maxDescriptionBytes] + "…[truncated]" + desc = tools.CutUTF8(desc, maxDescriptionBytes) + "…[truncated]" } schema := ti.InputSchema if len(schema) == 0 { @@ -109,6 +109,12 @@ func (m *Manager) RegisterInto(ctx context.Context, reg *tools.Registry) error { declName: name, description: desc, schema: schema, + redial: func(ctx context.Context) (Client, error) { + m.mu.Lock() + delete(m.clients, cfg.Name) + m.mu.Unlock() + return m.client(ctx, cfg) + }, }) } } @@ -130,7 +136,7 @@ func (m *Manager) Close() { m.clients = map[string]Client{} } -// bridgeTool 把一个 MCP 工具适配为 tools.Tool；调用失败重连一次后重试。 +// bridgeTool 把一个 MCP 工具适配为 tools.Tool。CallTool 失败时重连一次再试。 type bridgeTool struct { client Client serverCfg ServerConfig @@ -138,6 +144,7 @@ type bridgeTool struct { declName string description string schema json.RawMessage + redial func(ctx context.Context) (Client, error) } func (t *bridgeTool) Name() string { return t.declName } @@ -147,9 +154,15 @@ func (t *bridgeTool) Parameters() json.RawMessage { } func (t *bridgeTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) { - callCtx, cancel := context.WithTimeout(ctx, requestTimeout(t.serverCfg)) - defer cancel() - res, err := t.client.CallTool(callCtx, t.toolName, args) + res, err := t.call(ctx, args) + if err != nil && t.redial != nil { + redial := t.redial + t.redial = nil + if c, rerr := redial(ctx); rerr == nil { + t.client = c + res, err = t.call(ctx, args) + } + } if err != nil { return tools.Result{Content: fmt.Sprintf("mcp tool %s failed: %v", t.declName, err), IsError: true}, nil } @@ -162,10 +175,16 @@ func (t *bridgeTool) Execute(ctx context.Context, args json.RawMessage) (tools.R } out := b.String() if len(out) > maxTextResultBytes { - out = out[:maxTextResultBytes] + "\n[truncated]" + out = tools.CutUTF8(out, maxTextResultBytes) + "\n[truncated]" } if res.IsError { return tools.Result{Content: out, IsError: true}, nil } return tools.Result{Content: out}, nil } + +func (t *bridgeTool) call(ctx context.Context, args json.RawMessage) (callToolResult, error) { + callCtx, cancel := context.WithTimeout(ctx, requestTimeout(t.serverCfg)) + defer cancel() + return t.client.CallTool(callCtx, t.toolName, args) +} diff --git a/runtime/mcp/r2_regression_test.go b/runtime/mcp/r2_regression_test.go new file mode 100644 index 0000000..c8f4a46 --- /dev/null +++ b/runtime/mcp/r2_regression_test.go @@ -0,0 +1,66 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +type flakyClient struct { + fail bool + calls int +} + +func (f *flakyClient) Initialize(context.Context) (initializeResult, error) { + return initializeResult{}, nil +} +func (f *flakyClient) ListTools(context.Context) ([]toolInfo, error) { return nil, nil } +func (f *flakyClient) Close() error { return nil } +func (f *flakyClient) CallTool(context.Context, string, json.RawMessage) (callToolResult, error) { + f.calls++ + if f.fail { + return callToolResult{}, fmt.Errorf("down") + } + return callToolResult{Content: []contentBlock{{Type: "text", Text: "back"}}}, nil +} + +func TestR2MCPRetriesOnceAfterReconnect(t *testing.T) { + first := &flakyClient{fail: true} + second := &flakyClient{} + tool := &bridgeTool{ + client: first, + toolName: "echo", + declName: "mcp__s__echo", + redial: func(context.Context) (Client, error) { + return second, nil + }, + } + res, err := tool.Execute(context.Background(), []byte(`{}`)) + if err != nil || res.IsError { + t.Fatalf("retry: %+v %v", res, err) + } + if first.calls != 1 || second.calls != 1 || !strings.Contains(res.Content, "back") { + t.Fatalf("calls first=%d second=%d content=%q", first.calls, second.calls, res.Content) + } +} + +func TestR2HTTPCloseTerminatesSession(t *testing.T) { + var method, sid string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + method = r.Method + sid = r.Header.Get("Mcp-Session-Id") + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + c := &httpClient{cfg: ServerConfig{URL: srv.URL}, client: &http.Client{}, sessionID: "sess-1"} + if err := c.Close(); err != nil { + t.Fatal(err) + } + if method != http.MethodDelete || sid != "sess-1" { + t.Fatalf("terminate %s session %q", method, sid) + } +} diff --git a/runtime/r2_regression_test.go b/runtime/r2_regression_test.go new file mode 100644 index 0000000..e758fce --- /dev/null +++ b/runtime/r2_regression_test.go @@ -0,0 +1,130 @@ +package runtime + +import ( + "context" + "encoding/json" + "path/filepath" + "strings" + "testing" + + "agent-native/runtime/llm" + "agent-native/runtime/session" + "agent-native/runtime/tools" +) + +type countingTool struct { + name string + calls int + hook func() +} + +func (c *countingTool) Name() string { return c.name } +func (c *countingTool) Description() string { return c.name } +func (c *countingTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","properties":{}}`) +} +func (c *countingTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) { + c.calls++ + if c.hook != nil { + c.hook() + } + return tools.Result{Content: "ok"}, nil +} + +func TestR2TokenBudgetSkipsTools(t *testing.T) { + tool := &countingTool{name: "list_dir"} + reg := tools.NewRegistry(t.TempDir()) + reg.Register(tool) + res, err := Run(context.Background(), Config{ + Model: "fake", + Tools: reg, + Budget: &Budget{MaxPromptTokens: 10}, + Stream: usageScript(llm.Usage{PromptTokens: 500}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "list_dir", Arguments: `{}`, + }}}, + ), + }, llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + if tool.calls != 0 { + t.Fatalf("tool ran after token budget: %d", tool.calls) + } + if res.Err == nil || !strings.Contains(res.Err.Error(), "prompt tokens") { + t.Fatalf("budget error: %v", res.Err) + } +} + +func TestR2StopReplacementIsPersisted(t *testing.T) { + store, err := session.Open(filepath.Join(t.TempDir(), "t.jsonl")) + if err != nil { + t.Fatal(err) + } + _, err = Run(context.Background(), Config{ + Model: "fake", + MaxTurns: 3, + Session: store, + Stream: script( + llm.Message{Role: llm.RoleAssistant, Content: "draft", StopReason: llm.StopEnd}, + llm.Message{Role: llm.RoleAssistant, Content: "final", StopReason: llm.StopEnd}, + ), + Stop: func(last llm.Message, info StopInfo) StopDecision { + if info.Continuations == 0 { + return StopDecision{Continue: true, Replacement: "corrected-claim"} + } + return StopDecision{} + }, + }, llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + msgs, err := store.Load() + if err != nil { + t.Fatal(err) + } + var saw bool + for _, m := range msgs { + if m.Content == "corrected-claim" { + saw = true + } + } + if !saw { + t.Fatalf("replacement missing from session: %+v", msgs) + } +} + +func TestR2SequentialToolsStopAfterCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + first := &countingTool{name: "first", hook: cancel} + second := &countingTool{name: "second"} + reg := tools.NewRegistry(t.TempDir()) + reg.Register(first) + reg.Register(second) + res, err := Run(ctx, Config{ + Model: "fake", + Tools: reg, + Stream: script( + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{ + {ID: "1", Name: "first", Arguments: `{}`}, + {ID: "2", Name: "second", Arguments: `{}`}, + }}, + llm.Message{Role: llm.RoleAssistant, Content: "after", StopReason: llm.StopEnd}, + ), + }, llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + if first.calls != 1 || second.calls != 0 { + t.Fatalf("calls first=%d second=%d", first.calls, second.calls) + } + var paired int + for _, m := range res.Messages { + if m.Role == llm.RoleTool { + paired++ + } + } + if paired != 2 { + t.Fatalf("tool results = %d, want a pair for each call", paired) + } +} diff --git a/runtime/session/compact.go b/runtime/session/compact.go index f6014a0..90d4186 100644 --- a/runtime/session/compact.go +++ b/runtime/session/compact.go @@ -3,6 +3,7 @@ package session import ( "bufio" "encoding/json" + "fmt" "os" "agent-native/runtime/llm" @@ -72,13 +73,12 @@ func (s *Store) ReadMeta() map[string]string { func (s *Store) Compact(keepLastN int, summarize func(messages []llm.Message) (string, error)) (bool, error) { s.mu.Lock() entries, err := s.readEntriesLocked() - s.mu.Unlock() if err != nil { + s.mu.Unlock() return false, err } - leaf := "" - s.mu.Lock() - leaf = s.leafLocked(entries) + leaf := s.leafLocked(entries) + generation := len(entries) s.mu.Unlock() path := pathTo(entries, leaf) var messages []llm.Message @@ -102,6 +102,7 @@ func (s *Store) Compact(keepLastN int, summarize func(messages []llm.Message) (s if tailStart <= 0 { return false, nil } + tailStart = alignToolTail(rest, tailStart) head := rest[:tailStart] summary, err := summarize(head) if err != nil { @@ -113,6 +114,21 @@ func (s *Store) Compact(keepLastN int, summarize func(messages []llm.Message) (s } s.mu.Lock() defer s.mu.Unlock() + now, err := s.readEntriesLocked() + if err != nil { + return false, err + } + if len(now) != generation { + return false, fmt.Errorf("compaction conflict: session changed during summary") + } _, err = s.appendEntryLocked(Entry{Type: TypeCompaction, Summary: summary, TailIDs: tailIDs}) return err == nil, err } + +// alignToolTail 让保留尾巴从一条非 tool 消息开始，避免把 tool_calls 和结果拆开。 +func alignToolTail(messages []llm.Message, tailStart int) int { + for tailStart > 0 && messages[tailStart].Role == llm.RoleTool { + tailStart-- + } + return tailStart +} diff --git a/runtime/session/r2_regression_test.go b/runtime/session/r2_regression_test.go new file mode 100644 index 0000000..df0da96 --- /dev/null +++ b/runtime/session/r2_regression_test.go @@ -0,0 +1,85 @@ +package session + +import ( + "path/filepath" + "strings" + "testing" + + "agent-native/runtime/llm" +) + +func TestR2CompactKeepsToolPair(t *testing.T) { + s, err := Open(filepath.Join(t.TempDir(), "s.jsonl")) + if err != nil { + t.Fatal(err) + } + msgs := []llm.Message{ + {Role: llm.RoleUser, Content: "ask"}, + {Role: llm.RoleAssistant, Content: "", ToolCalls: []llm.ToolCall{{ID: "c", Name: "read"}}}, + {Role: llm.RoleTool, ToolCallID: "c", Content: "file-body"}, + } + for _, m := range msgs { + if err := s.Append(m); err != nil { + t.Fatal(err) + } + } + ok, err := s.Compact(1, func(messages []llm.Message) (string, error) { + return "summary", nil + }) + if err != nil || !ok { + t.Fatalf("compact: %v %v", ok, err) + } + loaded, err := s.Load() + if err != nil { + t.Fatal(err) + } + var sawTool bool + for _, m := range loaded { + if m.Role == llm.RoleTool && m.Content == "file-body" { + sawTool = true + } + } + if !sawTool { + t.Fatalf("tool result dropped: %+v", loaded) + } +} + +func TestR2CompactRejectsAppendDuringSummary(t *testing.T) { + s, err := Open(filepath.Join(t.TempDir(), "s.jsonl")) + if err != nil { + t.Fatal(err) + } + for _, text := range []string{"a", "b", "c"} { + if err := s.Append(llm.Message{Role: llm.RoleUser, Content: text}); err != nil { + t.Fatal(err) + } + } + _, err = s.Compact(1, func(messages []llm.Message) (string, error) { + return "summary", s.Append(llm.Message{Role: llm.RoleUser, Content: "raced"}) + }) + if err == nil || !strings.Contains(err.Error(), "compaction conflict") { + t.Fatalf("want conflict, got %v", err) + } +} + +func TestR2AppendVisibleAfterReopen(t *testing.T) { + path := filepath.Join(t.TempDir(), "s.jsonl") + s, err := Open(path) + if err != nil { + t.Fatal(err) + } + if err := s.Append(llm.Message{Role: llm.RoleUser, Content: "durable"}); err != nil { + t.Fatal(err) + } + again, err := Open(path) + if err != nil { + t.Fatal(err) + } + msgs, err := again.Load() + if err != nil { + t.Fatal(err) + } + if len(msgs) != 1 || msgs[0].Content != "durable" { + t.Fatalf("reopen: %+v", msgs) + } +} diff --git a/runtime/session/tree.go b/runtime/session/tree.go index 0417216..3c349a4 100644 --- a/runtime/session/tree.go +++ b/runtime/session/tree.go @@ -180,6 +180,10 @@ func (s *Store) appendEntryLocked(e Entry) (Entry, error) { if _, err := f.Write(append(raw, '\n')); err != nil { return Entry{}, err } + // 单写者：锁只覆盖本进程。Sync 把这一行落到存储，崩溃后重开仍能读到。 + if err := f.Sync(); err != nil { + return Entry{}, err + } if err := s.writeMetaKeyLocked("leaf", e.ID); err != nil { return Entry{}, err } diff --git a/runtime/tools/atomic.go b/runtime/tools/atomic.go new file mode 100644 index 0000000..ab00103 --- /dev/null +++ b/runtime/tools/atomic.go @@ -0,0 +1,30 @@ +package tools + +import "os" + +func atomicWriteFile(path string, data []byte, mode os.FileMode) error { + tmp := path + ".tmp" + f, err := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, mode) + if err != nil { + return err + } + if _, err := f.Write(data); err != nil { + _ = f.Close() + _ = os.Remove(tmp) + return err + } + if err := f.Sync(); err != nil { + _ = f.Close() + _ = os.Remove(tmp) + return err + } + if err := f.Close(); err != nil { + _ = os.Remove(tmp) + return err + } + if err := os.Rename(tmp, path); err != nil { + _ = os.Remove(tmp) + return err + } + return nil +} diff --git a/runtime/tools/builtin.go b/runtime/tools/builtin.go index d39b517..b1dd7f1 100644 --- a/runtime/tools/builtin.go +++ b/runtime/tools/builtin.go @@ -39,7 +39,7 @@ func (t *readFile) Execute(_ context.Context, args json.RawMessage) (Result, err if obj.Offset <= 0 && obj.Limit <= 0 { truncated := false if len(b) > maxReadBytes { - b = b[:maxReadBytes] + b = []byte(CutUTF8(string(b), maxReadBytes)) truncated = true } content := string(b) @@ -93,7 +93,7 @@ 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 } - if err := os.WriteFile(path, []byte(obj.Content), 0o644); err != nil { + if err := atomicWriteFile(path, []byte(obj.Content), 0o644); err != nil { return Result{Content: err.Error(), IsError: true}, nil } return Result{Content: "wrote " + obj.Path}, nil @@ -134,7 +134,7 @@ func (t *editFile) Execute(_ context.Context, args json.RawMessage) (Result, err return Result{Content: fmt.Sprintf("expected one match, found %d", n), IsError: true}, nil } updated := strings.Replace(text, obj.Old, obj.New, 1) - if err := os.WriteFile(path, []byte(updated), 0o644); err != nil { + if err := atomicWriteFile(path, []byte(updated), 0o644); err != nil { return Result{Content: err.Error(), IsError: true}, nil } return Result{Content: "edited " + obj.Path}, nil @@ -279,7 +279,7 @@ type multiEditSpan struct { func (t *multiEditTool) Execute(_ context.Context, args json.RawMessage) (Result, error) { var obj struct { - Path string `json:"path"` + Path string `json:"path"` Edits []multiEditSpan `json:"edits"` } if err := json.Unmarshal(args, &obj); err != nil { @@ -307,7 +307,7 @@ func (t *multiEditTool) Execute(_ context.Context, args json.RawMessage) (Result } text = strings.Replace(text, e.Old, e.New, 1) } - if err := os.WriteFile(path, []byte(text), 0o644); err != nil { + if err := atomicWriteFile(path, []byte(text), 0o644); err != nil { return Result{Content: err.Error(), IsError: true}, nil } return Result{Content: fmt.Sprintf("applied %d edits to %s", len(obj.Edits), obj.Path)}, nil diff --git a/runtime/tools/exec.go b/runtime/tools/exec.go index 410d2ce..416ac83 100644 --- a/runtime/tools/exec.go +++ b/runtime/tools/exec.go @@ -11,6 +11,7 @@ import ( "strings" "syscall" "time" + "unicode/utf8" ) // ExecConfig 声明执行类工具的运行条件，由组合根显式提供； @@ -176,7 +177,12 @@ func formatProcessOutput(out []byte, timedOut bool, err error) string { if len(out) > execMaxOutput { half := execMaxOutput / 2 dropped := len(out) - execMaxOutput - text = string(out[:half]) + fmt.Sprintf("\n[... truncated %d bytes ...]\n", dropped) + string(out[len(out)-half:]) + head := CutUTF8(string(out), half) + tail := string(out[len(out)-half:]) + for len(tail) > 0 && !utf8.RuneStart(tail[0]) { + tail = tail[1:] + } + text = head + fmt.Sprintf("\n[... truncated %d bytes ...]\n", dropped) + tail } switch { case timedOut: diff --git a/runtime/tools/r2_regression_test.go b/runtime/tools/r2_regression_test.go new file mode 100644 index 0000000..01f70bb --- /dev/null +++ b/runtime/tools/r2_regression_test.go @@ -0,0 +1,50 @@ +package tools + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + "unicode/utf8" +) + +func TestR2WriteFileIsAtomic(t *testing.T) { + dir := t.TempDir() + reg := NewRegistry(dir) + for _, tool := range []Tool{&writeFile{root: dir}} { + reg.Register(tool) + } + res := reg.Execute(context.Background(), "write_file", []byte(`{"path":"note.txt","content":"完整"}`)) + if res.IsError { + t.Fatal(res.Content) + } + if _, err := os.Stat(filepath.Join(dir, "note.txt.tmp")); !os.IsNotExist(err) { + t.Fatal("temp file left behind") + } + b, err := os.ReadFile(filepath.Join(dir, "note.txt")) + if err != nil || string(b) != "完整" { + t.Fatalf("file: %q %v", b, err) + } +} + +func TestR2CutUTF8DoesNotSplitRunes(t *testing.T) { + s := strings.Repeat("界", 5) + cut := CutUTF8(s, 4) + if !utf8.ValidString(cut) { + t.Fatalf("invalid utf-8 %q", cut) + } + if cut != "界" { + t.Fatalf("cut %q", cut) + } +} + +func TestR2SchemaRejectsUnknownKeys(t *testing.T) { + dir := t.TempDir() + reg := NewRegistry(dir) + reg.Register(&writeFile{root: dir}) + res := reg.Execute(context.Background(), "write_file", []byte(`{"path":"a.txt","content":"x","extra":1}`)) + if !res.IsError || !strings.Contains(res.Content, "unknown argument") { + t.Fatalf("result: %+v", res) + } +} diff --git a/runtime/tools/registry.go b/runtime/tools/registry.go index ab4a3c3..44c382c 100644 --- a/runtime/tools/registry.go +++ b/runtime/tools/registry.go @@ -238,6 +238,13 @@ func validateArgs(schema, args json.RawMessage) error { return fmt.Errorf("argument %s must be %s", key, prop.Type) } } + if len(spec.Properties) > 0 { + for key := range obj { + if _, ok := spec.Properties[key]; !ok { + return fmt.Errorf("unknown argument %s", key) + } + } + } return nil } diff --git a/runtime/tools/utf8.go b/runtime/tools/utf8.go new file mode 100644 index 0000000..99233d3 --- /dev/null +++ b/runtime/tools/utf8.go @@ -0,0 +1,15 @@ +package tools + +import "unicode/utf8" + +// CutUTF8 按字节上限截断，并回退到完整 rune，避免切开 UTF-8。 +func CutUTF8(s string, maxBytes int) string { + if maxBytes <= 0 || len(s) <= maxBytes { + return s + } + end := maxBytes + for end > 0 && !utf8.ValidString(s[:end]) { + end-- + } + return s[:end] +}