diff --git a/cmd/agent-native/main.go b/cmd/agent-native/main.go index 51748ea..e82185d 100644 --- a/cmd/agent-native/main.go +++ b/cmd/agent-native/main.go @@ -109,7 +109,15 @@ func liveStream(ref llm.ModelRef) llm.StreamFunc { if !ok { 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) } diff --git a/deps/agent-intent b/deps/agent-intent index 3674dc4..3130190 160000 --- a/deps/agent-intent +++ b/deps/agent-intent @@ -1 +1 @@ -Subproject commit 3674dc4bfd0b49191d21f0afab4002381e41a429 +Subproject commit 3130190acb396589d6022c243bbe07ce7c31c881 diff --git a/deps/error-code-system b/deps/error-code-system index dbeed2f..a3ee67b 160000 --- a/deps/error-code-system +++ b/deps/error-code-system @@ -1 +1 @@ -Subproject commit dbeed2f1b67eb417ce089e7033a228578fbbeb03 +Subproject commit a3ee67bba269e1276c14e22c026171ecf2cf6190 diff --git a/deps/logging-system b/deps/logging-system index 2c5716f..49b5d8b 160000 --- a/deps/logging-system +++ b/deps/logging-system @@ -1 +1 @@ -Subproject commit 2c5716fdcffbc4696fc0371afc70f39c44174201 +Subproject commit 49b5d8be604d05e4246dffa0d10fa1c2b7ad45d7 diff --git a/runtime/llm/providers.go b/runtime/llm/providers.go index 5121175..6d9a048 100644 --- a/runtime/llm/providers.go +++ b/runtime/llm/providers.go @@ -137,10 +137,10 @@ type Provider struct { // Presets are OpenAI-compatible or native endpoints that only need an API key. func Presets() []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-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: "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"}, @@ -193,6 +193,17 @@ func Lookup(name string, custom []Provider) (Provider, bool) { 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 { if len(m.Images) == 0 { return m.Content @@ -250,6 +261,9 @@ func BuildBody(protocol Protocol, model string, messages []Message, tools []Tool body["reasoning"] = map[string]any{"effort": effort} } return json.Marshal(body) + // ProtocolAnthropic 与 ProtocolGoogle 只供离线组包。 + // 当前 Client 只发 OpenAI chat completions，这两支没有接到客户端上。 + // 原生协议是后续阶段的事；预设表里的 anthropic / google 已改走兼容端点。 case ProtocolAnthropic: sys, msgs := anthropicMessages(messages) body := map[string]any{ diff --git a/runtime/llm/providers_test.go b/runtime/llm/providers_test.go index 5a2de99..27a184d 100644 --- a/runtime/llm/providers_test.go +++ b/runtime/llm/providers_test.go @@ -3,6 +3,8 @@ package llm import ( "context" "encoding/json" + "net/http" + "net/http/httptest" "os" "path/filepath" "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) { n := 0 base := StreamFunc(func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event { diff --git a/runtime/loop.go b/runtime/loop.go index 529dc70..b79b9bb 100644 --- a/runtime/loop.go +++ b/runtime/loop.go @@ -54,7 +54,9 @@ type Config struct { ToolExecution string GetSteering func() []llm.Message GetFollowUp func() []llm.Message - OnEvent func(Event) + // OnEvent 观察循环事件。同一次 Run 内不会并发调用 + // （与 BeforeTool、AfterTool 共用一把锁串行）。 + OnEvent func(Event) // ThinkingLevel is forwarded on each provider request. ThinkingLevel string // Resolve, when set, overrides Model and ThinkingLevel at the start of each turn. @@ -73,10 +75,12 @@ type Config struct { // nil 表示不询问。内核最多接受 3 次 Continue。 Stop func(last llm.Message, info StopInfo) StopDecision // BeforeTool 返回错误时不执行工具，但仍写一条 IsError 的 tool 结果。 + // 同一次 Run 内与 OnEvent、AfterTool 串行，不会并发调用。 BeforeTool func(call llm.ToolCall) error // UnknownTool 处理注册表里没有的调用。nil 时仍返回 unknown tool。 UnknownTool func(call llm.ToolCall) tools.Result // AfterTool 只观察工具结果，不能改写。 + // 同一次 Run 内与 OnEvent、BeforeTool 串行，不会并发调用。 AfterTool func(call llm.ToolCall, res llm.Message) // Reminders 在每次调用模型前提供提醒。正文包在 system-reminder 里， // 只进入本次请求，不写进 Session。 @@ -136,6 +140,7 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { if cfg.ToolExecution == "" { cfg.ToolExecution = ExecSequential } + serializeHostCallbacks(&cfg) var messages []llm.Message if cfg.Session != nil { 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 cfg.Budget != nil && cfg.Budget.MaxToolCalls > 0 && toolCalls+len(last.ToolCalls) > cfg.Budget.MaxToolCalls { reason := fmt.Sprintf("budget_exceeded: tool calls %d+%d > max %d", toolCalls, len(last.ToolCalls), cfg.Budget.MaxToolCalls) + var err error + messages, err = appendSkippedToolResults(cfg, last.ToolCalls, messages) + if err != nil { + return Result{}, err + } return budgetStop(cfg, last, messages, usage, reason), nil } toolCalls += len(last.ToolCalls) @@ -352,6 +362,56 @@ func consume(ctx context.Context, cfg Config, messages []llm.Message) (llm.Event 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 { if cfg.ToolExecution == ExecParallel && len(assistant.ToolCalls) > 1 { return executeParallel(ctx, cfg, assistant) diff --git a/runtime/loop_budget_test.go b/runtime/loop_budget_test.go index 46a800e..5da9702 100644 --- a/runtime/loop_budget_test.go +++ b/runtime/loop_budget_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "strings" + "sync" "testing" "time" @@ -66,15 +67,108 @@ func TestLoopBudgetToolCallCap(t *testing.T) { if !sawBudgetEvent { t.Fatal("expected budget_exceeded event") } - // 第二批工具调用不该被执行（预算在执行前拦截）：transcript 只含一次 tool 结果。 + // 第一批执行 1 次；第二批超限，两条合成错误结果，不执行。 toolResults := 0 + skipped := 0 for _, m := range res.Messages { - if m.Role == llm.RoleTool { - toolResults++ + if m.Role != llm.RoleTool { + continue + } + toolResults++ + if m.IsError && m.Content == toolBudgetSkipContent { + skipped++ } } - if toolResults != 1 { - t.Fatalf("expected exactly 1 tool result, got %d", toolResults) + if toolResults != 3 || skipped != 2 { + 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) } } diff --git a/runtime/mcp/fakeserver.go b/runtime/mcp/fakeserver.go index 1d6c81f..3b3880c 100644 --- a/runtime/mcp/fakeserver.go +++ b/runtime/mcp/fakeserver.go @@ -8,6 +8,7 @@ import ( "bufio" "encoding/json" "os" + "time" ) // RunFakeServer 在 stdin/stdout 上运行最小 MCP server（echo / boom / 未知方法）， @@ -67,6 +68,9 @@ func fakeServerHandle(req rpcRequest) rpcResponse { Content: []contentBlock{{Type: "text", Text: "boom failed"}}, IsError: true, }) + case "slow": + time.Sleep(2 * time.Second) + return respond(callToolResult{Content: []contentBlock{{Type: "text", Text: "slow done"}}}) default: return fail(-32602, "unknown tool "+p.Name) } @@ -87,6 +91,11 @@ func fakeServerTools() []toolInfo { Description: "Always fails", InputSchema: json.RawMessage(`{"type":"object","properties":{}}`), }, + { + Name: "slow", + Description: "Responds after a delay", + InputSchema: json.RawMessage(`{"type":"object","properties":{}}`), + }, } } diff --git a/runtime/mcp/mcp_test.go b/runtime/mcp/mcp_test.go index 92ed1e1..fd3a5ce 100644 --- a/runtime/mcp/mcp_test.go +++ b/runtime/mcp/mcp_test.go @@ -45,7 +45,7 @@ func TestStdioClientHandshakeListAndCall(t *testing.T) { defer c.Close() 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) } 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) { var req rpcRequest 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) } 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) } res, err := c.CallTool(ctx, "echo", json.RawMessage(`{"text":"net"}`)) @@ -130,7 +156,7 @@ func TestHTTPClientSSEMode(t *testing.T) { } defer c.Close() 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) } } diff --git a/runtime/mcp/stdio.go b/runtime/mcp/stdio.go index 2e5bcda..1eeabea 100644 --- a/runtime/mcp/stdio.go +++ b/runtime/mcp/stdio.go @@ -12,14 +12,23 @@ import ( ) // stdioClient 通过「子进程 stdin/stdout 按行分隔的 JSON-RPC」与 server 通信。 -// 同一时刻只允许一个在途请求（串行锁），匹配简单且足以支撑智能体工具调用。 +// 一个长驻读循环按 id 把响应分发给各请求；写 stdin 仍串行。 +// 取消会摘掉等待项，迟到的回复直接丢弃。读循环退出时，在等的请求全部失败，客户端标记为已死。 type stdioClient struct { - cmd *exec.Cmd - stdin io.WriteCloser - stdout *bufio.Reader - mu sync.Mutex - nextID int64 - closed bool + cmd *exec.Cmd + stdin io.WriteCloser + stdout *bufio.Reader + writeMu sync.Mutex + mu sync.Mutex + pending map[int64]chan stdioResult + nextID int64 + closed bool + readerOnce sync.Once +} + +type stdioResult struct { + resp rpcResponse + err error } // DialStdio 启动 server 子进程并完成 initialize 握手。 @@ -43,7 +52,13 @@ func DialStdio(ctx context.Context, cfg ServerConfig) (*stdioClient, error) { 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)} + 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 { _ = c.Close() return nil, err @@ -104,11 +119,12 @@ func (c *stdioClient) CallTool(ctx context.Context, name string, args json.RawMe func (c *stdioClient) Close() error { c.mu.Lock() - defer c.mu.Unlock() if c.closed { + c.mu.Unlock() return nil } c.closed = true + c.mu.Unlock() _ = c.stdin.Close() if c.cmd.Process != nil { _ = 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 { + c.startReader() + id := atomic.AddInt64(&c.nextID, 1) + ch := make(chan stdioResult, 1) c.mu.Lock() - defer c.mu.Unlock() if c.closed { + c.mu.Unlock() 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} body, err := json.Marshal(req) if err != nil { + c.dropPending(id) return err } 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) } - line, err := c.readResponseLine(ctx) - if err != nil { - return fmt.Errorf("mcp read %s: %w", method, err) + + var got stdioResult + select { + case got = <-ch: + case <-ctx.Done(): + c.dropPending(id) + return fmt.Errorf("mcp read %s: %w", method, ctx.Err()) } - var resp rpcResponse - if err := json.Unmarshal(line, &resp); err != nil { - return fmt.Errorf("mcp %s: bad response frame: %w", method, err) + if got.err != nil { + return fmt.Errorf("mcp read %s: %w", method, got.err) } + resp := got.resp if resp.ID == nil || *resp.ID != id { 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 } -// 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 - } - // 通知或日志行：忽略。 +func (c *stdioClient) dropPending(id int64) { + c.mu.Lock() + delete(c.pending, id) + c.mu.Unlock() +} + +func (c *stdioClient) startReader() { + c.readerOnce.Do(func() { + go c.readLoop() + }) +} + +func (c *stdioClient) readLoop() { + var fatal error + for { + line, err := c.stdout.ReadBytes('\n') + if err != nil { + fatal = err + break } - }() - select { - case r := <-done: - return r.line, r.err - case <-ctx.Done(): - return nil, ctx.Err() + var resp rpcResponse + if json.Unmarshal(line, &resp) != nil || resp.ID == nil { + continue + } + c.mu.Lock() + ch := c.pending[*resp.ID] + if ch != nil { + delete(c.pending, *resp.ID) + } + c.mu.Unlock() + if ch == nil { + continue + } + ch <- stdioResult{resp: resp} + } + 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 { c.mu.Lock() - defer c.mu.Unlock() - if c.closed { + closed := c.closed + c.mu.Unlock() + if closed { return fmt.Errorf("mcp client closed") } body, err := json.Marshal(rpcNotification{JSONRPC: "2.0", Method: method}) @@ -193,6 +239,8 @@ func (c *stdioClient) notify(method string) error { return err } body = append(body, '\n') + c.writeMu.Lock() + defer c.writeMu.Unlock() _, err = c.stdin.Write(body) return err } diff --git a/runtime/tools/fetch.go b/runtime/tools/fetch.go index 274bc1c..afc6f33 100644 --- a/runtime/tools/fetch.go +++ b/runtime/tools/fetch.go @@ -17,6 +17,10 @@ import ( // FetchConfig 声明 fetch_url 的运行条件；未注册（组合根未提供）则工具不存在。 type FetchConfig struct { // Client 可选注入（测试替身）；缺省构造 20s 超时客户端。 + // Transport 为 nil 或 *http.Transport 时，Clone 后替换 DialContext： + // 拨号时再次解析，只连接通过 SSRF 校验的 IP。 + // 自定义 RoundTripper 无法替换拨号，只做请求前预校验；实现不得在 + // RoundTrip 里重新解析到未校验的地址。 Client *http.Client // MaxBytes 响应体上限（默认 2 MiB）。 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") 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) - }, - } + hostRedirect := client.CheckRedirect + transport, _ := pinFetchTransport(client.Transport, resolve) + client2 := &http.Client{ + Transport: 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") + } + 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 { 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 } +// 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 硬边界（私网/回环/链路本地）。 func validateFetchTarget(u *url.URL, allowed map[string]bool, resolve func(string) ([]net.IP, error)) error { if u.Scheme != "http" && u.Scheme != "https" { diff --git a/runtime/tools/fetch_test.go b/runtime/tools/fetch_test.go index 842f8f8..7f1f6e5 100644 --- a/runtime/tools/fetch_test.go +++ b/runtime/tools/fetch_test.go @@ -144,6 +144,45 @@ func (s redirectTransport) RoundTrip(req *http.Request) (*http.Response, error) 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) { reg := WorkspaceTools(t.TempDir()) for _, d := range reg.Decls() { diff --git a/runtime/tools/registry.go b/runtime/tools/registry.go index f418af5..ab4a3c3 100644 --- a/runtime/tools/registry.go +++ b/runtime/tools/registry.go @@ -6,6 +6,7 @@ import ( "fmt" "io" "strings" + "sync" "time" ) @@ -37,6 +38,8 @@ type Registry struct { Before BeforeHook After AfterHook Audit io.Writer + // hookMu 串行 Before、After 与 Audit，避免并行工具调用交错写审计或重入钩子。 + hookMu sync.Mutex } func NewRegistry(root string) *Registry { @@ -159,7 +162,10 @@ func (r *Registry) Execute(ctx context.Context, name string, args json.RawMessag return res } 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 == "" { 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} } if r.After != nil { + r.hookMu.Lock() res = r.After(ctx, name, args, res) + r.hookMu.Unlock() } r.audit(name, false, res.IsError, "") return res } func (r *Registry) audit(name string, blocked, isError bool, note string) { + r.hookMu.Lock() + defer r.hookMu.Unlock() if r.Audit == nil { return } diff --git a/runtime/tools/tools_test.go b/runtime/tools/tools_test.go index 6f1f0c9..aa711c3 100644 --- a/runtime/tools/tools_test.go +++ b/runtime/tools/tools_test.go @@ -7,7 +7,9 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" + "time" ) 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) { dir := t.TempDir() reg := WorkspaceTools(dir)