fix(search): match non-contiguous underscore segments (#1291)

* Fix underscore segment tool search

* Address underscore search review feedback

Author: neylwalecki

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 fieldsContainExactToolName(hit.Fields, 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,64 @@ 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.
+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()

internal/index/bleve_test.go

@@ -217,6 +217,111 @@ 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)
+		}
+	})
+
+	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) {
 	// Test that all fields are properly indexed and searchable
 	tool := &config.ToolMetadata{