From a80335e6ed6e43938f274b34961e11971a2076db Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 16:44:27 +0000 Subject: [PATCH 01/48] feat(coderd/x/chatd): defer MCP tool schemas behind find_tools search --- cli/testdata/server-config.yaml.golden | 4 + coderd/apidoc/docs.go | 7 + coderd/apidoc/swagger.json | 7 + coderd/coderd.go | 1 + coderd/x/chatd/chatd.go | 5 + coderd/x/chatd/chatloop/chatloop.go | 15 +- .../chatloop/chatloop_run_internal_test.go | 34 +++ coderd/x/chatd/chatloop/metrics.go | 59 +++- coderd/x/chatd/chattool/findtools.go | 275 ++++++++++++++++++ .../chatd/chattool/findtools_internal_test.go | 100 +++++++ coderd/x/chatd/generation.go | 16 +- coderd/x/chatd/generation_preparer.go | 60 ++++ coderd/x/chatd/mcp_tool_search.go | 153 ++++++++++ .../x/chatd/mcp_tool_search_internal_test.go | 100 +++++++ codersdk/deployment.go | 28 +- docs/admin/integrations/prometheus.md | 5 + docs/reference/api/general.md | 3 +- docs/reference/api/schemas.md | 37 ++- scripts/metricsdocgen/generated_metrics | 15 + site/src/api/typesGenerated.ts | 3 + 20 files changed, 880 insertions(+), 47 deletions(-) create mode 100644 coderd/x/chatd/chattool/findtools.go create mode 100644 coderd/x/chatd/chattool/findtools_internal_test.go create mode 100644 coderd/x/chatd/mcp_tool_search.go create mode 100644 coderd/x/chatd/mcp_tool_search_internal_test.go diff --git a/cli/testdata/server-config.yaml.golden b/cli/testdata/server-config.yaml.golden index b2dca8d3ae7..dfc5bad74a8 100644 --- a/cli/testdata/server-config.yaml.golden +++ b/cli/testdata/server-config.yaml.golden @@ -799,6 +799,10 @@ chat: # How many pending chats a worker should acquire per polling cycle. # (default: 10, type: int) acquireBatchSize: 10 + # Force MCP tool schemas behind find_tools regardless of size. The mcp-tool-search + # experiment must also be enabled. + # (default: false, type: bool) + mcpToolSearchForceDefer: false # Force chat debug logging on for every chat, bypassing the runtime admin and user # opt-in settings. # (default: false, type: bool) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index e24d4196b51..b633631bee5 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -17618,6 +17618,9 @@ const docTemplate = `{ }, "hook_url": { "$ref": "#/definitions/serpent.URL" + }, + "mcp_tool_search_force_defer": { + "type": "boolean" } } }, @@ -20554,6 +20557,7 @@ const docTemplate = `{ "workspace-usage", "oauth2", "mcp-server-http", + "mcp-tool-search", "workspace-build-updates", "nats_pubsub", "workspace-capable-licensing", @@ -20570,6 +20574,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 +20589,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 +20605,7 @@ const docTemplate = `{ "ExperimentWorkspaceUsage", "ExperimentOAuth2", "ExperimentMCPServerHTTP", + "ExperimentMCPToolSearch", "ExperimentWorkspaceBuildUpdates", "ExperimentNATSPubsub", "ExperimentWorkspaceCapableLicensing", diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 0b18135d1cf..16c8c7fcfaf 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -15843,6 +15843,9 @@ }, "hook_url": { "$ref": "#/definitions/serpent.URL" + }, + "mcp_tool_search_force_defer": { + "type": "boolean" } } }, @@ -18676,6 +18679,7 @@ "workspace-usage", "oauth2", "mcp-server-http", + "mcp-tool-search", "workspace-build-updates", "nats_pubsub", "workspace-capable-licensing", @@ -18692,6 +18696,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 +18711,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 +18727,7 @@ "ExperimentWorkspaceUsage", "ExperimentOAuth2", "ExperimentMCPServerHTTP", + "ExperimentMCPToolSearch", "ExperimentWorkspaceBuildUpdates", "ExperimentNATSPubsub", "ExperimentWorkspaceCapableLicensing", diff --git a/coderd/coderd.go b/coderd/coderd.go index c7effe22b59..f4c68b87d8f 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -933,6 +933,7 @@ func New(options *Options) *API { AIBridgeTransportFactory: &api.AIBridgeTransportFactory, AlwaysEnableDebugLogs: options.DeploymentValues.AI.Chat.DebugLoggingEnabled.Value(), Experiments: experiments, + ForceMCPToolSearch: options.DeploymentValues.AI.Chat.MCPToolSearchForceDefer.Value(), AgentConn: api.agentProvider.AgentConn, AgentInactiveDisconnectTimeout: api.AgentInactiveDisconnectTimeout, InstructionLookupTimeout: options.ChatdInstructionLookupTimeout, diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 41966f45daa..c58a3ea5f61 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -200,6 +200,7 @@ type Server struct { aibridgeTransportFactory *atomic.Pointer[aibridge.TransportFactory] experiments codersdk.Experiments + forceMCPToolSearch bool // Configuration pendingChatAcquireInterval time.Duration @@ -3048,6 +3049,9 @@ type Config struct { Clock quartz.Clock AIBridgeTransportFactory *atomic.Pointer[aibridge.TransportFactory] Experiments codersdk.Experiments + // ForceMCPToolSearch ignores the schema-size threshold for development and tests. + // The mcp-tool-search experiment remains required. + ForceMCPToolSearch bool PrometheusRegistry prometheus.Registerer @@ -3155,6 +3159,7 @@ func New(ps pubsub.Pubsub, cfg Config) *Server { }, aibridgeTransportFactory: cfg.AIBridgeTransportFactory, experiments: cfg.Experiments, + forceMCPToolSearch: cfg.ForceMCPToolSearch, pendingChatAcquireInterval: pendingChatAcquireInterval, maxChatsPerAcquire: maxChatsPerAcquire, inFlightChatStaleAfter: inFlightChatStaleAfter, diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 89063ff0612..cd3fd8ed4f4 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -238,10 +238,11 @@ 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 ExclusiveToolNames map[string]bool BuiltinToolNames map[string]bool @@ -594,6 +595,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool opts.Clock, opts.Tools, opts.ActiveTools, + opts.AllowInactiveTools, opts.ProviderTools, localCalls, opts.Metrics, @@ -1057,6 +1059,7 @@ func executeTools( clock quartz.Clock, allTools []fantasy.AgentTool, activeTools []string, + allowInactiveTools map[string]bool, providerTools []ProviderTool, toolCalls []fantasy.ToolCallContent, metrics *Metrics, @@ -1141,6 +1144,7 @@ func executeTools( model, builtinToolNames, activeTools, + allowInactiveTools, providerRunnerNames, resultProviderMetadata, maxResultBytes, @@ -1262,6 +1266,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 +1299,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), } diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 71af2ca7f1b..89ae523aa5e 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -912,6 +912,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { "fake", "fake-model", map[string]bool{}, []string{"screenshot"}, + nil, map[string]struct{}{}, nil, defaultToolResultBytes, @@ -961,6 +962,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { "fake", "fake-model", map[string]bool{}, []string{"screenshot"}, + nil, map[string]struct{}{}, nil, defaultToolResultBytes, @@ -1005,6 +1007,7 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { "fake", "fake-model", map[string]bool{}, []string{"echo"}, + nil, map[string]struct{}{}, nil, defaultToolResultBytes, @@ -1053,6 +1056,7 @@ func TestExecuteSingleTool_ResolvesToolNameAlias(t *testing.T) { "fake", "fake-model", map[string]bool{}, []string{"interrupt_agent"}, + nil, map[string]struct{}{}, nil, defaultToolResultBytes, @@ -1093,6 +1097,7 @@ func TestExecuteSingleTool_UnknownAliasFallsThrough(t *testing.T) { "fake", "fake-model", map[string]bool{}, []string{"interrupt_agent"}, + nil, map[string]struct{}{}, nil, defaultToolResultBytes, @@ -1103,3 +1108,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 3263f01d87c..9e2193957d2 100644 --- a/coderd/x/chatd/chatloop/metrics.go +++ b/coderd/x/chatd/chatloop/metrics.go @@ -27,17 +27,22 @@ 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 + DeferredMCPToolTokens *prometheus.HistogramVec } // NewMetrics creates a new Metrics instance registered with the @@ -109,6 +114,38 @@ 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.", + }), + DeferredMCPToolTokens: factory.NewHistogramVec(prometheus.HistogramOpts{ + Namespace: metricsNamespace, + Subsystem: metricsSubsystem, + Name: "deferred_mcp_tool_tokens", + Help: "Estimated MCP tool schema tokens considered for deferral per generation.", + Buckets: prometheus.ExponentialBuckets(128, 2, 12), + }, []string{"provider", "model", "applied"}), 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 00000000000..72178f821a6 --- /dev/null +++ b/coderd/x/chatd/chattool/findtools.go @@ -0,0 +1,275 @@ +package chattool + +import ( + "context" + "regexp" + "slices" + "strconv" + "strings" + "unicode/utf8" + + "charm.land/fantasy" +) + +const ( + FindToolsName = "find_tools" + findToolsMaxMatches = 20 + findToolsCatalogTokens = 4000 +) + +var findToolsTokenSeparator = regexp.MustCompile(`[^a-z0-9]+`) + +// FindToolCatalogEntry is the searchable metadata for one deferred tool. +type FindToolCatalogEntry struct { + Name string + Description string + Server string + ServerDescription string + ParameterText string +} + +// FindToolsCall records one catalog search for logging and metrics. +type FindToolsCall struct { + Queries []string + Names []string + MatchCount int + Activated []string + TotalDeferred int +} + +// FindToolsOptions configures the find_tools tool. +type FindToolsOptions struct { + Entries []FindToolCatalogEntry + OnCall func(context.Context, FindToolsCall) +} + +// FindToolsArgs are the arguments for the find_tools tool. +type FindToolsArgs struct { + Queries []string `json:"queries"` + Names []string `json:"names"` +} + +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"` +} + +// FindTools returns the built-in used to discover deferred MCP tool schemas. +func FindTools(options FindToolsOptions) fantasy.AgentTool { + entries := slices.Clone(options.Entries) + return fantasy.NewAgentTool( + FindToolsName, + buildFindToolsDescription(entries), + func(ctx context.Context, args FindToolsArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + if len(args.Queries) == 0 && len(args.Names) == 0 { + return fantasy.NewTextErrorResponse("at least one query or name is required"), nil + } + result := SearchTools(entries, args) + 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 + }, + ) +} + +// SearchTools searches entries and returns the activation set. +func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs) FindToolsResult { + byName := make(map[string]FindToolCatalogEntry, len(entries)) + for _, entry := range entries { + byName[entry.Name] = entry + } + + type scoredEntry struct { + entry FindToolCatalogEntry + score int + } + scored := make([]scoredEntry, 0, len(entries)) + for _, entry := range entries { + score := 0 + for _, query := range args.Queries { + for _, token := range tokenizeFindTools(query) { + score += scoreFindToolToken(entry, 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) + }) + if len(scored) > findToolsMaxMatches { + scored = scored[:findToolsMaxMatches] + } + + matches := make([]FindToolsMatch, 0, len(scored)+len(args.Names)) + activatedSet := make(map[string]struct{}, len(scored)+len(args.Names)) + for _, item := range scored { + matches = append(matches, FindToolsMatch{Name: item.entry.Name, Description: item.entry.Description}) + activatedSet[item.entry.Name] = struct{}{} + } + for _, name := range args.Names { + entry, ok := byName[name] + if !ok { + continue + } + if _, exists := activatedSet[name]; !exists { + matches = append(matches, FindToolsMatch{Name: entry.Name, Description: entry.Description}) + activatedSet[name] = struct{}{} + } + } + 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)} +} + +func tokenizeFindTools(value string) []string { + parts := findToolsTokenSeparator.Split(strings.ToLower(value), -1) + return slices.DeleteFunc(parts, func(part string) bool { return part == "" }) +} + +func scoreFindToolToken(entry FindToolCatalogEntry, token string) int { + name := strings.ToLower(entry.Name) + nameTokens := tokenizeFindTools(name) + score := 0 + if slices.Contains(nameTokens, token) { + score += 8 + } else if strings.Contains(name, token) { + score += 5 + } + if slices.Contains(tokenizeFindTools(entry.Description), token) { + score += 2 + } + if slices.Contains(tokenizeFindTools(entry.ParameterText), token) { + score++ + } + return score +} + +func buildFindToolsDescription(entries []FindToolCatalogEntry) string { + const usage = "Search deferred MCP tools by keyword, activate exact tool names, or scope queries with a server prefix. Calling a cataloged tool directly by name is allowed and auto-loads its schema, but search first for unfamiliar tools.\n\n" + catalog := detailedFindToolsCatalog(entries) + if estimatedFindToolsTokens(usage+catalog) > findToolsCatalogTokens { + catalog = namesOnlyFindToolsCatalog(entries) + } + if estimatedFindToolsTokens(usage+catalog) > findToolsCatalogTokens { + catalog = countsOnlyFindToolsCatalog(entries) + } + return usage + catalog +} + +func detailedFindToolsCatalog(entries []FindToolCatalogEntry) string { + return renderFindToolsCatalog(entries, func(entry FindToolCatalogEntry) string { + return "- " + entry.Name + " - " + truncateFindToolsSummary(entry.Description, 80) + }) +} + +func namesOnlyFindToolsCatalog(entries []FindToolCatalogEntry) string { + return renderFindToolsCatalog(entries, func(entry FindToolCatalogEntry) string { return entry.Name }) +} + +func countsOnlyFindToolsCatalog(entries []FindToolCatalogEntry) string { + groups := groupFindToolsEntries(entries) + 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() +} + +type findToolsGroup struct { + server string + description string + entries []FindToolCatalogEntry +} + +func groupFindToolsEntries(entries []FindToolCatalogEntry) []findToolsGroup { + grouped := make(map[string]*findToolsGroup) + for _, entry := range entries { + server := entry.Server + if server == "" { + server = "workspace" + } + 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 renderFindToolsCatalog(entries []FindToolCatalogEntry, renderEntry func(FindToolCatalogEntry) string) string { + groups := groupFindToolsEntries(entries) + var b strings.Builder + for _, group := range groups { + _, _ = b.WriteString("## ") + _, _ = b.WriteString(group.server) + if summary := truncateFindToolsSummary(group.description, 60); summary != "" { + _, _ = b.WriteString(" - ") + _, _ = b.WriteString(summary) + } + _ = b.WriteByte('\n') + if len(group.entries) > 0 && !strings.HasPrefix(renderEntry(group.entries[0]), "-") { + names := make([]string, 0, len(group.entries)) + for _, entry := range group.entries { + names = append(names, renderEntry(entry)) + } + _, _ = b.WriteString(strings.Join(names, " ")) + _ = b.WriteByte('\n') + continue + } + for _, entry := range group.entries { + _, _ = b.WriteString(renderEntry(entry)) + _ = b.WriteByte('\n') + } + } + return b.String() +} + +func truncateFindToolsSummary(value string, maxRunes int) string { + value = strings.TrimSpace(strings.Split(strings.Split(value, "\n")[0], ". ")[0]) + 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 00000000000..428fcb3cf76 --- /dev/null +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -0,0 +1,100 @@ +package chattool + +import ( + "context" + "encoding/json" + "fmt" + "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"}}) + 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"}}) + 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"}}) + 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{}) + 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"}}) + require.Len(t, result.Matches, findToolsMaxMatches) + require.Equal(t, "server__tool_00", result.Matches[0].Name) + }) +} + +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 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) + 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) + require.Contains(t, degraded, "## server (300 tools)") + require.NotContains(t, degraded, "server__tool_000") +} diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 5fb7f00e6fd..de792a3db8d 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -40,13 +40,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. @@ -827,6 +828,7 @@ 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, ExclusiveToolNames: prepared.ExclusiveToolNames, diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index 6602d11c84b..7e05da460e1 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "slices" + "strconv" "strings" "sync" @@ -529,8 +530,30 @@ 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 := make([]deferredMCPTool, 0, len(mcpTools)+len(workspaceMCPTools)) + for _, tool := range mcpTools { + if !toolAllowedForTurn(tool, currentPlanMode, chat.ParentChatID, approvedPlanMCPConfigIDs) { + continue + } + candidate := deferredMCPTool{tool: tool} + if identified, ok := tool.(mcpclient.MCPToolIdentifier); ok { + if config, exists := mcpConfigByID[identified.MCPServerConfigID()]; exists { + candidate.server = config.Slug + candidate.serverDescription = config.Description + } + } + deferredCandidates = append(deferredCandidates, candidate) + } tools = append(tools, mcpTools...) if !isExploreSubagent { + for _, tool := range workspaceMCPTools { + serverName, _, _ := strings.Cut(tool.Info().Name, "__") + deferredCandidates = append(deferredCandidates, deferredMCPTool{tool: tool, server: serverName}) + } tools = append(tools, workspaceMCPTools...) } tools = filterToolsForTurn(tools, currentPlanMode, chat.ParentChatID, approvedPlanMCPConfigIDs) @@ -590,6 +613,42 @@ func (server *Server) prepareGeneration( if isExploreSubagent { activeToolNames = allowedExploreToolNames(tools) } + var allowInactiveTools map[string]bool + toolSearch := decideMCPToolSearch(mcpToolSearchInput{ + experimentEnabled: server.experiments.Enabled(codersdk.ExperimentMCPToolSearch), + forceDefer: server.forceMCPToolSearch, + contextWindow: modelConfig.ContextLimit, + candidates: deferredCandidates, + }) + server.metrics.DeferredMCPToolTokens.WithLabelValues( + model.Provider(), model.ModelID(), strconv.FormatBool(toolSearch.apply), + ).Observe(toolSearch.estimatedTokens) + if toolSearch.apply { + findTools := chattool.FindTools(chattool.FindToolsOptions{ + Entries: deferredMCPToolEntries(deferredCandidates), + OnCall: func(callCtx context.Context, call chattool.FindToolsCall) { + server.metrics.FindToolsCallsTotal.Inc() + server.metrics.FindToolsMatchCount.Observe(float64(call.MatchCount)) + server.metrics.FindToolsActivationsTotal.Add(float64(len(call.Activated))) + if call.MatchCount == 0 { + server.metrics.FindToolsEmptyTotal.Inc() + } + logger.Info(callCtx, "deferred MCP tool search", + slog.F("queries", call.Queries), + slog.F("names", call.Names), + slog.F("match_count", call.MatchCount), + slog.F("activated", call.Activated), + slog.F("total_deferred", call.TotalDeferred), + ) + }, + }) + tools = append(tools, findTools) + builtinToolNames[chattool.FindToolsName] = true + allowInactiveTools = deferredMCPToolNameSet(deferredCandidates) + activeToolNames = slices.DeleteFunc(activeToolNames, func(name string) bool { return allowInactiveTools[name] }) + activeToolNames = append(activeToolNames, chattool.FindToolsName) + activeToolNames = append(activeToolNames, deriveDeferredMCPActivations(promptRows, deferredCandidates)...) + } toolNameToConfigID := make(map[string]uuid.UUID) for _, t := range tools { @@ -675,6 +734,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 00000000000..336ead756cc --- /dev/null +++ b/coderd/x/chatd/mcp_tool_search.go @@ -0,0 +1,153 @@ +package chatd + +import ( + "encoding/json" + "slices" + "strings" + + "charm.land/fantasy" + + "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/codersdk" +) + +const mcpToolSearchThresholdDivisor = 10 + +type deferredMCPTool struct { + tool fantasy.AgentTool + server string + serverDescription string +} + +type mcpToolSearchDecision struct { + apply bool + estimatedTokens float64 +} + +type mcpToolSearchInput struct { + experimentEnabled bool + forceDefer bool + contextWindow int64 + candidates []deferredMCPTool +} + +func decideMCPToolSearch(input mcpToolSearchInput) mcpToolSearchDecision { + experimentEnabled, force, contextWindow, candidates := input.experimentEnabled, input.forceDefer, input.contextWindow, input.candidates + decision := mcpToolSearchDecision{estimatedTokens: estimateDeferredMCPToolTokens(candidates)} + if !experimentEnabled || len(candidates) == 0 { + return decision + } + for _, candidate := range candidates { + if candidate.tool.Info().Name == chattool.FindToolsName { + return decision + } + } + decision.apply = force || (contextWindow > 0 && decision.estimatedTokens > float64(contextWindow)/mcpToolSearchThresholdDivisor) + return decision +} + +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 { + entries := make([]chattool.FindToolCatalogEntry, 0, len(candidates)) + for _, candidate := range candidates { + info := candidate.tool.Info() + entries = append(entries, chattool.FindToolCatalogEntry{ + Name: info.Name, + Description: info.Description, + Server: candidate.server, + ServerDescription: candidate.serverDescription, + ParameterText: flattenMCPParameterText(info.Parameters), + }) + } + 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, " ") +} + +func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []deferredMCPTool) []string { + current := make(map[string]struct{}, len(candidates)) + for _, candidate := range candidates { + current[candidate.tool.Info().Name] = struct{}{} + } + seen := make(map[string]struct{}, len(candidates)) + activated := make([]string, 0, len(candidates)) + appendName := func(name string) { + if _, ok := current[name]; !ok { + return + } + if _, ok := seen[name]; ok { + return + } + seen[name] = struct{}{} + activated = append(activated, name) + } + for _, row := range rows { + parts, err := chatprompt.ParseContent(row) + if err != nil { + continue + } + for _, part := range parts { + switch { + case part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolName == chattool.FindToolsName: + var result chattool.FindToolsResult + if err := json.Unmarshal(part.Result, &result); err != nil { + continue + } + for _, name := range result.Activated { + appendName(name) + } + case part.Type == codersdk.ChatMessagePartTypeToolCall: + appendName(part.ToolName) + } + } + } + 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 00000000000..382c1777d35 --- /dev/null +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -0,0 +1,100 @@ +package chatd + +import ( + "context" + "strings" + "testing" + + "charm.land/fantasy" + "github.com/stretchr/testify/require" + + "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/codersdk" +) + +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() + small := []deferredMCPTool{testDeferredTool("server__small", "small", map[string]any{"value": map[string]any{"type": "string"}})} + large := []deferredMCPTool{testDeferredTool("server__large", strings.Repeat("large ", 2000), map[string]any{"value": map[string]any{"type": "string"}})} + + tests := []struct { + name string + experiment bool + force bool + window int64 + candidates []deferredMCPTool + want bool + }{ + {name: "below", experiment: true, window: 100_000, candidates: small}, + {name: "above", experiment: true, window: 10_000, candidates: large, want: true}, + {name: "forced", experiment: true, force: true, window: 100_000, candidates: small, want: true}, + {name: "experiment off", force: true, window: 10, candidates: large}, + {name: "empty", experiment: true, force: true}, + {name: "collision", experiment: true, force: true, candidates: []deferredMCPTool{testDeferredTool(chattool.FindToolsName, "collision", nil)}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tt.want, decideMCPToolSearch(mcpToolSearchInput{ + experimentEnabled: tt.experiment, + forceDefer: tt.force, + contextWindow: tt.window, + candidates: tt.candidates, + }).apply) + }) + } +} + +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__second", "server__first"}, deriveDeferredMCPActivations(rows, candidates)) +} + +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") +} diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 769fade108c..90588148884 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -4341,6 +4341,17 @@ Write out the current server config as YAML to stdout.`, YAML: "acquireBatchSize", Hidden: true, // Hidden because most operators should not need to modify this. }, + { + Name: "Chat: MCP Tool Search Force Defer", + Description: "Force MCP tool schemas behind find_tools regardless of size. The mcp-tool-search experiment must also be enabled.", + Flag: "chat-mcp-tool-search-force-defer", + Env: "CODER_CHAT_MCP_TOOL_SEARCH_FORCE_DEFER", + Value: &c.AI.Chat.MCPToolSearchForceDefer, + Default: "false", + Group: &deploymentGroupChat, + YAML: "mcpToolSearchForceDefer", + Hidden: true, + }, { Name: "Chat: Debug Logging Enabled", Description: "Force chat debug logging on for every chat, bypassing the runtime admin and user opt-in settings.", @@ -5089,13 +5100,14 @@ type AIBridgeProxyConfig struct { } type ChatConfig struct { - AcquireBatchSize serpent.Int64 `json:"acquire_batch_size" typescript:",notnull"` - DebugLoggingEnabled serpent.Bool `json:"debug_logging_enabled" typescript:",notnull"` - HookURL serpent.URL `json:"hook_url" typescript:",notnull"` - HookSecret serpent.String `json:"hook_secret" typescript:",notnull"` - HookTimeout serpent.Duration `json:"hook_timeout" typescript:",notnull"` - HookEnabled serpent.Bool `json:"hook_enabled" typescript:",notnull"` - HookAllowInsecure serpent.Bool `json:"hook_allow_insecure" typescript:",notnull"` + AcquireBatchSize serpent.Int64 `json:"acquire_batch_size" typescript:",notnull"` + DebugLoggingEnabled serpent.Bool `json:"debug_logging_enabled" typescript:",notnull"` + MCPToolSearchForceDefer serpent.Bool `json:"mcp_tool_search_force_defer" typescript:",notnull"` + HookURL serpent.URL `json:"hook_url" typescript:",notnull"` + HookSecret serpent.String `json:"hook_secret" typescript:",notnull"` + HookTimeout serpent.Duration `json:"hook_timeout" typescript:",notnull"` + HookEnabled serpent.Bool `json:"hook_enabled" typescript:",notnull"` + HookAllowInsecure serpent.Bool `json:"hook_allow_insecure" typescript:",notnull"` // Deprecated: AI Gateway routing is now the only routing path. Setting this // value has no effect. This option will be removed in a future release. AIGatewayRoutingEnabled serpent.Bool `json:"ai_gateway_routing_enabled" typescript:",notnull" swaggerignore:"true"` @@ -5398,6 +5410,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 +5464,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 fad03a94404..d240e9ec697 100644 --- a/docs/admin/integrations/prometheus.md +++ b/docs/admin/integrations/prometheus.md @@ -240,6 +240,11 @@ 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_deferred_mcp_tool_tokens` | histogram | Estimated MCP tool schema tokens considered for deferral per generation. | `applied` `model` `provider` | +| `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/general.md b/docs/reference/api/general.md index 982b5064099..3ea6e0e40d6 100644 --- a/docs/reference/api/general.md +++ b/docs/reference/api/general.md @@ -254,7 +254,8 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \ "rawQuery": "string", "scheme": "string", "user": {} - } + }, + "mcp_tool_search_force_defer": true } }, "allow_workspace_renames": true, diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index b00a4280943..27e87fac918 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -1103,7 +1103,8 @@ title: Schemas "rawQuery": "string", "scheme": "string", "user": {} - } + }, + "mcp_tool_search_force_defer": true } } ``` @@ -2532,21 +2533,23 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "rawQuery": "string", "scheme": "string", "user": {} - } + }, + "mcp_tool_search_force_defer": true } ``` ### Properties -| Name | Type | Required | Restrictions | Description | -|-------------------------|----------------------------|----------|--------------|-------------| -| `acquire_batch_size` | integer | false | | | -| `debug_logging_enabled` | boolean | false | | | -| `hook_allow_insecure` | boolean | false | | | -| `hook_enabled` | boolean | false | | | -| `hook_secret` | string | false | | | -| `hook_timeout` | integer | false | | | -| `hook_url` | [serpent.URL](#serpenturl) | false | | | +| Name | Type | Required | Restrictions | Description | +|-------------------------------|----------------------------|----------|--------------|-------------| +| `acquire_batch_size` | integer | false | | | +| `debug_logging_enabled` | boolean | false | | | +| `hook_allow_insecure` | boolean | false | | | +| `hook_enabled` | boolean | false | | | +| `hook_secret` | string | false | | | +| `hook_timeout` | integer | false | | | +| `hook_url` | [serpent.URL](#serpenturl) | false | | | +| `mcp_tool_search_force_defer` | boolean | false | | | ## codersdk.ChatContext @@ -6003,7 +6006,8 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "rawQuery": "string", "scheme": "string", "user": {} - } + }, + "mcp_tool_search_force_defer": true } }, "allow_workspace_renames": true, @@ -6631,7 +6635,8 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "rawQuery": "string", "scheme": "string", "user": {} - } + }, + "mcp_tool_search_force_defer": true } }, "allow_workspace_renames": true, @@ -7605,9 +7610,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 a0ad8d2a580..fe9a8ad331f 100644 --- a/scripts/metricsdocgen/generated_metrics +++ b/scripts/metricsdocgen/generated_metrics @@ -283,6 +283,21 @@ 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_deferred_mcp_tool_tokens Estimated MCP tool schema tokens considered for deferral per generation. +# TYPE coderd_chatd_deferred_mcp_tool_tokens histogram +coderd_chatd_deferred_mcp_tool_tokens{provider="",model="",applied=""} 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 da8a217888a..456f4d75f7e 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -2073,6 +2073,7 @@ export const ChatComputerUseProviders: ChatComputerUseProvider[] = [ export interface ChatConfig { readonly acquire_batch_size: number; readonly debug_logging_enabled: boolean; + readonly mcp_tool_search_force_defer: boolean; readonly hook_url: string; readonly hook_secret: string; readonly hook_timeout: number; @@ -4991,6 +4992,7 @@ export type Experiment = | "chat-virtual-desktop" | "example" | "mcp-server-http" + | "mcp-tool-search" | "nats_pubsub" | "notifications" | "oauth2" @@ -5006,6 +5008,7 @@ export const Experiments: Experiment[] = [ "chat-virtual-desktop", "example", "mcp-server-http", + "mcp-tool-search", "nats_pubsub", "notifications", "oauth2", From 80b3a4f90c07f5d2f29e821b27df3cae8774c182 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 16:55:41 +0000 Subject: [PATCH 02/48] test(coderd/x/chatd): cover deferred tool search generation flows --- coderd/x/chatd/chatd_internal_test.go | 2 + coderd/x/chatd/chatd_test.go | 215 ++++++++++++++++++ coderd/x/chatd/generation_preparer.go | 12 +- coderd/x/chatd/mcp_tool_search.go | 27 +++ .../x/chatd/mcp_tool_search_internal_test.go | 147 ++++++++++++ 5 files changed, 398 insertions(+), 5 deletions(-) diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index ae5b9bcd8fd..0025d6a2d62 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 53e4a4fd087..a5ca8fd1be5 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -10410,6 +10410,221 @@ 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)) + cfg.ForceMCPToolSearch = true + }) + 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)) + cfg.ForceMCPToolSearch = true + }) + 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("below threshold preserves wire tools", 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 { + cfg.Experiments = slices.DeleteFunc(slices.Clone(cfg.Experiments), func(experiment codersdk.Experiment) bool { + return experiment == codersdk.ExperimentMCPToolSearch + }) + } + }) + 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...) + } + + require.Equal(t, run(t, false), run(t, true)) + }) +} + // 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. diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index 7e05da460e1..69d1be97b2f 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -642,12 +642,14 @@ func (server *Server) prepareGeneration( ) }, }) - tools = append(tools, findTools) + tools, activeToolNames, allowInactiveTools = configureDeferredMCPToolSearch( + tools, + activeToolNames, + deferredCandidates, + findTools, + deriveDeferredMCPActivations(promptRows, deferredCandidates), + ) builtinToolNames[chattool.FindToolsName] = true - allowInactiveTools = deferredMCPToolNameSet(deferredCandidates) - activeToolNames = slices.DeleteFunc(activeToolNames, func(name string) bool { return allowInactiveTools[name] }) - activeToolNames = append(activeToolNames, chattool.FindToolsName) - activeToolNames = append(activeToolNames, deriveDeferredMCPActivations(promptRows, deferredCandidates)...) } toolNameToConfigID := make(map[string]uuid.UUID) diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index 336ead756cc..ca841655f10 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -48,6 +48,33 @@ func decideMCPToolSearch(input mcpToolSearchInput) mcpToolSearchDecision { return decision } +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 { diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index 382c1777d35..86bcd04a26a 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -2,14 +2,18 @@ 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" ) @@ -88,6 +92,8 @@ func TestDeriveDeferredMCPActivations(t *testing.T) { {Role: database.ChatMessageRoleTool, Content: malformed, ContentVersion: chatprompt.CurrentContentVersion}, } require.Equal(t, []string{"server__second", "server__first"}, deriveDeferredMCPActivations(rows, candidates)) + require.Equal(t, []string{"server__first"}, deriveDeferredMCPActivations(rows[1:], candidates), + "activations before a compaction summary are absent from the surviving prompt window") } func TestFlattenMCPParameterText(t *testing.T) { @@ -98,3 +104,144 @@ func TestFlattenMCPParameterText(t *testing.T) { 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 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}}) + 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) + 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)) +} + +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) + 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)) +} + +func TestMCPToolSearchBelowThresholdPreservesWireTools(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) + + decision := decideMCPToolSearch(mcpToolSearchInput{ + experimentEnabled: true, + contextWindow: 100_000, + candidates: []deferredMCPTool{candidate}, + }) + require.False(t, decision.apply) + 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 +} From 64eb04f128e16bc93ae8b86d136691c1bc51f9bf Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 17:10:43 +0000 Subject: [PATCH 03/48] feat(site/src/pages/AgentsPage): render find_tools calls in the conversation timeline --- coderd/x/chatd/ARCHITECTURE.md | 2 + .../ConversationTimeline.stories.tsx | 108 ++++++++++++++++++ .../ChatElements/tools/FindToolsTool.tsx | 54 +++++++++ .../ChatElements/tools/Tool.stories.tsx | 14 +++ .../components/ChatElements/tools/Tool.tsx | 93 +++++++++++++++ 5 files changed, 271 insertions(+) create mode 100644 site/src/pages/AgentsPage/components/ChatElements/tools/FindToolsTool.tsx diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index d9f1e16f8f6..b85fbc19a53 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -842,6 +842,8 @@ Tool calls have at least once semantics: if the goroutine executes a tool call, Parallel tool call results must be inserted in bulk after all parallel tool calls finish in a single `CommitStep` transition so that the generation goroutine only increments `history_version` once, since a change to the `history_version` interrupts the gorotuine. This is consistent with the existing chatd implementation. + + The generation goroutine supports: - chat compaction (automatic and manual, see [Manual compaction](#manual-compaction)) diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx index cb1b5d4d2d0..7d5f7f5cabd 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx @@ -571,6 +571,114 @@ 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"]), + }, + }, + ], + }, + { + ...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 -> 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 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 00000000000..dc1a8a5f82f --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/FindToolsTool.tsx @@ -0,0 +1,54 @@ +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[]; + matches: readonly FindToolsMatch[]; + status: ToolStatus; + isError: boolean; + errorMessage?: string; +}; + +export const FindToolsTool: FC = ({ + queries, + matches, + status, + isError, + errorMessage, +}) => { + const queryLabel = queries.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 3edecea02ca..ce3b7c1b23b 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 dc8f9e2757a..ff5256da420 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 parseStringList = (value: unknown): string[] | null => { + if (typeof value === "string") { + try { + return parseStringList(JSON.parse(value)); + } catch { + return null; + } + } + if (!Array.isArray(value)) { + return null; + } + const strings: string[] = []; + for (const item of value) { + if (typeof item !== "string") { + return null; + } + const trimmed = item.trim(); + if (trimmed) { + strings.push(trimmed); + } + } + return strings; +}; + +const parseFindToolsMatches = (value: unknown): FindToolsMatch[] | null => { + if (typeof value === "string") { + try { + return parseFindToolsMatches(JSON.parse(value)); + } catch { + return null; + } + } + if (!Array.isArray(value)) { + return null; + } + const matches: FindToolsMatch[] = []; + for (const item of value) { + const record = asRecord(item); + if ( + !record || + typeof record.name !== "string" || + typeof record.description !== "string" + ) { + return null; + } + matches.push({ + name: record.name, + description: record.description, + }); + } + return matches; +}; + +const FindToolsRenderer: FC = (props) => { + const parsedArgs = parseArgs(props.args); + const queries = parsedArgs + ? parsedArgs.queries === undefined + ? [] + : parseStringList(parsedArgs.queries) + : null; + const names = parsedArgs + ? parsedArgs.names === undefined + ? [] + : parseStringList(parsedArgs.names) + : null; + const searchTerms = queries && names ? [...queries, ...names] : null; + const parsedResult = parseArgs(props.result); + const matches = + props.status === "running" && props.result === undefined + ? [] + : parsedResult + ? parseFindToolsMatches(parsedResult.matches) + : null; + if (!searchTerms || !matches) { + return ; + } + + const errorMessage = parsedResult + ? asString(parsedResult.error || parsedResult.message) + : ""; + 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, From cc36157f57647690f702894a83f5f797ce16fe1b Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 17:23:58 +0000 Subject: [PATCH 04/48] refactor(coderd/x/chatd): apply cleanup gate findings to tool search --- coderd/x/chatd/chattool/findtools.go | 99 ++++++++++--------- .../chatd/chattool/findtools_internal_test.go | 1 + coderd/x/chatd/mcp_tool_search.go | 17 ++-- .../components/ChatElements/tools/Tool.tsx | 34 +++---- 4 files changed, 76 insertions(+), 75 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 72178f821a6..f4a2155790e 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -37,13 +37,11 @@ type FindToolsCall struct { TotalDeferred int } -// FindToolsOptions configures the find_tools tool. type FindToolsOptions struct { Entries []FindToolCatalogEntry OnCall func(context.Context, FindToolsCall) } -// FindToolsArgs are the arguments for the find_tools tool. type FindToolsArgs struct { Queries []string `json:"queries"` Names []string `json:"names"` @@ -86,13 +84,19 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { ) } -// SearchTools searches entries and returns the activation set. +// SearchTools scores entries against the query tokens, keeps the top +// matches, and always includes exact name activations. func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs) FindToolsResult { byName := make(map[string]FindToolCatalogEntry, len(entries)) for _, entry := range entries { byName[entry.Name] = entry } + var queryTokens []string + for _, query := range args.Queries { + queryTokens = append(queryTokens, tokenizeFindTools(query)...) + } + type scoredEntry struct { entry FindToolCatalogEntry score int @@ -100,10 +104,8 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs) FindToolsRe scored := make([]scoredEntry, 0, len(entries)) for _, entry := range entries { score := 0 - for _, query := range args.Queries { - for _, token := range tokenizeFindTools(query) { - score += scoreFindToolToken(entry, token) - } + for _, token := range queryTokens { + score += scoreFindToolToken(entry, token) } if score > 0 { scored = append(scored, scoredEntry{entry: entry, score: score}) @@ -168,28 +170,47 @@ func scoreFindToolToken(entry FindToolCatalogEntry, token string) int { func buildFindToolsDescription(entries []FindToolCatalogEntry) string { const usage = "Search deferred MCP tools by keyword, activate exact tool names, or scope queries with a server prefix. Calling a cataloged tool directly by name is allowed and auto-loads its schema, but search first for unfamiliar tools.\n\n" - catalog := detailedFindToolsCatalog(entries) + groups := groupFindToolsEntries(entries) + catalog := detailedFindToolsCatalog(groups) if estimatedFindToolsTokens(usage+catalog) > findToolsCatalogTokens { - catalog = namesOnlyFindToolsCatalog(entries) + catalog = namesOnlyFindToolsCatalog(groups) } if estimatedFindToolsTokens(usage+catalog) > findToolsCatalogTokens { - catalog = countsOnlyFindToolsCatalog(entries) + catalog = countsOnlyFindToolsCatalog(groups) } return usage + catalog } -func detailedFindToolsCatalog(entries []FindToolCatalogEntry) string { - return renderFindToolsCatalog(entries, func(entry FindToolCatalogEntry) string { - return "- " + entry.Name + " - " + truncateFindToolsSummary(entry.Description, 80) - }) +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(entries []FindToolCatalogEntry) string { - return renderFindToolsCatalog(entries, func(entry FindToolCatalogEntry) string { return entry.Name }) +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(entries []FindToolCatalogEntry) string { - groups := groupFindToolsEntries(entries) +func countsOnlyFindToolsCatalog(groups []findToolsGroup) string { var b strings.Builder for _, group := range groups { _, _ = b.WriteString("## ") @@ -201,6 +222,16 @@ func countsOnlyFindToolsCatalog(entries []FindToolCatalogEntry) string { 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 @@ -230,36 +261,10 @@ func groupFindToolsEntries(entries []FindToolCatalogEntry) []findToolsGroup { return groups } -func renderFindToolsCatalog(entries []FindToolCatalogEntry, renderEntry func(FindToolCatalogEntry) string) string { - groups := groupFindToolsEntries(entries) - var b strings.Builder - for _, group := range groups { - _, _ = b.WriteString("## ") - _, _ = b.WriteString(group.server) - if summary := truncateFindToolsSummary(group.description, 60); summary != "" { - _, _ = b.WriteString(" - ") - _, _ = b.WriteString(summary) - } - _ = b.WriteByte('\n') - if len(group.entries) > 0 && !strings.HasPrefix(renderEntry(group.entries[0]), "-") { - names := make([]string, 0, len(group.entries)) - for _, entry := range group.entries { - names = append(names, renderEntry(entry)) - } - _, _ = b.WriteString(strings.Join(names, " ")) - _ = b.WriteByte('\n') - continue - } - for _, entry := range group.entries { - _, _ = b.WriteString(renderEntry(entry)) - _ = b.WriteByte('\n') - } - } - return b.String() -} - func truncateFindToolsSummary(value string, maxRunes int) string { - value = strings.TrimSpace(strings.Split(strings.Split(value, "\n")[0], ". ")[0]) + line, _, _ := strings.Cut(value, "\n") + sentence, _, _ := strings.Cut(line, ". ") + value = strings.TrimSpace(sentence) if utf8.RuneCountInString(value) <= maxRunes { return value } diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 428fcb3cf76..ae0249d2baf 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -22,6 +22,7 @@ func TestSearchTools(t *testing.T) { t.Run("weights and tie break", func(t *testing.T) { t.Parallel() result := SearchTools(entries, FindToolsArgs{Queries: []string{"issue"}}) + 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) { diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index ca841655f10..bd37bcb9281 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -34,17 +34,17 @@ type mcpToolSearchInput struct { } func decideMCPToolSearch(input mcpToolSearchInput) mcpToolSearchDecision { - experimentEnabled, force, contextWindow, candidates := input.experimentEnabled, input.forceDefer, input.contextWindow, input.candidates - decision := mcpToolSearchDecision{estimatedTokens: estimateDeferredMCPToolTokens(candidates)} - if !experimentEnabled || len(candidates) == 0 { + decision := mcpToolSearchDecision{estimatedTokens: estimateDeferredMCPToolTokens(input.candidates)} + if !input.experimentEnabled || len(input.candidates) == 0 { return decision } - for _, candidate := range candidates { + for _, candidate := range input.candidates { if candidate.tool.Info().Name == chattool.FindToolsName { return decision } } - decision.apply = force || (contextWindow > 0 && decision.estimatedTokens > float64(contextWindow)/mcpToolSearchThresholdDivisor) + decision.apply = input.forceDefer || + (input.contextWindow > 0 && decision.estimatedTokens > float64(input.contextWindow)/mcpToolSearchThresholdDivisor) return decision } @@ -132,14 +132,11 @@ func flattenMCPParameterText(value any) string { } func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []deferredMCPTool) []string { - current := make(map[string]struct{}, len(candidates)) - for _, candidate := range candidates { - current[candidate.tool.Info().Name] = struct{}{} - } + current := deferredMCPToolNameSet(candidates) seen := make(map[string]struct{}, len(candidates)) activated := make([]string, 0, len(candidates)) appendName := func(name string) { - if _, ok := current[name]; !ok { + if !current[name] { return } if _, ok := seen[name]; ok { diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx index ff5256da420..1ee32098ef7 100644 --- a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx @@ -1041,25 +1041,23 @@ const parseFindToolsMatches = (value: unknown): FindToolsMatch[] | null => { const FindToolsRenderer: FC = (props) => { const parsedArgs = parseArgs(props.args); - const queries = parsedArgs - ? parsedArgs.queries === undefined - ? [] - : parseStringList(parsedArgs.queries) - : null; - const names = parsedArgs - ? parsedArgs.names === undefined - ? [] - : parseStringList(parsedArgs.names) - : null; - const searchTerms = queries && names ? [...queries, ...names] : null; + if (!parsedArgs) { + return ; + } + const queries = + parsedArgs.queries === undefined ? [] : parseStringList(parsedArgs.queries); + const names = + parsedArgs.names === undefined ? [] : parseStringList(parsedArgs.names); + if (!queries || !names) { + return ; + } + const searchTerms = [...queries, ...names]; const parsedResult = parseArgs(props.result); - const matches = - props.status === "running" && props.result === undefined - ? [] - : parsedResult - ? parseFindToolsMatches(parsedResult.matches) - : null; - if (!searchTerms || !matches) { + let matches: FindToolsMatch[] | null = []; + if (props.status !== "running" || props.result !== undefined) { + matches = parsedResult ? parseFindToolsMatches(parsedResult.matches) : null; + } + if (!matches) { return ; } From c08ef35bd87cd65d6e62447fffd4346e192fb090 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:15:34 +0000 Subject: [PATCH 05/48] fix(site/src/pages/AgentsPage): label exact-name activations separately in find_tools rendering --- .../ChatConversation/ConversationTimeline.stories.tsx | 3 ++- .../components/ChatElements/tools/FindToolsTool.tsx | 5 ++++- .../pages/AgentsPage/components/ChatElements/tools/Tool.tsx | 4 ++-- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx index 7d5f7f5cabd..7634817f4b9 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx @@ -586,6 +586,7 @@ export const FindToolsSearchResult: Story = { tool_name: "find_tools", args: { queries: JSON.stringify(["github issues", "pull requests"]), + names: JSON.stringify(["github__list_issues"]), }, }, ], @@ -624,7 +625,7 @@ export const FindToolsSearchResult: Story = { play: async ({ canvasElement }) => { const canvas = within(canvasElement); const summary = canvas.getByRole("button", { - name: "Searched tools: github issues, pull requests -> 2 matched", + name: "Searched tools: github issues, pull requests, name:github__list_issues -> 2 matched", }); expect(summary).toBeVisible(); expect(canvas.queryByText("github__list_issues")).not.toBeInTheDocument(); diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/FindToolsTool.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/FindToolsTool.tsx index dc1a8a5f82f..a0b86a993ff 100644 --- a/site/src/pages/AgentsPage/components/ChatElements/tools/FindToolsTool.tsx +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/FindToolsTool.tsx @@ -9,6 +9,7 @@ export type FindToolsMatch = { type FindToolsToolProps = { queries: readonly string[]; + names: readonly string[]; matches: readonly FindToolsMatch[]; status: ToolStatus; isError: boolean; @@ -17,12 +18,14 @@ type FindToolsToolProps = { export const FindToolsTool: FC = ({ queries, + names, matches, status, isError, errorMessage, }) => { - const queryLabel = queries.join(", ") || "tools"; + const queryLabel = + [...queries, ...names.map((name) => `name:${name}`)].join(", ") || "tools"; const label = status === "running" ? `Searching tools: ${queryLabel}` diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx index 1ee32098ef7..6a69475feb7 100644 --- a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx @@ -1051,7 +1051,6 @@ const FindToolsRenderer: FC = (props) => { if (!queries || !names) { return ; } - const searchTerms = [...queries, ...names]; const parsedResult = parseArgs(props.result); let matches: FindToolsMatch[] | null = []; if (props.status !== "running" || props.result !== undefined) { @@ -1066,7 +1065,8 @@ const FindToolsRenderer: FC = (props) => { : ""; return ( Date: Mon, 17 Aug 2026 19:26:18 +0000 Subject: [PATCH 06/48] fix(coderd/x/chatd/chattool): make find_tools queries and names optional in the schema --- coderd/x/chatd/chattool/findtools.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index f4a2155790e..8653022db11 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -43,8 +43,8 @@ type FindToolsOptions struct { } type FindToolsArgs struct { - Queries []string `json:"queries"` - Names []string `json:"names"` + Queries []string `json:"queries,omitempty"` + Names []string `json:"names,omitempty"` } type FindToolsMatch struct { From 9b8dc1fb44a9231f6ea5c39b87317dcd6aca845d Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:37:35 +0000 Subject: [PATCH 07/48] fix(coderd/x/chatd/chattool): bound find_tools results below generic tool truncation --- coderd/x/chatd/chattool/findtools.go | 44 +++++++++++-------- .../chatd/chattool/findtools_internal_test.go | 26 +++++++++++ 2 files changed, 51 insertions(+), 19 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 8653022db11..b12745023f3 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -84,8 +84,11 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { ) } -// SearchTools scores entries against the query tokens, keeps the top -// matches, and always includes exact name activations. +// 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. func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs) FindToolsResult { byName := make(map[string]FindToolCatalogEntry, len(entries)) for _, entry := range entries { @@ -117,26 +120,29 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs) FindToolsRe } return strings.Compare(a.entry.Name, b.entry.Name) }) - if len(scored) > findToolsMaxMatches { - scored = scored[:findToolsMaxMatches] - } - - matches := make([]FindToolsMatch, 0, len(scored)+len(args.Names)) - activatedSet := make(map[string]struct{}, len(scored)+len(args.Names)) - for _, item := range scored { - matches = append(matches, FindToolsMatch{Name: item.entry.Name, Description: item.entry.Description}) - activatedSet[item.entry.Name] = struct{}{} + matches := make([]FindToolsMatch, 0, findToolsMaxMatches) + activatedSet := make(map[string]struct{}, findToolsMaxMatches) + appendMatch := func(entry FindToolCatalogEntry) { + if _, exists := activatedSet[entry.Name]; exists { + return + } + if len(matches) >= findToolsMaxMatches { + return + } + matches = append(matches, FindToolsMatch{ + Name: entry.Name, + Description: truncateFindToolsSummary(entry.Description, 80), + }) + activatedSet[entry.Name] = struct{}{} } for _, name := range args.Names { - entry, ok := byName[name] - if !ok { - continue - } - if _, exists := activatedSet[name]; !exists { - matches = append(matches, FindToolsMatch{Name: entry.Name, Description: entry.Description}) - activatedSet[name] = struct{}{} + 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) @@ -169,7 +175,7 @@ func scoreFindToolToken(entry FindToolCatalogEntry, token string) int { } func buildFindToolsDescription(entries []FindToolCatalogEntry) string { - const usage = "Search deferred MCP tools by keyword, activate exact tool names, or scope queries with a server prefix. Calling a cataloged tool directly by name is allowed and auto-loads its schema, but search first for unfamiliar tools.\n\n" + const usage = "Search deferred MCP tools by keyword, activate exact tool names, or scope queries with a server 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" groups := groupFindToolsEntries(entries) catalog := detailedFindToolsCatalog(groups) if estimatedFindToolsTokens(usage+catalog) > findToolsCatalogTokens { diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index ae0249d2baf..2b5a22ce86e 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -52,6 +52,32 @@ func TestSearchTools(t *testing.T) { 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"}}) + 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}) + require.Len(t, capped.Matches, findToolsMaxMatches) + require.Len(t, capped.Activated, findToolsMaxMatches) + }) + 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"}}) + require.LessOrEqual(t, len([]rune(result.Matches[0].Description)), 80) + }) } func TestFindTools(t *testing.T) { From 3de6402554ce2907912bc43bcae3fce731dbb0a0 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:46:18 +0000 Subject: [PATCH 08/48] fix(coderd/x/chatd/chattool): tokenize find_tools search on unicode letters and digits --- coderd/x/chatd/chattool/findtools.go | 2 +- coderd/x/chatd/chattool/findtools_internal_test.go | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index b12745023f3..6fd75490b83 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -17,7 +17,7 @@ const ( findToolsCatalogTokens = 4000 ) -var findToolsTokenSeparator = regexp.MustCompile(`[^a-z0-9]+`) +var findToolsTokenSeparator = regexp.MustCompile(`[^\p{L}\p{N}]+`) // FindToolCatalogEntry is the searchable metadata for one deferred tool. type FindToolCatalogEntry struct { diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 2b5a22ce86e..a2ba3d9ba3a 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -69,6 +69,18 @@ func TestSearchTools(t *testing.T) { require.Len(t, capped.Matches, findToolsMaxMatches) require.Len(t, capped.Activated, findToolsMaxMatches) }) + 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{"検索"}}) + require.Equal(t, []string{"docs__検索"}, result.Activated) + + result = SearchTools(unicodeEntries, FindToolsArgs{Queries: []string{"Erstellen"}}) + require.Equal(t, []string{"docs__erstellen"}, result.Activated) + }) t.Run("result descriptions are summarized", func(t *testing.T) { t.Parallel() long := []FindToolCatalogEntry{{ From 697ee2ebb082e7799acbba9614dd7d37761d27b7 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:54:31 +0000 Subject: [PATCH 09/48] fix(coderd/x/chatd): fail open when a dynamic tool is named find_tools --- coderd/x/chatd/generation_preparer.go | 1 + coderd/x/chatd/mcp_tool_search.go | 7 +++++++ coderd/x/chatd/mcp_tool_search_internal_test.go | 16 ++++++++++------ 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index 69d1be97b2f..79ad5126e30 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -619,6 +619,7 @@ func (server *Server) prepareGeneration( forceDefer: server.forceMCPToolSearch, contextWindow: modelConfig.ContextLimit, candidates: deferredCandidates, + dynamicToolNames: dynamicToolNames, }) server.metrics.DeferredMCPToolTokens.WithLabelValues( model.Provider(), model.ModelID(), strconv.FormatBool(toolSearch.apply), diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index bd37bcb9281..b66303ccdc6 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -31,6 +31,7 @@ type mcpToolSearchInput struct { forceDefer bool contextWindow int64 candidates []deferredMCPTool + dynamicToolNames map[string]bool } func decideMCPToolSearch(input mcpToolSearchInput) mcpToolSearchDecision { @@ -38,6 +39,12 @@ func decideMCPToolSearch(input mcpToolSearchInput) mcpToolSearchDecision { if !input.experimentEnabled || len(input.candidates) == 0 { return decision } + // 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 decision + } for _, candidate := range input.candidates { if candidate.tool.Info().Name == chattool.FindToolsName { return decision diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index 86bcd04a26a..c0c7a41960d 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -41,12 +41,13 @@ func TestDecideMCPToolSearch(t *testing.T) { large := []deferredMCPTool{testDeferredTool("server__large", strings.Repeat("large ", 2000), map[string]any{"value": map[string]any{"type": "string"}})} tests := []struct { - name string - experiment bool - force bool - window int64 - candidates []deferredMCPTool - want bool + name string + experiment bool + force bool + window int64 + candidates []deferredMCPTool + dynamicNames map[string]bool + want bool }{ {name: "below", experiment: true, window: 100_000, candidates: small}, {name: "above", experiment: true, window: 10_000, candidates: large, want: true}, @@ -54,6 +55,8 @@ func TestDecideMCPToolSearch(t *testing.T) { {name: "experiment off", force: true, window: 10, candidates: large}, {name: "empty", experiment: true, force: true}, {name: "collision", experiment: true, force: true, candidates: []deferredMCPTool{testDeferredTool(chattool.FindToolsName, "collision", nil)}}, + {name: "dynamic collision", experiment: true, force: true, candidates: small, dynamicNames: map[string]bool{chattool.FindToolsName: true}}, + {name: "dynamic no collision", experiment: true, force: true, candidates: small, dynamicNames: map[string]bool{"other": true}, want: true}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -63,6 +66,7 @@ func TestDecideMCPToolSearch(t *testing.T) { forceDefer: tt.force, contextWindow: tt.window, candidates: tt.candidates, + dynamicToolNames: tt.dynamicNames, }).apply) }) } From 82644b3ebd671ea1dbdfa171bc31dc038fda6c4a Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:03:09 +0000 Subject: [PATCH 10/48] fix(coderd/x/chatd/chattool): score server metadata in find_tools search --- coderd/x/chatd/chattool/findtools.go | 7 +++++++ coderd/x/chatd/chattool/findtools_internal_test.go | 14 ++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 6fd75490b83..e8c66f9ee8b 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -171,6 +171,13 @@ func scoreFindToolToken(entry FindToolCatalogEntry, token string) int { if slices.Contains(tokenizeFindTools(entry.ParameterText), token) { score++ } + // 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. + if slices.Contains(tokenizeFindTools(entry.Server), token) || + slices.Contains(tokenizeFindTools(entry.ServerDescription), token) { + score++ + } return score } diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index a2ba3d9ba3a..c905a1f7a19 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -69,6 +69,20 @@ func TestSearchTools(t *testing.T) { require.Len(t, capped.Matches, findToolsMaxMatches) require.Len(t, capped.Activated, findToolsMaxMatches) }) + 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"}}) + require.Equal(t, []string{"tracker__create"}, result.Activated) + + result = SearchTools(serverEntries, FindToolsArgs{Queries: []string{"project"}}) + require.Equal(t, "docs__create", result.Matches[0].Name, + "tool description match outranks server metadata match") + require.Len(t, result.Matches, 2) + }) t.Run("unicode terms", func(t *testing.T) { t.Parallel() unicodeEntries := []FindToolCatalogEntry{ From a5fa358d3339d90bafd2b472d1693dd6721f97d2 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:12:08 +0000 Subject: [PATCH 11/48] fix(coderd/x/chatd/chattool): add a constant-size final fallback for the find_tools catalog --- coderd/x/chatd/chattool/findtools.go | 6 ++++++ coderd/x/chatd/chattool/findtools_internal_test.go | 12 ++++++++++++ 2 files changed, 18 insertions(+) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index e8c66f9ee8b..30443ec7307 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -2,6 +2,7 @@ package chattool import ( "context" + "fmt" "regexp" "slices" "strconv" @@ -191,6 +192,11 @@ func buildFindToolsDescription(entries []FindToolCatalogEntry) string { if estimatedFindToolsTokens(usage+catalog) > findToolsCatalogTokens { 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) > findToolsCatalogTokens { + catalog = fmt.Sprintf("%d deferred tools across %d servers.\n", len(entries), len(groups)) + } return usage + catalog } diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index c905a1f7a19..b4de4929e18 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -150,4 +150,16 @@ func TestBuildFindToolsDescription(t *testing.T) { degraded := buildFindToolsDescription(many) 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) + require.Contains(t, countsExceeded, "500 deferred tools across 500 servers.") + require.NotContains(t, countsExceeded, "## server_000") + require.LessOrEqual(t, estimatedFindToolsTokens(countsExceeded), float64(findToolsCatalogTokens)) } From 9ba563dbf7cc0d307f5c4849b6e264c579aa15ff Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:20:40 +0000 Subject: [PATCH 12/48] fix(coderd/x/chatd): apply the turn policy to workspace MCP deferred candidates --- coderd/x/chatd/generation_preparer.go | 27 ++++-------- coderd/x/chatd/mcp_tool_search.go | 43 ++++++++++++++++++ .../x/chatd/mcp_tool_search_internal_test.go | 44 +++++++++++++++++++ 3 files changed, 96 insertions(+), 18 deletions(-) diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index 79ad5126e30..fb9f38f6085 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -534,26 +534,17 @@ func (server *Server) prepareGeneration( for _, config := range mcpConnectConfigs { mcpConfigByID[config.ID] = config } - deferredCandidates := make([]deferredMCPTool, 0, len(mcpTools)+len(workspaceMCPTools)) - for _, tool := range mcpTools { - if !toolAllowedForTurn(tool, currentPlanMode, chat.ParentChatID, approvedPlanMCPConfigIDs) { - continue - } - candidate := deferredMCPTool{tool: tool} - if identified, ok := tool.(mcpclient.MCPToolIdentifier); ok { - if config, exists := mcpConfigByID[identified.MCPServerConfigID()]; exists { - candidate.server = config.Slug - candidate.serverDescription = config.Description - } - } - deferredCandidates = append(deferredCandidates, candidate) - } + deferredCandidates := collectDeferredMCPCandidates(deferredMCPCandidateInput{ + mcpTools: mcpTools, + workspaceMCPTools: workspaceMCPTools, + mcpConfigByID: mcpConfigByID, + planMode: currentPlanMode, + parentChatID: chat.ParentChatID, + approvedMCPConfigIDs: approvedPlanMCPConfigIDs, + includeWorkspaceTools: !isExploreSubagent, + }) tools = append(tools, mcpTools...) if !isExploreSubagent { - for _, tool := range workspaceMCPTools { - serverName, _, _ := strings.Cut(tool.Info().Name, "__") - deferredCandidates = append(deferredCandidates, deferredMCPTool{tool: tool, server: serverName}) - } tools = append(tools, workspaceMCPTools...) } tools = filterToolsForTurn(tools, currentPlanMode, chat.ParentChatID, approvedPlanMCPConfigIDs) diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index b66303ccdc6..150e329839c 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -6,10 +6,12 @@ import ( "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" ) @@ -21,6 +23,47 @@ type deferredMCPTool struct { 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 + } + for _, tool := range input.workspaceMCPTools { + if !toolAllowedForTurn(tool, input.planMode, input.parentChatID, input.approvedMCPConfigIDs) { + continue + } + serverName, _, _ := strings.Cut(tool.Info().Name, "__") + candidates = append(candidates, deferredMCPTool{tool: tool, server: serverName}) + } + return candidates +} + type mcpToolSearchDecision struct { apply bool estimatedTokens float64 diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index c0c7a41960d..04527e45d4f 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -116,6 +116,50 @@ type deferredExternalTestTool struct { 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))) +} + func TestConfigureDeferredMCPToolSearchGenerationFlows(t *testing.T) { t.Parallel() From c5d7a739c66065fa7fb0e8c8b3aeffd2ea50b156 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:31:04 +0000 Subject: [PATCH 13/48] fix(coderd/x/chatd): bound aggregate activated MCP schema weight by context budget --- coderd/x/chatd/generation_preparer.go | 3 ++- coderd/x/chatd/mcp_tool_search.go | 27 ++++++++++++++----- .../x/chatd/mcp_tool_search_internal_test.go | 17 +++++++----- 3 files changed, 34 insertions(+), 13 deletions(-) diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index fb9f38f6085..e9a733a234e 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -639,7 +639,8 @@ func (server *Server) prepareGeneration( activeToolNames, deferredCandidates, findTools, - deriveDeferredMCPActivations(promptRows, deferredCandidates), + deriveDeferredMCPActivations(promptRows, deferredCandidates, + float64(modelConfig.ContextLimit)/mcpToolSearchThresholdDivisor), ) builtinToolNames[chattool.FindToolsName] = true } diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index 150e329839c..4f51555268c 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -181,22 +181,37 @@ func flattenMCPParameterText(value any) string { return strings.Join(values, " ") } -func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []deferredMCPTool) []string { - current := deferredMCPToolNameSet(candidates) +// 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. Shed tools +// stay in the catalog and remain directly callable, which reactivates +// them as most recent. A tokenBudget <= 0 means unbounded. +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) { - if !current[name] { + candidate, ok := candidateByName[name] + if !ok { return } - if _, ok := seen[name]; ok { + if _, dup := seen[name]; dup { return } seen[name] = struct{}{} + weight := estimateDeferredMCPToolTokens([]deferredMCPTool{candidate}) + if tokenBudget > 0 && usedTokens+weight > tokenBudget { + return + } + usedTokens += weight activated = append(activated, name) } - for _, row := range rows { - parts, err := chatprompt.ParseContent(row) + for i := len(rows) - 1; i >= 0; i-- { + parts, err := chatprompt.ParseContent(rows[i]) if err != nil { continue } diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index 04527e45d4f..5356cc10f22 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -95,9 +95,14 @@ func TestDeriveDeferredMCPActivations(t *testing.T) { {Role: database.ChatMessageRoleAssistant, Content: directCall, ContentVersion: chatprompt.CurrentContentVersion}, {Role: database.ChatMessageRoleTool, Content: malformed, ContentVersion: chatprompt.CurrentContentVersion}, } - require.Equal(t, []string{"server__second", "server__first"}, deriveDeferredMCPActivations(rows, candidates)) - require.Equal(t, []string{"server__first"}, deriveDeferredMCPActivations(rows[1:], candidates), + 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") } func TestFlattenMCPParameterText(t *testing.T) { @@ -188,7 +193,7 @@ func TestConfigureDeferredMCPToolSearchGenerationFlows(t *testing.T) { history := []database.ChatMessage{{ Role: database.ChatMessageRoleTool, Content: resultContent, ContentVersion: chatprompt.CurrentContentVersion, }} - activations := deriveDeferredMCPActivations(history, candidates) + activations := deriveDeferredMCPActivations(history, candidates, 0) require.Equal(t, []string{second.tool.Info().Name}, activations) ordered, active, _ = configureDeferredMCPToolSearch(allTools, allActive, candidates, findTools, activations) @@ -198,7 +203,7 @@ func TestConfigureDeferredMCPToolSearchGenerationFlows(t *testing.T) { ) // 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)) + require.Equal(t, activations, deriveDeferredMCPActivations(history, candidates, 0)) } func TestConfigureDeferredMCPToolSearchDirectCallAndCompaction(t *testing.T) { @@ -213,7 +218,7 @@ func TestConfigureDeferredMCPToolSearchDirectCallAndCompaction(t *testing.T) { preSummary := []database.ChatMessage{{ Role: database.ChatMessageRoleAssistant, Content: directCall, ContentVersion: chatprompt.CurrentContentVersion, }} - activation := deriveDeferredMCPActivations(preSummary, candidates) + activation := deriveDeferredMCPActivations(preSummary, candidates, 0) require.Equal(t, []string{candidate.tool.Info().Name}, activation) findTools := chattool.FindTools(chattool.FindToolsOptions{Entries: deferredMCPToolEntries(candidates)}) @@ -231,7 +236,7 @@ func TestConfigureDeferredMCPToolSearchDirectCallAndCompaction(t *testing.T) { // 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)) + require.Empty(t, deriveDeferredMCPActivations(nil, candidates, 0)) } func TestMCPToolSearchBelowThresholdPreservesWireTools(t *testing.T) { From b35e464155fa5144592e2720fce1c32d3e3c93bf Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:39:58 +0000 Subject: [PATCH 14/48] fix(coderd/x/chatd): honor find_tools server prefixes and keep the newest activation --- coderd/x/chatd/chattool/findtools.go | 49 ++++++++++++++++--- .../chatd/chattool/findtools_internal_test.go | 18 +++++++ coderd/x/chatd/mcp_tool_search.go | 6 ++- .../x/chatd/mcp_tool_search_internal_test.go | 2 + 4 files changed, 66 insertions(+), 9 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 30443ec7307..73b78e8f18e 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -96,10 +96,7 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs) FindToolsRe byName[entry.Name] = entry } - var queryTokens []string - for _, query := range args.Queries { - queryTokens = append(queryTokens, tokenizeFindTools(query)...) - } + queries := parseFindToolsQueries(entries, args.Queries) type scoredEntry struct { entry FindToolCatalogEntry @@ -108,8 +105,17 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs) FindToolsRe scored := make([]scoredEntry, 0, len(entries)) for _, entry := range entries { score := 0 - for _, token := range queryTokens { - score += scoreFindToolToken(entry, token) + for _, query := range queries { + if query.server != "" && !strings.EqualFold(entry.Server, query.server) { + continue + } + if query.server != "" && len(query.tokens) == 0 { + score++ + continue + } + for _, token := range query.tokens { + score += scoreFindToolToken(entry, token) + } } if score > 0 { scored = append(scored, scoredEntry{entry: entry, score: score}) @@ -152,6 +158,35 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs) FindToolsRe return FindToolsResult{Matches: matches, Activated: activated, TotalDeferred: len(entries)} } +type scopedFindToolsQuery struct { + server string + 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. +func parseFindToolsQueries(entries []FindToolCatalogEntry, queries []string) []scopedFindToolsQuery { + servers := make(map[string]struct{}, len(entries)) + for _, entry := range entries { + if entry.Server != "" { + servers[strings.ToLower(entry.Server)] = struct{}{} + } + } + parsed := make([]scopedFindToolsQuery, 0, len(queries)) + for _, query := range queries { + if prefix, rest, ok := strings.Cut(query, ":"); ok { + server := strings.ToLower(strings.TrimSpace(prefix)) + if _, known := servers[server]; known { + parsed = append(parsed, scopedFindToolsQuery{server: server, tokens: tokenizeFindTools(rest)}) + continue + } + } + parsed = append(parsed, scopedFindToolsQuery{tokens: tokenizeFindTools(query)}) + } + return parsed +} + func tokenizeFindTools(value string) []string { parts := findToolsTokenSeparator.Split(strings.ToLower(value), -1) return slices.DeleteFunc(parts, func(part string) bool { return part == "" }) @@ -183,7 +218,7 @@ func scoreFindToolToken(entry FindToolCatalogEntry, token string) int { } func buildFindToolsDescription(entries []FindToolCatalogEntry) string { - const usage = "Search deferred MCP tools by keyword, activate exact tool names, or scope queries with a server 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" + 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" groups := groupFindToolsEntries(entries) catalog := detailedFindToolsCatalog(groups) if estimatedFindToolsTokens(usage+catalog) > findToolsCatalogTokens { diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index b4de4929e18..88bddf30581 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -83,6 +83,24 @@ func TestSearchTools(t *testing.T) { "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"}}) + 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:"}}) + 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"}}) + require.Len(t, result.Matches, 2, + "an unknown prefix is searched as plain keywords") + }) t.Run("unicode terms", func(t *testing.T) { t.Parallel() unicodeEntries := []FindToolCatalogEntry{ diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index 4f51555268c..ecab3fc5ead 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -183,7 +183,9 @@ func flattenMCPParameterText(value any) string { // 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. Shed tools +// 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. func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []deferredMCPTool, tokenBudget float64) []string { @@ -204,7 +206,7 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe } seen[name] = struct{}{} weight := estimateDeferredMCPToolTokens([]deferredMCPTool{candidate}) - if tokenBudget > 0 && usedTokens+weight > tokenBudget { + if len(activated) > 0 && tokenBudget > 0 && usedTokens+weight > tokenBudget { return } usedTokens += weight diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index 5356cc10f22..c390de22837 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -103,6 +103,8 @@ func TestDeriveDeferredMCPActivations(t *testing.T) { 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 TestFlattenMCPParameterText(t *testing.T) { From 007497c9602f65b0ec4c45de7a494fc6c343b606 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:48:31 +0000 Subject: [PATCH 15/48] fix(site/src/pages/AgentsPage): register a find_tools icon --- .../pages/AgentsPage/components/ChatElements/tools/ToolIcon.tsx | 2 ++ 1 file changed, 2 insertions(+) diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/ToolIcon.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/ToolIcon.tsx index 0d5cc6448b3..aa128484eff 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<{ From 1178b97d77c33af7302db6d6499452631be5cfc0 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:59:53 +0000 Subject: [PATCH 16/48] fix(coderd/x/chatd): budget find_tools results and derive workspace servers from routing names --- coderd/x/chatd/chattool/findtools.go | 22 +++++++-- .../chatd/chattool/findtools_internal_test.go | 45 ++++++++++++------- coderd/x/chatd/chattool/mcpworkspace.go | 11 +++++ coderd/x/chatd/generation_preparer.go | 7 +-- coderd/x/chatd/mcp_tool_search.go | 18 +++++++- .../x/chatd/mcp_tool_search_internal_test.go | 14 +++++- 6 files changed, 92 insertions(+), 25 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 73b78e8f18e..a1eaa0f67cf 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -27,6 +27,9 @@ type FindToolCatalogEntry struct { 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. @@ -40,7 +43,11 @@ type FindToolsCall struct { type FindToolsOptions struct { Entries []FindToolCatalogEntry - OnCall func(context.Context, FindToolsCall) + // SchemaTokenBudget caps the aggregate SchemaTokens one search may + // activate, so a result never reports activations that the + // activation budget would immediately shed. <= 0 means unbounded. + SchemaTokenBudget float64 + OnCall func(context.Context, FindToolsCall) } type FindToolsArgs struct { @@ -70,7 +77,7 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { if len(args.Queries) == 0 && len(args.Names) == 0 { return fantasy.NewTextErrorResponse("at least one query or name is required"), nil } - result := SearchTools(entries, args) + result := SearchTools(entries, args, options.SchemaTokenBudget) if options.OnCall != nil { options.OnCall(ctx, FindToolsCall{ Queries: args.Queries, @@ -89,8 +96,10 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { // 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. -func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs) FindToolsResult { +// activation JSON that later steps re-derive activations from. A +// positive schemaTokenBudget additionally stops admitting matches once +// their aggregate schema weight would exceed it, keeping at least one. +func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, schemaTokenBudget float64) FindToolsResult { byName := make(map[string]FindToolCatalogEntry, len(entries)) for _, entry := range entries { byName[entry.Name] = entry @@ -129,6 +138,7 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs) FindToolsRe }) matches := make([]FindToolsMatch, 0, findToolsMaxMatches) activatedSet := make(map[string]struct{}, findToolsMaxMatches) + usedSchemaTokens := 0.0 appendMatch := func(entry FindToolCatalogEntry) { if _, exists := activatedSet[entry.Name]; exists { return @@ -136,6 +146,10 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs) FindToolsRe if len(matches) >= findToolsMaxMatches { return } + if len(matches) > 0 && schemaTokenBudget > 0 && usedSchemaTokens+entry.SchemaTokens > schemaTokenBudget { + return + } + usedSchemaTokens += entry.SchemaTokens matches = append(matches, FindToolsMatch{ Name: entry.Name, Description: truncateFindToolsSummary(entry.Description, 80), diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 88bddf30581..e3944b7fdd9 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -21,24 +21,24 @@ func TestSearchTools(t *testing.T) { t.Run("weights and tie break", func(t *testing.T) { t.Parallel() - result := SearchTools(entries, FindToolsArgs{Queries: []string{"issue"}}) + result := SearchTools(entries, FindToolsArgs{Queries: []string{"issue"}}, 0) 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"}}) + result := SearchTools(entries, FindToolsArgs{Queries: []string{"channel"}}, 0) 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"}}) + result := SearchTools(entries, FindToolsArgs{Names: []string{"slack__post_message", "missing"}}, 0) 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{}) + result := SearchTools(entries, FindToolsArgs{}, 0) require.Empty(t, result.Matches) require.Empty(t, result.Activated) }) @@ -48,7 +48,7 @@ func TestSearchTools(t *testing.T) { for i := range many { many[i] = FindToolCatalogEntry{Name: fmt.Sprintf("server__tool_%02d", i), Description: "common"} } - result := SearchTools(many, FindToolsArgs{Queries: []string{"common"}}) + result := SearchTools(many, FindToolsArgs{Queries: []string{"common"}}, 0) require.Len(t, result.Matches, findToolsMaxMatches) require.Equal(t, "server__tool_00", result.Matches[0].Name) }) @@ -60,12 +60,12 @@ func TestSearchTools(t *testing.T) { 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"}}) + result := SearchTools(many, FindToolsArgs{Queries: []string{"common"}, Names: []string{"server__tool_24"}}, 0) 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}) + capped := SearchTools(many, FindToolsArgs{Names: names}, 0) require.Len(t, capped.Matches, findToolsMaxMatches) require.Len(t, capped.Activated, findToolsMaxMatches) }) @@ -75,10 +75,10 @@ func TestSearchTools(t *testing.T) { {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"}}) + result := SearchTools(serverEntries, FindToolsArgs{Queries: []string{"tracking"}}, 0) require.Equal(t, []string{"tracker__create"}, result.Activated) - result = SearchTools(serverEntries, FindToolsArgs{Queries: []string{"project"}}) + result = SearchTools(serverEntries, FindToolsArgs{Queries: []string{"project"}}, 0) require.Equal(t, "docs__create", result.Matches[0].Name, "tool description match outranks server metadata match") require.Len(t, result.Matches, 2) @@ -89,15 +89,15 @@ func TestSearchTools(t *testing.T) { {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"}}) + result := SearchTools(scopedEntries, FindToolsArgs{Queries: []string{"github: status"}}, 0) 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:"}}) + result = SearchTools(scopedEntries, FindToolsArgs{Queries: []string{"github:"}}, 0) 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"}}) + result = SearchTools(scopedEntries, FindToolsArgs{Queries: []string{"error: status"}}, 0) require.Len(t, result.Matches, 2, "an unknown prefix is searched as plain keywords") }) @@ -107,19 +107,34 @@ func TestSearchTools(t *testing.T) { {Name: "docs__検索", Description: "ドキュメント検索"}, {Name: "docs__erstellen", Description: "Dokument ERSTELLEN"}, } - result := SearchTools(unicodeEntries, FindToolsArgs{Queries: []string{"検索"}}) + result := SearchTools(unicodeEntries, FindToolsArgs{Queries: []string{"検索"}}, 0) require.Equal(t, []string{"docs__検索"}, result.Activated) - result = SearchTools(unicodeEntries, FindToolsArgs{Queries: []string{"Erstellen"}}) + result = SearchTools(unicodeEntries, FindToolsArgs{Queries: []string{"Erstellen"}}, 0) 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"}}, 100) + 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"}}, 100) + 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"}}) + result := SearchTools(long, FindToolsArgs{Names: []string{"server__verbose"}}, 0) require.LessOrEqual(t, len([]rune(result.Matches[0].Description)), 80) }) } diff --git a/coderd/x/chatd/chattool/mcpworkspace.go b/coderd/x/chatd/chattool/mcpworkspace.go index 8f851268826..7f9e90524f6 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/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index e9a733a234e..a083b9292ae 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -616,8 +616,10 @@ func (server *Server) prepareGeneration( model.Provider(), model.ModelID(), strconv.FormatBool(toolSearch.apply), ).Observe(toolSearch.estimatedTokens) if toolSearch.apply { + activationTokenBudget := float64(modelConfig.ContextLimit) / mcpToolSearchThresholdDivisor findTools := chattool.FindTools(chattool.FindToolsOptions{ - Entries: deferredMCPToolEntries(deferredCandidates), + Entries: deferredMCPToolEntries(deferredCandidates), + SchemaTokenBudget: activationTokenBudget, OnCall: func(callCtx context.Context, call chattool.FindToolsCall) { server.metrics.FindToolsCallsTotal.Inc() server.metrics.FindToolsMatchCount.Observe(float64(call.MatchCount)) @@ -639,8 +641,7 @@ func (server *Server) prepareGeneration( activeToolNames, deferredCandidates, findTools, - deriveDeferredMCPActivations(promptRows, deferredCandidates, - float64(modelConfig.ContextLimit)/mcpToolSearchThresholdDivisor), + deriveDeferredMCPActivations(promptRows, deferredCandidates, activationTokenBudget), ) builtinToolNames[chattool.FindToolsName] = true } diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index ecab3fc5ead..235f342865d 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -58,12 +58,25 @@ func collectDeferredMCPCandidates(input deferredMCPCandidateInput) []deferredMCP if !toolAllowedForTurn(tool, input.planMode, input.parentChatID, input.approvedMCPConfigIDs) { continue } - serverName, _, _ := strings.Cut(tool.Info().Name, "__") - candidates = append(candidates, deferredMCPTool{tool: tool, server: serverName}) + candidates = append(candidates, deferredMCPTool{tool: tool, server: workspaceMCPServerName(tool)}) } return candidates } +// 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 mcpToolSearchDecision struct { apply bool estimatedTokens float64 @@ -149,6 +162,7 @@ func deferredMCPToolEntries(candidates []deferredMCPTool) []chattool.FindToolCat Server: candidate.server, ServerDescription: candidate.serverDescription, ParameterText: flattenMCPParameterText(info.Parameters), + SchemaTokens: estimateDeferredMCPToolTokens([]deferredMCPTool{candidate}), }) } return entries diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index c390de22837..8918f72f3b1 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -16,6 +16,7 @@ import ( "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 { @@ -165,6 +166,17 @@ func TestCollectDeferredMCPCandidates(t *testing.T) { 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") } func TestConfigureDeferredMCPToolSearchGenerationFlows(t *testing.T) { @@ -185,7 +197,7 @@ func TestConfigureDeferredMCPToolSearchGenerationFlows(t *testing.T) { second.tool.Info().Name: true, }, allowInactive) - result := chattool.SearchTools(deferredMCPToolEntries(candidates), chattool.FindToolsArgs{Names: []string{second.tool.Info().Name}}) + result := chattool.SearchTools(deferredMCPToolEntries(candidates), chattool.FindToolsArgs{Names: []string{second.tool.Info().Name}}, 0) resultJSON, err := json.Marshal(result) require.NoError(t, err) resultContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ From 198d8991d7d3ca49da61e905ce24813ac4ac4651 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 21:11:09 +0000 Subject: [PATCH 17/48] fix(coderd/x/chatd): scale the find_tools catalog to the context window and share the activation budget across calls --- coderd/x/chatd/chattool/findtools.go | 51 +++++++++++++++---- .../chatd/chattool/findtools_internal_test.go | 36 +++++++++++-- coderd/x/chatd/generation_preparer.go | 5 +- 3 files changed, 77 insertions(+), 15 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index a1eaa0f67cf..862d83d6252 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -3,10 +3,12 @@ package chattool import ( "context" "fmt" + "math" "regexp" "slices" "strconv" "strings" + "sync" "unicode/utf8" "charm.land/fantasy" @@ -43,11 +45,17 @@ type FindToolsCall struct { type FindToolsOptions struct { Entries []FindToolCatalogEntry - // SchemaTokenBudget caps the aggregate SchemaTokens one search may - // activate, so a result never reports activations that the - // activation budget would immediately shed. <= 0 means unbounded. + // 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 - OnCall func(context.Context, FindToolsCall) + // 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 { @@ -70,14 +78,33 @@ type FindToolsResult struct { // 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 return fantasy.NewAgentTool( FindToolsName, - buildFindToolsDescription(entries), + buildFindToolsDescription(entries, options.CatalogTokenBudget), func(ctx context.Context, args FindToolsArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { if len(args.Queries) == 0 && len(args.Names) == 0 { return fantasy.NewTextErrorResponse("at least one query or name is required"), nil } - result := SearchTools(entries, args, options.SchemaTokenBudget) + budgetMu.Lock() + effectiveBudget := remainingBudget + if options.SchemaTokenBudget > 0 && effectiveBudget <= 0 { + // Keep SearchTools's first-match guarantee without + // flipping an exhausted budget into "unbounded". + effectiveBudget = math.SmallestNonzeroFloat64 + } + result := SearchTools(entries, args, effectiveBudget) + if options.SchemaTokenBudget > 0 { + for _, name := range result.Activated { + remainingBudget -= schemaTokensByName[name] + } + } + budgetMu.Unlock() if options.OnCall != nil { options.OnCall(ctx, FindToolsCall{ Queries: args.Queries, @@ -231,19 +258,23 @@ func scoreFindToolToken(entry FindToolCatalogEntry, token string) int { return score } -func buildFindToolsDescription(entries []FindToolCatalogEntry) string { +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) > findToolsCatalogTokens { + if estimatedFindToolsTokens(usage+catalog) > budget { catalog = namesOnlyFindToolsCatalog(groups) } - if estimatedFindToolsTokens(usage+catalog) > findToolsCatalogTokens { + 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) > findToolsCatalogTokens { + if estimatedFindToolsTokens(usage+catalog) > budget { catalog = fmt.Sprintf("%d deferred tools across %d servers.\n", len(entries), len(groups)) } return usage + catalog diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index e3944b7fdd9..7e2fc0ba221 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -158,6 +158,32 @@ func TestFindTools(t *testing.T) { require.True(t, resp.IsError) } +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") + require.Equal(t, []string{"server__d"}, activated(`{"names":["server__d"]}`), + "an exhausted budget still keeps the first match") +} + func TestBuildFindToolsDescription(t *testing.T) { t.Parallel() entries := []FindToolCatalogEntry{ @@ -165,7 +191,7 @@ func TestBuildFindToolsDescription(t *testing.T) { {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) + 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") @@ -180,7 +206,7 @@ func TestBuildFindToolsDescription(t *testing.T) { Server: "server", } } - degraded := buildFindToolsDescription(many) + degraded := buildFindToolsDescription(many, 0) require.Contains(t, degraded, "## server (300 tools)") require.NotContains(t, degraded, "server__tool_000") @@ -191,8 +217,12 @@ func TestBuildFindToolsDescription(t *testing.T) { Server: fmt.Sprintf("server_%03d_%s", i, strings.Repeat("s", 40)), } } - countsExceeded := buildFindToolsDescription(manyServers) + 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/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index a083b9292ae..eada5f55071 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -618,8 +618,9 @@ func (server *Server) prepareGeneration( if toolSearch.apply { activationTokenBudget := float64(modelConfig.ContextLimit) / mcpToolSearchThresholdDivisor findTools := chattool.FindTools(chattool.FindToolsOptions{ - Entries: deferredMCPToolEntries(deferredCandidates), - SchemaTokenBudget: activationTokenBudget, + Entries: deferredMCPToolEntries(deferredCandidates), + SchemaTokenBudget: activationTokenBudget, + CatalogTokenBudget: activationTokenBudget, OnCall: func(callCtx context.Context, call chattool.FindToolsCall) { server.metrics.FindToolsCallsTotal.Inc() server.metrics.FindToolsMatchCount.Observe(float64(call.MatchCount)) From 6a51807adcfe012f622dad44e0db235113a11dfa Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Mon, 17 Aug 2026 21:22:03 +0000 Subject: [PATCH 18/48] fix(coderd/x/chatd/chattool): reject find_tools over-claims once the shared budget is touched --- coderd/x/chatd/chattool/findtools.go | 26 +++++++++++++------ .../chatd/chattool/findtools_internal_test.go | 21 +++++++++++++-- 2 files changed, 37 insertions(+), 10 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 862d83d6252..5343d65d277 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -3,7 +3,6 @@ package chattool import ( "context" "fmt" - "math" "regexp" "slices" "strconv" @@ -22,6 +21,8 @@ const ( 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 @@ -92,17 +93,26 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { return fantasy.NewTextErrorResponse("at least one query or name is required"), nil } budgetMu.Lock() - effectiveBudget := remainingBudget - if options.SchemaTokenBudget > 0 && effectiveBudget <= 0 { - // Keep SearchTools's first-match guarantee without - // flipping an exhausted budget into "unbounded". - effectiveBudget = math.SmallestNonzeroFloat64 + if options.SchemaTokenBudget > 0 && remainingBudget <= 0 { + budgetMu.Unlock() + return fantasy.NewTextErrorResponse(findToolsBudgetExhausted), nil } - result := SearchTools(entries, args, effectiveBudget) + result := SearchTools(entries, args, remainingBudget) if options.SchemaTokenBudget > 0 { + admitted := 0.0 for _, name := range result.Activated { - remainingBudget -= schemaTokensByName[name] + admitted += schemaTokensByName[name] + } + // Derivation retains a single over-budget claim via its + // newest-keep rule, but only when it is the turn's sole + // claim, which is exactly when the budget is untouched. + // Any other over-claim would be silently shed on the + // next request, so fail it loudly instead. + if admitted > remainingBudget && remainingBudget < options.SchemaTokenBudget { + budgetMu.Unlock() + return fantasy.NewTextErrorResponse(findToolsBudgetExhausted), nil } + remainingBudget -= admitted } budgetMu.Unlock() if options.OnCall != nil { diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 7e2fc0ba221..d7cad90cdbc 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -180,8 +180,25 @@ func TestFindToolsSharedSchemaBudget(t *testing.T) { 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") - require.Equal(t, []string{"server__d"}, activated(`{"names":["server__d"]}`), - "an exhausted budget still keeps the first match") + + 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}}, + 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.True(t, resp.IsError, "the spent budget rejects further activations") } func TestBuildFindToolsDescription(t *testing.T) { From 14fe9e2ec2c785ead427cf766db3b3d0671ffcec Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 04:24:30 +0000 Subject: [PATCH 19/48] fix(coderd/x/chatd): admit direct calls before same-step search activations and serialize find_tools budget claims --- coderd/x/chatd/chatloop/chatloop.go | 101 ++++++++++++------ .../chatloop/chatloop_run_internal_test.go | 78 ++++++++++++++ coderd/x/chatd/chattool/findtools.go | 11 +- .../chatd/chattool/findtools_internal_test.go | 7 ++ coderd/x/chatd/mcp_tool_search.go | 34 +++++- .../x/chatd/mcp_tool_search_internal_test.go | 26 +++++ 6 files changed, 220 insertions(+), 37 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index cd3fd8ed4f4..0ed9edc599c 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -1114,42 +1114,63 @@ func executeTools( 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 on one + // goroutine in tool-call order, so order-sensitive shared state + // (for example the find_tools activation budget) is claimed + // deterministically. 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, - allowInactiveTools, - providerRunnerNames, - resultProviderMetadata, - maxResultBytes, - toolNameAliases, - ) + runCall(i, tc) + }() + } + if len(serialIndexes) > 0 { + wg.Add(1) + go func() { + defer wg.Done() + for _, i := range serialIndexes { + runCall(i, localToolCalls[i]) + } }() } wg.Wait() @@ -1460,6 +1481,22 @@ 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 } + +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 89ae523aa5e..67520d97ec4 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -5,6 +5,7 @@ import ( "encoding/base64" "errors" "iter" + "runtime" "sync" "sync/atomic" "testing" @@ -872,6 +873,83 @@ func TestSanitizeAnthropicProviderToolContent(t *testing.T) { } } +type serialMarkerTool struct{ fantasy.AgentTool } + +func (serialMarkerTool) SerialToolCalls() bool { return true } + +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, + 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() diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 5343d65d277..dc8f1e1c0ba 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -76,6 +76,13 @@ type FindToolsResult struct { TotalDeferred int `json:"total_deferred"` } +// serialCallsTool 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. +type serialCallsTool struct{ fantasy.AgentTool } + +func (serialCallsTool) SerialToolCalls() bool { return true } + // FindTools returns the built-in used to discover deferred MCP tool schemas. func FindTools(options FindToolsOptions) fantasy.AgentTool { entries := slices.Clone(options.Entries) @@ -85,7 +92,7 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { } var budgetMu sync.Mutex remainingBudget := options.SchemaTokenBudget - return fantasy.NewAgentTool( + return serialCallsTool{AgentTool: fantasy.NewAgentTool( FindToolsName, buildFindToolsDescription(entries, options.CatalogTokenBudget), func(ctx context.Context, args FindToolsArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { @@ -126,7 +133,7 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { } return marshalToolResponse(result), nil }, - ) + )} } // SearchTools includes exact name activations first, then fills the diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index d7cad90cdbc..1da6d2d3850 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -158,6 +158,13 @@ func TestFindTools(t *testing.T) { 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 TestFindToolsSharedSchemaBudget(t *testing.T) { t.Parallel() tool := FindTools(FindToolsOptions{ diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index 235f342865d..2e126eedc15 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -202,6 +202,11 @@ func flattenMCPParameterText(value any) string { // 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. func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []deferredMCPTool, tokenBudget float64) []string { candidateByName := make(map[string]deferredMCPTool, len(candidates)) for _, candidate := range candidates { @@ -226,23 +231,46 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe usedTokens += weight activated = append(activated, name) } - for i := len(rows) - 1; i >= 0; i-- { + parsedParts := make([][]codersdk.ChatMessagePart, len(rows)) + findToolsCallIDs := make(map[string]struct{}) + for i := range rows { parts, err := chatprompt.ParseContent(rows[i]) if err != nil { continue } + parsedParts[i] = parts for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName == chattool.FindToolsName && part.ToolCallID != "" { + findToolsCallIDs[part.ToolCallID] = struct{}{} + } + } + } + pendingSearch := make(map[string][]string) + for i := len(rows) - 1; i >= 0; i-- { + for _, part := range parsedParts[i] { + if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName != chattool.FindToolsName { + appendName(part.ToolName) + } + } + for _, 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 := findToolsCallIDs[part.ToolCallID]; paired { + pendingSearch[part.ToolCallID] = result.Activated + continue + } for _, name := range result.Activated { appendName(name) } - case part.Type == codersdk.ChatMessagePartTypeToolCall: - appendName(part.ToolName) + case part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName == chattool.FindToolsName: + for _, name := range pendingSearch[part.ToolCallID] { + appendName(name) + } + delete(pendingSearch, part.ToolCallID) } } } diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index 8918f72f3b1..8d7b2577b69 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -108,6 +108,32 @@ func TestDeriveDeferredMCPActivations(t *testing.T) { "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 TestFlattenMCPParameterText(t *testing.T) { t.Parallel() text := flattenMCPParameterText(map[string]any{ From f59c158357e1cc4347fe2fc8c610ad29de7ce012 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 04:49:42 +0000 Subject: [PATCH 20/48] fix(coderd/x/chatd): reserve same-step direct-call schema weight before find_tools admits activations --- coderd/x/chatd/chatloop/chatloop.go | 34 +++++++++ .../chatloop/chatloop_run_internal_test.go | 75 +++++++++++++++++++ coderd/x/chatd/chattool/findtools.go | 72 +++++++++++++++--- .../chatd/chattool/findtools_internal_test.go | 49 ++++++++++++ 4 files changed, 219 insertions(+), 11 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 0ed9edc599c..7e5b02dc965 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -1112,6 +1112,8 @@ func executeTools( } } + notifyStepToolCallObservers(toolMap, toolNameAliases, localToolCalls) + results := make([]fantasy.ToolResultContent, len(localToolCalls)) completedAt := make([]time.Time, len(localToolCalls)) runCall := func(i int, tc fantasy.ToolCallContent) { @@ -1485,6 +1487,38 @@ func isToolActive(name string, activeTools []string) bool { // 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) + } + } +} + func isSerialToolCall(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, name string) bool { if alias, ok := toolNameAliases[name]; ok { name = alias diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 67520d97ec4..783e5a8bb83 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -877,6 +877,81 @@ 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: "{}"}, + }, + 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"}, observedNames, + "a called observer sees every resolved tool-call name in the step") + 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") +} + func TestExecuteToolsSerialToolCallOrder(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index dc8f1e1c0ba..fb2f3338cef 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -17,6 +17,10 @@ const ( FindToolsName = "find_tools" findToolsMaxMatches = 20 findToolsCatalogTokens = 4000 + // 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}]+`) @@ -76,12 +80,19 @@ type FindToolsResult struct { TotalDeferred int `json:"total_deferred"` } -// serialCallsTool opts find_tools into in-order execution when one step +// 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. -type serialCallsTool struct{ fantasy.AgentTool } +// 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) +} + +func (findToolsTool) SerialToolCalls() bool { return true } -func (serialCallsTool) SerialToolCalls() bool { return true } +func (t findToolsTool) ObserveStepToolCalls(names []string) { t.reserveStepCalls(names) } // FindTools returns the built-in used to discover deferred MCP tool schemas. func FindTools(options FindToolsOptions) fantasy.AgentTool { @@ -92,7 +103,30 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { } var budgetMu sync.Mutex remainingBudget := options.SchemaTokenBudget - return serialCallsTool{AgentTool: fantasy.NewAgentTool( + // Direct calls to deferred tools in the same step are admitted by + // derivation before any search activations, so their schema weight + // is reserved out of the budget before searches run. Reserved names + // stay free to activate because derivation already retains them. + reserved := make(map[string]struct{}) + reserve := func(names []string) { + if options.SchemaTokenBudget <= 0 { + return + } + budgetMu.Lock() + defer budgetMu.Unlock() + for _, name := range names { + weight, ok := schemaTokensByName[name] + if !ok { + continue + } + if _, dup := reserved[name]; dup { + continue + } + reserved[name] = struct{}{} + remainingBudget -= weight + } + } + return findToolsTool{reserveStepCalls: reserve, AgentTool: fantasy.NewAgentTool( FindToolsName, buildFindToolsDescription(entries, options.CatalogTokenBudget), func(ctx context.Context, args FindToolsArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { @@ -100,22 +134,38 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { return fantasy.NewTextErrorResponse("at least one query or name is required"), nil } budgetMu.Lock() - if options.SchemaTokenBudget > 0 && remainingBudget <= 0 { - budgetMu.Unlock() - return fantasy.NewTextErrorResponse(findToolsBudgetExhausted), nil + searchEntries := entries + if len(reserved) > 0 { + searchEntries = slices.Clone(entries) + for i := range searchEntries { + if _, ok := reserved[searchEntries[i].Name]; ok { + searchEntries[i].SchemaTokens = 0 + } + } + } + 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 } - result := SearchTools(entries, args, remainingBudget) + result := SearchTools(searchEntries, args, searchBudget) if options.SchemaTokenBudget > 0 { admitted := 0.0 for _, name := range result.Activated { + if _, ok := reserved[name]; ok { + continue + } admitted += schemaTokensByName[name] } // Derivation retains a single over-budget claim via its // newest-keep rule, but only when it is the turn's sole // claim, which is exactly when the budget is untouched. // Any other over-claim would be silently shed on the - // next request, so fail it loudly instead. - if admitted > remainingBudget && remainingBudget < options.SchemaTokenBudget { + // next request, so fail it loudly instead. A zero new + // claim always succeeds: reserved names are already + // retained by derivation, whatever the budget holds. + if admitted > 0 && admitted > remainingBudget && remainingBudget < options.SchemaTokenBudget { budgetMu.Unlock() return fantasy.NewTextErrorResponse(findToolsBudgetExhausted), nil } diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 1da6d2d3850..e140cbe043c 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -165,6 +165,55 @@ func TestFindToolsSerialToolCalls(t *testing.T) { 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("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{ From fee1f79d641f4d3a30f02508212867fc02141351 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:00:19 +0000 Subject: [PATCH 21/48] fix(coderd/x/chatd): count budget-rejected find_tools calls and keep raw queries out of standard logs --- coderd/x/chatd/chattool/findtools.go | 12 ++++++++++ .../chatd/chattool/findtools_internal_test.go | 24 +++++++++++++++++++ coderd/x/chatd/generation_preparer.go | 21 ++++++++++------ 3 files changed, 50 insertions(+), 7 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index fb2f3338cef..642f4ab2bf6 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -46,6 +46,10 @@ type FindToolsCall struct { MatchCount int Activated []string TotalDeferred int + // BudgetRejected marks a call that returned the budget-exhausted + // error instead of results, so callers can count it without + // polluting match or activation statistics. + BudgetRejected bool } type FindToolsOptions struct { @@ -167,6 +171,14 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { // retained by derivation, whatever the budget holds. if admitted > 0 && admitted > remainingBudget && remainingBudget < options.SchemaTokenBudget { budgetMu.Unlock() + if options.OnCall != nil { + options.OnCall(ctx, FindToolsCall{ + Queries: args.Queries, + Names: args.Names, + TotalDeferred: len(entries), + BudgetRejected: true, + }) + } return fantasy.NewTextErrorResponse(findToolsBudgetExhausted), nil } remainingBudget -= admitted diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index e140cbe043c..bc585a7d73e 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -197,6 +197,30 @@ func TestFindToolsDirectCallReservation(t *testing.T) { require.Equal(t, []string{"server__c"}, result.Activated) }) + t.Run("budget-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) + require.Len(t, calls, 2, "rejected calls count toward call totals") + require.False(t, calls[0].BudgetRejected) + require.True(t, calls[1].BudgetRejected) + require.Equal(t, []string{"server__b"}, calls[1].Names) + require.Empty(t, calls[1].Activated, "a rejected call reports no activations") + }) + t.Run("reserved names stay activatable after the budget is spent", func(t *testing.T) { t.Parallel() tool, observer := newTool(50) diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index eada5f55071..2a7533804a3 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -623,17 +623,24 @@ func (server *Server) prepareGeneration( CatalogTokenBudget: activationTokenBudget, OnCall: func(callCtx context.Context, call chattool.FindToolsCall) { server.metrics.FindToolsCallsTotal.Inc() - server.metrics.FindToolsMatchCount.Observe(float64(call.MatchCount)) - server.metrics.FindToolsActivationsTotal.Add(float64(len(call.Activated))) - if call.MatchCount == 0 { - server.metrics.FindToolsEmptyTotal.Inc() + if !call.BudgetRejected { + 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("queries", call.Queries), - slog.F("names", call.Names), + slog.F("query_count", len(call.Queries)), + slog.F("name_count", len(call.Names)), slog.F("match_count", call.MatchCount), - slog.F("activated", call.Activated), + slog.F("activated_count", len(call.Activated)), slog.F("total_deferred", call.TotalDeferred), + slog.F("budget_rejected", call.BudgetRejected), ) }, }) From c57502b46a3a421b750cb0b9f4272e515a865249 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:09:13 +0000 Subject: [PATCH 22/48] fix(coderd/x/chatd/chattool): count empty-argument find_tools rejections --- coderd/x/chatd/chattool/findtools.go | 25 +++++++++++++------ .../chatd/chattool/findtools_internal_test.go | 13 +++++++--- coderd/x/chatd/generation_preparer.go | 4 +-- 3 files changed, 29 insertions(+), 13 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 642f4ab2bf6..b1268a1f30e 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -46,12 +46,17 @@ type FindToolsCall struct { MatchCount int Activated []string TotalDeferred int - // BudgetRejected marks a call that returned the budget-exhausted - // error instead of results, so callers can count it without + // Rejection is empty for successful searches. Rejected calls carry + // "budget" or "arguments" so callers can count them without // polluting match or activation statistics. - BudgetRejected bool + Rejection string } +const ( + findToolsRejectionBudget = "budget" + findToolsRejectionArguments = "arguments" +) + type FindToolsOptions struct { Entries []FindToolCatalogEntry // SchemaTokenBudget caps the aggregate SchemaTokens all searches on @@ -135,6 +140,12 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { 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() @@ -173,10 +184,10 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { budgetMu.Unlock() if options.OnCall != nil { options.OnCall(ctx, FindToolsCall{ - Queries: args.Queries, - Names: args.Names, - TotalDeferred: len(entries), - BudgetRejected: true, + Queries: args.Queries, + Names: args.Names, + TotalDeferred: len(entries), + Rejection: findToolsRejectionBudget, }) } return fantasy.NewTextErrorResponse(findToolsBudgetExhausted), nil diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index bc585a7d73e..55ca860522d 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -197,7 +197,7 @@ func TestFindToolsDirectCallReservation(t *testing.T) { require.Equal(t, []string{"server__c"}, result.Activated) }) - t.Run("budget-rejected calls still reach OnCall", func(t *testing.T) { + t.Run("rejected calls still reach OnCall", func(t *testing.T) { t.Parallel() var calls []FindToolsCall tool := FindTools(FindToolsOptions{ @@ -214,11 +214,16 @@ func TestFindToolsDirectCallReservation(t *testing.T) { resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{"names":["server__b"]}`}) require.NoError(t, err) require.True(t, resp.IsError) - require.Len(t, calls, 2, "rejected calls count toward call totals") - require.False(t, calls[0].BudgetRejected) - require.True(t, calls[1].BudgetRejected) + resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{}`}) + require.NoError(t, err) + require.True(t, resp.IsError) + require.Len(t, calls, 3, "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) }) t.Run("reserved names stay activatable after the budget is spent", func(t *testing.T) { diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index 2a7533804a3..45221f483f0 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -623,7 +623,7 @@ func (server *Server) prepareGeneration( CatalogTokenBudget: activationTokenBudget, OnCall: func(callCtx context.Context, call chattool.FindToolsCall) { server.metrics.FindToolsCallsTotal.Inc() - if !call.BudgetRejected { + if call.Rejection == "" { server.metrics.FindToolsMatchCount.Observe(float64(call.MatchCount)) server.metrics.FindToolsActivationsTotal.Add(float64(len(call.Activated))) if call.MatchCount == 0 { @@ -640,7 +640,7 @@ func (server *Server) prepareGeneration( slog.F("match_count", call.MatchCount), slog.F("activated_count", len(call.Activated)), slog.F("total_deferred", call.TotalDeferred), - slog.F("budget_rejected", call.BudgetRejected), + slog.F("rejection", call.Rejection), ) }, }) From 4abef9187b1aecc5a9305af3dd600ed1a74c6bfc Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:21:22 +0000 Subject: [PATCH 23/48] fix(coderd/x/chatd): reserve denied same-step direct calls in the find_tools budget --- coderd/x/chatd/chatloop/chatloop.go | 13 ++++++++++++- .../x/chatd/chatloop/chatloop_run_internal_test.go | 10 ++++++++-- coderd/x/chatd/generation.go | 1 + 3 files changed, 21 insertions(+), 3 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 7e5b02dc965..410bef41647 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -243,6 +243,11 @@ type ExecuteLocalToolsOptions struct { 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 @@ -598,6 +603,7 @@ func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (Tool opts.AllowInactiveTools, opts.ProviderTools, localCalls, + opts.ObservedToolCalls, opts.Metrics, opts.Logger, provider, @@ -1062,6 +1068,7 @@ func executeTools( allowInactiveTools map[string]bool, providerTools []ProviderTool, toolCalls []fantasy.ToolCallContent, + observedToolCalls []fantasy.ToolCallContent, metrics *Metrics, logger slog.Logger, provider, model string, @@ -1112,7 +1119,11 @@ func executeTools( } } - notifyStepToolCallObservers(toolMap, toolNameAliases, localToolCalls) + observed := observedToolCalls + if observed == nil { + observed = localToolCalls + } + notifyStepToolCallObservers(toolMap, toolNameAliases, observed) results := make([]fantasy.ToolResultContent, len(localToolCalls)) completedAt := make([]time.Time, len(localToolCalls)) diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 783e5a8bb83..51ee45c08ad 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -937,6 +937,11 @@ func TestExecuteToolsNotifiesStepToolCallObservers(t *testing.T) { {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", @@ -946,8 +951,8 @@ func TestExecuteToolsNotifiesStepToolCallObservers(t *testing.T) { nil, ) - require.Equal(t, []string{"observer_tool", "other_tool"}, observedNames, - "a called observer sees every resolved tool-call name in the step") + 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") } @@ -1002,6 +1007,7 @@ func TestExecuteToolsSerialToolCallOrder(t *testing.T) { nil, nil, calls, + nil, NewMetrics(prometheus.NewRegistry()), slog.Make(), "fake", "fake-model", diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index de792a3db8d..2050195cf2a 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -831,6 +831,7 @@ func (s *taskStarter) executeLocalTools( AllowInactiveTools: prepared.AllowInactiveTools, ProviderTools: prepared.ProviderTools, ToolCalls: allowed, + ObservedToolCalls: decision.localToolCalls, ExclusiveToolNames: prepared.ExclusiveToolNames, BuiltinToolNames: prepared.BuiltinToolNames, ModelProvider: provider, From 103b8bb64811bf282ddc42c3e3c4a0d23063a799 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:32:04 +0000 Subject: [PATCH 24/48] fix(coderd/x/chatd): exclude error-result direct calls from activation derivation --- coderd/x/chatd/mcp_tool_search.go | 13 +++++++++ .../x/chatd/mcp_tool_search_internal_test.go | 27 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index 2e126eedc15..3cb1160caf4 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -207,6 +207,12 @@ func flattenMCPParameterText(value any) string { // 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 do not activate: the call +// was denied before execution (hook policy, input validation) or +// failed, so inlining its schema would spend budget that surviving +// find_tools results have already promised elsewhere. Execution-time +// reservation may still charge such calls, which only under-claims. func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []deferredMCPTool, tokenBudget float64) []string { candidateByName := make(map[string]deferredMCPTool, len(candidates)) for _, candidate := range candidates { @@ -233,6 +239,7 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe } parsedParts := make([][]codersdk.ChatMessagePart, len(rows)) findToolsCallIDs := make(map[string]struct{}) + erroredCallIDs := make(map[string]struct{}) for i := range rows { parts, err := chatprompt.ParseContent(rows[i]) if err != nil { @@ -243,12 +250,18 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName == chattool.FindToolsName && part.ToolCallID != "" { findToolsCallIDs[part.ToolCallID] = struct{}{} } + if part.Type == codersdk.ChatMessagePartTypeToolResult && part.IsError && part.ToolCallID != "" { + erroredCallIDs[part.ToolCallID] = struct{}{} + } } } pendingSearch := make(map[string][]string) for i := len(rows) - 1; i >= 0; i-- { for _, part := range parsedParts[i] { if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName != chattool.FindToolsName { + if _, errored := erroredCallIDs[part.ToolCallID]; errored { + continue + } appendName(part.ToolName) } } diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index 8d7b2577b69..908474c3e1a 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -134,6 +134,33 @@ func TestDeriveDeferredMCPActivationsSameStepDirectCallPriority(t *testing.T) { "same-step search activations cannot shed a directly invoked tool's schema") } +func TestDeriveDeferredMCPActivationsExcludesDeniedDirectCalls(t *testing.T) { + t.Parallel() + candidates := []deferredMCPTool{ + testDeferredTool("server__denied", "denied", nil), + testDeferredTool("server__searched", "searched", nil), + } + assistantStep, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolCall("call-search", chattool.FindToolsName, []byte(`{"queries":["x"]}`)), + codersdk.ChatMessageToolCall("call-denied", "server__denied", []byte(`{}`)), + }) + require.NoError(t, err) + toolRow, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolResult("call-denied", "server__denied", []byte(`"blocked by policy"`), 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"}, deriveDeferredMCPActivations(rows, candidates, 0), + "a direct call with an error result does not activate its schema") + searchedWeight := estimateDeferredMCPToolTokens(candidates[1:]) + require.Equal(t, []string{"server__searched"}, deriveDeferredMCPActivations(rows, candidates, searchedWeight), + "a denied direct call cannot consume budget promised to the search's reported activations") +} + func TestFlattenMCPParameterText(t *testing.T) { t.Parallel() text := flattenMCPParameterText(map[string]any{ From 55cae23f43da2a48fb6b4aea59d95eec409284c8 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:48:24 +0000 Subject: [PATCH 25/48] fix(coderd/x/chatd/chattool): skip oversized find_tools matches once the shared budget is touched --- coderd/x/chatd/chattool/findtools.go | 57 ++++++++++++++----- .../chatd/chattool/findtools_internal_test.go | 56 ++++++++++++------ .../x/chatd/mcp_tool_search_internal_test.go | 2 +- 3 files changed, 83 insertions(+), 32 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index b1268a1f30e..04d6dca0b5b 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -164,8 +164,24 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { // so search with a floor instead of failing outright. searchBudget = findToolsSpentBudgetFloor } - result := SearchTools(searchEntries, args, searchBudget) + 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 { @@ -173,14 +189,11 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { } admitted += schemaTokensByName[name] } - // Derivation retains a single over-budget claim via its - // newest-keep rule, but only when it is the turn's sole - // claim, which is exactly when the budget is untouched. - // Any other over-claim would be silently shed on the - // next request, so fail it loudly instead. A zero new - // claim always succeeds: reserved names are already - // retained by derivation, whatever the budget holds. - if admitted > 0 && admitted > remainingBudget && remainingBudget < options.SchemaTokenBudget { + // 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{ @@ -209,14 +222,27 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { )} } +// 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 schemaTokenBudget additionally stops admitting matches once -// their aggregate schema weight would exceed it, keeping at least one. -func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, schemaTokenBudget float64) FindToolsResult { +// 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 @@ -256,6 +282,7 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, schemaToken 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 @@ -263,7 +290,9 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, schemaToken if len(matches) >= findToolsMaxMatches { return } - if len(matches) > 0 && schemaTokenBudget > 0 && usedSchemaTokens+entry.SchemaTokens > schemaTokenBudget { + overBudget := budget.SchemaTokens > 0 && usedSchemaTokens+entry.SchemaTokens > budget.SchemaTokens + if overBudget && (len(matches) > 0 || !budget.AllowFirstOverBudget) { + budgetSkipped++ return } usedSchemaTokens += entry.SchemaTokens @@ -286,7 +315,7 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, schemaToken activated = append(activated, name) } slices.Sort(activated) - return FindToolsResult{Matches: matches, Activated: activated, TotalDeferred: len(entries)} + return FindToolsResult{Matches: matches, Activated: activated, TotalDeferred: len(entries)}, budgetSkipped } type scopedFindToolsQuery struct { diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 55ca860522d..61d7d9e2460 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -21,24 +21,24 @@ func TestSearchTools(t *testing.T) { t.Run("weights and tie break", func(t *testing.T) { t.Parallel() - result := SearchTools(entries, FindToolsArgs{Queries: []string{"issue"}}, 0) + 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"}}, 0) + 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"}}, 0) + 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{}, 0) + result, _ := SearchTools(entries, FindToolsArgs{}, SearchBudget{}) require.Empty(t, result.Matches) require.Empty(t, result.Activated) }) @@ -48,7 +48,7 @@ func TestSearchTools(t *testing.T) { for i := range many { many[i] = FindToolCatalogEntry{Name: fmt.Sprintf("server__tool_%02d", i), Description: "common"} } - result := SearchTools(many, FindToolsArgs{Queries: []string{"common"}}, 0) + result, _ := SearchTools(many, FindToolsArgs{Queries: []string{"common"}}, SearchBudget{}) require.Len(t, result.Matches, findToolsMaxMatches) require.Equal(t, "server__tool_00", result.Matches[0].Name) }) @@ -60,12 +60,12 @@ func TestSearchTools(t *testing.T) { 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"}}, 0) + 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}, 0) + capped, _ := SearchTools(many, FindToolsArgs{Names: names}, SearchBudget{}) require.Len(t, capped.Matches, findToolsMaxMatches) require.Len(t, capped.Activated, findToolsMaxMatches) }) @@ -75,10 +75,10 @@ func TestSearchTools(t *testing.T) { {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"}}, 0) + result, _ := SearchTools(serverEntries, FindToolsArgs{Queries: []string{"tracking"}}, SearchBudget{}) require.Equal(t, []string{"tracker__create"}, result.Activated) - result = SearchTools(serverEntries, FindToolsArgs{Queries: []string{"project"}}, 0) + 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) @@ -89,15 +89,15 @@ func TestSearchTools(t *testing.T) { {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"}}, 0) + 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:"}}, 0) + 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"}}, 0) + result, _ = SearchTools(scopedEntries, FindToolsArgs{Queries: []string{"error: status"}}, SearchBudget{}) require.Len(t, result.Matches, 2, "an unknown prefix is searched as plain keywords") }) @@ -107,10 +107,10 @@ func TestSearchTools(t *testing.T) { {Name: "docs__検索", Description: "ドキュメント検索"}, {Name: "docs__erstellen", Description: "Dokument ERSTELLEN"}, } - result := SearchTools(unicodeEntries, FindToolsArgs{Queries: []string{"検索"}}, 0) + result, _ := SearchTools(unicodeEntries, FindToolsArgs{Queries: []string{"検索"}}, SearchBudget{}) require.Equal(t, []string{"docs__検索"}, result.Activated) - result = SearchTools(unicodeEntries, FindToolsArgs{Queries: []string{"Erstellen"}}, 0) + 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) { @@ -120,11 +120,11 @@ func TestSearchTools(t *testing.T) { {Name: "server__big_b", Description: "big", SchemaTokens: 60}, {Name: "server__huge", Description: "big", SchemaTokens: 500}, } - result := SearchTools(weighted, FindToolsArgs{Queries: []string{"big"}}, 100) + 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"}}, 100) + 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") }) @@ -134,7 +134,7 @@ func TestSearchTools(t *testing.T) { Name: "server__verbose", Description: strings.Repeat("word ", 100), }} - result := SearchTools(long, FindToolsArgs{Names: []string{"server__verbose"}}, 0) + result, _ := SearchTools(long, FindToolsArgs{Names: []string{"server__verbose"}}, SearchBudget{}) require.LessOrEqual(t, len([]rune(result.Matches[0].Description)), 80) }) } @@ -226,6 +226,28 @@ func TestFindToolsDirectCallReservation(t *testing.T) { require.Empty(t, calls[2].Activated) }) + 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("reserved names stay activatable after the budget is spent", func(t *testing.T) { t.Parallel() tool, observer := newTool(50) diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index 908474c3e1a..b3345a827f1 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -250,7 +250,7 @@ func TestConfigureDeferredMCPToolSearchGenerationFlows(t *testing.T) { second.tool.Info().Name: true, }, allowInactive) - result := chattool.SearchTools(deferredMCPToolEntries(candidates), chattool.FindToolsArgs{Names: []string{second.tool.Info().Name}}, 0) + 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{ From b259a3666ed4016f321c8cee51d25612c96d3679 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:59:03 +0000 Subject: [PATCH 26/48] fix(coderd/x/chatd/chattool): match full server names in find_tools scope prefixes --- coderd/x/chatd/chattool/findtools.go | 40 ++++++++++++++----- .../chatd/chattool/findtools_internal_test.go | 15 +++++++ 2 files changed, 46 insertions(+), 9 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 04d6dca0b5b..cf957de8015 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -325,24 +325,46 @@ type scopedFindToolsQuery struct { // parseFindToolsQueries treats "server: terms" as a scope only when the // prefix names a cataloged server, so queries like "error: timeout" -// still search normally. +// still search normally. Prefixes are matched against full cataloged +// server names, longest first, because workspace server names may +// themselves contain ":". func parseFindToolsQueries(entries []FindToolCatalogEntry, queries []string) []scopedFindToolsQuery { - servers := make(map[string]struct{}, len(entries)) + servers := make([]string, 0, len(entries)) + seen := make(map[string]struct{}, len(entries)) for _, entry := range entries { - if entry.Server != "" { - servers[strings.ToLower(entry.Server)] = struct{}{} + server := strings.ToLower(entry.Server) + if server == "" { + continue + } + if _, dup := seen[server]; dup { + continue } + seen[server] = struct{}{} + servers = append(servers, 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 { - if prefix, rest, ok := strings.Cut(query, ":"); ok { - server := strings.ToLower(strings.TrimSpace(prefix)) - if _, known := servers[server]; known { - parsed = append(parsed, scopedFindToolsQuery{server: server, tokens: tokenizeFindTools(rest)}) + scoped := false + trimmed := strings.ToLower(strings.TrimSpace(query)) + for _, server := range servers { + rest, ok := strings.CutPrefix(trimmed, server) + if !ok { continue } + rest, ok = strings.CutPrefix(strings.TrimLeft(rest, " "), ":") + if !ok { + continue + } + parsed = append(parsed, scopedFindToolsQuery{server: server, tokens: tokenizeFindTools(rest)}) + scoped = true + break + } + if !scoped { + parsed = append(parsed, scopedFindToolsQuery{tokens: tokenizeFindTools(query)}) } - parsed = append(parsed, scopedFindToolsQuery{tokens: tokenizeFindTools(query)}) } return parsed } diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 61d7d9e2460..c025521aba8 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -101,6 +101,21 @@ func TestSearchTools(t *testing.T) { require.Len(t, result.Matches, 2, "an unknown prefix is searched as plain keywords") }) + 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{ From c2fdd27eb4bfb1de7d73298cf13854f24fde26fa Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 06:13:45 +0000 Subject: [PATCH 27/48] fix(coderd/x/chatd): trim workspace server names for the find_tools catalog --- coderd/x/chatd/mcp_tool_search.go | 5 ++++- coderd/x/chatd/mcp_tool_search_internal_test.go | 8 ++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index 3cb1160caf4..6620f50bf42 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -67,9 +67,12 @@ func collectDeferredMCPCandidates(input deferredMCPCandidateInput) []deferredMCP // 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. +// Workspace config validation allows surrounding whitespace in server +// names, which find_tools trims from queries, so the catalog name is +// trimmed too; routing keeps the raw name. func workspaceMCPServerName(tool fantasy.AgentTool) string { if namer, ok := tool.(interface{ ServerName() string }); ok { - return namer.ServerName() + return strings.TrimSpace(namer.ServerName()) } if server, _, ok := strings.Cut(tool.Info().Name, "__"); ok { return server diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index b3345a827f1..e0605544709 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -230,6 +230,14 @@ func TestCollectDeferredMCPCandidates(t *testing.T) { } 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") } func TestConfigureDeferredMCPToolSearchGenerationFlows(t *testing.T) { From ff4c00dee26bb4d2b189af7ed0acfcd32f463c44 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 06:26:15 +0000 Subject: [PATCH 28/48] fix(coderd/x/chatd): activate errored direct calls last instead of excluding them --- coderd/x/chatd/mcp_tool_search.go | 17 +++++++---- .../x/chatd/mcp_tool_search_internal_test.go | 29 ++++++++++++++----- 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index 6620f50bf42..63627e30f5b 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -211,11 +211,13 @@ func flattenMCPParameterText(value any) string { // 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 do not activate: the call -// was denied before execution (hook policy, input validation) or -// failed, so inlining its schema would spend budget that surviving -// find_tools results have already promised elsewhere. Execution-time -// reservation may still charge such calls, which only under-claims. +// 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. func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []deferredMCPTool, tokenBudget float64) []string { candidateByName := make(map[string]deferredMCPTool, len(candidates)) for _, candidate := range candidates { @@ -259,10 +261,12 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe } } pendingSearch := make(map[string][]string) + var erroredNames []string for i := len(rows) - 1; i >= 0; i-- { for _, part := range parsedParts[i] { if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName != chattool.FindToolsName { if _, errored := erroredCallIDs[part.ToolCallID]; errored { + erroredNames = append(erroredNames, part.ToolName) continue } appendName(part.ToolName) @@ -290,6 +294,9 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe } } } + for _, name := range erroredNames { + appendName(name) + } return activated } diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index e0605544709..0684b12e0fc 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -134,19 +134,19 @@ func TestDeriveDeferredMCPActivationsSameStepDirectCallPriority(t *testing.T) { "same-step search activations cannot shed a directly invoked tool's schema") } -func TestDeriveDeferredMCPActivationsExcludesDeniedDirectCalls(t *testing.T) { +func TestDeriveDeferredMCPActivationsErroredDirectCallsActivateLast(t *testing.T) { t.Parallel() candidates := []deferredMCPTool{ - testDeferredTool("server__denied", "denied", nil), + 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-denied", "server__denied", []byte(`{}`)), + codersdk.ChatMessageToolCall("call-errored", "server__errored", []byte(`{}`)), }) require.NoError(t, err) toolRow, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageToolResult("call-denied", "server__denied", []byte(`"blocked by policy"`), true, false), + 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) @@ -154,11 +154,26 @@ func TestDeriveDeferredMCPActivationsExcludesDeniedDirectCalls(t *testing.T) { {Role: database.ChatMessageRoleAssistant, Content: assistantStep, ContentVersion: chatprompt.CurrentContentVersion}, {Role: database.ChatMessageRoleTool, Content: toolRow, ContentVersion: chatprompt.CurrentContentVersion}, } - require.Equal(t, []string{"server__searched"}, deriveDeferredMCPActivations(rows, candidates, 0), - "a direct call with an error result does not activate its schema") + 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), - "a denied direct call cannot consume budget promised to the search's reported activations") + "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 TestFlattenMCPParameterText(t *testing.T) { From e40c26c9d30eba3d11cce8f76419d87b01355b08 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 06:42:29 +0000 Subject: [PATCH 29/48] fix(coderd/x/chatd): refund find_tools reservations when a direct call errors --- coderd/x/chatd/chatloop/chatloop.go | 42 ++++++++++++ .../chatloop/chatloop_run_internal_test.go | 67 +++++++++++++++++++ coderd/x/chatd/chattool/findtools.go | 38 ++++++++++- .../chatd/chattool/findtools_internal_test.go | 35 ++++++++++ coderd/x/chatd/mcp_tool_search.go | 4 +- 5 files changed, 183 insertions(+), 3 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index 410bef41647..ab7aed560ba 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -1188,6 +1188,8 @@ func executeTools( } wg.Wait() + notifyStepToolResultObservers(toolMap, toolNameAliases, results) + // Publish results in the original tool-call order so SSE // subscribers see a deterministic event sequence. if onResult != nil { @@ -1530,6 +1532,46 @@ func notifyStepToolCallObservers(toolMap map[string]fantasy.AgentTool, toolNameA } } +// stepToolResultObserver is implemented by tools that need the step's +// execution outcomes, for example so find_tools can refund budget it +// reserved for a direct call whose execution errored. +type stepToolResultObserver interface { + ObserveStepToolResults(succeeded, errored []string) +} + +// notifyStepToolResultObservers passes the step's resolved result +// outcomes to each distinct called tool that observes them, after +// every call in the step has settled. +func notifyStepToolResultObservers(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, results []fantasy.ToolResultContent) { + succeeded := make([]string, 0, len(results)) + errored := make([]string, 0, len(results)) + for _, tr := range results { + name := tr.ToolName + if alias, ok := toolNameAliases[name]; ok { + name = alias + } + if _, isErr := tr.Result.(fantasy.ToolResultOutputContentError); isErr { + errored = append(errored, name) + } else { + succeeded = append(succeeded, name) + } + } + notified := make(map[string]struct{}, len(results)) + for _, name := range append(append([]string{}, succeeded...), errored...) { + if _, dup := notified[name]; dup { + continue + } + notified[name] = struct{}{} + tool, ok := toolMap[name] + if !ok { + continue + } + if observer, ok := tool.(stepToolResultObserver); ok { + observer.ObserveStepToolResults(succeeded, errored) + } + } +} + func isSerialToolCall(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, name string) bool { if alias, ok := toolNameAliases[name]; ok { name = alias diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 51ee45c08ad..ca70c0c4c99 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -957,6 +957,73 @@ func TestExecuteToolsNotifiesStepToolCallObservers(t *testing.T) { require.False(t, uncalledObserved.Load(), "tools not called this step are not notified") } +type resultObserverMarkerTool struct { + fantasy.AgentTool + observedResults func(succeeded, errored []string) +} + +func (t resultObserverMarkerTool) ObserveStepToolResults(succeeded, errored []string) { + t.observedResults(succeeded, errored) +} + +func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var gotSucceeded, gotErrored []string + 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(succeeded, errored []string) { + mu.Lock() + defer mu.Unlock() + notifications++ + gotSucceeded = append([]string{}, succeeded...) + gotErrored = append([]string{}, 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 + }, + ) + + executeTools( + context.Background(), + quartz.NewReal(), + []fantasy.AgentTool{observer, failing}, + nil, + nil, + nil, + []fantasy.ToolCallContent{ + {ToolCallID: "1", ToolName: "observer_alias", Input: "{}"}, + {ToolCallID: "2", ToolName: "failing_tool", Input: "{}"}, + {ToolCallID: "3", ToolName: "missing_tool", Input: "{}"}, + }, + nil, + 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"}, gotSucceeded) + require.Equal(t, []string{"failing_tool", "missing_tool"}, gotErrored, + "error results and unresolvable tools both settle as errored outcomes") +} + func TestExecuteToolsSerialToolCallOrder(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index cf957de8015..684856a0b35 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -96,13 +96,18 @@ type FindToolsResult struct { // charged against the budget before any search admits activations. type findToolsTool struct { fantasy.AgentTool - reserveStepCalls func(names []string) + reserveStepCalls func(names []string) + settleStepResults func(succeeded, errored []string) } func (findToolsTool) SerialToolCalls() bool { return true } func (t findToolsTool) ObserveStepToolCalls(names []string) { t.reserveStepCalls(names) } +func (t findToolsTool) ObserveStepToolResults(succeeded, errored []string) { + t.settleStepResults(succeeded, errored) +} + // FindTools returns the built-in used to discover deferred MCP tool schemas. func FindTools(options FindToolsOptions) fantasy.AgentTool { entries := slices.Clone(options.Entries) @@ -117,6 +122,7 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { // is reserved out of the budget before searches run. Reserved names // stay free to activate because derivation already retains them. reserved := make(map[string]struct{}) + executedOK := make(map[string]struct{}) reserve := func(names []string) { if options.SchemaTokenBudget <= 0 { return @@ -135,7 +141,35 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { remainingBudget -= weight } } - return findToolsTool{reserveStepCalls: reserve, AgentTool: fantasy.NewAgentTool( + // Derivation admits errored direct calls only with leftover budget, + // so once a step ends their pre-execution reservation is refunded + // and later searches regain the weight. Names that ever executed + // successfully keep their reservation: derivation admits them at + // full priority. Denied calls never produce step results here and + // stay charged, which only under-claims. + settle := func(succeeded, errored []string) { + if options.SchemaTokenBudget <= 0 { + return + } + budgetMu.Lock() + defer budgetMu.Unlock() + for _, name := range succeeded { + if _, ok := schemaTokensByName[name]; ok { + executedOK[name] = struct{}{} + } + } + for _, name := range errored { + if _, ok := executedOK[name]; ok { + continue + } + if _, ok := reserved[name]; !ok { + continue + } + delete(reserved, name) + remainingBudget += schemaTokensByName[name] + } + } + return findToolsTool{reserveStepCalls: reserve, settleStepResults: settle, AgentTool: fantasy.NewAgentTool( FindToolsName, buildFindToolsDescription(entries, options.CatalogTokenBudget), func(ctx context.Context, args FindToolsArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index c025521aba8..6497f1878eb 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -263,6 +263,41 @@ func TestFindToolsDirectCallReservation(t *testing.T) { "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(succeeded, errored []string) + }) + 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(nil, []string{"server__a", "unknown"}) + 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(succeeded, errored []string) + }) + require.True(t, ok) + observer.ObserveStepToolCalls([]string{"server__a"}) + settler.ObserveStepToolResults([]string{"server__a"}, []string{"server__a"}) + 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("reserved names stay activatable after the budget is spent", func(t *testing.T) { t.Parallel() tool, observer := newTool(50) diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index 63627e30f5b..4f96520be51 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -217,7 +217,9 @@ func flattenMCPParameterText(value any) string { // 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. +// or successful calls already claimed. Search-time reservations mirror +// this order because the step result observer refunds an errored +// call's reservation, so later 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 { From fcc8ecd9cc0ab23aafa2ea8ca558f6339e2959ca Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 06:51:26 +0000 Subject: [PATCH 30/48] fix(coderd/x/chatd): settle sibling outcomes before find_tools searches run --- coderd/x/chatd/chatloop/chatloop.go | 57 +++++++++------ .../chatloop/chatloop_run_internal_test.go | 72 +++++++++++++++++++ coderd/x/chatd/chattool/findtools.go | 11 +-- coderd/x/chatd/mcp_tool_search.go | 5 +- 4 files changed, 115 insertions(+), 30 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index ab7aed560ba..cf6bf2ed8cd 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -1160,10 +1160,12 @@ func executeTools( toolNameAliases, ) } - // Calls to tools that opt in via SerialToolCalls run on one - // goroutine in tool-call order, so order-sensitive shared state - // (for example the find_tools activation budget) is claimed - // deterministically. All other calls stay concurrent. + // 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 for i, tc := range localToolCalls { @@ -1177,18 +1179,22 @@ func executeTools( runCall(i, tc) }() } - if len(serialIndexes) > 0 { - wg.Add(1) - go func() { - defer wg.Done() - for _, i := range serialIndexes { - runCall(i, localToolCalls[i]) - } - }() - } wg.Wait() - notifyStepToolResultObservers(toolMap, toolNameAliases, results) + // 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, 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. @@ -1539,13 +1545,14 @@ type stepToolResultObserver interface { ObserveStepToolResults(succeeded, errored []string) } -// notifyStepToolResultObservers passes the step's resolved result -// outcomes to each distinct called tool that observes them, after -// every call in the step has settled. -func notifyStepToolResultObservers(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, results []fantasy.ToolResultContent) { - succeeded := make([]string, 0, len(results)) - errored := make([]string, 0, len(results)) - for _, tr := range results { +// notifyStepToolResultObservers passes the settled sibling outcomes to +// each distinct called tool that observes them. Serial calls have not +// run yet, so their own outcomes are absent; observers only need the +// concurrent siblings they share state with. +func notifyStepToolResultObservers(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, calls []fantasy.ToolCallContent, settled []fantasy.ToolResultContent) { + succeeded := make([]string, 0, len(settled)) + errored := make([]string, 0, len(settled)) + for _, tr := range settled { name := tr.ToolName if alias, ok := toolNameAliases[name]; ok { name = alias @@ -1556,8 +1563,12 @@ func notifyStepToolResultObservers(toolMap map[string]fantasy.AgentTool, toolNam succeeded = append(succeeded, name) } } - notified := make(map[string]struct{}, len(results)) - for _, name := range append(append([]string{}, succeeded...), errored...) { + 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 } diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index ca70c0c4c99..49e844122f5 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -1024,6 +1024,78 @@ func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { "error results and unresolvable tools both settle as errored outcomes") } +type serialResultObserverTool struct { + fantasy.AgentTool + observedResults func(succeeded, errored []string) +} + +func (serialResultObserverTool) SerialToolCalls() bool { return true } + +func (t serialResultObserverTool) ObserveStepToolResults(succeeded, errored []string) { + t.observedResults(succeeded, 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(_, errored []string) { + mu.Lock() + defer mu.Unlock() + notified = true + erroredAtNotify = append([]string{}, 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 + }, + ) + + 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() diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 684856a0b35..3e26925b821 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -142,11 +142,12 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { } } // Derivation admits errored direct calls only with leftover budget, - // so once a step ends their pre-execution reservation is refunded - // and later searches regain the weight. Names that ever executed - // successfully keep their reservation: derivation admits them at - // full priority. Denied calls never produce step results here and - // stay charged, which only under-claims. + // so their pre-execution reservation is refunded once the step's + // concurrent siblings settle, before this step's searches run. + // Names that ever executed successfully keep their reservation: + // derivation admits them at full priority. Hook-denied calls never + // produce step results here and stay charged, which only + // under-claims. settle := func(succeeded, errored []string) { if options.SchemaTokenBudget <= 0 { return diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index 4f96520be51..7f7eaa90c09 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -218,8 +218,9 @@ func flattenMCPParameterText(value any) string { // 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 the step result observer refunds an errored -// call's reservation, so later searches see the same leftover budget. +// 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 { From ded88f961a87b29c6d20dfcdae092273a71e850f Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:03:22 +0000 Subject: [PATCH 31/48] fix(coderd/x/chatd): settle pre-execution tool rejections as errored outcomes --- coderd/x/chatd/chatloop/chatloop.go | 32 +++++++++++++------ .../chatloop/chatloop_run_internal_test.go | 20 +++++++----- coderd/x/chatd/chattool/findtools.go | 8 ++--- 3 files changed, 39 insertions(+), 21 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index cf6bf2ed8cd..eeb47111067 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -1190,7 +1190,7 @@ func executeTools( settled = append(settled, results[i]) } } - notifyStepToolResultObservers(toolMap, toolNameAliases, localToolCalls, settled) + notifyStepToolResultObservers(toolMap, toolNameAliases, localToolCalls, observed, settled) for _, i := range serialIndexes { runCall(i, localToolCalls[i]) @@ -1548,19 +1548,33 @@ type stepToolResultObserver interface { // notifyStepToolResultObservers passes the settled sibling outcomes to // each distinct called tool that observes them. Serial calls have not // run yet, so their own outcomes are absent; observers only need the -// concurrent siblings they share state with. -func notifyStepToolResultObservers(toolMap map[string]fantasy.AgentTool, toolNameAliases map[string]string, calls []fantasy.ToolCallContent, settled []fantasy.ToolResultContent) { +// concurrent siblings they share state with. 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. +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 + } succeeded := make([]string, 0, len(settled)) errored := make([]string, 0, len(settled)) for _, tr := range settled { - name := tr.ToolName - if alias, ok := toolNameAliases[name]; ok { - name = alias - } if _, isErr := tr.Result.(fantasy.ToolResultOutputContentError); isErr { - errored = append(errored, name) + errored = append(errored, resolve(tr.ToolName)) } else { - succeeded = append(succeeded, name) + succeeded = append(succeeded, resolve(tr.ToolName)) + } + } + executedIDs := make(map[string]struct{}, len(calls)) + for _, tc := range calls { + executedIDs[tc.ToolCallID] = struct{}{} + } + for _, tc := range observed { + if _, ok := executedIDs[tc.ToolCallID]; !ok { + errored = append(errored, resolve(tc.ToolName)) } } notified := make(map[string]struct{}, len(calls)) diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 49e844122f5..cce9a050679 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -6,6 +6,7 @@ import ( "errors" "iter" "runtime" + "slices" "sync" "sync/atomic" "testing" @@ -996,6 +997,11 @@ func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { }, ) + 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(), @@ -1003,12 +1009,10 @@ func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { nil, nil, nil, - []fantasy.ToolCallContent{ - {ToolCallID: "1", ToolName: "observer_alias", Input: "{}"}, - {ToolCallID: "2", ToolName: "failing_tool", Input: "{}"}, - {ToolCallID: "3", ToolName: "missing_tool", Input: "{}"}, - }, - nil, + executed, + append(slices.Clone(executed), fantasy.ToolCallContent{ + ToolCallID: "4", ToolName: "rejected_tool", Input: "{not json", + }), NewMetrics(prometheus.NewRegistry()), slog.Make(), "fake", "fake-model", @@ -1020,8 +1024,8 @@ func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { require.Equal(t, 1, notifications, "each called observer is notified once per step") require.Equal(t, []string{"observer_tool"}, gotSucceeded) - require.Equal(t, []string{"failing_tool", "missing_tool"}, gotErrored, - "error results and unresolvable tools both settle as errored outcomes") + require.Equal(t, []string{"failing_tool", "missing_tool", "rejected_tool"}, gotErrored, + "error results, unresolvable tools, and observed calls rejected before execution all settle as errored outcomes") } type serialResultObserverTool struct { diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 3e26925b821..d7f7123fb05 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -144,10 +144,10 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { // Derivation admits errored direct calls only with leftover budget, // so their pre-execution reservation is refunded once the step's // concurrent siblings settle, before this step's searches run. - // Names that ever executed successfully keep their reservation: - // derivation admits them at full priority. Hook-denied calls never - // produce step results here and stay charged, which only - // under-claims. + // Calls rejected before execution settle as errored too, so their + // reservations refund the same way. Names that ever executed + // successfully keep their reservation: derivation admits them at + // full priority. settle := func(succeeded, errored []string) { if options.SchemaTokenBudget <= 0 { return From d8f5cc526fb7d6fef4760dc4cd483f9eb84dd965 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:20:00 +0000 Subject: [PATCH 32/48] fix(coderd/x/chatd/chattool): mirror derivation's retained prefix and exact-case scopes in find_tools --- coderd/x/chatd/chattool/findtools.go | 141 ++++++++++++------ .../chatd/chattool/findtools_internal_test.go | 79 ++++++++++ 2 files changed, 172 insertions(+), 48 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index d7f7123fb05..7e7a9f64082 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -118,11 +118,41 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { var budgetMu sync.Mutex remainingBudget := options.SchemaTokenBudget // Direct calls to deferred tools in the same step are admitted by - // derivation before any search activations, so their schema weight - // is reserved out of the budget before searches run. Reserved names - // stay free to activate because derivation already retains them. + // derivation before any search activations, 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 direct-call + // position, so no same-step search can inline its schema either. + // Errored calls (including calls rejected before execution) leave + // the prefix when siblings settle; a name that ever executed + // successfully stays, since derivation admits it at full priority. + var observedOrder []string reserved := make(map[string]struct{}) + unclaimable := make(map[string]struct{}) executedOK := make(map[string]struct{}) + erroredNames := make(map[string]struct{}) + searchClaimed := 0.0 + recompute := func() { + clear(reserved) + clear(unclaimable) + charged := 0.0 + for _, name := range observedOrder { + if _, errored := erroredNames[name]; errored { + if _, ok := executedOK[name]; !ok { + continue + } + } + weight := schemaTokensByName[name] + if len(reserved) > 0 && charged+weight > options.SchemaTokenBudget { + unclaimable[name] = struct{}{} + continue + } + reserved[name] = struct{}{} + charged += weight + } + remainingBudget = options.SchemaTokenBudget - charged - searchClaimed + } reserve := func(names []string) { if options.SchemaTokenBudget <= 0 { return @@ -130,24 +160,16 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { budgetMu.Lock() defer budgetMu.Unlock() for _, name := range names { - weight, ok := schemaTokensByName[name] - if !ok { + if _, ok := schemaTokensByName[name]; !ok { continue } - if _, dup := reserved[name]; dup { + if slices.Contains(observedOrder, name) { continue } - reserved[name] = struct{}{} - remainingBudget -= weight + observedOrder = append(observedOrder, name) } + recompute() } - // Derivation admits errored direct calls only with leftover budget, - // so their pre-execution reservation is refunded once the step's - // concurrent siblings settle, before this step's searches run. - // Calls rejected before execution settle as errored too, so their - // reservations refund the same way. Names that ever executed - // successfully keep their reservation: derivation admits them at - // full priority. settle := func(succeeded, errored []string) { if options.SchemaTokenBudget <= 0 { return @@ -160,15 +182,9 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { } } for _, name := range errored { - if _, ok := executedOK[name]; ok { - continue - } - if _, ok := reserved[name]; !ok { - continue - } - delete(reserved, name) - remainingBudget += schemaTokensByName[name] + erroredNames[name] = struct{}{} } + recompute() } return findToolsTool{reserveStepCalls: reserve, settleStepResults: settle, AgentTool: fantasy.NewAgentTool( FindToolsName, @@ -185,12 +201,16 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { } budgetMu.Lock() searchEntries := entries - if len(reserved) > 0 { - searchEntries = slices.Clone(entries) - for i := range searchEntries { - if _, ok := reserved[searchEntries[i].Name]; ok { - searchEntries[i].SchemaTokens = 0 + if len(reserved) > 0 || len(unclaimable) > 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 } + searchEntries = append(searchEntries, entry) } } searchBudget := remainingBudget @@ -240,9 +260,12 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { } return fantasy.NewTextErrorResponse(findToolsBudgetExhausted), nil } + searchClaimed += admitted remainingBudget -= admitted } 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, @@ -293,8 +316,13 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, budget Sear for _, entry := range entries { score := 0 for _, query := range queries { - if query.server != "" && !strings.EqualFold(entry.Server, query.server) { - continue + 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++ @@ -355,6 +383,10 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, budget Sear 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 } @@ -362,20 +394,21 @@ type scopedFindToolsQuery struct { // 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 ":". +// 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 { - server := strings.ToLower(entry.Server) - if server == "" { + if entry.Server == "" { continue } - if _, dup := seen[server]; dup { + if _, dup := seen[entry.Server]; dup { continue } - seen[server] = struct{}{} - servers = append(servers, server) + seen[entry.Server] = struct{}{} + servers = append(servers, entry.Server) } // Longest first, so a server named "jira:prod" wins over "jira" // when both are cataloged. @@ -383,19 +416,31 @@ func parseFindToolsQueries(entries []FindToolCatalogEntry, queries []string) []s parsed := make([]scopedFindToolsQuery, 0, len(queries)) for _, query := range queries { scoped := false - trimmed := strings.ToLower(strings.TrimSpace(query)) - for _, server := range servers { - rest, ok := strings.CutPrefix(trimmed, server) - if !ok { - continue - } - rest, ok = strings.CutPrefix(strings.TrimLeft(rest, " "), ":") - if !ok { - continue + trimmed := strings.TrimSpace(query) + for pass := 0; pass < 2 && !scoped; pass++ { + exact := pass == 0 + for _, server := range servers { + var rest string + if exact { + var ok bool + rest, ok = strings.CutPrefix(trimmed, server) + if !ok { + continue + } + } else { + if len(trimmed) < len(server) || !strings.EqualFold(trimmed[:len(server)], server) { + continue + } + rest = trimmed[len(server):] + } + rest, ok := strings.CutPrefix(strings.TrimLeft(rest, " "), ":") + if !ok { + continue + } + parsed = append(parsed, scopedFindToolsQuery{server: server, exact: exact, tokens: tokenizeFindTools(rest)}) + scoped = true + break } - parsed = append(parsed, scopedFindToolsQuery{server: server, tokens: tokenizeFindTools(rest)}) - scoped = true - break } if !scoped { parsed = append(parsed, scopedFindToolsQuery{tokens: tokenizeFindTools(query)}) diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 6497f1878eb..c8c35a68be2 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -100,6 +100,28 @@ func TestSearchTools(t *testing.T) { 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("server names containing colons", func(t *testing.T) { t.Parallel() @@ -298,6 +320,63 @@ func TestFindToolsDirectCallReservation(t *testing.T) { require.True(t, resp.IsError, "a successful execution pins the reservation even when a later call errors") }) + 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(succeeded, errored []string) + }) + require.True(t, ok) + observer.ObserveStepToolCalls([]string{"server__a", "server__b"}) + settler.ObserveStepToolResults([]string{"server__a", "server__b"}, nil) + + 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("an errored prefix call promotes the next observed name", func(t *testing.T) { + t.Parallel() + tool, observer := newTool(100) + settler, ok := tool.(interface { + ObserveStepToolResults(succeeded, errored []string) + }) + require.True(t, ok) + observer.ObserveStepToolCalls([]string{"server__a", "server__b"}) + settler.ObserveStepToolResults([]string{"server__b"}, []string{"server__a"}) + + 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) From afc7cbe4eb610ab9ea42e2e02994b4aa8baaea09 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:34:01 +0000 Subject: [PATCH 33/48] fix(coderd/x/chatd): keep raw workspace server names when trimming would collide --- .../chatd/chattool/findtools_internal_test.go | 10 ++++++ coderd/x/chatd/mcp_tool_search.go | 35 ++++++++++++++++--- .../x/chatd/mcp_tool_search_internal_test.go | 22 ++++++++++++ 3 files changed, 63 insertions(+), 4 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index c8c35a68be2..32167252e3d 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -123,6 +123,16 @@ func TestSearchTools(t *testing.T) { require.Len(t, result.Activated, 2, "a prefix matching no exact-case name falls back to spanning the case-colliding servers") }) + 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") + }) t.Run("server names containing colons", func(t *testing.T) { t.Parallel() colonEntries := []FindToolCatalogEntry{ diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index 7f7eaa90c09..fc757796a9d 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -54,25 +54,52 @@ func collectDeferredMCPCandidates(input deferredMCPCandidateInput) []deferredMCP 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. -// Workspace config validation allows surrounding whitespace in server -// names, which find_tools trims from queries, so the catalog name is -// trimmed too; routing keeps the raw name. func workspaceMCPServerName(tool fantasy.AgentTool) string { if namer, ok := tool.(interface{ ServerName() string }); ok { - return strings.TrimSpace(namer.ServerName()) + return namer.ServerName() } if server, _, ok := strings.Cut(tool.Info().Name, "__"); ok { return server diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index 0684b12e0fc..a424c45090f 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -253,6 +253,28 @@ func TestCollectDeferredMCPCandidates(t *testing.T) { } 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) { From 910fbaa21f072c77d6bc18e096b094c6ea54ebbb Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:44:13 +0000 Subject: [PATCH 34/48] fix(coderd/x/chatd/chattool): dedup repeat search claims and match raw padded server scopes --- coderd/x/chatd/chattool/findtools.go | 44 +++++++++++++++---- .../chatd/chattool/findtools_internal_test.go | 33 +++++++++++++- 2 files changed, 67 insertions(+), 10 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 7e7a9f64082..cce4af38e8a 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -132,6 +132,9 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { unclaimable := make(map[string]struct{}) executedOK := make(map[string]struct{}) erroredNames := 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) @@ -201,7 +204,7 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { } budgetMu.Lock() searchEntries := entries - if len(reserved) > 0 || len(unclaimable) > 0 { + 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 { @@ -210,6 +213,9 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { if _, ok := reserved[entry.Name]; ok { entry.SchemaTokens = 0 } + if _, ok := claimedBySearch[entry.Name]; ok { + entry.SchemaTokens = 0 + } searchEntries = append(searchEntries, entry) } } @@ -242,6 +248,9 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { if _, ok := reserved[name]; ok { continue } + if _, ok := claimedBySearch[name]; ok { + continue + } admitted += schemaTokensByName[name] } // Defensive invariant: with allowFirstOverBudget off, @@ -262,6 +271,11 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { } 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. @@ -417,27 +431,41 @@ func parseFindToolsQueries(entries []FindToolCatalogEntry, queries []string) []s for _, query := range queries { scoped := false trimmed := strings.TrimSpace(query) - for pass := 0; pass < 2 && !scoped; pass++ { - exact := pass == 0 + // 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 exact { + if pass.exact { var ok bool - rest, ok = strings.CutPrefix(trimmed, server) + rest, ok = strings.CutPrefix(pass.text, server) if !ok { continue } } else { - if len(trimmed) < len(server) || !strings.EqualFold(trimmed[:len(server)], server) { + if len(pass.text) < len(server) || !strings.EqualFold(pass.text[:len(server)], server) { continue } - rest = trimmed[len(server):] + rest = pass.text[len(server):] } rest, ok := strings.CutPrefix(strings.TrimLeft(rest, " "), ":") if !ok { continue } - parsed = append(parsed, scopedFindToolsQuery{server: server, exact: exact, tokens: tokenizeFindTools(rest)}) + parsed = append(parsed, scopedFindToolsQuery{server: server, exact: pass.exact, tokens: tokenizeFindTools(rest)}) scoped = true break } diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 32167252e3d..300aaef57a9 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -132,6 +132,10 @@ func TestSearchTools(t *testing.T) { 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() @@ -363,6 +367,25 @@ func TestFindToolsDirectCallReservation(t *testing.T) { "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) @@ -433,7 +456,10 @@ func TestFindToolsSharedSchemaBudget(t *testing.T) { "a call whose claims cannot fit the remaining budget errors instead of over-claiming") huge := FindTools(FindToolsOptions{ - Entries: []FindToolCatalogEntry{{Name: "server__huge", SchemaTokens: 500}}, + 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"]}`}) @@ -444,7 +470,10 @@ func TestFindToolsSharedSchemaBudget(t *testing.T) { "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.True(t, resp.IsError, "the spent budget rejects further activations") + 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) { From 167a1131e280a1bfbd452d1b5a4c3085b7c08457 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:00:57 +0000 Subject: [PATCH 35/48] fix(coderd/x/chatd): track find_tools reservation outcomes per call instead of per name --- coderd/x/chatd/chatloop/chatloop.go | 44 ++++++---- .../chatloop/chatloop_run_internal_test.go | 35 +++++--- coderd/x/chatd/chattool/findtools.go | 88 ++++++++++--------- .../chatd/chattool/findtools_internal_test.go | 54 ++++++++++-- 4 files changed, 136 insertions(+), 85 deletions(-) diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index eeb47111067..e7586520f1e 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -1539,19 +1539,23 @@ func notifyStepToolCallObservers(toolMap map[string]fantasy.AgentTool, toolNameA } // stepToolResultObserver is implemented by tools that need the step's -// execution outcomes, for example so find_tools can refund budget it -// reserved for a direct call whose execution errored. +// 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(succeeded, errored []string) + ObserveStepToolResults(names []string, errored []bool) } // notifyStepToolResultObservers passes the settled sibling outcomes to -// each distinct called tool that observes them. Serial calls have not -// run yet, so their own outcomes are absent; observers only need the -// concurrent siblings they share state with. 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. +// 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 { @@ -1559,23 +1563,25 @@ func notifyStepToolResultObservers(toolMap map[string]fantasy.AgentTool, toolNam } return name } - succeeded := make([]string, 0, len(settled)) - errored := make([]string, 0, len(settled)) + erroredByID := make(map[string]bool, len(settled)) for _, tr := range settled { - if _, isErr := tr.Result.(fantasy.ToolResultOutputContentError); isErr { - errored = append(errored, resolve(tr.ToolName)) - } else { - succeeded = append(succeeded, resolve(tr.ToolName)) - } + _, 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 { - if _, ok := executedIDs[tc.ToolCallID]; !ok { - errored = append(errored, resolve(tc.ToolName)) + 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 { @@ -1592,7 +1598,7 @@ func notifyStepToolResultObservers(toolMap map[string]fantasy.AgentTool, toolNam continue } if observer, ok := tool.(stepToolResultObserver); ok { - observer.ObserveStepToolResults(succeeded, errored) + observer.ObserveStepToolResults(names, errored) } } } diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index cce9a050679..7d35d2588e7 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -960,18 +960,19 @@ func TestExecuteToolsNotifiesStepToolCallObservers(t *testing.T) { type resultObserverMarkerTool struct { fantasy.AgentTool - observedResults func(succeeded, errored []string) + observedResults func(names []string, errored []bool) } -func (t resultObserverMarkerTool) ObserveStepToolResults(succeeded, errored []string) { - t.observedResults(succeeded, errored) +func (t resultObserverMarkerTool) ObserveStepToolResults(names []string, errored []bool) { + t.observedResults(names, errored) } func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { t.Parallel() var mu sync.Mutex - var gotSucceeded, gotErrored []string + var gotNames []string + var gotErrored []bool notifications := 0 observer := resultObserverMarkerTool{ AgentTool: fantasy.NewAgentTool( @@ -981,12 +982,12 @@ func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { return fantasy.NewTextResponse("ok"), nil }, ), - observedResults: func(succeeded, errored []string) { + observedResults: func(names []string, errored []bool) { mu.Lock() defer mu.Unlock() notifications++ - gotSucceeded = append([]string{}, succeeded...) - gotErrored = append([]string{}, errored...) + gotNames = append([]string{}, names...) + gotErrored = append([]bool{}, errored...) }, } failing := fantasy.NewAgentTool( @@ -1023,20 +1024,21 @@ func TestExecuteToolsNotifiesStepToolResultObservers(t *testing.T) { ) require.Equal(t, 1, notifications, "each called observer is notified once per step") - require.Equal(t, []string{"observer_tool"}, gotSucceeded) - require.Equal(t, []string{"failing_tool", "missing_tool", "rejected_tool"}, gotErrored, + 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(succeeded, errored []string) + observedResults func(names []string, errored []bool) } func (serialResultObserverTool) SerialToolCalls() bool { return true } -func (t serialResultObserverTool) ObserveStepToolResults(succeeded, errored []string) { - t.observedResults(succeeded, errored) +func (t serialResultObserverTool) ObserveStepToolResults(names []string, errored []bool) { + t.observedResults(names, errored) } func TestExecuteToolsReconcilesResultsBeforeSerialCalls(t *testing.T) { @@ -1057,11 +1059,16 @@ func TestExecuteToolsReconcilesResultsBeforeSerialCalls(t *testing.T) { return fantasy.NewTextResponse("ok"), nil }, ), - observedResults: func(_, errored []string) { + observedResults: func(names []string, errored []bool) { mu.Lock() defer mu.Unlock() notified = true - erroredAtNotify = append([]string{}, errored...) + erroredAtNotify = nil + for i, name := range names { + if errored[i] { + erroredAtNotify = append(erroredAtNotify, name) + } + } }, } failing := fantasy.NewAgentTool( diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index cce4af38e8a..5216195b2c8 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -97,15 +97,15 @@ type FindToolsResult struct { type findToolsTool struct { fantasy.AgentTool reserveStepCalls func(names []string) - settleStepResults func(succeeded, errored []string) + settleStepResults func(names []string, errored []bool) } func (findToolsTool) SerialToolCalls() bool { return true } func (t findToolsTool) ObserveStepToolCalls(names []string) { t.reserveStepCalls(names) } -func (t findToolsTool) ObserveStepToolResults(succeeded, errored []string) { - t.settleStepResults(succeeded, errored) +func (t findToolsTool) ObserveStepToolResults(names []string, errored []bool) { + t.settleStepResults(names, errored) } // FindTools returns the built-in used to discover deferred MCP tool schemas. @@ -118,20 +118,23 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { var budgetMu sync.Mutex remainingBudget := options.SchemaTokenBudget // Direct calls to deferred tools in the same step are admitted by - // derivation before any search activations, 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 direct-call - // position, so no same-step search can inline its schema either. - // Errored calls (including calls rejected before execution) leave - // the prefix when siblings settle; a name that ever executed - // successfully stays, since derivation admits it at full priority. - var observedOrder []string + // 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{}) - executedOK := make(map[string]struct{}) - erroredNames := 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{}) @@ -140,54 +143,53 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { clear(reserved) clear(unclaimable) charged := 0.0 - for _, name := range observedOrder { - if _, errored := erroredNames[name]; errored { - if _, ok := executedOK[name]; !ok { - continue - } + seen := make(map[string]struct{}, len(stepCalls)) + for _, call := range stepCalls { + if call.errored { + continue + } + if _, dup := seen[call.name]; dup { + continue } - weight := schemaTokensByName[name] + seen[call.name] = struct{}{} + weight := schemaTokensByName[call.name] if len(reserved) > 0 && charged+weight > options.SchemaTokenBudget { - unclaimable[name] = struct{}{} + unclaimable[call.name] = struct{}{} continue } - reserved[name] = struct{}{} + reserved[call.name] = struct{}{} charged += weight } remainingBudget = options.SchemaTokenBudget - charged - searchClaimed } - reserve := func(names []string) { - if options.SchemaTokenBudget <= 0 { - return - } - budgetMu.Lock() - defer budgetMu.Unlock() - for _, name := range names { + rebuild := func(names []string, errored []bool) { + stepCalls = stepCalls[:0] + for i, name := range names { if _, ok := schemaTokensByName[name]; !ok { continue } - if slices.Contains(observedOrder, name) { - continue - } - observedOrder = append(observedOrder, name) + stepCalls = append(stepCalls, stepToolCall{name: name, errored: len(errored) > i && errored[i]}) } recompute() } - settle := func(succeeded, errored []string) { + reserve := func(names []string) { if options.SchemaTokenBudget <= 0 { return } budgetMu.Lock() defer budgetMu.Unlock() - for _, name := range succeeded { - if _, ok := schemaTokensByName[name]; ok { - executedOK[name] = struct{}{} - } - } - for _, name := range errored { - erroredNames[name] = struct{}{} + // 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 } - recompute() + budgetMu.Lock() + defer budgetMu.Unlock() + rebuild(names, errored) } return findToolsTool{reserveStepCalls: reserve, settleStepResults: settle, AgentTool: fantasy.NewAgentTool( FindToolsName, diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 300aaef57a9..19a002aba94 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -303,7 +303,7 @@ func TestFindToolsDirectCallReservation(t *testing.T) { t.Parallel() tool, observer := newTool(100) settler, ok := tool.(interface { - ObserveStepToolResults(succeeded, errored []string) + ObserveStepToolResults(names []string, errored []bool) }) require.True(t, ok, "find_tools must observe step results to refund errored reservations") observer.ObserveStepToolCalls([]string{"server__a"}) @@ -311,7 +311,7 @@ func TestFindToolsDirectCallReservation(t *testing.T) { require.NoError(t, err) require.True(t, resp.IsError, "the pre-execution reservation holds while the outcome is unknown") - settler.ObserveStepToolResults(nil, []string{"server__a", "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") @@ -324,25 +324,61 @@ func TestFindToolsDirectCallReservation(t *testing.T) { t.Parallel() tool, observer := newTool(100) settler, ok := tool.(interface { - ObserveStepToolResults(succeeded, errored []string) + ObserveStepToolResults(names []string, errored []bool) }) require.True(t, ok) - observer.ObserveStepToolCalls([]string{"server__a"}) - settler.ObserveStepToolResults([]string{"server__a"}, []string{"server__a"}) + 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(succeeded, errored []string) + ObserveStepToolResults(names []string, errored []bool) }) require.True(t, ok) observer.ObserveStepToolCalls([]string{"server__a", "server__b"}) - settler.ObserveStepToolResults([]string{"server__a", "server__b"}, nil) + 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) @@ -390,11 +426,11 @@ func TestFindToolsDirectCallReservation(t *testing.T) { t.Parallel() tool, observer := newTool(100) settler, ok := tool.(interface { - ObserveStepToolResults(succeeded, errored []string) + ObserveStepToolResults(names []string, errored []bool) }) require.True(t, ok) observer.ObserveStepToolCalls([]string{"server__a", "server__b"}) - settler.ObserveStepToolResults([]string{"server__b"}, []string{"server__a"}) + 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) From b57393a673c228620be0bb52d200bff4bfdfd716 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:39:03 +0000 Subject: [PATCH 36/48] docs(coderd/x/chatd): drop the architecture TODO marker --- coderd/x/chatd/ARCHITECTURE.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index b85fbc19a53..d9f1e16f8f6 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -842,8 +842,6 @@ Tool calls have at least once semantics: if the goroutine executes a tool call, Parallel tool call results must be inserted in bulk after all parallel tool calls finish in a single `CommitStep` transition so that the generation goroutine only increments `history_version` once, since a change to the `history_version` interrupts the gorotuine. This is consistent with the existing chatd implementation. - - The generation goroutine supports: - chat compaction (automatic and manual, see [Manual compaction](#manual-compaction)) From 84a6c584e3a56a29a919156f69d0a67addee0a86 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 09:04:19 +0000 Subject: [PATCH 37/48] feat(coderd/x/chatd): defer all MCP tool schemas whenever the experiment is enabled Removes the ContextLimit/10 deferral threshold, the estimated-token decision, the deferred_mcp_tool_tokens metric, and the --chat-mcp-tool-search-force-defer deployment flag. With the mcp-tool-search experiment enabled, every generation with deferrable MCP candidates now defers behind find_tools. Also pairs activation-derivation error results with their own tool call, so a provider reusing a tool-call ID cannot demote a later successful call. --- cli/testdata/server-config.yaml.golden | 4 -- coderd/apidoc/docs.go | 3 - coderd/apidoc/swagger.json | 3 - coderd/coderd.go | 1 - coderd/x/chatd/chatd.go | 8 +-- coderd/x/chatd/chatd_test.go | 48 +++++++++++-- coderd/x/chatd/chatloop/metrics.go | 8 --- coderd/x/chatd/forced_mcp_test.go | 3 + coderd/x/chatd/generation_preparer.go | 13 +--- coderd/x/chatd/mcp_tool_search.go | 57 ++++++++++------ .../x/chatd/mcp_tool_search_internal_test.go | 67 +++++++++++++------ codersdk/deployment.go | 26 ++----- docs/admin/integrations/prometheus.md | 1 - docs/reference/api/general.md | 3 +- docs/reference/api/schemas.md | 31 ++++----- scripts/metricsdocgen/generated_metrics | 3 - site/src/api/typesGenerated.ts | 1 - 17 files changed, 150 insertions(+), 130 deletions(-) diff --git a/cli/testdata/server-config.yaml.golden b/cli/testdata/server-config.yaml.golden index dfc5bad74a8..b2dca8d3ae7 100644 --- a/cli/testdata/server-config.yaml.golden +++ b/cli/testdata/server-config.yaml.golden @@ -799,10 +799,6 @@ chat: # How many pending chats a worker should acquire per polling cycle. # (default: 10, type: int) acquireBatchSize: 10 - # Force MCP tool schemas behind find_tools regardless of size. The mcp-tool-search - # experiment must also be enabled. - # (default: false, type: bool) - mcpToolSearchForceDefer: false # Force chat debug logging on for every chat, bypassing the runtime admin and user # opt-in settings. # (default: false, type: bool) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index b633631bee5..27ed0964edd 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -17618,9 +17618,6 @@ const docTemplate = `{ }, "hook_url": { "$ref": "#/definitions/serpent.URL" - }, - "mcp_tool_search_force_defer": { - "type": "boolean" } } }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 16c8c7fcfaf..53df566fc82 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -15843,9 +15843,6 @@ }, "hook_url": { "$ref": "#/definitions/serpent.URL" - }, - "mcp_tool_search_force_defer": { - "type": "boolean" } } }, diff --git a/coderd/coderd.go b/coderd/coderd.go index f4c68b87d8f..c7effe22b59 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -933,7 +933,6 @@ func New(options *Options) *API { AIBridgeTransportFactory: &api.AIBridgeTransportFactory, AlwaysEnableDebugLogs: options.DeploymentValues.AI.Chat.DebugLoggingEnabled.Value(), Experiments: experiments, - ForceMCPToolSearch: options.DeploymentValues.AI.Chat.MCPToolSearchForceDefer.Value(), AgentConn: api.agentProvider.AgentConn, AgentInactiveDisconnectTimeout: api.AgentInactiveDisconnectTimeout, InstructionLookupTimeout: options.ChatdInstructionLookupTimeout, diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index c58a3ea5f61..c5df8e0a04b 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -200,7 +200,6 @@ type Server struct { aibridgeTransportFactory *atomic.Pointer[aibridge.TransportFactory] experiments codersdk.Experiments - forceMCPToolSearch bool // Configuration pendingChatAcquireInterval time.Duration @@ -3049,11 +3048,7 @@ type Config struct { Clock quartz.Clock AIBridgeTransportFactory *atomic.Pointer[aibridge.TransportFactory] Experiments codersdk.Experiments - // ForceMCPToolSearch ignores the schema-size threshold for development and tests. - // The mcp-tool-search experiment remains required. - ForceMCPToolSearch bool - - PrometheusRegistry prometheus.Registerer + PrometheusRegistry prometheus.Registerer AgentCapacityUnlock AgentCapacityUnlock @@ -3159,7 +3154,6 @@ func New(ps pubsub.Pubsub, cfg Config) *Server { }, aibridgeTransportFactory: cfg.AIBridgeTransportFactory, experiments: cfg.Experiments, - forceMCPToolSearch: cfg.ForceMCPToolSearch, pendingChatAcquireInterval: pendingChatAcquireInterval, maxChatsPerAcquire: maxChatsPerAcquire, inFlightChatStaleAfter: inFlightChatStaleAfter, diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index a5ca8fd1be5..03116099830 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -802,6 +802,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 +1076,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 +1239,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 +8497,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 @@ -10455,7 +10466,6 @@ func TestMCPToolSearchGenerationFlows(t *testing.T) { }) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) - cfg.ForceMCPToolSearch = true }) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -10538,7 +10548,6 @@ func TestMCPToolSearchGenerationFlows(t *testing.T) { }) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) - cfg.ForceMCPToolSearch = true }) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -10562,7 +10571,7 @@ func TestMCPToolSearchGenerationFlows(t *testing.T) { require.True(t, openAIMessagesContain(recorded[1].Messages, "echo: hello")) }) - t.Run("below threshold preserves wire tools", func(t *testing.T) { + t.Run("experiment gates deferral regardless of catalog size", func(t *testing.T) { t.Parallel() run := func(t *testing.T, experimentEnabled bool) []byte { @@ -10599,9 +10608,7 @@ func TestMCPToolSearchGenerationFlows(t *testing.T) { server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) if !experimentEnabled { - cfg.Experiments = slices.DeleteFunc(slices.Clone(cfg.Experiments), func(experiment codersdk.Experiment) bool { - return experiment == codersdk.ExperimentMCPToolSearch - }) + withoutMCPToolSearch(cfg) } }) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ @@ -10621,7 +10628,30 @@ func TestMCPToolSearchGenerationFlows(t *testing.T) { return append([]byte(nil), toolsJSON...) } - require.Equal(t, run(t, false), run(t, true)) + 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") }) } @@ -10717,6 +10747,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) @@ -10867,6 +10898,7 @@ func TestPlanModeRootChatApprovedExternalMCPToolInvocation(t *testing.T) { }) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + withoutMCPToolSearch(cfg) cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) }) @@ -10981,6 +11013,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) @@ -11188,6 +11221,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/metrics.go b/coderd/x/chatd/chatloop/metrics.go index 9e2193957d2..874a0ca52a0 100644 --- a/coderd/x/chatd/chatloop/metrics.go +++ b/coderd/x/chatd/chatloop/metrics.go @@ -42,7 +42,6 @@ type Metrics struct { FindToolsEmptyTotal prometheus.Counter FindToolsMatchCount prometheus.Histogram FindToolsActivationsTotal prometheus.Counter - DeferredMCPToolTokens *prometheus.HistogramVec } // NewMetrics creates a new Metrics instance registered with the @@ -139,13 +138,6 @@ func NewMetrics(reg prometheus.Registerer) *Metrics { Name: "find_tools_activations_total", Help: "Total deferred tool activations returned by find_tools.", }), - DeferredMCPToolTokens: factory.NewHistogramVec(prometheus.HistogramOpts{ - Namespace: metricsNamespace, - Subsystem: metricsSubsystem, - Name: "deferred_mcp_tool_tokens", - Help: "Estimated MCP tool schema tokens considered for deferral per generation.", - Buckets: prometheus.ExponentialBuckets(128, 2, 12), - }, []string{"provider", "model", "applied"}), StreamBufferDroppedTotal: factory.NewCounter(prometheus.CounterOpts{ Namespace: metricsNamespace, Subsystem: metricsSubsystem, diff --git a/coderd/x/chatd/forced_mcp_test.go b/coderd/x/chatd/forced_mcp_test.go index 2447b203203..7ef075ac563 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_preparer.go b/coderd/x/chatd/generation_preparer.go index 45221f483f0..c63cfbe2fe6 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -4,7 +4,6 @@ import ( "context" "encoding/json" "slices" - "strconv" "strings" "sync" @@ -605,18 +604,12 @@ func (server *Server) prepareGeneration( activeToolNames = allowedExploreToolNames(tools) } var allowInactiveTools map[string]bool - toolSearch := decideMCPToolSearch(mcpToolSearchInput{ + if decideMCPToolSearch(mcpToolSearchInput{ experimentEnabled: server.experiments.Enabled(codersdk.ExperimentMCPToolSearch), - forceDefer: server.forceMCPToolSearch, - contextWindow: modelConfig.ContextLimit, candidates: deferredCandidates, dynamicToolNames: dynamicToolNames, - }) - server.metrics.DeferredMCPToolTokens.WithLabelValues( - model.Provider(), model.ModelID(), strconv.FormatBool(toolSearch.apply), - ).Observe(toolSearch.estimatedTokens) - if toolSearch.apply { - activationTokenBudget := float64(modelConfig.ContextLimit) / mcpToolSearchThresholdDivisor + }) { + activationTokenBudget := float64(modelConfig.ContextLimit) / mcpToolSearchBudgetDivisor findTools := chattool.FindTools(chattool.FindToolsOptions{ Entries: deferredMCPToolEntries(deferredCandidates), SchemaTokenBudget: activationTokenBudget, diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index fc757796a9d..e7acdd35ca4 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -15,7 +15,10 @@ import ( "github.com/coder/coder/v2/codersdk" ) -const mcpToolSearchThresholdDivisor = 10 +// 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 @@ -107,38 +110,31 @@ func workspaceMCPServerName(tool fantasy.AgentTool) string { return "" } -type mcpToolSearchDecision struct { - apply bool - estimatedTokens float64 -} - type mcpToolSearchInput struct { experimentEnabled bool - forceDefer bool - contextWindow int64 candidates []deferredMCPTool dynamicToolNames map[string]bool } -func decideMCPToolSearch(input mcpToolSearchInput) mcpToolSearchDecision { - decision := mcpToolSearchDecision{estimatedTokens: estimateDeferredMCPToolTokens(input.candidates)} +// 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 decision + 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 decision + return false } for _, candidate := range input.candidates { if candidate.tool.Info().Name == chattool.FindToolsName { - return decision + return false } } - decision.apply = input.forceDefer || - (input.contextWindow > 0 && decision.estimatedTokens > float64(input.contextWindow)/mcpToolSearchThresholdDivisor) - return decision + return true } func configureDeferredMCPToolSearch( @@ -274,7 +270,6 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe } parsedParts := make([][]codersdk.ChatMessagePart, len(rows)) findToolsCallIDs := make(map[string]struct{}) - erroredCallIDs := make(map[string]struct{}) for i := range rows { parts, err := chatprompt.ParseContent(rows[i]) if err != nil { @@ -285,17 +280,37 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName == chattool.FindToolsName && part.ToolCallID != "" { findToolsCallIDs[part.ToolCallID] = struct{}{} } - if part.Type == codersdk.ChatMessagePartTypeToolResult && part.IsError && part.ToolCallID != "" { - erroredCallIDs[part.ToolCallID] = struct{}{} + } + } + // Providers may reuse a tool-call ID in a later step, so each call + // pairs with its own result: walking oldest first, a result settles + // the oldest still-unpaired call with its ID. A call whose result + // was compacted away stays unpaired and counts as successful. + type callRef struct{ row, part int } + callErrored := make(map[callRef]bool) + pendingByID := make(map[string][]callRef) + for i := range rows { + for j, part := range parsedParts[i] { + if part.ToolCallID == "" { + continue + } + switch part.Type { + case codersdk.ChatMessagePartTypeToolCall: + pendingByID[part.ToolCallID] = append(pendingByID[part.ToolCallID], callRef{row: i, part: j}) + case codersdk.ChatMessagePartTypeToolResult: + if refs := pendingByID[part.ToolCallID]; len(refs) > 0 { + callErrored[refs[0]] = part.IsError + pendingByID[part.ToolCallID] = refs[1:] + } } } } pendingSearch := make(map[string][]string) var erroredNames []string for i := len(rows) - 1; i >= 0; i-- { - for _, part := range parsedParts[i] { + for j, part := range parsedParts[i] { if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName != chattool.FindToolsName { - if _, errored := erroredCallIDs[part.ToolCallID]; errored { + if callErrored[callRef{row: i, part: j}] { erroredNames = append(erroredNames, part.ToolName) continue } diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index a424c45090f..66650ff312e 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -38,37 +38,30 @@ func testDeferredTool(name, description string, parameters map[string]any) defer func TestDecideMCPToolSearch(t *testing.T) { t.Parallel() - small := []deferredMCPTool{testDeferredTool("server__small", "small", map[string]any{"value": map[string]any{"type": "string"}})} - large := []deferredMCPTool{testDeferredTool("server__large", strings.Repeat("large ", 2000), map[string]any{"value": map[string]any{"type": "string"}})} + candidates := []deferredMCPTool{testDeferredTool("server__small", "small", map[string]any{"value": map[string]any{"type": "string"}})} tests := []struct { name string experiment bool - force bool - window int64 candidates []deferredMCPTool dynamicNames map[string]bool want bool }{ - {name: "below", experiment: true, window: 100_000, candidates: small}, - {name: "above", experiment: true, window: 10_000, candidates: large, want: true}, - {name: "forced", experiment: true, force: true, window: 100_000, candidates: small, want: true}, - {name: "experiment off", force: true, window: 10, candidates: large}, - {name: "empty", experiment: true, force: true}, - {name: "collision", experiment: true, force: true, candidates: []deferredMCPTool{testDeferredTool(chattool.FindToolsName, "collision", nil)}}, - {name: "dynamic collision", experiment: true, force: true, candidates: small, dynamicNames: map[string]bool{chattool.FindToolsName: true}}, - {name: "dynamic no collision", experiment: true, force: true, candidates: small, dynamicNames: map[string]bool{"other": true}, want: true}, + {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, - forceDefer: tt.force, - contextWindow: tt.window, candidates: tt.candidates, dynamicToolNames: tt.dynamicNames, - }).apply) + })) }) } } @@ -176,6 +169,39 @@ func TestDeriveDeferredMCPActivationsErroredDirectCallsActivateLast(t *testing.T "the newest errored call keeps the first-activation allowance when nothing else activates") } +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") +} + func TestFlattenMCPParameterText(t *testing.T) { t.Parallel() text := flattenMCPParameterText(map[string]any{ @@ -351,7 +377,7 @@ func TestConfigureDeferredMCPToolSearchDirectCallAndCompaction(t *testing.T) { require.Empty(t, deriveDeferredMCPActivations(nil, candidates, 0)) } -func TestMCPToolSearchBelowThresholdPreservesWireTools(t *testing.T) { +func TestMCPToolSearchExperimentDisabledPreservesWireTools(t *testing.T) { t.Parallel() hot := deferredTestAgentTool{info: fantasy.ToolInfo{Name: "read_file"}} @@ -360,12 +386,9 @@ func TestMCPToolSearchBelowThresholdPreservesWireTools(t *testing.T) { active := []string{"read_file", candidate.tool.Info().Name} withoutExperiment := captureWireToolNames(t, tools, active) - decision := decideMCPToolSearch(mcpToolSearchInput{ - experimentEnabled: true, - contextWindow: 100_000, - candidates: []deferredMCPTool{candidate}, - }) - require.False(t, decision.apply) + require.False(t, decideMCPToolSearch(mcpToolSearchInput{ + candidates: []deferredMCPTool{candidate}, + })) require.Equal(t, withoutExperiment, captureWireToolNames(t, tools, active)) } diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 90588148884..18174128299 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -4341,17 +4341,6 @@ Write out the current server config as YAML to stdout.`, YAML: "acquireBatchSize", Hidden: true, // Hidden because most operators should not need to modify this. }, - { - Name: "Chat: MCP Tool Search Force Defer", - Description: "Force MCP tool schemas behind find_tools regardless of size. The mcp-tool-search experiment must also be enabled.", - Flag: "chat-mcp-tool-search-force-defer", - Env: "CODER_CHAT_MCP_TOOL_SEARCH_FORCE_DEFER", - Value: &c.AI.Chat.MCPToolSearchForceDefer, - Default: "false", - Group: &deploymentGroupChat, - YAML: "mcpToolSearchForceDefer", - Hidden: true, - }, { Name: "Chat: Debug Logging Enabled", Description: "Force chat debug logging on for every chat, bypassing the runtime admin and user opt-in settings.", @@ -5100,14 +5089,13 @@ type AIBridgeProxyConfig struct { } type ChatConfig struct { - AcquireBatchSize serpent.Int64 `json:"acquire_batch_size" typescript:",notnull"` - DebugLoggingEnabled serpent.Bool `json:"debug_logging_enabled" typescript:",notnull"` - MCPToolSearchForceDefer serpent.Bool `json:"mcp_tool_search_force_defer" typescript:",notnull"` - HookURL serpent.URL `json:"hook_url" typescript:",notnull"` - HookSecret serpent.String `json:"hook_secret" typescript:",notnull"` - HookTimeout serpent.Duration `json:"hook_timeout" typescript:",notnull"` - HookEnabled serpent.Bool `json:"hook_enabled" typescript:",notnull"` - HookAllowInsecure serpent.Bool `json:"hook_allow_insecure" typescript:",notnull"` + AcquireBatchSize serpent.Int64 `json:"acquire_batch_size" typescript:",notnull"` + DebugLoggingEnabled serpent.Bool `json:"debug_logging_enabled" typescript:",notnull"` + HookURL serpent.URL `json:"hook_url" typescript:",notnull"` + HookSecret serpent.String `json:"hook_secret" typescript:",notnull"` + HookTimeout serpent.Duration `json:"hook_timeout" typescript:",notnull"` + HookEnabled serpent.Bool `json:"hook_enabled" typescript:",notnull"` + HookAllowInsecure serpent.Bool `json:"hook_allow_insecure" typescript:",notnull"` // Deprecated: AI Gateway routing is now the only routing path. Setting this // value has no effect. This option will be removed in a future release. AIGatewayRoutingEnabled serpent.Bool `json:"ai_gateway_routing_enabled" typescript:",notnull" swaggerignore:"true"` diff --git a/docs/admin/integrations/prometheus.md b/docs/admin/integrations/prometheus.md index d240e9ec697..fe092ff5c88 100644 --- a/docs/admin/integrations/prometheus.md +++ b/docs/admin/integrations/prometheus.md @@ -240,7 +240,6 @@ 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_deferred_mcp_tool_tokens` | histogram | Estimated MCP tool schema tokens considered for deferral per generation. | `applied` `model` `provider` | | `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. | | diff --git a/docs/reference/api/general.md b/docs/reference/api/general.md index 3ea6e0e40d6..982b5064099 100644 --- a/docs/reference/api/general.md +++ b/docs/reference/api/general.md @@ -254,8 +254,7 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \ "rawQuery": "string", "scheme": "string", "user": {} - }, - "mcp_tool_search_force_defer": true + } } }, "allow_workspace_renames": true, diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 27e87fac918..e5ce947f4e7 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -1103,8 +1103,7 @@ title: Schemas "rawQuery": "string", "scheme": "string", "user": {} - }, - "mcp_tool_search_force_defer": true + } } } ``` @@ -2533,23 +2532,21 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "rawQuery": "string", "scheme": "string", "user": {} - }, - "mcp_tool_search_force_defer": true + } } ``` ### Properties -| Name | Type | Required | Restrictions | Description | -|-------------------------------|----------------------------|----------|--------------|-------------| -| `acquire_batch_size` | integer | false | | | -| `debug_logging_enabled` | boolean | false | | | -| `hook_allow_insecure` | boolean | false | | | -| `hook_enabled` | boolean | false | | | -| `hook_secret` | string | false | | | -| `hook_timeout` | integer | false | | | -| `hook_url` | [serpent.URL](#serpenturl) | false | | | -| `mcp_tool_search_force_defer` | boolean | false | | | +| Name | Type | Required | Restrictions | Description | +|-------------------------|----------------------------|----------|--------------|-------------| +| `acquire_batch_size` | integer | false | | | +| `debug_logging_enabled` | boolean | false | | | +| `hook_allow_insecure` | boolean | false | | | +| `hook_enabled` | boolean | false | | | +| `hook_secret` | string | false | | | +| `hook_timeout` | integer | false | | | +| `hook_url` | [serpent.URL](#serpenturl) | false | | | ## codersdk.ChatContext @@ -6006,8 +6003,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "rawQuery": "string", "scheme": "string", "user": {} - }, - "mcp_tool_search_force_defer": true + } } }, "allow_workspace_renames": true, @@ -6635,8 +6631,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "rawQuery": "string", "scheme": "string", "user": {} - }, - "mcp_tool_search_force_defer": true + } } }, "allow_workspace_renames": true, diff --git a/scripts/metricsdocgen/generated_metrics b/scripts/metricsdocgen/generated_metrics index fe9a8ad331f..bb52806447d 100644 --- a/scripts/metricsdocgen/generated_metrics +++ b/scripts/metricsdocgen/generated_metrics @@ -283,9 +283,6 @@ 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_deferred_mcp_tool_tokens Estimated MCP tool schema tokens considered for deferral per generation. -# TYPE coderd_chatd_deferred_mcp_tool_tokens histogram -coderd_chatd_deferred_mcp_tool_tokens{provider="",model="",applied=""} 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 diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 456f4d75f7e..e532692338a 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -2073,7 +2073,6 @@ export const ChatComputerUseProviders: ChatComputerUseProvider[] = [ export interface ChatConfig { readonly acquire_batch_size: number; readonly debug_logging_enabled: boolean; - readonly mcp_tool_search_force_defer: boolean; readonly hook_url: string; readonly hook_secret: string; readonly hook_timeout: number; From 1210f0ee9ae20e7c147ae6131af53a70ef18468c Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 09:23:56 +0000 Subject: [PATCH 38/48] fix(coderd/x/chatd): pair reused tool-call IDs with their newest call A result now settles the newest unpaired call with its ID, so an older call whose result was lost cannot adopt a later call's error. Also folds the duplicated find_tools array parsers in Tool.tsx into one shared parseArray helper. --- coderd/x/chatd/mcp_tool_search.go | 19 ++--- .../x/chatd/mcp_tool_search_internal_test.go | 14 ++++ .../components/ChatElements/tools/Tool.tsx | 69 ++++++++----------- 3 files changed, 54 insertions(+), 48 deletions(-) diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index e7acdd35ca4..c9e33599a39 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -282,13 +282,14 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe } } } - // Providers may reuse a tool-call ID in a later step, so each call - // pairs with its own result: walking oldest first, a result settles - // the oldest still-unpaired call with its ID. A call whose result - // was compacted away stays unpaired and counts as successful. + // 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. type callRef struct{ row, part int } callErrored := make(map[callRef]bool) - pendingByID := make(map[string][]callRef) + pendingByID := make(map[string]callRef) for i := range rows { for j, part := range parsedParts[i] { if part.ToolCallID == "" { @@ -296,11 +297,11 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe } switch part.Type { case codersdk.ChatMessagePartTypeToolCall: - pendingByID[part.ToolCallID] = append(pendingByID[part.ToolCallID], callRef{row: i, part: j}) + pendingByID[part.ToolCallID] = callRef{row: i, part: j} case codersdk.ChatMessagePartTypeToolResult: - if refs := pendingByID[part.ToolCallID]; len(refs) > 0 { - callErrored[refs[0]] = part.IsError - pendingByID[part.ToolCallID] = refs[1:] + if ref, ok := pendingByID[part.ToolCallID]; ok { + callErrored[ref] = part.IsError + delete(pendingByID, part.ToolCallID) } } } diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index 66650ff312e..cb297502a1c 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -200,6 +200,20 @@ func TestDeriveDeferredMCPActivationsReusedCallIDs(t *testing.T) { 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") } func TestFlattenMCPParameterText(t *testing.T) { diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx index 6a69475feb7..c8e4bfea32d 100644 --- a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx @@ -986,58 +986,49 @@ const GenericToolRenderer: FC = ({ ); }; -const parseStringList = (value: unknown): string[] | null => { - if (typeof value === "string") { +// parseArray narrows a value (or its JSON-encoded string form) to an +// array and maps every element with parseItem, returning null when the +// value is not an array or any element fails to parse. +const parseArray = ( + value: unknown, + parseItem: (item: unknown) => T | null, +): T[] | null => { + let array = value; + if (typeof array === "string") { try { - return parseStringList(JSON.parse(value)); + array = JSON.parse(array); } catch { return null; } } - if (!Array.isArray(value)) { + if (!Array.isArray(array)) { return null; } - const strings: string[] = []; - for (const item of value) { - if (typeof item !== "string") { + const items: T[] = []; + for (const item of array) { + const parsed = parseItem(item); + if (parsed === null) { return null; } - const trimmed = item.trim(); - if (trimmed) { - strings.push(trimmed); - } + items.push(parsed); } - return strings; + return items; }; -const parseFindToolsMatches = (value: unknown): FindToolsMatch[] | null => { - if (typeof value === "string") { - try { - return parseFindToolsMatches(JSON.parse(value)); - } catch { - return null; - } - } - if (!Array.isArray(value)) { - return null; - } - const matches: FindToolsMatch[] = []; - for (const item of value) { +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); - if ( - !record || - typeof record.name !== "string" || - typeof record.description !== "string" - ) { - return null; - } - matches.push({ - name: record.name, - description: record.description, - }); - } - return matches; -}; + 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); From fb2f2a8426529bc052d2555a30e1759004e324cc Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 09:52:43 +0000 Subject: [PATCH 39/48] fix(coderd/x/chatd): admit orphan find_tools results at their own row under reused IDs --- coderd/x/chatd/mcp_tool_search.go | 27 +++++++++---------- .../x/chatd/mcp_tool_search_internal_test.go | 12 +++++++++ 2 files changed, 25 insertions(+), 14 deletions(-) diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index c9e33599a39..95d01aa1bce 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -269,27 +269,25 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe activated = append(activated, name) } parsedParts := make([][]codersdk.ChatMessagePart, len(rows)) - findToolsCallIDs := make(map[string]struct{}) for i := range rows { parts, err := chatprompt.ParseContent(rows[i]) if err != nil { continue } parsedParts[i] = parts - for _, part := range parts { - if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName == chattool.FindToolsName && part.ToolCallID != "" { - findToolsCallIDs[part.ToolCallID] = struct{}{} - } - } } // 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. - type callRef struct{ row, part int } - callErrored := make(map[callRef]bool) - pendingByID := make(map[string]callRef) + // 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 == "" { @@ -297,10 +295,11 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe } switch part.Type { case codersdk.ChatMessagePartTypeToolCall: - pendingByID[part.ToolCallID] = callRef{row: i, part: j} + 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) } } @@ -311,21 +310,21 @@ func deriveDeferredMCPActivations(rows []database.ChatMessage, candidates []defe 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[callRef{row: i, part: j}] { + if callErrored[partRef{row: i, part: j}] { erroredNames = append(erroredNames, part.ToolName) continue } appendName(part.ToolName) } } - for _, part := range parsedParts[i] { + 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 := findToolsCallIDs[part.ToolCallID]; paired { + if _, paired := resultPaired[partRef{row: i, part: j}]; paired { pendingSearch[part.ToolCallID] = result.Activated continue } diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index cb297502a1c..498830bc738 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -214,6 +214,18 @@ func TestDeriveDeferredMCPActivationsReusedCallIDs(t *testing.T) { 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) { From 0bca0480cafaafc62d4c598cd73f4477c4e89b6e Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:03:34 +0000 Subject: [PATCH 40/48] fix(coderd/x/chatd/chattool): count find_tools calls rejected during argument decoding --- coderd/x/chatd/chattool/findtools.go | 24 ++++++++++++++++++- .../chatd/chattool/findtools_internal_test.go | 7 +++++- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 5216195b2c8..5096bca43e7 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -2,6 +2,7 @@ package chattool import ( "context" + "encoding/json" "fmt" "regexp" "slices" @@ -98,6 +99,7 @@ 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 } @@ -108,6 +110,18 @@ 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) @@ -191,7 +205,15 @@ func FindTools(options FindToolsOptions) fantasy.AgentTool { defer budgetMu.Unlock() rebuild(names, errored) } - return findToolsTool{reserveStepCalls: reserve, settleStepResults: settle, AgentTool: fantasy.NewAgentTool( + 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) { diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 19a002aba94..0f61570a332 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -268,13 +268,18 @@ func TestFindToolsDirectCallReservation(t *testing.T) { resp, err = tool.Run(context.Background(), fantasy.ToolCall{Input: `{}`}) require.NoError(t, err) require.True(t, resp.IsError) - require.Len(t, calls, 3, "rejected calls count toward call totals") + 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) { From 6d73c4a58380ae2df837119cd5c1e790f242b118 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:15:11 +0000 Subject: [PATCH 41/48] fix(coderd/x/chatd): label empty server identities at catalog construction --- coderd/x/chatd/chattool/findtools.go | 7 ++--- coderd/x/chatd/mcp_tool_search.go | 26 +++++++++++++++++- .../x/chatd/mcp_tool_search_internal_test.go | 27 +++++++++++++++++++ 3 files changed, 56 insertions(+), 4 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 5096bca43e7..1e5546ca044 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -613,10 +613,11 @@ type findToolsGroup struct { 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 - if server == "" { - server = "workspace" - } group := grouped[server] if group == nil { group = &findToolsGroup{server: server, description: entry.ServerDescription} diff --git a/coderd/x/chatd/mcp_tool_search.go b/coderd/x/chatd/mcp_tool_search.go index 95d01aa1bce..a62906ffb1f 100644 --- a/coderd/x/chatd/mcp_tool_search.go +++ b/coderd/x/chatd/mcp_tool_search.go @@ -3,6 +3,7 @@ package chatd import ( "encoding/json" "slices" + "strconv" "strings" "charm.land/fantasy" @@ -179,13 +180,36 @@ func estimateDeferredMCPToolTokens(candidates []deferredMCPTool) float64 { } 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: candidate.server, + Server: server, ServerDescription: candidate.serverDescription, ParameterText: flattenMCPParameterText(info.Parameters), SchemaTokens: estimateDeferredMCPToolTokens([]deferredMCPTool{candidate}), diff --git a/coderd/x/chatd/mcp_tool_search_internal_test.go b/coderd/x/chatd/mcp_tool_search_internal_test.go index 498830bc738..7c2834290b4 100644 --- a/coderd/x/chatd/mcp_tool_search_internal_test.go +++ b/coderd/x/chatd/mcp_tool_search_internal_test.go @@ -169,6 +169,33 @@ func TestDeriveDeferredMCPActivationsErroredDirectCallsActivateLast(t *testing.T "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{ From bac9da4ab2f765462985c3ae96cbd008446e3ba1 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:26:02 +0000 Subject: [PATCH 42/48] fix(coderd/x/chatd/chattool): bound find_tools query scoring work --- coderd/x/chatd/chattool/findtools.go | 78 +++++++++++++++---- .../chatd/chattool/findtools_internal_test.go | 25 ++++++ 2 files changed, 87 insertions(+), 16 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 1e5546ca044..12d083af36e 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -18,6 +18,11 @@ 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 // 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. @@ -344,7 +349,11 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, budget Sear byName[entry.Name] = entry } - queries := parseFindToolsQueries(entries, args.Queries) + queryArgs := args.Queries + if len(queryArgs) > findToolsMaxQueries { + queryArgs = queryArgs[:findToolsMaxQueries] + } + queries := parseFindToolsQueries(entries, queryArgs) type scoredEntry struct { entry FindToolCatalogEntry @@ -352,6 +361,7 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, budget Sear } scored := make([]scoredEntry, 0, len(entries)) for _, entry := range entries { + tokens := tokenizeFindToolsEntry(entry) score := 0 for _, query := range queries { if query.server != "" { @@ -367,7 +377,7 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, budget Sear continue } for _, token := range query.tokens { - score += scoreFindToolToken(entry, token) + score += tokens.score(token) } } if score > 0 { @@ -489,13 +499,13 @@ func parseFindToolsQueries(entries []FindToolCatalogEntry, queries []string) []s if !ok { continue } - parsed = append(parsed, scopedFindToolsQuery{server: server, exact: pass.exact, tokens: tokenizeFindTools(rest)}) + parsed = append(parsed, scopedFindToolsQuery{server: server, exact: pass.exact, tokens: tokenizeFindToolsQuery(rest)}) scoped = true break } } if !scoped { - parsed = append(parsed, scopedFindToolsQuery{tokens: tokenizeFindTools(query)}) + parsed = append(parsed, scopedFindToolsQuery{tokens: tokenizeFindToolsQuery(query)}) } } return parsed @@ -506,26 +516,62 @@ func tokenizeFindTools(value string) []string { return slices.DeleteFunc(parts, func(part string) bool { return part == "" }) } -func scoreFindToolToken(entry FindToolCatalogEntry, token string) int { - name := strings.ToLower(entry.Name) - nameTokens := tokenizeFindTools(name) +// 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 slices.Contains(nameTokens, token) { + if _, ok := t.nameTokens[token]; ok { score += 8 - } else if strings.Contains(name, token) { + } else if strings.Contains(t.name, token) { score += 5 } - if slices.Contains(tokenizeFindTools(entry.Description), token) { + if _, ok := t.description[token]; ok { score += 2 } - if slices.Contains(tokenizeFindTools(entry.ParameterText), token) { + if _, ok := t.parameters[token]; ok { score++ } - // 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. - if slices.Contains(tokenizeFindTools(entry.Server), token) || - slices.Contains(tokenizeFindTools(entry.ServerDescription), token) { + if _, ok := t.server[token]; ok { score++ } return score diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 0f61570a332..bd0f9a7d0ce 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -123,6 +123,31 @@ func TestSearchTools(t *testing.T) { require.Len(t, result.Activated, 2, "a prefix matching no exact-case name falls back to spanning the case-colliding servers") }) + 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{ From d87988f547a0ef84307c15a56957a2bf5b5e24b4 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:37:19 +0000 Subject: [PATCH 43/48] fix(coderd/x/chatd): count find_tools calls at the execution choke point --- coderd/x/chatd/chatd_test.go | 56 +++++++++++++++++++++++++++ coderd/x/chatd/generation.go | 12 ++++++ coderd/x/chatd/generation_preparer.go | 4 +- 3 files changed, 71 insertions(+), 1 deletion(-) diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 03116099830..f3b298bcff6 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -10571,6 +10571,62 @@ func TestMCPToolSearchGenerationFlows(t *testing.T) { 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("experiment gates deferral regardless of catalog size", func(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 2050195cf2a..84661849728 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" @@ -811,6 +812,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) diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index c63cfbe2fe6..f9085157de8 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -614,8 +614,10 @@ func (server *Server) prepareGeneration( 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) { - server.metrics.FindToolsCallsTotal.Inc() if call.Rejection == "" { server.metrics.FindToolsMatchCount.Observe(float64(call.MatchCount)) server.metrics.FindToolsActivationsTotal.Add(float64(len(call.Activated))) From eda35d205c2c19ab1a2d1d84652a0b119e9afc55 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:55:08 +0000 Subject: [PATCH 44/48] fix(coderd/x/chatd): count hook-denied find_tools calls and cover error and empty render states --- coderd/x/chatd/chatd_test.go | 59 ++++++++++ coderd/x/chatd/generation.go | 10 ++ .../ConversationTimeline.stories.tsx | 103 ++++++++++++++++++ .../components/ChatElements/tools/Tool.tsx | 24 +++- 4 files changed, 191 insertions(+), 5 deletions(-) diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index f3b298bcff6..7c6f86a5ae5 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" @@ -10627,6 +10628,64 @@ func TestMCPToolSearchGenerationFlows(t *testing.T) { 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("experiment gates deferral regardless of catalog size", func(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 84661849728..e921d230960 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -797,6 +797,16 @@ func (s *taskStarter) admitStepToolCalls( 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 } diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx index 7634817f4b9..abde9acd32d 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx @@ -641,6 +641,109 @@ export const FindToolsSearchResult: Story = { }, }; +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(); + // A valid empty result has no match list to expand and no error + // indicator. + 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(); + // The failure indicator carries the persisted error message. + 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, diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx index c8e4bfea32d..0cf85e0fa68 100644 --- a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx @@ -1043,6 +1043,24 @@ const FindToolsRenderer: FC = (props) => { 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; @@ -1051,17 +1069,13 @@ const FindToolsRenderer: FC = (props) => { return ; } - const errorMessage = parsedResult - ? asString(parsedResult.error || parsedResult.message) - : ""; return ( ); }; From 2098335ad6e7e1617c24e0359aa6f81474fb6574 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:03:06 +0000 Subject: [PATCH 45/48] refactor(site/src/pages/AgentsPage): drop assertion-narrating story comments --- .../ChatConversation/ConversationTimeline.stories.tsx | 3 --- 1 file changed, 3 deletions(-) diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx index abde9acd32d..8dfad522e1b 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx @@ -683,8 +683,6 @@ export const FindToolsEmptyResult: Story = { "Searched tools: nonexistent capability -> 0 matched", ); expect(summary).toBeVisible(); - // A valid empty result has no match list to expand and no error - // indicator. expect( canvas.queryByRole("button", { name: "Searched tools: nonexistent capability -> 0 matched", @@ -735,7 +733,6 @@ export const FindToolsErrorResult: Story = { expect( canvas.getByText("Searched tools: github issues -> 0 matched"), ).toBeVisible(); - // The failure indicator carries the persisted error message. expect( canvas.getByRole("img", { name: "The schema budget for this step is exhausted; call the tools already activated or retry next step.", From 3a26aec26e76e9800faee8205a6fe381254285aa Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:15:00 +0000 Subject: [PATCH 46/48] fix(coderd/x/chatd/chattool): cut folded scope prefixes at rune boundaries --- coderd/x/chatd/chattool/findtools.go | 35 +++++++++++++++++-- .../chatd/chattool/findtools_internal_test.go | 21 +++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 12d083af36e..9b08519c3d3 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -9,6 +9,7 @@ import ( "strconv" "strings" "sync" + "unicode" "unicode/utf8" "charm.land/fantasy" @@ -490,10 +491,11 @@ func parseFindToolsQueries(entries []FindToolCatalogEntry, queries []string) []s continue } } else { - if len(pass.text) < len(server) || !strings.EqualFold(pass.text[:len(server)], server) { + var ok bool + rest, ok = cutPrefixFold(pass.text, server) + if !ok { continue } - rest = pass.text[len(server):] } rest, ok := strings.CutPrefix(strings.TrimLeft(rest, " "), ":") if !ok { @@ -511,6 +513,35 @@ func parseFindToolsQueries(entries []FindToolCatalogEntry, queries []string) []s 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 == "" }) diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index bd0f9a7d0ce..56e4ba91b6a 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -123,6 +123,27 @@ func TestSearchTools(t *testing.T) { 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{ From 065fdb52329e3512c8ae6b7341c5aa22548db5bb Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:25:22 +0000 Subject: [PATCH 47/48] fix: bound the find_tools names list and drop a narrating comment --- coderd/x/chatd/chattool/findtools.go | 9 ++++++++- .../chatd/chattool/findtools_internal_test.go | 18 ++++++++++++++++++ .../components/ChatElements/tools/Tool.tsx | 3 --- 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/coderd/x/chatd/chattool/findtools.go b/coderd/x/chatd/chattool/findtools.go index 9b08519c3d3..8338df54787 100644 --- a/coderd/x/chatd/chattool/findtools.go +++ b/coderd/x/chatd/chattool/findtools.go @@ -24,6 +24,9 @@ const ( // 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. @@ -414,7 +417,11 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, budget Sear }) activatedSet[entry.Name] = struct{}{} } - for _, name := range args.Names { + nameArgs := args.Names + if len(nameArgs) > findToolsMaxNames { + nameArgs = nameArgs[:findToolsMaxNames] + } + for _, name := range nameArgs { if entry, ok := byName[name]; ok { appendMatch(entry) } diff --git a/coderd/x/chatd/chattool/findtools_internal_test.go b/coderd/x/chatd/chattool/findtools_internal_test.go index 56e4ba91b6a..0a7af3fdc47 100644 --- a/coderd/x/chatd/chattool/findtools_internal_test.go +++ b/coderd/x/chatd/chattool/findtools_internal_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "slices" "strings" "testing" @@ -69,6 +70,23 @@ func TestSearchTools(t *testing.T) { 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{ diff --git a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx index 0cf85e0fa68..3d553fa5706 100644 --- a/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx +++ b/site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx @@ -986,9 +986,6 @@ const GenericToolRenderer: FC = ({ ); }; -// parseArray narrows a value (or its JSON-encoded string form) to an -// array and maps every element with parseItem, returning null when the -// value is not an array or any element fails to parse. const parseArray = ( value: unknown, parseItem: (item: unknown) => T | null, From af9fc38d8132694163190934dc256ef712bf1580 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Tue, 18 Aug 2026 11:43:51 +0000 Subject: [PATCH 48/48] fix(coderd/x/chatd): count find_tools calls discarded by admission errors --- coderd/x/chatd/chatd_test.go | 61 ++++++++++++++++++++++++++++++++++++ coderd/x/chatd/generation.go | 16 ++++++++++ 2 files changed, 77 insertions(+) diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 7c6f86a5ae5..11865529fb1 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -10686,6 +10686,67 @@ func TestMCPToolSearchGenerationFlows(t *testing.T) { 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() diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index e921d230960..8b1bb881458 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -783,17 +783,33 @@ 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...)