Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 89 additions & 0 deletions internal/index/bleve.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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
Expand Down Expand Up @@ -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 {
Expand All @@ -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()
Expand Down
105 changes: 105 additions & 0 deletions internal/index/bleve_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down
Loading