Close the native P1 gaps the audit listed for this batch
Why: token budget let the current tool batch run, compaction could split a tool pair or race an append, writes were not crash-atomic, and several MCP and session edges only existed as comments. Evidence: GOWORK=off go test -count=1 -race ./runtime/ ./runtime/llm/ ./runtime/session/ ./runtime/tools/ ./runtime/mcp/ ./adapters/native/ ./adapters/intentgate/ passed, including the TestR2* cases for budget, replacement, cancel, compaction, SSE, UTF-8, schema, MCP retry, session terminate, poison reload, and registry reload. Impact: over-budget tool calls are skipped and paired; Stop.Replacement is appended to the session; the intent gate can follow Registry.Reload. Refs: N5, N6, N7, N8, N9, N11, N12, N15, N16, N17, N18, N19, N21, I9 Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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},
|
||||
)})
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
+57
-9
@@ -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":
|
||||
|
||||
@@ -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 {
|
||||
|
||||
+25
-1
@@ -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()
|
||||
|
||||
+25
-6
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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]
|
||||
}
|
||||
Reference in New Issue
Block a user