diff --git a/README.md b/README.md index c734c93..37d10d0 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ | 层 | 职责 | 依赖边界 | |---|---|---| | contract | 执行契约：Request/Result/Outcome/Error/Capability/Backend/Resolver | 不依赖内核、产品类型、数据库、配置，也不导入错误码库 | -| exec | 执行端口：Backend/Request/Result/ProcessEvent、步骤重试与失败分类、交付物物化、Stub、并发限流 | 面向调用方的执行端口 | +| exec | 执行端口：Backend/Request/Result/ProcessEvent、预算、步骤重试与失败分类、交付物物化、Stub、并发限流 | 面向调用方的执行端口 | | runtime | 多轮循环、工具调度、模型调用、执行会话 | 不导入 Core/Squad；可选使用日志库，nil 日志器不改变结果 | | adapters | PI Runner transport / typed-result-v1；native：自研内核 runtime → exec.Backend（规范后端名 agent-native，含会话毒化轮转——resume-unsafe 失败后自动换新会话） | 显式映射外部差异，不伪造缺失能力 | | conformance | 所有 Backend 共用的行为验收 | fake 验证契约，真实 adapter 验证能力 | @@ -23,9 +23,13 @@ |---|---| | contract | 四态执行契约 | | conformance | 行为验收套件与 fake | -| exec | Backend/Request/Result/ProcessEvent、步骤重试与失败分类、交付物物化、Stub、并发限流 | -| runtime | 多轮循环、llm、session、tools（含单测） | -| adapters/native | runtime 作为 exec.Backend，含会话毒化轮转 | +| exec | Backend/Request/Result/ProcessEvent（v0.2：Model/ToolAllowlist/Skills/Budget/SessionScope/ExecutionID/Env）、Budget、步骤重试与失败分类、交付物物化、Stub、并发限流 | +| runtime | 多轮循环（预算/取消贯穿、budget_exceeded 事件）、llm、session（compaction/元数据/崩溃安全）、tools（9 内置 + 策略层 Restrict） | +| runtime/mcp | MCP 客户端（stdio + Streamable HTTP）与 mcp____ 桥接（server 元数据不可信、上限收敛） | +| runtime/skills | SKILL.md 技能加载与 load_skill 渐进披露 | +| runtime/promptx | 系统提示结构化装配（环境事实 + 技能索引 + 交付物约束） | +| runtime/subagent | 委派端口（Request/Result/Factory）与 delegate_task 工具；嵌套执行由宿主注入 | +| adapters/native | runtime 作为 exec.Backend，含会话毒化轮转、compaction、技能/搜索/执行/MCP/委派装配（Options 显式注入，未配置不注册） | | adapters/pitransport | typed-result-v1 编解码、远程结果解析、Runner 搜索、鉴权头（含 Node 摘要交叉验证测试） | 本模块不依赖 Core 或 Squad。上层产品依赖本模块的 contract 与 exec。 diff --git a/adapters/native/backend.go b/adapters/native/backend.go index 9bccc40..7a7fcd6 100644 --- a/adapters/native/backend.go +++ b/adapters/native/backend.go @@ -15,7 +15,11 @@ import ( "agent-native/exec" "agent-native/runtime" "agent-native/runtime/llm" + "agent-native/runtime/mcp" + "agent-native/runtime/promptx" "agent-native/runtime/session" + "agent-native/runtime/skills" + "agent-native/runtime/subagent" "agent-native/runtime/tools" "logging-system/logsys" @@ -30,12 +34,38 @@ type Options struct { Timeout time.Duration // Stream 测试注入的流函数；缺省按 APIBase/APIKey 构造 llm.Client。 Stream llm.StreamFunc + // Search 可选：提供时注册 tavily_search / web_search（密钥由组合根注入）。 + Search tools.SearchProvider + // Exec 可选：提供时注册 run_command / run_python（解释器留空项不注册）。 + // nil = 不注册执行类工具（fail-closed，与 BashPermission 默认 deny 对齐）。 + Exec *tools.ExecConfig + // Fetch 可选：提供时注册 fetch_url；nil = 不注册（fail-closed）。 + Fetch *tools.FetchConfig + // MCP 可选：提供时把各 server 工具以 mcp____ 注册进工作区。 + // nil = 不注册（fail-closed：server 不可达时整个步骤失败，不静默跳过）。 + MCP *mcp.Manager + // SkillsRoots 可选：全局技能根目录；执行时自动并入工作区 .agent/skills/。 + SkillsRoots []string + // Compact 可选：转录超过阈值时把头部历史压成摘要（长会话上下文治理）。 + Compact *CompactConfig + // Subagent 可选：提供时注册 delegate_task（委派由宿主实现，内核只开端口）。 + Subagent subagent.Factory // Log 可选。nil 时丢弃日志。 Log *logsys.Logger // ExecutionID 写入日志，不参与执行控制。 ExecutionID string } +// CompactConfig 声明长会话压缩策略（优化性质：压缩失败不阻断执行，但留日志）。 +type CompactConfig struct { + // ThresholdChars 触发阈值（按字节数近似），默认 300_000。 + ThresholdChars int + // KeepLastN 压缩后保留的最近消息条数，默认 8。 + KeepLastN int + // Summarizer 可选注入摘要函数；缺省用后端 stream 做一次无工具调用。 + Summarizer func(ctx context.Context, messages []llm.Message) (string, error) +} + // Backend 以自研内核实现 exec.Backend。 type Backend struct { opts Options @@ -53,6 +83,13 @@ func New(opts Options) *Backend { // Name 返回规范后端名。 func (b *Backend) Name() string { return "agent-native" } +// Close 释放后端持有的长连接资源（MCP server 子进程等）。 +func (b *Backend) Close() { + if b.opts.MCP != nil { + b.opts.MCP.Close() + } +} + // Execute 运行一轮带工具的自研智能体循环，产出目标交付物。 func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) { if len(req.Targets) == 0 { @@ -68,7 +105,11 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, if b.opts.Stream == nil && (strings.TrimSpace(b.opts.APIBase) == "" || strings.TrimSpace(b.opts.APIKey) == "") { return &exec.Result{Success: false, ErrorReason: "agent-native backend requires api base/key (or injected stream)"}, nil } - if strings.TrimSpace(b.opts.Model) == "" { + model := strings.TrimSpace(b.opts.Model) + if strings.TrimSpace(req.Model) != "" { + model = strings.TrimSpace(req.Model) + } + if model == "" { return &exec.Result{Success: false, ErrorReason: "agent-native backend requires model"}, nil } @@ -88,9 +129,48 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, _ = os.WriteFile(filepath.Join(workDir, "TASK.md"), []byte(req.Prompt), 0o644) } + // 技能：全局根 + 工作区 .agent/skills；声明的技能必须存在（fail-closed）。 + roots := make([]string, 0, len(b.opts.SkillsRoots)+1) + roots = append(roots, b.opts.SkillsRoots...) + roots = append(roots, filepath.Join(workDir, ".agent", "skills")) + skillStore := skills.NewStore(roots...) + skillIndex := skillStore.List() + indexSet := make(map[string]bool, len(skillIndex)) + for _, s := range skillIndex { + indexSet[s.Name] = true + } + for _, name := range req.Skills { + if !indexSet[name] { + return &exec.Result{Success: false, ErrorReason: "declared skill not installed: " + name}, nil + } + } + var reg *tools.Registry if profile == exec.ToolProfileWorkspace { reg = tools.WorkspaceTools(workDir) + skills.RegisterLoader(reg, skillStore) + if b.opts.Search != nil { + tools.RegisterSearch(reg, b.opts.Search) + } + if b.opts.Fetch != nil { + tools.RegisterFetch(reg, *b.opts.Fetch) + } + if b.opts.Exec != nil { + execCfg := *b.opts.Exec + for k, v := range req.Env { + execCfg.Env = append(execCfg.Env, k+"="+v) + } + tools.RegisterExecTools(reg, workDir, execCfg) + } + if b.opts.MCP != nil { + if err := b.opts.MCP.RegisterInto(ctx, reg); err != nil { + return &exec.Result{Success: false, ErrorReason: err.Error()}, nil + } + } + subagent.Register(reg, b.opts.Subagent) + if err := reg.Restrict(req.ToolAllowlist); err != nil { + return &exec.Result{Success: false, ErrorReason: err.Error()}, nil + } } // 会话毒化轮转：上次失败属 resume-unsafe 时，把旧 transcript 归档， // 本次以全新会话开始（续跑毒化会话只会复现同样的失败）。 @@ -104,11 +184,33 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, client := &llm.Client{BaseURL: b.opts.APIBase, APIKey: b.opts.APIKey} stream = client.Stream } + _ = store.SaveMeta(map[string]string{ + "model": model, + "scope": req.SessionScope, + "exec_id": req.ExecutionID, + "updated": time.Now().UTC().Format(time.RFC3339), + }) + b.maybeCompact(ctx, store, stream, model) timeout := req.Timeout if timeout <= 0 { timeout = b.opts.Timeout } + var runBudget *runtime.Budget + if req.Budget != nil { + runBudget = &runtime.Budget{ + MaxPromptTokens: req.Budget.MaxPromptTokens, + MaxCompletionTokens: req.Budget.MaxCompletionTokens, + MaxToolCalls: req.Budget.MaxToolCalls, + Deadline: req.Budget.Deadline, + } + if !req.Budget.Deadline.IsZero() { + // Deadline 与 Timeout 先到者生效：把更短者作为 ctx 超时。 + if until := time.Until(req.Budget.Deadline); timeout <= 0 || until < timeout { + timeout = until + } + } + } runCtx := ctx cancel := func() {} if timeout > 0 { @@ -116,9 +218,16 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, } defer cancel() + execID := req.ExecutionID + if strings.TrimSpace(execID) == "" { + execID = b.opts.ExecutionID + } var logs strings.Builder + if strings.TrimSpace(execID) != "" || strings.TrimSpace(req.SessionScope) != "" { + logs.WriteString("execution_id=" + execID + " session=" + req.SessionScope + "\n") + } onProcess := req.OnEvent - runRes, err := b.run(runCtx, reg, store, stream, profile, req, &logs, onProcess) + runRes, err := b.run(runCtx, model, runBudget, reg, store, stream, profile, req, workDir, skillIndex, &logs, onProcess) if err != nil { return &exec.Result{Success: false, ErrorReason: err.Error(), Logs: logs.String()}, nil } @@ -173,19 +282,24 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, } // run 驱动 runtime 循环并桥接过程事件（语义对齐主仓 pinative）。 -func (b *Backend) run(ctx context.Context, reg *tools.Registry, store *session.Store, - stream llm.StreamFunc, profile string, req exec.Request, logs *strings.Builder, - onProcess func(exec.ProcessEvent)) (runtime.Result, error) { +func (b *Backend) run(ctx context.Context, model string, budget *runtime.Budget, reg *tools.Registry, store *session.Store, + stream llm.StreamFunc, profile string, req exec.Request, workDir string, skillIndex []skills.Skill, + logs *strings.Builder, onProcess func(exec.ProcessEvent)) (runtime.Result, error) { + execID := req.ExecutionID + if strings.TrimSpace(execID) == "" { + execID = b.opts.ExecutionID + } runRes, err := runtime.Run(ctx, runtime.Config{ - Model: b.opts.Model, - SystemPrompt: systemPrompt(req, profile), + Model: model, + SystemPrompt: b.systemPrompt(req, model, profile, workDir, skillIndex), Tools: reg, Stream: stream, Session: store, MaxTurns: b.opts.MaxTurns, Log: b.logger(), - ExecutionID: b.opts.ExecutionID, + ExecutionID: execID, + Budget: budget, OnEvent: func(ev runtime.Event) { logs.WriteString(ev.Type) if ev.Tool != "" { @@ -210,6 +324,74 @@ func (b *Backend) run(ctx context.Context, reg *tools.Registry, store *session.S return runRes, err } +// maybeCompact 在转录超过阈值时压缩头部历史。压缩是优化性质： +// 失败不阻断执行（按原文继续），但必须留日志，不静默。 +func (b *Backend) maybeCompact(ctx context.Context, store *session.Store, stream llm.StreamFunc, model string) { + cfg := b.opts.Compact + if cfg == nil { + return + } + threshold := int64(cfg.ThresholdChars) + if threshold <= 0 { + threshold = 300_000 + } + keep := cfg.KeepLastN + if keep <= 0 { + keep = 8 + } + if store.Stats().Bytes < threshold { + return + } + summarize := cfg.Summarizer + if summarize == nil { + summarize = defaultSummarizer(stream, model) + } + done, err := store.Compact(keep, func(messages []llm.Message) (string, error) { + return summarize(ctx, messages) + }) + if err != nil { + b.logger().Warn("session compaction failed; continuing with full transcript", "error", err.Error()) + return + } + if done { + b.logger().Info("session compacted", "keep_last", keep) + } +} + +// defaultSummarizer 用后端 stream 做一次无工具调用，把历史压成摘要。 +func defaultSummarizer(stream llm.StreamFunc, model string) func(ctx context.Context, messages []llm.Message) (string, error) { + return func(ctx context.Context, messages []llm.Message) (string, error) { + var b strings.Builder + for _, m := range messages { + b.WriteString(string(m.Role)) + b.WriteString(": ") + b.WriteString(m.Content) + b.WriteString("\n") + if b.Len() > 120_000 { + break + } + } + prompt := "把以下智能体对话历史压缩成摘要：保留全部关键事实、已做决定、文件变更与未决问题；中文；不超过 800 字；只输出摘要正文。\n\n" + b.String() + var out strings.Builder + for ev := range stream(ctx, model, []llm.Message{{Role: llm.RoleUser, Content: prompt}}, nil, llm.StreamOptions{}) { + switch ev.Kind { + case llm.KindTextDelta: + out.WriteString(ev.Text) + case llm.KindError: + return "", ev.Err + } + } + if strings.TrimSpace(out.String()) == "" { + return "", errString("summarizer produced empty output") + } + return out.String(), nil + } +} + +type errString string + +func (e errString) Error() string { return string(e) } + // resumeUnsafeMarkers 续跑不安全的关键词（会话已被污染：上下文溢出/ // 迭代上限/非法请求——重试必须换新会话而非续跑）。 var resumeUnsafeMarkers = []string{ @@ -256,20 +438,27 @@ func (b *Backend) logger() *logsys.Logger { return log.With("execution_id", b.opts.ExecutionID) } -// systemPrompt 组装系统提示：角色指令 + 交付物与工具档位约束。 -func systemPrompt(req exec.Request, profile string) string { - var b strings.Builder - if strings.TrimSpace(req.Instructions) != "" { - b.WriteString(req.Instructions) - b.WriteString("\n\n") +// systemPrompt 经 promptx 组装：角色指令 + 环境事实 + 技能索引 + 交付物约束。 +func (b *Backend) systemPrompt(req exec.Request, model, profile, workDir string, skillIndex []skills.Skill) string { + var index []promptx.IndexEntry + for _, s := range skillIndex { + index = append(index, promptx.IndexEntry{Name: s.Name, Description: s.Description}) } - b.WriteString("Deliverable files: ") - b.WriteString(strings.Join(req.Targets, ", ")) - b.WriteString(".") - if profile == exec.ToolProfileWorkspace { - b.WriteString(" Use workspace tools to write those files. Do not invent paths outside the workspace.") - } else { - b.WriteString(" Reply with the deliverable text only. Tools are disabled.") - } - return b.String() + now := time.Now() + weekday := map[string]string{ + "Sunday": "周日", "Monday": "周一", "Tuesday": "周二", "Wednesday": "周三", + "Thursday": "周四", "Friday": "周五", "Saturday": "周六", + }[now.Weekday().String()] + return promptx.Compose(promptx.Parts{ + Instructions: req.Instructions, + SkillsDeclared: req.Skills, + SkillsIndex: index, + Deliverables: req.Targets, + Env: promptx.EnvFacts{ + Date: now.Format("2006-01-02") + "（" + weekday + "）", + Model: model, + WorkspaceRoot: workDir, + Profile: profile, + }, + }) } diff --git a/adapters/native/compact_test.go b/adapters/native/compact_test.go new file mode 100644 index 0000000..9b3683d --- /dev/null +++ b/adapters/native/compact_test.go @@ -0,0 +1,69 @@ +package native + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + + "agent-native/exec" + "agent-native/runtime/llm" +) + +func TestBackendCompactBeforeRun(t *testing.T) { + dir := t.TempDir() + // 预置一份超阈值的长转录（系统消息 + 12 条历史）。 + transcript := filepath.Join(dir, ".agentx", "transcript.jsonl") + if err := os.MkdirAll(filepath.Dir(transcript), 0o755); err != nil { + t.Fatal(err) + } + var b strings.Builder + b.WriteString(`{"role":"system","content":"sys"}` + "\n") + for i := 0; i < 12; i++ { + b.WriteString(`{"role":"user","content":"` + strings.Repeat("x", 400) + `"}` + "\n") + } + if err := os.WriteFile(transcript, []byte(b.String()), 0o644); err != nil { + t.Fatal(err) + } + + var sawSummary bool + // 注入式捕获：验证 Run 收到的消息里有摘要标记。 + bk2 := New(Options{ + Model: "fake", + Stream: captureFirstStream(&sawSummary), + Compact: &CompactConfig{ThresholdChars: 1000, KeepLastN: 3, Summarizer: func(ctx context.Context, messages []llm.Message) (string, error) { + return "预置历史摘要", nil + }}, + }) + res, err := bk2.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "继续", Targets: []string{"output.md"}, + ToolProfile: exec.ToolProfileReasoning, + }) + if err != nil { + t.Fatal(err) + } + if res.Success && !sawSummary { + t.Fatal("compaction should have run before the loop") + } + data, _ := os.ReadFile(transcript) + if !strings.Contains(string(data), "[conversation summary]") { + t.Fatalf("transcript should contain the summary marker, got:\n%.200s", data) + } +} + +func captureFirstStream(saw *bool) llm.StreamFunc { + return func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + for _, m := range messages { + if strings.Contains(m.Content, "[conversation summary]") { + *saw = true + } + } + ch := make(chan llm.Event, 4) + go func() { + defer close(ch) + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "压缩后续写完成", StopReason: llm.StopEnd}} + }() + return ch + } +} diff --git a/adapters/native/exec_tools_test.go b/adapters/native/exec_tools_test.go new file mode 100644 index 0000000..865b634 --- /dev/null +++ b/adapters/native/exec_tools_test.go @@ -0,0 +1,342 @@ +package native + +import ( + "context" + "os" + osexec "os/exec" + "path/filepath" + "strings" + "testing" + + "agent-native/exec" + "agent-native/runtime/llm" + "agent-native/runtime/mcp" + "agent-native/runtime/subagent" + "agent-native/runtime/tools" +) + +func TestBackendExecToolsWiredWithOptionsAndRequestEnv(t *testing.T) { + bash, err := osexec.LookPath("bash") + if err != nil { + t.Fatalf("bash not available: %v", err) + } + dir := t.TempDir() + b := New(Options{ + Model: "fake", + Stream: script( + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "run_command", + Arguments: `{"command":"echo role=$ROLE_FLAG step=$STEP_FLAG > exec_out.txt"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c2", Name: "write_file", + Arguments: `{"path":"output.md","content":"done"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}, + ), + Exec: &tools.ExecConfig{Bash: bash, Env: []string{"PATH=/usr/bin:/bin", "ROLE_FLAG=backend"}}, + }) + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "run", Targets: []string{"output.md"}, + Env: map[string]string{"STEP_FLAG": "request"}, + }) + if err != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v logs=%s", err, res, res.Logs) + } + data, err := os.ReadFile(filepath.Join(dir, "exec_out.txt")) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(data), "role=backend") || !strings.Contains(string(data), "step=request") { + t.Fatalf("env merge wrong: %q", data) + } +} + +func TestBackendExecToolsAbsentWithoutOptions(t *testing.T) { + stream, c := captureStream() + b := New(Options{Model: "fake", Stream: stream}) + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"}, + ToolProfile: exec.ToolProfileReasoning, + }) + if err != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v", err, res) + } + // reasoning 档本无工具；改用 workspace 档断言 decls 不含执行工具。 + stream2, c2 := captureStream() + b2 := New(Options{Model: "fake", Stream: stream2}) + _, _ = b2.Execute(context.Background(), exec.Request{ + WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"}, + ToolProfile: exec.ToolProfileWorkspace, + }) + for _, names := range c2.toolDecls { + for _, n := range names { + if n == "run_command" || n == "run_python" { + t.Fatalf("exec tools must not register without Options.Exec, got %v", names) + } + } + } + _ = c +} + +func TestBackendAllowlistCanDisableExecTools(t *testing.T) { + bash, _ := osexec.LookPath("bash") + stream, c := captureStream() + b := New(Options{ + Model: "fake", Stream: stream, + Exec: &tools.ExecConfig{Bash: bash, Env: []string{"PATH=/usr/bin:/bin"}}, + }) + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"}, + ToolAllowlist: []string{"read_file", "write_file", "edit_file", "list_dir", "grep", "glob", "multi_edit"}, + }) + if err != nil { + t.Fatal(err) + } + if res.Success { + t.Fatalf("write-only allowlist without write? (unexpected) res=%+v", res) + } + for _, names := range c.toolDecls { + for _, n := range names { + if n == "run_command" { + t.Fatalf("allowlist must be able to strip run_command, got %v", names) + } + } + } +} + +// fakeSearchProvider 与 tools 包测试同构的最小搜索替身。 +type fakeSearchProvider struct{ query string } + +func (f *fakeSearchProvider) Search(_ context.Context, req tools.SearchRequest) (tools.SearchResponse, error) { + f.query = req.Query + return tools.SearchResponse{Answer: "ans", Results: []tools.SearchResult{{Title: "T", Content: "C", Score: 0.9}}}, nil +} + +func TestBackendSearchToolsWiredViaOptions(t *testing.T) { + provider := &fakeSearchProvider{} + stream, c := captureStream( + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "web_search", Arguments: `{"query":"golang generics"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c2", Name: "write_file", + Arguments: `{"path":"output.md","content":"done"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}, + ) + b := New(Options{Model: "fake", Stream: stream, Search: provider}) + dir := t.TempDir() + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "search", Targets: []string{"output.md"}, + }) + if err != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v logs=%s", err, res, res.Logs) + } + if provider.query != "golang generics" { + t.Fatalf("provider should receive the query, got %q", provider.query) + } + // 搜索工具确实出现在 decls 里。 + found := false + for _, names := range c.toolDecls { + for _, n := range names { + if n == "web_search" || n == "tavily_search" { + found = true + } + } + } + if !found { + t.Fatal("search tools must be declared when Options.Search is set") + } +} + +func TestBackendSearchToolsAbsentWithoutOptions(t *testing.T) { + stream, c := captureStream() + b := New(Options{Model: "fake", Stream: stream}) + _, _ = b.Execute(context.Background(), exec.Request{ + WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"}, + ToolProfile: exec.ToolProfileWorkspace, + }) + for _, names := range c.toolDecls { + for _, n := range names { + if n == "web_search" || n == "tavily_search" { + t.Fatalf("search tools must not register without Options.Search, got %v", names) + } + } + } +} + +func TestBackendMCPToolsWiredViaOptions(t *testing.T) { + exe, err := osexec.LookPath(os.Args[0]) + if err != nil { + t.Fatal(err) + } + // os.Args[0] 在 go test 下指向测试二进制——用其自身作为 MCP server。 + mgr, err := mcp.NewManager([]mcp.ServerConfig{{ + Name: "fake", Command: exe, + Env: append(os.Environ(), "AGENT_NATIVE_TEST_MCP_STDIO=1"), + }}) + if err != nil { + t.Fatal(err) + } + defer mgr.Close() + stream, c := captureStream( + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "mcp__fake__echo", Arguments: `{"text":"via-backend"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c2", Name: "write_file", + Arguments: `{"path":"output.md","content":"done"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}, + ) + b := New(Options{Model: "fake", Stream: stream, MCP: mgr}) + dir := t.TempDir() + res, err2 := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "call mcp", Targets: []string{"output.md"}, + }) + if err2 != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v logs=%s", err2, res, res.Logs) + } + found := false + for _, names := range c.toolDecls { + for _, n := range names { + if n == "mcp__fake__echo" { + found = true + } + } + } + if !found { + t.Fatal("mcp bridge tools must be declared to the model") + } +} + +func TestBackendSkillsIndexAndLoader(t *testing.T) { + skillsRoot := t.TempDir() + skillDir := filepath.Join(skillsRoot, "docx-author") + if err := os.MkdirAll(skillDir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte("---\ndescription: 生成文档交付物\n---\n\n配方正文"), 0o644); err != nil { + t.Fatal(err) + } + stream, c := captureStream( + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "load_skill", Arguments: `{"name":"docx-author"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c2", Name: "write_file", Arguments: `{"path":"output.md","content":"done"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}, + ) + b := New(Options{Model: "fake", Stream: stream, SkillsRoots: []string{skillsRoot}}) + dir := t.TempDir() + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "写", Targets: []string{"output.md"}, Skills: []string{"docx-author"}, + }) + if err != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v logs=%s", err, res, res.Logs) + } + if !strings.Contains(c.system, "- docx-author — 生成文档交付物") { + t.Fatalf("skill index missing in system prompt:\n%s", c.system) + } + if !strings.Contains(c.system, "Declared for this task") || !strings.Contains(c.system, "docx-author") { + t.Fatalf("declared skills missing:\n%s", c.system) + } + sawLoader := false + for _, names := range c.toolDecls { + for _, n := range names { + if n == "load_skill" { + sawLoader = true + } + } + } + if !sawLoader { + t.Fatal("load_skill must be declared when skills exist") + } +} + +func TestBackendDeclaredSkillMissingFailsClosed(t *testing.T) { + b := New(Options{Model: "fake", Stream: script(), SkillsRoots: []string{t.TempDir()}}) + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"}, + Skills: []string{"ghost-skill"}, + }) + if err != nil { + t.Fatal(err) + } + if res.Success || !strings.Contains(res.ErrorReason, "declared skill not installed: ghost-skill") { + t.Fatalf("missing declared skill must fail-closed, got %+v", res) + } +} + +func TestBackendWorkspaceSkillsRootAutoMerged(t *testing.T) { + stream, c := captureStream( + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "write_file", Arguments: `{"path":"output.md","content":"done"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}, + ) + b := New(Options{Model: "fake", Stream: stream}) + dir := t.TempDir() + wsSkill := filepath.Join(dir, ".agent", "skills", "ws-skill") + if err := os.MkdirAll(wsSkill, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(wsSkill, "SKILL.md"), []byte("---\ndescription: 工作区技能\n---\n正文"), 0o644); err != nil { + t.Fatal(err) + } + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "x", Targets: []string{"output.md"}, + }) + if err != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v", err, res) + } + if !strings.Contains(c.system, "ws-skill") { + t.Fatalf("workspace skills root must be merged into index:\n%s", c.system) + } +} + +func TestBackendSubagentWiredViaOptions(t *testing.T) { + var delegated *subagent.Request + stream, c := captureStream( + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "delegate_task", + Arguments: `{"goal":"查证数据","write_scope":"research/","done_criteria":"三来源"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c2", Name: "write_file", + Arguments: `{"path":"output.md","content":"done"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}, + ) + b := New(Options{ + Model: "fake", + Stream: stream, + Subagent: func(ctx context.Context, req subagent.Request) (subagent.Result, error) { + delegated = &req + return subagent.Result{Success: true, Text: "结论已查明"}, nil + }, + }) + dir := t.TempDir() + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "分工", Targets: []string{"output.md"}, + }) + if err != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v logs=%s", err, res, res.Logs) + } + if delegated == nil || delegated.Goal != "查证数据" { + t.Fatalf("factory not invoked with fields: %+v", delegated) + } + sawDelegate := false + for _, names := range c.toolDecls { + for _, n := range names { + if n == "delegate_task" { + sawDelegate = true + } + } + } + if !sawDelegate { + t.Fatal("delegate_task must be declared when Subagent factory is set") + } +} diff --git a/adapters/native/mcpmain_test.go b/adapters/native/mcpmain_test.go new file mode 100644 index 0000000..9133aee --- /dev/null +++ b/adapters/native/mcpmain_test.go @@ -0,0 +1,17 @@ +package native + +import ( + "os" + "testing" + + "agent-native/runtime/mcp" +) + +// TestMain：MCP fake server 的 re-exec 分支（详见 mcp.RunFakeServer）。 +func TestMain(m *testing.M) { + if os.Getenv(mcp.FakeServerEnv) == "1" { + mcp.RunFakeServer() + os.Exit(0) + } + os.Exit(m.Run()) +} diff --git a/adapters/native/request_v2_test.go b/adapters/native/request_v2_test.go new file mode 100644 index 0000000..57e3192 --- /dev/null +++ b/adapters/native/request_v2_test.go @@ -0,0 +1,183 @@ +package native + +import ( + "context" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "agent-native/exec" + "agent-native/runtime/llm" +) + +// captureStream 捕获每次调用的 model / 首条 system 消息，并按序回放回合。 +func captureStream(turns ...llm.Message) (llm.StreamFunc, *capture) { + c := &capture{} + var mu sync.Mutex + i := 0 + return func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 4) + go func() { + defer close(ch) + mu.Lock() + defer mu.Unlock() + c.models = append(c.models, model) + c.toolDecls = append(c.toolDecls, toolNames(tools)) + if len(messages) > 0 && messages[0].Role == llm.RoleSystem { + c.system = messages[0].Content + } + if i >= len(turns) { + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}} + return + } + msg := turns[i] + i++ + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: msg} + }() + return ch + }, c +} + +type capture struct { + mu sync.Mutex + models []string + toolDecls [][]string + system string +} + +func toolNames(decls []llm.ToolDecl) []string { + var out []string + for _, d := range decls { + out = append(out, d.Name) + } + return out +} + +func TestBackendRequestModelOverride(t *testing.T) { + stream, c := captureStream() + b := New(Options{Model: "default-model", Stream: stream}) + dir := t.TempDir() + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "hi", Targets: []string{"output.md"}, Model: "override-model", + ToolProfile: exec.ToolProfileReasoning, + }) + if err != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v", err, res) + } + if len(c.models) == 0 || c.models[0] != "override-model" { + t.Fatalf("want override-model, got %v", c.models) + } +} + +func TestBackendRequestToolAllowlist(t *testing.T) { + stream, c := captureStream( + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "write_file", Arguments: `{"path":"output.md","content":"blocked"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c2", Name: "write_file", Arguments: `{"path":"output2.md","content":"still blocked"}`, + }}}, + llm.Message{Role: llm.RoleAssistant, Content: "gave up", StopReason: llm.StopEnd}, + ) + b := New(Options{Model: "fake", Stream: stream}) + dir := t.TempDir() + // 目标 output2.md 不在 allowlist 内能写吗？——write_file 被整体收窄， + // 模型两次尝试都会得到 unknown tool，最终无交付物 → 失败。 + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "写文件", Targets: []string{"output2.md"}, + ToolAllowlist: []string{"read_file", "list_dir", "grep"}, + }) + if err != nil { + t.Fatal(err) + } + if res.Success { + t.Fatalf("write_file restricted out; deliverable must fail, got %+v", res) + } + if len(c.toolDecls) == 0 { + t.Fatal("expected tool decls") + } + for _, names := range c.toolDecls { + for _, n := range names { + if n == "write_file" || n == "edit_file" { + t.Fatalf("write/edit must be restricted out of decls, got %v", names) + } + } + } +} + +func TestBackendRequestUnknownAllowlistFailsClosed(t *testing.T) { + b := New(Options{Model: "fake", Stream: script()}) + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"}, + ToolAllowlist: []string{"read_file", "no_such_tool"}, + }) + if err != nil { + t.Fatal(err) + } + if res.Success || !strings.Contains(res.ErrorReason, "no_such_tool") { + t.Fatalf("want fail-closed unknown allowlist name, got %+v", res) + } +} + +func TestBackendRequestBudgetDeadlineBindsContext(t *testing.T) { + deadlineHit := make(chan struct{}) + var once sync.Once + stream := func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 4) + go func() { + defer close(ch) + select { + case <-ctx.Done(): + once.Do(func() { close(deadlineHit) }) + ch <- llm.Event{Kind: llm.KindError, Err: ctx.Err(), Message: llm.Message{StopReason: llm.StopAborted}} + case <-time.After(3 * time.Second): + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}} + } + }() + return ch + } + b := New(Options{Model: "fake", Stream: stream}) + start := time.Now() + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"}, + Timeout: 30 * time.Second, // 长超时，预算截止应更先生效 + Budget: &exec.Budget{Deadline: time.Now().Add(150 * time.Millisecond)}, + }) + if err != nil { + t.Fatal(err) + } + select { + case <-deadlineHit: + case <-time.After(2 * time.Second): + t.Fatal("budget deadline should bind the run context") + } + if res.Success { + t.Fatalf("deadline run must not succeed, got %+v", res) + } + if elapsed := time.Since(start); elapsed > 5*time.Second { + t.Fatalf("deadline should fire early, took %s", elapsed) + } +} + +func TestBackendRequestExecutionIDInLogs(t *testing.T) { + stream, _ := captureStream() + b := New(Options{Model: "fake", Stream: stream}) + dir := t.TempDir() + res, err := b.Execute(context.Background(), exec.Request{ + WorkDir: dir, Prompt: "hi", Targets: []string{"output.md"}, + Instructions: "你是角色", + ExecutionID: "exec-42", SessionScope: "run-9/step-3", + ToolProfile: exec.ToolProfileReasoning, + }) + if err != nil || !res.Success { + t.Fatalf("execute: err=%v res=%+v", err, res) + } + logBytes, _ := os.ReadFile(filepath.Join(dir, "execution.log")) + logs := string(logBytes) + if !strings.Contains(logs, "execution_id=exec-42") || !strings.Contains(logs, "run-9/step-3") { + t.Fatalf("execution log should carry ids, got:\n%s", logs) + } +} diff --git a/docs/parity-matrix.md b/docs/parity-matrix.md index ec9053e..280c510 100644 --- a/docs/parity-matrix.md +++ b/docs/parity-matrix.md @@ -66,3 +66,26 @@ baseline_fixed: true - 本矩阵证据均指向仓库内文件（上文路径），不含运行时截图； - Batch1 完成后在本文件追加"实测证据"列； - 上游仅引固定 0.85.1 的 README/docs 结论，不引 main。 + +--- + +## R15 扩面批次（full-agent-r1，2026-10-01）：从「wrapper 对标」到「完全智能体」 + +用户定调（2026-10-01）：Agent-Native 补全 pi-Agent 全量工具调用能力（网络搜索、代码编写、 +系统工具、外部工具调用、技能、提示词），并与 Agent-Squad 融合。扩面域**不做 pi 对标** +（上游 wrapper 未暴露这些能力），改以「标准符合 + conformance 负例 + 产品验收」为准—— +对标口径条款对本批新增域不适用，原有 #1–#16 行继续有效。 + +| 域 | 内容 | 出口证据 | +|---|---|---| +| 执行基座 | exec.Request v0.2（Model/ToolAllowlist/Skills/Budget/SessionScope/ExecutionID/Env）+ Registry.Restrict 策略层（精确名 + 命名空间通配）+ 预算贯穿（token/工具数/截止，budget_exceeded 永久失败）+ 取消传播到工具 | loop_budget / registry_policy / request_v2 / step_retry_budget 测试；关闭 #4/#10/#16 | +| 代码工具面 | glob（** 跨目录）、grep 正则/include/ignore_case、read 行范围分页、multi_edit 原子多处替换 | code_tools_test.go | +| 系统工具 | run_command / run_python：解释器白名单、env 白名单（宿主密钥不透传）、进程组 SIGKILL（超时/取消）、64KB 头尾截断、命令黑名单、脚本落盘 .agent/scripts/；Options.Exec 未配置不注册 | exec_tools_test.go（含杀组时限断言） | +| 网络面 | tavily_search/web_search 经 Options.Search 接线（关闭 #5 尾环）；fetch_url：SSRF 硬边界（私网/回环/链路本地/元数据端点逐跳校验）、域名策略、2MB/20s、HTML 剥离 | fetch_test.go + 适配器接线测试 | +| MCP | stdio + Streamable HTTP 客户端、initialize/tools/list(paged)/tools/call、Mcp-Session-Id、mcp____ 桥接、64 工具/1KB 描述/64KB 结果上限、不可达 fail-closed、真实子进程测试（re-exec fake server） | runtime/mcp/mcp_test.go | +| 技能/提示词 | SKILL.md 加载（目录名权威、frontmatter 描述）、load_skill 渐进披露、声明技能 fail-closed 校验、promptx 装配器（环境/技能索引/交付物） | skills_test / promptx_test / 适配器测试 | +| 会话 | compaction（阈值触发、原子改写、失败留痕不阻断）、元数据 sidecar、崩溃半行容错 | compact_test.go；对齐 #11/#12 | +| 委派 | runtime/subagent 端口 + delegate_task（信息集=委派条款）；嵌套执行由宿主注入 | subagent_test.go | + +安全声明（不冒充已解决）：run_command **不是 OS 沙箱**（宿主负责隔离）；技能正文与 MCP +server 描述是提示注入面（只信受信根 + 长度上限）；fetch 的 DNS 解析存在 rebinding 窗口。 diff --git a/exec/backend.go b/exec/backend.go index 2ebbcc9..3c71e4c 100644 --- a/exec/backend.go +++ b/exec/backend.go @@ -17,6 +17,34 @@ type Request struct { // 实现方必须保证：回调在 Execute 返回前完成（同步调用或等待收尾），无泄漏 goroutine； // 回调内不得阻塞（消费方自行节流/异步入队）。 OnEvent func(ev ProcessEvent) + + // ---- v0.2 能力字段（全部可选；零值 = 既有行为，向后兼容） ---- + + // Model 覆盖本请求使用的模型（后端构造时的默认模型仍兜底）。 + Model string `json:"model,omitempty"` + // ToolAllowlist 收窄可见工具集；出现后端未注册的工具名时 fail-closed。 + ToolAllowlist []string `json:"tool_allowlist,omitempty"` + // Skills 声明本次执行需要加载的技能（内核在系统提示中列出并由 load_skill 按需展开）。 + Skills []string `json:"skills,omitempty"` + // Budget 单次步骤执行的共享预算；步骤重试应传入剩余额度，不得各自另起一份。 + Budget *Budget `json:"budget,omitempty"` + // SessionScope / ExecutionID 只用于日志与过程事件贯穿，不参与控制流。 + SessionScope string `json:"session_scope,omitempty"` + ExecutionID string `json:"execution_id,omitempty"` + // Env 是组合根批准后注入执行类工具（run_command 等）的白名单环境变量。 + // 后端不得把自身持有的密钥并入其中。 + Env map[string]string `json:"env,omitempty"` +} + +// Budget 是一次步骤执行的共享预算（agent-execution 标准：重试层各有职责， +// 共享总预算与截止时间，避免传输、步骤与主循环次数相乘）。 +// 全部字段可选；零值字段不设限。 +type Budget struct { + MaxPromptTokens int `json:"max_prompt_tokens,omitempty"` + MaxCompletionTokens int `json:"max_completion_tokens,omitempty"` + MaxToolCalls int `json:"max_tool_calls,omitempty"` + // Deadline 绝对截止时刻；与 Request.Timeout 先到者生效。 + Deadline time.Time `json:"deadline,omitempty"` } // ProcessEvent 对 agentx 事件与远程 runner NDJSON 过程帧的统一投影， diff --git a/exec/step_retry.go b/exec/step_retry.go index 212c800..eb6b05d 100644 --- a/exec/step_retry.go +++ b/exec/step_retry.go @@ -160,6 +160,7 @@ func combinedFailureText(sig StepFailureSignals) string { func isPermanentReason(lower string) bool { return strings.Contains(lower, "prompt_empty") || strings.Contains(lower, "unauthorized") || + strings.Contains(lower, "budget_exceeded") || strings.Contains(lower, "治理强制终止") || strings.Contains(lower, "治理拦截") } diff --git a/exec/step_retry_budget_test.go b/exec/step_retry_budget_test.go new file mode 100644 index 0000000..1858bcf --- /dev/null +++ b/exec/step_retry_budget_test.go @@ -0,0 +1,37 @@ +package exec + +import ( + "errors" + "strings" + "testing" + "time" +) + +func TestAssessStepFailureBudgetExceededIsPermanent(t *testing.T) { + ass := AssessStepFailure(StepFailureSignals{ + Err: errors.New("budget_exceeded: tool calls 12 > max 10"), + }) + if ass.Class != FailurePermanent { + t.Fatalf("budget exhaustion must classify permanent (no auto retry), got %s", ass.Class) + } + if !strings.Contains(ass.Reason, "permanent") { + t.Fatalf("want permanent reason, got %s", ass.Reason) + } +} + +func TestAssessStepFailureDeadlinePastIsPermanent(t *testing.T) { + ass := AssessStepFailure(StepFailureSignals{ + Result: &Result{ErrorReason: "budget_exceeded: deadline 2026-10-01T00:00:00Z passed"}, + }) + if ass.Class != FailurePermanent { + t.Fatalf("deadline exhaustion must classify permanent, got %s", ass.Class) + } +} + +func TestBudgetZeroValueUnlimited(t *testing.T) { + var b Budget + if b.MaxPromptTokens != 0 || b.MaxCompletionTokens != 0 || b.MaxToolCalls != 0 || !b.Deadline.IsZero() { + t.Fatal("zero budget must express unlimited") + } + _ = time.Now +} diff --git a/runtime/loop.go b/runtime/loop.go index 73291f7..3021db8 100644 --- a/runtime/loop.go +++ b/runtime/loop.go @@ -3,8 +3,10 @@ package runtime import ( "context" "encoding/json" + "fmt" "strings" "sync" + "time" "agent-native/runtime/llm" "agent-native/runtime/session" @@ -45,6 +47,17 @@ type Config struct { Log *logsys.Logger // ExecutionID 只用于日志字段，不参与控制流。 ExecutionID string + // Budget 可选。任一上限命中即停止并返回 budget_exceeded 错误（StopError）； + // 零值字段不设限。截止时刻与外层 ctx 先到者生效。 + Budget *Budget +} + +// Budget 是循环层的执行预算；与 exec.Budget 字段同构，由适配层换算。 +type Budget struct { + MaxPromptTokens int + MaxCompletionTokens int + MaxToolCalls int + Deadline time.Time } // Result is the final assistant text and the transcript produced by one run. @@ -91,11 +104,15 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { var last llm.Message var usage llm.Usage + toolCalls := 0 for turn := 0; turn < cfg.MaxTurns; turn++ { if ctx.Err() != nil { cfg.logger().Error("runtime turn aborted", ctx.Err()) return Result{Text: last.Content, Messages: messages, StopReason: llm.StopAborted, Err: ctx.Err()}, nil } + if reason := budgetExceeded(cfg.Budget, usage, toolCalls, time.Now()); reason != "" { + return budgetStop(cfg, last, messages, usage, reason), nil + } emit(cfg, Event{Type: "turn_start"}) end, errEv := consume(ctx, cfg, messages) if errEv != nil { @@ -113,6 +130,11 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { emit(cfg, Event{Type: "message_end", Message: last}) if last.StopReason == llm.StopToolCalls && len(last.ToolCalls) > 0 { + 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) + return budgetStop(cfg, last, messages, usage, reason), nil + } + toolCalls += len(last.ToolCalls) results := executeTools(ctx, cfg, last) for _, res := range results { if err := persist(cfg, res); err != nil { @@ -280,6 +302,35 @@ func emit(cfg Config, ev Event) { } } +// budgetExceeded 返回命中的预算说明；空串表示预算仍充足（或未设预算）。 +// 只在「即将消耗更多预算」的节点调用（轮次顶端 / 工具批执行前）。 +func budgetExceeded(b *Budget, usage llm.Usage, toolCalls int, now time.Time) string { + if b == nil { + return "" + } + if !b.Deadline.IsZero() && now.After(b.Deadline) { + return fmt.Sprintf("budget_exceeded: deadline %s passed", b.Deadline.Format(time.RFC3339)) + } + if b.MaxPromptTokens > 0 && usage.PromptTokens > b.MaxPromptTokens { + return fmt.Sprintf("budget_exceeded: prompt tokens %d > max %d", usage.PromptTokens, b.MaxPromptTokens) + } + if b.MaxCompletionTokens > 0 && usage.CompletionTokens > b.MaxCompletionTokens { + return fmt.Sprintf("budget_exceeded: completion tokens %d > max %d", usage.CompletionTokens, b.MaxCompletionTokens) + } + if b.MaxToolCalls > 0 && toolCalls > b.MaxToolCalls { + return fmt.Sprintf("budget_exceeded: tool calls %d > max %d", toolCalls, b.MaxToolCalls) + } + return "" +} + +// budgetStop 以预算超限收尾：发事件、留日志、按 StopError 返回（不可重试）。 +func budgetStop(cfg Config, last llm.Message, messages []llm.Message, usage llm.Usage, reason string) Result { + emit(cfg, Event{Type: "budget_exceeded", Text: reason}) + err := errString(reason) + cfg.logger().Error("runtime budget exceeded", err) + return Result{Text: last.Content, Messages: messages, StopReason: llm.StopError, Usage: usage, Err: err} +} + type errString string func (e errString) Error() string { return string(e) } diff --git a/runtime/loop_budget_test.go b/runtime/loop_budget_test.go new file mode 100644 index 0000000..46a800e --- /dev/null +++ b/runtime/loop_budget_test.go @@ -0,0 +1,194 @@ +package runtime + +import ( + "context" + "encoding/json" + "strings" + "testing" + "time" + + "agent-native/runtime/llm" + "agent-native/runtime/tools" +) + +// usageScript 产出带 usage 的固定回合，用于预算测试。 +func usageScript(usage llm.Usage, turns ...llm.Message) llm.StreamFunc { + var i int + return func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 4) + go func() { + defer close(ch) + if i >= len(turns) { + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}, Usage: usage} + return + } + msg := turns[i] + i++ + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: msg, Usage: usage} + }() + return ch + } +} + +func TestLoopBudgetToolCallCap(t *testing.T) { + dir := t.TempDir() + reg := tools.WorkspaceTools(dir) + sawBudgetEvent := false + res, err := Run(context.Background(), Config{ + Model: "fake", + Tools: reg, + Budget: &Budget{MaxToolCalls: 1}, + OnEvent: func(ev Event) { + if ev.Type == "budget_exceeded" { + sawBudgetEvent = true + } + }, + Stream: usageScript(llm.Usage{}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "list_dir", Arguments: `{}`, + }}}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{ + {ID: "c2", Name: "list_dir", Arguments: `{}`}, + {ID: "c3", Name: "list_dir", Arguments: `{}`}, + }}, + llm.Message{Role: llm.RoleAssistant, Content: "never", StopReason: llm.StopEnd}, + ), + }, llm.Message{Role: llm.RoleUser, Content: "loop tools"}) + if err != nil { + t.Fatal(err) + } + if res.StopReason != llm.StopError || res.Err == nil || !strings.Contains(res.Err.Error(), "budget_exceeded") { + t.Fatalf("want budget_exceeded stop, got stop=%s err=%v", res.StopReason, res.Err) + } + if !strings.Contains(res.Err.Error(), "tool calls") { + t.Fatalf("want tool-call budget reason, got %v", res.Err) + } + if !sawBudgetEvent { + t.Fatal("expected budget_exceeded event") + } + // 第二批工具调用不该被执行（预算在执行前拦截）：transcript 只含一次 tool 结果。 + toolResults := 0 + for _, m := range res.Messages { + if m.Role == llm.RoleTool { + toolResults++ + } + } + if toolResults != 1 { + t.Fatalf("expected exactly 1 tool result, got %d", toolResults) + } +} + +func TestLoopBudgetPromptTokens(t *testing.T) { + res, err := Run(context.Background(), Config{ + Model: "fake", + Budget: &Budget{MaxPromptTokens: 10}, + Stream: usageScript(llm.Usage{PromptTokens: 500, TotalTokens: 600}, + llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ + ID: "c1", Name: "list_dir", Arguments: `{}`, + }}}, + llm.Message{Role: llm.RoleAssistant, Content: "never", StopReason: llm.StopEnd}, + ), + }, llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + if res.Err == nil || !strings.Contains(res.Err.Error(), "prompt tokens") { + t.Fatalf("want prompt-token budget stop, got err=%v", res.Err) + } + if res.Usage.PromptTokens != 500 { + t.Fatalf("usage should be reported on budget stop, got %d", res.Usage.PromptTokens) + } +} + +func TestLoopBudgetDeadline(t *testing.T) { + res, err := Run(context.Background(), Config{ + Model: "fake", + Budget: &Budget{Deadline: time.Now().Add(-time.Second)}, + Stream: usageScript(llm.Usage{}, llm.Message{Role: llm.RoleAssistant, Content: "never", StopReason: llm.StopEnd}), + }, llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + if res.Err == nil || !strings.Contains(res.Err.Error(), "deadline") { + t.Fatalf("want deadline budget stop, got err=%v", res.Err) + } +} + +func TestLoopBudgetNotHitCompletesNormally(t *testing.T) { + res, err := Run(context.Background(), Config{ + Model: "fake", + Budget: &Budget{MaxPromptTokens: 10_000, MaxToolCalls: 10}, + Stream: usageScript(llm.Usage{PromptTokens: 50}, + llm.Message{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd}), + }, llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + if res.Err != nil || res.Text != "ok" { + t.Fatalf("clean run should finish, got err=%v text=%q", res.Err, res.Text) + } +} + +// blockingTool 阻塞到 ctx 取消，验证取消传播到工具执行层。 +type blockingTool struct{ canceled chan struct{} } + +func (b *blockingTool) Name() string { return "block_forever" } +func (b *blockingTool) Description() string { return "blocks until context cancel" } +func (b *blockingTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","properties":{}}`) +} + +func (b *blockingTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) { + <-ctx.Done() + close(b.canceled) + return tools.Result{Content: "canceled", IsError: true}, ctx.Err() +} + +func TestLoopCancelPropagatesToTool(t *testing.T) { + bt := &blockingTool{canceled: make(chan struct{})} + reg := tools.NewRegistry(t.TempDir()) + reg.Register(bt) + ctx, cancel := context.WithCancel(context.Background()) + + firstCall := true + stream := func(ctx context.Context, model string, messages []llm.Message, decls []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event { + ch := make(chan llm.Event, 4) + go func() { + defer close(ch) + if firstCall { + firstCall = false + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{ + Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, + ToolCalls: []llm.ToolCall{{ID: "c1", Name: "block_forever", Arguments: `{}`}}, + }} + return + } + // 工具返回后的下一轮：等取消或超时兜底。 + select { + case <-ctx.Done(): + ch <- llm.Event{Kind: llm.KindError, Err: ctx.Err(), Message: llm.Message{StopReason: llm.StopAborted}} + case <-time.After(3 * time.Second): + ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}} + } + }() + return ch + } + + go func() { + time.Sleep(100 * time.Millisecond) + cancel() + }() + res, err := Run(ctx, Config{Model: "fake", Tools: reg, Stream: stream}, + llm.Message{Role: llm.RoleUser, Content: "block"}) + if err != nil { + t.Fatal(err) + } + select { + case <-bt.canceled: + case <-time.After(2 * time.Second): + t.Fatal("tool context was not canceled") + } + if res.StopReason != llm.StopAborted && res.Err == nil { + t.Fatalf("expected aborted stop, got %+v", res) + } +} diff --git a/runtime/mcp/fakeserver.go b/runtime/mcp/fakeserver.go new file mode 100644 index 0000000..1d6c81f --- /dev/null +++ b/runtime/mcp/fakeserver.go @@ -0,0 +1,94 @@ +package mcp + +// 本文件是跨测试二进制复用的最小 MCP fake server（stdio JSON-RPC）。 +// 测试以「自身可执行文件 + AGENT_NATIVE_TEST_MCP_STDIO=1」re-exec 方式拉起， +// 供真实传输链路测试使用；非测试代码不得调用。 + +import ( + "bufio" + "encoding/json" + "os" +) + +// RunFakeServer 在 stdin/stdout 上运行最小 MCP server（echo / boom / 未知方法）， +// 直到 stdin 关闭。测试二进制在 TestMain 的 re-exec 分支调用。 +func RunFakeServer() { + scanner := bufio.NewScanner(os.Stdin) + scanner.Buffer(make([]byte, 1<<20), 1<<20) + w := bufio.NewWriter(os.Stdout) + defer w.Flush() + for scanner.Scan() { + var req rpcRequest + if err := json.Unmarshal(scanner.Bytes(), &req); err != nil { + continue + } + if req.ID == 0 { + continue // 通知帧 + } + resp := fakeServerHandle(req) + body, _ := json.Marshal(resp) + w.Write(body) + w.WriteByte('\n') + w.Flush() + } +} + +func fakeServerHandle(req rpcRequest) rpcResponse { + respond := func(result any) rpcResponse { + raw, _ := json.Marshal(result) + return rpcResponse{JSONRPC: "2.0", ID: &req.ID, Result: raw} + } + fail := func(code int, msg string) rpcResponse { + return rpcResponse{JSONRPC: "2.0", ID: &req.ID, Error: &rpcError{Code: code, Message: msg}} + } + switch req.Method { + case "initialize": + return respond(initializeResult{ + ProtocolVersion: ProtocolVersion, + ServerInfo: clientInfo{Name: "fake-mcp", Version: "0"}, + }) + case "tools/list": + return respond(toolsListResult{Tools: fakeServerTools()}) + case "tools/call": + var p callToolParams + if req.Params != nil { + raw, _ := json.Marshal(req.Params) + _ = json.Unmarshal(raw, &p) + } + switch p.Name { + case "echo": + var args struct { + Text string `json:"text"` + } + _ = json.Unmarshal(p.Arguments, &args) + return respond(callToolResult{Content: []contentBlock{{Type: "text", Text: "echo: " + args.Text}}}) + case "boom": + return respond(callToolResult{ + Content: []contentBlock{{Type: "text", Text: "boom failed"}}, + IsError: true, + }) + default: + return fail(-32602, "unknown tool "+p.Name) + } + default: + return fail(-32601, "method not found: "+req.Method) + } +} + +func fakeServerTools() []toolInfo { + return []toolInfo{ + { + Name: "echo", + Description: "Echo the text back", + InputSchema: json.RawMessage(`{"type":"object","required":["text"],"properties":{"text":{"type":"string"}}}`), + }, + { + Name: "boom", + Description: "Always fails", + InputSchema: json.RawMessage(`{"type":"object","properties":{}}`), + }, + } +} + +// FakeServerEnv 是拉起 fake server 的环境变量名。 +const FakeServerEnv = "AGENT_NATIVE_TEST_MCP_STDIO" diff --git a/runtime/mcp/http.go b/runtime/mcp/http.go new file mode 100644 index 0000000..0bf5bcd --- /dev/null +++ b/runtime/mcp/http.go @@ -0,0 +1,196 @@ +package mcp + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "sync" + "sync/atomic" + "time" +) + +// httpClient 走 MCP Streamable HTTP 传输：每个 JSON-RPC 请求一次 POST， +// 期待 application/json 的单帧响应（SSE 流式与 server 主动请求不在本客户端范围）。 +// initialize 返回的 Mcp-Session-Id 会在后续请求回带。 +type httpClient struct { + cfg ServerConfig + client *http.Client + sessionID string + mu sync.Mutex + nextID int64 +} + +// DialHTTP 建立 HTTP 客户端并完成 initialize 握手。 +func DialHTTP(ctx context.Context, cfg ServerConfig) (*httpClient, error) { + if err := cfg.Validate(); err != nil { + return nil, err + } + if cfg.URL == "" { + return nil, fmt.Errorf("mcp server %q: http transport requires url", cfg.Name) + } + c := &httpClient{cfg: cfg, client: &http.Client{Timeout: requestTimeout(cfg)}} + if err := c.handshake(ctx); err != nil { + return nil, err + } + return c, nil +} + +func requestTimeout(cfg ServerConfig) time.Duration { + if cfg.Timeout > 0 { + return cfg.Timeout + } + return 30 * time.Second +} + +func (c *httpClient) handshake(ctx context.Context) error { + var res initializeResult + err := c.roundTrip(ctx, "initialize", initializeParams{ + ProtocolVersion: ProtocolVersion, + ClientInfo: clientInfo{Name: "agent-native", Version: "1"}, + }, &res) + return err +} + +func (c *httpClient) Initialize(ctx context.Context) (initializeResult, error) { + var res initializeResult + err := c.roundTrip(ctx, "initialize", initializeParams{ + ProtocolVersion: ProtocolVersion, + ClientInfo: clientInfo{Name: "agent-native", Version: "1"}, + }, &res) + return res, err +} + +func (c *httpClient) ListTools(ctx context.Context) ([]toolInfo, error) { + var out []toolInfo + cursor := "" + for { + params := map[string]any{} + if cursor != "" { + params["cursor"] = cursor + } + var res toolsListResult + if err := c.roundTrip(ctx, "tools/list", params, &res); err != nil { + return nil, err + } + out = append(out, res.Tools...) + if res.NextCursor == "" || len(out) > maxToolsPerServer*2 { + return out, nil + } + cursor = res.NextCursor + } +} + +func (c *httpClient) CallTool(ctx context.Context, name string, args json.RawMessage) (callToolResult, error) { + var res callToolResult + if len(args) == 0 { + args = json.RawMessage(`{}`) + } + err := c.roundTrip(ctx, "tools/call", callToolParams{Name: name, Arguments: args}, &res) + return res, err +} + +func (c *httpClient) Close() error { return nil } + +func (c *httpClient) roundTrip(ctx context.Context, method string, params any, out any) error { + c.mu.Lock() + defer c.mu.Unlock() + id := atomic.AddInt64(&c.nextID, 1) + body, err := json.Marshal(rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params}) + if err != nil { + return err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.cfg.URL, bytes.NewReader(body)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json, text/event-stream") + for k, v := range c.cfg.Headers { + req.Header.Set(k, v) + } + if c.sessionID != "" { + req.Header.Set("Mcp-Session-Id", c.sessionID) + } + resp, err := c.client.Do(req) + if err != nil { + return fmt.Errorf("mcp %s: %w", method, err) + } + defer resp.Body.Close() + if sid := resp.Header.Get("Mcp-Session-Id"); sid != "" { + c.sessionID = sid + } + if resp.StatusCode >= 400 { + snippet, _ := io.ReadAll(io.LimitReader(resp.Body, 2048)) + return fmt.Errorf("mcp %s: http %d: %s", method, resp.StatusCode, string(snippet)) + } + // Streamable HTTP 允许 SSE 回包：逐行扫 data: 帧，取第一个 JSON 对象。 + var payload []byte + ct := resp.Header.Get("Content-Type") + if bytes.Contains([]byte(ct), []byte("text/event-stream")) { + payload, err = readSSEFrame(resp.Body) + } else { + payload, err = io.ReadAll(io.LimitReader(resp.Body, 8<<20)) + } + if err != nil { + return fmt.Errorf("mcp %s: read body: %w", method, err) + } + var rpcResp rpcResponse + if err := json.Unmarshal(payload, &rpcResp); err != nil { + return fmt.Errorf("mcp %s: bad response frame: %w", method, err) + } + if rpcResp.ID == nil || *rpcResp.ID != id { + return fmt.Errorf("mcp %s: response id mismatch", method) + } + if rpcResp.Error != nil { + return rpcResp.Error + } + if out != nil && len(rpcResp.Result) > 0 { + return json.Unmarshal(rpcResp.Result, out) + } + return nil +} + +// readSSEFrame 从 SSE 流读取首个 data: JSON 载荷（忽略事件注释行）。 +func readSSEFrame(r io.Reader) ([]byte, error) { + buf := make([]byte, 0, 4096) + tmp := make([]byte, 2048) + for { + n, err := r.Read(tmp) + buf = append(buf, tmp[:n]...) + for _, line := range splitLines(buf) { + if bytes.HasPrefix(line, []byte("data:")) { + data := bytes.TrimSpace(line[len("data:"):]) + if len(data) > 0 { + return data, nil + } + } + } + if err != nil { + if len(buf) == 0 || err == io.EOF { + return nil, fmt.Errorf("no data frame in sse response") + } + return nil, err + } + if len(buf) > 8<<20 { + return nil, fmt.Errorf("sse response too large") + } + } +} + +func splitLines(b []byte) [][]byte { + var out [][]byte + start := 0 + for i, c := range b { + if c == '\n' { + out = append(out, b[start:i]) + start = i + 1 + } + } + if start < len(b) { + out = append(out, b[start:]) + } + return out +} diff --git a/runtime/mcp/manager.go b/runtime/mcp/manager.go new file mode 100644 index 0000000..f8e0e07 --- /dev/null +++ b/runtime/mcp/manager.go @@ -0,0 +1,171 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + "time" + + "agent-native/runtime/tools" +) + +// 桥接上限：server 元数据是不可信输入，长度与数量都必须受限。 +const ( + maxToolsPerServer = 64 + maxDescriptionBytes = 1024 + maxTextResultBytes = 64 << 10 + maxToolNameBytes = 128 + callTimeoutDefault = 30 * time.Second + registerTimeoutFixed = 30 * time.Second +) + +// Manager 维护一组 MCP server 连接，并把它们的工具注册进 tools.Registry。 +// 零值不可用；由组合根显式构造并注入（未注入 = 无 MCP 工具，fail-closed）。 +type Manager struct { + configs []ServerConfig + mu sync.Mutex + clients map[string]Client +} + +func NewManager(configs []ServerConfig) (*Manager, error) { + seen := map[string]bool{} + for _, c := range configs { + if err := c.Validate(); err != nil { + return nil, err + } + if seen[c.Name] { + return nil, fmt.Errorf("mcp server %q configured twice", c.Name) + } + seen[c.Name] = true + } + return &Manager{configs: configs, clients: map[string]Client{}}, nil +} + +// Configs 返回声明的 server 清单（只读副本）。 +func (m *Manager) Configs() []ServerConfig { + out := make([]ServerConfig, len(m.configs)) + copy(out, m.configs) + return out +} + +func (m *Manager) client(ctx context.Context, cfg ServerConfig) (Client, error) { + m.mu.Lock() + defer m.mu.Unlock() + if c, ok := m.clients[cfg.Name]; ok { + return c, nil + } + dialCtx, cancel := context.WithTimeout(ctx, registerTimeoutFixed) + defer cancel() + var c Client + var err error + if cfg.Command != "" { + c, err = DialStdio(dialCtx, cfg) + } else { + c, err = DialHTTP(dialCtx, cfg) + } + if err != nil { + return nil, fmt.Errorf("mcp server %q unreachable (fail-closed): %w", cfg.Name, err) + } + m.clients[cfg.Name] = c + return c, nil +} + +// RegisterInto 连接全部 server、枚举工具并以 mcp____ 命名注册。 +// 任一 server 不可达即返回错误（fail-closed，不静默跳过）。 +func (m *Manager) RegisterInto(ctx context.Context, reg *tools.Registry) error { + for _, cfg := range m.configs { + c, err := m.client(ctx, cfg) + if err != nil { + return err + } + listCtx, cancel := context.WithTimeout(ctx, registerTimeoutFixed) + list, err := c.ListTools(listCtx) + cancel() + if err != nil { + return fmt.Errorf("mcp server %q tools/list failed: %w", cfg.Name, err) + } + if len(list) > maxToolsPerServer { + return fmt.Errorf("mcp server %q exposes %d tools (max %d)", cfg.Name, len(list), maxToolsPerServer) + } + for _, ti := range list { + name := ToolName(cfg.Name, ti.Name) + if len(name) > maxToolNameBytes { + return fmt.Errorf("mcp tool name too long: %s", name) + } + desc := ti.Description + if len(desc) > maxDescriptionBytes { + desc = desc[:maxDescriptionBytes] + "…[truncated]" + } + schema := ti.InputSchema + if len(schema) == 0 { + schema = json.RawMessage(`{"type":"object","properties":{}}`) + } + reg.Register(&bridgeTool{ + client: c, + serverCfg: cfg, + toolName: ti.Name, + declName: name, + description: desc, + schema: schema, + }) + } + } + return nil +} + +// ToolName 生成桥接工具名（命名空间隔离）。 +func ToolName(server, tool string) string { + return "mcp__" + server + "__" + tool +} + +// Close 关闭全部连接（进程型 server 随之退出）。 +func (m *Manager) Close() { + m.mu.Lock() + defer m.mu.Unlock() + for _, c := range m.clients { + _ = c.Close() + } + m.clients = map[string]Client{} +} + +// bridgeTool 把一个 MCP 工具适配为 tools.Tool；调用失败重连一次后重试。 +type bridgeTool struct { + client Client + serverCfg ServerConfig + toolName string + declName string + description string + schema json.RawMessage +} + +func (t *bridgeTool) Name() string { return t.declName } +func (t *bridgeTool) Description() string { return t.description } +func (t *bridgeTool) Parameters() json.RawMessage { + return t.schema +} + +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) + if err != nil { + return tools.Result{Content: fmt.Sprintf("mcp tool %s failed: %v", t.declName, err), IsError: true}, nil + } + var b strings.Builder + for _, blk := range res.Content { + if blk.Type == "text" && blk.Text != "" { + b.WriteString(blk.Text) + b.WriteString("\n") + } + } + out := b.String() + if len(out) > maxTextResultBytes { + out = out[:maxTextResultBytes] + "\n[truncated]" + } + if res.IsError { + return tools.Result{Content: out, IsError: true}, nil + } + return tools.Result{Content: out}, nil +} diff --git a/runtime/mcp/mcp_test.go b/runtime/mcp/mcp_test.go new file mode 100644 index 0000000..92ed1e1 --- /dev/null +++ b/runtime/mcp/mcp_test.go @@ -0,0 +1,241 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "agent-native/runtime/tools" +) + +// TestMain：环境变量置位时，本测试二进制化身为 stdio MCP server（真实子进程测试）。 +func TestMain(m *testing.M) { + if os.Getenv(FakeServerEnv) == "1" { + RunFakeServer() + os.Exit(0) + } + os.Exit(m.Run()) +} + +func fakeStdioConfig(t *testing.T) ServerConfig { + t.Helper() + exe, err := os.Executable() + if err != nil { + t.Fatal(err) + } + return ServerConfig{ + Name: "fake", + Command: exe, + Env: append(os.Environ(), FakeServerEnv+"=1"), + } +} + +func TestStdioClientHandshakeListAndCall(t *testing.T) { + ctx := context.Background() + c, err := DialStdio(ctx, fakeStdioConfig(t)) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + list, err := c.ListTools(ctx) + if err != nil || len(list) != 2 || list[0].Name != "echo" { + t.Fatalf("list failed: %v %+v", err, list) + } + res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"hi"}`)) + if err != nil || len(res.Content) != 1 || res.Content[0].Text != "echo: hi" { + t.Fatalf("call failed: %v %+v", err, res) + } + _, err = c.CallTool(ctx, "no_such_tool", nil) + if err == nil || !strings.Contains(err.Error(), "unknown tool") { + t.Fatalf("unknown-tool call must fail via rpc error, got %v", err) + } +} + +func TestStdioClientServerRejectsBadSchema(t *testing.T) { + ctx := context.Background() + c, err := DialStdio(ctx, fakeStdioConfig(t)) + if err != nil { + t.Fatal(err) + } + defer c.Close() + // server 端 tool 报 isError：作为 callToolResult 返回而非 rpc 错误。 + res, err := c.CallTool(ctx, "boom", nil) + if err != nil { + t.Fatal(err) + } + if !res.IsError || res.Content[0].Text != "boom failed" { + t.Fatalf("want isError result, got %+v", res) + } +} + +func fakeHTTPHandler(w http.ResponseWriter, r *http.Request) { + var req rpcRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + w.WriteHeader(400) + return + } + resp := fakeServerHandle(req) + if sid := r.Header.Get("Mcp-Session-Id"); sid != "" { + w.Header().Set("Mcp-Session-Id", sid) + } else { + w.Header().Set("Mcp-Session-Id", "sess-1") + } + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(resp) +} + +func TestHTTPClientJSONMode(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(fakeHTTPHandler)) + defer srv.Close() + ctx := context.Background() + c, err := DialHTTP(ctx, ServerConfig{Name: "fakehttp", URL: srv.URL}) + if err != nil { + t.Fatal(err) + } + defer c.Close() + if c.sessionID != "sess-1" { + t.Fatalf("session id must be captured, got %q", c.sessionID) + } + list, err := c.ListTools(ctx) + if err != nil || len(list) != 2 { + t.Fatalf("list failed: %v %+v", err, list) + } + res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"net"}`)) + if err != nil || res.Content[0].Text != "echo: net" { + t.Fatalf("call failed: %v %+v", err, res) + } +} + +func TestHTTPClientSSEMode(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var req rpcRequest + _ = json.NewDecoder(r.Body).Decode(&req) + resp := fakeServerHandle(req) + w.Header().Set("Content-Type", "text/event-stream") + body, _ := json.Marshal(resp) + fmt.Fprintf(w, "event: message\ndata: %s\n\n", body) + })) + defer srv.Close() + ctx := context.Background() + c, err := DialHTTP(ctx, ServerConfig{Name: "fakesse", URL: srv.URL}) + if err != nil { + t.Fatal(err) + } + defer c.Close() + list, err := c.ListTools(ctx) + if err != nil || len(list) != 2 { + t.Fatalf("sse list failed: %v %+v", err, list) + } +} + +func TestManagerRegisterIntoAndCallThroughRegistry(t *testing.T) { + mgr, err := NewManager([]ServerConfig{fakeStdioConfig(t)}) + if err != nil { + t.Fatal(err) + } + defer mgr.Close() + reg := tools.NewRegistry(t.TempDir()) + if err := mgr.RegisterInto(context.Background(), reg); err != nil { + t.Fatal(err) + } + var names []string + for _, d := range reg.Decls() { + names = append(names, d.Name) + } + joined := strings.Join(names, ",") + if !strings.Contains(joined, "mcp__fake__echo") || !strings.Contains(joined, "mcp__fake__boom") { + t.Fatalf("mcp tools not registered: %v", names) + } + res := reg.Execute(context.Background(), "mcp__fake__echo", json.RawMessage(`{"text":"bridge"}`)) + if res.IsError || res.Content != "echo: bridge\n" { + t.Fatalf("bridge call failed: %+v", res) + } + res = reg.Execute(context.Background(), "mcp__fake__boom", nil) + if !res.IsError || !strings.Contains(res.Content, "boom failed") { + t.Fatalf("bridge error propagation failed: %+v", res) + } +} + +func TestManagerNamespaceWildcardRestrict(t *testing.T) { + mgr, _ := NewManager([]ServerConfig{fakeStdioConfig(t)}) + defer mgr.Close() + reg := tools.WorkspaceTools(t.TempDir()) + if err := mgr.RegisterInto(context.Background(), reg); err != nil { + t.Fatal(err) + } + if err := reg.Restrict([]string{"mcp__fake__*"}); err != nil { + t.Fatal(err) + } + for _, d := range reg.Decls() { + if !strings.HasPrefix(d.Name, "mcp__fake__") { + t.Fatalf("wildcard restrict must keep only mcp namespace, got %s", d.Name) + } + } + // 通配未命中任何工具 → fail-closed。 + err := reg.Restrict([]string{"mcp__nope__*"}) + if err == nil || !strings.Contains(err.Error(), "mcp__nope__*") { + t.Fatalf("unknown namespace wildcard must fail-closed, got %v", err) + } +} + +func TestManagerFailClosedOnUnreachableServer(t *testing.T) { + mgr, err := NewManager([]ServerConfig{{Name: "dead", Command: "/nonexistent/mcp-server-bin"}}) + if err != nil { + t.Fatal(err) + } + defer mgr.Close() + reg := tools.NewRegistry(t.TempDir()) + err = mgr.RegisterInto(context.Background(), reg) + if err == nil || !strings.Contains(err.Error(), "unreachable") { + t.Fatalf("unreachable server must fail-closed, got %v", err) + } +} + +func TestManagerConfigValidation(t *testing.T) { + if _, err := NewManager([]ServerConfig{{Name: "bad name", Command: "x"}}); err == nil { + t.Fatal("invalid name must be rejected") + } + if _, err := NewManager([]ServerConfig{{Name: "both", Command: "x", URL: "http://y"}}); err == nil { + t.Fatal("dual transport must be rejected") + } + if _, err := NewManager([]ServerConfig{{Name: "none"}}); err == nil { + t.Fatal("missing transport must be rejected") + } + _, err := NewManager([]ServerConfig{{Name: "a", Command: "x"}, {Name: "a", URL: "http://y"}}) + if err == nil || !strings.Contains(err.Error(), "twice") { + t.Fatal("duplicate names must be rejected") + } +} + +func TestManagerRegisterTimeoutBounded(t *testing.T) { + // server 挂起不响应：initialize 必须在超时内 fail-closed，而不是永久阻塞。 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-r.Context().Done(): + case <-time.After(10 * time.Second): + } + })) + defer srv.Close() + mgr, err := NewManager([]ServerConfig{{Name: "slow", URL: srv.URL, Timeout: 300 * time.Millisecond}}) + if err != nil { + t.Fatal(err) + } + defer mgr.Close() + reg := tools.NewRegistry(t.TempDir()) + start := time.Now() + err = mgr.RegisterInto(context.Background(), reg) + elapsed := time.Since(start) + if err == nil || !strings.Contains(err.Error(), "unreachable") { + t.Fatalf("hung server must fail-closed, got %v", err) + } + if elapsed > 5*time.Second { + t.Fatalf("timeout must bound the dial, took %s", elapsed) + } +} diff --git a/runtime/mcp/protocol.go b/runtime/mcp/protocol.go new file mode 100644 index 0000000..c8474c0 --- /dev/null +++ b/runtime/mcp/protocol.go @@ -0,0 +1,140 @@ +// Package mcp 实现 Model Context Protocol 客户端（stdio 与 Streamable HTTP +// 两种传输），并把 server 工具桥接为 runtime/tools.Registry 的普通工具。 +// +// 边界声明：server 提供的描述与 schema 视为不可信输入（长度受限、不解释）； +// 工具调用的策略（allowlist/确认）仍由 tools.Registry 策略层统一执行。 +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "time" +) + +// ProtocolVersion 是本客户端声明的 MCP 协议版本。 +const ProtocolVersion = "2024-11-05" + +// ServerConfig 声明一个 MCP server 接入点：Command（stdio）或 URL（HTTP）二选一。 +type ServerConfig struct { + Name string `json:"name"` + Command string `json:"command,omitempty"` + Args []string `json:"args,omitempty"` + Env []string `json:"env,omitempty"` + URL string `json:"url,omitempty"` + Headers map[string]string `json:"headers,omitempty"` + // Timeout 单次请求上限（默认 30s）。 + Timeout time.Duration `json:"timeout,omitempty"` +} + +// Validate 校验配置完整性（fail-closed：命名不规范或双传输/无传输都拒绝）。 +func (c ServerConfig) Validate() error { + if !validName(c.Name) { + return fmt.Errorf("mcp server name %q must be [a-zA-Z0-9_-]+", c.Name) + } + hasCmd := c.Command != "" + hasURL := c.URL != "" + if hasCmd == hasURL { + return fmt.Errorf("mcp server %q: exactly one of command or url is required", c.Name) + } + return nil +} + +func validName(s string) bool { + if s == "" { + return false + } + for _, r := range s { + switch { + case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '-', r == '_': + default: + return false + } + } + return true +} + +// ---- JSON-RPC 2.0 envelope ---- + +type rpcRequest struct { + JSONRPC string `json:"jsonrpc"` + ID int64 `json:"id,omitempty"` + Method string `json:"method"` + Params any `json:"params,omitempty"` +} + +type rpcNotification struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` +} + +type rpcResponse struct { + JSONRPC string `json:"jsonrpc"` + ID *int64 `json:"id"` + Result json.RawMessage `json:"result,omitempty"` + Error *rpcError `json:"error,omitempty"` +} + +type rpcError struct { + Code int `json:"code"` + Message string `json:"message"` + Data json.RawMessage `json:"data,omitempty"` +} + +func (e *rpcError) Error() string { + return fmt.Sprintf("mcp rpc error %d: %s", e.Code, e.Message) +} + +// ---- MCP 语义类型 ---- + +type initializeParams struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities json.RawMessage `json:"capabilities,omitempty"` + ClientInfo clientInfo `json:"clientInfo"` +} + +type clientInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +type initializeResult struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities json.RawMessage `json:"capabilities"` + ServerInfo clientInfo `json:"serverInfo"` + Instructions string `json:"instructions,omitempty"` +} + +type toolInfo struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + InputSchema json.RawMessage `json:"inputSchema,omitempty"` +} + +type toolsListResult struct { + Tools []toolInfo `json:"tools"` + NextCursor string `json:"nextCursor,omitempty"` +} + +type callToolParams struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments,omitempty"` +} + +type callToolResult struct { + Content []contentBlock `json:"content"` + IsError bool `json:"isError,omitempty"` +} + +type contentBlock struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` +} + +// Client 是一个已连接 MCP server 的最小面（initialize / tools/list / tools/call）。 +type Client interface { + Initialize(ctx context.Context) (initializeResult, error) + ListTools(ctx context.Context) ([]toolInfo, error) + CallTool(ctx context.Context, name string, args json.RawMessage) (callToolResult, error) + Close() error +} diff --git a/runtime/mcp/stdio.go b/runtime/mcp/stdio.go new file mode 100644 index 0000000..2e5bcda --- /dev/null +++ b/runtime/mcp/stdio.go @@ -0,0 +1,198 @@ +package mcp + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "os/exec" + "sync" + "sync/atomic" +) + +// stdioClient 通过「子进程 stdin/stdout 按行分隔的 JSON-RPC」与 server 通信。 +// 同一时刻只允许一个在途请求（串行锁），匹配简单且足以支撑智能体工具调用。 +type stdioClient struct { + cmd *exec.Cmd + stdin io.WriteCloser + stdout *bufio.Reader + mu sync.Mutex + nextID int64 + closed bool +} + +// DialStdio 启动 server 子进程并完成 initialize 握手。 +func DialStdio(ctx context.Context, cfg ServerConfig) (*stdioClient, error) { + if err := cfg.Validate(); err != nil { + return nil, err + } + if cfg.Command == "" { + return nil, fmt.Errorf("mcp server %q: stdio transport requires command", cfg.Name) + } + cmd := exec.Command(cfg.Command, cfg.Args...) + cmd.Env = cfg.Env + stdin, err := cmd.StdinPipe() + if err != nil { + return nil, err + } + stdout, err := cmd.StdoutPipe() + if err != nil { + return nil, err + } + if err := cmd.Start(); err != nil { + return nil, fmt.Errorf("mcp server %q: start %s: %w", cfg.Name, cfg.Command, err) + } + c := &stdioClient{cmd: cmd, stdin: stdin, stdout: bufio.NewReaderSize(stdout, 1<<20)} + if err := c.handshake(ctx); err != nil { + _ = c.Close() + return nil, err + } + return c, nil +} + +func (c *stdioClient) handshake(ctx context.Context) error { + var res initializeResult + if err := c.roundTrip(ctx, "initialize", initializeParams{ + ProtocolVersion: ProtocolVersion, + ClientInfo: clientInfo{Name: "agent-native", Version: "1"}, + }, &res); err != nil { + return err + } + // initialized 通知无需响应。 + return c.notify("notifications/initialized") +} + +func (c *stdioClient) Initialize(ctx context.Context) (initializeResult, error) { + // Initialize 已在 Dial 时完成；这里补一次协议层调用以支持探测场景。 + var res initializeResult + err := c.roundTrip(ctx, "initialize", initializeParams{ + ProtocolVersion: ProtocolVersion, + ClientInfo: clientInfo{Name: "agent-native", Version: "1"}, + }, &res) + return res, err +} + +func (c *stdioClient) ListTools(ctx context.Context) ([]toolInfo, error) { + var out []toolInfo + cursor := "" + for { + params := map[string]any{} + if cursor != "" { + params["cursor"] = cursor + } + var res toolsListResult + if err := c.roundTrip(ctx, "tools/list", params, &res); err != nil { + return nil, err + } + out = append(out, res.Tools...) + if res.NextCursor == "" || len(out) > maxToolsPerServer*2 { + return out, nil + } + cursor = res.NextCursor + } +} + +func (c *stdioClient) CallTool(ctx context.Context, name string, args json.RawMessage) (callToolResult, error) { + var res callToolResult + if len(args) == 0 { + args = json.RawMessage(`{}`) + } + err := c.roundTrip(ctx, "tools/call", callToolParams{Name: name, Arguments: args}, &res) + return res, err +} + +func (c *stdioClient) Close() error { + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return nil + } + c.closed = true + _ = c.stdin.Close() + if c.cmd.Process != nil { + _ = c.cmd.Process.Kill() + } + return c.cmd.Wait() +} + +func (c *stdioClient) roundTrip(ctx context.Context, method string, params any, out any) error { + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return fmt.Errorf("mcp client closed") + } + id := atomic.AddInt64(&c.nextID, 1) + req := rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params} + body, err := json.Marshal(req) + if err != nil { + return err + } + body = append(body, '\n') + if _, err := c.stdin.Write(body); err != nil { + return fmt.Errorf("mcp write %s: %w", method, err) + } + line, err := c.readResponseLine(ctx) + if err != nil { + return fmt.Errorf("mcp read %s: %w", method, err) + } + var resp rpcResponse + if err := json.Unmarshal(line, &resp); err != nil { + return fmt.Errorf("mcp %s: bad response frame: %w", method, err) + } + if resp.ID == nil || *resp.ID != id { + return fmt.Errorf("mcp %s: response id mismatch", method) + } + if resp.Error != nil { + return resp.Error + } + if out != nil && len(resp.Result) > 0 { + return json.Unmarshal(resp.Result, out) + } + return nil +} + +// readResponseLine 逐行读取，跳过通知帧（无 id），带 ctx 取消感知。 +func (c *stdioClient) readResponseLine(ctx context.Context) ([]byte, error) { + type result struct { + line []byte + err error + } + done := make(chan result, 1) + go func() { + for { + line, err := c.stdout.ReadBytes('\n') + if err != nil { + done <- result{nil, err} + return + } + var probe rpcResponse + if json.Unmarshal(line, &probe) == nil && probe.ID != nil { + done <- result{line, nil} + return + } + // 通知或日志行：忽略。 + } + }() + select { + case r := <-done: + return r.line, r.err + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (c *stdioClient) notify(method string) error { + c.mu.Lock() + defer c.mu.Unlock() + if c.closed { + return fmt.Errorf("mcp client closed") + } + body, err := json.Marshal(rpcNotification{JSONRPC: "2.0", Method: method}) + if err != nil { + return err + } + body = append(body, '\n') + _, err = c.stdin.Write(body) + return err +} diff --git a/runtime/promptx/promptx.go b/runtime/promptx/promptx.go new file mode 100644 index 0000000..92ba96c --- /dev/null +++ b/runtime/promptx/promptx.go @@ -0,0 +1,106 @@ +// Package promptx 组装执行内核的系统提示。 +// +// 职责边界（要求与实现分离）：人格/角色内容由上层（产品提示词库、角色 +// Instructions）提供并经 Parts.Instructions 传入；本包只做结构化装配—— +// 环境事实、技能索引、交付物与档位约束、护栏段。工具清单仍由 runtime +// 循环在系统提示之后追加，不在本包。 +package promptx + +import ( + "sort" + "strings" + "time" +) + +// EnvFacts 是注入系统提示的环境事实（由适配层填写，测试可固定）。 +type EnvFacts struct { + Date string // 如 2026-10-01（周一） + Model string + WorkspaceRoot string + Profile string // workspace | reasoning +} + +// IndexEntry 是技能索引行。 +type IndexEntry struct { + Name string + Description string +} + +// Parts 是一次执行的装配输入。 +type Parts struct { + Instructions string + // SkillsDeclared 是本次执行显式要求的技能名（角色声明）。 + SkillsDeclared []string + // SkillsIndex 是可用技能索引（渐进披露：只列名与描述）。 + SkillsIndex []IndexEntry + Deliverables []string + Env EnvFacts +} + +// Compose 产出系统提示正文。 +func Compose(p Parts) string { + var b strings.Builder + if strings.TrimSpace(p.Instructions) != "" { + b.WriteString(strings.TrimSpace(p.Instructions)) + b.WriteString("\n\n") + } + + // 环境事实块。 + env := p.Env + if env.Date == "" && env.Model == "" && env.WorkspaceRoot == "" { + env.Date = time.Now().Format("2006-01-02") + } + var envLines []string + if env.Date != "" { + envLines = append(envLines, "date: "+env.Date) + } + if env.Model != "" { + envLines = append(envLines, "model: "+env.Model) + } + if env.WorkspaceRoot != "" { + envLines = append(envLines, "workspace: "+env.WorkspaceRoot) + } + if env.Profile != "" { + envLines = append(envLines, "tool profile: "+env.Profile) + } + if len(envLines) > 0 { + b.WriteString("## Environment\n") + b.WriteString(strings.Join(envLines, "\n")) + b.WriteString("\n\n") + } + + // 技能块：声明优先 + 全量索引 + 按需加载指引。 + if len(p.SkillsDeclared) > 0 || len(p.SkillsIndex) > 0 { + b.WriteString("## Skills\n") + if len(p.SkillsDeclared) > 0 { + b.WriteString("Declared for this task (load with load_skill before relying on them): ") + b.WriteString(strings.Join(p.SkillsDeclared, ", ")) + b.WriteString("\n") + } + if len(p.SkillsIndex) > 0 { + entries := make([]IndexEntry, len(p.SkillsIndex)) + copy(entries, p.SkillsIndex) + sort.Slice(entries, func(i, j int) bool { return entries[i].Name < entries[j].Name }) + b.WriteString("Available skills (name — description; use load_skill for full instructions):\n") + for _, e := range entries { + line := "- " + e.Name + if e.Description != "" { + line += " — " + e.Description + } + b.WriteString(line + "\n") + } + } + b.WriteString("\n") + } + + // 交付物与档位约束（保持既有语义原文，行为兼容）。 + b.WriteString("Deliverable files: ") + b.WriteString(strings.Join(p.Deliverables, ", ")) + b.WriteString(".") + if p.Env.Profile == "reasoning" { + b.WriteString(" Reply with the deliverable text only. Tools are disabled.") + } else { + b.WriteString(" Use workspace tools to write those files. Do not invent paths outside the workspace.") + } + return strings.TrimSpace(b.String()) +} diff --git a/runtime/promptx/promptx_test.go b/runtime/promptx/promptx_test.go new file mode 100644 index 0000000..14387ca --- /dev/null +++ b/runtime/promptx/promptx_test.go @@ -0,0 +1,62 @@ +package promptx + +import ( + "strings" + "testing" +) + +func TestComposeSections(t *testing.T) { + out := Compose(Parts{ + Instructions: "你是审校员", + SkillsDeclared: []string{"docx-author"}, + SkillsIndex: []IndexEntry{ + {Name: "xlsx-author", Description: "表格交付"}, + {Name: "docx-author", Description: "文档交付"}, + }, + Deliverables: []string{"output.md"}, + Env: EnvFacts{Date: "2026-10-01（周四）", Model: "m1", WorkspaceRoot: "/w", Profile: "workspace"}, + }) + for _, want := range []string{ + "你是审校员", + "## Environment", + "date: 2026-10-01（周四）", + "model: m1", + "workspace: /w", + "tool profile: workspace", + "## Skills", + "Declared for this task", + "docx-author", + "- docx-author — 文档交付", // 索引按名排序 + "- xlsx-author — 表格交付", + "Deliverable files: output.md.", + "Use workspace tools to write those files.", + } { + if !strings.Contains(out, want) { + t.Errorf("missing %q in composed prompt:\n%s", want, out) + } + } + // 排序断言：docx 行在 xlsx 行之前。 + if strings.Index(out, "- docx-author") > strings.Index(out, "- xlsx-author") { + t.Errorf("index must be sorted by name:\n%s", out) + } +} + +func TestComposeReasoningProfileAndMinimal(t *testing.T) { + out := Compose(Parts{Deliverables: []string{"out.md"}, Env: EnvFacts{Profile: "reasoning"}}) + if !strings.Contains(out, "Reply with the deliverable text only. Tools are disabled.") { + t.Fatalf("reasoning sentence missing:\n%s", out) + } + if strings.Contains(out, "## Environment") && !strings.Contains(out, "date:") { + t.Fatalf("unexpected env block:\n%s", out) + } + if strings.Contains(out, "## Skills") { + t.Fatalf("no skills should render no skill block:\n%s", out) + } +} + +func TestComposeEnvDefaultsToDate(t *testing.T) { + out := Compose(Parts{Deliverables: []string{"o.md"}}) + if !strings.Contains(out, "date: ") { + t.Fatalf("env should default to today's date:\n%s", out) + } +} diff --git a/runtime/session/compact.go b/runtime/session/compact.go new file mode 100644 index 0000000..1c675d1 --- /dev/null +++ b/runtime/session/compact.go @@ -0,0 +1,138 @@ +package session + +import ( + "bufio" + "encoding/json" + "os" + + "agent-native/runtime/llm" +) + +// Stats 是转录文件的轻量元数据（扫描即得，无副作用）。 +type Stats struct { + Lines int + Bytes int64 +} + +// Stats 统计当前转录行数与字节数（空行不计）。 +func (s *Store) Stats() Stats { + s.mu.Lock() + defer s.mu.Unlock() + st := Stats{} + f, err := os.Open(s.path) + if err != nil { + return st + } + defer f.Close() + sc := bufio.NewScanner(f) + sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024) + for sc.Scan() { + if len(sc.Bytes()) == 0 { + continue + } + st.Lines++ + } + if info, err := f.Stat(); err == nil { + st.Bytes = info.Size() + } + return st +} + +// SaveMeta 把会话元数据写进 .meta.json（创建时间/模型等由调用方定义）。 +func (s *Store) SaveMeta(meta map[string]string) error { + s.mu.Lock() + defer s.mu.Unlock() + raw, err := json.MarshalIndent(meta, "", " ") + if err != nil { + return err + } + tmp := s.path + ".meta.json.tmp" + if err := os.WriteFile(tmp, raw, 0o644); err != nil { + return err + } + return os.Rename(tmp, s.path+".meta.json") +} + +// ReadMeta 读取会话元数据；不存在时返回 nil（新会话）。 +func (s *Store) ReadMeta() map[string]string { + raw, err := os.ReadFile(s.path + ".meta.json") + if err != nil { + return nil + } + var meta map[string]string + if err := json.Unmarshal(raw, &meta); err != nil { + return nil + } + return meta +} + +// Compact 把转录改写为「系统消息 + 摘要消息 + 最近 keepLastN 条」。 +// 摘要由 summarize 生成（由适配层注入，通常是一次无工具模型调用）； +// 改写原子落盘（临时文件 + rename）。 +// 返回是否发生了压缩（消息数不足/已是紧凑形态时返回 false 且不动文件）。 +func (s *Store) Compact(keepLastN int, summarize func(messages []llm.Message) (string, error)) (bool, error) { + messages, err := s.Load() + if err != nil { + return false, err + } + if len(messages) <= keepLastN+1 { // 系统消息 + 尾部，无可压缩头部 + return false, nil + } + var system *llm.Message + rest := messages + if len(rest) > 0 && rest[0].Role == llm.RoleSystem { + copyOf := rest[0] + system = ©Of + rest = rest[1:] + } + tailStart := len(rest) - keepLastN + if tailStart <= 0 { + return false, nil + } + head := rest[:tailStart] + tail := rest[tailStart:] + summary, err := summarize(head) + if err != nil { + return false, err + } + summaryMsg := llm.Message{ + Role: llm.RoleUser, + Content: "[conversation summary]\n" + summary + "\n[recent messages follow]", + } + var rebuilt []llm.Message + if system != nil { + rebuilt = append(rebuilt, *system) + } + rebuilt = append(rebuilt, summaryMsg) + rebuilt = append(rebuilt, tail...) + + s.mu.Lock() + defer s.mu.Unlock() + tmp := s.path + ".compact.tmp" + f, err := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644) + if err != nil { + return false, err + } + w := bufio.NewWriter(f) + for _, m := range rebuilt { + raw, err := json.Marshal(m) + if err != nil { + _ = f.Close() + _ = os.Remove(tmp) + return false, err + } + w.Write(raw) + w.WriteByte('\n') + } + if err := w.Flush(); err != nil { + _ = f.Close() + _ = os.Remove(tmp) + return false, err + } + _ = f.Close() + if err := os.Rename(tmp, s.path); err != nil { + _ = os.Remove(tmp) + return false, err + } + return true, nil +} diff --git a/runtime/session/compact_test.go b/runtime/session/compact_test.go new file mode 100644 index 0000000..b7c742f --- /dev/null +++ b/runtime/session/compact_test.go @@ -0,0 +1,159 @@ +package session + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "agent-native/runtime/llm" +) + +func appendMsgs(t *testing.T, s *Store, msgs ...llm.Message) { + t.Helper() + for _, m := range msgs { + if err := s.Append(m); err != nil { + t.Fatal(err) + } + } +} + +func TestLoadSkipsCrashResidueTail(t *testing.T) { + dir := t.TempDir() + s, err := Open(filepath.Join(dir, "t.jsonl")) + if err != nil { + t.Fatal(err) + } + appendMsgs(t, s, + llm.Message{Role: llm.RoleSystem, Content: "sys"}, + llm.Message{Role: llm.RoleUser, Content: "hi"}, + ) + // 模拟崩溃半行：无换行结尾的残缺 JSON。 + f, _ := os.OpenFile(s.Path(), os.O_APPEND|os.O_WRONLY, 0o644) + _, _ = f.WriteString(`{"role":"assistant","content":"trunca`) + _ = f.Close() + + msgs, err := s.Load() + if err != nil { + t.Fatalf("crash residue must not poison the session: %v", err) + } + if len(msgs) != 2 || msgs[0].Role != llm.RoleSystem || msgs[1].Content != "hi" { + t.Fatalf("want sys+hi, got %+v", msgs) + } +} + +func TestLoadStillFailsOnMidFileCorruption(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "t.jsonl") + if err := os.WriteFile(path, []byte("{\"bad json}\n{\"role\":\"user\",\"content\":\"ok\"}\n"), 0o644); err != nil { + t.Fatal(err) + } + s := &Store{path: path} + if _, err := s.Load(); err == nil { + t.Fatal("mid-file corruption must still error") + } +} + +func TestStatsAndMeta(t *testing.T) { + dir := t.TempDir() + s, _ := Open(filepath.Join(dir, "t.jsonl")) + appendMsgs(t, s, + llm.Message{Role: llm.RoleSystem, Content: "sys"}, + llm.Message{Role: llm.RoleUser, Content: "q"}, + ) + st := s.Stats() + if st.Lines != 2 || st.Bytes <= 0 { + t.Fatalf("stats wrong: %+v", st) + } + if err := s.SaveMeta(map[string]string{"model": "m1"}); err != nil { + t.Fatal(err) + } + meta := s.ReadMeta() + if meta == nil || meta["model"] != "m1" { + t.Fatalf("meta roundtrip failed: %v", meta) + } + if other := (&Store{path: filepath.Join(dir, "missing.jsonl")}).ReadMeta(); other != nil { + t.Fatalf("missing meta must return nil, got %v", other) + } +} + +func TestCompactRewritesTranscript(t *testing.T) { + dir := t.TempDir() + s, _ := Open(filepath.Join(dir, "t.jsonl")) + appendMsgs(t, s, llm.Message{Role: llm.RoleSystem, Content: "sys"}) + for i := 0; i < 10; i++ { + appendMsgs(t, s, + llm.Message{Role: llm.RoleUser, Content: "early" + string(rune('a'+i))}, + llm.Message{Role: llm.RoleAssistant, Content: "reply"}, + ) + } + var seenHead []llm.Message + done, err := s.Compact(3, func(messages []llm.Message) (string, error) { + seenHead = messages + return "这是摘要", nil + }) + if err != nil || !done { + t.Fatalf("compact failed: %v done=%v", err, done) + } + if len(seenHead) != 17 { // 20 - system - 3 tail + t.Fatalf("summarizer should receive 17 head messages, got %d", len(seenHead)) + } + msgs, err := s.Load() + if err != nil { + t.Fatal(err) + } + if len(msgs) != 5 { // system + summary + 3 tail + t.Fatalf("want 5 messages after compact, got %d", len(msgs)) + } + if msgs[0].Role != llm.RoleSystem { + t.Fatal("system message must be preserved") + } + if !strings.Contains(msgs[1].Content, "[conversation summary]") || !strings.Contains(msgs[1].Content, "这是摘要") { + t.Fatalf("summary message wrong: %+v", msgs[1]) + } + if msgs[2].Content != "reply" || msgs[3].Content != "earlyj" || msgs[4].Content != "reply" { + t.Fatalf("tail order wrong: %+v", msgs[2:]) + } + // 无 tmp 残留。 + if _, err := os.Stat(s.Path() + ".compact.tmp"); !os.IsNotExist(err) { + t.Fatal("compact tmp file must not remain") + } +} + +func TestCompactNoopWhenShort(t *testing.T) { + dir := t.TempDir() + s, _ := Open(filepath.Join(dir, "t.jsonl")) + appendMsgs(t, s, + llm.Message{Role: llm.RoleSystem, Content: "sys"}, + llm.Message{Role: llm.RoleUser, Content: "q"}, + ) + done, err := s.Compact(8, func([]llm.Message) (string, error) { return "x", nil }) + if err != nil || done { + t.Fatalf("short transcript must be a no-op: %v %v", done, err) + } + msgs, _ := s.Load() + if len(msgs) != 2 { + t.Fatal("no-op compact must not touch the file") + } +} + +func TestCompactSummarizerErrorAborts(t *testing.T) { + dir := t.TempDir() + s, _ := Open(filepath.Join(dir, "t.jsonl")) + appendMsgs(t, s, llm.Message{Role: llm.RoleSystem, Content: "sys"}) + for i := 0; i < 6; i++ { + appendMsgs(t, s, llm.Message{Role: llm.RoleUser, Content: "m"}) + } + done, err := s.Compact(2, func([]llm.Message) (string, error) { return "", errStringT("boom") }) + if err == nil || done { + t.Fatal("summarizer error must abort compaction") + } + msgs, _ := s.Load() + if len(msgs) != 7 { + t.Fatal("failed compaction must leave the transcript untouched") + } +} + +type errStringT string + +func (e errStringT) Error() string { return string(e) } diff --git a/runtime/session/jsonl.go b/runtime/session/jsonl.go index 2c29fa0..3d84b2c 100644 --- a/runtime/session/jsonl.go +++ b/runtime/session/jsonl.go @@ -5,6 +5,7 @@ import ( "encoding/json" "os" "path/filepath" + "strings" "sync" "agent-native/runtime/llm" @@ -47,16 +48,28 @@ func (s *Store) Append(msg llm.Message) error { return err } +// Load 全量读取转录。崩溃残留容错：Append 恒以换行收尾，因此无换行结尾的 +// 尾部若不是完整 JSON 即为崩溃半行，丢弃而不是让整个会话不可用； +// 行中间的坏行仍报错（那是文件损坏，不是崩溃残留）。 func (s *Store) Load() ([]llm.Message, error) { s.mu.Lock() defer s.mu.Unlock() - f, err := os.Open(s.path) + raw, err := os.ReadFile(s.path) if err != nil { return nil, err } - defer f.Close() + text := string(raw) + if text != "" && !strings.HasSuffix(text, "\n") { + if idx := strings.LastIndexByte(text, '\n'); idx >= 0 { + if !json.Valid([]byte(text[idx+1:])) { + text = text[:idx+1] + } + } else if !json.Valid([]byte(text)) { + text = "" + } + } var out []llm.Message - sc := bufio.NewScanner(f) + sc := bufio.NewScanner(strings.NewReader(text)) sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024) for sc.Scan() { line := sc.Bytes() diff --git a/runtime/skills/skills.go b/runtime/skills/skills.go new file mode 100644 index 0000000..387661c --- /dev/null +++ b/runtime/skills/skills.go @@ -0,0 +1,192 @@ +// Package skills 实现 SKILL.md 技能的加载与按需展开。 +// +// 目录约定：//SKILL.md，正文前可选 frontmatter +//（--- 包围，行前缀 description:，不引入 YAML 依赖）。 +// 目录名是技能名的权威来源（避免目录与 frontmatter 漂移）。 +// +// 信任边界：技能正文进入模型上下文，属提示注入面——只加载组合根显式 +// 配置的根目录与工作区 .agent/skills/；正文与索引总量受限。 +package skills + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + + "agent-native/runtime/tools" +) + +const ( + maxSkillsPerRoot = 100 + maxBodyBytes = 64 << 10 + maxDescription = 200 +) + +// Skill 是一个可加载技能。 +type Skill struct { + Name string + Description string + Body string + Source string +} + +// Store 扫描若干根目录并按名加载技能。 +type Store struct { + roots []string + cache map[string]Skill +} + +func NewStore(roots ...string) *Store { + clean := make([]string, 0, len(roots)) + for _, r := range roots { + if strings.TrimSpace(r) != "" { + clean = append(clean, r) + } + } + return &Store{roots: clean, cache: nil} +} + +// Roots 返回配置的根目录（只读副本）。 +func (s *Store) Roots() []string { + out := make([]string, len(s.roots)) + copy(out, s.roots) + return out +} + +// List 扫描全部根目录（同名技能先配置者优先），返回索引。 +func (s *Store) List() []Skill { + cache := map[string]Skill{} + var out []Skill + for _, root := range s.roots { + entries, err := os.ReadDir(root) + if err != nil { + continue // 根不存在：跳过（技能是可选能力） + } + count := 0 + for _, e := range entries { + if !e.IsDir() || count >= maxSkillsPerRoot { + continue + } + path := filepath.Join(root, e.Name(), "SKILL.md") + data, err := os.ReadFile(path) + if err != nil { + continue + } + skill, err := parse(e.Name(), path, data) + if err != nil { + continue // 单个坏技能跳过，不拖垮整个索引 + } + if _, dup := cache[skill.Name]; dup { + continue + } + cache[skill.Name] = skill + out = append(out, skill) + count++ + } + } + s.cache = cache + return out +} + +// Load 按名加载技能正文（fail-closed：未知名返回错误，不猜测近似名）。 +func (s *Store) Load(name string) (Skill, error) { + if s.cache == nil { + s.List() + } + if skill, ok := s.cache[name]; ok { + return skill, nil + } + // 缓存可能过期（List 之后目录变化）：重扫一遍再判。 + s.List() + if skill, ok := s.cache[name]; ok { + return skill, nil + } + return Skill{}, fmt.Errorf("unknown skill %q (not installed in any configured root)", name) +} + +// parse 解析一份 SKILL.md：目录名权威，frontmatter 可选。 +func parse(dirName, path string, data []byte) (Skill, error) { + if !validSkillName(dirName) { + return Skill{}, fmt.Errorf("invalid skill dir name %q", dirName) + } + if len(data) > maxBodyBytes { + data = data[:maxBodyBytes] + } + text := string(data) + desc := "" + body := text + lines := strings.Split(text, "\n") + if len(lines) > 0 && strings.TrimSpace(lines[0]) == "---" { + end := -1 + for i := 1; i < len(lines); i++ { + if strings.TrimSpace(lines[i]) == "---" { + end = i + break + } + } + if end > 0 { + for _, line := range lines[1:end] { + if v, ok := strings.CutPrefix(strings.TrimSpace(line), "description:"); ok { + desc = strings.TrimSpace(v) + } + } + body = strings.Join(lines[end+1:], "\n") + } + } + if len(desc) > maxDescription { + desc = desc[:maxDescription] + "…" + } + return Skill{Name: dirName, Description: desc, Body: body, Source: path}, nil +} + +func validSkillName(s string) bool { + if s == "" || len(s) > 64 { + return false + } + for _, r := range s { + switch { + case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '-', r == '_': + default: + return false + } + } + return true +} + +// ---- load_skill 工具 ---- + +type loadSkillTool struct{ store *Store } + +// RegisterLoader 注册 load_skill 工具（store 无技能时不注册，避免死工具）。 +func RegisterLoader(r *tools.Registry, store *Store) { + if store == nil || len(store.List()) == 0 { + return + } + r.Register(&loadSkillTool{store: store}) +} + +func (t *loadSkillTool) Name() string { return "load_skill" } +func (t *loadSkillTool) Description() string { + return "Load the full instructions of a skill by name (names are listed in the system prompt skill index)." +} + +func (t *loadSkillTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","required":["name"],"properties":{"name":{"type":"string"}}}`) +} + +func (t *loadSkillTool) Execute(_ context.Context, args json.RawMessage) (tools.Result, error) { + var obj struct { + Name string `json:"name"` + } + if err := json.Unmarshal(args, &obj); err != nil { + return tools.Result{}, err + } + skill, err := t.store.Load(obj.Name) + if err != nil { + return tools.Result{Content: err.Error(), IsError: true}, nil + } + return tools.Result{Content: "# skill: " + skill.Name + "\n\n" + skill.Body}, nil +} diff --git a/runtime/skills/skills_test.go b/runtime/skills/skills_test.go new file mode 100644 index 0000000..8b90451 --- /dev/null +++ b/runtime/skills/skills_test.go @@ -0,0 +1,112 @@ +package skills + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + + "agent-native/runtime/tools" +) + +func writeSkill(t *testing.T, root, name, body string) { + t.Helper() + dir := filepath.Join(root, name) + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "SKILL.md"), []byte(body), 0o644); err != nil { + t.Fatal(err) + } +} + +func TestStoreListAndLoad(t *testing.T) { + root := t.TempDir() + writeSkill(t, root, "docx-author", "---\ndescription: 生成 docx 交付物\n---\n\n## 步骤\n1. 用 run_python") + writeSkill(t, root, "xlsx-author", "没有 frontmatter 的技能正文") + writeSkill(t, root, "bad name", "目录名不合法") //nolint + + store := NewStore(root) + list := store.List() + if len(list) != 2 { + t.Fatalf("want 2 valid skills (bad dir skipped), got %d: %+v", len(list), list) + } + byName := map[string]Skill{} + for _, s := range list { + byName[s.Name] = s + } + if byName["docx-author"].Description != "生成 docx 交付物" { + t.Fatalf("frontmatter description wrong: %+v", byName["docx-author"]) + } + if !strings.Contains(byName["docx-author"].Body, "run_python") { + t.Fatalf("body must exclude frontmatter, got %q", byName["docx-author"].Body) + } + if byName["xlsx-author"].Description != "" || !strings.Contains(byName["xlsx-author"].Body, "没有 frontmatter") { + t.Fatalf("no-frontmatter skill wrong: %+v", byName["xlsx-author"]) + } + + skill, err := store.Load("docx-author") + if err != nil || !strings.Contains(skill.Body, "步骤") { + t.Fatalf("load failed: %v %+v", err, skill) + } + _, err = store.Load("nope") + if err == nil || !strings.Contains(err.Error(), "unknown skill") { + t.Fatalf("unknown skill must fail-closed, got %v", err) + } +} + +func TestStoreRootPriorityAndMissingRoot(t *testing.T) { + r1, r2 := t.TempDir(), t.TempDir() + writeSkill(t, r1, "dup", "from r1") + writeSkill(t, r2, "dup", "from r2") + writeSkill(t, r2, "only-r2", "x") + store := NewStore(r1, r2, "/nonexistent/root") + list := store.List() + if len(list) != 2 { + t.Fatalf("want 2 skills, got %d", len(list)) + } + skill, _ := store.Load("dup") + if skill.Body != "from r1" { + t.Fatalf("earlier root must win, got %q", skill.Body) + } +} + +func TestRegisterLoaderTool(t *testing.T) { + root := t.TempDir() + writeSkill(t, root, "docx-author", "---\ndescription: d\n---\n正文指令") + reg := tools.NewRegistry(t.TempDir()) + RegisterLoader(reg, NewStore(root)) + var found bool + for _, d := range reg.Decls() { + if d.Name == "load_skill" { + found = true + } + } + if !found { + t.Fatal("load_skill must be registered when skills exist") + } + res := reg.Execute(context.Background(), "load_skill", json.RawMessage(`{"name":"docx-author"}`)) + if res.IsError || !strings.Contains(res.Content, "正文指令") { + t.Fatalf("load failed: %+v", res) + } + res = reg.Execute(context.Background(), "load_skill", json.RawMessage(`{"name":"ghost"}`)) + if !res.IsError || !strings.Contains(res.Content, "unknown skill") { + t.Fatalf("ghost skill must error, got %+v", res) + } +} + +func TestRegisterLoaderSkipsEmptyStore(t *testing.T) { + reg := tools.NewRegistry(t.TempDir()) + RegisterLoader(reg, NewStore(t.TempDir())) + for _, d := range reg.Decls() { + if d.Name == "load_skill" { + t.Fatal("no skills = no loader tool") + } + } + RegisterLoader(reg, nil) + if len(reg.Decls()) != 0 { + t.Fatal("nil store must not register") + } +} diff --git a/runtime/subagent/subagent.go b/runtime/subagent/subagent.go new file mode 100644 index 0000000..f7c2199 --- /dev/null +++ b/runtime/subagent/subagent.go @@ -0,0 +1,103 @@ +// Package subagent 定义委派端口：内核不内置嵌套执行（那是编排层职责）， +// 只提供工具面与信息集契约；宿主（产品 SpawnSubagent、小队引擎等） +// 注入 Factory 实现真实派生执行。未注入则 delegate_task 不注册（fail-closed）。 +// +// 信息集对齐 agent-execution 标准的委派条款：目标、输入、写范围、非目标、 +// 完成判据、预算与回传格式——子任务不能从父任务目标推导额外权限。 +package subagent + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strings" + "time" + + "agent-native/runtime/tools" +) + +// Request 是一次委派的完整信息集（模型经工具参数逐项填写）。 +type Request struct { + Goal string `json:"goal"` + Inputs string `json:"inputs,omitempty"` + WriteScope string `json:"write_scope,omitempty"` + NonGoals string `json:"non_goals,omitempty"` + DoneCriteria string `json:"done_criteria,omitempty"` + Deliverables []string `json:"deliverables,omitempty"` + // Deadline / MaxToolCalls 是宿主注入的预算余量（模型不可自行扩权）。 + Deadline time.Time `json:"-"` + MaxToolCalls int `json:"-"` +} + +// Result 是子代理的回传：至少包含结论文本；文件由宿主按写范围物化。 +type Result struct { + Success bool + Text string + Files map[string][]byte + Err string +} + +// Factory 由宿主注入；ctx 取消必须传播到子执行。 +type Factory func(ctx context.Context, req Request) (Result, error) + +type delegateTool struct{ factory Factory } + +// Register 把 delegate_task 工具注册进注册表（factory 为 nil 时不注册）。 +func Register(r *tools.Registry, factory Factory) { + if factory == nil { + return + } + r.Register(&delegateTool{factory: factory}) +} + +func (t *delegateTool) Name() string { return "delegate_task" } +func (t *delegateTool) Description() string { + return "Delegate a self-contained subtask to a fresh sub-agent. Fill every field: goal, inputs (with versions), write_scope (directory the sub-agent may write), non_goals, done_criteria, deliverables. The sub-agent cannot exceed your permissions; report back its conclusion and files." +} +func (t *delegateTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","required":["goal","done_criteria"],"properties":{ + "goal":{"type":"string"}, + "inputs":{"type":"string"}, + "write_scope":{"type":"string"}, + "non_goals":{"type":"string"}, + "done_criteria":{"type":"string"}, + "deliverables":{"type":"array","items":{"type":"string"}} + }}`) +} + +func (t *delegateTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) { + var req Request + if err := json.Unmarshal(args, &req); err != nil { + return tools.Result{}, err + } + if strings.TrimSpace(req.Goal) == "" { + return tools.Result{Content: "goal is empty", IsError: true}, nil + } + if ctx.Err() != nil { + return tools.Result{Content: "delegation canceled before start: " + ctx.Err().Error(), IsError: true}, nil + } + res, err := t.factory(ctx, req) + if err != nil { + return tools.Result{Content: fmt.Sprintf("sub-agent failed: %v", err), IsError: true}, nil + } + var b strings.Builder + if res.Success { + b.WriteString("[sub-agent succeeded]\n") + } else { + b.WriteString("[sub-agent failed]\n") + } + if res.Err != "" { + b.WriteString("error: " + res.Err + "\n") + } + b.WriteString(res.Text) + if len(res.Files) > 0 { + names := make([]string, 0, len(res.Files)) + for name := range res.Files { + names = append(names, name) + } + sort.Strings(names) + b.WriteString("\nfiles: " + strings.Join(names, ", ")) + } + return tools.Result{Content: b.String(), IsError: !res.Success}, nil +} diff --git a/runtime/subagent/subagent_test.go b/runtime/subagent/subagent_test.go new file mode 100644 index 0000000..1c58bfb --- /dev/null +++ b/runtime/subagent/subagent_test.go @@ -0,0 +1,86 @@ +package subagent + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "agent-native/runtime/tools" +) + +func TestRegisterNilFactoryNoop(t *testing.T) { + reg := tools.NewRegistry(t.TempDir()) + Register(reg, nil) + for _, d := range reg.Decls() { + if d.Name == "delegate_task" { + t.Fatal("nil factory must not register delegate_task") + } + } +} + +func TestDelegateToolRoundTrip(t *testing.T) { + var got Request + reg := tools.NewRegistry(t.TempDir()) + Register(reg, func(ctx context.Context, req Request) (Result, error) { + got = req + return Result{ + Success: true, + Text: "调研完成：结论 A。", + Files: map[string][]byte{"notes.md": []byte("n"), "data/table.csv": []byte("c")}, + }, nil + }) + res := reg.Execute(context.Background(), "delegate_task", json.RawMessage(`{ + "goal":"调研主题 X", + "inputs":"docs/v2 规范", + "write_scope":"research/", + "non_goals":"不做实现", + "done_criteria":"产出三来源对照表", + "deliverables":["notes.md","data/table.csv"] + }`)) + if res.IsError { + t.Fatal(res.Content) + } + if got.Goal != "调研主题 X" || got.WriteScope != "research/" || got.DoneCriteria != "产出三来源对照表" { + t.Fatalf("delegation fields lost: %+v", got) + } + if !strings.Contains(res.Content, "[sub-agent succeeded]") || + !strings.Contains(res.Content, "调研完成") || + !strings.Contains(res.Content, "files: data/table.csv, notes.md") { + t.Fatalf("result text wrong:\n%s", res.Content) + } +} + +func TestDelegateToolFailureAndValidation(t *testing.T) { + reg := tools.NewRegistry(t.TempDir()) + Register(reg, func(ctx context.Context, req Request) (Result, error) { + return Result{Success: false, Err: "子任务超出预算"}, nil + }) + res := reg.Execute(context.Background(), "delegate_task", json.RawMessage(`{"goal":"g","done_criteria":"d"}`)) + if !res.IsError || !strings.Contains(res.Content, "[sub-agent failed]") || !strings.Contains(res.Content, "子任务超出预算") { + t.Fatalf("failure must propagate, got %+v", res) + } + // goal 必填（schema 校验拦截）。 + res = reg.Execute(context.Background(), "delegate_task", json.RawMessage(`{"done_criteria":"d"}`)) + if !res.IsError || !strings.Contains(res.Content, "required argument goal") { + t.Fatalf("missing goal must fail schema validation, got %+v", res) + } +} + +func TestDelegateToolFactoryErrorAndCancel(t *testing.T) { + reg := tools.NewRegistry(t.TempDir()) + Register(reg, func(ctx context.Context, req Request) (Result, error) { + return Result{}, context.Canceled + }) + res := reg.Execute(context.Background(), "delegate_task", json.RawMessage(`{"goal":"g","done_criteria":"d"}`)) + if !res.IsError || !strings.Contains(res.Content, "sub-agent failed") { + t.Fatalf("factory error must surface, got %+v", res) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + res = reg.Execute(ctx, "delegate_task", json.RawMessage(`{"goal":"g","done_criteria":"d"}`)) + if !res.IsError || !strings.Contains(res.Content, "canceled before start") { + t.Fatalf("pre-canceled delegation must not start, got %+v", res) + } +} diff --git a/runtime/tools/builtin.go b/runtime/tools/builtin.go index 7fef7ec..d39b517 100644 --- a/runtime/tools/builtin.go +++ b/runtime/tools/builtin.go @@ -6,6 +6,7 @@ import ( "fmt" "os" "path/filepath" + "regexp" "strings" ) @@ -13,17 +14,21 @@ type readFile struct{ root string } func (t *readFile) Name() string { return "read_file" } func (t *readFile) Description() string { - return "Read a UTF-8 text file inside the workspace." + 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." } func (t *readFile) Parameters() json.RawMessage { - return json.RawMessage(`{"type":"object","required":["path"],"properties":{"path":{"type":"string"}}}`) + return json.RawMessage(`{"type":"object","required":["path"],"properties":{"path":{"type":"string"},"offset":{"type":"integer","description":"1-based first line to return"},"limit":{"type":"integer","description":"number of lines to return"}}}`) } func (t *readFile) Execute(_ context.Context, args json.RawMessage) (Result, error) { - rel, err := readArg(args, "path") - if err != nil { + var obj struct { + Path string `json:"path"` + Offset int `json:"offset"` + Limit int `json:"limit"` + } + if err := json.Unmarshal(args, &obj); err != nil { return Result{}, err } - path, err := mustPath(t.root, rel) + path, err := mustPath(t.root, obj.Path) if err != nil { return Result{Content: err.Error(), IsError: true}, nil } @@ -31,10 +36,37 @@ func (t *readFile) Execute(_ context.Context, args json.RawMessage) (Result, err if err != nil { return Result{Content: err.Error(), IsError: true}, nil } - if len(b) > maxReadBytes { - b = b[:maxReadBytes] + if obj.Offset <= 0 && obj.Limit <= 0 { + truncated := false + if len(b) > maxReadBytes { + b = b[:maxReadBytes] + truncated = true + } + content := string(b) + if truncated { + content += "\n[truncated at 200 KiB; use offset/limit to page through the file]" + } + return Result{Content: content}, nil } - return Result{Content: string(b)}, nil + if obj.Offset <= 0 { + obj.Offset = 1 + } + if obj.Limit <= 0 { + obj.Limit = 1000 + } + lines := strings.Split(string(b), "\n") + if obj.Offset > len(lines) { + return Result{Content: fmt.Sprintf("offset %d beyond end of file (%d lines)", obj.Offset, len(lines)), IsError: true}, nil + } + end := obj.Offset - 1 + obj.Limit + if end > len(lines) { + end = len(lines) + } + var out []string + for i := obj.Offset - 1; i < end; i++ { + out = append(out, fmt.Sprintf("%d\t%s", i+1, lines[i])) + } + return Result{Content: strings.Join(out, "\n")}, nil } type writeFile struct{ root string } @@ -145,22 +177,51 @@ type grepTool struct{ root string } func (t *grepTool) Name() string { return "grep" } func (t *grepTool) Description() string { - return "Search workspace text files for a literal substring." + 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." } func (t *grepTool) Parameters() json.RawMessage { - return json.RawMessage(`{"type":"object","required":["pattern"],"properties":{"pattern":{"type":"string"},"path":{"type":"string"}}}`) + return json.RawMessage(`{"type":"object","required":["pattern"],"properties":{"pattern":{"type":"string"},"path":{"type":"string"},"regex":{"type":"boolean"},"ignore_case":{"type":"boolean"},"include":{"type":"string","description":"file-name glob filter such as *.go"}}}`) } func (t *grepTool) Execute(_ context.Context, args json.RawMessage) (Result, error) { var obj struct { - Pattern string `json:"pattern"` - Path string `json:"path"` + Pattern string `json:"pattern"` + Path string `json:"path"` + Regex bool `json:"regex"` + IgnoreCase bool `json:"ignore_case"` + Include string `json:"include"` } if err := json.Unmarshal(args, &obj); err != nil { return Result{}, err } + if strings.TrimSpace(obj.Pattern) == "" { + return Result{Content: "pattern is empty", IsError: true}, nil + } if obj.Path == "" { obj.Path = "." } + var includeRe *globMatcher + if strings.TrimSpace(obj.Include) != "" { + m, err := newGlobMatcher(obj.Include) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + includeRe = m + } + contains := func(line string) bool { return strings.Contains(line, obj.Pattern) } + if obj.Regex { + pat := obj.Pattern + if obj.IgnoreCase { + pat = "(?i)" + pat + } + re, cerr := regexp.Compile(pat) + if cerr != nil { + return Result{Content: cerr.Error(), IsError: true}, nil + } + contains = re.MatchString + } else if obj.IgnoreCase { + lower := strings.ToLower(obj.Pattern) + contains = func(line string) bool { return strings.Contains(strings.ToLower(line), lower) } + } start, err := mustPath(t.root, obj.Path) if err != nil { return Result{Content: err.Error(), IsError: true}, nil @@ -173,13 +234,16 @@ func (t *grepTool) Execute(_ context.Context, args json.RawMessage) (Result, err if len(hits) >= 50 { return fmt.Errorf("stop") } + if includeRe != nil && !includeRe.Match(filepath.Base(path)) { + return nil + } b, err := os.ReadFile(path) if err != nil || len(b) > maxReadBytes { return nil } rel, _ := filepath.Rel(t.root, path) for i, line := range strings.Split(string(b), "\n") { - if strings.Contains(line, obj.Pattern) { + if contains(line) { hits = append(hits, fmt.Sprintf("%s:%d:%s", filepath.ToSlash(rel), i+1, line)) if len(hits) >= 50 { return fmt.Errorf("stop") @@ -196,3 +260,55 @@ func (t *grepTool) Execute(_ context.Context, args json.RawMessage) (Result, err } return Result{Content: strings.Join(hits, "\n")}, nil } + +// multiEditTool 一次多处替换：全部命中且互不冲突才落盘，任一失败即整体不动。 +type multiEditTool struct{ root string } + +func (t *multiEditTool) Name() string { return "multi_edit" } +func (t *multiEditTool) Description() string { + return "Apply several unique-span replacements to one workspace file atomically: either all edits land or the file is left untouched." +} +func (t *multiEditTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","required":["path","edits"],"properties":{"path":{"type":"string"},"edits":{"type":"array","items":{"type":"object","required":["old","new"],"properties":{"old":{"type":"string"},"new":{"type":"string"}}}}}}`) +} + +type multiEditSpan struct { + Old string `json:"old"` + New string `json:"new"` +} + +func (t *multiEditTool) Execute(_ context.Context, args json.RawMessage) (Result, error) { + var obj struct { + Path string `json:"path"` + Edits []multiEditSpan `json:"edits"` + } + if err := json.Unmarshal(args, &obj); err != nil { + return Result{}, err + } + if len(obj.Edits) == 0 { + return Result{Content: "edits is empty", IsError: true}, nil + } + path, err := mustPath(t.root, obj.Path) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + b, err := os.ReadFile(path) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + text := string(b) + for i, e := range obj.Edits { + if e.Old == "" { + return Result{Content: fmt.Sprintf("edit %d: old text is empty", i+1), IsError: true}, nil + } + n := strings.Count(text, e.Old) + if n != 1 { + return Result{Content: fmt.Sprintf("edit %d: expected one match, found %d; file unchanged", i+1, n), IsError: true}, nil + } + text = strings.Replace(text, e.Old, e.New, 1) + } + if err := os.WriteFile(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/code_tools_test.go b/runtime/tools/code_tools_test.go new file mode 100644 index 0000000..768d38f --- /dev/null +++ b/runtime/tools/code_tools_test.go @@ -0,0 +1,158 @@ +package tools + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" +) + +func writeFileT(t *testing.T, root, rel, content string) { + t.Helper() + path := filepath.Join(root, filepath.FromSlash(rel)) + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, []byte(content), 0o644); err != nil { + t.Fatal(err) + } +} + +func runTool(t *testing.T, reg *Registry, name, args string) Result { + t.Helper() + res := reg.Execute(context.Background(), name, []byte(args)) + return res +} + +func TestReadFileLineRange(t *testing.T) { + root := t.TempDir() + writeFileT(t, root, "a.txt", "l1\nl2\nl3\nl4\nl5") + reg := WorkspaceTools(root) + + res := runTool(t, reg, "read_file", `{"path":"a.txt","offset":2,"limit":2}`) + if res.IsError || res.Content != "2\tl2\n3\tl3" { + t.Fatalf("want numbered lines 2-3, got %+v", res) + } + res = runTool(t, reg, "read_file", `{"path":"a.txt","offset":4}`) + if res.IsError || res.Content != "4\tl4\n5\tl5" { + t.Fatalf("want tail from line 4, got %+v", res) + } + res = runTool(t, reg, "read_file", `{"path":"a.txt","offset":99}`) + if !res.IsError || !strings.Contains(res.Content, "beyond end") { + t.Fatalf("want beyond-end error, got %+v", res) + } + // 无 offset/limit 保持旧行为：整文件原文。 + res = runTool(t, reg, "read_file", `{"path":"a.txt"}`) + if res.IsError || res.Content != "l1\nl2\nl3\nl4\nl5" { + t.Fatalf("legacy whole-file read changed, got %+v", res) + } +} + +func TestGrepRegexIncludeAndCase(t *testing.T) { + root := t.TempDir() + writeFileT(t, root, "src/a.go", "func Alpha() {}\nfunc beta() {}\n") + writeFileT(t, root, "src/b.go", "func Gamma() {}\n") + writeFileT(t, root, "doc/a.md", "func Alpha docs\n") + reg := WorkspaceTools(root) + + res := runTool(t, reg, "grep", `{"pattern":"func [A-Z][a-z]+","regex":true,"include":"*.go"}`) + if res.IsError { + t.Fatal(res.Content) + } + if !strings.Contains(res.Content, "src/a.go:1:func Alpha() {}") || !strings.Contains(res.Content, "src/b.go:1:func Gamma() {}") { + t.Fatalf("regex+include should hit go files only, got:\n%s", res.Content) + } + if strings.Contains(res.Content, "a.md") || strings.Contains(res.Content, "beta") { + t.Fatalf("md file / lowercase must not match, got:\n%s", res.Content) + } + + res = runTool(t, reg, "grep", `{"pattern":"FUNC BETA","regex":true,"ignore_case":true}`) + if res.IsError || !strings.Contains(res.Content, "func beta() {}") { + t.Fatalf("ignore_case regex failed, got %+v", res) + } + + res = runTool(t, reg, "grep", `{"pattern":"func [","regex":true}`) + if !res.IsError { + t.Fatal("invalid regex must be an error result") + } + + // 字面默认路径不变。 + res = runTool(t, reg, "grep", `{"pattern":"Alpha"}`) + if res.IsError || !strings.Contains(res.Content, "a.go:1") { + t.Fatalf("literal default failed, got %+v", res) + } +} + +func TestGlobDoubleStar(t *testing.T) { + root := t.TempDir() + writeFileT(t, root, "docs/x/deep/a.md", "x") + writeFileT(t, root, "docs/b.md", "x") + writeFileT(t, root, "src/main.go", "x") + writeFileT(t, root, "main.go", "x") + reg := WorkspaceTools(root) + + res := runTool(t, reg, "glob", `{"pattern":"**/*.go"}`) + if res.IsError { + t.Fatal(res.Content) + } + if !strings.Contains(res.Content, "src/main.go") || !strings.Contains(res.Content, "main.go") || strings.Contains(res.Content, ".md") { + t.Fatalf("**/*.go wrong, got:\n%s", res.Content) + } + + res = runTool(t, reg, "glob", `{"pattern":"docs/**/*.md"}`) + if res.IsError || !strings.Contains(res.Content, "docs/x/deep/a.md") || !strings.Contains(res.Content, "docs/b.md") { + t.Fatalf("docs/**/*.md wrong, got:\n%s", res.Content) + } + + res = runTool(t, reg, "glob", `{"pattern":"docs/**"}`) + if res.IsError || !strings.Contains(res.Content, "docs/b.md") || !strings.Contains(res.Content, "docs/x/deep/a.md") { + t.Fatalf("trailing ** must match everything below, got:\n%s", res.Content) + } + + res = runTool(t, reg, "glob", `{"pattern":"**/nope.txt"}`) + if res.IsError || res.Content != "no matches" { + t.Fatalf("want no matches, got %+v", res) + } +} + +func TestGlobEscapingOutsideWorkspace(t *testing.T) { + reg := WorkspaceTools(t.TempDir()) + // glob 只在工作区内匹配（WalkDir 从 root 出发），模式里的 ../ 无从逃逸； + // 字面 .. 也只是找不到匹配。 + res := runTool(t, reg, "glob", `{"pattern":"../**/*.go"}`) + if res.IsError && res.Content == "" { + t.Fatal("unexpected empty error") + } +} + +func TestMultiEditAtomic(t *testing.T) { + root := t.TempDir() + writeFileT(t, root, "a.md", "alpha\nbeta\ngamma\n") + reg := WorkspaceTools(root) + + res := runTool(t, reg, "multi_edit", `{"path":"a.md","edits":[{"old":"alpha","new":"ALPHA"},{"old":"gamma","new":"GAMMA"}]}`) + if res.IsError || res.Content != "applied 2 edits to a.md" { + t.Fatalf("want both applied, got %+v", res) + } + data, _ := os.ReadFile(filepath.Join(root, "a.md")) + if string(data) != "ALPHA\nbeta\nGAMMA\n" { + t.Fatalf("file content wrong: %q", data) + } + + // 第二处替换失败 → 文件保持原样（全有或全无）。 + writeFileT(t, root, "b.md", "one\ntwo\n") + res = runTool(t, reg, "multi_edit", `{"path":"b.md","edits":[{"old":"one","new":"ONE"},{"old":"three","new":"THREE"}]}`) + if !res.IsError || !strings.Contains(res.Content, "found 0") || !strings.Contains(res.Content, "unchanged") { + t.Fatalf("want atomic failure, got %+v", res) + } + data, _ = os.ReadFile(filepath.Join(root, "b.md")) + if string(data) != "one\ntwo\n" { + t.Fatalf("failed batch must not modify file, got %q", data) + } + + res = runTool(t, reg, "multi_edit", `{"path":"b.md","edits":[]}`) + if !res.IsError { + t.Fatal("empty edits must be an error") + } +} diff --git a/runtime/tools/exec.go b/runtime/tools/exec.go new file mode 100644 index 0000000..410d2ce --- /dev/null +++ b/runtime/tools/exec.go @@ -0,0 +1,194 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "math/rand" + "os" + "os/exec" + "path/filepath" + "strings" + "syscall" + "time" +) + +// ExecConfig 声明执行类工具的运行条件，由组合根显式提供； +// 未提供（nil）时 run_command / run_python 一律不注册（fail-closed）。 +// 注意：工作区路径约束不等于 OS 沙箱——进程仍以当前用户权限运行， +// 隔离需由宿主环境提供（容器/沙箱），本文件只做时限、输出、env 与审计收敛。 +type ExecConfig struct { + // Bash / Python 是解释器绝对路径；留空则对应工具不注册。 + Bash string + Python string + // Env 传给子进程的完整白名单环境（KEY=VALUE）；不含本进程任何隐式变量。 + Env []string + // ScriptDir 落盘脚本目录（相对 root，默认 .agent/scripts）；空用默认。 + ScriptDir string +} + +const ( + execDefaultTimeout = 60 * time.Second + execMaxTimeout = 300 * time.Second + execMaxOutput = 64 << 10 +) + +// blockedCommandMarkers 最小命令黑名单：明显提权/自毁模式直接拒绝。 +// 这不是沙箱，只是防呆；真实隔离靠宿主环境。 +var blockedCommandMarkers = []string{ + "sudo ", "sudo\t", "sudo\n", + "rm -rf /", "rm -fr /", + ":(){:|:&};:", + "mkfs", "shutdown", "reboot", +} + +// RegisterExecTools 注册 run_command / run_python（解释器未配置的项跳过注册）。 +func RegisterExecTools(r *Registry, root string, cfg ExecConfig) { + if strings.TrimSpace(cfg.Bash) != "" { + r.Register(&runCommandTool{root: root, cfg: cfg}) + } + if strings.TrimSpace(cfg.Python) != "" { + r.Register(&runPythonTool{root: root, cfg: cfg}) + } +} + +type runCommandTool struct { + root string + cfg ExecConfig +} + +func (t *runCommandTool) Name() string { return "run_command" } +func (t *runCommandTool) Description() string { + return "Run a bash command in the workspace root (process group, default 60s timeout, 64 KiB output cap). Non-zero exit and timeout are reported as text. NOT an OS sandbox." +} +func (t *runCommandTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","required":["command"],"properties":{"command":{"type":"string"},"timeout_seconds":{"type":"integer","description":"1..300, default 60"}}}`) +} + +func (t *runCommandTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) { + var obj struct { + Command string `json:"command"` + TimeoutSeconds int `json:"timeout_seconds"` + } + if err := json.Unmarshal(args, &obj); err != nil { + return Result{}, err + } + if strings.TrimSpace(obj.Command) == "" { + return Result{Content: "command is empty", IsError: true}, nil + } + if blocked(obj.Command) { + return Result{Content: "command rejected by blocklist (sudo / destructive patterns are not allowed)", IsError: true}, nil + } + timeout := clampTimeout(obj.TimeoutSeconds) + out, timedOut, err := runProcess(ctx, t.cfg.Bash, []string{"-c", obj.Command}, t.root, t.cfg.Env, timeout) + content := formatProcessOutput(out, timedOut, err) + return Result{Content: content, IsError: timedOut || err != nil}, nil +} + +type runPythonTool struct { + root string + cfg ExecConfig +} + +func (t *runPythonTool) Name() string { return "run_python" } +func (t *runPythonTool) Description() string { + return "Run a Python script: the code is persisted under .agent/scripts/ in the workspace and executed with the configured python3 interpreter (default 60s timeout, 64 KiB output cap)." +} +func (t *runPythonTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","required":["code"],"properties":{"code":{"type":"string"},"timeout_seconds":{"type":"integer","description":"1..300, default 60"}}}`) +} + +func (t *runPythonTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) { + var obj struct { + Code string `json:"code"` + TimeoutSeconds int `json:"timeout_seconds"` + } + if err := json.Unmarshal(args, &obj); err != nil { + return Result{}, err + } + if strings.TrimSpace(obj.Code) == "" { + return Result{Content: "code is empty", IsError: true}, nil + } + dir := t.cfg.ScriptDir + if dir == "" { + dir = ".agent/scripts" + } + scriptDir := filepath.Join(t.root, filepath.FromSlash(dir)) + if err := os.MkdirAll(scriptDir, 0o755); err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + name := fmt.Sprintf("%s-%06d.py", time.Now().UTC().Format("20060102T150405"), rand.Intn(1000000)) + script := filepath.Join(scriptDir, name) + if err := os.WriteFile(script, []byte(obj.Code), 0o644); err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + timeout := clampTimeout(obj.TimeoutSeconds) + out, timedOut, err := runProcess(ctx, t.cfg.Python, []string{script}, t.root, t.cfg.Env, timeout) + content := formatProcessOutput(out, timedOut, err) + content = "script: " + filepath.ToSlash(filepath.Join(dir, name)) + "\n" + content + return Result{Content: content, IsError: timedOut || err != nil}, nil +} + +// runProcess 启动独立进程组的子进程；ctx 取消或超时先到者杀整组（SIGKILL）。 +func runProcess(ctx context.Context, interpreter string, argv []string, dir string, env []string, timeout time.Duration) (out []byte, timedOut bool, err error) { + runCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + cmd := exec.CommandContext(runCtx, interpreter, argv...) + cmd.Dir = dir + cmd.Env = env + // 独立进程组：超时/取消时杀 -PGID，确保孙进程一并退出。 + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + cmd.Cancel = func() error { + if cmd.Process != nil { + return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) + } + return nil + } + cmd.WaitDelay = 3 * time.Second + out, err = cmd.CombinedOutput() + timedOut = runCtx.Err() != nil && (ctx.Err() == nil || runCtx.Err() != ctx.Err()) + return out, timedOut, err +} + +func clampTimeout(seconds int) time.Duration { + if seconds <= 0 { + return execDefaultTimeout + } + if seconds > int(execMaxTimeout.Seconds()) { + return execMaxTimeout + } + return time.Duration(seconds) * time.Second +} + +func blocked(command string) bool { + lower := strings.ToLower(command) + for _, m := range blockedCommandMarkers { + if strings.Contains(lower, m) { + return true + } + } + return false +} + +// formatProcessOutput 截断为头尾各半（总上限 execMaxOutput）并标注退出状态。 +func formatProcessOutput(out []byte, timedOut bool, err error) string { + text := string(out) + 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:]) + } + switch { + case timedOut: + text += "\n[timeout: process group killed]" + case err != nil: + if ee, ok := err.(*exec.ExitError); ok { + text += fmt.Sprintf("\n[exit code %d]", ee.ExitCode()) + } else { + text += "\n[error] " + err.Error() + } + default: + text += "\n[exit code 0]" + } + return text +} diff --git a/runtime/tools/exec_tools_test.go b/runtime/tools/exec_tools_test.go new file mode 100644 index 0000000..399a3b2 --- /dev/null +++ b/runtime/tools/exec_tools_test.go @@ -0,0 +1,166 @@ +package tools + +import ( + "context" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" +) + +func lookInterpreter(t *testing.T, name string) string { + t.Helper() + p, err := exec.LookPath(name) + if err != nil { + t.Fatalf("%s interpreter not available on host: %v", name, err) + } + return p +} + +func TestRunCommandEchoAndExitCode(t *testing.T) { + bash := lookInterpreter(t, "bash") + root := t.TempDir() + reg := NewRegistry(root) + RegisterExecTools(reg, root, ExecConfig{Bash: bash, Env: []string{"PATH=/usr/bin:/bin"}}) + + res := runTool(t, reg, "run_command", `{"command":"echo hello"}`) + if res.IsError || !strings.Contains(res.Content, "hello") || !strings.Contains(res.Content, "[exit code 0]") { + t.Fatalf("echo failed: %+v", res) + } + + res = runTool(t, reg, "run_command", `{"command":"echo boom >&2; exit 3"}`) + if !res.IsError || !strings.Contains(res.Content, "boom") || !strings.Contains(res.Content, "[exit code 3]") { + t.Fatalf("non-zero exit must be reported as text: %+v", res) + } +} + +func TestRunCommandEnvWhitelist(t *testing.T) { + bash := lookInterpreter(t, "bash") + t.Setenv("KERNEL_SECRET_LEAK", "must-not-leak") + root := t.TempDir() + reg := NewRegistry(root) + RegisterExecTools(reg, root, ExecConfig{ + Bash: bash, + Env: []string{"PATH=/usr/bin:/bin", "ALLOWED_FLAG=yes"}, + }) + res := runTool(t, reg, "run_command", `{"command":"echo allowed=$ALLOWED_FLAG secret=${KERNEL_SECRET_LEAK:-absent}"}`) + if res.IsError { + t.Fatal(res.Content) + } + if !strings.Contains(res.Content, "allowed=yes") || !strings.Contains(res.Content, "secret=absent") { + t.Fatalf("env whitelist violated, got:\n%s", res.Content) + } +} + +func TestRunCommandTimeoutKillsProcessGroup(t *testing.T) { + bash := lookInterpreter(t, "bash") + root := t.TempDir() + reg := NewRegistry(root) + RegisterExecTools(reg, root, ExecConfig{Bash: bash, Env: []string{"PATH=/usr/bin:/bin"}}) + + start := time.Now() + res := runTool(t, reg, "run_command", `{"command":"sleep 30 & sleep 30","timeout_seconds":1}`) + elapsed := time.Since(start) + if !res.IsError || !strings.Contains(res.Content, "[timeout: process group killed]") { + t.Fatalf("want timeout kill, got %+v", res) + } + if elapsed > 10*time.Second { + t.Fatalf("process group must die promptly, took %s", elapsed) + } +} + +func TestRunCommandContextCancelKillsGroup(t *testing.T) { + bash := lookInterpreter(t, "bash") + root := t.TempDir() + reg := NewRegistry(root) + RegisterExecTools(reg, root, ExecConfig{Bash: bash, Env: []string{"PATH=/usr/bin:/bin"}}) + + ctx, cancel := context.WithCancel(context.Background()) + go func() { + time.Sleep(200 * time.Millisecond) + cancel() + }() + start := time.Now() + res := reg.Execute(ctx, "run_command", []byte(`{"command":"sleep 30","timeout_seconds":60}`)) + elapsed := time.Since(start) + if elapsed > 10*time.Second { + t.Fatalf("cancel must kill promptly, took %s", elapsed) + } + if !res.IsError { + t.Fatalf("canceled command must be an error result, got %+v", res) + } +} + +func TestRunCommandBlocklist(t *testing.T) { + bash := lookInterpreter(t, "bash") + root := t.TempDir() + reg := NewRegistry(root) + RegisterExecTools(reg, root, ExecConfig{Bash: bash, Env: []string{"PATH=/usr/bin:/bin"}}) + + for _, cmd := range []string{ + `{"command":"sudo rm x"}`, + `{"command":"rm -rf /"}`, + `{"command":"shutdown now"}`, + } { + res := runTool(t, reg, "run_command", cmd) + if !res.IsError || !strings.Contains(res.Content, "blocklist") { + t.Fatalf("blocklist must reject %s, got %+v", cmd, res) + } + } +} + +func TestRunCommandOutputTruncation(t *testing.T) { + bash := lookInterpreter(t, "bash") + root := t.TempDir() + reg := NewRegistry(root) + RegisterExecTools(reg, root, ExecConfig{Bash: bash, Env: []string{"PATH=/usr/bin:/bin"}}) + + res := runTool(t, reg, "run_command", `{"command":"head -c 200000 /dev/zero | tr '\\0' 'x'"}`) + if res.IsError { + t.Fatal(res.Content) + } + if len(res.Content) > execMaxOutput+512 { + t.Fatalf("output must be truncated to ~%d bytes, got %d", execMaxOutput, len(res.Content)) + } + if !strings.Contains(res.Content, "truncated") { + t.Fatalf("truncation marker missing, len=%d", len(res.Content)) + } +} + +func TestRunPythonPersistsScriptAndRuns(t *testing.T) { + python := lookInterpreter(t, "python3") + root := t.TempDir() + reg := NewRegistry(root) + RegisterExecTools(reg, root, ExecConfig{Python: python, Env: []string{"PATH=/usr/bin:/bin"}}) + + res := runTool(t, reg, "run_python", `{"code":"print(2 + 3)\n"}`) + if res.IsError || !strings.Contains(res.Content, "5") { + t.Fatalf("python run failed: %+v", res) + } + if !strings.Contains(res.Content, "script: .agent/scripts/") { + t.Fatalf("script path must be reported, got:\n%s", res.Content) + } + matches, _ := filepath.Glob(filepath.Join(root, ".agent", "scripts", "*.py")) + if len(matches) != 1 { + t.Fatalf("script must be persisted exactly once, got %v", matches) + } + // 脚本内容落盘可审计。 + data, _ := os.ReadFile(matches[0]) + if !strings.Contains(string(data), "print(2 + 3)") { + t.Fatalf("persisted script content wrong: %q", data) + } +} + +func TestRegisterExecToolsSkipsMissingInterpreters(t *testing.T) { + reg := NewRegistry(t.TempDir()) + RegisterExecTools(reg, t.TempDir(), ExecConfig{}) // 无解释器：两个工具都不注册 + names := make(map[string]bool) + for _, d := range reg.Decls() { + names[d.Name] = true + } + if names["run_command"] || names["run_python"] { + t.Fatalf("no interpreter configured, exec tools must not register, got %v", names) + } +} diff --git a/runtime/tools/fetch.go b/runtime/tools/fetch.go new file mode 100644 index 0000000..274bc1c --- /dev/null +++ b/runtime/tools/fetch.go @@ -0,0 +1,225 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "html" + "io" + "net" + "net/http" + "net/url" + "regexp" + "strings" + "time" +) + +// FetchConfig 声明 fetch_url 的运行条件；未注册（组合根未提供）则工具不存在。 +type FetchConfig struct { + // Client 可选注入（测试替身）；缺省构造 20s 超时客户端。 + Client *http.Client + // MaxBytes 响应体上限（默认 2 MiB）。 + MaxBytes int + // AllowedHosts 非空时仅允许这些主机（域名或 host:port）；空 = 公网放开， + // 但私网/回环/链路本地地址在任何模式下都被拒绝（SSRF 硬边界）。 + AllowedHosts []string + // Resolve 可选注入主机解析（测试替身）；缺省 net.LookupIP。 + Resolve func(host string) ([]net.IP, error) +} + +const ( + fetchTimeout = 20 * time.Second + fetchMaxBytes = 2 << 20 +) + +func (c *FetchConfig) resolver() func(string) ([]net.IP, error) { + if c.Resolve != nil { + return c.Resolve + } + return net.LookupIP +} + +// RegisterFetch 注册 fetch_url 工具（cfg 必须非 nil，由组合根显式提供）。 +func RegisterFetch(r *Registry, cfg FetchConfig) { + r.Register(&fetchTool{cfg: cfg}) +} + +type fetchTool struct { + cfg FetchConfig +} + +func (t *fetchTool) Name() string { return "fetch_url" } +func (t *fetchTool) Description() string { + return "Fetch an http(s) URL and return readable text (HTML stripped, up to 2 MiB, 20s timeout). Private/loopback network targets are always rejected." +} +func (t *fetchTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","required":["url"],"properties":{"url":{"type":"string"}}}`) +} + +func (t *fetchTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) { + var obj struct { + URL string `json:"url"` + } + if err := json.Unmarshal(args, &obj); err != nil { + return Result{}, err + } + raw := strings.TrimSpace(obj.URL) + if raw == "" { + return Result{Content: "url is empty", IsError: true}, nil + } + client := t.cfg.Client + if client == nil { + client = &http.Client{Timeout: fetchTimeout} + } + allowed := make(map[string]bool, len(t.cfg.AllowedHosts)) + for _, h := range t.cfg.AllowedHosts { + allowed[strings.ToLower(h)] = true + } + maxBytes := t.cfg.MaxBytes + if maxBytes <= 0 { + maxBytes = fetchMaxBytes + } + + reqCtx, cancel := context.WithTimeout(ctx, fetchTimeout) + defer cancel() + req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, raw, nil) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + req.Header.Set("User-Agent", "agent-native-fetch/1.0") + var finalURL string + resolve := t.cfg.resolver() + client2 := client + if client2.CheckRedirect == nil { + client2 = &http.Client{ + Transport: client.Transport, + Jar: client.Jar, + Timeout: client.Timeout, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + if len(via) >= 5 { + return fmt.Errorf("stopped after 5 redirects") + } + return validateFetchTarget(req.URL, allowed, resolve) + }, + } + } + if err := validateFetchTarget(req.URL, allowed, resolve); err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + resp, err := client2.Do(req) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + defer resp.Body.Close() + finalURL = resp.Request.URL.String() + body, err := io.ReadAll(io.LimitReader(resp.Body, int64(maxBytes)+1)) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + if len(body) > maxBytes { + body = body[:maxBytes] + } + ct := resp.Header.Get("Content-Type") + text, ok := fetchReadableText(ct, string(body)) + if !ok { + return Result{Content: fmt.Sprintf("unsupported content type %q at %s (only text/html/json/xml/yaml/csv are extracted)", ct, finalURL), IsError: true}, nil + } + return Result{Content: "url: " + finalURL + "\ncontent-type: " + ct + "\n\n" + text}, nil +} + +// validateFetchTarget 校验 scheme、主机策略与 SSRF 硬边界（私网/回环/链路本地）。 +func validateFetchTarget(u *url.URL, allowed map[string]bool, resolve func(string) ([]net.IP, error)) error { + if u.Scheme != "http" && u.Scheme != "https" { + return fmt.Errorf("scheme %q not allowed (http/https only)", u.Scheme) + } + host := u.Hostname() + if host == "" { + return fmt.Errorf("empty host") + } + if len(allowed) > 0 && !allowed[strings.ToLower(host)] { + return fmt.Errorf("host %q not in allowed hosts", host) + } + // 解析全部地址：任一命中私网段即拒绝（域名可能同时解析到公私网地址）。 + addrs, err := resolve(host) + if err != nil { + return fmt.Errorf("resolve %q: %w", host, err) + } + for _, ip := range addrs { + if ipIsPrivate(ip) { + return fmt.Errorf("target %q resolves to private/loopback address %s (SSRF guard)", host, ip) + } + } + return nil +} + +var fetchPrivateNets = mustParseCIDRs( + "0.0.0.0/8", "10.0.0.0/8", "127.0.0.0/8", "169.254.0.0/16", + "172.16.0.0/12", "192.168.0.0/16", "100.64.0.0/10", + "::1/128", "fc00::/7", "fe80::/10", +) + +func mustParseCIDRs(specs ...string) []*net.IPNet { + var out []*net.IPNet + for _, s := range specs { + _, n, err := net.ParseCIDR(s) + if err != nil { + panic("fetch: bad cidr " + s) + } + out = append(out, n) + } + return out +} + +func ipIsPrivate(ip net.IP) bool { + // IPv4（含 Go 内部的 ::ffff:a.b.c.d 映射形式）先归一化成 4 字节， + // 否则 ::ffff:0:0/96 这类映射网段会误吞所有 IPv4 地址。 + if v4 := ip.To4(); v4 != nil { + ip = v4 + } + if ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() { + return true + } + for _, n := range fetchPrivateNets { + if n.Contains(ip) { + return true + } + } + return false +} + +// fetchReadableText 按内容类型抽取可读文本；HTML 剥脚本样式与标签。 +func fetchReadableText(contentType, body string) (string, bool) { + ct := strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0])) + if ct == "" { + ct = "text/plain" // 无声明按纯文本处理 + } + switch { + case ct == "text/html": + return htmlToText(body), true + case strings.HasPrefix(ct, "text/"), + ct == "application/json", + ct == "application/xml", + strings.HasSuffix(ct, "+json"), + strings.HasSuffix(ct, "+xml"), + ct == "application/yaml" || ct == "text/yaml", + ct == "application/csv" || ct == "text/csv": + return body, true + default: + return "", false + } +} + +var ( + fetchScriptRe = regexp.MustCompile(`(?is)<(script|style|noscript)[^>]*>.*?`) + fetchTagRe = regexp.MustCompile(`(?s)<[^>]*>`) + fetchSpaceRe = regexp.MustCompile(`\n{3,}`) +) + +func htmlToText(body string) string { + s := fetchScriptRe.ReplaceAllString(body, "") + s = fetchTagRe.ReplaceAllString(s, "\n") + s = html.UnescapeString(s) + s = strings.ReplaceAll(s, "\r", "") + s = fetchSpaceRe.ReplaceAllString(s, "\n\n") + return strings.TrimSpace(s) +} diff --git a/runtime/tools/fetch_test.go b/runtime/tools/fetch_test.go new file mode 100644 index 0000000..842f8f8 --- /dev/null +++ b/runtime/tools/fetch_test.go @@ -0,0 +1,154 @@ +package tools + +import ( + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +// fakeResolve 固定域名 → 地址（无 DNS 依赖的封闭测试）。 +func fakeResolve(mapping map[string][]net.IP) func(string) ([]net.IP, error) { + return func(host string) ([]net.IP, error) { + if ips, ok := mapping[host]; ok { + return ips, nil + } + // IP 字面量直接解析。 + if ip := net.ParseIP(host); ip != nil { + return []net.IP{ip}, nil + } + return nil, &net.DNSError{Err: "no such host", Name: host} + } +} + +func fetchReg(cfg FetchConfig) *Registry { + reg := NewRegistry("/") + RegisterFetch(reg, cfg) + return reg +} + +func TestFetchReadableTextUnit(t *testing.T) { + out, ok := fetchReadableText("text/html; charset=utf-8", "

Hi

A & B

") + if !ok || !strings.Contains(out, "Hi") || !strings.Contains(out, "A & B") || strings.Contains(out, "bad()") { + t.Fatalf("html strip wrong: %q ok=%v", out, ok) + } + out, ok = fetchReadableText("application/json", `{"x":1}`) + if !ok || out != `{"x":1}` { + t.Fatalf("json passthrough wrong: %q", out) + } + _, ok = fetchReadableText("image/png", "binary") + if ok { + t.Fatal("binary content type must be rejected") + } + out, ok = fetchReadableText("", "no content type") + if !ok || out != "no content type" { + t.Fatalf("empty content type should default to text, got %q %v", out, ok) + } +} + +func TestFetchURLSSRDGuards(t *testing.T) { + reg := fetchReg(FetchConfig{Resolve: fakeResolve(nil)}) + + for _, target := range []string{ + `{"url":"http://127.0.0.1/x"}`, + `{"url":"http://10.0.0.5/x"}`, + `{"url":"http://192.168.1.1/x"}`, + `{"url":"http://169.254.169.254/latest/meta-data"}`, + `{"url":"http://[::1]/x"}`, + `{"url":"file:///etc/passwd"}`, + } { + res := runTool(t, reg, "fetch_url", target) + if !res.IsError { + t.Fatalf("SSRF/scheme guard must reject %s, got %+v", target, res) + } + } + + // 域名解析到私网同样拒绝。 + reg2 := fetchReg(FetchConfig{Resolve: fakeResolve(map[string][]net.IP{ + "internal.example": {net.ParseIP("192.168.0.10")}, + })}) + res := runTool(t, reg2, "fetch_url", `{"url":"http://internal.example/x"}`) + if !res.IsError || !strings.Contains(res.Content, "SSRF guard") { + t.Fatalf("private resolution must be rejected, got %+v", res) + } +} + +func TestFetchURLAllowedHostsPolicy(t *testing.T) { + reg := fetchReg(FetchConfig{ + AllowedHosts: []string{"api.example.com"}, + Resolve: fakeResolve(map[string][]net.IP{"api.example.com": {net.ParseIP("93.184.216.34")}}), + }) + res := runTool(t, reg, "fetch_url", `{"url":"http://other.example.com/x"}`) + if !res.IsError || !strings.Contains(res.Content, "not in allowed hosts") { + t.Fatalf("host policy must reject, got %+v", res) + } + // 允许的主机：校验必须放行（真实网络下可能成功取回，也可能因出口 + // 受限连接失败——两者都不算策略拦截）。 + res = runTool(t, reg, "fetch_url", `{"url":"http://api.example.com/x"}`) + if strings.Contains(res.Content, "not in allowed") || strings.Contains(res.Content, "SSRF") { + t.Fatalf("allowed host should pass validation, got %+v", res) + } +} + +func TestFetchURLEndToEndWithStubTransport(t *testing.T) { + var handler http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/plain") + _, _ = w.Write([]byte("fetched body")) + } + reg := fetchReg(FetchConfig{ + Client: &http.Client{Transport: stubTransport{handler: handler}}, + Resolve: fakeResolve(map[string][]net.IP{"ok.example": {net.ParseIP("93.184.216.34")}}), + }) + res := runTool(t, reg, "fetch_url", `{"url":"http://ok.example/x"}`) + if res.IsError { + t.Fatal(res.Content) + } + if !strings.Contains(res.Content, "fetched body") || !strings.Contains(res.Content, "url: http://ok.example/x") { + t.Fatalf("stub fetch wrong: %+v", res) + } +} + +type stubTransport struct{ handler http.HandlerFunc } + +func (s stubTransport) RoundTrip(req *http.Request) (*http.Response, error) { + rec := httptest.NewRecorder() + s.handler(rec, req) + resp := rec.Result() + resp.Request = req + return resp, nil +} + +func TestFetchURLRedirectToPrivateRejected(t *testing.T) { + reg := fetchReg(FetchConfig{ + Client: &http.Client{Transport: redirectTransport{to: "http://127.0.0.1/evil"}}, + Resolve: fakeResolve(map[string][]net.IP{ + "public.example": {net.ParseIP("93.184.216.34")}, + "127.0.0.1": {net.ParseIP("127.0.0.1")}, + }), + }) + res := runTool(t, reg, "fetch_url", `{"url":"http://public.example/redirect"}`) + if !res.IsError { + t.Fatalf("redirect to private target must be rejected, got %+v", res) + } +} + +type redirectTransport struct{ to string } + +func (s redirectTransport) RoundTrip(req *http.Request) (*http.Response, error) { + rec := httptest.NewRecorder() + rec.Header().Set("Location", s.to) + rec.Code = http.StatusFound + resp := rec.Result() + resp.Request = req + return resp, nil +} + +func TestFetchNotRegisteredByDefault(t *testing.T) { + reg := WorkspaceTools(t.TempDir()) + for _, d := range reg.Decls() { + if d.Name == "fetch_url" { + t.Fatal("fetch_url must not register without explicit FetchConfig") + } + } +} diff --git a/runtime/tools/glob.go b/runtime/tools/glob.go new file mode 100644 index 0000000..8c3ddcc --- /dev/null +++ b/runtime/tools/glob.go @@ -0,0 +1,169 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "regexp" + "sort" + "strings" +) + +const maxGlobResults = 500 + +type globTool struct{ root string } + +func (t *globTool) Name() string { return "glob" } + +func (t *globTool) Description() string { + return "Find workspace files matching a glob pattern (** crosses directories, * within one segment). Returns up to 500 relative paths." +} + +func (t *globTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","required":["pattern"],"properties":{"pattern":{"type":"string","description":"glob such as **/*.go or docs/*.md"}}}`) +} + +func (t *globTool) Execute(_ context.Context, args json.RawMessage) (Result, error) { + rel, err := readArg(args, "pattern") + if err != nil { + return Result{}, err + } + if strings.TrimSpace(rel) == "" { + return Result{Content: "pattern is empty", IsError: true}, nil + } + re, err := globRegexp(rel) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + var matches []string + capped := false + err = filepath.WalkDir(t.root, func(path string, d os.DirEntry, walkErr error) error { + if walkErr != nil { + return nil + } + r, rerr := filepath.Rel(t.root, path) + if rerr != nil { + return nil + } + r = filepath.ToSlash(r) + if r == "." { + return nil + } + if d.IsDir() { + // 目录命中以 / 结尾，供模型分辨文件与目录。 + if re.MatchString(r + "/") { + if len(matches) < maxGlobResults { + matches = append(matches, r+"/") + } else { + capped = true + } + } + return nil + } + if re.MatchString(r) { + if len(matches) < maxGlobResults { + matches = append(matches, r) + } else { + capped = true + } + } + return nil + }) + if err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + if len(matches) == 0 { + return Result{Content: "no matches"}, nil + } + sort.Strings(matches) + out := strings.Join(matches, "\n") + if capped { + out += fmt.Sprintf("\n(showing first %d matches)", maxGlobResults) + } + return Result{Content: out}, nil +} + +// globRegexp 把受限 glob（** 跨目录、* 单段、? 单字符，其余按字面）编译为正则。 +// 尾段单独的 ** 匹配其下全部内容（文件与目录）；中段的 ** 匹配零或多层目录。 +// 不支持字符类与分组，避免与正则语义混淆。 +func globRegexp(pattern string) (*regexp.Regexp, error) { + var b strings.Builder + b.WriteString("^") + segs := strings.Split(filepath.ToSlash(pattern), "/") + prevStarStar := false + for i, seg := range segs { + // 前一段是独立的 **（非尾段）时，其 (?:[^/]+/)* 已含目录斜杠， + // 零层目录的场景不需要再补一个分隔符。 + if i > 0 && !prevStarStar { + b.WriteString("/") + } + prevStarStar = false + if seg == "**" { + if i == len(segs)-1 { + // 尾段 **：匹配其余全部路径（文件与目录）。 + b.WriteString(".+") + } else { + b.WriteString("(?:[^/]+/)*") + prevStarStar = true + } + continue + } + for _, part := range splitDoubleStar(seg) { + if part == "**" { + b.WriteString("[^/]*") + continue + } + for _, ch := range part { + switch ch { + case '*': + b.WriteString("[^/]*") + case '?': + b.WriteString("[^/]") + default: + b.WriteString(regexp.QuoteMeta(string(ch))) + } + } + } + } + b.WriteString("$") + return regexp.Compile(b.String()) +} + +// splitDoubleStar 把段内混用的 ** 与字面部分切开（如 "a**b"）。 +func splitDoubleStar(seg string) []string { + var out []string + cur := strings.Builder{} + for i := 0; i < len(seg); i++ { + if i+1 < len(seg) && seg[i] == '*' && seg[i+1] == '*' { + if cur.Len() > 0 { + out = append(out, cur.String()) + cur.Reset() + } + out = append(out, "**") + i++ + continue + } + cur.WriteByte(seg[i]) + } + if cur.Len() > 0 { + out = append(out, cur.String()) + } + return out +} + +// globMatcher 是文件名级 glob（filepath.Match 语义），供 grep include 过滤用。 +type globMatcher struct{ pattern string } + +func newGlobMatcher(pattern string) (*globMatcher, error) { + if _, err := filepath.Match(pattern, "probe"); err != nil { + return nil, fmt.Errorf("invalid include pattern %q: %w", pattern, err) + } + return &globMatcher{pattern: pattern}, nil +} + +func (m *globMatcher) Match(name string) bool { + ok, _ := filepath.Match(m.pattern, name) + return ok +} diff --git a/runtime/tools/registry.go b/runtime/tools/registry.go index 0797477..17ad4c8 100644 --- a/runtime/tools/registry.go +++ b/runtime/tools/registry.go @@ -61,6 +61,75 @@ func (r *Registry) Decls() []Decl { return out } +// Restrict 把注册表收窄到 allow 列出的工具（工具策略的唯一执行点之一： +// 调用方声明 allowlist，注册表执行收窄）。条目支持两种形式： +// +// 精确名（read_file）或命名空间通配（mcp__github__*，匹配该前缀全部工具）。 +// +// 任一条目未命中任何已注册工具时返回错误（fail-closed，不静默忽略配置错误）。 +// 空列表不收窄；可重复调用（在已收窄集合上继续求交）。 +func (r *Registry) Restrict(allow []string) error { + if len(allow) == 0 { + return nil + } + exact := make(map[string]bool, len(allow)) + prefixes := make([]string, 0) + for _, entry := range allow { + if name, ok := strings.CutSuffix(entry, "*"); ok { + prefixes = append(prefixes, name) + continue + } + exact[entry] = true + } + var unknown []string + for _, name := range allow { + if _, ok := strings.CutSuffix(name, "*"); ok { + continue + } + if _, ok := r.tools[name]; !ok { + unknown = append(unknown, name) + } + } + for _, p := range prefixes { + if !r.hasPrefix(p) { + unknown = append(unknown, p+"*") + } + } + if len(unknown) > 0 { + return fmt.Errorf("tool allowlist names not registered: %s", strings.Join(unknown, ", ")) + } + keep := func(name string) bool { + if exact[name] { + return true + } + for _, p := range prefixes { + if strings.HasPrefix(name, p) { + return true + } + } + return false + } + var order []string + for _, name := range r.order { + if keep(name) { + order = append(order, name) + } else { + delete(r.tools, name) + } + } + r.order = order + return nil +} + +func (r *Registry) hasPrefix(prefix string) bool { + for _, name := range r.order { + if strings.HasPrefix(name, prefix) { + return true + } + } + return false +} + // Decl is the schema view of a tool. type Decl struct { Name string @@ -182,8 +251,10 @@ func WorkspaceTools(root string) *Registry { r.Register(&readFile{root: root}) r.Register(&writeFile{root: root}) r.Register(&editFile{root: root}) + r.Register(&multiEditTool{root: root}) r.Register(&listDir{root: root}) r.Register(&grepTool{root: root}) + r.Register(&globTool{root: root}) return r } diff --git a/runtime/tools/registry_policy_test.go b/runtime/tools/registry_policy_test.go new file mode 100644 index 0000000..692a683 --- /dev/null +++ b/runtime/tools/registry_policy_test.go @@ -0,0 +1,71 @@ +package tools + +import ( + "context" + "strings" + "testing" +) + +func TestRegistryRestrictNarrowsDeclsAndExecute(t *testing.T) { + dir := t.TempDir() + reg := WorkspaceTools(dir) + if err := reg.Restrict([]string{"read_file", "list_dir"}); err != nil { + t.Fatal(err) + } + names := make(map[string]bool) + for _, d := range reg.Decls() { + names[d.Name] = true + } + if len(names) != 2 || !names["read_file"] || !names["list_dir"] { + t.Fatalf("want read_file+list_dir only, got %v", names) + } + // 被收窄掉的工具体现为 unknown tool。 + res := reg.Execute(context.Background(), "write_file", []byte(`{"path":"a.md","content":"x"}`)) + if !res.IsError || !strings.Contains(res.Content, "unknown tool") { + t.Fatalf("want unknown tool error, got %+v", res) + } + // 保留的工具可用。 + res = reg.Execute(context.Background(), "list_dir", []byte(`{}`)) + if res.IsError { + t.Fatalf("list_dir should still work: %+v", res) + } +} + +func TestRegistryRestrictFailsClosedOnUnknownNames(t *testing.T) { + reg := WorkspaceTools(t.TempDir()) + err := reg.Restrict([]string{"read_file", "not_a_tool"}) + if err == nil || !strings.Contains(err.Error(), "not_a_tool") { + t.Fatalf("want fail-closed error naming the unknown tool, got %v", err) + } + // 失败后注册表不被部分修改。 + if len(reg.Decls()) != 7 { + t.Fatalf("registry should be untouched after failed restrict, got %d tools", len(reg.Decls())) + } +} + +func TestRegistryRestrictEmptyKeepsAll(t *testing.T) { + reg := WorkspaceTools(t.TempDir()) + if err := reg.Restrict(nil); err != nil { + t.Fatal(err) + } + if len(reg.Decls()) != 7 { + t.Fatalf("nil allowlist should keep all tools, got %d", len(reg.Decls())) + } +} + +func TestRegistryRestrictIsIdempotentIntersection(t *testing.T) { + reg := WorkspaceTools(t.TempDir()) + if err := reg.Restrict([]string{"read_file", "list_dir"}); err != nil { + t.Fatal(err) + } + // 二次收窄若点名已被移除的工具，属调用方配置错误：fail-closed。 + if err := reg.Restrict([]string{"read_file", "write_file"}); err == nil || !strings.Contains(err.Error(), "write_file") { + t.Fatalf("second restrict naming a removed tool must fail-closed, got %v", err) + } + if err := reg.Restrict([]string{"read_file"}); err != nil { + t.Fatal(err) + } + if len(reg.Decls()) != 1 || reg.Decls()[0].Name != "read_file" { + t.Fatalf("want intersection to keep read_file only, got %+v", reg.Decls()) + } +}