From e4020d297c168347626f9fd75456a71374a0a6c8 Mon Sep 17 00:00:00 2001 From: weichen Date: Sat, 3 Oct 2026 18:02:46 +0000 Subject: [PATCH] Add a script host, deferred tools, and background MCP connect Why: models need a capped way to call several tools from one script, and a slow MCP server must not block the first prompt. Evidence: GOWORK=off go test ./runtime/codeorch/ ./runtime/tools/ ./runtime/mcp/ ./runtime/policy/ ./adapters/native/ Impact: run_tool_script stays unregistered until the caller turns it on. Deferred tools stay out of the model list. An unready MCP server connects in the background. Refs: feature/code-orchestration Co-authored-by: Cursor --- adapters/native/backend.go | 21 +- adapters/native/exec_tools_test.go | 3 + runtime/codeorch/agentwf.py | 51 ++++ runtime/codeorch/forbid.go | 12 + runtime/codeorch/host.go | 380 +++++++++++++++++++++++++++++ runtime/codeorch/host_test.go | 89 +++++++ runtime/mcp/ask.go | 27 ++ runtime/mcp/manager.go | 192 +++++++++++---- runtime/mcp/mcp_test.go | 113 +++++++++ runtime/mcp/protocol.go | 14 ++ runtime/mcp/stdio.go | 5 + runtime/policy/policy.go | 4 +- runtime/tools/code_tool.go | 59 +++++ runtime/tools/deferred.go | 164 +++++++++++++ runtime/tools/deferred_test.go | 75 ++++++ runtime/tools/registry.go | 78 +++++- 16 files changed, 1223 insertions(+), 64 deletions(-) create mode 100644 runtime/codeorch/agentwf.py create mode 100644 runtime/codeorch/forbid.go create mode 100644 runtime/codeorch/host.go create mode 100644 runtime/codeorch/host_test.go create mode 100644 runtime/mcp/ask.go create mode 100644 runtime/tools/code_tool.go create mode 100644 runtime/tools/deferred.go create mode 100644 runtime/tools/deferred_test.go diff --git a/adapters/native/backend.go b/adapters/native/backend.go index d5aeab0..8999e66 100644 --- a/adapters/native/backend.go +++ b/adapters/native/backend.go @@ -49,6 +49,8 @@ type ToolOptions struct { HostTools []tools.Tool Checkpoints *checkpoint.Log CheckpointAnchor func() string + // CodeOrchestration 为真时注册 run_tool_script。默认关。 + CodeOrchestration bool } // SessionOptions 是单次执行的会话预算与压缩。 @@ -73,6 +75,8 @@ type GateOptions struct { Reminders func(turn int) []string RepeatGuard bool Hooks *hooks.Bus + // Ask 在 MCP 工具要求用户输入时调用。nil 则把问题作为错误返回。 + Ask func(ctx context.Context, question string) (string, error) } // Options 构造后端所需的模型接入与执行参数（由组合根注入，无产品依赖）。 @@ -147,6 +151,9 @@ 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 b.opts.Ask != nil { + ctx = mcp.WithAsk(ctx, b.opts.Ask) + } model := strings.TrimSpace(b.opts.Model) if strings.TrimSpace(req.Model) != "" { model = strings.TrimSpace(req.Model) @@ -204,9 +211,19 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result, } tools.RegisterExecTools(reg, workDir, execCfg) } + if b.opts.CodeOrchestration { + tools.RegisterCodeTool(reg) + } + reg.SetActivationFile(filepath.Join(workDir, ".agentx", "activated-tools.json")) + _ = reg.RestoreActivated() if b.opts.MCP != nil { - if err := b.opts.MCP.RegisterInto(ctx, reg); err != nil { - return &exec.Result{Success: false, ErrorReason: err.Error()}, nil + switch { + case len(req.ToolAllowlist) > 0 || b.opts.MCP.Ready(): + if err := b.opts.MCP.RegisterInto(ctx, reg); err != nil { + return &exec.Result{Success: false, ErrorReason: err.Error()}, nil + } + default: + b.opts.MCP.StartBackground(ctx, reg) } } subagent.Register(reg, b.opts.Subagent) diff --git a/adapters/native/exec_tools_test.go b/adapters/native/exec_tools_test.go index 9cf7bc5..9536072 100644 --- a/adapters/native/exec_tools_test.go +++ b/adapters/native/exec_tools_test.go @@ -212,6 +212,9 @@ func TestBackendMCPToolsWiredViaOptions(t *testing.T) { t.Fatal(err) } defer mgr.Close() + if err := mgr.RegisterInto(context.Background(), tools.NewRegistry(t.TempDir())); err != nil { + t.Fatal(err) + } stream, c := captureStream( llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ ID: "c1", Name: "mcp__fake__echo", Arguments: `{"text":"via-backend"}`, diff --git a/runtime/codeorch/agentwf.py b/runtime/codeorch/agentwf.py new file mode 100644 index 0000000..d31d477 --- /dev/null +++ b/runtime/codeorch/agentwf.py @@ -0,0 +1,51 @@ +"""工作流脚本回调。用法：import agentwf as wf。 + +请求写到 fd 3，响应从 fd 4 按行读回。宿主一次只处理一条。 +""" + +import json +import os + +_seq = 0 +_out = os.fdopen(3, "w", encoding="utf-8", buffering=1) +_inp = os.fdopen(4, "r", encoding="utf-8", buffering=1) + + +def _rpc(op, **fields): + global _seq + _seq += 1 + msg = {"id": str(_seq), "op": op} + msg.update(fields) + _out.write(json.dumps(msg, ensure_ascii=False) + "\n") + _out.flush() + line = _inp.readline() + if not line: + raise RuntimeError("workflow host closed") + resp = json.loads(line) + if not resp.get("ok"): + raise RuntimeError(resp.get("error") or "workflow call failed") + return resp.get("output", "") + + +def tool(name, args=None): + return _rpc("tool.call", name=name, args=args or {}) + + +def ask(question, choices=None): + return _rpc("hitl.ask", question=question, choices=choices or []) + + +def agent(name, task): + return _rpc("agent.spawn", agent=name, task=task) + + +def artifact(name, content): + return _rpc("artifact.register", name=name, content=content) + + +def card(title, kind="result"): + return _rpc("card.create", title=title, kind=kind) + + +def step(step_id, title, status): + return _rpc("step", stepId=step_id, title=title, status=status) diff --git a/runtime/codeorch/forbid.go b/runtime/codeorch/forbid.go new file mode 100644 index 0000000..da7c739 --- /dev/null +++ b/runtime/codeorch/forbid.go @@ -0,0 +1,12 @@ +package codeorch + +// Forbidden 是脚本不能再调用的工具：解释器、自身、派发和工作流。 +func Forbidden(name string) bool { + switch name { + case "run_python", "run_shell", "run_command", "run_tool_script", + "delegate_task", "squad_orchestrate", "run_workflow", "review_findings": + return true + default: + return false + } +} diff --git a/runtime/codeorch/host.go b/runtime/codeorch/host.go new file mode 100644 index 0000000..c2deb7c --- /dev/null +++ b/runtime/codeorch/host.go @@ -0,0 +1,380 @@ +// Package codeorch 在会话目录里跑一段 Python，并按行把工具调用交回宿主。 +// 工作流脚本和 run_tool_script 共用这一套进程与限额。 +package codeorch + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "time" + "unicode/utf8" + + _ "embed" +) + +//go:embed agentwf.py +var agentwfSource string + +const ( + defaultTimeout = 120 * time.Second + defaultMaxWait = 600 * time.Second + defaultCalls = 50 + defaultResult = 1 << 20 + defaultModel = 16000 +) + +var ( + helperOnce sync.Once + helperDir string + helperErr error +) + +// Call 是脚本发来的一行请求。Args 与 Choices 保持原始 JSON。 +type Call struct { + ID string `json:"id"` + Op string `json:"op"` + Name string `json:"name"` + Args json.RawMessage `json:"args"` + Question string `json:"question"` + Choices json.RawMessage `json:"choices"` + Agent string `json:"agent"` + Task string `json:"task"` + Content string `json:"content"` + Title string `json:"title"` + Kind string `json:"kind"` + Status string `json:"status"` + StepID string `json:"stepId"` +} + +// Limits 控制一次脚本运行。零值使用工具路径的缺省：120 秒、最长 600 秒、50 次、1 MiB、16k 字符。 +type Limits struct { + Timeout time.Duration + MaxTimeout time.Duration + MaxCalls int + MaxResultBytes int + MaxModelChars int + FullDir string +} + +// Outcome 是给模型的文本、已成功的调用，以及超限时的全文路径。 +type Outcome struct { + Output string + FullPath string + Calls []string + Err error +} + +// Starter 按工作目录和脚本路径准备命令。空则用 python3。 +type Starter func(ctx context.Context, workDir, path string) (*exec.Cmd, error) + +// Dispatch 处理一条回调。返回的错误会写回脚本并记入最终结果。 +type Dispatch func(ctx context.Context, call Call) (string, error) + +// ModelText 把截断后的输出、已执行调用和错误拼成给模型的文本。 +func (o Outcome) ModelText() string { + text := o.Output + if o.Err != nil { + msg := o.Err.Error() + if len(o.Calls) > 0 { + msg = "已执行：" + strings.Join(o.Calls, "、") + "\n" + msg + } + if text != "" { + text += "\n" + } + text += msg + } + if o.FullPath != "" { + if text != "" { + text += "\n" + } + text += "全文在 " + o.FullPath + } + return text +} + +// RunFile 执行已经落盘的脚本。 +func RunFile(ctx context.Context, workDir, path string, lim Limits, start Starter, dispatch Dispatch) Outcome { + if ctx == nil { + ctx = context.Background() + } + timeout := lim.bounded() + runCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + helper, err := helperPath() + if err != nil { + return Outcome{Err: err} + } + if start == nil { + start = defaultStart + } + cmd, err := start(runCtx, workDir, path) + if err != nil { + return Outcome{Err: err} + } + cmd.Env = prependPythonPath(cmd.Env, helper) + reqR, reqW, err := os.Pipe() + if err != nil { + return Outcome{Err: err} + } + respR, respW, err := os.Pipe() + if err != nil { + reqR.Close() + reqW.Close() + return Outcome{Err: err} + } + cmd.ExtraFiles = []*os.File{reqW, respR} + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + if err := cmd.Start(); err != nil { + reqR.Close() + reqW.Close() + respR.Close() + respW.Close() + return Outcome{Err: err} + } + reqW.Close() + respR.Close() + defer reqR.Close() + defer respW.Close() + calls, callErr := serve(runCtx, reqR, respW, lim, dispatch) + if callErr != nil { + cancel() + } + waitErr := cmd.Wait() + out := Outcome{Output: clipJoined(&stdout, &stderr), Calls: calls} + out = clipModel(out, lim) + switch { + case callErr != nil: + out.Err = callErr + case runCtx.Err() == context.DeadlineExceeded: + out.Err = fmt.Errorf("脚本超时") + case runCtx.Err() == context.Canceled: + out.Err = fmt.Errorf("脚本已取消") + case waitErr != nil: + out.Err = fmt.Errorf("脚本退出：%v", waitErr) + } + return out +} + +// Run 把脚本体写入工作目录后再执行。 +func Run(ctx context.Context, workDir, body string, lim Limits, start Starter, dispatch Dispatch) Outcome { + dir := filepath.Join(workDir, ".agent") + if err := os.MkdirAll(dir, 0o755); err != nil { + return Outcome{Err: err} + } + path := filepath.Join(dir, fmt.Sprintf("orch-%d.py", time.Now().UnixNano())) + if err := os.WriteFile(path, []byte(body), 0o600); err != nil { + return Outcome{Err: err} + } + return RunFile(ctx, workDir, path, lim, start, dispatch) +} + +func (l Limits) bounded() time.Duration { + d := l.Timeout + if d <= 0 { + d = defaultTimeout + } + max := l.MaxTimeout + if max <= 0 { + max = defaultMaxWait + } + if d > max { + return max + } + return d +} + +func (l Limits) calls() int { + if l.MaxCalls <= 0 { + return defaultCalls + } + return l.MaxCalls +} + +func (l Limits) resultBytes() int { + if l.MaxResultBytes <= 0 { + return defaultResult + } + return l.MaxResultBytes +} + +func (l Limits) modelChars() int { + if l.MaxModelChars <= 0 { + return defaultModel + } + return l.MaxModelChars +} + +func serve(ctx context.Context, reqR, respW *os.File, lim Limits, dispatch Dispatch) ([]string, error) { + sc := bufio.NewScanner(reqR) + sc.Buffer(make([]byte, 64*1024), lim.resultBytes()) + var calls []string + n := 0 + for sc.Scan() { + line := strings.TrimSpace(sc.Text()) + if line == "" { + continue + } + n++ + var call Call + if err := json.Unmarshal([]byte(line), &call); err != nil { + _ = writeResp(respW, call.ID, false, "", "请求不是合法 JSON") + continue + } + if n > lim.calls() { + _ = writeResp(respW, call.ID, false, "", "脚本回调超过 50 次") + return calls, fmt.Errorf("脚本回调超过 50 次") + } + if ctx.Err() != nil { + return calls, ctx.Err() + } + if dispatch == nil { + _ = writeResp(respW, call.ID, false, "", "没有回调处理") + continue + } + output, err := dispatch(ctx, call) + if err != nil { + _ = writeResp(respW, call.ID, false, "", err.Error()) + if ctx.Err() != nil { + return calls, ctx.Err() + } + continue + } + if call.Op == "tool.call" && call.Name != "" { + calls = append(calls, call.Name) + } + if err := writeResp(respW, call.ID, true, clipBytes(output, lim.resultBytes()), ""); err != nil { + return calls, err + } + } + if err := sc.Err(); err != nil { + if err == bufio.ErrTooLong { + return calls, fmt.Errorf("单行请求超过 1MB") + } + if ctx.Err() != nil { + return calls, ctx.Err() + } + return calls, err + } + return calls, nil +} + +func writeResp(w io.Writer, id string, ok bool, output, errText string) error { + raw, err := json.Marshal(struct { + ID string `json:"id"` + OK bool `json:"ok"` + Output string `json:"output,omitempty"` + Error string `json:"error,omitempty"` + }{ID: id, OK: ok, Output: output, Error: errText}) + if err != nil { + return err + } + _, err = w.Write(append(raw, '\n')) + return err +} + +func clipBytes(s string, n int) string { + if n <= 0 || len(s) <= n { + return s + } + cut := s[:n] + for !utf8.ValidString(cut) && len(cut) > 0 { + cut = cut[:len(cut)-1] + } + return cut +} + +func clipJoined(stdout, stderr *bytes.Buffer) string { + text := strings.TrimSpace(stdout.String()) + errText := strings.TrimSpace(stderr.String()) + if errText != "" { + if text != "" { + text += "\n" + } + text += errText + } + return text +} + +func clipModel(out Outcome, lim Limits) Outcome { + limit := lim.modelChars() + if utf8.RuneCountInString(out.Output) <= limit { + return out + } + dir := lim.FullDir + if dir == "" { + dir = os.TempDir() + } + if err := os.MkdirAll(dir, 0o755); err != nil { + out.Output = clipRunes(out.Output, limit) + return out + } + path := filepath.Join(dir, fmt.Sprintf("script-output-%d.txt", time.Now().UnixNano())) + if err := os.WriteFile(path, []byte(out.Output), 0o644); err != nil { + out.Output = clipRunes(out.Output, limit) + return out + } + out.FullPath = path + out.Output = clipRunes(out.Output, limit) + return out +} + +func clipRunes(s string, n int) string { + if utf8.RuneCountInString(s) <= n { + return s + } + i := 0 + for count := 0; count < n && i < len(s); count++ { + _, size := utf8.DecodeRuneInString(s[i:]) + i += size + } + return s[:i] + "…（超长截断）" +} + +func defaultStart(ctx context.Context, workDir, path string) (*exec.Cmd, error) { + bin, err := exec.LookPath("python3") + if err != nil { + return nil, fmt.Errorf("找不到 python3: %w", err) + } + cmd := exec.CommandContext(ctx, bin, path) + cmd.Dir = workDir + return cmd, nil +} + +func helperPath() (string, error) { + helperOnce.Do(func() { + dir, err := os.MkdirTemp("", "agentwf-*") + if err != nil { + helperErr = err + return + } + if err := os.WriteFile(filepath.Join(dir, "agentwf.py"), []byte(agentwfSource), 0o644); err != nil { + helperErr = err + return + } + helperDir = dir + }) + return helperDir, helperErr +} + +func prependPythonPath(env []string, dir string) []string { + if env == nil { + env = os.Environ() + } + for i, kv := range env { + if strings.HasPrefix(kv, "PYTHONPATH=") { + env[i] = "PYTHONPATH=" + dir + string(os.PathListSeparator) + strings.TrimPrefix(kv, "PYTHONPATH=") + return env + } + } + return append(env, "PYTHONPATH="+dir) +} diff --git a/runtime/codeorch/host_test.go b/runtime/codeorch/host_test.go new file mode 100644 index 0000000..12071f3 --- /dev/null +++ b/runtime/codeorch/host_test.go @@ -0,0 +1,89 @@ +package codeorch + +import ( + "context" + "errors" + "os" + "strings" + "testing" + "time" + "unicode/utf8" +) + +func TestScriptCallsTwoToolsAndReportsTheOneThatRan(t *testing.T) { + body := "import agentwf as wf\nprint(wf.tool('alpha', {}))\nprint(wf.tool('beta', {}))\n" + out := Run(context.Background(), t.TempDir(), body, Limits{}, nil, func(_ context.Context, call Call) (string, error) { + if call.Name == "beta" { + return "", errors.New("beta 失败") + } + return "A", nil + }) + text := out.ModelText() + if len(out.Calls) != 1 || out.Calls[0] != "alpha" { + t.Fatalf("calls %+v", out.Calls) + } + if !strings.Contains(text, "已执行：alpha") || !strings.Contains(text, "beta 失败") { + t.Fatalf("text %q err %v", text, out.Err) + } +} + +func TestForbiddenToolIsRejectedBeforeDispatchRunsIt(t *testing.T) { + body := "import agentwf as wf\nprint(wf.tool('read_file', {}))\nprint(wf.tool('run_shell', {}))\n" + var ran []string + out := Run(context.Background(), t.TempDir(), body, Limits{}, nil, func(_ context.Context, call Call) (string, error) { + if Forbidden(call.Name) { + return "", errors.New("不能调用 " + call.Name) + } + ran = append(ran, call.Name) + return "ok", nil + }) + if len(ran) != 1 || ran[0] != "read_file" { + t.Fatalf("ran %+v", ran) + } + if !strings.Contains(out.ModelText(), "不能调用 run_shell") { + t.Fatalf("text %q", out.ModelText()) + } +} + +func TestLongOutputIsClippedAndStored(t *testing.T) { + body := "import agentwf as wf\nprint(wf.tool('big', {}))\n" + dir := t.TempDir() + out := Run(context.Background(), dir, body, Limits{MaxModelChars: 8, FullDir: dir}, nil, func(_ context.Context, call Call) (string, error) { + return strings.Repeat("测", 20), nil + }) + if out.FullPath == "" { + t.Fatal("expected a full-text file") + } + full, err := os.ReadFile(out.FullPath) + if err != nil { + t.Fatal(err) + } + if utf8.RuneCount(full) < 20 { + t.Fatalf("full text too short: %q", full) + } + if !strings.HasSuffix(out.Output, "…（超长截断）") || utf8.RuneCountInString(strings.TrimSuffix(out.Output, "…（超长截断）")) != 8 { + t.Fatalf("model output not clipped: %q", out.Output) + } + if !strings.Contains(out.ModelText(), out.FullPath) { + t.Fatalf("model text missing path: %q", out.ModelText()) + } +} + +func TestResultGivenToScriptIsCapped(t *testing.T) { + got := clipBytes(strings.Repeat("x", 20), 8) + if got != "xxxxxxxx" { + t.Fatalf("clip %q", got) + } + if !utf8.ValidString(clipBytes("测测测", 4)) { + t.Fatal("clip broke utf-8") + } +} + +func TestTimeoutIsCapped(t *testing.T) { + if (Limits{Timeout: 900 * time.Second, MaxTimeout: 600 * time.Second}).bounded() != 600*time.Second { + t.Fatal("timeout must clamp to the max") + } + if (Limits{}).bounded() != defaultTimeout { + t.Fatal("empty limits must use the 120s default") + } +} diff --git a/runtime/mcp/ask.go b/runtime/mcp/ask.go new file mode 100644 index 0000000..308d901 --- /dev/null +++ b/runtime/mcp/ask.go @@ -0,0 +1,27 @@ +package mcp + +import "context" + +type askKey struct{} + +// WithAsk 把当前回合的 AskUser 放进上下文。MCP 返回 input_required 时使用它。 +func WithAsk(ctx context.Context, fn func(context.Context, string) (string, error)) context.Context { + return context.WithValue(ctx, askKey{}, fn) +} + +func askFrom(ctx context.Context) func(context.Context, string) (string, error) { + if ctx == nil { + return nil + } + fn, _ := ctx.Value(askKey{}).(func(context.Context, string) (string, error)) + return fn +} + +func inputQuestion(res callToolResult) (string, bool) { + for _, blk := range res.Content { + if blk.Type == "input_required" { + return blk.Text, true + } + } + return "", false +} diff --git a/runtime/mcp/manager.go b/runtime/mcp/manager.go index c5eedb6..e25f98c 100644 --- a/runtime/mcp/manager.go +++ b/runtime/mcp/manager.go @@ -19,6 +19,8 @@ const ( maxToolNameBytes = 128 callTimeoutDefault = 30 * time.Second registerTimeoutFixed = 30 * time.Second + // 工具数超过这个值且没有显式暴露级别时，整台服务器改为延迟暴露。 + deferredToolThreshold = 8 ) // Manager 维护一组 MCP server 连接，并把它们的工具注册进 tools.Registry。 @@ -72,55 +74,138 @@ func (m *Manager) client(ctx context.Context, cfg ServerConfig) (Client, error) 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 { +// ApplyOverrides 用项目配置覆盖启用状态和暴露级别。空字符串不改。 +func (m *Manager) ApplyOverrides(raw string) error { + if m == nil || strings.TrimSpace(raw) == "" { + return nil + } + var over map[string]struct { + Enabled *bool `json:"enabled"` + Exposure string `json:"exposure"` + } + if err := json.Unmarshal([]byte(raw), &over); err != nil { + return err + } + for i, cfg := range m.configs { + o, ok := over[cfg.Name] + if !ok { + continue + } + if o.Enabled != nil { + m.configs[i].Enabled = o.Enabled + } + if o.Exposure != "" { + m.configs[i].Exposure = o.Exposure + } + if err := m.configs[i].Validate(); 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 = tools.CutUTF8(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, - redial: func(ctx context.Context) (Client, error) { - m.mu.Lock() - delete(m.clients, cfg.Name) - m.mu.Unlock() - return m.client(ctx, cfg) - }, - }) - } } return nil } +// Ready 报告已启用的服务器是否都已经连上。已连上时注册只是枚举工具，不再握手。 +func (m *Manager) Ready() bool { + if m == nil { + return true + } + m.mu.Lock() + defer m.mu.Unlock() + for _, cfg := range m.configs { + if !cfg.IsEnabled() { + continue + } + if _, ok := m.clients[cfg.Name]; !ok { + return false + } + } + return true +} +func (m *Manager) StartBackground(ctx context.Context, reg *tools.Registry) { + if m == nil || reg == nil { + return + } + for _, cfg := range m.configs { + cfg := cfg + if !cfg.IsEnabled() { + continue + } + go func() { _ = m.registerServer(ctx, reg, cfg) }() + } +} + +// RegisterInto 连接全部已启用的 server 并注册工具。未启用的跳过。 +func (m *Manager) RegisterInto(ctx context.Context, reg *tools.Registry) error { + for _, cfg := range m.configs { + if !cfg.IsEnabled() { + continue + } + if err := m.registerServer(ctx, reg, cfg); err != nil { + return err + } + } + return nil +} + +func (m *Manager) registerServer(ctx context.Context, reg *tools.Registry, cfg ServerConfig) error { + 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) + } + exposure := chooseExposure(cfg.Exposure, len(list)) + 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 = tools.CutUTF8(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, + exposure: exposure, + redial: func(ctx context.Context) (Client, error) { + m.mu.Lock() + delete(m.clients, cfg.Name) + m.mu.Unlock() + return m.client(ctx, cfg) + }, + }) + } + return nil +} + +func chooseExposure(explicit string, n int) string { + switch explicit { + case tools.ExposureDirect, tools.ExposureDeferred: + return explicit + default: + if n > deferredToolThreshold { + return tools.ExposureDeferred + } + return tools.ExposureDirect + } +} + // ToolName 生成桥接工具名（命名空间隔离）。 func ToolName(server, tool string) string { return "mcp__" + server + "__" + tool @@ -144,14 +229,14 @@ type bridgeTool struct { declName string description string schema json.RawMessage + exposure string redial func(ctx context.Context) (Client, error) } -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) 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) Exposure() string { return t.exposure } func (t *bridgeTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) { res, err := t.call(ctx, args) @@ -166,6 +251,17 @@ func (t *bridgeTool) Execute(ctx context.Context, args json.RawMessage) (tools.R if err != nil { return tools.Result{Content: fmt.Sprintf("mcp tool %s failed: %v", t.declName, err), IsError: true}, nil } + if question, ok := inputQuestion(res); ok { + ask := askFrom(ctx) + if ask == nil { + return tools.Result{Content: "需要用户输入：" + question, IsError: true}, nil + } + answer, aerr := ask(ctx, question) + if aerr != nil { + return tools.Result{Content: aerr.Error(), IsError: true}, nil + } + return tools.Result{Content: answer}, nil + } var b strings.Builder for _, blk := range res.Content { if blk.Type == "text" && blk.Text != "" { diff --git a/runtime/mcp/mcp_test.go b/runtime/mcp/mcp_test.go index fd3a5ce..8ce52e2 100644 --- a/runtime/mcp/mcp_test.go +++ b/runtime/mcp/mcp_test.go @@ -99,6 +99,9 @@ func TestStdioSlowCancelThenEcho(t *testing.T) { if len(res.Content) != 1 || res.Content[0].Text != "echo: after-slow" { t.Fatalf("echo must receive its own result, got %+v", res) } + if c.Readers() != 1 { + t.Fatalf("cancel must leave a single reader, got %d", c.Readers()) + } } func fakeHTTPHandler(w http.ResponseWriter, r *http.Request) { @@ -265,3 +268,113 @@ func TestManagerRegisterTimeoutBounded(t *testing.T) { t.Fatalf("timeout must bound the dial, took %s", elapsed) } } + +func TestBackgroundConnectDoesNotBlockTheCaller(t *testing.T) { + started := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-started: + default: + close(started) + } + time.Sleep(400 * time.Millisecond) + fakeHTTPHandler(w, r) + })) + defer srv.Close() + mgr, err := NewManager([]ServerConfig{{Name: "late", URL: srv.URL}}) + if err != nil { + t.Fatal(err) + } + defer mgr.Close() + reg := tools.NewRegistry(t.TempDir()) + t0 := time.Now() + mgr.StartBackground(context.Background(), reg) + if time.Since(t0) > 150*time.Millisecond { + t.Fatalf("background connect blocked the caller for %s", time.Since(t0)) + } + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if reg.Has("mcp__late__echo") { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Fatal("tools did not appear after the server became ready") +} + +func TestDeferredMCPToolsUseSearchAndCall(t *testing.T) { + mgr, err := NewManager([]ServerConfig{{Name: "fake", Command: fakeStdioConfig(t).Command, Args: fakeStdioConfig(t).Args, Env: fakeStdioConfig(t).Env, Exposure: "deferred"}}) + 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) + } + for _, d := range reg.Decls() { + if strings.HasPrefix(d.Name, "mcp__fake__") { + t.Fatalf("deferred mcp tool listed: %s", d.Name) + } + } + if !strings.Contains(reg.Search("echo"), "mcp__fake__echo") { + t.Fatalf("search %q", reg.Search("echo")) + } + res := reg.Execute(context.Background(), "call_tool", json.RawMessage(`{"name":"mcp__fake__echo","arguments":{"text":"via"}}`)) + if res.IsError || !strings.Contains(res.Content, "echo: via") { + t.Fatalf("call %+v", res) + } +} + +func TestDisabledServerIsSkipped(t *testing.T) { + off := false + mgr, err := NewManager([]ServerConfig{{Name: "fake", Command: "x", Enabled: &off}}) + if err != nil { + t.Fatal(err) + } + reg := tools.NewRegistry(t.TempDir()) + if err := mgr.RegisterInto(context.Background(), reg); err != nil { + t.Fatal(err) + } + if reg.Has("mcp__fake__echo") { + t.Fatal("disabled server registered a tool") + } + if err := mgr.ApplyOverrides(`{"fake":{"enabled":true,"exposure":"direct"}}`); err != nil { + t.Fatal(err) + } + if !mgr.Configs()[0].IsEnabled() || mgr.Configs()[0].Exposure != "direct" { + t.Fatalf("override %+v", mgr.Configs()[0]) + } +} + +func TestManyToolsDefaultToDeferred(t *testing.T) { + if chooseExposure("", 9) != tools.ExposureDeferred { + t.Fatal("many tools should default to deferred") + } + if chooseExposure("direct", 9) != tools.ExposureDirect { + t.Fatal("explicit direct must win") + } +} + +type askClient struct{} + +func (askClient) Initialize(context.Context) (initializeResult, error) { + return initializeResult{}, nil +} +func (askClient) ListTools(context.Context) ([]toolInfo, error) { return nil, nil } +func (askClient) Close() error { return nil } +func (askClient) CallTool(context.Context, string, json.RawMessage) (callToolResult, error) { + return callToolResult{Content: []contentBlock{{Type: "input_required", Text: "选一个颜色"}}}, nil +} + +func TestInputRequiredUsesAsk(t *testing.T) { + tool := &bridgeTool{ + client: askClient{}, declName: "mcp__x__ask", toolName: "ask", + schema: json.RawMessage(`{"type":"object","properties":{}}`), + } + ctx := WithAsk(context.Background(), func(context.Context, string) (string, error) { return "蓝色", nil }) + res, err := tool.Execute(ctx, json.RawMessage(`{}`)) + if err != nil || res.IsError || res.Content != "蓝色" { + t.Fatalf("ask result %+v %v", res, err) + } +} diff --git a/runtime/mcp/protocol.go b/runtime/mcp/protocol.go index c8474c0..bd4ea15 100644 --- a/runtime/mcp/protocol.go +++ b/runtime/mcp/protocol.go @@ -25,6 +25,10 @@ type ServerConfig struct { Headers map[string]string `json:"headers,omitempty"` // Timeout 单次请求上限（默认 30s）。 Timeout time.Duration `json:"timeout,omitempty"` + // Enabled 为 false 时跳过这台服务器。nil 表示启用。 + Enabled *bool `json:"enabled,omitempty"` + // Exposure 为 direct 或 deferred。空则工具多时延迟暴露。 + Exposure string `json:"exposure,omitempty"` } // Validate 校验配置完整性（fail-closed：命名不规范或双传输/无传输都拒绝）。 @@ -37,9 +41,19 @@ func (c ServerConfig) Validate() error { if hasCmd == hasURL { return fmt.Errorf("mcp server %q: exactly one of command or url is required", c.Name) } + switch c.Exposure { + case "", "direct", "deferred": + default: + return fmt.Errorf("mcp server %q: exposure must be direct or deferred", c.Name) + } return nil } +// IsEnabled 报告这台服务器是否要连接。 +func (c ServerConfig) IsEnabled() bool { + return c.Enabled == nil || *c.Enabled +} + func validName(s string) bool { if s == "" { return false diff --git a/runtime/mcp/stdio.go b/runtime/mcp/stdio.go index 1eeabea..2da8d2c 100644 --- a/runtime/mcp/stdio.go +++ b/runtime/mcp/stdio.go @@ -22,6 +22,7 @@ type stdioClient struct { mu sync.Mutex pending map[int64]chan stdioResult nextID int64 + readers atomic.Int32 closed bool readerOnce sync.Once } @@ -195,6 +196,8 @@ func (c *stdioClient) startReader() { } func (c *stdioClient) readLoop() { + c.readers.Add(1) + defer c.readers.Add(-1) var fatal error for { line, err := c.stdout.ReadBytes('\n') @@ -227,6 +230,8 @@ func (c *stdioClient) readLoop() { } } +func (c *stdioClient) Readers() int { return int(c.readers.Load()) } + func (c *stdioClient) notify(method string) error { c.mu.Lock() closed := c.closed diff --git a/runtime/policy/policy.go b/runtime/policy/policy.go index 2cf00f8..fb11b98 100644 --- a/runtime/policy/policy.go +++ b/runtime/policy/policy.go @@ -116,11 +116,11 @@ func TraitsFor(name string) Traits { "fs_read_file", "fs_list_dir", "time_now", "knowledge_search", "list_conversations", "read_conversation", "quality_review", "web_search", "tavily_search", "fetch_url", "ask_user", "exit_plan_mode", - "update_plan", "update_todos": + "update_plan", "update_todos", "tool_search": return Traits{ReadOnly: true, Risk: "low"} case "write_file", "edit_file", "multi_edit", "fs_write_file", "write_artifact", "register_artifact": return Traits{Edits: true, Risk: "medium"} - case "run_shell", "run_command", "run_python": + case "run_shell", "run_command", "run_python", "run_tool_script", "call_tool": return Traits{Risk: "medium"} default: return Traits{} diff --git a/runtime/tools/code_tool.go b/runtime/tools/code_tool.go new file mode 100644 index 0000000..e181e14 --- /dev/null +++ b/runtime/tools/code_tool.go @@ -0,0 +1,59 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + "agent-native/runtime/codeorch" +) + +// RegisterCodeTool 注册 run_tool_script。调用方在开关打开时才注册。 +func RegisterCodeTool(reg *Registry) { + if reg == nil { + return + } + reg.Register(&codeScriptTool{reg: reg}) +} + +type codeScriptTool struct{ reg *Registry } + +func (t *codeScriptTool) Name() string { return "run_tool_script" } +func (t *codeScriptTool) Description() string { + return "用一段 Python 依次调用已有工具。禁止再调用解释器、本工具、派发和工作流。失败时会列出已经执行的调用" +} +func (t *codeScriptTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","properties":{"script":{"type":"string"},"timeout_sec":{"type":"integer"}},"required":["script"]}`) +} +func (t *codeScriptTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) { + var req struct { + Script string `json:"script"` + Timeout int `json:"timeout_sec"` + } + if err := json.Unmarshal(args, &req); err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + if strings.TrimSpace(req.Script) == "" { + return Result{Content: "script 不能为空", IsError: true}, nil + } + lim := codeorch.Limits{FullDir: t.reg.Root()} + if req.Timeout > 0 { + lim.Timeout = time.Duration(req.Timeout) * time.Second + } + out := codeorch.Run(ctx, t.reg.Root(), req.Script, lim, nil, func(ctx context.Context, call codeorch.Call) (string, error) { + if call.Op != "tool.call" { + return "", fmt.Errorf("不能调用 %s", call.Op) + } + if codeorch.Forbidden(call.Name) { + return "", fmt.Errorf("不能调用 %s", call.Name) + } + res := t.reg.Execute(ctx, call.Name, call.Args) + if res.IsError { + return "", fmt.Errorf("%s", res.Content) + } + return res.Content, nil + }) + return Result{Content: out.ModelText(), IsError: out.Err != nil}, nil +} diff --git a/runtime/tools/deferred.go b/runtime/tools/deferred.go new file mode 100644 index 0000000..2601aa6 --- /dev/null +++ b/runtime/tools/deferred.go @@ -0,0 +1,164 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "os" + "strings" + + "agent-native/runtime/codeorch" +) + +type searchTool struct{ reg *Registry } + +func (t *searchTool) Name() string { return "tool_search" } +func (t *searchTool) Description() string { + return "按名称或描述查找延迟暴露的工具，并以文本返回它们的 schema" +} +func (t *searchTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}`) +} +func (t *searchTool) Execute(_ context.Context, args json.RawMessage) (Result, error) { + var q struct { + Query string `json:"query"` + } + if len(args) > 0 { + _ = json.Unmarshal(args, &q) + } + return Result{Content: t.reg.Search(q.Query)}, nil +} +func (t *searchTool) Replay() string { return ReplaySafe } + +type callTool struct{ reg *Registry } + +func (t *callTool) Name() string { return "call_tool" } +func (t *callTool) Description() string { + return "调用一个延迟暴露的工具。参数 schema 固定，不随工具清单变化" +} +func (t *callTool) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","properties":{"name":{"type":"string"},"arguments":{"type":"object"}},"required":["name"]}`) +} +func (t *callTool) Execute(ctx context.Context, args json.RawMessage) (Result, error) { + var req struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` + } + if err := json.Unmarshal(args, &req); err != nil { + return Result{Content: err.Error(), IsError: true}, nil + } + if req.Name == "" || req.Name == "call_tool" || req.Name == "tool_search" || codeorch.Forbidden(req.Name) { + return Result{Content: "不能调用 " + req.Name, IsError: true}, nil + } + res := t.reg.Execute(ctx, req.Name, req.Arguments) + if !res.IsError { + t.reg.markActivated(req.Name) + } + return res, nil +} + +func (r *Registry) ensureMetaLocked() { + if _, ok := r.tools["tool_search"]; ok { + return + } + r.tools["tool_search"] = &searchTool{reg: r} + r.order = append(r.order, "tool_search") + r.tools["call_tool"] = &callTool{reg: r} + r.order = append(r.order, "call_tool") +} + +// Search 把名称或描述命中的延迟工具 schema 写成文本。 +func (r *Registry) Search(query string) string { + r.mu.Lock() + defer r.mu.Unlock() + q := strings.ToLower(strings.TrimSpace(query)) + var b strings.Builder + for _, name := range r.order { + t := r.tools[name] + if ExposureOf(t) != ExposureDeferred { + continue + } + blob := strings.ToLower(t.Name() + " " + t.Description()) + if q != "" && !strings.Contains(blob, q) { + continue + } + fmt.Fprintf(&b, "name: %s\ndescription: %s\nparameters: %s\n\n", t.Name(), t.Description(), t.Parameters()) + } + if b.Len() == 0 { + return "没有匹配的延迟工具" + } + return strings.TrimSpace(b.String()) +} + +// SetActivationFile 指定已激活工具的持久文件。 +func (r *Registry) SetActivationFile(path string) { + r.mu.Lock() + r.activatePath = path + r.mu.Unlock() +} + +// RestoreActivated 读回上次激活的延迟工具名。清单仍然不包含它们。 +func (r *Registry) RestoreActivated() error { + r.mu.Lock() + path := r.activatePath + r.mu.Unlock() + if path == "" { + return nil + } + raw, err := os.ReadFile(path) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return err + } + var names []string + if err := json.Unmarshal(raw, &names); err != nil { + return err + } + r.mu.Lock() + defer r.mu.Unlock() + if r.activated == nil { + r.activated = map[string]bool{} + } + for _, name := range names { + r.activated[name] = true + } + return nil +} + +// Activated 返回已恢复或本次调用过的延迟工具名。 +func (r *Registry) Activated() []string { + r.mu.Lock() + defer r.mu.Unlock() + var out []string + for name := range r.activated { + out = append(out, name) + } + return out +} + +func (r *Registry) markActivated(name string) { + r.mu.Lock() + defer r.mu.Unlock() + t, ok := r.tools[name] + if !ok || ExposureOf(t) != ExposureDeferred { + return + } + if r.activated == nil { + r.activated = map[string]bool{} + } + r.activated[name] = true + if r.activatePath == "" { + return + } + var names []string + for n := range r.activated { + names = append(names, n) + } + raw, err := json.Marshal(names) + if err != nil { + return + } + _ = os.WriteFile(r.activatePath, raw, 0o644) +} diff --git a/runtime/tools/deferred_test.go b/runtime/tools/deferred_test.go new file mode 100644 index 0000000..d2b08dd --- /dev/null +++ b/runtime/tools/deferred_test.go @@ -0,0 +1,75 @@ +package tools + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" +) + +type deferredEcho struct{} + +func (deferredEcho) Name() string { return "mcp__late__echo" } +func (deferredEcho) Description() string { return "回声延迟工具" } +func (deferredEcho) Parameters() json.RawMessage { + return json.RawMessage(`{"type":"object","properties":{"text":{"type":"string"}}}`) +} +func (deferredEcho) Exposure() string { return ExposureDeferred } +func (deferredEcho) Execute(_ context.Context, args json.RawMessage) (Result, error) { + var a struct { + Text string `json:"text"` + } + _ = json.Unmarshal(args, &a) + return Result{Content: "echo " + a.Text}, nil +} + +func TestDeferredToolStaysOutOfDeclsUntilCalledByName(t *testing.T) { + reg := NewRegistry(t.TempDir()) + reg.Register(deferredEcho{}) + for _, d := range reg.Decls() { + if d.Name == "mcp__late__echo" { + t.Fatal("deferred tool must not be listed") + } + } + found := reg.Search("回声") + if !strings.Contains(found, "mcp__late__echo") || !strings.Contains(found, `"text"`) { + t.Fatalf("search %q", found) + } + res := reg.Execute(context.Background(), "call_tool", json.RawMessage(`{"name":"mcp__late__echo","arguments":{"text":"hi"}}`)) + if res.IsError || res.Content != "echo hi" { + t.Fatalf("call %+v", res) + } + denied := reg.Execute(context.Background(), "call_tool", json.RawMessage(`{"name":"run_shell"}`)) + if !denied.IsError { + t.Fatalf("forbidden call should fail: %+v", denied) + } +} + +func TestActivatedNamesRestoreWithoutEnteringDecls(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "activated.json") + reg := NewRegistry(dir) + reg.SetActivationFile(path) + reg.Register(deferredEcho{}) + _ = reg.Execute(context.Background(), "call_tool", json.RawMessage(`{"name":"mcp__late__echo","arguments":{"text":"x"}}`)) + if _, err := os.Stat(path); err != nil { + t.Fatal(err) + } + next := NewRegistry(dir) + next.SetActivationFile(path) + if err := next.RestoreActivated(); err != nil { + t.Fatal(err) + } + next.Register(deferredEcho{}) + names := next.Activated() + if len(names) != 1 || names[0] != "mcp__late__echo" { + t.Fatalf("activated %+v", names) + } + for _, d := range next.Decls() { + if d.Name == "mcp__late__echo" { + t.Fatal("restored tool must stay out of the list") + } + } +} diff --git a/runtime/tools/registry.go b/runtime/tools/registry.go index d90e72b..820d2be 100644 --- a/runtime/tools/registry.go +++ b/runtime/tools/registry.go @@ -53,6 +53,25 @@ type BeforeHook func(ctx context.Context, name string, args json.RawMessage) (bl // AfterHook can rewrite a tool result before it is returned. type AfterHook func(ctx context.Context, name string, args json.RawMessage, result Result) Result +const ( + // ExposureDirect 进入模型看到的工具清单。ExposureDeferred 只通过 tool_search 与 call_tool。 + ExposureDirect = "direct" + ExposureDeferred = "deferred" +) + +// Exposer 声明工具是否直接出现在清单里。未实现视为直接。 +type Exposer interface { + Exposure() string +} + +// ExposureOf 读取暴露级别。 +func ExposureOf(t Tool) string { + if e, ok := t.(Exposer); ok && e.Exposure() == ExposureDeferred { + return ExposureDeferred + } + return ExposureDirect +} + // Registry executes jailed tools and appends an audit line per call. type Registry struct { root string @@ -63,6 +82,13 @@ type Registry struct { Audit io.Writer // hookMu 串行 Before、After 与 Audit，避免并行工具调用交错写审计或重入钩子。 hookMu sync.Mutex + mu sync.Mutex + // restricted 为真时，后注册的工具也必须命中 allowExact / allowPrefix。 + restricted bool + allowExact map[string]bool + allowPrefix []string + activated map[string]bool + activatePath string } func NewRegistry(root string) *Registry { @@ -72,6 +98,8 @@ func NewRegistry(root string) *Registry { func (r *Registry) Root() string { return r.root } func (r *Registry) Has(name string) bool { + r.mu.Lock() + defer r.mu.Unlock() _, ok := r.tools[name] return ok } @@ -81,7 +109,9 @@ func (r *Registry) Replay(name string) string { if r == nil { return ReplayNever } + r.mu.Lock() t, ok := r.tools[name] + r.mu.Unlock() if !ok { return ReplayNever } @@ -89,16 +119,29 @@ func (r *Registry) Replay(name string) string { } func (r *Registry) Register(t Tool) { + r.mu.Lock() + defer r.mu.Unlock() + if r.restricted && !r.keepLocked(t.Name()) { + return + } if _, ok := r.tools[t.Name()]; !ok { r.order = append(r.order, t.Name()) } r.tools[t.Name()] = t + if ExposureOf(t) == ExposureDeferred { + r.ensureMetaLocked() + } } func (r *Registry) Decls() []Decl { + r.mu.Lock() + defer r.mu.Unlock() out := make([]Decl, 0, len(r.order)) for _, name := range r.order { t := r.tools[name] + if ExposureOf(t) == ExposureDeferred { + continue + } out = append(out, Decl{Name: t.Name(), Description: t.Description(), Parameters: t.Parameters()}) } return out @@ -115,6 +158,8 @@ func (r *Registry) Restrict(allow []string) error { if len(allow) == 0 { return nil } + r.mu.Lock() + defer r.mu.Unlock() exact := make(map[string]bool, len(allow)) prefixes := make([]string, 0) for _, entry := range allow { @@ -141,20 +186,12 @@ func (r *Registry) Restrict(allow []string) error { 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 - } + r.restricted = true + r.allowExact = exact + r.allowPrefix = prefixes var order []string for _, name := range r.order { - if keep(name) { + if r.keepLocked(name) { order = append(order, name) } else { delete(r.tools, name) @@ -164,6 +201,21 @@ func (r *Registry) Restrict(allow []string) error { return nil } +func (r *Registry) keepLocked(name string) bool { + if name == "tool_search" || name == "call_tool" { + return true + } + if r.allowExact[name] { + return true + } + for _, p := range r.allowPrefix { + if strings.HasPrefix(name, p) { + return true + } + } + return false +} + func (r *Registry) hasPrefix(prefix string) bool { for _, name := range r.order { if strings.HasPrefix(name, prefix) { @@ -182,7 +234,9 @@ type Decl struct { // Execute validates arguments, runs hooks, and audits the outcome. func (r *Registry) Execute(ctx context.Context, name string, args json.RawMessage) Result { + r.mu.Lock() t, ok := r.tools[name] + r.mu.Unlock() if !ok { res := Result{Content: "unknown tool " + name, IsError: true} r.audit(name, false, true, "unknown")