diff --git a/mcp/zep-mcp-server/internal/handlers/search.go b/mcp/zep-mcp-server/internal/handlers/search.go index f8e3aacb..b10e909f 100644 --- a/mcp/zep-mcp-server/internal/handlers/search.go +++ b/mcp/zep-mcp-server/internal/handlers/search.go @@ -14,7 +14,11 @@ import ( // HandleSearchGraph handles the search_graph tool using the new MCP SDK signature func HandleSearchGraph(client *zepclient.Client) mcp.ToolHandlerFor[SearchGraphInput, any] { return func(ctx context.Context, req *mcp.CallToolRequest, input SearchGraphInput) (*mcp.CallToolResult, any, error) { - // Apply defaults + if err := input.Validate(); err != nil { + return nil, nil, err + } + + // Apply defaults (after validation so explicit out-of-range limits are rejected) if input.Scope == "" { input.Scope = "edges" } diff --git a/mcp/zep-mcp-server/internal/handlers/types.go b/mcp/zep-mcp-server/internal/handlers/types.go index 7bfb1b68..7793b64e 100644 --- a/mcp/zep-mcp-server/internal/handlers/types.go +++ b/mcp/zep-mcp-server/internal/handlers/types.go @@ -1,5 +1,20 @@ package handlers +import ( + "fmt" + "strings" + "unicode/utf8" +) + +// Search input validation bounds (enforced before calling the Zep API). +const ( + searchQueryMaxLen = 2000 + searchLimitMin = 1 + searchLimitMax = 100 + searchMinScoreMin = 0.0 + searchMinScoreMax = 1.0 +) + // Input and output types for MCP tool handlers // These types are used to automatically generate JSON schemas @@ -17,6 +32,31 @@ type SearchGraphInput struct { EdgeTypes []string `json:"edge_types,omitempty" jsonschema:"Filter results by edge types"` } +// Validate checks SearchGraphInput before it is sent to the downstream API. +// Errors use the same fmt.Errorf patterns as internal/transform.ValidateRequired. +func (input SearchGraphInput) Validate() error { + query := strings.TrimSpace(input.Query) + if query == "" { + return fmt.Errorf("query cannot be empty") + } + if utf8.RuneCountInString(input.Query) > searchQueryMaxLen { + return fmt.Errorf("query exceeds maximum length of %d characters", searchQueryMaxLen) + } + + // MinFactRating is the minimum score field (min_fact_rating). Zero means omitted. + // Values outside [0, 1] are rejected (covers negative and >1.0). + if input.MinFactRating < searchMinScoreMin || input.MinFactRating > searchMinScoreMax { + return fmt.Errorf("min_fact_rating must be between %v and %v", searchMinScoreMin, searchMinScoreMax) + } + + // Limit 0 means use handler default; an explicit limit must be in [1, 100]. + if input.Limit != 0 && (input.Limit < searchLimitMin || input.Limit > searchLimitMax) { + return fmt.Errorf("limit must be between %d and %d", searchLimitMin, searchLimitMax) + } + + return nil +} + // GetUserContextInput defines the input parameters for get_user_context type GetUserContextInput struct { ThreadID string `json:"thread_id" jsonschema:"The thread ID for which to retrieve context"` diff --git a/mcp/zep-mcp-server/internal/handlers/types_test.go b/mcp/zep-mcp-server/internal/handlers/types_test.go index ccf6f488..f3af6245 100644 --- a/mcp/zep-mcp-server/internal/handlers/types_test.go +++ b/mcp/zep-mcp-server/internal/handlers/types_test.go @@ -1,6 +1,7 @@ package handlers import ( + "strings" "testing" ) @@ -92,3 +93,136 @@ func TestListThreadsInput(t *testing.T) { t.Errorf("Expected UserID 'user123', got '%s'", input.UserID) } } + +// TestSearchGraphInputValidate verifies search input validation rules. +func TestSearchGraphInputValidate(t *testing.T) { + longQuery := strings.Repeat("a", searchQueryMaxLen+1) + maxQuery := strings.Repeat("b", searchQueryMaxLen) + + tests := []struct { + name string + input SearchGraphInput + wantErr bool + errSub string // substring expected in error when wantErr is true + }{ + // valid cases + { + name: "valid minimal query", + input: SearchGraphInput{UserID: "u1", Query: "hello"}, + wantErr: false, + }, + { + name: "valid query with surrounding whitespace", + input: SearchGraphInput{UserID: "u1", Query: " hello "}, + wantErr: false, + }, + { + name: "valid query at max length", + input: SearchGraphInput{UserID: "u1", Query: maxQuery}, + wantErr: false, + }, + { + name: "valid limit omitted (zero uses handler default)", + input: SearchGraphInput{UserID: "u1", Query: "q", Limit: 0}, + wantErr: false, + }, + { + name: "valid limit at minimum boundary", + input: SearchGraphInput{UserID: "u1", Query: "q", Limit: searchLimitMin}, + wantErr: false, + }, + { + name: "valid limit at maximum boundary", + input: SearchGraphInput{UserID: "u1", Query: "q", Limit: searchLimitMax}, + wantErr: false, + }, + { + name: "valid min_fact_rating omitted (zero)", + input: SearchGraphInput{UserID: "u1", Query: "q", MinFactRating: 0}, + wantErr: false, + }, + { + name: "valid min_fact_rating at minimum boundary", + input: SearchGraphInput{UserID: "u1", Query: "q", MinFactRating: searchMinScoreMin}, + wantErr: false, + }, + { + name: "valid min_fact_rating at maximum boundary", + input: SearchGraphInput{UserID: "u1", Query: "q", MinFactRating: searchMinScoreMax}, + wantErr: false, + }, + { + name: "valid min_fact_rating mid range", + input: SearchGraphInput{UserID: "u1", Query: "q", MinFactRating: 0.5}, + wantErr: false, + }, + // invalid query + { + name: "empty query", + input: SearchGraphInput{UserID: "u1", Query: ""}, + wantErr: true, + errSub: "query cannot be empty", + }, + { + name: "whitespace only query", + input: SearchGraphInput{UserID: "u1", Query: " \t\n "}, + wantErr: true, + errSub: "query cannot be empty", + }, + { + name: "query exceeds max length", + input: SearchGraphInput{UserID: "u1", Query: longQuery}, + wantErr: true, + errSub: "query exceeds maximum length", + }, + // invalid limit (explicit values only; 0 is allowed as default sentinel) + { + name: "limit below minimum", + input: SearchGraphInput{UserID: "u1", Query: "q", Limit: -1}, + wantErr: true, + errSub: "limit must be between", + }, + { + name: "limit zero is default sentinel not error", + input: SearchGraphInput{UserID: "u1", Query: "q", Limit: 0}, + wantErr: false, + }, + { + name: "limit above maximum", + input: SearchGraphInput{UserID: "u1", Query: "q", Limit: searchLimitMax + 1}, + wantErr: true, + errSub: "limit must be between", + }, + // invalid min_fact_rating (MinScore-equivalent field on SearchGraphInput) + { + name: "min_fact_rating below minimum", + input: SearchGraphInput{UserID: "u1", Query: "q", MinFactRating: -0.01}, + wantErr: true, + errSub: "min_fact_rating must be between", + }, + { + name: "min_fact_rating above maximum", + input: SearchGraphInput{UserID: "u1", Query: "q", MinFactRating: 1.01}, + wantErr: true, + errSub: "min_fact_rating must be between", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.input.Validate() + if tt.wantErr { + if err == nil { + t.Fatalf("Validate() error = nil, wantErr true") + } + if tt.errSub != "" && !strings.Contains(err.Error(), tt.errSub) { + t.Errorf("Validate() error = %q, want substring %q", err.Error(), tt.errSub) + } + return + } + if err != nil { + t.Errorf("Validate() unexpected error = %v", err) + } + }) + } +}