Pair skipped tool calls and close the stdio and fetch gaps
ci / vet-and-test (pull_request) Failing after 2s
ci / docker-checkout-hostgateway (pull_request) Successful in 2s

Why: a tool-budget stop left tool calls without results, host callbacks raced, a canceled MCP read stole the next response, and fetch re-resolved after the SSRF check.
Evidence: GOWORK=off go vet ./... && go test -count=1 ./... && go test -race ./runtime/... ./adapters/native/... passed, including TestLoopBudgetToolCallCap, TestParallelHostCallbacksAreSerial, TestStdioSlowCancelThenEcho, TestPresetsAndClientForChatCompletions, TestFetchResolveChangesToLoopbackRejected, and TestFetchHostCheckRedirectStillRejectsPrivate.
Impact: over-budget calls get an error result and no tool event; one Run serializes host callbacks; MCP dispatches by id; anthropic and google presets use OpenAI-compatible chat endpoints; fetch dials the address it just checked. deps/agent-intent is 3130190; error-code-system and logging-system are a3ee67b and 49b5d8b.
Refs: N1, N2, N3, N4, N13, N14
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
2026-10-03 01:10:29 +00:00
co-authored by Cursor
parent 7e0374ff40
commit 694b08cdbe
15 changed files with 522 additions and 75 deletions
+9 -1
View File
@@ -109,7 +109,15 @@ func liveStream(ref llm.ModelRef) llm.StreamFunc {
if !ok { if !ok {
provider = llm.Provider{Protocol: llm.ProtocolOpenAIChat, BaseURL: os.Getenv("OPENAI_BASE_URL"), APIKey: os.Getenv("OPENAI_API_KEY")} provider = llm.Provider{Protocol: llm.ProtocolOpenAIChat, BaseURL: os.Getenv("OPENAI_BASE_URL"), APIKey: os.Getenv("OPENAI_API_KEY")}
} }
client := &llm.Client{BaseURL: provider.BaseURL, APIKey: provider.APIKey} client, err := llm.ClientFor(provider)
if err != nil {
return func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
ch := make(chan llm.Event, 1)
ch <- llm.Event{Kind: llm.KindError, Err: err, Message: llm.Message{StopReason: llm.StopError}}
close(ch)
return ch
}
}
return llm.WithRetry(client.Stream, 2, nil) return llm.WithRetry(client.Stream, 2, nil)
} }
+16 -2
View File
@@ -137,10 +137,10 @@ type Provider struct {
// Presets are OpenAI-compatible or native endpoints that only need an API key. // Presets are OpenAI-compatible or native endpoints that only need an API key.
func Presets() []Provider { func Presets() []Provider {
return []Provider{ return []Provider{
{Name: "anthropic", Protocol: ProtocolAnthropic, BaseURL: "https://api.anthropic.com", APIKeyEnv: "ANTHROPIC_API_KEY"}, {Name: "anthropic", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.anthropic.com/v1", APIKeyEnv: "ANTHROPIC_API_KEY"},
{Name: "openai", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.openai.com/v1", APIKeyEnv: "OPENAI_API_KEY"}, {Name: "openai", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.openai.com/v1", APIKeyEnv: "OPENAI_API_KEY"},
{Name: "openai-responses", Protocol: ProtocolOpenAIResponses, BaseURL: "https://api.openai.com/v1", APIKeyEnv: "OPENAI_API_KEY"}, {Name: "openai-responses", Protocol: ProtocolOpenAIResponses, BaseURL: "https://api.openai.com/v1", APIKeyEnv: "OPENAI_API_KEY"},
{Name: "google", Protocol: ProtocolGoogle, BaseURL: "https://generativelanguage.googleapis.com/v1beta", APIKeyEnv: "GEMINI_API_KEY"}, {Name: "google", Protocol: ProtocolOpenAIChat, BaseURL: "https://generativelanguage.googleapis.com/v1beta/openai", APIKeyEnv: "GEMINI_API_KEY"},
{Name: "azure", Protocol: ProtocolAzure, BaseURL: "", APIKeyEnv: "AZURE_OPENAI_API_KEY"}, {Name: "azure", Protocol: ProtocolAzure, BaseURL: "", APIKeyEnv: "AZURE_OPENAI_API_KEY"},
{Name: "deepseek", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.deepseek.com/v1", APIKeyEnv: "DEEPSEEK_API_KEY"}, {Name: "deepseek", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.deepseek.com/v1", APIKeyEnv: "DEEPSEEK_API_KEY"},
{Name: "minimax", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.minimaxi.com/v1", APIKeyEnv: "MINIMAX_API_KEY"}, {Name: "minimax", Protocol: ProtocolOpenAIChat, BaseURL: "https://api.minimaxi.com/v1", APIKeyEnv: "MINIMAX_API_KEY"},
@@ -193,6 +193,17 @@ func Lookup(name string, custom []Provider) (Provider, bool) {
return Provider{}, false return Provider{}, false
} }
// ClientFor 返回会说该协议的客户端。responses、azure 以及原生 Anthropic/Google
// 协议当前没有接到客户端上,返回明确错误而不是发到错误的路径。
func ClientFor(p Provider) (*Client, error) {
switch p.Protocol {
case ProtocolOpenAIChat:
return &Client{BaseURL: p.BaseURL, APIKey: p.APIKey}, nil
default:
return nil, fmt.Errorf("llm: protocol %q is not implemented", p.Protocol)
}
}
func messageContent(m Message) any { func messageContent(m Message) any {
if len(m.Images) == 0 { if len(m.Images) == 0 {
return m.Content return m.Content
@@ -250,6 +261,9 @@ func BuildBody(protocol Protocol, model string, messages []Message, tools []Tool
body["reasoning"] = map[string]any{"effort": effort} body["reasoning"] = map[string]any{"effort": effort}
} }
return json.Marshal(body) return json.Marshal(body)
// ProtocolAnthropic 与 ProtocolGoogle 只供离线组包。
// 当前 Client 只发 OpenAI chat completions,这两支没有接到客户端上。
// 原生协议是后续阶段的事;预设表里的 anthropic / google 已改走兼容端点。
case ProtocolAnthropic: case ProtocolAnthropic:
sys, msgs := anthropicMessages(messages) sys, msgs := anthropicMessages(messages)
body := map[string]any{ body := map[string]any{
+46
View File
@@ -3,6 +3,8 @@ package llm
import ( import (
"context" "context"
"encoding/json" "encoding/json"
"net/http"
"net/http/httptest"
"os" "os"
"path/filepath" "path/filepath"
"strings" "strings"
@@ -65,6 +67,50 @@ func TestModelsJSONAndBodies(t *testing.T) {
} }
} }
func TestPresetsAndClientForChatCompletions(t *testing.T) {
want := map[string]string{
"anthropic": "https://api.anthropic.com/v1",
"google": "https://generativelanguage.googleapis.com/v1beta/openai",
}
for _, p := range Presets() {
base, ok := want[p.Name]
if !ok {
continue
}
if p.Protocol != ProtocolOpenAIChat || p.BaseURL != base {
t.Fatalf("%s preset: %+v", p.Name, p)
}
}
for _, name := range []string{"openai-responses", "azure"} {
p, ok := Lookup(name, nil)
if !ok {
t.Fatal(name)
}
if _, err := ClientFor(p); err == nil {
t.Fatalf("%s should fail closed", name)
}
}
var path string
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
path = r.URL.Path
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte("data: [DONE]\n\n"))
}))
defer srv.Close()
client, err := ClientFor(Provider{Protocol: ProtocolOpenAIChat, BaseURL: srv.URL, APIKey: "k"})
if err != nil {
t.Fatal(err)
}
for ev := range client.Stream(context.Background(), "m", []Message{{Role: RoleUser, Content: "hi"}}, nil, StreamOptions{}) {
if ev.Kind == KindError {
t.Fatal(ev.Err)
}
}
if path != "/chat/completions" {
t.Fatalf("path %q", path)
}
}
func TestRetryThenSuccess(t *testing.T) { func TestRetryThenSuccess(t *testing.T) {
n := 0 n := 0
base := StreamFunc(func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event { base := StreamFunc(func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
+60
View File
@@ -54,6 +54,8 @@ type Config struct {
ToolExecution string ToolExecution string
GetSteering func() []llm.Message GetSteering func() []llm.Message
GetFollowUp func() []llm.Message GetFollowUp func() []llm.Message
// OnEvent 观察循环事件。同一次 Run 内不会并发调用
// (与 BeforeTool、AfterTool 共用一把锁串行)。
OnEvent func(Event) OnEvent func(Event)
// ThinkingLevel is forwarded on each provider request. // ThinkingLevel is forwarded on each provider request.
ThinkingLevel string ThinkingLevel string
@@ -73,10 +75,12 @@ type Config struct {
// nil 表示不询问。内核最多接受 3 次 Continue。 // nil 表示不询问。内核最多接受 3 次 Continue。
Stop func(last llm.Message, info StopInfo) StopDecision Stop func(last llm.Message, info StopInfo) StopDecision
// BeforeTool 返回错误时不执行工具,但仍写一条 IsError 的 tool 结果。 // BeforeTool 返回错误时不执行工具,但仍写一条 IsError 的 tool 结果。
// 同一次 Run 内与 OnEvent、AfterTool 串行,不会并发调用。
BeforeTool func(call llm.ToolCall) error BeforeTool func(call llm.ToolCall) error
// UnknownTool 处理注册表里没有的调用。nil 时仍返回 unknown tool。 // UnknownTool 处理注册表里没有的调用。nil 时仍返回 unknown tool。
UnknownTool func(call llm.ToolCall) tools.Result UnknownTool func(call llm.ToolCall) tools.Result
// AfterTool 只观察工具结果,不能改写。 // AfterTool 只观察工具结果,不能改写。
// 同一次 Run 内与 OnEvent、BeforeTool 串行,不会并发调用。
AfterTool func(call llm.ToolCall, res llm.Message) AfterTool func(call llm.ToolCall, res llm.Message)
// Reminders 在每次调用模型前提供提醒。正文包在 system-reminder 里, // Reminders 在每次调用模型前提供提醒。正文包在 system-reminder 里,
// 只进入本次请求,不写进 Session。 // 只进入本次请求,不写进 Session。
@@ -136,6 +140,7 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
if cfg.ToolExecution == "" { if cfg.ToolExecution == "" {
cfg.ToolExecution = ExecSequential cfg.ToolExecution = ExecSequential
} }
serializeHostCallbacks(&cfg)
var messages []llm.Message var messages []llm.Message
if cfg.Session != nil { if cfg.Session != nil {
loaded, err := cfg.Session.Load() loaded, err := cfg.Session.Load()
@@ -196,6 +201,11 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
if last.StopReason == llm.StopToolCalls && len(last.ToolCalls) > 0 { if last.StopReason == llm.StopToolCalls && len(last.ToolCalls) > 0 {
if cfg.Budget != nil && cfg.Budget.MaxToolCalls > 0 && toolCalls+len(last.ToolCalls) > cfg.Budget.MaxToolCalls { 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) reason := fmt.Sprintf("budget_exceeded: tool calls %d+%d > max %d", toolCalls, len(last.ToolCalls), cfg.Budget.MaxToolCalls)
var err error
messages, err = appendSkippedToolResults(cfg, last.ToolCalls, messages)
if err != nil {
return Result{}, err
}
return budgetStop(cfg, last, messages, usage, reason), nil return budgetStop(cfg, last, messages, usage, reason), nil
} }
toolCalls += len(last.ToolCalls) toolCalls += len(last.ToolCalls)
@@ -352,6 +362,56 @@ func consume(ctx context.Context, cfg Config, messages []llm.Message) (llm.Event
return end, nil return end, nil
} }
const toolBudgetSkipContent = "not executed: tool budget exceeded"
// appendSkippedToolResults 为未执行的 tool_call 各写一条错误结果,保持调用与结果成对。
// 不发工具事件:这些调用没有进入执行。
func appendSkippedToolResults(cfg Config, calls []llm.ToolCall, messages []llm.Message) ([]llm.Message, error) {
for _, call := range calls {
msg := llm.Message{
Role: llm.RoleTool,
Name: call.Name,
ToolCallID: call.ID,
Content: toolBudgetSkipContent,
IsError: true,
}
if err := persist(cfg, msg); err != nil {
return nil, err
}
messages = append(messages, msg)
}
return messages, nil
}
// serializeHostCallbacks 让同一次 Run 的宿主回调不会并发进入。
func serializeHostCallbacks(cfg *Config) {
var mu sync.Mutex
if cfg.OnEvent != nil {
inner := cfg.OnEvent
cfg.OnEvent = func(ev Event) {
mu.Lock()
defer mu.Unlock()
inner(ev)
}
}
if cfg.BeforeTool != nil {
inner := cfg.BeforeTool
cfg.BeforeTool = func(call llm.ToolCall) error {
mu.Lock()
defer mu.Unlock()
return inner(call)
}
}
if cfg.AfterTool != nil {
inner := cfg.AfterTool
cfg.AfterTool = func(call llm.ToolCall, res llm.Message) {
mu.Lock()
defer mu.Unlock()
inner(call, res)
}
}
}
func executeTools(ctx context.Context, cfg Config, assistant llm.Message) []llm.Message { func executeTools(ctx context.Context, cfg Config, assistant llm.Message) []llm.Message {
if cfg.ToolExecution == ExecParallel && len(assistant.ToolCalls) > 1 { if cfg.ToolExecution == ExecParallel && len(assistant.ToolCalls) > 1 {
return executeParallel(ctx, cfg, assistant) return executeParallel(ctx, cfg, assistant)
+98 -4
View File
@@ -4,6 +4,7 @@ import (
"context" "context"
"encoding/json" "encoding/json"
"strings" "strings"
"sync"
"testing" "testing"
"time" "time"
@@ -66,15 +67,108 @@ func TestLoopBudgetToolCallCap(t *testing.T) {
if !sawBudgetEvent { if !sawBudgetEvent {
t.Fatal("expected budget_exceeded event") t.Fatal("expected budget_exceeded event")
} }
// 第二批工具调用不该被执行(预算在执行前拦截):transcript 只含一次 tool 结果。 // 第一批执行 1 次;第二批超限,两条合成错误结果,不执行。
toolResults := 0 toolResults := 0
skipped := 0
for _, m := range res.Messages { for _, m := range res.Messages {
if m.Role == llm.RoleTool { if m.Role != llm.RoleTool {
continue
}
toolResults++ toolResults++
if m.IsError && m.Content == toolBudgetSkipContent {
skipped++
} }
} }
if toolResults != 1 { if toolResults != 3 || skipped != 2 {
t.Fatalf("expected exactly 1 tool result, got %d", toolResults) t.Fatalf("expected 3 tool results including 2 skipped, got results=%d skipped=%d", toolResults, skipped)
}
assertToolCallsPaired(t, res.Messages)
}
func assertToolCallsPaired(t *testing.T, messages []llm.Message) {
t.Helper()
pending := map[string]string{}
for _, m := range messages {
if m.Role == llm.RoleAssistant {
for _, call := range m.ToolCalls {
if _, ok := pending[call.ID]; ok {
t.Fatalf("duplicate tool call id %s", call.ID)
}
pending[call.ID] = call.Name
}
}
if m.Role == llm.RoleTool {
name, ok := pending[m.ToolCallID]
if !ok {
t.Fatalf("tool result %s has no matching call", m.ToolCallID)
}
if name != m.Name {
t.Fatalf("tool result %s name %q != call %q", m.ToolCallID, m.Name, name)
}
delete(pending, m.ToolCallID)
}
}
if len(pending) != 0 {
t.Fatalf("unpaired tool calls: %v", pending)
}
}
type pauseTool struct{}
func (pauseTool) Name() string { return "pause" }
func (pauseTool) Description() string { return "slow tool" }
func (pauseTool) Parameters() json.RawMessage {
return json.RawMessage(`{"type":"object","properties":{}}`)
}
func (pauseTool) Execute(context.Context, json.RawMessage) (tools.Result, error) {
time.Sleep(40 * time.Millisecond)
return tools.Result{Content: "ok"}, nil
}
func TestParallelHostCallbacksAreSerial(t *testing.T) {
reg := tools.NewRegistry(t.TempDir())
reg.Register(pauseTool{})
var mu sync.Mutex
inFlight, maxInFlight := 0, 0
enter := func() {
mu.Lock()
inFlight++
if inFlight > maxInFlight {
maxInFlight = inFlight
}
mu.Unlock()
time.Sleep(20 * time.Millisecond)
mu.Lock()
inFlight--
mu.Unlock()
}
_, err := Run(context.Background(), Config{
Model: "fake",
Tools: reg,
ToolExecution: ExecParallel,
OnEvent: func(Event) {
enter()
},
BeforeTool: func(llm.ToolCall) error {
enter()
return nil
},
AfterTool: func(llm.ToolCall, llm.Message) {
enter()
},
Stream: usageScript(llm.Usage{},
llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{
{ID: "a", Name: "pause", Arguments: `{}`},
{ID: "b", Name: "pause", Arguments: `{}`},
}},
llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd},
),
}, llm.Message{Role: llm.RoleUser, Content: "go"})
if err != nil {
t.Fatal(err)
}
if maxInFlight != 1 {
t.Fatalf("host callbacks overlapped: max in flight %d", maxInFlight)
} }
} }
+9
View File
@@ -8,6 +8,7 @@ import (
"bufio" "bufio"
"encoding/json" "encoding/json"
"os" "os"
"time"
) )
// RunFakeServer 在 stdin/stdout 上运行最小 MCP server(echo / boom / 未知方法), // RunFakeServer 在 stdin/stdout 上运行最小 MCP server(echo / boom / 未知方法),
@@ -67,6 +68,9 @@ func fakeServerHandle(req rpcRequest) rpcResponse {
Content: []contentBlock{{Type: "text", Text: "boom failed"}}, Content: []contentBlock{{Type: "text", Text: "boom failed"}},
IsError: true, IsError: true,
}) })
case "slow":
time.Sleep(2 * time.Second)
return respond(callToolResult{Content: []contentBlock{{Type: "text", Text: "slow done"}}})
default: default:
return fail(-32602, "unknown tool "+p.Name) return fail(-32602, "unknown tool "+p.Name)
} }
@@ -87,6 +91,11 @@ func fakeServerTools() []toolInfo {
Description: "Always fails", Description: "Always fails",
InputSchema: json.RawMessage(`{"type":"object","properties":{}}`), InputSchema: json.RawMessage(`{"type":"object","properties":{}}`),
}, },
{
Name: "slow",
Description: "Responds after a delay",
InputSchema: json.RawMessage(`{"type":"object","properties":{}}`),
},
} }
} }
+29 -3
View File
@@ -45,7 +45,7 @@ func TestStdioClientHandshakeListAndCall(t *testing.T) {
defer c.Close() defer c.Close()
list, err := c.ListTools(ctx) list, err := c.ListTools(ctx)
if err != nil || len(list) != 2 || list[0].Name != "echo" { if err != nil || len(list) != 3 || list[0].Name != "echo" {
t.Fatalf("list failed: %v %+v", err, list) t.Fatalf("list failed: %v %+v", err, list)
} }
res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"hi"}`)) res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"hi"}`))
@@ -75,6 +75,32 @@ func TestStdioClientServerRejectsBadSchema(t *testing.T) {
} }
} }
func TestStdioSlowCancelThenEcho(t *testing.T) {
ctx := context.Background()
c, err := DialStdio(ctx, fakeStdioConfig(t))
if err != nil {
t.Fatal(err)
}
defer c.Close()
slowCtx, cancel := context.WithTimeout(ctx, 80*time.Millisecond)
defer cancel()
_, err = c.CallTool(slowCtx, "slow", nil)
if err == nil {
t.Fatal("slow call should time out")
}
echoCtx, echoCancel := context.WithTimeout(ctx, 3*time.Second)
defer echoCancel()
res, err := c.CallTool(echoCtx, "echo", json.RawMessage(`{"text":"after-slow"}`))
if err != nil {
t.Fatal(err)
}
if len(res.Content) != 1 || res.Content[0].Text != "echo: after-slow" {
t.Fatalf("echo must receive its own result, got %+v", res)
}
}
func fakeHTTPHandler(w http.ResponseWriter, r *http.Request) { func fakeHTTPHandler(w http.ResponseWriter, r *http.Request) {
var req rpcRequest var req rpcRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil { if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
@@ -104,7 +130,7 @@ func TestHTTPClientJSONMode(t *testing.T) {
t.Fatalf("session id must be captured, got %q", c.sessionID) t.Fatalf("session id must be captured, got %q", c.sessionID)
} }
list, err := c.ListTools(ctx) list, err := c.ListTools(ctx)
if err != nil || len(list) != 2 { if err != nil || len(list) != 3 {
t.Fatalf("list failed: %v %+v", err, list) t.Fatalf("list failed: %v %+v", err, list)
} }
res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"net"}`)) res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"net"}`))
@@ -130,7 +156,7 @@ func TestHTTPClientSSEMode(t *testing.T) {
} }
defer c.Close() defer c.Close()
list, err := c.ListTools(ctx) list, err := c.ListTools(ctx)
if err != nil || len(list) != 2 { if err != nil || len(list) != 3 {
t.Fatalf("sse list failed: %v %+v", err, list) t.Fatalf("sse list failed: %v %+v", err, list)
} }
} }
+82 -34
View File
@@ -12,14 +12,23 @@ import (
) )
// stdioClient 通过「子进程 stdin/stdout 按行分隔的 JSON-RPC」与 server 通信。 // stdioClient 通过「子进程 stdin/stdout 按行分隔的 JSON-RPC」与 server 通信。
// 同一时刻只允许一个在途请求(串行锁),匹配简单且足以支撑智能体工具调用。 // 一个长驻读循环按 id 把响应分发给各请求;写 stdin 仍串行。
// 取消会摘掉等待项,迟到的回复直接丢弃。读循环退出时,在等的请求全部失败,客户端标记为已死。
type stdioClient struct { type stdioClient struct {
cmd *exec.Cmd cmd *exec.Cmd
stdin io.WriteCloser stdin io.WriteCloser
stdout *bufio.Reader stdout *bufio.Reader
writeMu sync.Mutex
mu sync.Mutex mu sync.Mutex
pending map[int64]chan stdioResult
nextID int64 nextID int64
closed bool closed bool
readerOnce sync.Once
}
type stdioResult struct {
resp rpcResponse
err error
} }
// DialStdio 启动 server 子进程并完成 initialize 握手。 // DialStdio 启动 server 子进程并完成 initialize 握手。
@@ -43,7 +52,13 @@ func DialStdio(ctx context.Context, cfg ServerConfig) (*stdioClient, error) {
if err := cmd.Start(); err != nil { if err := cmd.Start(); err != nil {
return nil, fmt.Errorf("mcp server %q: start %s: %w", cfg.Name, cfg.Command, err) 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)} c := &stdioClient{
cmd: cmd,
stdin: stdin,
stdout: bufio.NewReaderSize(stdout, 1<<20),
pending: map[int64]chan stdioResult{},
}
c.startReader()
if err := c.handshake(ctx); err != nil { if err := c.handshake(ctx); err != nil {
_ = c.Close() _ = c.Close()
return nil, err return nil, err
@@ -104,11 +119,12 @@ func (c *stdioClient) CallTool(ctx context.Context, name string, args json.RawMe
func (c *stdioClient) Close() error { func (c *stdioClient) Close() error {
c.mu.Lock() c.mu.Lock()
defer c.mu.Unlock()
if c.closed { if c.closed {
c.mu.Unlock()
return nil return nil
} }
c.closed = true c.closed = true
c.mu.Unlock()
_ = c.stdin.Close() _ = c.stdin.Close()
if c.cmd.Process != nil { if c.cmd.Process != nil {
_ = c.cmd.Process.Kill() _ = c.cmd.Process.Kill()
@@ -117,29 +133,43 @@ func (c *stdioClient) Close() error {
} }
func (c *stdioClient) roundTrip(ctx context.Context, method string, params any, out any) error { func (c *stdioClient) roundTrip(ctx context.Context, method string, params any, out any) error {
c.startReader()
id := atomic.AddInt64(&c.nextID, 1)
ch := make(chan stdioResult, 1)
c.mu.Lock() c.mu.Lock()
defer c.mu.Unlock()
if c.closed { if c.closed {
c.mu.Unlock()
return fmt.Errorf("mcp client closed") return fmt.Errorf("mcp client closed")
} }
id := atomic.AddInt64(&c.nextID, 1) c.pending[id] = ch
c.mu.Unlock()
req := rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params} req := rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params}
body, err := json.Marshal(req) body, err := json.Marshal(req)
if err != nil { if err != nil {
c.dropPending(id)
return err return err
} }
body = append(body, '\n') body = append(body, '\n')
if _, err := c.stdin.Write(body); err != nil { c.writeMu.Lock()
_, err = c.stdin.Write(body)
c.writeMu.Unlock()
if err != nil {
c.dropPending(id)
return fmt.Errorf("mcp write %s: %w", method, err) return fmt.Errorf("mcp write %s: %w", method, err)
} }
line, err := c.readResponseLine(ctx)
if err != nil { var got stdioResult
return fmt.Errorf("mcp read %s: %w", method, err) select {
case got = <-ch:
case <-ctx.Done():
c.dropPending(id)
return fmt.Errorf("mcp read %s: %w", method, ctx.Err())
} }
var resp rpcResponse if got.err != nil {
if err := json.Unmarshal(line, &resp); err != nil { return fmt.Errorf("mcp read %s: %w", method, got.err)
return fmt.Errorf("mcp %s: bad response frame: %w", method, err)
} }
resp := got.resp
if resp.ID == nil || *resp.ID != id { if resp.ID == nil || *resp.ID != id {
return fmt.Errorf("mcp %s: response id mismatch", method) return fmt.Errorf("mcp %s: response id mismatch", method)
} }
@@ -152,40 +182,56 @@ func (c *stdioClient) roundTrip(ctx context.Context, method string, params any,
return nil return nil
} }
// readResponseLine 逐行读取,跳过通知帧(无 id),带 ctx 取消感知。 func (c *stdioClient) dropPending(id int64) {
func (c *stdioClient) readResponseLine(ctx context.Context) ([]byte, error) { c.mu.Lock()
type result struct { delete(c.pending, id)
line []byte c.mu.Unlock()
err error
} }
done := make(chan result, 1)
go func() { func (c *stdioClient) startReader() {
c.readerOnce.Do(func() {
go c.readLoop()
})
}
func (c *stdioClient) readLoop() {
var fatal error
for { for {
line, err := c.stdout.ReadBytes('\n') line, err := c.stdout.ReadBytes('\n')
if err != nil { if err != nil {
done <- result{nil, err} fatal = err
return break
} }
var probe rpcResponse var resp rpcResponse
if json.Unmarshal(line, &probe) == nil && probe.ID != nil { if json.Unmarshal(line, &resp) != nil || resp.ID == nil {
done <- result{line, nil} continue
return
} }
// 通知或日志行:忽略。 c.mu.Lock()
ch := c.pending[*resp.ID]
if ch != nil {
delete(c.pending, *resp.ID)
} }
}() c.mu.Unlock()
select { if ch == nil {
case r := <-done: continue
return r.line, r.err }
case <-ctx.Done(): ch <- stdioResult{resp: resp}
return nil, ctx.Err() }
c.mu.Lock()
c.closed = true
pending := c.pending
c.pending = map[int64]chan stdioResult{}
c.mu.Unlock()
for _, ch := range pending {
ch <- stdioResult{err: fatal}
} }
} }
func (c *stdioClient) notify(method string) error { func (c *stdioClient) notify(method string) error {
c.mu.Lock() c.mu.Lock()
defer c.mu.Unlock() closed := c.closed
if c.closed { c.mu.Unlock()
if closed {
return fmt.Errorf("mcp client closed") return fmt.Errorf("mcp client closed")
} }
body, err := json.Marshal(rpcNotification{JSONRPC: "2.0", Method: method}) body, err := json.Marshal(rpcNotification{JSONRPC: "2.0", Method: method})
@@ -193,6 +239,8 @@ func (c *stdioClient) notify(method string) error {
return err return err
} }
body = append(body, '\n') body = append(body, '\n')
c.writeMu.Lock()
defer c.writeMu.Unlock()
_, err = c.stdin.Write(body) _, err = c.stdin.Write(body)
return err return err
} }
+56 -6
View File
@@ -17,6 +17,10 @@ import (
// FetchConfig 声明 fetch_url 的运行条件;未注册(组合根未提供)则工具不存在。 // FetchConfig 声明 fetch_url 的运行条件;未注册(组合根未提供)则工具不存在。
type FetchConfig struct { type FetchConfig struct {
// Client 可选注入(测试替身);缺省构造 20s 超时客户端。 // Client 可选注入(测试替身);缺省构造 20s 超时客户端。
// Transport 为 nil 或 *http.Transport 时,Clone 后替换 DialContext:
// 拨号时再次解析,只连接通过 SSRF 校验的 IP。
// 自定义 RoundTripper 无法替换拨号,只做请求前预校验;实现不得在
// RoundTrip 里重新解析到未校验的地址。
Client *http.Client Client *http.Client
// MaxBytes 响应体上限(默认 2 MiB)。 // MaxBytes 响应体上限(默认 2 MiB)。
MaxBytes int MaxBytes int
@@ -89,19 +93,24 @@ func (t *fetchTool) Execute(ctx context.Context, args json.RawMessage) (Result,
req.Header.Set("User-Agent", "agent-native-fetch/1.0") req.Header.Set("User-Agent", "agent-native-fetch/1.0")
var finalURL string var finalURL string
resolve := t.cfg.resolver() resolve := t.cfg.resolver()
client2 := client hostRedirect := client.CheckRedirect
if client2.CheckRedirect == nil { transport, _ := pinFetchTransport(client.Transport, resolve)
client2 = &http.Client{ client2 := &http.Client{
Transport: client.Transport, Transport: transport,
Jar: client.Jar, Jar: client.Jar,
Timeout: client.Timeout, Timeout: client.Timeout,
CheckRedirect: func(req *http.Request, via []*http.Request) error { CheckRedirect: func(req *http.Request, via []*http.Request) error {
if len(via) >= 5 { if len(via) >= 5 {
return fmt.Errorf("stopped after 5 redirects") return fmt.Errorf("stopped after 5 redirects")
} }
return validateFetchTarget(req.URL, allowed, resolve) if err := validateFetchTarget(req.URL, allowed, resolve); err != nil {
}, return err
} }
if hostRedirect != nil {
return hostRedirect(req, via)
}
return nil
},
} }
if err := validateFetchTarget(req.URL, allowed, resolve); err != nil { if err := validateFetchTarget(req.URL, allowed, resolve); err != nil {
return Result{Content: err.Error(), IsError: true}, nil return Result{Content: err.Error(), IsError: true}, nil
@@ -127,6 +136,47 @@ func (t *fetchTool) Execute(ctx context.Context, args json.RawMessage) (Result,
return Result{Content: "url: " + finalURL + "\ncontent-type: " + ct + "\n\n" + text}, nil return Result{Content: "url: " + finalURL + "\ncontent-type: " + ct + "\n\n" + text}, nil
} }
// pinFetchTransport 在默认或 *http.Transport 上安装「解析并拨号已校验 IP」的 DialContext。
// 自定义 RoundTripper 原样返回,调用方保留预校验。
func pinFetchTransport(base http.RoundTripper, resolve func(string) ([]net.IP, error)) (http.RoundTripper, bool) {
var tr *http.Transport
switch t := base.(type) {
case nil:
tr = http.DefaultTransport.(*http.Transport).Clone()
case *http.Transport:
tr = t.Clone()
default:
return base, false
}
// 不走环境变量里的代理:代理拨号会绕开「只连接已校验 IP」。
tr.Proxy = nil
dialer := &net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
host, port, err := net.SplitHostPort(addr)
if err != nil {
return nil, err
}
ips, err := resolve(host)
if err != nil {
return nil, fmt.Errorf("resolve %q: %w", host, err)
}
var chosen net.IP
for _, ip := range ips {
if ipIsPrivate(ip) {
return nil, fmt.Errorf("target %q resolves to private/loopback address %s (SSRF guard)", host, ip)
}
if chosen == nil {
chosen = ip
}
}
if chosen == nil {
return nil, fmt.Errorf("resolve %q: no addresses", host)
}
return dialer.DialContext(ctx, network, net.JoinHostPort(chosen.String(), port))
}
return tr, true
}
// validateFetchTarget 校验 scheme、主机策略与 SSRF 硬边界(私网/回环/链路本地)。 // validateFetchTarget 校验 scheme、主机策略与 SSRF 硬边界(私网/回环/链路本地)。
func validateFetchTarget(u *url.URL, allowed map[string]bool, resolve func(string) ([]net.IP, error)) error { func validateFetchTarget(u *url.URL, allowed map[string]bool, resolve func(string) ([]net.IP, error)) error {
if u.Scheme != "http" && u.Scheme != "https" { if u.Scheme != "http" && u.Scheme != "https" {
+39
View File
@@ -144,6 +144,45 @@ func (s redirectTransport) RoundTrip(req *http.Request) (*http.Response, error)
return resp, nil return resp, nil
} }
func TestFetchResolveChangesToLoopbackRejected(t *testing.T) {
var n int
reg := fetchReg(FetchConfig{
Resolve: func(host string) ([]net.IP, error) {
n++
if n == 1 {
return []net.IP{net.ParseIP("93.184.216.34")}, nil
}
return []net.IP{net.ParseIP("127.0.0.1")}, nil
},
})
res := runTool(t, reg, "fetch_url", `{"url":"http://flip.example/x"}`)
if !res.IsError || !strings.Contains(res.Content, "private/loopback") {
t.Fatalf("second resolve to loopback must be rejected, got %+v", res)
}
if n < 2 {
t.Fatalf("dial should resolve again, got %d calls", n)
}
}
func TestFetchHostCheckRedirectStillRejectsPrivate(t *testing.T) {
reg := fetchReg(FetchConfig{
Client: &http.Client{
Transport: redirectTransport{to: "http://127.0.0.1/evil"},
CheckRedirect: func(req *http.Request, via []*http.Request) error {
return nil
},
},
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 || !strings.Contains(res.Content, "private/loopback") {
t.Fatalf("host CheckRedirect must not bypass SSRF guard, got %+v", res)
}
}
func TestFetchNotRegisteredByDefault(t *testing.T) { func TestFetchNotRegisteredByDefault(t *testing.T) {
reg := WorkspaceTools(t.TempDir()) reg := WorkspaceTools(t.TempDir())
for _, d := range reg.Decls() { for _, d := range reg.Decls() {
+11 -1
View File
@@ -6,6 +6,7 @@ import (
"fmt" "fmt"
"io" "io"
"strings" "strings"
"sync"
"time" "time"
) )
@@ -37,6 +38,8 @@ type Registry struct {
Before BeforeHook Before BeforeHook
After AfterHook After AfterHook
Audit io.Writer Audit io.Writer
// hookMu 串行 Before、After 与 Audit,避免并行工具调用交错写审计或重入钩子。
hookMu sync.Mutex
} }
func NewRegistry(root string) *Registry { func NewRegistry(root string) *Registry {
@@ -159,7 +162,10 @@ func (r *Registry) Execute(ctx context.Context, name string, args json.RawMessag
return res return res
} }
if r.Before != nil { if r.Before != nil {
if block, reason := r.Before(ctx, name, args); block { r.hookMu.Lock()
block, reason := r.Before(ctx, name, args)
r.hookMu.Unlock()
if block {
if reason == "" { if reason == "" {
reason = "blocked" reason = "blocked"
} }
@@ -173,13 +179,17 @@ func (r *Registry) Execute(ctx context.Context, name string, args json.RawMessag
res = Result{Content: err.Error(), IsError: true} res = Result{Content: err.Error(), IsError: true}
} }
if r.After != nil { if r.After != nil {
r.hookMu.Lock()
res = r.After(ctx, name, args, res) res = r.After(ctx, name, args, res)
r.hookMu.Unlock()
} }
r.audit(name, false, res.IsError, "") r.audit(name, false, res.IsError, "")
return res return res
} }
func (r *Registry) audit(name string, blocked, isError bool, note string) { func (r *Registry) audit(name string, blocked, isError bool, note string) {
r.hookMu.Lock()
defer r.hookMu.Unlock()
if r.Audit == nil { if r.Audit == nil {
return return
} }
+43
View File
@@ -7,7 +7,9 @@ import (
"os" "os"
"path/filepath" "path/filepath"
"strings" "strings"
"sync"
"testing" "testing"
"time"
) )
func TestSchemaRejectsMissingAndWrongType(t *testing.T) { func TestSchemaRejectsMissingAndWrongType(t *testing.T) {
@@ -91,6 +93,47 @@ func TestBeforeHookBlocksAndAudits(t *testing.T) {
} }
} }
type slowHookTool struct{}
func (slowHookTool) Name() string { return "slow_hook" }
func (slowHookTool) Description() string { return "noop" }
func (slowHookTool) Parameters() json.RawMessage { return json.RawMessage(`{"type":"object","properties":{}}`) }
func (slowHookTool) Execute(context.Context, json.RawMessage) (Result, error) {
return Result{Content: "ok"}, nil
}
func TestHooksDoNotRunConcurrently(t *testing.T) {
reg := NewRegistry(t.TempDir())
reg.Register(slowHookTool{})
var mu sync.Mutex
inFlight, maxInFlight := 0, 0
reg.Before = func(context.Context, string, json.RawMessage) (bool, string) {
mu.Lock()
inFlight++
if inFlight > maxInFlight {
maxInFlight = inFlight
}
mu.Unlock()
time.Sleep(20 * time.Millisecond)
mu.Lock()
inFlight--
mu.Unlock()
return false, ""
}
var wg sync.WaitGroup
for i := 0; i < 2; i++ {
wg.Add(1)
go func() {
defer wg.Done()
reg.Execute(context.Background(), "slow_hook", []byte(`{}`))
}()
}
wg.Wait()
if maxInFlight != 1 {
t.Fatalf("hooks overlapped: max in flight %d", maxInFlight)
}
}
func TestAfterHookRewrites(t *testing.T) { func TestAfterHookRewrites(t *testing.T) {
dir := t.TempDir() dir := t.TempDir()
reg := WorkspaceTools(dir) reg := WorkspaceTools(dir)