diff --git a/internal/claudebridge/mcp.go b/internal/claudebridge/mcp.go index da8dd4a6..c372c531 100644 --- a/internal/claudebridge/mcp.go +++ b/internal/claudebridge/mcp.go @@ -14,7 +14,9 @@ import ( "io" "net/http" "os" + "reflect" "sync" + "time" ) type tool struct { @@ -39,8 +41,15 @@ type callbackRequest struct { type callbackResponse struct { Content any `json:"content"` IsError bool `json:"is_error,omitempty"` + // Tools, when set, replace the tools offered: the client loaded one + // mid-turn (Claude Code's ToolSearch) + Tools []tool `json:"tools,omitempty"` } +// listWait is how long a call's answer waits for Claude Code to list the +// tools again, once told they changed. +var listWait = 5 * time.Second + // RunMCP runs the hidden stdio MCP subprocess. args are callback URL and tools // JSON path, or, for an agent whose MCP servers are set once for every run // (Grok), MAGPIE_MCP_CALLBACK and MAGPIE_MCP_TOOLS in the environment it @@ -72,6 +81,25 @@ func RunMCP(args []string) error { _ = enc.Encode(v) outMu.Unlock() } + // listed is closed by the next tools/list, once the tools changed + var toolsMu sync.Mutex + var listed chan struct{} + retool := func(next []tool) { + toolsMu.Lock() + if reflect.DeepEqual(next, tools) { + toolsMu.Unlock() + return + } + tools = next + done := make(chan struct{}) + listed = done + toolsMu.Unlock() + respond(map[string]any{"jsonrpc": "2.0", "method": "notifications/tools/list_changed"}) + select { + case <-done: + case <-time.After(listWait): + } + } for in.Scan() { var req rpcRequest @@ -89,11 +117,18 @@ func RunMCP(args []string) error { case "initialize": result = map[string]any{ "protocolVersion": "2025-06-18", - "capabilities": map[string]any{"tools": map[string]any{}}, + "capabilities": map[string]any{"tools": map[string]any{"listChanged": true}}, "serverInfo": map[string]any{"name": "magpie", "version": "1"}, } case "tools/list": + toolsMu.Lock() result = map[string]any{"tools": tools} + if listed != nil { + // answered before the calls waiting on it go on + defer close(listed) + listed = nil + } + toolsMu.Unlock() case "tools/call": var p struct { Name string `json:"name"` @@ -129,6 +164,11 @@ func RunMCP(args []string) error { rpcErr = map[string]any{"code": -32000, "message": err.Error()} break } + if len(cb.Tools) > 0 { + // Claude Code takes up the tools before the result: its + // next request offers them + retool(cb.Tools) + } result = map[string]any{"content": cb.Content, "isError": cb.IsError} default: rpcErr = map[string]any{"code": -32601, "message": "method not found"} diff --git a/internal/claudebridge/mcp_test.go b/internal/claudebridge/mcp_test.go new file mode 100644 index 00000000..9313330b --- /dev/null +++ b/internal/claudebridge/mcp_test.go @@ -0,0 +1,88 @@ +package claudebridge + +import ( + "bufio" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" + "time" +) + +// A call's answer that brings new tools (the client loaded one with Claude +// Code's ToolSearch) has the helper tell Claude Code its tools changed, and +// is given only once Claude Code listed them again: its next request then +// offers the tool, as Claude Code 2.1 does with a list_changed mid-turn. +func TestNewToolsListedBeforeTheAnswer(t *testing.T) { + cb := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]any{ + "content": []map[string]any{{"type": "text", "text": "Tool WebFetch is loaded and can be called now."}}, + "tools": []map[string]any{ + {"name": "ToolSearch", "inputSchema": map[string]any{"type": "object"}}, + {"name": "WebFetch", "inputSchema": map[string]any{"type": "object"}}, + }, + }) + })) + defer cb.Close() + dir := t.TempDir() + toolsPath := filepath.Join(dir, "tools.json") + os.WriteFile(toolsPath, []byte(`[{"name":"ToolSearch","inputSchema":{"type":"object"}}]`), 0o600) + + inR, inW, _ := os.Pipe() + outR, outW, _ := os.Pipe() + oldIn, oldOut := os.Stdin, os.Stdout + os.Stdin, os.Stdout = inR, outW + defer func() { os.Stdin, os.Stdout = oldIn, oldOut }() + go func() { _ = RunMCP([]string{cb.URL, toolsPath}); outW.Close() }() + lines := make(chan map[string]any, 16) + go func() { + sc := bufio.NewScanner(outR) + sc.Buffer(make([]byte, 64<<10), 1<<20) + for sc.Scan() { + var m map[string]any + _ = json.Unmarshal(sc.Bytes(), &m) + lines <- m + } + close(lines) + }() + send := func(s string) { inW.Write([]byte(s + "\n")) } + next := func() map[string]any { + t.Helper() + select { + case m := <-lines: + return m + case <-time.After(3 * time.Second): + t.Fatal("no message from the helper") + return nil + } + } + + send(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}`) + caps, _ := next()["result"].(map[string]any)["capabilities"].(map[string]any)["tools"].(map[string]any) + if caps["listChanged"] != true { + t.Errorf("tools capability = %v, no listChanged", caps) + } + send(`{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"ToolSearch","arguments":{"query":"select:WebFetch"},"_meta":{"claudecode/toolUseId":"toolu_1"}}}`) + if m := next(); m["method"] != "notifications/tools/list_changed" { + t.Fatalf("first after the call = %v, not list_changed", m) + } + select { + case m := <-lines: + t.Fatalf("answered before the tools were listed again: %v", m) + case <-time.After(200 * time.Millisecond): + } + send(`{"jsonrpc":"2.0","id":3,"method":"tools/list"}`) + list := next() + if list["id"] != float64(3) { + t.Fatalf("expected the tools/list answer first, got %v", list) + } + if tools := list["result"].(map[string]any)["tools"].([]any); len(tools) != 2 { + t.Fatalf("tools listed = %v", tools) + } + if m := next(); m["id"] != float64(2) || m["result"] == nil { + t.Fatalf("call answer = %v", m) + } + inW.Close() +} diff --git a/internal/gateway/claude_subscription.go b/internal/gateway/claude_subscription.go index 5d6d64b5..9cdae0ac 100644 --- a/internal/gateway/claude_subscription.go +++ b/internal/gateway/claude_subscription.go @@ -93,7 +93,7 @@ type subscriptionRun struct { tree *proc.Tree // cmd once started, with all it starts tmp string - // the tools it was started with: its agent is told them once + // the tools its agent was told of, as it started and since tools map[string]bool mu sync.Mutex @@ -148,6 +148,10 @@ const waitTool = "magpie_wait" type mcpToolResult struct { Content []map[string]any `json:"content"` IsError bool `json:"is_error,omitempty"` + // Tools, when set, is the run's tools from now on: the MCP helper tells + // Claude Code its list changed, and answers the call once Claude Code + // has listed them again + Tools []bridgeTool `json:"tools,omitempty"` } type bridgeTool struct { @@ -1046,6 +1050,8 @@ func (b *subscriptionBridge) findRun(req *Request) (*subscriptionRun, []Part) { // offers says the run's agent was told every one of tools. func (r *subscriptionRun) offers(tools []Tool) bool { + r.mu.Lock() + defer r.mu.Unlock() if r.tools == nil { return true } @@ -1058,8 +1064,9 @@ func (r *subscriptionRun) offers(tools []Tool) bool { } // continueWith hands the agent its tool results: at once for the calls it -// is waiting on, the others kept until it makes them. -func (r *subscriptionRun) continueWith(results []Part) (<-chan Event, error) { +// is waiting on, the others kept until it makes them. tools, when set, are +// the agent's tools from now on, handed over with the results. +func (r *subscriptionRun) continueWith(results []Part, tools []bridgeTool) (<-chan Event, error) { if r.timer != nil { r.timer.Reset(30 * time.Minute) } @@ -1073,6 +1080,7 @@ func (r *subscriptionRun) continueWith(results []Part) (<-chan Event, error) { delivered := 0 for _, p := range results { result := mcpResult(p) + result.Tools = tools r.mu.Lock() waiter := r.pending[p.CallID] delete(r.pending, p.CallID) @@ -1385,19 +1393,27 @@ func (s *Server) serveSubscription(w http.ResponseWriter, r *http.Request, from } run, results := s.subscription.findRun(req) + // the client offers a tool the run's agent was never told of, as Claude + // Code's ToolSearch loads a deferred one mid-turn: the run is handed it + // with the results, as its MCP server's tools changed. A run started + // anew would be told the conversation in one message, its tool calls as + // text, and its model went on to write its own calls as text. + var more []bridgeTool if run != nil && !run.offers(req.Tools) { - // the client offers a tool the run's agent was never told of, as - // Claude Code's ToolSearch loads a deferred one (WebSearch) mid-turn: - // a run started now has it, the conversation told over - run.abort() - run = nil + more = bridgeTools(req) } var events <-chan Event if run != nil { - if events, err = run.continueWith(results); err != nil { + if events, err = run.continueWith(results, more); err != nil { // it ended while it waited: a new one is told the whole // conversation run = nil + } else if more != nil { + run.mu.Lock() + for _, t := range req.Tools { + run.tools[t.Name] = true + } + run.mu.Unlock() } } if run == nil { diff --git a/internal/gateway/claude_subscription_calls_test.go b/internal/gateway/claude_subscription_calls_test.go index 57bf4ed3..db7cfa73 100644 --- a/internal/gateway/claude_subscription_calls_test.go +++ b/internal/gateway/claude_subscription_calls_test.go @@ -58,7 +58,7 @@ func TestResultForACallNotMadeYetWaitsForIt(t *testing.T) { if found != run || len(results) != 2 { t.Fatalf("findRun = %v, %d results", found == run, len(results)) } - if _, err := run.continueWith(results); err != nil { + if _, err := run.continueWith(results, nil); err != nil { t.Fatal(err) } if got := <-first; got != "one" { diff --git a/internal/gateway/claude_subscription_images_test.go b/internal/gateway/claude_subscription_images_test.go index 7a9ba37f..9a28b21b 100644 --- a/internal/gateway/claude_subscription_images_test.go +++ b/internal/gateway/claude_subscription_images_test.go @@ -83,7 +83,7 @@ func TestSubscriptionToolResultImages(t *testing.T) { if found != run { t.Fatal("no run for call_1") } - if _, err := run.continueWith(results); err != nil { + if _, err := run.continueWith(results, nil); err != nil { t.Fatal(err) } var got mcpToolResult diff --git a/internal/gateway/claude_subscription_toolsearch_test.go b/internal/gateway/claude_subscription_toolsearch_test.go new file mode 100644 index 00000000..78e8a0c7 --- /dev/null +++ b/internal/gateway/claude_subscription_toolsearch_test.go @@ -0,0 +1,102 @@ +package gateway + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/yetone/magpie/internal/provider" +) + +// Claude Code with ENABLE_TOOL_SEARCH loads a deferred tool by a ToolSearch +// whose result comes back with the tool now offered. The run waiting on the +// ToolSearch call is handed the tool with the result, and goes on; it isn't +// replaced by one told the conversation in one message, its tool calls as +// text, from which the model went on to write its own calls as text. +func TestToolSearchLoadKeepsTheRun(t *testing.T) { + s := New() + b := s.subscription + run := &subscriptionRun{bridge: b, token: "tok", pending: map[string]chan mcpToolResult{}, tools: map[string]bool{"ToolSearch": true}} + b.runs["tok"] = run + mux := http.NewServeMux() + mux.HandleFunc("POST /cb/{token}", b.mcpCall) + srv := httptest.NewServer(mux) + defer srv.Close() + + answer := make(chan mcpToolResult, 1) + go func() { + res, err := http.Post(srv.URL+"/cb/tok", "application/json", strings.NewReader(`{"tool_call_id":"toolu_1","name":"ToolSearch","arguments":{"query":"select:WebFetch"}}`)) + if err != nil { + answer <- mcpToolResult{} + return + } + defer res.Body.Close() + var r mcpToolResult + _ = json.NewDecoder(res.Body).Decode(&r) + answer <- r + }() + for deadline := time.Now().Add(2 * time.Second); ; { + b.mu.Lock() + n := len(b.calls) + b.mu.Unlock() + if n == 1 { + break + } + if time.Now().After(deadline) { + t.Fatal("the ToolSearch call was not registered") + } + time.Sleep(5 * time.Millisecond) + } + + body := `{"model":"claude-sonnet-5","max_tokens":100,"tools":[ + {"name":"ToolSearch","input_schema":{"type":"object"}}, + {"name":"WebFetch","description":"fetch a page","input_schema":{"type":"object","properties":{"url":{"type":"string"}}}}, + {"name":"DeferredToolPlaceholder","input_schema":{"type":"object"},"defer_loading":true}], + "messages":[{"role":"user","content":"read example.com"}, + {"role":"assistant","content":[{"type":"tool_use","id":"toolu_1","name":"ToolSearch","input":{"query":"select:WebFetch"}}]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":[{"type":"tool_reference","tool_name":"WebFetch"}]}]}]}` + started := false + start := func(ctx context.Context, req *Request) (*subscriptionRun, <-chan Event, error) { + started = true + return nil, nil, errors.New("a new run was started") + } + done := make(chan struct{}) + go func() { + defer close(done) + w := httptest.NewRecorder() + r := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(body)) + s.serveSubscription(w, r, provider.Anthropic, "Claude Code", "claude-sonnet-5", []byte(body), &Usage{}, start) + }() + + var got mcpToolResult + select { + case got = <-answer: + case <-done: + t.Fatalf("the request ended before the run had its result (started anew: %v)", started) + case <-time.After(3 * time.Second): + t.Fatal("the ToolSearch call never had its result") + } + var names []string + for _, tl := range got.Tools { + names = append(names, tl.Name) + } + if strings.Join(names, ",") != "ToolSearch,WebFetch" { + t.Errorf("tools handed over = %v", names) + } + if len(got.Content) == 0 || !strings.Contains(got.Content[0]["text"].(string), "WebFetch is loaded") { + t.Errorf("result = %v", got.Content) + } + if !run.offers([]Tool{{Name: "WebFetch"}}) { + t.Error("the run doesn't know it has WebFetch now") + } + run.finish() + <-done + if started { + t.Error("a new run was started for the loaded tool") + } +}