From f648969626bd26e193e5470b49ffe2d9f0362308 Mon Sep 17 00:00:00 2001 From: Neyl Walecki Date: Wed, 16 Sep 2026 18:13:11 -0300 Subject: [PATCH 1/2] Fix underscore segment tool search --- internal/index/bleve.go | 80 ++++++++++++++++++++++++++++++++++++ internal/index/bleve_test.go | 44 ++++++++++++++++++++ 2 files changed, 124 insertions(+) diff --git a/internal/index/bleve.go b/internal/index/bleve.go index 3868178ae..11d12ac56 100644 --- a/internal/index/bleve.go +++ b/internal/index/bleve.go @@ -9,6 +9,7 @@ import ( "github.com/blevesearch/bleve/v2" "github.com/blevesearch/bleve/v2/analysis/analyzer/keyword" "github.com/blevesearch/bleve/v2/analysis/analyzer/standard" + "github.com/blevesearch/bleve/v2/search/query" "go.uber.org/zap" "github.com/smart-mcp-proxy/mcpproxy-go/internal/config" @@ -20,6 +21,13 @@ import ( // needed, so total coverage is never bounded by this value (MCP-3319). const defaultSearchPageSize = 10000 +const ( + // maxUnderscoreSearchSegments bounds the extra wildcard clauses generated for + // identifier-style queries. Typical MCP tool names stay well below this cap. + maxUnderscoreSearchSegments = 16 + underscoreSegmentBoost = 5.0 +) + // BleveIndex wraps Bleve index operations type BleveIndex struct { index bleve.Index @@ -350,6 +358,29 @@ func (b *BleveIndex) SearchTools(queryStr string, limit int) ([]*config.SearchRe return nil, fmt.Errorf("search failed: %w", err) } + // Identifier queries often include only the meaningful segments of a + // longer tool name. If the legacy query did not find a canonical exact + // match, repeat it with an additional segment-aware signal. This preserves + // exact-name scores while letting boundary matches outrank substring hits. + segmentQuery := underscoreSegmentQuery(queryStr) + if segmentQuery != nil { + hasExactToolName := false + for _, hit := range searchResult.Hits { + if toolName, ok := hit.Fields["tool_name"].(string); ok && toolName == queryStr { + hasExactToolName = true + break + } + } + + if !hasExactToolName { + boolQuery.AddShould(segmentQuery) + searchResult, err = b.index.Search(searchReq) + if err != nil { + return nil, fmt.Errorf("underscore segment search failed: %w", err) + } + } + } + // Convert results var results []*config.SearchResult for _, hit := range searchResult.Hits { @@ -363,6 +394,55 @@ func (b *BleveIndex) SearchTools(queryStr string, limit int) ([]*config.SearchRe return results, nil } +// underscoreSegmentQuery matches each underscore-delimited query segment at a +// complete segment boundary in the keyword-indexed tool_name field. Segment +// order is intentionally irrelevant, but every segment is required. +func underscoreSegmentQuery(queryStr string) query.Query { + segments := strings.Split(queryStr, "_") + if len(segments) < 2 || len(segments) > maxUnderscoreSearchSegments { + return nil + } + + segmentQueries := make([]query.Query, 0, len(segments)) + for _, segment := range segments { + if !isASCIIAlphanumeric(segment) { + return nil + } + + exact := bleve.NewTermQuery(segment) + exact.SetField("tool_name") + exact.SetBoost(underscoreSegmentBoost) + + prefix := bleve.NewPrefixQuery(segment + "_") + prefix.SetField("tool_name") + prefix.SetBoost(underscoreSegmentBoost) + + middle := bleve.NewWildcardQuery("*_" + segment + "_*") + middle.SetField("tool_name") + middle.SetBoost(underscoreSegmentBoost) + + suffix := bleve.NewWildcardQuery("*_" + segment) + suffix.SetField("tool_name") + suffix.SetBoost(underscoreSegmentBoost) + + segmentQueries = append(segmentQueries, bleve.NewDisjunctionQuery(exact, prefix, middle, suffix)) + } + + return bleve.NewConjunctionQuery(segmentQueries...) +} + +func isASCIIAlphanumeric(value string) bool { + if value == "" { + return false + } + for _, r := range value { + if (r < 'a' || r > 'z') && (r < 'A' || r > 'Z') && (r < '0' || r > '9') { + return false + } + } + return true +} + // GetDocumentCount returns the number of documents in the index func (b *BleveIndex) GetDocumentCount() (uint64, error) { return b.index.DocCount() diff --git a/internal/index/bleve_test.go b/internal/index/bleve_test.go index 8f9c5a26b..6e6cc1c45 100644 --- a/internal/index/bleve_test.go +++ b/internal/index/bleve_test.go @@ -217,6 +217,50 @@ func TestBleveIndex_SearchTokenization(t *testing.T) { } } +func TestBleveIndex_SearchUnderscoreSegments(t *testing.T) { + idx, err := NewBleveIndex(t.TempDir(), zap.NewNop()) + require.NoError(t, err) + defer idx.Close() + + require.NoError(t, idx.BatchIndex([]*config.ToolMetadata{ + { + Name: "work_start_task_attachment_upload", + ServerName: "fixture", + Description: "Create a signed request for a local file.", + ParamsJSON: `{"type":"object","properties":{}}`, + Hash: "target", + }, + { + Name: "network_upload_attachment", + ServerName: "fixture", + Description: "Create a signed request for a remote file.", + ParamsJSON: `{"type":"object","properties":{}}`, + Hash: "substring-decoy", + }, + })) + + t.Run("matches non-contiguous segments in any order", func(t *testing.T) { + for _, query := range []string{ + "work_upload_attachment", + "attachment_work_upload", + "work_start_task_attachment_upload", + } { + results, err := idx.SearchTools(query, 10) + require.NoError(t, err) + require.NotEmpty(t, results, query) + assert.Equal(t, "fixture:work_start_task_attachment_upload", results[0].Tool.Name, query) + } + }) + + t.Run("requires segment boundaries and every query segment", func(t *testing.T) { + for _, query := range []string{"network_work_upload_attachment", "work_missing_attachment"} { + results, err := idx.SearchTools(query, 10) + require.NoError(t, err) + assert.Empty(t, results, query) + } + }) +} + func TestBleveIndex_FieldMapping(t *testing.T) { // Test that all fields are properly indexed and searchable tool := &config.ToolMetadata{ From 29e0722925ccbec8214639d9663a57dbd5e9af0c Mon Sep 17 00:00:00 2001 From: Neyl Walecki Date: Thu, 17 Sep 2026 01:15:22 -0300 Subject: [PATCH 2/2] Address underscore search review feedback --- internal/index/bleve.go | 11 ++++++- internal/index/bleve_test.go | 61 ++++++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/internal/index/bleve.go b/internal/index/bleve.go index 11d12ac56..2637e6da0 100644 --- a/internal/index/bleve.go +++ b/internal/index/bleve.go @@ -366,7 +366,7 @@ func (b *BleveIndex) SearchTools(queryStr string, limit int) ([]*config.SearchRe if segmentQuery != nil { hasExactToolName := false for _, hit := range searchResult.Hits { - if toolName, ok := hit.Fields["tool_name"].(string); ok && toolName == queryStr { + if fieldsContainExactToolName(hit.Fields, queryStr) { hasExactToolName = true break } @@ -394,6 +394,15 @@ func (b *BleveIndex) SearchTools(queryStr string, limit int) ([]*config.SearchRe return results, nil } +func fieldsContainExactToolName(fields map[string]interface{}, queryStr string) bool { + for _, field := range []string{"tool_name", "full_tool_name"} { + if value, ok := fields[field].(string); ok && value == queryStr { + return true + } + } + return false +} + // underscoreSegmentQuery matches each underscore-delimited query segment at a // complete segment boundary in the keyword-indexed tool_name field. Segment // order is intentionally irrelevant, but every segment is required. diff --git a/internal/index/bleve_test.go b/internal/index/bleve_test.go index 6e6cc1c45..e1499f8de 100644 --- a/internal/index/bleve_test.go +++ b/internal/index/bleve_test.go @@ -259,6 +259,67 @@ func TestBleveIndex_SearchUnderscoreSegments(t *testing.T) { assert.Empty(t, results, query) } }) + + t.Run("preserves a legacy prefix match ahead of an added segment match", func(t *testing.T) { + prefixIdx, err := NewBleveIndex(t.TempDir(), zap.NewNop()) + require.NoError(t, err) + defer prefixIdx.Close() + + require.NoError(t, prefixIdx.BatchIndex([]*config.ToolMetadata{ + { + Name: "work_upload_attachment_preview", + ServerName: "fixture", + Description: "Preview an attachment upload for a Work task.", + Hash: "legacy-prefix", + }, + { + Name: "work_start_task_attachment_upload", + ServerName: "fixture", + Description: "Create a signed request for a local file.", + Hash: "segment-match", + }, + })) + + results, err := prefixIdx.SearchTools("work_upload_attachment", 10) + require.NoError(t, err) + require.Len(t, results, 2) + assert.Equal(t, "fixture:work_upload_attachment_preview", results[0].Tool.Name) + assert.Equal(t, "fixture:work_start_task_attachment_upload", results[1].Tool.Name) + }) +} + +func TestFieldsContainExactToolName(t *testing.T) { + tests := []struct { + name string + fields map[string]interface{} + query string + want bool + }{ + { + name: "tool name", + fields: map[string]interface{}{"tool_name": "work_upload_attachment", "full_tool_name": "fixture:work_upload_attachment"}, + query: "work_upload_attachment", + want: true, + }, + { + name: "full tool name", + fields: map[string]interface{}{"tool_name": "work_upload_attachment", "full_tool_name": "fixture:work_upload_attachment"}, + query: "fixture:work_upload_attachment", + want: true, + }, + { + name: "no exact name", + fields: map[string]interface{}{"tool_name": "work_start_task_attachment_upload", "full_tool_name": "fixture:work_start_task_attachment_upload"}, + query: "work_upload_attachment", + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + assert.Equal(t, tt.want, fieldsContainExactToolName(tt.fields, tt.query)) + }) + } } func TestBleveIndex_FieldMapping(t *testing.T) {