Thanks to visit codestin.com
Credit goes to github.com

Skip to content
70 changes: 45 additions & 25 deletions coderd/x/chatd/chattool/findtools.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,11 @@ import (
)

const (
FindToolsName = "find_tools"
findToolsMaxMatches = 20
findToolsCatalogTokens = 4000
FindToolsName = "find_tools"
// Keep broad query results concise while allowing higher explicit limits.
findToolsDefaultMatches = 10
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.
Expand Down Expand Up @@ -83,8 +85,9 @@ type FindToolsOptions struct {
}

type FindToolsArgs struct {
Queries []string `json:"queries,omitempty"`
Names []string `json:"names,omitempty"`
Queries []string `json:"queries,omitempty" description:"Task or capability keywords, matched against tool names, descriptions, parameters, and server metadata. Prefer a few specific keywords over sentences."`
Names []string `json:"names,omitempty" description:"Exact cataloged tool names to activate directly."`
Limit int `json:"limit,omitempty" description:"Cap on total tools returned and activated per call (default 10, max 20). Exact names are always included and may exceed it."`
}

type FindToolsMatch struct {
Expand Down Expand Up @@ -338,15 +341,11 @@ type SearchBudget struct {
AllowFirstOverBudget bool
}

// SearchTools includes exact name activations first, then fills the
// remaining match slots with the top-scored keyword matches. The shared
// cap and summary-length descriptions keep the persisted result small
// enough that generic tool-result truncation can never corrupt the
// activation JSON that later steps re-derive activations from. A
// positive budget additionally skips matches whose schema weight would
// push the aggregate past it, admitting later matches that still fit.
// The second result counts matches skipped for budget, so callers can
// tell an exhausted budget from no matches.
// SearchTools prioritizes exact names, then keyword matches ranked by
// distinct query terms matched before raw score. Exact names bypass the
// per-call limit, but a hard cap keeps the persisted result safe from
// generic tool-result truncation; the second return counts
// budget-skipped matches.
func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, budget SearchBudget) (FindToolsResult, int) {
byName := make(map[string]FindToolCatalogEntry, len(entries))
for _, entry := range entries {
Expand All @@ -360,13 +359,19 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, budget Sear
queries := parseFindToolsQueries(entries, queryArgs)

type scoredEntry struct {
entry FindToolCatalogEntry
score int
entry FindToolCatalogEntry
coverage int
score int
}
entryTokens := make([]findToolsEntryTokens, len(entries))
for i, entry := range entries {
entryTokens[i] = tokenizeFindToolsEntry(entry)
}
scored := make([]scoredEntry, 0, len(entries))
for _, entry := range entries {
tokens := tokenizeFindToolsEntry(entry)
for i, entry := range entries {
tokens := entryTokens[i]
score := 0
matched := make(map[string]struct{})
for _, query := range queries {
if query.server != "" {
if query.exact && entry.Server != query.server {
Expand All @@ -378,31 +383,46 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, budget Sear
}
if query.server != "" && len(query.tokens) == 0 {
score++
// A scope-only hit counts as one covered term; ":" cannot occur in tokens.
matched[":"+query.server] = struct{}{}
continue
}
for _, token := range query.tokens {
score += tokens.score(token)
if tokenScore := tokens.score(token); tokenScore > 0 {
score += tokenScore
matched[token] = struct{}{}
}
}
}
if score > 0 {
scored = append(scored, scoredEntry{entry: entry, score: score})
scored = append(scored, scoredEntry{entry: entry, coverage: len(matched), score: score})
}
}
slices.SortFunc(scored, func(a, b scoredEntry) int {
if a.coverage != b.coverage {
return b.coverage - a.coverage
}
if a.score != b.score {
return b.score - a.score
}
return strings.Compare(a.entry.Name, b.entry.Name)
})
matchLimit := args.Limit
if matchLimit <= 0 {
matchLimit = findToolsDefaultMatches
}
if matchLimit > findToolsMaxMatches {
matchLimit = findToolsMaxMatches
}
matches := make([]FindToolsMatch, 0, findToolsMaxMatches)
activatedSet := make(map[string]struct{}, findToolsMaxMatches)
usedSchemaTokens := 0.0
budgetSkipped := 0
appendMatch := func(entry FindToolCatalogEntry) {
appendMatch := func(entry FindToolCatalogEntry, limit int) {
if _, exists := activatedSet[entry.Name]; exists {
return
}
if len(matches) >= findToolsMaxMatches {
if len(matches) >= limit {
return
}
overBudget := budget.SchemaTokens > 0 && usedSchemaTokens+entry.SchemaTokens > budget.SchemaTokens
Expand All @@ -423,11 +443,11 @@ func SearchTools(entries []FindToolCatalogEntry, args FindToolsArgs, budget Sear
}
for _, name := range nameArgs {
if entry, ok := byName[name]; ok {
appendMatch(entry)
appendMatch(entry, findToolsMaxMatches)
}
}
for _, item := range scored {
appendMatch(item.entry)
appendMatch(item.entry, matchLimit)
}
activated := make([]string, 0, len(activatedSet))
for name := range activatedSet {
Expand Down Expand Up @@ -616,7 +636,7 @@ func (t findToolsEntryTokens) score(token string) int {
}

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"
const usage = "The MCP tools cataloged below are deferred: not in your tool list until activated. Search by keyword, activate exact tool names, or scope with a \"server: terms\" prefix; matches activate and become callable on the next step. Direct calls to cataloged tools also work, but search first for unfamiliar tools. limit caps total results per call (default 10, max 20); exact names bypass it but still spend the shared schema budget. Narrow the query or raise limit for more.\n\n"
budget := float64(findToolsCatalogTokens)
if catalogTokenBudget > 0 && catalogTokenBudget < budget {
budget = catalogTokenBudget
Expand Down
48 changes: 45 additions & 3 deletions coderd/x/chatd/chattool/findtools_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,20 @@ func TestSearchTools(t *testing.T) {
many[i] = FindToolCatalogEntry{Name: fmt.Sprintf("server__tool_%02d", i), Description: "common"}
}
result, _ := SearchTools(many, FindToolsArgs{Queries: []string{"common"}}, SearchBudget{})
require.Len(t, result.Matches, findToolsMaxMatches)
require.Len(t, result.Matches, findToolsDefaultMatches,
"an omitted limit returns the default match count")
require.Equal(t, "server__tool_00", result.Matches[0].Name)

raised, _ := SearchTools(many, FindToolsArgs{Queries: []string{"common"}, Limit: 15}, SearchBudget{})
require.Len(t, raised.Matches, 15)

clamped, _ := SearchTools(many, FindToolsArgs{Queries: []string{"common"}, Limit: 25}, SearchBudget{})
require.Len(t, clamped.Matches, findToolsMaxMatches,
"a limit above the hard cap clamps to it")

invalid, _ := SearchTools(many, FindToolsArgs{Queries: []string{"common"}, Limit: -1}, SearchBudget{})
require.Len(t, invalid.Matches, findToolsDefaultMatches,
"a non-positive limit falls back to the default")
})
t.Run("names capped and prioritized over queries", func(t *testing.T) {
t.Parallel()
Expand All @@ -62,13 +74,21 @@ func TestSearchTools(t *testing.T) {
names = append(names, many[i].Name)
}
result, _ := SearchTools(many, FindToolsArgs{Queries: []string{"common"}, Names: []string{"server__tool_24"}}, SearchBudget{})
require.Len(t, result.Matches, findToolsMaxMatches)
require.Len(t, result.Matches, findToolsDefaultMatches)
require.Equal(t, "server__tool_24", result.Matches[0].Name)
require.Contains(t, result.Activated, "server__tool_24")

capped, _ := SearchTools(many, FindToolsArgs{Names: names}, SearchBudget{})
require.Len(t, capped.Matches, findToolsMaxMatches)
require.Len(t, capped.Matches, findToolsMaxMatches,
"exact names bypass the default limit up to the hard cap")
require.Len(t, capped.Activated, findToolsMaxMatches)

bypassed, _ := SearchTools(many, FindToolsArgs{Queries: []string{"common"}, Names: names[:12], Limit: 5}, SearchBudget{})
require.Len(t, bypassed.Matches, 12,
"exact names bypass an explicit lower limit and leave no keyword slots")
for i, name := range names[:12] {
require.Equal(t, name, bypassed.Matches[i].Name)
}
})
t.Run("names list is bounded", func(t *testing.T) {
t.Parallel()
Expand Down Expand Up @@ -101,6 +121,17 @@ func TestSearchTools(t *testing.T) {
"tool description match outranks server metadata match")
require.Len(t, result.Matches, 2)
})
t.Run("coverage outranks concentrated score", func(t *testing.T) {
t.Parallel()
coverageEntries := []FindToolCatalogEntry{
{Name: "tracker__update", Description: "Assign labels and update status"},
{Name: "labels__tool", Description: "unrelated"},
}
result, _ := SearchTools(coverageEntries, FindToolsArgs{Queries: []string{"labels status"}}, SearchBudget{})
require.Len(t, result.Matches, 2)
require.Equal(t, "tracker__update", result.Matches[0].Name,
"an entry matching more distinct query terms outranks a higher single-term score")
})
t.Run("server prefix scope", func(t *testing.T) {
t.Parallel()
scopedEntries := []FindToolCatalogEntry{
Expand Down Expand Up @@ -273,6 +304,17 @@ func TestFindTools(t *testing.T) {
require.True(t, resp.IsError)
}

func TestFindToolsArgDescriptions(t *testing.T) {
t.Parallel()
info := FindTools(FindToolsOptions{}).Info()
for _, name := range []string{"queries", "names", "limit"} {
property, ok := info.Parameters[name].(map[string]any)
require.True(t, ok, "parameter %q must exist in the schema", name)
description, _ := property["description"].(string)
require.NotEmpty(t, description, "parameter %q needs model guidance in its schema description", name)
}
}

func TestFindToolsSerialToolCalls(t *testing.T) {
t.Parallel()
serial, ok := FindTools(FindToolsOptions{}).(interface{ SerialToolCalls() bool })
Expand Down
Loading