diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index e24d4196b5164..27ed0964edd60 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -20554,6 +20554,7 @@ const docTemplate = `{ "workspace-usage", "oauth2", "mcp-server-http", + "mcp-tool-search", "workspace-build-updates", "nats_pubsub", "workspace-capable-licensing", @@ -20570,6 +20571,7 @@ const docTemplate = `{ "ExperimentChatVirtualDesktop": "Enables virtual desktop and computer use provider for agents.", "ExperimentExample": "This isn't used for anything.", "ExperimentMCPServerHTTP": "Enables the MCP HTTP server functionality.", + "ExperimentMCPToolSearch": "Defers MCP tool schemas behind a searchable catalog in agent chats.", "ExperimentNATSPubsub": "Enables embedded NATS pubsub.", "ExperimentNotifications": "Sends notifications via SMTP and webhooks following certain events.", "ExperimentOAuth2": "Enables OAuth2 provider functionality.", @@ -20584,6 +20586,7 @@ const docTemplate = `{ "Enables the new workspace usage tracking.", "Enables OAuth2 provider functionality.", "Enables the MCP HTTP server functionality.", + "Defers MCP tool schemas behind a searchable catalog in agent chats.", "Enables publishing workspace build updates to the all builds pubsub channel.", "Enables embedded NATS pubsub.", "Counts only users holding the workspace-create permission toward the license seat limit.", @@ -20599,6 +20602,7 @@ const docTemplate = `{ "ExperimentWorkspaceUsage", "ExperimentOAuth2", "ExperimentMCPServerHTTP", + "ExperimentMCPToolSearch", "ExperimentWorkspaceBuildUpdates", "ExperimentNATSPubsub", "ExperimentWorkspaceCapableLicensing", diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 0b18135d1cff4..53df566fc82c5 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -18676,6 +18676,7 @@ "workspace-usage", "oauth2", "mcp-server-http", + "mcp-tool-search", "workspace-build-updates", "nats_pubsub", "workspace-capable-licensing", @@ -18692,6 +18693,7 @@ "ExperimentChatVirtualDesktop": "Enables virtual desktop and computer use provider for agents.", "ExperimentExample": "This isn't used for anything.", "ExperimentMCPServerHTTP": "Enables the MCP HTTP server functionality.", + "ExperimentMCPToolSearch": "Defers MCP tool schemas behind a searchable catalog in agent chats.", "ExperimentNATSPubsub": "Enables embedded NATS pubsub.", "ExperimentNotifications": "Sends notifications via SMTP and webhooks following certain events.", "ExperimentOAuth2": "Enables OAuth2 provider functionality.", @@ -18706,6 +18708,7 @@ "Enables the new workspace usage tracking.", "Enables OAuth2 provider functionality.", "Enables the MCP HTTP server functionality.", + "Defers MCP tool schemas behind a searchable catalog in agent chats.", "Enables publishing workspace build updates to the all builds pubsub channel.", "Enables embedded NATS pubsub.", "Counts only users holding the workspace-create permission toward the license seat limit.", @@ -18721,6 +18724,7 @@ "ExperimentWorkspaceUsage", "ExperimentOAuth2", "ExperimentMCPServerHTTP", + "ExperimentMCPToolSearch", "ExperimentWorkspaceBuildUpdates", "ExperimentNATSPubsub", "ExperimentWorkspaceCapableLicensing", diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 41966f45daa17..c5df8e0a04baf 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -3048,8 +3048,7 @@ type Config struct { Clock quartz.Clock AIBridgeTransportFactory *atomic.Pointer[aibridge.TransportFactory] Experiments codersdk.Experiments - - PrometheusRegistry prometheus.Registerer + PrometheusRegistry prometheus.Registerer AgentCapacityUnlock AgentCapacityUnlock diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index ae5b9bcd8fd39..0025d6a2d6205 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -798,6 +798,7 @@ func TestAllowedExploreToolNames(t *testing.T) { newTestAgentTool("read_skill"), newTestAgentTool("read_skill_file"), newTestAgentTool("ask_user_question"), + newTestAgentTool(chattool.FindToolsName), }) require.Equal(t, []string{ @@ -812,6 +813,7 @@ func TestAllowedExploreToolNames(t *testing.T) { require.NotContains(t, got, "start_workspace") require.NotContains(t, got, "stop_workspace") require.NotContains(t, got, "ask_user_question") + require.NotContains(t, got, chattool.FindToolsName) } func TestAllowedBehaviorToolNames(t *testing.T) { diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 53e4a4fd08743..11865529fb139 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -60,6 +60,7 @@ import ( "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/workspacesdk" "github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock" + "github.com/coder/coder/v2/codersdk/x/agenthooks" "github.com/coder/coder/v2/provisioner/echo" proto "github.com/coder/coder/v2/provisionersdk/proto" "github.com/coder/coder/v2/testutil" @@ -802,6 +803,7 @@ func TestExploreChatUsesPersistedMCPSnapshot(t *testing.T) { factory := chattest.NewMockAIBridgeTransport(t, openAIURL) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + withoutMCPToolSearch(cfg) cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(factory) cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { require.Equal(t, dbAgent.ID, agentID) @@ -1075,6 +1077,7 @@ func TestExploreChatSendMessageCannotMutateMCPSnapshot(t *testing.T) { factory := chattest.NewMockAIBridgeTransport(t, openAIURL) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + withoutMCPToolSearch(cfg) cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(factory) }) @@ -1237,6 +1240,7 @@ func TestPlanModeRootChatAllowsApprovedExternalMCPTools(t *testing.T) { factory := chattest.NewMockAIBridgeTransport(t, openAIURL) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + withoutMCPToolSearch(cfg) cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(factory) cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { require.Equal(t, dbAgent.ID, agentID) @@ -8494,6 +8498,14 @@ func newDebugEnabledTestServer( return server } +// withoutMCPToolSearch disables the mcp-tool-search experiment so a +// test exercises direct MCP tool advertisement instead of deferral. +func withoutMCPToolSearch(cfg *chatd.Config) { + cfg.Experiments = slices.DeleteFunc(slices.Clone(cfg.Experiments), func(experiment codersdk.Experiment) bool { + return experiment == codersdk.ExperimentMCPToolSearch + }) +} + // newActiveTestServer creates a chatd server that actively polls for // and processes pending chats. Use this instead of newTestServer when // the test needs the chat loop to actually run. Optional config @@ -10410,6 +10422,415 @@ func (d *panicOnInTxDB) InTx(f func(database.Store) error, opts *database.TxOpti return d.Store.InTx(f, opts) } +func TestMCPToolSearchGenerationFlows(t *testing.T) { + t.Parallel() + + t.Run("search activates tools within and across turns", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + mcpSrv := newTestMCPServer("search-mcp") + addTestMCPTextTool(mcpSrv, "alpha", "Alpha deferred action", "alpha: ") + addTestMCPTextTool(mcpSrv, "beta", "Beta deferred action", "beta: ") + mcpTS := httptest.NewServer(testMCPHTTPHandler(mcpSrv)) + t.Cleanup(mcpTS.Close) + + var ( + streamCount atomic.Int32 + requestsMu sync.Mutex + requests []recordedOpenAIRequest + ) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + requestsMu.Lock() + requests = append(requests, recordOpenAIRequest(req)) + requestsMu.Unlock() + switch streamCount.Add(1) { + case 1: + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk(chattool.FindToolsName, `{"names":["search-mcp__alpha","search-mcp__beta"]}`), + ) + default: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + } + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Search MCP", + Slug: "search-mcp", + Url: mcpTS.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + }) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "deferred search", + ModelConfigID: model.ID, + MCPServerIDs: []uuid.UUID{mcpConfig.ID}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("find deferred actions"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + requestsMu.Lock() + recorded := append([]recordedOpenAIRequest(nil), requests...) + requestsMu.Unlock() + require.Len(t, recorded, 3) + require.Contains(t, recorded[0].Tools, "read_file") + require.Contains(t, recorded[0].Tools, chattool.FindToolsName) + require.NotContains(t, recorded[0].Tools, "search-mcp__alpha") + require.NotContains(t, recorded[0].Tools, "search-mcp__beta") + for _, request := range recorded[1:] { + require.Contains(t, request.Tools, chattool.FindToolsName) + require.Contains(t, request.Tools, "search-mcp__alpha") + require.Contains(t, request.Tools, "search-mcp__beta") + require.Less(t, + slices.Index(request.Tools, "search-mcp__alpha"), + slices.Index(request.Tools, "search-mcp__beta"), + ) + } + }) + + t.Run("direct call activates schema on next step", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + mcpSrv := newTestMCPServer("direct-mcp") + addTestMCPTextTool(mcpSrv, "echo", "Echo deferred input", "echo: ") + mcpTS := httptest.NewServer(testMCPHTTPHandler(mcpSrv)) + t.Cleanup(mcpTS.Close) + + var ( + streamCount atomic.Int32 + requestsMu sync.Mutex + requests []recordedOpenAIRequest + ) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + requestsMu.Lock() + requests = append(requests, recordOpenAIRequest(req)) + requestsMu.Unlock() + if streamCount.Add(1) == 1 { + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("direct-mcp__echo", `{"input":"hello"}`), + ) + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Direct MCP", + Slug: "direct-mcp", + Url: mcpTS.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + }) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "direct deferred call", + ModelConfigID: model.ID, + MCPServerIDs: []uuid.UUID{mcpConfig.ID}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("call the deferred tool directly"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + requestsMu.Lock() + recorded := append([]recordedOpenAIRequest(nil), requests...) + requestsMu.Unlock() + require.Len(t, recorded, 2) + require.NotContains(t, recorded[0].Tools, "direct-mcp__echo") + require.Contains(t, recorded[1].Tools, "direct-mcp__echo") + require.True(t, openAIMessagesContain(recorded[1].Messages, "echo: hello")) + }) + + t.Run("partition-denied find_tools calls count toward call totals", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + mcpSrv := newTestMCPServer("count-mcp") + addTestMCPTextTool(mcpSrv, "echo", "Echo input", "echo: ") + mcpTS := httptest.NewServer(testMCPHTTPHandler(mcpSrv)) + t.Cleanup(mcpTS.Close) + + var streamCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + switch streamCount.Add(1) { + case 1: + // Malformed JSON input: partitioned into a synthetic + // denial before ExecuteLocalTools, so the tool's own + // handler and decode never see this call. + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk(chattool.FindToolsName, `{"queries":["echo"`), + ) + default: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + } + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Count MCP", + Slug: "count-mcp", + Url: mcpTS.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + }) + reg := prometheus.NewRegistry() + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.PrometheusRegistry = reg + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "count denied find_tools", + ModelConfigID: model.ID, + MCPServerIDs: []uuid.UUID{mcpConfig.ID}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("search"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + requireChatdMetricCounter(t, reg, "coderd_chatd_find_tools_calls_total", 1, nil) + }) + + t.Run("hook-denied find_tools calls count toward call totals", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + mcpSrv := newTestMCPServer("hooked-mcp") + addTestMCPTextTool(mcpSrv, "echo", "Echo input", "echo: ") + mcpTS := httptest.NewServer(testMCPHTTPHandler(mcpSrv)) + t.Cleanup(mcpTS.Close) + + var streamCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + switch streamCount.Add(1) { + case 1: + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk(chattool.FindToolsName, `{"queries":["echo"]}`), + ) + default: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + } + }) + consumer := preToolUseConsumer(t, func(data agenthooks.PreToolUseData) string { + require.Equal(t, chattool.FindToolsName, data.ToolName) + return `{"permission":{"decision":"deny","reason":"blocked by policy"}}` + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Hooked MCP", + Slug: "hooked-mcp", + Url: mcpTS.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + }) + reg := prometheus.NewRegistry() + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.PrometheusRegistry = reg + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "count hook-denied find_tools", + ModelConfigID: model.ID, + MCPServerIDs: []uuid.UUID{mcpConfig.ID}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("search"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + requireChatdMetricCounter(t, reg, "coderd_chatd_find_tools_calls_total", 1, nil) + }) + + t.Run("admission-failed find_tools calls count toward call totals", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + mcpSrv := newTestMCPServer("failing-mcp") + addTestMCPTextTool(mcpSrv, "echo", "Echo input", "echo: ") + mcpTS := httptest.NewServer(testMCPHTTPHandler(mcpSrv)) + t.Cleanup(mcpTS.Close) + + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk(chattool.FindToolsName, `{"queries":["echo"]}`), + ) + }) + // A pre_tool_use dispatch failure errors admission before the + // step commits, so the call never reaches executeLocalTools. + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventPreToolUse { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(consumer.Close) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Failing MCP", + Slug: "failing-mcp", + Url: mcpTS.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + }) + reg := prometheus.NewRegistry() + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.PrometheusRegistry = reg + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "count admission-failed find_tools", + ModelConfigID: model.ID, + MCPServerIDs: []uuid.UUID{mcpConfig.ID}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("search"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusError) + + requireChatdMetricCounter(t, reg, "coderd_chatd_find_tools_calls_total", 1, nil) + }) + + t.Run("experiment gates deferral regardless of catalog size", func(t *testing.T) { + t.Parallel() + + run := func(t *testing.T, experimentEnabled bool) []byte { + t.Helper() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + mcpSrv := newTestMCPServer("small-mcp") + addTestMCPTextTool(mcpSrv, "echo", "Echo input", "echo: ") + mcpTS := httptest.NewServer(testMCPHTTPHandler(mcpSrv)) + t.Cleanup(mcpTS.Close) + var toolsJSON []byte + var toolsMu sync.Mutex + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + encoded, err := json.Marshal(req.Tools) + require.NoError(t, err) + toolsMu.Lock() + toolsJSON = encoded + toolsMu.Unlock() + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + model.ContextLimit = 100_000 + model = updateChatModelContextLimit(t, db, model) + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Small MCP", + Slug: "small-mcp", + Url: mcpTS.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + }) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + if !experimentEnabled { + withoutMCPToolSearch(cfg) + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "small deferred catalog", + ModelConfigID: model.ID, + MCPServerIDs: []uuid.UUID{mcpConfig.ID}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("finish"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + toolsMu.Lock() + defer toolsMu.Unlock() + return append([]byte(nil), toolsJSON...) + } + + toolNames := func(t *testing.T, toolsJSON []byte) []string { + t.Helper() + var tools []struct { + Function struct { + Name string `json:"name"` + } `json:"function"` + } + require.NoError(t, json.Unmarshal(toolsJSON, &tools)) + names := make([]string, 0, len(tools)) + for _, tool := range tools { + names = append(names, tool.Function.Name) + } + return names + } + + withoutExperiment := toolNames(t, run(t, false)) + require.Contains(t, withoutExperiment, "small-mcp__echo", + "without the experiment the MCP schema is advertised directly") + require.NotContains(t, withoutExperiment, chattool.FindToolsName) + + withExperiment := toolNames(t, run(t, true)) + require.Contains(t, withExperiment, chattool.FindToolsName, + "the experiment defers every MCP schema behind find_tools, even a small catalog") + require.NotContains(t, withExperiment, "small-mcp__echo") + }) +} + // TestMCPServerToolInvocation verifies that when a chat has // mcp_server_ids set, the chat loop connects to those MCP servers, // discovers their tools, and the LLM can invoke them. @@ -10502,6 +10923,7 @@ func TestMCPServerToolInvocation(t *testing.T) { Return(io.NopCloser(strings.NewReader("")), "", nil).AnyTimes() server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + withoutMCPToolSearch(cfg) cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { require.Equal(t, dbAgent.ID, agentID) @@ -10652,6 +11074,7 @@ func TestPlanModeRootChatApprovedExternalMCPToolInvocation(t *testing.T) { }) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + withoutMCPToolSearch(cfg) cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) }) @@ -10766,6 +11189,7 @@ func TestPlanModeRootChatApprovedExternalMCPWorkflowCanReachProposePlan(t *testi }).AnyTimes() server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + withoutMCPToolSearch(cfg) cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { require.Equal(t, dbAgent.ID, agentID) @@ -10973,6 +11397,7 @@ func TestMCPServerOAuth2TokenRefresh(t *testing.T) { mockConn.EXPECT().ReadFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). Return(io.NopCloser(strings.NewReader("")), "", nil).AnyTimes() server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + withoutMCPToolSearch(cfg) cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { require.Equal(t, dbAgent.ID, agentID) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 89063ff061280..e7586520f1ea0 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -238,10 +238,16 @@ type AssistantOutcome struct { // ExecuteLocalToolsOptions configures one local tool execution batch. type ExecuteLocalToolsOptions struct { - Tools []fantasy.AgentTool - ActiveTools []string - ProviderTools []ProviderTool - ToolCalls []fantasy.ToolCallContent + Tools []fantasy.AgentTool + ActiveTools []string + AllowInactiveTools map[string]bool + ProviderTools []ProviderTool + ToolCalls []fantasy.ToolCallContent + // ObservedToolCalls optionally carries the step's full assistant + // tool-call batch, including calls denied before execution, so + // step observers account for denied siblings that derivation will + // still count. Defaults to ToolCalls. + ObservedToolCalls []fantasy.ToolCallContent ExclusiveToolNames map[string]bool BuiltinToolNames map[string]bool @@ -594,8 +600,10 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool opts.Clock, opts.Tools, opts.ActiveTools, + opts.AllowInactiveTools, opts.ProviderTools, localCalls, + opts.ObservedToolCalls, opts.Metrics, opts.Logger, provider, @@ -1057,8 +1065,10 @@ func executeTools( clock quartz.Clock, allTools []fantasy.AgentTool, activeTools []string, + allowInactiveTools map[string]bool, providerTools []ProviderTool, toolCalls []fantasy.ToolCallContent, + observedToolCalls []fantasy.ToolCallContent, metrics *Metrics, logger slog.Logger, provider, model string, @@ -1109,47 +1119,83 @@ func executeTools( } } + observed := observedToolCalls + if observed == nil { + observed = localToolCalls + } + notifyStepToolCallObservers(toolMap, toolNameAliases, observed) + results := make([]fantasy.ToolResultContent, len(localToolCalls)) completedAt := make([]time.Time, len(localToolCalls)) + runCall := func(i int, tc fantasy.ToolCallContent) { + defer func() { + if r := recover(); r != nil { + results[i] = fantasy.ToolResultContent{ + ToolCallID: tc.ToolCallID, + ToolName: tc.ToolName, + Result: fantasy.ToolResultOutputContentError{ + Error: xerrors.Errorf("tool panicked: %v", r), + }, + } + } + // Record when this tool completed (or panicked). + // Captured per call so parallel tools get + // accurate individual completion times. + completedAt[i] = clockNow(clock) + }() + results[i] = executeSingleTool( + ctx, + toolMap, + tc, + metrics, + logger, + provider, + model, + builtinToolNames, + activeTools, + allowInactiveTools, + providerRunnerNames, + resultProviderMetadata, + maxResultBytes, + toolNameAliases, + ) + } + // Calls to tools that opt in via SerialToolCalls run in tool-call + // order after every concurrent sibling has settled. The step waits + // for all calls anyway, so sequencing them last costs nothing, and + // order-sensitive shared state (for example the find_tools + // activation budget) is claimed deterministically after sibling + // outcomes are known. All other calls stay concurrent. + var serialIndexes []int var wg sync.WaitGroup - wg.Add(len(localToolCalls)) for i, tc := range localToolCalls { + if isSerialToolCall(toolMap, toolNameAliases, tc.ToolName) { + serialIndexes = append(serialIndexes, i) + continue + } + wg.Add(1) go func() { defer wg.Done() - defer func() { - if r := recover(); r != nil { - results[i] = fantasy.ToolResultContent{ - ToolCallID: tc.ToolCallID, - ToolName: tc.ToolName, - Result: fantasy.ToolResultOutputContentError{ - Error: xerrors.Errorf("tool panicked: %v", r), - }, - } - } - // Record when this tool completed (or panicked). - // Captured per-goroutine so parallel tools get - // accurate individual completion times. - completedAt[i] = clockNow(clock) - }() - results[i] = executeSingleTool( - ctx, - toolMap, - tc, - metrics, - logger, - provider, - model, - builtinToolNames, - activeTools, - providerRunnerNames, - resultProviderMetadata, - maxResultBytes, - toolNameAliases, - ) + runCall(i, tc) }() } wg.Wait() + // Reconcile settled sibling outcomes before serial tools run, so + // for example find_tools refunds reservations of errored direct + // calls before its searches admit activations. + settled := make([]fantasy.ToolResultContent, 0, len(results)) + for i := range results { + if !slices.Contains(serialIndexes, i) { + settled = append(settled, results[i]) + } + } + notifyStepToolResultObservers(toolMap, toolNameAliases, localToolCalls, observed, settled) + + for _, i := range serialIndexes { + runCall(i, localToolCalls[i]) + } + // Publish results in the original tool-call order so SSE // subscribers see a deterministic event sequence. if onResult != nil { @@ -1262,6 +1308,7 @@ func executeSingleTool( provider, model string, builtinToolNames map[string]bool, activeTools []string, + allowInactiveTools map[string]bool, providerRunnerNames map[string]struct{}, resultProviderMetadata map[string]func(fantasy.ToolResponse) fantasy.ProviderMetadata, maxResultBytes int, @@ -1294,7 +1341,7 @@ func executeSingleTool( } _, isProviderRunner := providerRunnerNames[resolvedName] - if !isProviderRunner && !isToolActive(resolvedName, activeTools) { + if !isProviderRunner && !isToolActive(resolvedName, activeTools) && !allowInactiveTools[resolvedName] { result.Result = fantasy.ToolResultOutputContentError{ Error: xerrors.New("Tool not active in this turn: " + resolvedName), } @@ -1455,6 +1502,119 @@ func isToolActive(name string, activeTools []string) bool { return len(activeTools) == 0 || slices.Contains(activeTools, name) } +// serialToolCaller is implemented by tools whose calls within one step +// must execute in tool-call order because they claim from shared state. +type serialToolCaller interface{ SerialToolCalls() bool } + +// stepToolCallObserver is implemented by tools that need to see every +// tool-call name in the step before any call executes, for example so +// find_tools can charge same-step direct calls against its budget. +type stepToolCallObserver interface{ ObserveStepToolCalls(names []string) } + +// notifyStepToolCallObservers passes the step's resolved tool-call +// names to each distinct called tool that observes them. +func notifyStepToolCallObservers(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, calls []fantasy.ToolCallContent) { + names := make([]string, 0, len(calls)) + for _, tc := range calls { + name := tc.ToolName + if alias, ok := toolNameAliases[name]; ok { + name = alias + } + names = append(names, name) + } + notified := make(map[string]struct{}, len(names)) + for _, name := range names { + if _, dup := notified[name]; dup { + continue + } + notified[name] = struct{}{} + tool, ok := toolMap[name] + if !ok { + continue + } + if observer, ok := tool.(stepToolCallObserver); ok { + observer.ObserveStepToolCalls(names) + } + } +} + +// stepToolResultObserver is implemented by tools that need the step's +// per-call execution outcomes, for example so find_tools can refund +// budget it reserved for a direct call whose execution errored. names +// and errored are parallel slices in the observed tool-call order; +// outcomes are kept per call because one tool can be called several +// times in a step with different results. +type stepToolResultObserver interface { + ObserveStepToolResults(names []string, errored []bool) +} + +// notifyStepToolResultObservers passes the settled sibling outcomes to +// each distinct called tool that observes them, per call in observed +// order. Observed calls missing from the executed batch were rejected +// before execution (for example malformed JSON partitioned into +// synthetic denials) and settle as errored, since their persisted +// results always carry IsError. Serial calls have not run yet, so +// their own outcomes are reported as not errored; observers only need +// the concurrent siblings they share state with. +func notifyStepToolResultObservers(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, calls, observed []fantasy.ToolCallContent, settled []fantasy.ToolResultContent) { + resolve := func(name string) string { + if alias, ok := toolNameAliases[name]; ok { + return alias + } + return name + } + erroredByID := make(map[string]bool, len(settled)) + for _, tr := range settled { + _, isErr := tr.Result.(fantasy.ToolResultOutputContentError) + erroredByID[tr.ToolCallID] = isErr + } + executedIDs := make(map[string]struct{}, len(calls)) + for _, tc := range calls { + executedIDs[tc.ToolCallID] = struct{}{} + } + names := make([]string, 0, len(observed)) + errored := make([]bool, 0, len(observed)) + for _, tc := range observed { + names = append(names, resolve(tc.ToolName)) + if isErr, ok := erroredByID[tc.ToolCallID]; ok { + errored = append(errored, isErr) + continue + } + _, executed := executedIDs[tc.ToolCallID] + errored = append(errored, !executed) + } + notified := make(map[string]struct{}, len(calls)) + for _, tc := range calls { + name := tc.ToolName + if alias, ok := toolNameAliases[name]; ok { + name = alias + } + if _, dup := notified[name]; dup { + continue + } + notified[name] = struct{}{} + tool, ok := toolMap[name] + if !ok { + continue + } + if observer, ok := tool.(stepToolResultObserver); ok { + observer.ObserveStepToolResults(names, errored) + } + } +} + +func isSerialToolCall(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, name string) bool { + if alias, ok := toolNameAliases[name]; ok { + name = alias + } + tool, ok := toolMap[name] + if !ok { + return false + } + serial, ok := tool.(serialToolCaller) + return ok && serial.SerialToolCalls() +} + // buildToolDefinitions converts AgentTool definitions into the // fantasy.Tool slice expected by fantasy.Call. When activeTools // is non-empty, only function tools whose name appears in the diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 71af2ca7f1bc2..7d35d2588e708 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -5,6 +5,8 @@ import ( "encoding/base64" "errors" "iter" + "runtime" + "slices" "sync" "sync/atomic" "testing" @@ -872,6 +874,313 @@ func TestSanitizeAnthropicProviderToolContent(t *testing.T) { } } +type serialMarkerTool struct{ fantasy.AgentTool } + +func (serialMarkerTool) SerialToolCalls() bool { return true } + +type observerMarkerTool struct { + fantasy.AgentTool + observed func(names []string) +} + +func (t observerMarkerTool) ObserveStepToolCalls(names []string) { t.observed(names) } + +func TestExecuteToolsNotifiesStepToolCallObservers(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var observedNames []string + observedBeforeRun := false + observer := observerMarkerTool{ + AgentTool: fantasy.NewAgentTool( + "observer_tool", + "records sibling calls", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + mu.Lock() + observedBeforeRun = observedNames != nil + mu.Unlock() + return fantasy.NewTextResponse("ok"), nil + }, + ), + observed: func(names []string) { + mu.Lock() + defer mu.Unlock() + observedNames = append([]string{}, names...) + }, + } + var uncalledObserved atomic.Bool + uncalledObserver := observerMarkerTool{ + AgentTool: fantasy.NewAgentTool( + "uncalled_observer", + "never called this step", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + return fantasy.NewTextResponse("ok"), nil + }, + ), + observed: func([]string) { uncalledObserved.Store(true) }, + } + other := fantasy.NewAgentTool( + "other_tool", + "plain tool", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + return fantasy.NewTextResponse("ok"), nil + }, + ) + + executeTools( + context.Background(), + quartz.NewReal(), + []fantasy.AgentTool{observer, uncalledObserver, other}, + nil, + nil, + nil, + []fantasy.ToolCallContent{ + {ToolCallID: "1", ToolName: "observer_alias", Input: "{}"}, + {ToolCallID: "2", ToolName: "other_tool", Input: "{}"}, + }, + []fantasy.ToolCallContent{ + {ToolCallID: "1", ToolName: "observer_alias", Input: "{}"}, + {ToolCallID: "2", ToolName: "other_tool", Input: "{}"}, + {ToolCallID: "3", ToolName: "denied_tool", Input: "{}"}, + }, + NewMetrics(prometheus.NewRegistry()), + slog.Make(), + "fake", "fake-model", + map[string]bool{}, + defaultToolResultBytes, + map[string]string{"observer_alias": "observer_tool"}, + nil, + ) + + require.Equal(t, []string{"observer_tool", "other_tool", "denied_tool"}, observedNames, + "a called observer sees every observed tool-call name, including calls denied before execution") + require.True(t, observedBeforeRun, "observers are notified before any tool call executes") + require.False(t, uncalledObserved.Load(), "tools not called this step are not notified") +} + +type resultObserverMarkerTool struct { + fantasy.AgentTool + observedResults func(names []string, errored []bool) +} + +func (t resultObserverMarkerTool) ObserveStepToolResults(names []string, errored []bool) { + t.observedResults(names, errored) +} + +func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var gotNames []string + var gotErrored []bool + notifications := 0 + observer := resultObserverMarkerTool{ + AgentTool: fantasy.NewAgentTool( + "observer_tool", + "records sibling outcomes", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + return fantasy.NewTextResponse("ok"), nil + }, + ), + observedResults: func(names []string, errored []bool) { + mu.Lock() + defer mu.Unlock() + notifications++ + gotNames = append([]string{}, names...) + gotErrored = append([]bool{}, errored...) + }, + } + failing := fantasy.NewAgentTool( + "failing_tool", + "returns an error result", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + return fantasy.NewTextErrorResponse("remote error"), nil + }, + ) + + executed := []fantasy.ToolCallContent{ + {ToolCallID: "1", ToolName: "observer_alias", Input: "{}"}, + {ToolCallID: "2", ToolName: "failing_tool", Input: "{}"}, + {ToolCallID: "3", ToolName: "missing_tool", Input: "{}"}, + } + executeTools( + context.Background(), + quartz.NewReal(), + []fantasy.AgentTool{observer, failing}, + nil, + nil, + nil, + executed, + append(slices.Clone(executed), fantasy.ToolCallContent{ + ToolCallID: "4", ToolName: "rejected_tool", Input: "{not json", + }), + NewMetrics(prometheus.NewRegistry()), + slog.Make(), + "fake", "fake-model", + map[string]bool{}, + defaultToolResultBytes, + map[string]string{"observer_alias": "observer_tool"}, + nil, + ) + + require.Equal(t, 1, notifications, "each called observer is notified once per step") + require.Equal(t, []string{"observer_tool", "failing_tool", "missing_tool", "rejected_tool"}, gotNames, + "outcomes are reported per call in observed order with aliases resolved") + require.Equal(t, []bool{false, true, true, true}, gotErrored, + "error results, unresolvable tools, and observed calls rejected before execution all settle as errored outcomes") +} + +type serialResultObserverTool struct { + fantasy.AgentTool + observedResults func(names []string, errored []bool) +} + +func (serialResultObserverTool) SerialToolCalls() bool { return true } + +func (t serialResultObserverTool) ObserveStepToolResults(names []string, errored []bool) { + t.observedResults(names, errored) +} + +func TestExecuteToolsReconcilesResultsBeforeSerialCalls(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var erroredAtNotify []string + var erroredAtRun []string + notified := false + serial := serialResultObserverTool{ + AgentTool: fantasy.NewAgentTool( + "serial_observer", + "observes sibling outcomes before running", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + mu.Lock() + erroredAtRun = append([]string{}, erroredAtNotify...) + mu.Unlock() + return fantasy.NewTextResponse("ok"), nil + }, + ), + observedResults: func(names []string, errored []bool) { + mu.Lock() + defer mu.Unlock() + notified = true + erroredAtNotify = nil + for i, name := range names { + if errored[i] { + erroredAtNotify = append(erroredAtNotify, name) + } + } + }, + } + failing := fantasy.NewAgentTool( + "failing_tool", + "returns an error result", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + return fantasy.NewTextErrorResponse("remote error"), nil + }, + ) + + results := executeTools( + context.Background(), + quartz.NewReal(), + []fantasy.AgentTool{serial, failing}, + nil, + nil, + nil, + []fantasy.ToolCallContent{ + {ToolCallID: "1", ToolName: "serial_observer", Input: "{}"}, + {ToolCallID: "2", ToolName: "failing_tool", Input: "{}"}, + }, + nil, + NewMetrics(prometheus.NewRegistry()), + slog.Make(), + "fake", "fake-model", + map[string]bool{}, + defaultToolResultBytes, + nil, + nil, + ) + + require.True(t, notified) + require.Equal(t, []string{"failing_tool"}, erroredAtRun, + "a serial tool must see settled sibling outcomes before it executes") + require.Len(t, results, 2) + require.Equal(t, "1", results[0].ToolCallID, "results keep original call order") +} + +func TestExecuteToolsSerialToolCallOrder(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var events []string + inFlight := 0 + maxInFlight := 0 + record := func(event string, delta int) { + mu.Lock() + defer mu.Unlock() + inFlight += delta + if inFlight > maxInFlight { + maxInFlight = inFlight + } + events = append(events, event) + } + serial := serialMarkerTool{AgentTool: fantasy.NewAgentTool( + "serial_tool", + "records call order", + func(_ context.Context, _ struct{}, call fantasy.ToolCall) (fantasy.ToolResponse, error) { + record(call.ID+":start", 1) + runtime.Gosched() + record(call.ID+":end", -1) + return fantasy.NewTextResponse("ok"), nil + }, + )} + parallelRan := make(chan struct{}) + parallel := fantasy.NewAgentTool( + "parallel_tool", + "plain tool", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + close(parallelRan) + return fantasy.NewTextResponse("ok"), nil + }, + ) + + calls := []fantasy.ToolCallContent{ + {ToolCallID: "a", ToolName: "serial_tool", Input: "{}"}, + {ToolCallID: "p", ToolName: "parallel_tool", Input: "{}"}, + {ToolCallID: "b", ToolName: "serial_tool", Input: "{}"}, + {ToolCallID: "c", ToolName: "serial_tool", Input: "{}"}, + } + results := executeTools( + context.Background(), + quartz.NewReal(), + []fantasy.AgentTool{serial, parallel}, + nil, + nil, + nil, + calls, + nil, + NewMetrics(prometheus.NewRegistry()), + slog.Make(), + "fake", "fake-model", + map[string]bool{}, + defaultToolResultBytes, + nil, + nil, + ) + + require.Equal(t, []string{"a:start", "a:end", "b:start", "b:end", "c:start", "c:end"}, events, + "serial tool calls must run one at a time in tool-call order") + require.Equal(t, 1, maxInFlight) + select { + case <-parallelRan: + default: + t.Fatal("parallel tool call did not run") + } + require.Len(t, results, len(calls)) + for i, tc := range calls { + require.Equal(t, tc.ToolCallID, results[i].ToolCallID, "results keep original call order") + } +} + func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { t.Parallel() @@ -912,6 +1221,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { "fake", "fake-model", map[string]bool{}, []string{"screenshot"}, + nil, map[string]struct{}{}, nil, defaultToolResultBytes, @@ -961,6 +1271,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { "fake", "fake-model", map[string]bool{}, []string{"screenshot"}, + nil, map[string]struct{}{}, nil, defaultToolResultBytes, @@ -1005,6 +1316,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { "fake", "fake-model", map[string]bool{}, []string{"echo"}, + nil, map[string]struct{}{}, nil, defaultToolResultBytes, @@ -1053,6 +1365,7 @@ func TestExecuteSingleTool_ResolvesToolNameAlias(t *testing.T) { "fake", "fake-model", map[string]bool{}, []string{"interrupt_agent"}, + nil, map[string]struct{}{}, nil, defaultToolResultBytes, @@ -1093,6 +1406,7 @@ func TestExecuteSingleTool_UnknownAliasFallsThrough(t *testing.T) { "fake", "fake-model", map[string]bool{}, []string{"interrupt_agent"}, + nil, map[string]struct{}{}, nil, defaultToolResultBytes, @@ -1103,3 +1417,32 @@ func TestExecuteSingleTool_UnknownAliasFallsThrough(t *testing.T) { require.True(t, ok, "expected error output, got %T", result.Result) require.Contains(t, errOutput.Error.Error(), "close_agent") } + +func TestExecuteSingleTool_AllowsDeferredDirectCall(t *testing.T) { + t.Parallel() + tool := fantasy.NewAgentTool( + "server__direct", + "direct", + func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + return fantasy.NewTextResponse("ok"), nil + }, + ) + result := executeSingleTool( + context.Background(), + map[string]fantasy.AgentTool{"server__direct": tool}, + fantasy.ToolCallContent{ToolCallID: "call-direct", ToolName: "server__direct", Input: "{}"}, + NewMetrics(prometheus.NewRegistry()), + slog.Make(), + "fake", "fake-model", + map[string]bool{}, + []string{"find_tools"}, + map[string]bool{"server__direct": true}, + map[string]struct{}{}, + nil, + defaultToolResultBytes, + nil, + ) + text, ok := result.Result.(fantasy.ToolResultOutputContentText) + require.True(t, ok) + require.Equal(t, "ok", text.Text) +} diff --git a/coderd/x/chatd/chatloop/metrics.go b/coderd/x/chatd/chatloop/metrics.go index 3263f01d87c59..874a0ca52a0ae 100644 --- a/coderd/x/chatd/chatloop/metrics.go +++ b/coderd/x/chatd/chatloop/metrics.go @@ -27,17 +27,21 @@ const ( // Metrics holds Prometheus metrics for the chatd subsystem. type Metrics struct { - Chats *prometheus.GaugeVec - MessageCount *prometheus.HistogramVec - PromptSizeBytes *prometheus.HistogramVec - ToolResultSizeBytes *prometheus.HistogramVec - ToolResultTruncatedTotal *prometheus.CounterVec - ToolErrorsTotal *prometheus.CounterVec - TTFTSeconds *prometheus.HistogramVec - CompactionTotal *prometheus.CounterVec - StepsTotal *prometheus.CounterVec - StreamRetriesTotal *prometheus.CounterVec - StreamBufferDroppedTotal prometheus.Counter + Chats *prometheus.GaugeVec + MessageCount *prometheus.HistogramVec + PromptSizeBytes *prometheus.HistogramVec + ToolResultSizeBytes *prometheus.HistogramVec + ToolResultTruncatedTotal *prometheus.CounterVec + ToolErrorsTotal *prometheus.CounterVec + TTFTSeconds *prometheus.HistogramVec + CompactionTotal *prometheus.CounterVec + StepsTotal *prometheus.CounterVec + StreamRetriesTotal *prometheus.CounterVec + StreamBufferDroppedTotal prometheus.Counter + FindToolsCallsTotal prometheus.Counter + FindToolsEmptyTotal prometheus.Counter + FindToolsMatchCount prometheus.Histogram + FindToolsActivationsTotal prometheus.Counter } // NewMetrics creates a new Metrics instance registered with the @@ -109,6 +113,31 @@ func NewMetrics(reg prometheus.Registerer) *Metrics { Name: "stream_retries_total", Help: "Total LLM stream retries.", }, []string{"provider", "model", "kind"}), + FindToolsCallsTotal: factory.NewCounter(prometheus.CounterOpts{ + Namespace: metricsNamespace, + Subsystem: metricsSubsystem, + Name: "find_tools_calls_total", + Help: "Total find_tools calls.", + }), + FindToolsEmptyTotal: factory.NewCounter(prometheus.CounterOpts{ + Namespace: metricsNamespace, + Subsystem: metricsSubsystem, + Name: "find_tools_empty_total", + Help: "Total find_tools calls with no matches.", + }), + FindToolsMatchCount: factory.NewHistogram(prometheus.HistogramOpts{ + Namespace: metricsNamespace, + Subsystem: metricsSubsystem, + Name: "find_tools_match_count", + Help: "Number of matches returned by find_tools calls.", + Buckets: prometheus.LinearBuckets(0, 2, 11), + }), + FindToolsActivationsTotal: factory.NewCounter(prometheus.CounterOpts{ + Namespace: metricsNamespace, + Subsystem: metricsSubsystem, + Name: "find_tools_activations_total", + Help: "Total deferred tool activations returned by find_tools.", + }), StreamBufferDroppedTotal: factory.NewCounter(prometheus.CounterOpts{ Namespace: metricsNamespace, Subsystem: metricsSubsystem, diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go new file mode 100644 index 0000000000000..8338df5478743 --- /dev/null +++ b/coderd/x/chatd/chattool/findtools.go @@ -0,0 +1,737 @@ +package chattool + +import ( + "context" + "encoding/json" + "fmt" + "regexp" + "slices" + "strconv" + "strings" + "sync" + "unicode" + "unicode/utf8" + + "charm.land/fantasy" +) + +const ( + FindToolsName = "find_tools" + findToolsMaxMatches = 20 + findToolsCatalogTokens = 4000 + // findToolsMaxQueries and findToolsMaxQueryTokens bound scoring + // work: queries are model output, so one call could otherwise + // carry arbitrarily many tokens scored against every entry. + findToolsMaxQueries = 10 + findToolsMaxQueryTokens = 16 + // findToolsMaxNames bounds exact-name lookups the same way. Twice + // the match cap leaves room for unknown or duplicate names. + findToolsMaxNames = 2 * findToolsMaxMatches + // findToolsSpentBudgetFloor replaces a spent or over-reserved budget + // for searches so zero-cost reserved names remain activatable while + // any real schema weight still exceeds it. + findToolsSpentBudgetFloor = 0.000001 +) + +var findToolsTokenSeparator = regexp.MustCompile(`[^\p{L}\p{N}]+`) + +const findToolsBudgetExhausted = "the schema activation budget for this conversation is exhausted; call a cataloged tool directly by name to activate it in place of the least recently used schema" + +// FindToolCatalogEntry is the searchable metadata for one deferred tool. +type FindToolCatalogEntry struct { + Name string + Description string + Server string + ServerDescription string + ParameterText string + // SchemaTokens is the estimated prompt weight of the tool's full + // definition, used to cap how much one search may activate. + SchemaTokens float64 +} + +// FindToolsCall records one catalog search for logging and metrics. +type FindToolsCall struct { + Queries []string + Names []string + MatchCount int + Activated []string + TotalDeferred int + // Rejection is empty for successful searches. Rejected calls carry + // "budget" or "arguments" so callers can count them without + // polluting match or activation statistics. + Rejection string +} + +const ( + findToolsRejectionBudget = "budget" + findToolsRejectionArguments = "arguments" +) + +type FindToolsOptions struct { + Entries []FindToolCatalogEntry + // SchemaTokenBudget caps the aggregate SchemaTokens all searches on + // this tool instance may activate, so results never report + // activations that the activation budget would immediately shed. + // The budget is shared across calls because one step can execute + // several searches concurrently. <= 0 means unbounded. + SchemaTokenBudget float64 + // CatalogTokenBudget lowers the default catalog size cap so small + // context windows are not consumed by the catalog itself. <= 0 or + // values above the default keep the default. + CatalogTokenBudget float64 + OnCall func(context.Context, FindToolsCall) +} + +type FindToolsArgs struct { + Queries []string `json:"queries,omitempty"` + Names []string `json:"names,omitempty"` +} + +type FindToolsMatch struct { + Name string `json:"name"` + Description string `json:"description"` +} + +// FindToolsResult is persisted as the tool result and re-read on later steps. +type FindToolsResult struct { + Matches []FindToolsMatch `json:"matches"` + Activated []string `json:"activated"` + TotalDeferred int `json:"total_deferred"` +} + +// findToolsTool opts find_tools into in-order execution when one step +// contains several calls, so the shared schema budget is claimed in +// tool-call order rather than scheduler order. It also observes the +// step's sibling tool-call names so direct calls to deferred tools are +// charged against the budget before any search admits activations. +type findToolsTool struct { + fantasy.AgentTool + reserveStepCalls func(names []string) + settleStepResults func(names []string, errored []bool) + onDecodeRejected func(ctx context.Context) +} + +func (findToolsTool) SerialToolCalls() bool { return true } + +func (t findToolsTool) ObserveStepToolCalls(names []string) { t.reserveStepCalls(names) } + +func (t findToolsTool) ObserveStepToolResults(names []string, errored []bool) { + t.settleStepResults(names, errored) +} + +// Run counts calls the typed wrapper rejects during argument decoding, +// which never reach the handler and would otherwise be missing from +// call metrics. The response itself still comes from the wrapper's own +// decode so its wording stays canonical. +func (t findToolsTool) Run(ctx context.Context, call fantasy.ToolCall) (fantasy.ToolResponse, error) { + var args FindToolsArgs + if err := json.Unmarshal([]byte(call.Input), &args); err != nil && t.onDecodeRejected != nil { + t.onDecodeRejected(ctx) + } + return t.AgentTool.Run(ctx, call) +} + +// FindTools returns the built-in used to discover deferred MCP tool schemas. +func FindTools(options FindToolsOptions) fantasy.AgentTool { + entries := slices.Clone(options.Entries) + schemaTokensByName := make(map[string]float64, len(entries)) + for _, entry := range entries { + schemaTokensByName[entry.Name] = entry.SchemaTokens + } + var budgetMu sync.Mutex + remainingBudget := options.SchemaTokenBudget + // Direct calls to deferred tools in the same step are admitted by + // derivation before any search activations, per call in call order, + // while their cumulative weight fits the budget (the first always + // fits, mirroring derivation's newest-keep rule). Only that + // retained prefix is free to activate; a call past it is + // unclaimable this step because derivation marks it seen at its + // rejected position, so no same-step search can inline its schema + // either. Errored calls (including calls rejected before execution) + // are skipped per call when siblings settle, exactly as derivation + // postpones them by call ID, so one tool called several times with + // mixed outcomes admits at its first successful call's position. + type stepToolCall struct { + name string + errored bool + } + var stepCalls []stepToolCall + reserved := make(map[string]struct{}) + unclaimable := make(map[string]struct{}) + // Derivation deduplicates activations by name, so a name an earlier + // search already claimed is free for later searches in the step. + claimedBySearch := make(map[string]struct{}) + searchClaimed := 0.0 + recompute := func() { + clear(reserved) + clear(unclaimable) + charged := 0.0 + seen := make(map[string]struct{}, len(stepCalls)) + for _, call := range stepCalls { + if call.errored { + continue + } + if _, dup := seen[call.name]; dup { + continue + } + seen[call.name] = struct{}{} + weight := schemaTokensByName[call.name] + if len(reserved) > 0 && charged+weight > options.SchemaTokenBudget { + unclaimable[call.name] = struct{}{} + continue + } + reserved[call.name] = struct{}{} + charged += weight + } + remainingBudget = options.SchemaTokenBudget - charged - searchClaimed + } + rebuild := func(names []string, errored []bool) { + stepCalls = stepCalls[:0] + for i, name := range names { + if _, ok := schemaTokensByName[name]; !ok { + continue + } + stepCalls = append(stepCalls, stepToolCall{name: name, errored: len(errored) > i && errored[i]}) + } + recompute() + } + reserve := func(names []string) { + if options.SchemaTokenBudget <= 0 { + return + } + budgetMu.Lock() + defer budgetMu.Unlock() + // Outcomes are unknown before execution, so every call charges; + // settle rebuilds with real per-call outcomes before searches + // run. + rebuild(names, nil) + } + settle := func(names []string, errored []bool) { + if options.SchemaTokenBudget <= 0 { + return + } + budgetMu.Lock() + defer budgetMu.Unlock() + rebuild(names, errored) + } + onDecodeRejected := func(ctx context.Context) { + if options.OnCall != nil { + options.OnCall(ctx, FindToolsCall{ + TotalDeferred: len(entries), + Rejection: findToolsRejectionArguments, + }) + } + } + return findToolsTool{reserveStepCalls: reserve, settleStepResults: settle, onDecodeRejected: onDecodeRejected, AgentTool: fantasy.NewAgentTool( + FindToolsName, + buildFindToolsDescription(entries, options.CatalogTokenBudget), + func(ctx context.Context, args FindToolsArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + if len(args.Queries) == 0 && len(args.Names) == 0 { + if options.OnCall != nil { + options.OnCall(ctx, FindToolsCall{ + TotalDeferred: len(entries), + Rejection: findToolsRejectionArguments, + }) + } + return fantasy.NewTextErrorResponse("at least one query or name is required"), nil + } + budgetMu.Lock() + searchEntries := entries + if len(reserved) > 0 || len(unclaimable) > 0 || len(claimedBySearch) > 0 { + searchEntries = make([]FindToolCatalogEntry, 0, len(entries)) + for _, entry := range entries { + if _, ok := unclaimable[entry.Name]; ok { + continue + } + if _, ok := reserved[entry.Name]; ok { + entry.SchemaTokens = 0 + } + if _, ok := claimedBySearch[entry.Name]; ok { + entry.SchemaTokens = 0 + } + searchEntries = append(searchEntries, entry) + } + } + searchBudget := remainingBudget + if options.SchemaTokenBudget > 0 && searchBudget <= 0 { + // A spent budget still admits zero-cost reserved names, + // so search with a floor instead of failing outright. + searchBudget = findToolsSpentBudgetFloor + } + budgetTouched := options.SchemaTokenBudget > 0 && remainingBudget < options.SchemaTokenBudget + result, budgetSkipped := SearchTools(searchEntries, args, SearchBudget{ + SchemaTokens: searchBudget, + AllowFirstOverBudget: !budgetTouched, + }) + if options.SchemaTokenBudget > 0 { + if len(result.Activated) == 0 && budgetSkipped > 0 { + budgetMu.Unlock() + if options.OnCall != nil { + options.OnCall(ctx, FindToolsCall{ + Queries: args.Queries, + Names: args.Names, + TotalDeferred: len(entries), + Rejection: findToolsRejectionBudget, + }) + } + return fantasy.NewTextErrorResponse(findToolsBudgetExhausted), nil + } + admitted := 0.0 + for _, name := range result.Activated { + if _, ok := reserved[name]; ok { + continue + } + if _, ok := claimedBySearch[name]; ok { + continue + } + admitted += schemaTokensByName[name] + } + // Defensive invariant: with allowFirstOverBudget off, + // a touched budget can never admit an over-claim. If + // bookkeeping ever drifts, fail loudly rather than + // report activations derivation would shed. + if admitted > 0 && admitted > remainingBudget && budgetTouched { + budgetMu.Unlock() + if options.OnCall != nil { + options.OnCall(ctx, FindToolsCall{ + Queries: args.Queries, + Names: args.Names, + TotalDeferred: len(entries), + Rejection: findToolsRejectionBudget, + }) + } + return fantasy.NewTextErrorResponse(findToolsBudgetExhausted), nil + } + searchClaimed += admitted + remainingBudget -= admitted + for _, name := range result.Activated { + if _, ok := reserved[name]; !ok { + claimedBySearch[name] = struct{}{} + } + } + } + budgetMu.Unlock() + // Unclaimable entries stay deferred; report the full count. + result.TotalDeferred = len(entries) + if options.OnCall != nil { + options.OnCall(ctx, FindToolsCall{ + Queries: args.Queries, + Names: args.Names, + MatchCount: len(result.Matches), + Activated: result.Activated, + TotalDeferred: result.TotalDeferred, + }) + } + return marshalToolResponse(result), nil + }, + )} +} + +// SearchBudget bounds the schema weight one search may activate. +type SearchBudget struct { + // SchemaTokens is the remaining activation budget. <= 0 means + // unbounded. + SchemaTokens float64 + // AllowFirstOverBudget admits the first match even over budget. + // Callers set it only while the shared budget is untouched, where + // derivation's newest-keep rule retains a sole over-budget claim. + AllowFirstOverBudget bool +} + +// SearchTools includes exact name activations first, then fills the +// remaining match slots with the top-scored keyword matches. The shared +// cap and summary-length descriptions keep the persisted result small +// enough that generic tool-result truncation can never corrupt the +// activation JSON that later steps re-derive activations from. A +// positive budget additionally skips matches whose schema weight would +// push the aggregate past it, admitting later matches that still fit. +// The second result counts matches skipped for budget, so callers can +// tell an exhausted budget from no matches. +func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, budget SearchBudget) (FindToolsResult, int) { + byName := make(map[string]FindToolCatalogEntry, len(entries)) + for _, entry := range entries { + byName[entry.Name] = entry + } + + queryArgs := args.Queries + if len(queryArgs) > findToolsMaxQueries { + queryArgs = queryArgs[:findToolsMaxQueries] + } + queries := parseFindToolsQueries(entries, queryArgs) + + type scoredEntry struct { + entry FindToolCatalogEntry + score int + } + scored := make([]scoredEntry, 0, len(entries)) + for _, entry := range entries { + tokens := tokenizeFindToolsEntry(entry) + score := 0 + for _, query := range queries { + if query.server != "" { + if query.exact && entry.Server != query.server { + continue + } + if !query.exact && !strings.EqualFold(entry.Server, query.server) { + continue + } + } + if query.server != "" && len(query.tokens) == 0 { + score++ + continue + } + for _, token := range query.tokens { + score += tokens.score(token) + } + } + if score > 0 { + scored = append(scored, scoredEntry{entry: entry, score: score}) + } + } + slices.SortFunc(scored, func(a, b scoredEntry) int { + if a.score != b.score { + return b.score - a.score + } + return strings.Compare(a.entry.Name, b.entry.Name) + }) + matches := make([]FindToolsMatch, 0, findToolsMaxMatches) + activatedSet := make(map[string]struct{}, findToolsMaxMatches) + usedSchemaTokens := 0.0 + budgetSkipped := 0 + appendMatch := func(entry FindToolCatalogEntry) { + if _, exists := activatedSet[entry.Name]; exists { + return + } + if len(matches) >= findToolsMaxMatches { + return + } + overBudget := budget.SchemaTokens > 0 && usedSchemaTokens+entry.SchemaTokens > budget.SchemaTokens + if overBudget && (len(matches) > 0 || !budget.AllowFirstOverBudget) { + budgetSkipped++ + return + } + usedSchemaTokens += entry.SchemaTokens + matches = append(matches, FindToolsMatch{ + Name: entry.Name, + Description: truncateFindToolsSummary(entry.Description, 80), + }) + activatedSet[entry.Name] = struct{}{} + } + nameArgs := args.Names + if len(nameArgs) > findToolsMaxNames { + nameArgs = nameArgs[:findToolsMaxNames] + } + for _, name := range nameArgs { + if entry, ok := byName[name]; ok { + appendMatch(entry) + } + } + for _, item := range scored { + appendMatch(item.entry) + } + activated := make([]string, 0, len(activatedSet)) + for name := range activatedSet { + activated = append(activated, name) + } + slices.Sort(activated) + return FindToolsResult{Matches: matches, Activated: activated, TotalDeferred: len(entries)}, budgetSkipped +} + +type scopedFindToolsQuery struct { + server string + // exact scopes to the one server whose name matched byte-for-byte; + // otherwise the scope folds case and may span case-colliding + // servers. + exact bool + tokens []string +} + +// parseFindToolsQueries treats "server: terms" as a scope only when the +// prefix names a cataloged server, so queries like "error: timeout" +// still search normally. Prefixes are matched against full cataloged +// server names, longest first, because workspace server names may +// themselves contain ":". An exact-case prefix wins before the +// case-insensitive fallback, so servers whose names differ only by +// case each stay reachable by their advertised catalog name. +func parseFindToolsQueries(entries []FindToolCatalogEntry, queries []string) []scopedFindToolsQuery { + servers := make([]string, 0, len(entries)) + seen := make(map[string]struct{}, len(entries)) + for _, entry := range entries { + if entry.Server == "" { + continue + } + if _, dup := seen[entry.Server]; dup { + continue + } + seen[entry.Server] = struct{}{} + servers = append(servers, entry.Server) + } + // Longest first, so a server named "jira:prod" wins over "jira" + // when both are cataloged. + slices.SortFunc(servers, func(a, b string) int { return len(b) - len(a) }) + parsed := make([]scopedFindToolsQuery, 0, len(queries)) + for _, query := range queries { + scoped := false + trimmed := strings.TrimSpace(query) + // The raw query is matched before whitespace normalization so + // a whitespace-padded server name retained by collision + // handling stays selectable; then the trimmed exact and + // case-insensitive passes run as fallbacks. + passes := []struct { + text string + exact bool + }{ + {text: query, exact: true}, + {text: trimmed, exact: true}, + {text: trimmed, exact: false}, + } + for _, pass := range passes { + if scoped { + break + } + for _, server := range servers { + var rest string + if pass.exact { + var ok bool + rest, ok = strings.CutPrefix(pass.text, server) + if !ok { + continue + } + } else { + var ok bool + rest, ok = cutPrefixFold(pass.text, server) + if !ok { + continue + } + } + rest, ok := strings.CutPrefix(strings.TrimLeft(rest, " "), ":") + if !ok { + continue + } + parsed = append(parsed, scopedFindToolsQuery{server: server, exact: pass.exact, tokens: tokenizeFindToolsQuery(rest)}) + scoped = true + break + } + } + if !scoped { + parsed = append(parsed, scopedFindToolsQuery{tokens: tokenizeFindToolsQuery(query)}) + } + } + return parsed +} + +// cutPrefixFold is a case-insensitive strings.CutPrefix. It compares +// rune by rune with the same simple folding as strings.EqualFold, so a +// prefix whose folded form differs in UTF-8 byte length (like S and +// the long s) still matches and the cut lands on a rune boundary, +// which byte-length slicing cannot guarantee. +func cutPrefixFold(s, prefix string) (string, bool) { + rest := s + for _, prefixRune := range prefix { + restRune, size := utf8.DecodeRuneInString(rest) + if size == 0 || !runesFoldEqual(restRune, prefixRune) { + return "", false + } + rest = rest[size:] + } + return rest, true +} + +func runesFoldEqual(a, b rune) bool { + if a == b { + return true + } + for r := unicode.SimpleFold(a); r != a; r = unicode.SimpleFold(r) { + if r == b { + return true + } + } + return false +} + +func tokenizeFindTools(value string) []string { + parts := findToolsTokenSeparator.Split(strings.ToLower(value), -1) + return slices.DeleteFunc(parts, func(part string) bool { return part == "" }) +} + +// tokenizeFindToolsQuery caps model-supplied query tokens; catalog +// fields are tokenized uncapped so every term stays searchable. +func tokenizeFindToolsQuery(value string) []string { + tokens := tokenizeFindTools(value) + if len(tokens) > findToolsMaxQueryTokens { + tokens = tokens[:findToolsMaxQueryTokens] + } + return tokens +} + +// findToolsEntryTokens holds an entry's fields tokenized once per +// search, so scoring a token is a set lookup instead of re-splitting +// name, description, parameter, and server text for every query token. +type findToolsEntryTokens struct { + name string + nameTokens map[string]struct{} + description map[string]struct{} + parameters map[string]struct{} + server map[string]struct{} +} + +func tokenizeFindToolsEntry(entry FindToolCatalogEntry) findToolsEntryTokens { + toSet := func(value string) map[string]struct{} { + tokens := tokenizeFindTools(value) + set := make(map[string]struct{}, len(tokens)) + for _, token := range tokens { + set[token] = struct{}{} + } + return set + } + return findToolsEntryTokens{ + name: strings.ToLower(entry.Name), + nameTokens: toSet(entry.Name), + description: toSet(entry.Description), + parameters: toSet(entry.ParameterText), + // Server metadata is shown in catalog headers, so its terms + // must be searchable too. It applies to every tool on the + // server, so it scores below tool-specific matches. + server: toSet(entry.Server + " " + entry.ServerDescription), + } +} + +func (t findToolsEntryTokens) score(token string) int { + score := 0 + if _, ok := t.nameTokens[token]; ok { + score += 8 + } else if strings.Contains(t.name, token) { + score += 5 + } + if _, ok := t.description[token]; ok { + score += 2 + } + if _, ok := t.parameters[token]; ok { + score++ + } + if _, ok := t.server[token]; ok { + score++ + } + return score +} + +func buildFindToolsDescription(entries []FindToolCatalogEntry, catalogTokenBudget float64) string { + const usage = "Search deferred MCP tools by keyword, activate exact tool names, or scope a query to one server with a \"server: terms\" prefix. Calling a cataloged tool directly by name is allowed and auto-loads its schema, but search first for unfamiliar tools. At most 20 tools are returned and activated per call; call again for more.\n\n" + budget := float64(findToolsCatalogTokens) + if catalogTokenBudget > 0 && catalogTokenBudget < budget { + budget = catalogTokenBudget + } + groups := groupFindToolsEntries(entries) + catalog := detailedFindToolsCatalog(groups) + if estimatedFindToolsTokens(usage+catalog) > budget { + catalog = namesOnlyFindToolsCatalog(groups) + } + if estimatedFindToolsTokens(usage+catalog) > budget { + catalog = countsOnlyFindToolsCatalog(groups) + } + // Server count and slug length are unbounded, so even the per-server + // counts catalog needs a final constant-size fallback. + if estimatedFindToolsTokens(usage+catalog) > budget { + catalog = fmt.Sprintf("%d deferred tools across %d servers.\n", len(entries), len(groups)) + } + return usage + catalog +} + +func detailedFindToolsCatalog(groups []findToolsGroup) string { + var b strings.Builder + for _, group := range groups { + writeFindToolsGroupHeader(&b, group) + for _, entry := range group.entries { + _, _ = b.WriteString("- ") + _, _ = b.WriteString(entry.Name) + _, _ = b.WriteString(" - ") + _, _ = b.WriteString(truncateFindToolsSummary(entry.Description, 80)) + _ = b.WriteByte('\n') + } + } + return b.String() +} + +func namesOnlyFindToolsCatalog(groups []findToolsGroup) string { + var b strings.Builder + for _, group := range groups { + writeFindToolsGroupHeader(&b, group) + names := make([]string, 0, len(group.entries)) + for _, entry := range group.entries { + names = append(names, entry.Name) + } + _, _ = b.WriteString(strings.Join(names, " ")) + _ = b.WriteByte('\n') + } + return b.String() +} + +func countsOnlyFindToolsCatalog(groups []findToolsGroup) string { + var b strings.Builder + for _, group := range groups { + _, _ = b.WriteString("## ") + _, _ = b.WriteString(group.server) + _, _ = b.WriteString(" (") + _, _ = b.WriteString(strconv.Itoa(len(group.entries))) + _, _ = b.WriteString(" tools)\n") + } + return b.String() +} + +func writeFindToolsGroupHeader(b *strings.Builder, group findToolsGroup) { + _, _ = b.WriteString("## ") + _, _ = b.WriteString(group.server) + if summary := truncateFindToolsSummary(group.description, 60); summary != "" { + _, _ = b.WriteString(" - ") + _, _ = b.WriteString(summary) + } + _ = b.WriteByte('\n') +} + +type findToolsGroup struct { + server string + description string + entries []FindToolCatalogEntry +} + +func groupFindToolsEntries(entries []FindToolCatalogEntry) []findToolsGroup { + grouped := make(map[string]*findToolsGroup) + for _, entry := range entries { + // Callers assign every entry a non-empty Server so display, + // scope matching, and scoring share one identity; an empty + // value groups as-is rather than under a label scopes cannot + // reach. + server := entry.Server + group := grouped[server] + if group == nil { + group = &findToolsGroup{server: server, description: entry.ServerDescription} + grouped[server] = group + } + group.entries = append(group.entries, entry) + } + groups := make([]findToolsGroup, 0, len(grouped)) + for _, group := range grouped { + slices.SortFunc(group.entries, func(a, b FindToolCatalogEntry) int { return strings.Compare(a.Name, b.Name) }) + groups = append(groups, *group) + } + slices.SortFunc(groups, func(a, b findToolsGroup) int { return strings.Compare(a.server, b.server) }) + return groups +} + +func truncateFindToolsSummary(value string, maxRunes int) string { + line, _, _ := strings.Cut(value, "\n") + sentence, _, _ := strings.Cut(line, ". ") + value = strings.TrimSpace(sentence) + if utf8.RuneCountInString(value) <= maxRunes { + return value + } + runes := []rune(value) + if maxRunes <= 3 { + return string(runes[:maxRunes]) + } + return strings.TrimSpace(string(runes[:maxRunes-3])) + "..." +} + +func estimatedFindToolsTokens(value string) float64 { + return float64(len(value)) / 2.5 +} diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go new file mode 100644 index 0000000000000..0a7af3fdc4786 --- /dev/null +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -0,0 +1,625 @@ +package chattool + +import ( + "context" + "encoding/json" + "fmt" + "slices" + "strings" + "testing" + + "charm.land/fantasy" + "github.com/stretchr/testify/require" +) + +func TestSearchTools(t *testing.T) { + t.Parallel() + entries := []FindToolCatalogEntry{ + {Name: "github__create_issue", Description: "Create an issue", ParameterText: "repository title body"}, + {Name: "github__search_issues", Description: "Search issue descriptions", ParameterText: "repository query"}, + {Name: "slack__post_message", Description: "Post a message", ParameterText: "channel text"}, + } + + t.Run("weights and tie break", func(t *testing.T) { + t.Parallel() + result, _ := SearchTools(entries, FindToolsArgs{Queries: []string{"issue"}}, SearchBudget{}) + require.Len(t, result.Matches, 2) + require.Equal(t, []string{"github__create_issue", "github__search_issues"}, []string{result.Matches[0].Name, result.Matches[1].Name}) + }) + t.Run("parameter text", func(t *testing.T) { + t.Parallel() + result, _ := SearchTools(entries, FindToolsArgs{Queries: []string{"channel"}}, SearchBudget{}) + require.Equal(t, "slack__post_message", result.Matches[0].Name) + }) + t.Run("exact names", func(t *testing.T) { + t.Parallel() + result, _ := SearchTools(entries, FindToolsArgs{Names: []string{"slack__post_message", "missing"}}, SearchBudget{}) + require.Equal(t, []string{"slack__post_message"}, result.Activated) + require.Equal(t, "slack__post_message", result.Matches[0].Name) + }) + t.Run("empty queries", func(t *testing.T) { + t.Parallel() + result, _ := SearchTools(entries, FindToolsArgs{}, SearchBudget{}) + require.Empty(t, result.Matches) + require.Empty(t, result.Activated) + }) + t.Run("cap", func(t *testing.T) { + t.Parallel() + many := make([]FindToolCatalogEntry, 25) + for i := range many { + many[i] = FindToolCatalogEntry{Name: fmt.Sprintf("server__tool_%02d", i), Description: "common"} + } + result, _ := SearchTools(many, FindToolsArgs{Queries: []string{"common"}}, SearchBudget{}) + require.Len(t, result.Matches, findToolsMaxMatches) + require.Equal(t, "server__tool_00", result.Matches[0].Name) + }) + t.Run("names capped and prioritized over queries", func(t *testing.T) { + t.Parallel() + many := make([]FindToolCatalogEntry, 25) + names := make([]string, 0, len(many)) + for i := range many { + many[i] = FindToolCatalogEntry{Name: fmt.Sprintf("server__tool_%02d", i), Description: "common"} + names = append(names, many[i].Name) + } + result, _ := SearchTools(many, FindToolsArgs{Queries: []string{"common"}, Names: []string{"server__tool_24"}}, SearchBudget{}) + require.Len(t, result.Matches, findToolsMaxMatches) + require.Equal(t, "server__tool_24", result.Matches[0].Name) + require.Contains(t, result.Activated, "server__tool_24") + + capped, _ := SearchTools(many, FindToolsArgs{Names: names}, SearchBudget{}) + require.Len(t, capped.Matches, findToolsMaxMatches) + require.Len(t, capped.Activated, findToolsMaxMatches) + }) + t.Run("names list is bounded", func(t *testing.T) { + t.Parallel() + entries := []FindToolCatalogEntry{ + {Name: "server__target", Description: "does things"}, + } + unknown := make([]string, findToolsMaxNames) + for i := range unknown { + unknown[i] = fmt.Sprintf("missing_%02d", i) + } + result, _ := SearchTools(entries, FindToolsArgs{Names: append(slices.Clone(unknown), "server__target")}, SearchBudget{}) + require.Empty(t, result.Activated, + "a name past the inspection cap is not looked up") + + result, _ = SearchTools(entries, FindToolsArgs{Names: append(unknown[:findToolsMaxNames-1], "server__target")}, SearchBudget{}) + require.Equal(t, []string{"server__target"}, result.Activated, + "a name within the inspection cap still activates") + }) + t.Run("server metadata", func(t *testing.T) { + t.Parallel() + serverEntries := []FindToolCatalogEntry{ + {Name: "tracker__create", Description: "Create an item", Server: "tracker", ServerDescription: "Project tracking"}, + {Name: "docs__create", Description: "Create a project document", Server: "docs", ServerDescription: "Documentation"}, + } + result, _ := SearchTools(serverEntries, FindToolsArgs{Queries: []string{"tracking"}}, SearchBudget{}) + require.Equal(t, []string{"tracker__create"}, result.Activated) + + result, _ = SearchTools(serverEntries, FindToolsArgs{Queries: []string{"project"}}, SearchBudget{}) + require.Equal(t, "docs__create", result.Matches[0].Name, + "tool description match outranks server metadata match") + require.Len(t, result.Matches, 2) + }) + t.Run("server prefix scope", func(t *testing.T) { + t.Parallel() + scopedEntries := []FindToolCatalogEntry{ + {Name: "ci__status", Description: "Pipeline status", Server: "ci"}, + {Name: "github__get_commit", Description: "Get commit status", Server: "github"}, + } + result, _ := SearchTools(scopedEntries, FindToolsArgs{Queries: []string{"github: status"}}, SearchBudget{}) + require.Equal(t, []string{"github__get_commit"}, result.Activated, + "a known server prefix restricts matches to that server") + + result, _ = SearchTools(scopedEntries, FindToolsArgs{Queries: []string{"github:"}}, SearchBudget{}) + require.Equal(t, []string{"github__get_commit"}, result.Activated, + "a bare server prefix lists that server's tools") + + result, _ = SearchTools(scopedEntries, FindToolsArgs{Queries: []string{"error: status"}}, SearchBudget{}) + require.Len(t, result.Matches, 2, + "an unknown prefix is searched as plain keywords") + + result, _ = SearchTools(scopedEntries, FindToolsArgs{Queries: []string{"GitHub: status"}}, SearchBudget{}) + require.Equal(t, []string{"github__get_commit"}, result.Activated, + "a case-variant prefix still scopes to its server when no exact-case name collides") + }) + t.Run("case-colliding server names", func(t *testing.T) { + t.Parallel() + caseEntries := []FindToolCatalogEntry{ + {Name: "GitHub__enterprise_status", Description: "Enterprise status", Server: "GitHub"}, + {Name: "github__get_commit", Description: "Get commit status", Server: "github"}, + } + result, _ := SearchTools(caseEntries, FindToolsArgs{Queries: []string{"GitHub: status"}}, SearchBudget{}) + require.Equal(t, []string{"GitHub__enterprise_status"}, result.Activated, + "an exact-case prefix scopes only to its own server") + + result, _ = SearchTools(caseEntries, FindToolsArgs{Queries: []string{"github: status"}}, SearchBudget{}) + require.Equal(t, []string{"github__get_commit"}, result.Activated, + "the case-colliding sibling stays reachable by its own exact name") + + result, _ = SearchTools(caseEntries, FindToolsArgs{Queries: []string{"GITHUB: status"}}, SearchBudget{}) + require.Len(t, result.Activated, 2, + "a prefix matching no exact-case name falls back to spanning the case-colliding servers") + }) + t.Run("folded scopes with different byte lengths", func(t *testing.T) { + t.Parallel() + // The long s folds with S and s but is two UTF-8 bytes, so a + // byte-length prefix slice can never line the two forms up. + // Scope-only queries keep the assertion sharp: an unscoped + // fallback tokenizes to a term that matches nothing because + // ToLower does not case-fold the long s. + foldedEntries := []FindToolCatalogEntry{ + {Name: "ſerver__tool", Description: "does things", Server: "ſerver"}, + } + result, _ := SearchTools(foldedEntries, FindToolsArgs{Queries: []string{"Server:"}}, SearchBudget{}) + require.Equal(t, []string{"ſerver__tool"}, result.Activated, + "a folded scope with fewer bytes than the server name still scopes") + + asciiEntries := []FindToolCatalogEntry{ + {Name: "server__tool", Description: "does things", Server: "server"}, + } + result, _ = SearchTools(asciiEntries, FindToolsArgs{Queries: []string{"ſerver:"}}, SearchBudget{}) + require.Equal(t, []string{"server__tool"}, result.Activated, + "a folded scope with more bytes than the server name still scopes") + }) + t.Run("bounded query work", func(t *testing.T) { + t.Parallel() + entries := []FindToolCatalogEntry{ + {Name: "server__match", Description: "Matches the last token", Server: "server"}, + } + // The matching term is placed beyond both caps, so a match + // proves the caps were not applied. + overflowQuery := strings.Repeat("filler ", findToolsMaxQueryTokens) + "matches" + result, _ := SearchTools(entries, FindToolsArgs{Queries: []string{overflowQuery}}, SearchBudget{}) + require.Empty(t, result.Activated, + "tokens beyond the per-query cap are not scored") + + queries := make([]string, findToolsMaxQueries+1) + for i := range queries { + queries[i] = "filler" + } + queries[len(queries)-1] = "matches" + result, _ = SearchTools(entries, FindToolsArgs{Queries: queries}, SearchBudget{}) + require.Empty(t, result.Activated, + "queries beyond the per-call cap are not scored") + + result, _ = SearchTools(entries, FindToolsArgs{Queries: []string{"matches"}}, SearchBudget{}) + require.Equal(t, []string{"server__match"}, result.Activated, + "capped search still scores in-bound tokens") + }) + t.Run("whitespace-colliding server names", func(t *testing.T) { + t.Parallel() + paddedEntries := []FindToolCatalogEntry{ + {Name: "_everything___ping", Description: "Ping status", Server: " everything "}, + {Name: "everything__status", Description: "Get status", Server: "everything"}, + } + result, _ := SearchTools(paddedEntries, FindToolsArgs{Queries: []string{"everything: status"}}, SearchBudget{}) + require.Equal(t, []string{"everything__status"}, result.Activated, + "the exact-form scope matches only its own server, not a whitespace-padded sibling") + + result, _ = SearchTools(paddedEntries, FindToolsArgs{Queries: []string{" everything : ping"}}, SearchBudget{}) + require.Equal(t, []string{"_everything___ping"}, result.Activated, + "the raw query prefix is matched before trimming, so the padded server stays selectable") + }) + t.Run("server names containing colons", func(t *testing.T) { + t.Parallel() + colonEntries := []FindToolCatalogEntry{ + {Name: "jira_prod__list_issues", Description: "List issues", Server: "jira:prod"}, + {Name: "jira__list_issues", Description: "List issues", Server: "jira"}, + {Name: "ci__status", Description: "Issue pipeline status", Server: "ci"}, + } + result, _ := SearchTools(colonEntries, FindToolsArgs{Queries: []string{"jira:prod: issues"}}, SearchBudget{}) + require.Equal(t, []string{"jira_prod__list_issues"}, result.Activated, + "the longest cataloged server name wins over its colon-split prefix") + + result, _ = SearchTools(colonEntries, FindToolsArgs{Queries: []string{"jira: issues"}}, SearchBudget{}) + require.Equal(t, []string{"jira__list_issues"}, result.Activated, + "the shorter server still scopes its own queries") + }) + t.Run("unicode terms", func(t *testing.T) { + t.Parallel() + unicodeEntries := []FindToolCatalogEntry{ + {Name: "docs__検索", Description: "ドキュメント検索"}, + {Name: "docs__erstellen", Description: "Dokument ERSTELLEN"}, + } + result, _ := SearchTools(unicodeEntries, FindToolsArgs{Queries: []string{"検索"}}, SearchBudget{}) + require.Equal(t, []string{"docs__検索"}, result.Activated) + + result, _ = SearchTools(unicodeEntries, FindToolsArgs{Queries: []string{"Erstellen"}}, SearchBudget{}) + require.Equal(t, []string{"docs__erstellen"}, result.Activated) + }) + t.Run("schema token budget", func(t *testing.T) { + t.Parallel() + weighted := []FindToolCatalogEntry{ + {Name: "server__big_a", Description: "big", SchemaTokens: 60}, + {Name: "server__big_b", Description: "big", SchemaTokens: 60}, + {Name: "server__huge", Description: "big", SchemaTokens: 500}, + } + result, _ := SearchTools(weighted, FindToolsArgs{Queries: []string{"big"}}, SearchBudget{SchemaTokens: 100, AllowFirstOverBudget: true}) + require.Equal(t, []string{"server__big_a"}, result.Activated, + "matches stop once the schema budget is spent") + + result, _ = SearchTools(weighted, FindToolsArgs{Names: []string{"server__huge"}}, SearchBudget{SchemaTokens: 100, AllowFirstOverBudget: true}) + require.Equal(t, []string{"server__huge"}, result.Activated, + "the first match is kept even when it alone exceeds the budget") + }) + t.Run("result descriptions are summarized", func(t *testing.T) { + t.Parallel() + long := []FindToolCatalogEntry{{ + Name: "server__verbose", + Description: strings.Repeat("word ", 100), + }} + result, _ := SearchTools(long, FindToolsArgs{Names: []string{"server__verbose"}}, SearchBudget{}) + require.LessOrEqual(t, len([]rune(result.Matches[0].Description)), 80) + }) +} + +func TestFindTools(t *testing.T) { + t.Parallel() + var recorded FindToolsCall + tool := FindTools(FindToolsOptions{ + Entries: []FindToolCatalogEntry{{Name: "github__create_issue", Description: "Create an issue"}}, + OnCall: func(_ context.Context, call FindToolsCall) { recorded = call }, + }) + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"queries":["issue"]}`}) + require.NoError(t, err) + var result FindToolsResult + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Equal(t, []string{"github__create_issue"}, result.Activated) + require.Equal(t, 1, recorded.MatchCount) + + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{}`}) + require.NoError(t, err) + require.True(t, resp.IsError) +} + +func TestFindToolsSerialToolCalls(t *testing.T) { + t.Parallel() + serial, ok := FindTools(FindToolsOptions{}).(interface{ SerialToolCalls() bool }) + require.True(t, ok, "find_tools must opt into serial execution so shared-budget admission follows tool-call order") + require.True(t, serial.SerialToolCalls()) +} + +func TestFindToolsDirectCallReservation(t *testing.T) { + t.Parallel() + newTool := func(budget float64) (fantasy.AgentTool, interface{ ObserveStepToolCalls([]string) }) { + tool := FindTools(FindToolsOptions{ + Entries: []FindToolCatalogEntry{ + {Name: "server__a", SchemaTokens: 60}, + {Name: "server__b", SchemaTokens: 50}, + {Name: "server__c", SchemaTokens: 30}, + }, + SchemaTokenBudget: budget, + }) + observer, ok := tool.(interface{ ObserveStepToolCalls([]string) }) + require.True(t, ok, "find_tools must observe step tool calls to reserve direct-call schema weight") + return tool, observer + } + + t.Run("direct calls charge the budget before searches", func(t *testing.T) { + t.Parallel() + tool, observer := newTool(100) + observer.ObserveStepToolCalls([]string{"server__a", "server__a", "unknown", FindToolsName}) + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__b"]}`}) + require.NoError(t, err) + require.True(t, resp.IsError, "a search claim exceeding the budget left by direct calls must fail loudly") + + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__c"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError, "duplicate direct-call names are charged once") + var result FindToolsResult + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Equal(t, []string{"server__c"}, result.Activated) + }) + + t.Run("rejected calls still reach OnCall", func(t *testing.T) { + t.Parallel() + var calls []FindToolsCall + tool := FindTools(FindToolsOptions{ + Entries: []FindToolCatalogEntry{ + {Name: "server__a", SchemaTokens: 60}, + {Name: "server__b", SchemaTokens: 50}, + }, + SchemaTokenBudget: 60, + OnCall: func(_ context.Context, call FindToolsCall) { calls = append(calls, call) }, + }) + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__a"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError) + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__b"]}`}) + require.NoError(t, err) + require.True(t, resp.IsError) + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{}`}) + require.NoError(t, err) + require.True(t, resp.IsError) + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"queries":"github"}`}) + require.NoError(t, err) + require.True(t, resp.IsError, "a type mismatch is rejected by the argument decoder") + require.Len(t, calls, 4, "rejected calls count toward call totals") + require.Empty(t, calls[0].Rejection) + require.Equal(t, "budget", calls[1].Rejection) + require.Equal(t, []string{"server__b"}, calls[1].Names) + require.Empty(t, calls[1].Activated, "a rejected call reports no activations") + require.Equal(t, "arguments", calls[2].Rejection, "empty-argument calls are counted as rejected") + require.Empty(t, calls[2].Activated) + require.Equal(t, "arguments", calls[3].Rejection, + "calls rejected during argument decoding are counted before the handler is reached") + }) + + t.Run("a touched budget skips oversized matches and admits later fits", func(t *testing.T) { + t.Parallel() + tool := FindTools(FindToolsOptions{ + Entries: []FindToolCatalogEntry{ + {Name: "server__a", SchemaTokens: 60}, + {Name: "server__b", SchemaTokens: 50}, + {Name: "server__c", SchemaTokens: 30}, + }, + SchemaTokenBudget: 100, + }) + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__a"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError) + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__b","server__c"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError, "an oversized top match must not fail the call when a later match fits") + var result FindToolsResult + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Equal(t, []string{"server__c"}, result.Activated, + "the oversized first match is skipped and the fitting later match admitted") + }) + + t.Run("an errored direct call refunds its reservation", func(t *testing.T) { + t.Parallel() + tool, observer := newTool(100) + settler, ok := tool.(interface { + ObserveStepToolResults(names []string, errored []bool) + }) + require.True(t, ok, "find_tools must observe step results to refund errored reservations") + observer.ObserveStepToolCalls([]string{"server__a"}) + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__b"]}`}) + require.NoError(t, err) + require.True(t, resp.IsError, "the pre-execution reservation holds while the outcome is unknown") + + settler.ObserveStepToolResults([]string{"server__a", "unknown"}, []bool{true, true}) + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__b"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError, "the refunded reservation admits later searches") + var result FindToolsResult + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Equal(t, []string{"server__b"}, result.Activated) + }) + + t.Run("a name that executed successfully keeps its reservation", func(t *testing.T) { + t.Parallel() + tool, observer := newTool(100) + settler, ok := tool.(interface { + ObserveStepToolResults(names []string, errored []bool) + }) + require.True(t, ok) + observer.ObserveStepToolCalls([]string{"server__a", "server__a"}) + settler.ObserveStepToolResults([]string{"server__a", "server__a"}, []bool{false, true}) + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__b"]}`}) + require.NoError(t, err) + require.True(t, resp.IsError, "a successful execution pins the reservation even when a later call errors") + }) + + t.Run("mixed outcomes admit at the first successful call position", func(t *testing.T) { + t.Parallel() + tool, observer := newTool(100) + settler, ok := tool.(interface { + ObserveStepToolResults(names []string, errored []bool) + }) + require.True(t, ok) + // A errors, B succeeds, then A succeeds: derivation postpones + // the errored A by call ID, admits B, and budget-rejects the + // later A (60 over the 50 already charged), so only B's schema + // reaches the next request. + observer.ObserveStepToolCalls([]string{"server__a", "server__b", "server__a"}) + settler.ObserveStepToolResults([]string{"server__a", "server__b", "server__a"}, []bool{true, false, false}) + + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__a"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError) + var result FindToolsResult + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Empty(t, result.Activated, + "a name derivation budget-rejects at its successful call position is unclaimable, not free") + + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__b"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError) + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Equal(t, []string{"server__b"}, result.Activated, "the admitted call stays free") + + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__c"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError) + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Equal(t, []string{"server__c"}, result.Activated, + "only the admitted prefix is charged, so the leftover budget admits new claims") + }) + + t.Run("aggregate overflow frees only the prefix derivation retains", func(t *testing.T) { + t.Parallel() + tool, observer := newTool(100) + settler, ok := tool.(interface { + ObserveStepToolResults(names []string, errored []bool) + }) + require.True(t, ok) + observer.ObserveStepToolCalls([]string{"server__a", "server__b"}) + settler.ObserveStepToolResults([]string{"server__a", "server__b"}, []bool{false, false}) + + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__b"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError) + var result FindToolsResult + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Empty(t, result.Activated, + "a direct call past the retained prefix cannot be reported activated: derivation sheds it") + require.Equal(t, 3, result.TotalDeferred, "unclaimable entries still count as deferred") + + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__a"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError) + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Equal(t, []string{"server__a"}, result.Activated, "the retained prefix stays free") + + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__c"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError) + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Equal(t, []string{"server__c"}, result.Activated, + "the skipped call's weight is not charged, so later searches keep the leftover budget") + }) + + t.Run("a name claimed by an earlier search is free for later searches", func(t *testing.T) { + t.Parallel() + tool, _ := newTool(60) + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__a"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError) + + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__a"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError, "derivation deduplicates by name, so a repeated claim costs nothing") + var result FindToolsResult + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Equal(t, []string{"server__a"}, result.Activated) + + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__c"]}`}) + require.NoError(t, err) + require.True(t, resp.IsError, "the repeated claim must not have refunded the spent budget") + }) + + t.Run("an errored prefix call promotes the next observed name", func(t *testing.T) { + t.Parallel() + tool, observer := newTool(100) + settler, ok := tool.(interface { + ObserveStepToolResults(names []string, errored []bool) + }) + require.True(t, ok) + observer.ObserveStepToolCalls([]string{"server__a", "server__b"}) + settler.ObserveStepToolResults([]string{"server__a", "server__b"}, []bool{true, false}) + + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__b"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError) + var result FindToolsResult + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Equal(t, []string{"server__b"}, result.Activated, + "the errored call leaves the prefix, so the succeeding call becomes free") + + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__a"]}`}) + require.NoError(t, err) + require.True(t, resp.IsError, + "the errored call is claimable at full weight, which exceeds the leftover budget") + }) + + t.Run("reserved names stay activatable after the budget is spent", func(t *testing.T) { + t.Parallel() + tool, observer := newTool(50) + observer.ObserveStepToolCalls([]string{"server__a"}) + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__a"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError, "derivation retains direct calls, so reporting them activated is free") + var result FindToolsResult + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + require.Equal(t, []string{"server__a"}, result.Activated) + + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__c"]}`}) + require.NoError(t, err) + require.True(t, resp.IsError, "an over-reserved budget admits no new schema weight") + }) +} + +func TestFindToolsSharedSchemaBudget(t *testing.T) { + t.Parallel() + tool := FindTools(FindToolsOptions{ + Entries: []FindToolCatalogEntry{ + {Name: "server__a", SchemaTokens: 60}, + {Name: "server__b", SchemaTokens: 60}, + {Name: "server__c", SchemaTokens: 60}, + {Name: "server__d", SchemaTokens: 60}, + }, + SchemaTokenBudget: 200, + }) + activated := func(input string) []string { + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: input}) + require.NoError(t, err) + var result FindToolsResult + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + return result.Activated + } + + require.Equal(t, []string{"server__a", "server__b"}, activated(`{"names":["server__a","server__b"]}`)) + require.Equal(t, []string{"server__c"}, activated(`{"names":["server__c","server__d"]}`), + "the second call spends the remaining shared budget, not a fresh one") + + resp, err := tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__d"]}`}) + require.NoError(t, err) + require.True(t, resp.IsError, + "a call whose claims cannot fit the remaining budget errors instead of over-claiming") + + huge := FindTools(FindToolsOptions{ + Entries: []FindToolCatalogEntry{ + {Name: "server__huge", SchemaTokens: 500}, + {Name: "server__other", SchemaTokens: 60}, + }, + SchemaTokenBudget: 200, + }) + resp, err = huge.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__huge"]}`}) + require.NoError(t, err) + var hugeResult FindToolsResult + require.NoError(t, json.Unmarshal([]byte(resp.Content), &hugeResult)) + require.Equal(t, []string{"server__huge"}, hugeResult.Activated, + "an untouched budget may over-claim once; derivation's newest-keep retains the sole claim") + resp, err = huge.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__huge"]}`}) + require.NoError(t, err) + require.False(t, resp.IsError, "repeating an already claimed name costs nothing") + resp, err = huge.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__other"]}`}) + require.NoError(t, err) + require.True(t, resp.IsError, "the spent budget rejects further new activations") +} + +func TestBuildFindToolsDescription(t *testing.T) { + t.Parallel() + entries := []FindToolCatalogEntry{ + {Name: "zeta__last", Description: "Last tool. More detail", Server: "zeta", ServerDescription: strings.Repeat("z", 80)}, + {Name: "alpha__second", Description: strings.Repeat("x", 100), Server: "alpha", ServerDescription: "Alpha server"}, + {Name: "alpha__first", Description: "First tool\nmore detail", Server: "alpha", ServerDescription: "Alpha server"}, + } + description := buildFindToolsDescription(entries, 0) + require.Less(t, strings.Index(description, "## alpha"), strings.Index(description, "## zeta")) + require.Less(t, strings.Index(description, "alpha__first"), strings.Index(description, "alpha__second")) + require.Contains(t, description, "First tool") + require.NotContains(t, description, "more detail") + require.Contains(t, description, "...") + + many := make([]FindToolCatalogEntry, 300) + for i := range many { + many[i] = FindToolCatalogEntry{ + Name: fmt.Sprintf("server__tool_%03d_%s", i, strings.Repeat("n", 40)), + Description: strings.Repeat("description ", 20), + Server: "server", + } + } + degraded := buildFindToolsDescription(many, 0) + require.Contains(t, degraded, "## server (300 tools)") + require.NotContains(t, degraded, "server__tool_000") + + manyServers := make([]FindToolCatalogEntry, 500) + for i := range manyServers { + manyServers[i] = FindToolCatalogEntry{ + Name: fmt.Sprintf("server_%03d_%s__tool", i, strings.Repeat("s", 40)), + Server: fmt.Sprintf("server_%03d_%s", i, strings.Repeat("s", 40)), + } + } + countsExceeded := buildFindToolsDescription(manyServers, 0) + require.Contains(t, countsExceeded, "500 deferred tools across 500 servers.") + require.NotContains(t, countsExceeded, "## server_000") + require.LessOrEqual(t, estimatedFindToolsTokens(countsExceeded), float64(findToolsCatalogTokens)) + + smallWindow := buildFindToolsDescription(entries, 150) + require.NotContains(t, smallWindow, "First tool", + "a small context window budget forces catalog degradation below the 4000-token default") +} diff --git a/coderd/x/chatd/chattool/mcpworkspace.go b/coderd/x/chatd/chattool/mcpworkspace.go index 8f85126882643..7f9e90524f6d2 100644 --- a/coderd/x/chatd/chattool/mcpworkspace.go +++ b/coderd/x/chatd/chattool/mcpworkspace.go @@ -120,6 +120,17 @@ func buildWorkspaceMCPTool( } } +// ServerName returns the originating MCP server name from the unsanitized +// routing name. The model-facing info.Name can lose the "__" separator to +// sanitization or length capping, so it cannot be parsed for the server. +func (t *WorkspaceMCPTool) ServerName() string { + server, _, ok := strings.Cut(t.routingName, "__") + if !ok { + return "" + } + return server +} + // sanitizeModelToolName returns the provider-safe form of a workspace MCP tool // name: characters outside [a-zA-Z0-9_-] become "_" and the result is capped // at maxModelToolNameLen. The "__" server/tool separator survives because diff --git a/coderd/x/chatd/forced_mcp_test.go b/coderd/x/chatd/forced_mcp_test.go index 2447b20320303..7ef075ac563a3 100644 --- a/coderd/x/chatd/forced_mcp_test.go +++ b/coderd/x/chatd/forced_mcp_test.go @@ -90,6 +90,7 @@ func TestCreateChat_ForceOnMCPServerEnforced(t *testing.T) { }) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + withoutMCPToolSearch(cfg) cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) }) @@ -159,6 +160,7 @@ func TestSendMessage_ForceOnMCPServerEnforced(t *testing.T) { }) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + withoutMCPToolSearch(cfg) cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) }) @@ -231,6 +233,7 @@ func TestGeneration_ForceOnMCPServerEnforcedForExistingChats(t *testing.T) { user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + withoutMCPToolSearch(cfg) cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) }) diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 5fb7f00e6fd21..8b1bb8814582f 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -23,6 +23,7 @@ import ( "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" "github.com/coder/coder/v2/coderd/x/chatd/chatretry" "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/coderd/x/chatd/chattool" "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/x/agenthooks" @@ -40,13 +41,14 @@ type generationPrepared struct { Chat database.Chat Messages []database.ChatMessage - Model chatprovider.Model - Prompt []fantasy.Message - Tools []fantasy.AgentTool - ActiveTools []string - ProviderTools []chatloop.ProviderTool - ModelRoute aiGatewayModelRoute - ModelBuildOptions modelBuildOptions + Model chatprovider.Model + Prompt []fantasy.Message + Tools []fantasy.AgentTool + ActiveTools []string + AllowInactiveTools map[string]bool + ProviderTools []chatloop.ProviderTool + ModelRoute aiGatewayModelRoute + ModelBuildOptions modelBuildOptions // ResolvedProvider is the configured provider identity used to label // user-facing errors. See chatloop.GenerateAssistantOptions.ErrorProvider. @@ -781,20 +783,46 @@ func (s *taskStarter) admitStepToolCalls( if len(toolCalls) == 0 || exclusiveBatchRejected(toolCalls, prepared.ExclusiveToolNames) { return chathooks.PreToolUseExecutionResult{}, nil } + // An admission error discards the whole batch before it can be + // committed, so its find_tools calls would otherwise never reach + // the executeLocalTools counter; count them at each error exit. + countBatch := func() { + if !prepared.BuiltinToolNames[chattool.FindToolsName] { + return + } + for _, toolCall := range toolCalls { + if toolCall.ToolName == chattool.FindToolsName { + s.server.metrics.FindToolsCallsTotal.Inc() + } + } + } // Check the full batch first: a call removed below still occupies its ID // in the step, so filtering before this would hide the collision. if err := chathooks.RejectDuplicateToolUseIDs(toolCalls); err != nil { + countBatch() return chathooks.PreToolUseExecutionResult{}, chathooks.GenerationDispatchError(agenthooks.EventPreToolUse, err) } unambiguous, ambiguous := partitionAmbiguousToolCalls(prepared, toolCalls) preflight, err := s.server.hooks.PreflightPendingToolCalls(ctx, chathooks.ChatFor(prepared.Chat, input.hookTurnID()), unambiguous) if err != nil { + countBatch() return chathooks.PreToolUseExecutionResult{}, chathooks.GenerationDispatchError(agenthooks.EventPreToolUse, err) } if err := validateOverriddenToolInputs(prepared, preflight); err != nil { + countBatch() return chathooks.PreToolUseExecutionResult{}, chathooks.GenerationDispatchError(agenthooks.EventPreToolUse, err) } preflight.Denied = append(preflight.Denied, ambiguous...) + // Calls denied at admission persist synthetic results with the + // assistant step, so they never surface as unresolved calls where + // executeLocalTools counts find_tools invocations; count them here. + if prepared.BuiltinToolNames[chattool.FindToolsName] { + for _, result := range preflight.Denied { + if result.ToolName == chattool.FindToolsName { + s.server.metrics.FindToolsCallsTotal.Inc() + } + } + } return preflight, nil } @@ -810,6 +838,17 @@ func (s *taskStarter) executeLocalTools( if !exclusiveBatchRejected(decision.localToolCalls, prepared.ExclusiveToolNames) { allowed, denied = partitionAmbiguousToolCalls(prepared, decision.localToolCalls) } + // find_tools calls are counted here, at the single point every + // model-emitted call passes through, because rejections upstream of + // the tool (partition denials, hook denials, exclusive-policy + // batches) never reach its handler or OnCall. + if prepared.BuiltinToolNames[chattool.FindToolsName] { + for _, toolCall := range decision.localToolCalls { + if toolCall.ToolName == chattool.FindToolsName { + s.server.metrics.FindToolsCallsTotal.Inc() + } + } + } attempt, err := s.beginGenerationAttempt(ctx, machine, input) if err != nil { return xerrors.Errorf("beginGenerationAttempt: %w", err) @@ -827,8 +866,10 @@ func (s *taskStarter) executeLocalTools( outcome, err = chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ Tools: prepared.Tools, ActiveTools: prepared.ActiveTools, + AllowInactiveTools: prepared.AllowInactiveTools, ProviderTools: prepared.ProviderTools, ToolCalls: allowed, + ObservedToolCalls: decision.localToolCalls, ExclusiveToolNames: prepared.ExclusiveToolNames, BuiltinToolNames: prepared.BuiltinToolNames, ModelProvider: provider, diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index 6602d11c84b2c..f9085157de85b 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -529,6 +529,19 @@ func (server *Server) prepareGeneration( builtinToolNames[t.Info().Name] = true } + mcpConfigByID := make(map[uuid.UUID]database.MCPServerConfig, len(mcpConnectConfigs)) + for _, config := range mcpConnectConfigs { + mcpConfigByID[config.ID] = config + } + deferredCandidates := collectDeferredMCPCandidates(deferredMCPCandidateInput{ + mcpTools: mcpTools, + workspaceMCPTools: workspaceMCPTools, + mcpConfigByID: mcpConfigByID, + planMode: currentPlanMode, + parentChatID: chat.ParentChatID, + approvedMCPConfigIDs: approvedPlanMCPConfigIDs, + includeWorkspaceTools: !isExploreSubagent, + }) tools = append(tools, mcpTools...) if !isExploreSubagent { tools = append(tools, workspaceMCPTools...) @@ -590,6 +603,51 @@ func (server *Server) prepareGeneration( if isExploreSubagent { activeToolNames = allowedExploreToolNames(tools) } + var allowInactiveTools map[string]bool + if decideMCPToolSearch(mcpToolSearchInput{ + experimentEnabled: server.experiments.Enabled(codersdk.ExperimentMCPToolSearch), + candidates: deferredCandidates, + dynamicToolNames: dynamicToolNames, + }) { + activationTokenBudget := float64(modelConfig.ContextLimit) / mcpToolSearchBudgetDivisor + findTools := chattool.FindTools(chattool.FindToolsOptions{ + Entries: deferredMCPToolEntries(deferredCandidates), + SchemaTokenBudget: activationTokenBudget, + CatalogTokenBudget: activationTokenBudget, + // Calls total is counted in executeLocalTools, which also + // sees calls rejected before the tool runs; OnCall covers + // only calls that reach the handler or its decode. + OnCall: func(callCtx context.Context, call chattool.FindToolsCall) { + if call.Rejection == "" { + server.metrics.FindToolsMatchCount.Observe(float64(call.MatchCount)) + server.metrics.FindToolsActivationsTotal.Add(float64(len(call.Activated))) + if call.MatchCount == 0 { + server.metrics.FindToolsEmptyTotal.Inc() + } + } + // Queries and names are model output that can echo + // prompt content, so standard logs carry only + // aggregate fields; raw values are visible through + // the opt-in chat debug logging path. + logger.Info(callCtx, "deferred MCP tool search", + slog.F("query_count", len(call.Queries)), + slog.F("name_count", len(call.Names)), + slog.F("match_count", call.MatchCount), + slog.F("activated_count", len(call.Activated)), + slog.F("total_deferred", call.TotalDeferred), + slog.F("rejection", call.Rejection), + ) + }, + }) + tools, activeToolNames, allowInactiveTools = configureDeferredMCPToolSearch( + tools, + activeToolNames, + deferredCandidates, + findTools, + deriveDeferredMCPActivations(promptRows, deferredCandidates, activationTokenBudget), + ) + builtinToolNames[chattool.FindToolsName] = true + } toolNameToConfigID := make(map[string]uuid.UUID) for _, t := range tools { @@ -675,6 +733,7 @@ func (server *Server) prepareGeneration( Prompt: prompt, Tools: tools, ActiveTools: activeToolNames, + AllowInactiveTools: allowInactiveTools, ProviderTools: providerTools, ModelRoute: modelRoute, ModelBuildOptions: modelOpts, diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go new file mode 100644 index 0000000000000..a62906ffb1f98 --- /dev/null +++ b/coderd/x/chatd/mcp_tool_search.go @@ -0,0 +1,378 @@ +package chatd + +import ( + "encoding/json" + "slices" + "strconv" + "strings" + + "charm.land/fantasy" + "github.com/google/uuid" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chattool" + "github.com/coder/coder/v2/coderd/x/chatd/mcpclient" + "github.com/coder/coder/v2/codersdk" +) + +// mcpToolSearchBudgetDivisor scales the activation and catalog budgets +// to the model context window: activated schemas may re-inline up to +// ContextLimit / 10 estimated tokens per generation. +const mcpToolSearchBudgetDivisor = 10 + +type deferredMCPTool struct { + tool fantasy.AgentTool + server string + serverDescription string +} + +type deferredMCPCandidateInput struct { + mcpTools []fantasy.AgentTool + workspaceMCPTools []fantasy.AgentTool + mcpConfigByID map[uuid.UUID]database.MCPServerConfig + planMode database.NullChatPlanMode + parentChatID uuid.NullUUID + approvedMCPConfigIDs map[uuid.UUID]struct{} + includeWorkspaceTools bool +} + +// collectDeferredMCPCandidates applies the same turn policy that +// filterToolsForTurn later applies to the executable tool set, so the +// find_tools catalog never advertises tools the turn cannot run. +func collectDeferredMCPCandidates(input deferredMCPCandidateInput) []deferredMCPTool { + candidates := make([]deferredMCPTool, 0, len(input.mcpTools)+len(input.workspaceMCPTools)) + for _, tool := range input.mcpTools { + if !toolAllowedForTurn(tool, input.planMode, input.parentChatID, input.approvedMCPConfigIDs) { + continue + } + candidate := deferredMCPTool{tool: tool} + if identified, ok := tool.(mcpclient.MCPToolIdentifier); ok { + if config, exists := input.mcpConfigByID[identified.MCPServerConfigID()]; exists { + candidate.server = config.Slug + candidate.serverDescription = config.Description + } + } + candidates = append(candidates, candidate) + } + if !input.includeWorkspaceTools { + return candidates + } + wsStart := len(candidates) + for _, tool := range input.workspaceMCPTools { + if !toolAllowedForTurn(tool, input.planMode, input.parentChatID, input.approvedMCPConfigIDs) { + continue + } + candidates = append(candidates, deferredMCPTool{tool: tool, server: workspaceMCPServerName(tool)}) + } + trimWorkspaceServerNames(candidates, wsStart) + return candidates +} + +// trimWorkspaceServerNames trims surrounding whitespace from workspace +// server names, which config validation permits but find_tools strips +// from queries, so a padded server stays reachable by scope. Trimming +// is skipped when it would collapse distinct servers into one catalog +// identity (a padded and an unpadded sibling, or a collision with an +// external slug): those keep their raw names, so each server keeps its +// own catalog group and an exact-form scope matches only its own +// server. +func trimWorkspaceServerNames(candidates []deferredMCPTool, wsStart int) { + sources := make(map[string]map[string]struct{}, len(candidates)) + for i, candidate := range candidates { + key := candidate.server + if i >= wsStart { + key = strings.TrimSpace(key) + } + if sources[key] == nil { + sources[key] = make(map[string]struct{}, 1) + } + sources[key][candidate.server] = struct{}{} + } + for i := wsStart; i < len(candidates); i++ { + trimmed := strings.TrimSpace(candidates[i].server) + if len(sources[trimmed]) == 1 { + candidates[i].server = trimmed + } + } +} + +// workspaceMCPServerName prefers the wrapper's unsanitized routing name +// because sanitization can truncate the model-facing name before the +// "__" separator, which would otherwise catalog each such tool under a +// fake single-tool server that prefix scoping cannot reach. +func workspaceMCPServerName(tool fantasy.AgentTool) string { + if namer, ok := tool.(interface{ ServerName() string }); ok { + return namer.ServerName() + } + if server, _, ok := strings.Cut(tool.Info().Name, "__"); ok { + return server + } + return "" +} + +type mcpToolSearchInput struct { + experimentEnabled bool + candidates []deferredMCPTool + dynamicToolNames map[string]bool +} + +// decideMCPToolSearch reports whether MCP tool schemas are deferred +// behind find_tools. With the experiment enabled, every generation with +// deferrable candidates defers. +func decideMCPToolSearch(input mcpToolSearchInput) bool { + if !input.experimentEnabled || len(input.candidates) == 0 { + return false + } + // A client-executed dynamic tool named find_tools would otherwise be + // advertised alongside the built-in and capture its calls as + // requires_action, so a collision on either surface fails open. + if input.dynamicToolNames[chattool.FindToolsName] { + return false + } + for _, candidate := range input.candidates { + if candidate.tool.Info().Name == chattool.FindToolsName { + return false + } + } + return true +} + +func configureDeferredMCPToolSearch( + tools []fantasy.AgentTool, + activeToolNames []string, + candidates []deferredMCPTool, + findTools fantasy.AgentTool, + activations []string, +) ([]fantasy.AgentTool, []string, map[string]bool) { + candidateNames := deferredMCPToolNameSet(candidates) + ordered := make([]fantasy.AgentTool, 0, len(tools)+1) + for _, tool := range tools { + if !candidateNames[tool.Info().Name] { + ordered = append(ordered, tool) + } + } + ordered = append(ordered, findTools) + for _, tool := range tools { + if candidateNames[tool.Info().Name] { + ordered = append(ordered, tool) + } + } + + activeToolNames = slices.DeleteFunc(activeToolNames, func(name string) bool { return candidateNames[name] }) + activeToolNames = append(activeToolNames, chattool.FindToolsName) + activeToolNames = append(activeToolNames, activations...) + return ordered, activeToolNames, candidateNames +} + +func estimateDeferredMCPToolTokens(candidates []deferredMCPTool) float64 { + chars := 0 + for _, candidate := range candidates { + info := candidate.tool.Info() + schema := map[string]any{"type": "object", "properties": info.Parameters} + if len(info.Required) > 0 { + schema["required"] = info.Required + } + serialized, _ := json.Marshal(schema) + chars += len(info.Name) + len(info.Description) + len(serialized) + } + return float64(chars) / 2.5 +} + +func deferredMCPToolEntries(candidates []deferredMCPTool) []chattool.FindToolCatalogEntry { + // Workspace config validation permits an empty server key, and + // candidates whose config lookup failed also carry no server, so + // empty identities get a real label here. The label is the entry's + // Server for grouping, scope matching, and scoring alike, and it is + // collision-safe: a literal server with the same name keeps its own + // group and scope. + fallback := "workspace" + taken := make(map[string]struct{}, len(candidates)) + for _, candidate := range candidates { + if candidate.server != "" { + taken[candidate.server] = struct{}{} + } + } + for suffix := 2; ; suffix++ { + if _, collides := taken[fallback]; !collides { + break + } + fallback = "workspace-" + strconv.Itoa(suffix) + } + entries := make([]chattool.FindToolCatalogEntry, 0, len(candidates)) + for _, candidate := range candidates { + info := candidate.tool.Info() + server := candidate.server + if server == "" { + server = fallback + } + entries = append(entries, chattool.FindToolCatalogEntry{ + Name: info.Name, + Description: info.Description, + Server: server, + ServerDescription: candidate.serverDescription, + ParameterText: flattenMCPParameterText(info.Parameters), + SchemaTokens: estimateDeferredMCPToolTokens([]deferredMCPTool{candidate}), + }) + } + return entries +} + +func flattenMCPParameterText(value any) string { + var values []string + var walk func(any) + walk = func(value any) { + switch typed := value.(type) { + case map[string]any: + keys := make([]string, 0, len(typed)) + for key := range typed { + keys = append(keys, key) + } + slices.Sort(keys) + for _, key := range keys { + values = append(values, key) + walk(typed[key]) + } + case []any: + for _, item := range typed { + walk(item) + } + case string: + values = append(values, typed) + } + } + walk(value) + return strings.Join(values, " ") +} + +// deriveDeferredMCPActivations walks the surviving history newest first +// so that when the aggregate schema weight of activations exceeds +// tokenBudget, the least recently activated schemas are shed. The newest +// activation is always kept even when its schema alone exceeds the +// budget, so the tool the model just requested stays usable. Shed tools +// stay in the catalog and remain directly callable, which reactivates +// them as most recent. A tokenBudget <= 0 means unbounded. +// +// find_tools results are admitted at their tool-call row, after that +// row's direct tool calls, so a step's own search activations cannot +// shed the schema of a tool the model invoked directly. Results whose +// call row was compacted away are admitted at the result row. +// +// Direct calls whose tool result is an error are admitted last, newest +// first. History cannot distinguish a pre-execution denial (hook +// policy, input validation) from an executed call whose MCP server +// returned an error, so errored calls activate only with budget left +// after every other activation: the schema stays available for a +// corrected retry without displacing schemas that find_tools results +// or successful calls already claimed. Search-time reservations mirror +// this order because sibling calls settle before a step's searches run +// and the step result observer refunds errored reservations first, so +// searches see the same leftover budget. +func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []deferredMCPTool, tokenBudget float64) []string { + candidateByName := make(map[string]deferredMCPTool, len(candidates)) + for _, candidate := range candidates { + candidateByName[candidate.tool.Info().Name] = candidate + } + seen := make(map[string]struct{}, len(candidates)) + activated := make([]string, 0, len(candidates)) + usedTokens := 0.0 + appendName := func(name string) { + candidate, ok := candidateByName[name] + if !ok { + return + } + if _, dup := seen[name]; dup { + return + } + seen[name] = struct{}{} + weight := estimateDeferredMCPToolTokens([]deferredMCPTool{candidate}) + if len(activated) > 0 && tokenBudget > 0 && usedTokens+weight > tokenBudget { + return + } + usedTokens += weight + activated = append(activated, name) + } + parsedParts := make([][]codersdk.ChatMessagePart, len(rows)) + for i := range rows { + parts, err := chatprompt.ParseContent(rows[i]) + if err != nil { + continue + } + parsedParts[i] = parts + } + // Providers may reuse a tool-call ID in a later step, so a result + // settles the newest unpaired call with its ID: a new call abandons + // any older unpaired call, whose own result was lost or compacted + // away, and abandoned calls count as successful rather than + // adopting a later call's result. Results paired to a call + // occurrence are recorded so orphan results, whose call row was + // compacted away, are admitted at their own row even when a later + // step reuses their ID. + type partRef struct{ row, part int } + callErrored := make(map[partRef]bool) + resultPaired := make(map[partRef]struct{}) + pendingByID := make(map[string]partRef) + for i := range rows { + for j, part := range parsedParts[i] { + if part.ToolCallID == "" { + continue + } + switch part.Type { + case codersdk.ChatMessagePartTypeToolCall: + pendingByID[part.ToolCallID] = partRef{row: i, part: j} + case codersdk.ChatMessagePartTypeToolResult: + if ref, ok := pendingByID[part.ToolCallID]; ok { + callErrored[ref] = part.IsError + resultPaired[partRef{row: i, part: j}] = struct{}{} + delete(pendingByID, part.ToolCallID) + } + } + } + } + pendingSearch := make(map[string][]string) + var erroredNames []string + for i := len(rows) - 1; i >= 0; i-- { + for j, part := range parsedParts[i] { + if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName != chattool.FindToolsName { + if callErrored[partRef{row: i, part: j}] { + erroredNames = append(erroredNames, part.ToolName) + continue + } + appendName(part.ToolName) + } + } + for j, part := range parsedParts[i] { + switch { + case part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolName == chattool.FindToolsName: + var result chattool.FindToolsResult + if err := json.Unmarshal(part.Result, &result); err != nil { + continue + } + if _, paired := resultPaired[partRef{row: i, part: j}]; paired { + pendingSearch[part.ToolCallID] = result.Activated + continue + } + for _, name := range result.Activated { + appendName(name) + } + case part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName == chattool.FindToolsName: + for _, name := range pendingSearch[part.ToolCallID] { + appendName(name) + } + delete(pendingSearch, part.ToolCallID) + } + } + } + for _, name := range erroredNames { + appendName(name) + } + return activated +} + +func deferredMCPToolNameSet(candidates []deferredMCPTool) map[string]bool { + names := make(map[string]bool, len(candidates)) + for _, candidate := range candidates { + names[candidate.tool.Info().Name] = true + } + return names +} diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go new file mode 100644 index 0000000000000..7c2834290b4a9 --- /dev/null +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -0,0 +1,488 @@ +package chatd + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "charm.land/fantasy" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatloop" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chattest" + "github.com/coder/coder/v2/coderd/x/chatd/chattool" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/workspacesdk" +) + +type deferredTestAgentTool struct { + info fantasy.ToolInfo +} + +func (t deferredTestAgentTool) Info() fantasy.ToolInfo { return t.info } +func (deferredTestAgentTool) ProviderOptions() fantasy.ProviderOptions { return nil } +func (deferredTestAgentTool) SetProviderOptions(fantasy.ProviderOptions) {} +func (deferredTestAgentTool) Run(context.Context, fantasy.ToolCall) (fantasy.ToolResponse, error) { + return fantasy.NewTextResponse("ok"), nil +} + +func testDeferredTool(name, description string, parameters map[string]any) deferredMCPTool { + return deferredMCPTool{tool: deferredTestAgentTool{info: fantasy.ToolInfo{ + Name: name, Description: description, Parameters: parameters, + }}} +} + +func TestDecideMCPToolSearch(t *testing.T) { + t.Parallel() + candidates := []deferredMCPTool{testDeferredTool("server__small", "small", map[string]any{"value": map[string]any{"type": "string"}})} + + tests := []struct { + name string + experiment bool + candidates []deferredMCPTool + dynamicNames map[string]bool + want bool + }{ + {name: "experiment on", experiment: true, candidates: candidates, want: true}, + {name: "experiment off", candidates: candidates}, + {name: "empty", experiment: true}, + {name: "collision", experiment: true, candidates: []deferredMCPTool{testDeferredTool(chattool.FindToolsName, "collision", nil)}}, + {name: "dynamic collision", experiment: true, candidates: candidates, dynamicNames: map[string]bool{chattool.FindToolsName: true}}, + {name: "dynamic no collision", experiment: true, candidates: candidates, dynamicNames: map[string]bool{"other": true}, want: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tt.want, decideMCPToolSearch(mcpToolSearchInput{ + experimentEnabled: tt.experiment, + candidates: tt.candidates, + dynamicToolNames: tt.dynamicNames, + })) + }) + } +} + +func TestDeriveDeferredMCPActivations(t *testing.T) { + t.Parallel() + candidates := []deferredMCPTool{ + testDeferredTool("server__first", "first", nil), + testDeferredTool("server__second", "second", nil), + } + findResult, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolResult("call-1", chattool.FindToolsName, []byte(`{"activated":["server__second","disconnected"]}`), false, false), + }) + require.NoError(t, err) + directCall, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolCall("call-2", "server__first", []byte(`{}`)), + }) + require.NoError(t, err) + malformed, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolResult("call-3", chattool.FindToolsName, []byte(`"not-json"`), false, false), + }) + require.NoError(t, err) + rows := []database.ChatMessage{ + {Role: database.ChatMessageRoleTool, Content: findResult, ContentVersion: chatprompt.CurrentContentVersion}, + {Role: database.ChatMessageRoleAssistant, Content: directCall, ContentVersion: chatprompt.CurrentContentVersion}, + {Role: database.ChatMessageRoleTool, Content: malformed, ContentVersion: chatprompt.CurrentContentVersion}, + } + require.Equal(t, []string{"server__first", "server__second"}, deriveDeferredMCPActivations(rows, candidates, 0), + "newest activations first") + require.Equal(t, []string{"server__first"}, deriveDeferredMCPActivations(rows[1:], candidates, 0), + "activations before a compaction summary are absent from the surviving prompt window") + + firstWeight := estimateDeferredMCPToolTokens(candidates[:1]) + require.Equal(t, []string{"server__first"}, deriveDeferredMCPActivations(rows, candidates, firstWeight), + "a token budget sheds the least recent activations") + require.Equal(t, []string{"server__first"}, deriveDeferredMCPActivations(rows, candidates, 0.001), + "the newest activation survives a budget smaller than its own schema") +} + +func TestDeriveDeferredMCPActivationsSameStepDirectCallPriority(t *testing.T) { + t.Parallel() + candidates := []deferredMCPTool{ + testDeferredTool("server__direct", "direct", nil), + testDeferredTool("server__searched", "searched", nil), + } + assistantStep, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolCall("call-search", chattool.FindToolsName, []byte(`{"queries":["direct"]}`)), + codersdk.ChatMessageToolCall("call-direct", "server__direct", []byte(`{}`)), + }) + require.NoError(t, err) + searchResult, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolResult("call-search", chattool.FindToolsName, []byte(`{"activated":["server__searched"]}`), false, false), + }) + require.NoError(t, err) + rows := []database.ChatMessage{ + {Role: database.ChatMessageRoleAssistant, Content: assistantStep, ContentVersion: chatprompt.CurrentContentVersion}, + {Role: database.ChatMessageRoleTool, Content: searchResult, ContentVersion: chatprompt.CurrentContentVersion}, + } + require.Equal(t, []string{"server__direct", "server__searched"}, deriveDeferredMCPActivations(rows, candidates, 0), + "a step's direct calls outrank its own search activations") + directWeight := estimateDeferredMCPToolTokens(candidates[:1]) + require.Equal(t, []string{"server__direct"}, deriveDeferredMCPActivations(rows, candidates, directWeight), + "same-step search activations cannot shed a directly invoked tool's schema") +} + +func TestDeriveDeferredMCPActivationsErroredDirectCallsActivateLast(t *testing.T) { + t.Parallel() + candidates := []deferredMCPTool{ + testDeferredTool("server__errored", "errored", nil), + testDeferredTool("server__searched", "searched", nil), + } + assistantStep, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolCall("call-search", chattool.FindToolsName, []byte(`{"queries":["x"]}`)), + codersdk.ChatMessageToolCall("call-errored", "server__errored", []byte(`{}`)), + }) + require.NoError(t, err) + toolRow, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolResult("call-errored", "server__errored", []byte(`"remote tool error"`), true, false), + codersdk.ChatMessageToolResult("call-search", chattool.FindToolsName, []byte(`{"activated":["server__searched"]}`), false, false), + }) + require.NoError(t, err) + rows := []database.ChatMessage{ + {Role: database.ChatMessageRoleAssistant, Content: assistantStep, ContentVersion: chatprompt.CurrentContentVersion}, + {Role: database.ChatMessageRoleTool, Content: toolRow, ContentVersion: chatprompt.CurrentContentVersion}, + } + require.Equal(t, []string{"server__searched", "server__errored"}, deriveDeferredMCPActivations(rows, candidates, 0), + "an errored direct call activates last so the model keeps the schema for a corrected retry") + searchedWeight := estimateDeferredMCPToolTokens(candidates[1:]) + require.Equal(t, []string{"server__searched"}, deriveDeferredMCPActivations(rows, candidates, searchedWeight), + "an errored direct call cannot consume budget promised to the search's reported activations") + + erroredCall, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolCall("call-errored", "server__errored", []byte(`{}`)), + }) + require.NoError(t, err) + erroredResult, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolResult("call-errored", "server__errored", []byte(`"remote tool error"`), true, false), + }) + require.NoError(t, err) + erroredOnly := []database.ChatMessage{ + {Role: database.ChatMessageRoleAssistant, Content: erroredCall, ContentVersion: chatprompt.CurrentContentVersion}, + {Role: database.ChatMessageRoleTool, Content: erroredResult, ContentVersion: chatprompt.CurrentContentVersion}, + } + require.Equal(t, []string{"server__errored"}, deriveDeferredMCPActivations(erroredOnly, candidates, 0.001), + "the newest errored call keeps the first-activation allowance when nothing else activates") +} + +func TestDeferredMCPToolEntriesEmptyServerLabel(t *testing.T) { + t.Parallel() + empty := deferredMCPTool{tool: deferredTestAgentTool{info: fantasy.ToolInfo{Name: "_orphan__echo"}}} + + entries := deferredMCPToolEntries([]deferredMCPTool{empty}) + require.Equal(t, "workspace", entries[0].Server, + "an empty server identity gets the workspace label at construction, so scopes and grouping agree") + result, _ := chattool.SearchTools(entries, chattool.FindToolsArgs{Queries: []string{"workspace:"}}, chattool.SearchBudget{}) + require.Equal(t, []string{"_orphan__echo"}, result.Activated, + "the advertised workspace scope reaches the relabeled entries") + + literal := deferredMCPTool{ + tool: deferredTestAgentTool{info: fantasy.ToolInfo{Name: "workspace__run"}}, + server: "workspace", + } + colliding := deferredMCPToolEntries([]deferredMCPTool{empty, literal}) + require.Equal(t, "workspace-2", colliding[0].Server, + "a literal workspace server keeps its own identity; the fallback label steps aside") + require.Equal(t, "workspace", colliding[1].Server) + result, _ = chattool.SearchTools(colliding, chattool.FindToolsArgs{Queries: []string{"workspace:"}}, chattool.SearchBudget{}) + require.Equal(t, []string{"workspace__run"}, result.Activated, + "the workspace scope matches only the literal server") + result, _ = chattool.SearchTools(colliding, chattool.FindToolsArgs{Queries: []string{"workspace-2:"}}, chattool.SearchBudget{}) + require.Equal(t, []string{"_orphan__echo"}, result.Activated, + "the suffixed label scopes the empty-identity server") +} + +func TestDeriveDeferredMCPActivationsReusedCallIDs(t *testing.T) { + t.Parallel() + candidates := []deferredMCPTool{ + testDeferredTool("server__a", "a", nil), + testDeferredTool("server__b", "b", nil), + testDeferredTool("server__c", "c", nil), + } + row := func(t *testing.T, role database.ChatMessageRole, parts ...codersdk.ChatMessagePart) database.ChatMessage { + t.Helper() + content, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + return database.ChatMessage{Role: role, Content: content, ContentVersion: chatprompt.CurrentContentVersion} + } + // call-1 errors for server__a, is reused for a successful + // server__b call, then server__c errors under its own ID. Only + // per-call pairing keeps server__b a success: a history-wide + // errored-ID set would demote it behind the newer errored + // server__c. + rows := []database.ChatMessage{ + row(t, database.ChatMessageRoleAssistant, codersdk.ChatMessageToolCall("call-1", "server__a", []byte(`{}`))), + row(t, database.ChatMessageRoleTool, codersdk.ChatMessageToolResult("call-1", "server__a", []byte(`"boom"`), true, false)), + row(t, database.ChatMessageRoleAssistant, codersdk.ChatMessageToolCall("call-1", "server__b", []byte(`{}`))), + row(t, database.ChatMessageRoleTool, codersdk.ChatMessageToolResult("call-1", "server__b", []byte(`"ok"`), false, false)), + row(t, database.ChatMessageRoleAssistant, codersdk.ChatMessageToolCall("call-2", "server__c", []byte(`{}`))), + row(t, database.ChatMessageRoleTool, codersdk.ChatMessageToolResult("call-2", "server__c", []byte(`"boom"`), true, false)), + } + require.Equal(t, []string{"server__b", "server__c", "server__a"}, deriveDeferredMCPActivations(rows, candidates, 0), + "a reused tool-call ID pairs each call with its own result, so the later success outranks errored calls") + bWeight := estimateDeferredMCPToolTokens(candidates[1:2]) + require.Equal(t, []string{"server__b"}, deriveDeferredMCPActivations(rows, candidates, bWeight), + "under budget the reused-ID success wins over newer errored calls") + + // server__a's result is missing entirely; a later step reuses its + // ID for server__b whose result errors. The error belongs to the + // newer call, and the abandoned older call counts as successful. + missingResult := []database.ChatMessage{ + row(t, database.ChatMessageRoleAssistant, codersdk.ChatMessageToolCall("call-1", "server__a", []byte(`{}`))), + row(t, database.ChatMessageRoleAssistant, codersdk.ChatMessageToolCall("call-1", "server__b", []byte(`{}`))), + row(t, database.ChatMessageRoleTool, codersdk.ChatMessageToolResult("call-1", "server__b", []byte(`"boom"`), true, false)), + } + require.Equal(t, []string{"server__a", "server__b"}, deriveDeferredMCPActivations(missingResult, candidates, 0), + "a reused ID must not assign the newer call's result to the older missing-result call") + aWeight := estimateDeferredMCPToolTokens(candidates[:1]) + require.Equal(t, []string{"server__a"}, deriveDeferredMCPActivations(missingResult, candidates, aWeight), + "under budget the abandoned call keeps its successful position") + + // Compaction removed the orphan find_tools result's call row, and a + // later step reuses its ID for a fresh find_tools call. The orphan + // must be admitted at its own row, not stashed for a call that was + // already visited. + orphanSearch := []database.ChatMessage{ + row(t, database.ChatMessageRoleTool, codersdk.ChatMessageToolResult("call-1", chattool.FindToolsName, []byte(`{"activated":["server__a"]}`), false, false)), + row(t, database.ChatMessageRoleAssistant, codersdk.ChatMessageToolCall("call-1", chattool.FindToolsName, []byte(`{"queries":["b"]}`))), + row(t, database.ChatMessageRoleTool, codersdk.ChatMessageToolResult("call-1", chattool.FindToolsName, []byte(`{"activated":["server__b"]}`), false, false)), + } + require.Equal(t, []string{"server__b", "server__a"}, deriveDeferredMCPActivations(orphanSearch, candidates, 0), + "an orphan search result with a reused ID is admitted at its own row") +} + +func TestFlattenMCPParameterText(t *testing.T) { + t.Parallel() + text := flattenMCPParameterText(map[string]any{ + "repository": map[string]any{"type": "string", "description": "Repository name"}, + }) + require.Contains(t, text, "repository") + require.Contains(t, text, "Repository name") +} + +type deferredExternalTestTool struct { + deferredTestAgentTool + configID uuid.UUID +} + +func (t deferredExternalTestTool) MCPServerConfigID() uuid.UUID { return t.configID } + +func TestCollectDeferredMCPCandidates(t *testing.T) { + t.Parallel() + approvedID := uuid.New() + unapprovedID := uuid.New() + external := deferredExternalTestTool{ + deferredTestAgentTool: deferredTestAgentTool{info: fantasy.ToolInfo{Name: "github__create_issue"}}, + configID: approvedID, + } + unapproved := deferredExternalTestTool{ + deferredTestAgentTool: deferredTestAgentTool{info: fantasy.ToolInfo{Name: "linear__create_issue"}}, + configID: unapprovedID, + } + workspace := deferredTestAgentTool{info: fantasy.ToolInfo{Name: "everything__echo"}} + input := deferredMCPCandidateInput{ + mcpTools: []fantasy.AgentTool{external, unapproved}, + workspaceMCPTools: []fantasy.AgentTool{workspace}, + mcpConfigByID: map[uuid.UUID]database.MCPServerConfig{approvedID: {Slug: "github", Description: "GitHub"}}, + approvedMCPConfigIDs: map[uuid.UUID]struct{}{approvedID: {}}, + includeWorkspaceTools: true, + } + + names := func(candidates []deferredMCPTool) []string { + out := make([]string, 0, len(candidates)) + for _, candidate := range candidates { + out = append(out, candidate.tool.Info().Name) + } + return out + } + + all := collectDeferredMCPCandidates(input) + require.Equal(t, []string{"github__create_issue", "linear__create_issue", "everything__echo"}, names(all)) + require.Equal(t, "github", all[0].server) + require.Equal(t, "everything", all[2].server) + + planInput := input + planInput.planMode = database.NullChatPlanMode{Valid: true, ChatPlanMode: database.ChatPlanModePlan} + require.Equal(t, []string{"github__create_issue"}, names(collectDeferredMCPCandidates(planInput)), + "plan mode keeps only approved external tools, matching filterToolsForTurn") + + noWorkspace := input + noWorkspace.includeWorkspaceTools = false + require.Equal(t, []string{"github__create_issue", "linear__create_issue"}, names(collectDeferredMCPCandidates(noWorkspace))) + + longServer := strings.Repeat("s", 70) + truncated := chattool.NewWorkspaceMCPTool(workspacesdk.MCPToolInfo{Name: longServer + "__echo"}, nil, nil) + require.NotContains(t, truncated.Info().Name, "__", + "sanitization must drop the separator for this scenario to be meaningful") + truncatedInput := deferredMCPCandidateInput{ + workspaceMCPTools: []fantasy.AgentTool{truncated}, + includeWorkspaceTools: true, + } + require.Equal(t, longServer, collectDeferredMCPCandidates(truncatedInput)[0].server, + "the server comes from the unsanitized routing name, not the capped model name") + + padded := chattool.NewWorkspaceMCPTool(workspacesdk.MCPToolInfo{Name: " everything __echo"}, nil, nil) + paddedInput := deferredMCPCandidateInput{ + workspaceMCPTools: []fantasy.AgentTool{padded}, + includeWorkspaceTools: true, + } + require.Equal(t, "everything", collectDeferredMCPCandidates(paddedInput)[0].server, + "surrounding whitespace is trimmed so scope matching and catalog display see the canonical name") + + unpadded := chattool.NewWorkspaceMCPTool(workspacesdk.MCPToolInfo{Name: "everything__ping"}, nil, nil) + collidingInput := deferredMCPCandidateInput{ + workspaceMCPTools: []fantasy.AgentTool{padded, unpadded}, + includeWorkspaceTools: true, + } + colliding := collectDeferredMCPCandidates(collidingInput) + require.Equal(t, " everything ", colliding[0].server, + "trimming must not collapse distinct servers into one catalog identity") + require.Equal(t, "everything", colliding[1].server) + + slugColliding := deferredMCPCandidateInput{ + mcpTools: []fantasy.AgentTool{external}, + mcpConfigByID: map[uuid.UUID]database.MCPServerConfig{approvedID: {Slug: "everything"}}, + approvedMCPConfigIDs: map[uuid.UUID]struct{}{approvedID: {}}, + workspaceMCPTools: []fantasy.AgentTool{padded}, + includeWorkspaceTools: true, + } + slugCands := collectDeferredMCPCandidates(slugColliding) + require.Equal(t, "everything", slugCands[0].server) + require.Equal(t, " everything ", slugCands[1].server, + "a workspace server whose trimmed name collides with an external slug keeps its raw name") +} + +func TestConfigureDeferredMCPToolSearchGenerationFlows(t *testing.T) { + t.Parallel() + + hot := deferredTestAgentTool{info: fantasy.ToolInfo{Name: "read_file", Description: "Read a file"}} + first := testDeferredTool("github__create_issue", "Create an issue", nil) + second := testDeferredTool("github__list_issues", "List issues", nil) + candidates := []deferredMCPTool{first, second} + findTools := chattool.FindTools(chattool.FindToolsOptions{Entries: deferredMCPToolEntries(candidates)}) + allTools := []fantasy.AgentTool{hot, first.tool, second.tool} + allActive := []string{"read_file", first.tool.Info().Name, second.tool.Info().Name} + + ordered, active, allowInactive := configureDeferredMCPToolSearch(allTools, allActive, candidates, findTools, nil) + require.Equal(t, []string{"read_file", chattool.FindToolsName}, captureWireToolNames(t, ordered, active)) + require.Equal(t, map[string]bool{ + first.tool.Info().Name: true, + second.tool.Info().Name: true, + }, allowInactive) + + result, _ := chattool.SearchTools(deferredMCPToolEntries(candidates), chattool.FindToolsArgs{Names: []string{second.tool.Info().Name}}, chattool.SearchBudget{}) + resultJSON, err := json.Marshal(result) + require.NoError(t, err) + resultContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolResult("find-1", chattool.FindToolsName, resultJSON, false, false), + }) + require.NoError(t, err) + history := []database.ChatMessage{{ + Role: database.ChatMessageRoleTool, Content: resultContent, ContentVersion: chatprompt.CurrentContentVersion, + }} + activations := deriveDeferredMCPActivations(history, candidates, 0) + require.Equal(t, []string{second.tool.Info().Name}, activations) + + ordered, active, _ = configureDeferredMCPToolSearch(allTools, allActive, candidates, findTools, activations) + require.Equal(t, + []string{"read_file", chattool.FindToolsName, second.tool.Info().Name}, + captureWireToolNames(t, ordered, active), + ) + // Re-preparing the following turn from the same surviving history produces + // the same activation set without separate persisted state. + require.Equal(t, activations, deriveDeferredMCPActivations(history, candidates, 0)) +} + +func TestConfigureDeferredMCPToolSearchDirectCallAndCompaction(t *testing.T) { + t.Parallel() + + candidate := testDeferredTool("github__create_issue", "Create an issue", nil) + candidates := []deferredMCPTool{candidate} + directCall, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolCall("direct-1", candidate.tool.Info().Name, []byte(`{"title":"bug"}`)), + }) + require.NoError(t, err) + preSummary := []database.ChatMessage{{ + Role: database.ChatMessageRoleAssistant, Content: directCall, ContentVersion: chatprompt.CurrentContentVersion, + }} + activation := deriveDeferredMCPActivations(preSummary, candidates, 0) + require.Equal(t, []string{candidate.tool.Info().Name}, activation) + + findTools := chattool.FindTools(chattool.FindToolsOptions{Entries: deferredMCPToolEntries(candidates)}) + ordered, active, _ := configureDeferredMCPToolSearch( + []fantasy.AgentTool{candidate.tool}, + []string{candidate.tool.Info().Name}, + candidates, + findTools, + activation, + ) + require.Equal(t, + []string{chattool.FindToolsName, candidate.tool.Info().Name}, + captureWireToolNames(t, ordered, active), + ) + + // Prompt preparation passes only the post-summary history window, so an + // activation before chat_summarized naturally lapses after compaction. + require.Empty(t, deriveDeferredMCPActivations(nil, candidates, 0)) +} + +func TestMCPToolSearchExperimentDisabledPreservesWireTools(t *testing.T) { + t.Parallel() + + hot := deferredTestAgentTool{info: fantasy.ToolInfo{Name: "read_file"}} + candidate := testDeferredTool("github__list_issues", "List issues", nil) + tools := []fantasy.AgentTool{hot, candidate.tool} + active := []string{"read_file", candidate.tool.Info().Name} + withoutExperiment := captureWireToolNames(t, tools, active) + + require.False(t, decideMCPToolSearch(mcpToolSearchInput{ + candidates: []deferredMCPTool{candidate}, + })) + require.Equal(t, withoutExperiment, captureWireToolNames(t, tools, active)) +} + +func TestMCPToolSearchExploreAllowlist(t *testing.T) { + t.Parallel() + + hot := deferredTestAgentTool{info: fantasy.ToolInfo{Name: "read_file"}} + external := deferredExternalTestTool{ + deferredTestAgentTool: deferredTestAgentTool{info: fantasy.ToolInfo{Name: "github__list_issues"}}, + configID: uuid.New(), + } + tools := []fantasy.AgentTool{hot, external} + exploreActive := allowedExploreToolNames(tools) + require.Equal(t, []string{"read_file", external.Info().Name}, exploreActive) + + candidate := deferredMCPTool{tool: external} + findTools := chattool.FindTools(chattool.FindToolsOptions{Entries: deferredMCPToolEntries([]deferredMCPTool{candidate})}) + _, deferredActive, _ := configureDeferredMCPToolSearch(tools, exploreActive, []deferredMCPTool{candidate}, findTools, nil) + require.Equal(t, []string{"read_file", chattool.FindToolsName}, deferredActive) +} + +func captureWireToolNames(t *testing.T, tools []fantasy.AgentTool, active []string) []string { + t.Helper() + var names []string + model := &chattest.FakeModel{ + ProviderName: "test", + ModelName: "test", + StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { + for _, tool := range call.Tools { + names = append(names, tool.GetName()) + } + return func(yield func(fantasy.StreamPart) bool) { + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}) + }, nil + }, + } + _, err := chatloop.GenerateAssistant(context.Background(), chatloop.GenerateAssistantOptions{ + Model: model, + Tools: tools, + ActiveTools: active, + }) + require.NoError(t, err) + return names +} diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 769fade108cab..1817412829947 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -5398,6 +5398,7 @@ const ( ExperimentWorkspaceUsage Experiment = "workspace-usage" // Enables the new workspace usage tracking. ExperimentOAuth2 Experiment = "oauth2" // Enables OAuth2 provider functionality. ExperimentMCPServerHTTP Experiment = "mcp-server-http" // Enables the MCP HTTP server functionality. + ExperimentMCPToolSearch Experiment = "mcp-tool-search" // Defers MCP tool schemas behind a searchable catalog in agent chats. ExperimentWorkspaceBuildUpdates Experiment = "workspace-build-updates" // Enables publishing workspace build updates to the all builds pubsub channel. ExperimentNATSPubsub Experiment = "nats_pubsub" // Enables embedded NATS pubsub. ExperimentWorkspaceCapableLicensing Experiment = "workspace-capable-licensing" // Counts only users holding the workspace-create permission toward the license seat limit. @@ -5451,6 +5452,7 @@ var ExperimentsKnown = Experiments{ ExperimentWorkspaceUsage, ExperimentOAuth2, ExperimentMCPServerHTTP, + ExperimentMCPToolSearch, ExperimentNATSPubsub, ExperimentWorkspaceBuildUpdates, ExperimentWorkspaceCapableLicensing, diff --git a/docs/admin/integrations/prometheus.md b/docs/admin/integrations/prometheus.md index fad03a9440400..fe092ff5c88da 100644 --- a/docs/admin/integrations/prometheus.md +++ b/docs/admin/integrations/prometheus.md @@ -240,6 +240,10 @@ The `coder_ai_gateway_cost_control_*` metrics are exported only by `coderd`. | `coderd_chatd_agents_queued_for_capacity` | gauge | Deployment-wide number of chats waiting for a concurrent-agent capacity slot. Every replica reports the same database-derived value; aggregate with max, not sum. | `pool` | | `coderd_chatd_chats` | gauge | Number of chats being processed, by state. | `state` | | `coderd_chatd_compaction_total` | counter | Total compaction outcomes (only recorded when compaction was triggered or failed). | `model` `provider` `result` | +| `coderd_chatd_find_tools_activations_total` | counter | Total deferred tool activations returned by find_tools. | | +| `coderd_chatd_find_tools_calls_total` | counter | Total find_tools calls. | | +| `coderd_chatd_find_tools_empty_total` | counter | Total find_tools calls with no matches. | | +| `coderd_chatd_find_tools_match_count` | histogram | Number of matches returned by find_tools calls. | | | `coderd_chatd_hook_context_size_bytes` | histogram | Lifecycle hook model context response size in bytes. | `event` | | `coderd_chatd_hook_decisions_total` | counter | Total lifecycle hook permission decisions by event and decision. | `decision` `event` | | `coderd_chatd_hook_dispatch_seconds` | histogram | Lifecycle hook dispatch duration in seconds. | `event` | diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index b00a428094326..e5ce947f4e73e 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -7605,9 +7605,9 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o #### Enumerated Values -| Value(s) | -|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `agent-lifecycle-hooks`, `ai-gateway-seat-exclusion`, `auto-fill-parameters`, `chat-advisor`, `chat-virtual-desktop`, `example`, `mcp-server-http`, `nats_pubsub`, `notifications`, `oauth2`, `workspace-build-updates`, `workspace-capable-licensing`, `workspace-usage` | +| Value(s) | +|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `agent-lifecycle-hooks`, `ai-gateway-seat-exclusion`, `auto-fill-parameters`, `chat-advisor`, `chat-virtual-desktop`, `example`, `mcp-server-http`, `mcp-tool-search`, `nats_pubsub`, `notifications`, `oauth2`, `workspace-build-updates`, `workspace-capable-licensing`, `workspace-usage` | ## codersdk.ExternalAPIKeyScopes diff --git a/scripts/metricsdocgen/generated_metrics b/scripts/metricsdocgen/generated_metrics index a0ad8d2a5809e..bb52806447dfa 100644 --- a/scripts/metricsdocgen/generated_metrics +++ b/scripts/metricsdocgen/generated_metrics @@ -283,6 +283,18 @@ coderd_chatd_chats{state=""} 0 # HELP coderd_chatd_compaction_total Total compaction outcomes (only recorded when compaction was triggered or failed). # TYPE coderd_chatd_compaction_total counter coderd_chatd_compaction_total{provider="",model="",result=""} 0 +# HELP coderd_chatd_find_tools_activations_total Total deferred tool activations returned by find_tools. +# TYPE coderd_chatd_find_tools_activations_total counter +coderd_chatd_find_tools_activations_total 0 +# HELP coderd_chatd_find_tools_calls_total Total find_tools calls. +# TYPE coderd_chatd_find_tools_calls_total counter +coderd_chatd_find_tools_calls_total 0 +# HELP coderd_chatd_find_tools_empty_total Total find_tools calls with no matches. +# TYPE coderd_chatd_find_tools_empty_total counter +coderd_chatd_find_tools_empty_total 0 +# HELP coderd_chatd_find_tools_match_count Number of matches returned by find_tools calls. +# TYPE coderd_chatd_find_tools_match_count histogram +coderd_chatd_find_tools_match_count 0 # HELP coderd_chatd_hook_context_size_bytes Lifecycle hook model context response size in bytes. # TYPE coderd_chatd_hook_context_size_bytes histogram coderd_chatd_hook_context_size_bytes{event=""} 0 diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index da8a217888a87..e532692338ad7 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -4991,6 +4991,7 @@ export type Experiment = | "chat-virtual-desktop" | "example" | "mcp-server-http" + | "mcp-tool-search" | "nats_pubsub" | "notifications" | "oauth2" @@ -5006,6 +5007,7 @@ export const Experiments: Experiment[] = [ "chat-virtual-desktop", "example", "mcp-server-http", + "mcp-tool-search", "nats_pubsub", "notifications", "oauth2", diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx index cb1b5d4d2d09e..8dfad522e1bed 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx @@ -571,6 +571,215 @@ export const LifecycleHookNoticeAfterEditedMessage: Story = { }, }; +export const FindToolsSearchResult: Story = { + args: { + ...defaultArgs, + parsedMessages: buildMessages([ + { + ...baseMessage, + id: 1, + role: "assistant", + content: [ + { + type: "tool-call", + tool_call_id: "find-tools-1", + tool_name: "find_tools", + args: { + queries: JSON.stringify(["github issues", "pull requests"]), + names: JSON.stringify(["github__list_issues"]), + }, + }, + ], + }, + { + ...baseMessage, + id: 2, + role: "tool", + content: [ + { + type: "tool-result", + tool_call_id: "find-tools-1", + tool_name: "find_tools", + result: { + matches: JSON.stringify([ + { + name: "github__list_issues", + description: "List issues in a GitHub repository.", + }, + { + name: "github__list_pull_requests", + description: "List pull requests in a GitHub repository.", + }, + ]), + activated: JSON.stringify([ + "github__list_issues", + "github__list_pull_requests", + ]), + total_deferred: "24", + }, + }, + ], + }, + ]), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const summary = canvas.getByRole("button", { + name: "Searched tools: github issues, pull requests, name:github__list_issues -> 2 matched", + }); + expect(summary).toBeVisible(); + expect(canvas.queryByText("github__list_issues")).not.toBeInTheDocument(); + await userEvent.click(summary); + expect(canvas.getByText("github__list_issues")).toBeVisible(); + expect( + canvas.getByText("List issues in a GitHub repository."), + ).toBeVisible(); + expect(canvas.getByText("github__list_pull_requests")).toBeVisible(); + expect( + canvas.getByText("List pull requests in a GitHub repository."), + ).toBeVisible(); + }, +}; + +export const FindToolsEmptyResult: Story = { + args: { + ...defaultArgs, + parsedMessages: buildMessages([ + { + ...baseMessage, + id: 1, + role: "assistant", + content: [ + { + type: "tool-call", + tool_call_id: "find-tools-empty", + tool_name: "find_tools", + args: { queries: JSON.stringify(["nonexistent capability"]) }, + }, + ], + }, + { + ...baseMessage, + id: 2, + role: "tool", + content: [ + { + type: "tool-result", + tool_call_id: "find-tools-empty", + tool_name: "find_tools", + result: { + matches: JSON.stringify([]), + activated: JSON.stringify([]), + total_deferred: "24", + }, + }, + ], + }, + ]), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const summary = canvas.getByText( + "Searched tools: nonexistent capability -> 0 matched", + ); + expect(summary).toBeVisible(); + expect( + canvas.queryByRole("button", { + name: "Searched tools: nonexistent capability -> 0 matched", + }), + ).not.toBeInTheDocument(); + expect(canvas.queryByRole("img")).not.toBeInTheDocument(); + }, +}; + +export const FindToolsErrorResult: Story = { + args: { + ...defaultArgs, + parsedMessages: buildMessages([ + { + ...baseMessage, + id: 1, + role: "assistant", + content: [ + { + type: "tool-call", + tool_call_id: "find-tools-error", + tool_name: "find_tools", + args: { queries: JSON.stringify(["github issues"]) }, + }, + ], + }, + { + ...baseMessage, + id: 2, + role: "tool", + content: [ + { + type: "tool-result", + tool_call_id: "find-tools-error", + tool_name: "find_tools", + is_error: true, + result: { + error: + "The schema budget for this step is exhausted; call the tools already activated or retry next step.", + }, + }, + ], + }, + ]), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.getByText("Searched tools: github issues -> 0 matched"), + ).toBeVisible(); + expect( + canvas.getByRole("img", { + name: "The schema budget for this step is exhausted; call the tools already activated or retry next step.", + }), + ).toBeVisible(); + }, +}; + +export const FindToolsMalformedResultUsesDefaultRenderer: Story = { + args: { + ...defaultArgs, + parsedMessages: buildMessages([ + { + ...baseMessage, + id: 1, + role: "assistant", + content: [ + { + type: "tool-call", + tool_call_id: "find-tools-invalid", + tool_name: "find_tools", + args: { queries: JSON.stringify(["github"]) }, + }, + ], + }, + { + ...baseMessage, + id: 2, + role: "tool", + content: [ + { + type: "tool-result", + tool_call_id: "find-tools-invalid", + tool_name: "find_tools", + result: { matches: "not-json" }, + }, + ], + }, + ]), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect(canvas.queryByText(/Searched tools:/)).not.toBeInTheDocument(); + expect(canvas.getByRole("button", { name: "find_tools" })).toBeVisible(); + }, +}; + export const DurableListTemplatesToolLifecycle: Story = { args: { ...defaultArgs, diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/FindToolsTool.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/FindToolsTool.tsx new file mode 100644 index 0000000000000..a0b86a993ffde --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/FindToolsTool.tsx @@ -0,0 +1,57 @@ +import type { FC } from "react"; +import { ToolCall } from "./ToolCall"; +import type { ToolStatus } from "./utils"; + +export type FindToolsMatch = { + name: string; + description: string; +}; + +type FindToolsToolProps = { + queries: readonly string[]; + names: readonly string[]; + matches: readonly FindToolsMatch[]; + status: ToolStatus; + isError: boolean; + errorMessage?: string; +}; + +export const FindToolsTool: FC = ({ + queries, + names, + matches, + status, + isError, + errorMessage, +}) => { + const queryLabel = + [...queries, ...names.map((name) => `name:${name}`)].join(", ") || "tools"; + const label = + status === "running" + ? `Searching tools: ${queryLabel}` + : `Searched tools: ${queryLabel} -> ${matches.length} matched`; + + return ( + 0} + > + + +
    + {matches.map((match) => ( +
  • +
    + {match.name} +
    + {match.description ?
    {match.description}
    : null} +
    • + ))} +
+ + + ); +}; diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.stories.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.stories.tsx index 3edecea02caff..ce3b7c1b23b51 100644 --- a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.stories.tsx @@ -264,6 +264,20 @@ const allToolShowcaseItems: ToolShowcaseItem[] = [ build_id: "a1b2c3d4-e5f6-7890-abcd-ef1234567890", }, }, + { + name: "find_tools", + args: { queries: ["github issues"] }, + result: { + matches: [ + { + name: "github__list_issues", + description: "List issues in a GitHub repository.", + }, + ], + activated: ["github__list_issues"], + total_deferred: 12, + }, + }, { name: "unknown_tool", args: { example: true }, diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx index dc8f9e2757a05..3d553fa5706ce 100644 --- a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx @@ -15,6 +15,7 @@ import { CreateWorkspaceTool } from "./CreateWorkspaceTool"; import { DiffFileHeader } from "./DiffFileHeader"; import { EditFilesTool } from "./EditFilesTool"; import { ExecuteTool as ExecuteToolComponent } from "./ExecuteTool"; +import { type FindToolsMatch, FindToolsTool } from "./FindToolsTool"; import { ListAgentsTool } from "./ListAgentsTool"; import { ListSubagentModelsTool } from "./ListSubagentModelsTool"; import { ListTemplatesTool } from "./ListTemplatesTool"; @@ -985,6 +986,97 @@ const GenericToolRenderer: FC = ({ ); }; +const parseArray = ( + value: unknown, + parseItem: (item: unknown) => T | null, +): T[] | null => { + let array = value; + if (typeof array === "string") { + try { + array = JSON.parse(array); + } catch { + return null; + } + } + if (!Array.isArray(array)) { + return null; + } + const items: T[] = []; + for (const item of array) { + const parsed = parseItem(item); + if (parsed === null) { + return null; + } + items.push(parsed); + } + return items; +}; + +const parseStringList = (value: unknown): string[] | null => + parseArray(value, (item) => + typeof item === "string" ? item.trim() : null, + )?.filter(Boolean) ?? null; + +const parseFindToolsMatches = (value: unknown): FindToolsMatch[] | null => + parseArray(value, (item) => { + const record = asRecord(item); + return record && + typeof record.name === "string" && + typeof record.description === "string" + ? { name: record.name, description: record.description } + : null; + }); + +const FindToolsRenderer: FC = (props) => { + const parsedArgs = parseArgs(props.args); + if (!parsedArgs) { + return ; + } + const queries = + parsedArgs.queries === undefined ? [] : parseStringList(parsedArgs.queries); + const names = + parsedArgs.names === undefined ? [] : parseStringList(parsedArgs.names); + if (!queries || !names) { + return ; + } + const parsedResult = parseArgs(props.result); + if (props.isError) { + // Error results carry plain text or an error record instead of + // matches, so they render through the specialized error state + // rather than the malformed-result fallback. + const errorMessage = parsedResult + ? asString(parsedResult.error || parsedResult.message) + : asString(props.result); + return ( + + ); + } + let matches: FindToolsMatch[] | null = []; + if (props.status !== "running" || props.result !== undefined) { + matches = parsedResult ? parseFindToolsMatches(parsedResult.matches) : null; + } + if (!matches) { + return ; + } + + return ( + + ); +}; + // --------------------------------------------------------------------------- // process_signal promotes soft failures (success=false // in the result body, isError=false at protocol level) so the generic @@ -1033,6 +1125,7 @@ const StartWorkspaceRenderer: FC = ({ // --------------------------------------------------------------------------- export const toolRenderers: Record> = { + find_tools: FindToolsRenderer, execute: ExecuteRenderer, process_output: ProcessOutputRenderer, process_signal: ProcessSignalRenderer, diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/ToolIcon.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/ToolIcon.tsx index 0d5cc6448b32d..aa128484eff51 100644 --- a/site/src/pages/AgentsPage/components/ChatElements/tools/ToolIcon.tsx +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/ToolIcon.tsx @@ -9,6 +9,7 @@ import { MonitorIcon, PowerIcon, RouteIcon, + SearchIcon, ServerIcon, TerminalIcon, WrenchIcon, @@ -44,6 +45,7 @@ export const toolIcons: Partial> = { ask_user_question: BadgeQuestionMarkIcon, advisor: CompassIcon, computer: MonitorIcon, + find_tools: SearchIcon, }; export const ToolIcon: React.FC<{