diff --git a/codecov.yaml b/codecov.yaml index 1a8032e484..410f9ae7ee 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -13,6 +13,8 @@ coverage: - "**/mocks/**/*" - "**/mock_*.go" - "**/zz_generated.deepcopy.go" + - "**/*_test.go" + - "**/*_test_coverage.go" status: project: default: diff --git a/examples/vmcp-config-optimizer.yaml b/examples/vmcp-config-optimizer.yaml index 5b20b074d9..7687dabb7d 100644 --- a/examples/vmcp-config-optimizer.yaml +++ b/examples/vmcp-config-optimizer.yaml @@ -45,11 +45,11 @@ optimizer: # Enable the optimizer enabled: true - # Embedding backend: "ollama", "openai-compatible", or "placeholder" - # - "ollama": Uses local Ollama HTTP API for embeddings + # Embedding backend: "ollama" (default), "openai-compatible", or "vllm" + # - "ollama": Uses local Ollama HTTP API for embeddings (default, requires 'ollama serve') # - "openai-compatible": Uses OpenAI-compatible API (vLLM, OpenAI, etc.) - # - "placeholder": Uses deterministic hash-based embeddings (for testing) - embeddingBackend: placeholder + # - "vllm": Alias for OpenAI-compatible API + embeddingBackend: ollama # Embedding dimension (common values: 384, 768, 1536) # 384 is standard for all-MiniLM-L6-v2 and nomic-embed-text @@ -75,7 +75,7 @@ optimizer: # Option 1: Local Ollama (good for development/testing) # embeddingBackend: ollama # embeddingURL: http://localhost:11434 - # embeddingModel: nomic-embed-text + # embeddingModel: all-minilm # Default model (all-MiniLM-L6-v2) # embeddingDimension: 384 # Option 2: vLLM (recommended for production with GPU acceleration) diff --git a/pkg/optimizer/INTEGRATION.md b/pkg/optimizer/INTEGRATION.md index 4d2db78b59..e1cbd4d2df 100644 --- a/pkg/optimizer/INTEGRATION.md +++ b/pkg/optimizer/INTEGRATION.md @@ -93,7 +93,10 @@ func TestOptimizerIntegration(t *testing.T) { optimizerSvc, err := ingestion.NewService(&ingestion.Config{ DBConfig: &db.Config{Path: "/tmp/test-optimizer.db"}, EmbeddingConfig: &embeddings.Config{ - BackendType: "placeholder", + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, Dimension: 384, }, }) diff --git a/pkg/optimizer/README.md b/pkg/optimizer/README.md index 2984f2697a..f1a14938aa 100644 --- a/pkg/optimizer/README.md +++ b/pkg/optimizer/README.md @@ -132,9 +132,11 @@ func main() { panic(err) } - // Initialize embedding manager with placeholder (no external dependencies) + // Initialize embedding manager with Ollama (default) embeddingMgr, err := embeddings.NewManager(&embeddings.Config{ - BackendType: "placeholder", + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", Dimension: 384, }) if err != nil { @@ -201,7 +203,7 @@ spec: ollama serve # Pull an embedding model -ollama pull nomic-embed-text +ollama pull all-minilm ``` Configure vMCP: @@ -211,7 +213,7 @@ optimizer: enabled: true embeddingBackend: ollama embeddingURL: http://localhost:11434 - embeddingModel: nomic-embed-text + embeddingModel: all-minilm embeddingDimension: 384 ``` diff --git a/pkg/optimizer/db/backend_server.go b/pkg/optimizer/db/backend_server.go index 8685d4c47d..84ae5a3742 100644 --- a/pkg/optimizer/db/backend_server.go +++ b/pkg/optimizer/db/backend_server.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" "fmt" + "time" "github.com/philippgille/chromem-go" @@ -64,8 +65,13 @@ func (ops *BackendServerOps) Create(ctx context.Context, server *models.BackendS } // Also add to FTS5 database if available (for keyword filtering) + // Use background context to avoid cancellation issues - FTS5 is supplementary if ftsDB := ops.db.GetFTSDB(); ftsDB != nil { - if err := ftsDB.UpsertServer(ctx, server); err != nil { + // Use background context with timeout for FTS operations + // This ensures FTS operations complete even if the original context is canceled + ftsCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := ftsDB.UpsertServer(ftsCtx, server); err != nil { // Log but don't fail - FTS5 is supplementary logger.Warnf("Failed to upsert server to FTS5: %v", err) } diff --git a/pkg/optimizer/db/backend_server_test_coverage.go b/pkg/optimizer/db/backend_server_test_coverage.go new file mode 100644 index 0000000000..411be12673 --- /dev/null +++ b/pkg/optimizer/db/backend_server_test_coverage.go @@ -0,0 +1,94 @@ +package db + +import ( + "context" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/optimizer/models" +) + +// TestBackendServerOps_Create_FTS tests FTS integration in Create +func TestBackendServerOps_Create_FTS(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + config := &Config{ + PersistPath: filepath.Join(tmpDir, "test-db"), + FTSDBPath: filepath.Join(tmpDir, "fts.db"), + } + + db, err := NewDB(config) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + embeddingFunc := func(_ context.Context, _ string) ([]float32, error) { + return []float32{0.1, 0.2, 0.3}, nil + } + + ops := NewBackendServerOps(db, embeddingFunc) + + server := &models.BackendServer{ + ID: "server-1", + Name: "Test Server", + Description: stringPtr("A test server"), + Group: "default", + CreatedAt: time.Now(), + LastUpdated: time.Now(), + } + + // Create should also update FTS + err = ops.Create(ctx, server) + require.NoError(t, err) + + // Verify FTS was updated by checking FTS DB directly + ftsDB := db.GetFTSDB() + require.NotNil(t, ftsDB) + + // FTS should have the server + // We can't easily query FTS directly, but we can verify it doesn't error +} + +// TestBackendServerOps_Delete_FTS tests FTS integration in Delete +func TestBackendServerOps_Delete_FTS(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + config := &Config{ + PersistPath: filepath.Join(tmpDir, "test-db"), + FTSDBPath: filepath.Join(tmpDir, "fts.db"), + } + + db, err := NewDB(config) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + embeddingFunc := func(_ context.Context, _ string) ([]float32, error) { + return []float32{0.1, 0.2, 0.3}, nil + } + + ops := NewBackendServerOps(db, embeddingFunc) + + desc := "A test server" + server := &models.BackendServer{ + ID: "server-1", + Name: "Test Server", + Description: &desc, + Group: "default", + CreatedAt: time.Now(), + LastUpdated: time.Now(), + } + + // Create server + err = ops.Create(ctx, server) + require.NoError(t, err) + + // Delete should also delete from FTS + err = ops.Delete(ctx, server.ID) + require.NoError(t, err) +} diff --git a/pkg/optimizer/db/backend_tool.go b/pkg/optimizer/db/backend_tool.go index 909779edb8..3197428663 100644 --- a/pkg/optimizer/db/backend_tool.go +++ b/pkg/optimizer/db/backend_tool.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "time" "github.com/philippgille/chromem-go" @@ -63,8 +64,13 @@ func (ops *BackendToolOps) Create(ctx context.Context, tool *models.BackendTool, } // Also add to FTS5 database if available (for BM25 search) + // Use background context to avoid cancellation issues - FTS5 is supplementary if ops.db.fts != nil { - if err := ops.db.fts.UpsertToolMeta(ctx, tool, serverName); err != nil { + // Use background context with timeout for FTS operations + // This ensures FTS operations complete even if the original context is canceled + ftsCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if err := ops.db.fts.UpsertToolMeta(ftsCtx, tool, serverName); err != nil { // Log but don't fail - FTS5 is supplementary logger.Warnf("Failed to upsert tool to FTS5: %v", err) } diff --git a/pkg/optimizer/db/backend_tool_test.go b/pkg/optimizer/db/backend_tool_test.go index 557e5ca5f5..95d2d5330b 100644 --- a/pkg/optimizer/db/backend_tool_test.go +++ b/pkg/optimizer/db/backend_tool_test.go @@ -12,7 +12,7 @@ import ( "github.com/stacklok/toolhive/pkg/optimizer/models" ) -// createTestDB creates a test database with placeholder embeddings +// createTestDB creates a test database func createTestDB(t *testing.T) *DB { t.Helper() tmpDir := t.TempDir() @@ -27,18 +27,23 @@ func createTestDB(t *testing.T) *DB { return db } -// createTestEmbeddingFunc creates a test embedding function using placeholder embeddings +// createTestEmbeddingFunc creates a test embedding function using Ollama embeddings func createTestEmbeddingFunc(t *testing.T) func(ctx context.Context, text string) ([]float32, error) { t.Helper() - // Create placeholder embedding manager + // Try to use Ollama if available, otherwise skip test config := &embeddings.Config{ - BackendType: "placeholder", + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", Dimension: 384, } manager, err := embeddings.NewManager(config) - require.NoError(t, err) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull all-minilm'", err) + return nil + } t.Cleanup(func() { _ = manager.Close() }) return func(_ context.Context, text string) ([]float32, error) { @@ -454,9 +459,12 @@ func TestBackendToolOps_Search(t *testing.T) { require.NoError(t, err) assert.NotEmpty(t, results, "Should find tools") - // With placeholder embeddings, we just verify we get results - // Semantic similarity isn't guaranteed with hash-based embeddings - assert.Len(t, results, 2, "Should return both tools") + // Weather tool should be most similar to weather query + assert.NotEmpty(t, results, "Should find at least one tool") + if len(results) > 0 { + assert.Equal(t, "get_weather", results[0].ToolName, + "Weather tool should be most similar to weather query") + } } // TestBackendToolOps_Search_WithServerFilter tests search with server ID filter diff --git a/pkg/optimizer/db/backend_tool_test_coverage.go b/pkg/optimizer/db/backend_tool_test_coverage.go new file mode 100644 index 0000000000..a8766c302b --- /dev/null +++ b/pkg/optimizer/db/backend_tool_test_coverage.go @@ -0,0 +1,96 @@ +package db + +import ( + "context" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/optimizer/models" +) + +// TestBackendToolOps_Create_FTS tests FTS integration in Create +func TestBackendToolOps_Create_FTS(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + config := &Config{ + PersistPath: filepath.Join(tmpDir, "test-db"), + FTSDBPath: filepath.Join(tmpDir, "fts.db"), + } + + db, err := NewDB(config) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + embeddingFunc := func(_ context.Context, _ string) ([]float32, error) { + return []float32{0.1, 0.2, 0.3}, nil + } + + ops := NewBackendToolOps(db, embeddingFunc) + + desc := "A test tool" + tool := &models.BackendTool{ + ID: "tool-1", + MCPServerID: "server-1", + ToolName: "test_tool", + Description: &desc, + InputSchema: []byte(`{"type": "object"}`), + TokenCount: 10, + CreatedAt: time.Now(), + LastUpdated: time.Now(), + } + + // Create should also update FTS + err = ops.Create(ctx, tool, "TestServer") + require.NoError(t, err) + + // Verify FTS was updated + ftsDB := db.GetFTSDB() + require.NotNil(t, ftsDB) +} + +// TestBackendToolOps_DeleteByServer_FTS tests FTS integration in DeleteByServer +func TestBackendToolOps_DeleteByServer_FTS(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + config := &Config{ + PersistPath: filepath.Join(tmpDir, "test-db"), + FTSDBPath: filepath.Join(tmpDir, "fts.db"), + } + + db, err := NewDB(config) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + embeddingFunc := func(_ context.Context, _ string) ([]float32, error) { + return []float32{0.1, 0.2, 0.3}, nil + } + + ops := NewBackendToolOps(db, embeddingFunc) + + desc := "A test tool" + tool := &models.BackendTool{ + ID: "tool-1", + MCPServerID: "server-1", + ToolName: "test_tool", + Description: &desc, + InputSchema: []byte(`{"type": "object"}`), + TokenCount: 10, + CreatedAt: time.Now(), + LastUpdated: time.Now(), + } + + // Create tool + err = ops.Create(ctx, tool, "TestServer") + require.NoError(t, err) + + // DeleteByServer should also delete from FTS + err = ops.DeleteByServer(ctx, "server-1") + require.NoError(t, err) +} diff --git a/pkg/optimizer/db/db.go b/pkg/optimizer/db/db.go index f7e7df5bb8..2e1b88a24f 100644 --- a/pkg/optimizer/db/db.go +++ b/pkg/optimizer/db/db.go @@ -3,6 +3,8 @@ package db import ( "context" "fmt" + "os" + "strings" "sync" "github.com/philippgille/chromem-go" @@ -54,7 +56,35 @@ func NewDB(config *Config) (*DB, error) { logger.Infof("Creating chromem-go database with persistence at: %s", config.PersistPath) chromemDB, err = chromem.NewPersistentDB(config.PersistPath, false) if err != nil { - return nil, fmt.Errorf("failed to create persistent database: %w", err) + // Check if error is due to corrupted database (missing collection metadata) + if strings.Contains(err.Error(), "collection metadata file not found") { + logger.Warnf("Database appears corrupted, attempting to remove and recreate: %v", err) + // Try to remove corrupted database directory + // Use RemoveAll which should handle directories recursively + // If it fails, we'll try to create with a new path or fall back to in-memory + if removeErr := os.RemoveAll(config.PersistPath); removeErr != nil { + logger.Warnf("Failed to remove corrupted database directory (may be in use): %v. Will try to recreate anyway.", removeErr) + // Try to rename the corrupted directory and create a new one + backupPath := config.PersistPath + ".corrupted" + if renameErr := os.Rename(config.PersistPath, backupPath); renameErr != nil { + logger.Warnf("Failed to rename corrupted database: %v. Attempting to create database anyway.", renameErr) + // Continue and let chromem-go handle it - it might work if the corruption is partial + } else { + logger.Infof("Renamed corrupted database to: %s", backupPath) + } + } + // Retry creating the database + chromemDB, err = chromem.NewPersistentDB(config.PersistPath, false) + if err != nil { + // If still failing, return the error but suggest manual cleanup + return nil, fmt.Errorf( + "failed to create persistent database after cleanup attempt. Please manually remove %s and try again: %w", + config.PersistPath, err) + } + logger.Info("Successfully recreated database after cleanup") + } else { + return nil, fmt.Errorf("failed to create persistent database: %w", err) + } } } else { logger.Info("Creating in-memory chromem-go database") @@ -160,7 +190,7 @@ func (db *DB) GetFTSDB() *FTSDatabase { return db.fts } -// Reset clears all collections and FTS tables (useful for testing) +// Reset clears all collections and FTS tables (useful for testing and startup) func (db *DB) Reset() { db.mu.Lock() defer db.mu.Unlock() diff --git a/pkg/optimizer/db/db_test.go b/pkg/optimizer/db/db_test.go new file mode 100644 index 0000000000..2da34c214a --- /dev/null +++ b/pkg/optimizer/db/db_test.go @@ -0,0 +1,302 @@ +package db + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestNewDB_CorruptedDatabase tests database recovery from corruption +func TestNewDB_CorruptedDatabase(t *testing.T) { + t.Parallel() + tmpDir := t.TempDir() + dbPath := filepath.Join(tmpDir, "corrupted-db") + + // Create a directory that looks like a corrupted database + err := os.MkdirAll(dbPath, 0755) + require.NoError(t, err) + + // Create a file that might cause issues + err = os.WriteFile(filepath.Join(dbPath, "some-file"), []byte("corrupted"), 0644) + require.NoError(t, err) + + config := &Config{ + PersistPath: dbPath, + } + + // Should recover from corruption + db, err := NewDB(config) + require.NoError(t, err) + require.NotNil(t, db) + defer func() { _ = db.Close() }() +} + +// TestNewDB_CorruptedDatabase_RecoveryFailure tests when recovery fails +func TestNewDB_CorruptedDatabase_RecoveryFailure(t *testing.T) { + t.Parallel() + tmpDir := t.TempDir() + dbPath := filepath.Join(tmpDir, "corrupted-db") + + // Create a directory that looks like a corrupted database + err := os.MkdirAll(dbPath, 0755) + require.NoError(t, err) + + // Create a file that might cause issues + err = os.WriteFile(filepath.Join(dbPath, "some-file"), []byte("corrupted"), 0644) + require.NoError(t, err) + + // Make directory read-only to simulate recovery failure + // Note: This might not work on all systems, so we'll test the error path differently + // Instead, we'll test with an invalid path that can't be created + config := &Config{ + PersistPath: "/invalid/path/that/does/not/exist", + } + + _, err = NewDB(config) + // Should return error for invalid path + assert.Error(t, err) +} + +// TestDB_GetOrCreateCollection tests collection creation and retrieval +func TestDB_GetOrCreateCollection(t *testing.T) { + t.Parallel() + ctx := context.Background() + + config := &Config{ + PersistPath: "", // In-memory + } + + db, err := NewDB(config) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + // Create a simple embedding function + embeddingFunc := func(_ context.Context, _ string) ([]float32, error) { + return []float32{0.1, 0.2, 0.3}, nil + } + + // Get or create collection + collection, err := db.GetOrCreateCollection(ctx, "test-collection", embeddingFunc) + require.NoError(t, err) + require.NotNil(t, collection) + + // Get existing collection + collection2, err := db.GetOrCreateCollection(ctx, "test-collection", embeddingFunc) + require.NoError(t, err) + require.NotNil(t, collection2) + assert.Equal(t, collection, collection2) +} + +// TestDB_GetCollection tests collection retrieval +func TestDB_GetCollection(t *testing.T) { + t.Parallel() + ctx := context.Background() + + config := &Config{ + PersistPath: "", // In-memory + } + + db, err := NewDB(config) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + embeddingFunc := func(_ context.Context, _ string) ([]float32, error) { + return []float32{0.1, 0.2, 0.3}, nil + } + + // Get non-existent collection should fail + _, err = db.GetCollection("non-existent", embeddingFunc) + assert.Error(t, err) + + // Create collection first + _, err = db.GetOrCreateCollection(ctx, "test-collection", embeddingFunc) + require.NoError(t, err) + + // Now get it + collection, err := db.GetCollection("test-collection", embeddingFunc) + require.NoError(t, err) + require.NotNil(t, collection) +} + +// TestDB_DeleteCollection tests collection deletion +func TestDB_DeleteCollection(t *testing.T) { + t.Parallel() + ctx := context.Background() + + config := &Config{ + PersistPath: "", // In-memory + } + + db, err := NewDB(config) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + embeddingFunc := func(_ context.Context, _ string) ([]float32, error) { + return []float32{0.1, 0.2, 0.3}, nil + } + + // Create collection + _, err = db.GetOrCreateCollection(ctx, "test-collection", embeddingFunc) + require.NoError(t, err) + + // Delete collection + db.DeleteCollection("test-collection") + + // Verify it's deleted + _, err = db.GetCollection("test-collection", embeddingFunc) + assert.Error(t, err) +} + +// TestDB_Reset tests database reset +func TestDB_Reset(t *testing.T) { + t.Parallel() + ctx := context.Background() + + config := &Config{ + PersistPath: "", // In-memory + } + + db, err := NewDB(config) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + embeddingFunc := func(_ context.Context, _ string) ([]float32, error) { + return []float32{0.1, 0.2, 0.3}, nil + } + + // Create collections + _, err = db.GetOrCreateCollection(ctx, BackendServerCollection, embeddingFunc) + require.NoError(t, err) + + _, err = db.GetOrCreateCollection(ctx, BackendToolCollection, embeddingFunc) + require.NoError(t, err) + + // Reset database + db.Reset() + + // Verify collections are deleted + _, err = db.GetCollection(BackendServerCollection, embeddingFunc) + assert.Error(t, err) + + _, err = db.GetCollection(BackendToolCollection, embeddingFunc) + assert.Error(t, err) +} + +// TestDB_GetChromemDB tests chromem DB accessor +func TestDB_GetChromemDB(t *testing.T) { + t.Parallel() + + config := &Config{ + PersistPath: "", // In-memory + } + + db, err := NewDB(config) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + chromemDB := db.GetChromemDB() + require.NotNil(t, chromemDB) +} + +// TestDB_GetFTSDB tests FTS DB accessor +func TestDB_GetFTSDB(t *testing.T) { + t.Parallel() + + config := &Config{ + PersistPath: "", // In-memory + } + + db, err := NewDB(config) + require.NoError(t, err) + defer func() { _ = db.Close() }() + + ftsDB := db.GetFTSDB() + require.NotNil(t, ftsDB) +} + +// TestDB_Close tests database closing +func TestDB_Close(t *testing.T) { + t.Parallel() + + config := &Config{ + PersistPath: "", // In-memory + } + + db, err := NewDB(config) + require.NoError(t, err) + + err = db.Close() + require.NoError(t, err) + + // Multiple closes should be safe + err = db.Close() + require.NoError(t, err) +} + +// TestNewDB_FTSDBPath tests FTS database path configuration +func TestNewDB_FTSDBPath(t *testing.T) { + t.Parallel() + tmpDir := t.TempDir() + + tests := []struct { + name string + config *Config + wantErr bool + }{ + { + name: "in-memory FTS with persistent chromem", + config: &Config{ + PersistPath: filepath.Join(tmpDir, "db"), + FTSDBPath: ":memory:", + }, + wantErr: false, + }, + { + name: "persistent FTS with persistent chromem", + config: &Config{ + PersistPath: filepath.Join(tmpDir, "db2"), + FTSDBPath: filepath.Join(tmpDir, "fts.db"), + }, + wantErr: false, + }, + { + name: "default FTS path with persistent chromem", + config: &Config{ + PersistPath: filepath.Join(tmpDir, "db3"), + // FTSDBPath not set, should default to {PersistPath}/fts.db + }, + wantErr: false, + }, + { + name: "in-memory FTS with in-memory chromem", + config: &Config{ + PersistPath: "", + FTSDBPath: ":memory:", + }, + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + db, err := NewDB(tt.config) + if tt.wantErr { + assert.Error(t, err) + } else { + require.NoError(t, err) + require.NotNil(t, db) + defer func() { _ = db.Close() }() + + // Verify FTS DB is accessible + ftsDB := db.GetFTSDB() + require.NotNil(t, ftsDB) + } + }) + } +} diff --git a/pkg/optimizer/db/fts.go b/pkg/optimizer/db/fts.go index 8dde0b2aa3..e9cecd7a09 100644 --- a/pkg/optimizer/db/fts.go +++ b/pkg/optimizer/db/fts.go @@ -316,6 +316,22 @@ func (fts *FTSDatabase) SearchBM25( return results, nil } +// GetTotalToolTokens returns the sum of token_count across all tools +func (fts *FTSDatabase) GetTotalToolTokens(ctx context.Context) (int, error) { + fts.mu.RLock() + defer fts.mu.RUnlock() + + var totalTokens int + query := "SELECT COALESCE(SUM(token_count), 0) FROM backend_tools_fts" + + err := fts.db.QueryRowContext(ctx, query).Scan(&totalTokens) + if err != nil { + return 0, fmt.Errorf("failed to get total tool tokens: %w", err) + } + + return totalTokens, nil +} + // Close closes the FTS database connection func (fts *FTSDatabase) Close() error { return fts.db.Close() diff --git a/pkg/optimizer/db/fts_test_coverage.go b/pkg/optimizer/db/fts_test_coverage.go new file mode 100644 index 0000000000..b6a7fe2321 --- /dev/null +++ b/pkg/optimizer/db/fts_test_coverage.go @@ -0,0 +1,159 @@ +package db + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/optimizer/models" +) + +// stringPtr returns a pointer to the given string +func stringPtr(s string) *string { + return &s +} + +// TestFTSDatabase_GetTotalToolTokens tests token counting +func TestFTSDatabase_GetTotalToolTokens(t *testing.T) { + t.Parallel() + ctx := context.Background() + + config := &FTSConfig{ + DBPath: ":memory:", + } + + ftsDB, err := NewFTSDatabase(config) + require.NoError(t, err) + defer func() { _ = ftsDB.Close() }() + + // Initially should be 0 + totalTokens, err := ftsDB.GetTotalToolTokens(ctx) + require.NoError(t, err) + assert.Equal(t, 0, totalTokens) + + // Add a tool + tool := &models.BackendTool{ + ID: "tool-1", + MCPServerID: "server-1", + ToolName: "test_tool", + Description: stringPtr("Test tool"), + TokenCount: 100, + CreatedAt: time.Now(), + LastUpdated: time.Now(), + } + + err = ftsDB.UpsertToolMeta(ctx, tool, "TestServer") + require.NoError(t, err) + + // Should now have tokens + totalTokens, err = ftsDB.GetTotalToolTokens(ctx) + require.NoError(t, err) + assert.Equal(t, 100, totalTokens) + + // Add another tool + tool2 := &models.BackendTool{ + ID: "tool-2", + MCPServerID: "server-1", + ToolName: "test_tool2", + Description: stringPtr("Test tool 2"), + TokenCount: 50, + CreatedAt: time.Now(), + LastUpdated: time.Now(), + } + + err = ftsDB.UpsertToolMeta(ctx, tool2, "TestServer") + require.NoError(t, err) + + // Should sum tokens + totalTokens, err = ftsDB.GetTotalToolTokens(ctx) + require.NoError(t, err) + assert.Equal(t, 150, totalTokens) +} + +// TestSanitizeFTS5Query tests query sanitization +func TestSanitizeFTS5Query(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input string + expected string + }{ + { + name: "remove quotes", + input: `"test query"`, + expected: "test query", + }, + { + name: "remove wildcards", + input: "test*query", + expected: "test query", + }, + { + name: "remove parentheses", + input: "test(query)", + expected: "test query", + }, + { + name: "remove multiple spaces", + input: "test query", + expected: "test query", + }, + { + name: "trim whitespace", + input: " test query ", + expected: "test query", + }, + { + name: "empty string", + input: "", + expected: "", + }, + { + name: "only special characters", + input: `"*()`, + expected: "", + }, + { + name: "mixed special characters", + input: `test"query*with(special)chars`, + expected: "test query with special chars", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + result := sanitizeFTS5Query(tt.input) + assert.Equal(t, tt.expected, result) + }) + } +} + +// TestFTSDatabase_SearchBM25_EmptyQuery tests empty query handling +func TestFTSDatabase_SearchBM25_EmptyQuery(t *testing.T) { + t.Parallel() + ctx := context.Background() + + config := &FTSConfig{ + DBPath: ":memory:", + } + + ftsDB, err := NewFTSDatabase(config) + require.NoError(t, err) + defer func() { _ = ftsDB.Close() }() + + // Empty query should return empty results + results, err := ftsDB.SearchBM25(ctx, "", 10, nil) + require.NoError(t, err) + assert.Empty(t, results) + + // Query with only special characters should return empty results + results, err = ftsDB.SearchBM25(ctx, `"*()`, 10, nil) + require.NoError(t, err) + assert.Empty(t, results) +} diff --git a/pkg/optimizer/doc.go b/pkg/optimizer/doc.go index 0808bb76b2..549bf23900 100644 --- a/pkg/optimizer/doc.go +++ b/pkg/optimizer/doc.go @@ -69,7 +69,9 @@ // // // Create embedding manager // embMgr, err := embeddings.NewManager(embeddings.Config{ -// BackendType: "placeholder", // or "ollama" or "openai-compatible" +// BackendType: "ollama", // or "openai-compatible" or "vllm" +// BaseURL: "http://localhost:11434", +// Model: "all-minilm", // Dimension: 384, // }) // diff --git a/pkg/optimizer/embeddings/manager.go b/pkg/optimizer/embeddings/manager.go index 9ccc94fca3..70ac838492 100644 --- a/pkg/optimizer/embeddings/manager.go +++ b/pkg/optimizer/embeddings/manager.go @@ -8,17 +8,19 @@ import ( ) const ( - // BackendTypePlaceholder is the placeholder backend type - BackendTypePlaceholder = "placeholder" + // DefaultModelAllMiniLM is the default Ollama model name + DefaultModelAllMiniLM = "all-minilm" + // BackendTypeOllama is the Ollama backend type + BackendTypeOllama = "ollama" ) // Config holds configuration for the embedding manager type Config struct { // BackendType specifies which backend to use: - // - "ollama": Ollama native API + // - "ollama": Ollama native API (default) // - "vllm": vLLM OpenAI-compatible API // - "unified": Generic OpenAI-compatible API (works with both) - // - "placeholder": Hash-based embeddings for testing + // - "openai": OpenAI-compatible API BackendType string // BaseURL is the base URL for the embedding service @@ -27,7 +29,7 @@ type Config struct { BaseURL string // Model is the model name to use - // - Ollama: "nomic-embed-text", "all-minilm" + // - Ollama: "all-minilm" (default), "nomic-embed-text" // - vLLM: "sentence-transformers/all-MiniLM-L6-v2", "intfloat/e5-mistral-7b-instruct" Model string @@ -68,9 +70,9 @@ func NewManager(config *Config) (*Manager, error) { config.MaxCacheSize = 1000 } - // Default to placeholder (zero dependencies) + // Default to Ollama if config.BackendType == "" { - config.BackendType = "placeholder" + config.BackendType = BackendTypeOllama } // Initialize backend based on configuration @@ -78,7 +80,7 @@ func NewManager(config *Config) (*Manager, error) { var err error switch config.BackendType { - case "ollama": + case BackendTypeOllama: // Use Ollama native API (requires ollama serve) baseURL := config.BaseURL if baseURL == "" { @@ -86,13 +88,17 @@ func NewManager(config *Config) (*Manager, error) { } model := config.Model if model == "" { - model = "nomic-embed-text" + model = DefaultModelAllMiniLM // Default: all-MiniLM-L6-v2 + } + // Update dimension if not set and using default model + if config.Dimension == 0 && model == DefaultModelAllMiniLM { + config.Dimension = 384 } backend, err = NewOllamaBackend(baseURL, model) if err != nil { - logger.Warnf("Failed to initialize Ollama backend: %v", err) - logger.Info("Falling back to placeholder embeddings. To use Ollama: ollama serve && ollama pull nomic-embed-text") - backend = &PlaceholderBackend{dimension: config.Dimension} + return nil, fmt.Errorf( + "failed to initialize Ollama backend: %w (ensure 'ollama serve' is running and 'ollama pull %s' has been executed)", + err, DefaultModelAllMiniLM) } case "vllm", "unified", "openai": @@ -107,17 +113,11 @@ func NewManager(config *Config) (*Manager, error) { } backend, err = NewOpenAICompatibleBackend(config.BaseURL, config.Model, config.Dimension) if err != nil { - logger.Warnf("Failed to initialize %s backend: %v", config.BackendType, err) - logger.Infof("Falling back to placeholder embeddings") - backend = &PlaceholderBackend{dimension: config.Dimension} + return nil, fmt.Errorf("failed to initialize %s backend: %w", config.BackendType, err) } - case BackendTypePlaceholder: - // Use placeholder for testing - backend = &PlaceholderBackend{dimension: config.Dimension} - default: - return nil, fmt.Errorf("unknown backend type: %s (supported: ollama, vllm, unified, placeholder)", config.BackendType) + return nil, fmt.Errorf("unknown backend type: %s (supported: ollama, vllm, unified, openai)", config.BackendType) } m := &Manager{ @@ -154,17 +154,7 @@ func (m *Manager) GenerateEmbedding(texts []string) ([][]float32, error) { // Use backend to generate embeddings embeddings, err := m.backend.EmbedBatch(texts) if err != nil { - // If backend fails, fall back to placeholder for non-placeholder backends - if m.config.BackendType != "placeholder" { - logger.Warnf("%s backend failed: %v, falling back to placeholder", m.config.BackendType, err) - placeholder := &PlaceholderBackend{dimension: m.config.Dimension} - embeddings, err = placeholder.EmbedBatch(texts) - if err != nil { - return nil, fmt.Errorf("failed to generate embeddings: %w", err) - } - } else { - return nil, fmt.Errorf("failed to generate embeddings: %w", err) - } + return nil, fmt.Errorf("failed to generate embeddings: %w", err) } // Cache single embeddings @@ -176,65 +166,6 @@ func (m *Manager) GenerateEmbedding(texts []string) ([][]float32, error) { return embeddings, nil } -// PlaceholderBackend is a simple backend for testing -type PlaceholderBackend struct { - dimension int -} - -// Embed generates a deterministic hash-based embedding for the given text. -func (p *PlaceholderBackend) Embed(text string) ([]float32, error) { - return p.generatePlaceholderEmbedding(text), nil -} - -// EmbedBatch generates embeddings for multiple texts. -func (p *PlaceholderBackend) EmbedBatch(texts []string) ([][]float32, error) { - embeddings := make([][]float32, len(texts)) - for i, text := range texts { - embeddings[i] = p.generatePlaceholderEmbedding(text) - } - return embeddings, nil -} - -// Dimension returns the embedding dimension. -func (p *PlaceholderBackend) Dimension() int { - return p.dimension -} - -// Close closes the backend (no-op for placeholder). -func (*PlaceholderBackend) Close() error { - return nil -} - -func (p *PlaceholderBackend) generatePlaceholderEmbedding(text string) []float32 { - embedding := make([]float32, p.dimension) - - // Simple hash-based generation for testing - hash := 0 - for _, c := range text { - hash = (hash*31 + int(c)) % 1000000 - } - - // Generate deterministic values - for i := range embedding { - hash = (hash*1103515245 + 12345) % 1000000 - embedding[i] = float32(hash) / 1000000.0 - } - - // Normalize the embedding (L2 normalization) - var norm float32 - for _, v := range embedding { - norm += v * v - } - if norm > 0 { - norm = float32(1.0 / float64(norm)) - for i := range embedding { - embedding[i] *= norm - } - } - - return embedding -} - // GetCacheStats returns cache statistics func (m *Manager) GetCacheStats() map[string]interface{} { if !m.config.EnableCache || m.cache == nil { diff --git a/pkg/optimizer/embeddings/manager_test_coverage.go b/pkg/optimizer/embeddings/manager_test_coverage.go new file mode 100644 index 0000000000..98eb4a9eec --- /dev/null +++ b/pkg/optimizer/embeddings/manager_test_coverage.go @@ -0,0 +1,155 @@ +package embeddings + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestManager_GetCacheStats tests cache statistics +func TestManager_GetCacheStats(t *testing.T) { + t.Parallel() + + config := &Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + EnableCache: true, + MaxCacheSize: 100, + } + + manager, err := NewManager(config) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + defer func() { _ = manager.Close() }() + + stats := manager.GetCacheStats() + require.NotNil(t, stats) + assert.True(t, stats["enabled"].(bool)) + assert.Contains(t, stats, "hits") + assert.Contains(t, stats, "misses") + assert.Contains(t, stats, "size") + assert.Contains(t, stats, "maxsize") +} + +// TestManager_GetCacheStats_Disabled tests cache statistics when cache is disabled +func TestManager_GetCacheStats_Disabled(t *testing.T) { + t.Parallel() + + config := &Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + EnableCache: false, + } + + manager, err := NewManager(config) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + defer func() { _ = manager.Close() }() + + stats := manager.GetCacheStats() + require.NotNil(t, stats) + assert.False(t, stats["enabled"].(bool)) +} + +// TestManager_ClearCache tests cache clearing +func TestManager_ClearCache(t *testing.T) { + t.Parallel() + + config := &Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + EnableCache: true, + MaxCacheSize: 100, + } + + manager, err := NewManager(config) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + defer func() { _ = manager.Close() }() + + // Clear cache should not panic + manager.ClearCache() + + // Multiple clears should be safe + manager.ClearCache() +} + +// TestManager_ClearCache_Disabled tests cache clearing when cache is disabled +func TestManager_ClearCache_Disabled(t *testing.T) { + t.Parallel() + + config := &Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + EnableCache: false, + } + + manager, err := NewManager(config) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + defer func() { _ = manager.Close() }() + + // Clear cache should not panic even when disabled + manager.ClearCache() +} + +// TestManager_Dimension tests dimension accessor +func TestManager_Dimension(t *testing.T) { + t.Parallel() + + config := &Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + manager, err := NewManager(config) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + defer func() { _ = manager.Close() }() + + dimension := manager.Dimension() + assert.Equal(t, 384, dimension) +} + +// TestManager_Dimension_Default tests default dimension +func TestManager_Dimension_Default(t *testing.T) { + t.Parallel() + + config := &Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + // Dimension not set, should default to 384 + } + + manager, err := NewManager(config) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + defer func() { _ = manager.Close() }() + + dimension := manager.Dimension() + assert.Equal(t, 384, dimension) +} diff --git a/pkg/optimizer/embeddings/ollama.go b/pkg/optimizer/embeddings/ollama.go index d6f4874375..a05af2af11 100644 --- a/pkg/optimizer/embeddings/ollama.go +++ b/pkg/optimizer/embeddings/ollama.go @@ -31,21 +31,27 @@ type ollamaEmbedResponse struct { // NewOllamaBackend creates a new Ollama backend // Requires Ollama to be running locally: ollama serve -// Default model: nomic-embed-text (768 dimensions) +// Default model: all-minilm (all-MiniLM-L6-v2, 384 dimensions) func NewOllamaBackend(baseURL, model string) (*OllamaBackend, error) { if baseURL == "" { baseURL = "http://localhost:11434" } if model == "" { - model = "nomic-embed-text" // Default embedding model + model = "all-minilm" // Default embedding model (all-MiniLM-L6-v2) } logger.Infof("Initializing Ollama backend (model: %s, url: %s)", model, baseURL) + // Determine dimension based on model + dimension := 384 // Default for all-minilm + if model == "nomic-embed-text" { + dimension = 768 + } + backend := &OllamaBackend{ baseURL: baseURL, model: model, - dimension: 768, // nomic-embed-text dimension + dimension: dimension, client: &http.Client{}, } diff --git a/pkg/optimizer/embeddings/ollama_test.go b/pkg/optimizer/embeddings/ollama_test.go index 5254b7c072..83594863e5 100644 --- a/pkg/optimizer/embeddings/ollama_test.go +++ b/pkg/optimizer/embeddings/ollama_test.go @@ -4,13 +4,12 @@ import ( "testing" ) -func TestOllamaBackend_Placeholder(t *testing.T) { +func TestOllamaBackend_ConnectionFailure(t *testing.T) { t.Parallel() - // This test verifies that Ollama backend is properly structured - // Actual Ollama tests require ollama to be running + // This test verifies that Ollama backend handles connection failures gracefully // Test that NewOllamaBackend handles connection failure gracefully - _, err := NewOllamaBackend("http://localhost:99999", "nomic-embed-text") + _, err := NewOllamaBackend("http://localhost:99999", "all-minilm") if err == nil { t.Error("Expected error when connecting to invalid Ollama URL") } @@ -18,68 +17,36 @@ func TestOllamaBackend_Placeholder(t *testing.T) { func TestManagerWithOllama(t *testing.T) { t.Parallel() - // Test that Manager falls back to placeholder when Ollama is not available or model not pulled + // Test that Manager works with Ollama when available config := &Config{ - BackendType: "ollama", - Dimension: 384, + BackendType: BackendTypeOllama, + BaseURL: "http://localhost:11434", + Model: DefaultModelAllMiniLM, + Dimension: 768, EnableCache: true, MaxCacheSize: 100, } manager, err := NewManager(config) if err != nil { - t.Fatalf("Failed to create manager: %v", err) - } - defer manager.Close() - - // Should work with placeholder backend fallback - // (Ollama might not have model pulled, so it falls back to placeholder) - embeddings, err := manager.GenerateEmbedding([]string{"test text"}) - - // If Ollama is available with the model, great! - // If not, it should have fallen back to placeholder - if err != nil { - // Check if it's a "model not found" error - this is expected - if embeddings == nil { - t.Skip("Ollama not available or model not pulled (expected in CI/test environments)") - } - } - - if len(embeddings) != 1 { - t.Errorf("Expected 1 embedding, got %d", len(embeddings)) - } - - // Dimension could be 384 (placeholder) or 768 (Ollama nomic-embed-text) - if len(embeddings[0]) != 384 && len(embeddings[0]) != 768 { - t.Errorf("Expected dimension 384 or 768, got %d", len(embeddings[0])) - } -} - -func TestManagerWithPlaceholder(t *testing.T) { - t.Parallel() - // Test explicit placeholder backend - config := &Config{ - BackendType: "placeholder", - Dimension: 384, - EnableCache: false, - } - - manager, err := NewManager(config) - if err != nil { - t.Fatalf("Failed to create manager: %v", err) + t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull all-minilm'", err) + return } defer manager.Close() // Test single embedding - embeddings, err := manager.GenerateEmbedding([]string{"hello world"}) + embeddings, err := manager.GenerateEmbedding([]string{"test text"}) if err != nil { - t.Fatalf("Failed to generate embedding: %v", err) + // Model might not be pulled - skip gracefully + t.Skipf("Skipping test: Failed to generate embedding. Error: %v. Run 'ollama pull nomic-embed-text'", err) + return } if len(embeddings) != 1 { t.Errorf("Expected 1 embedding, got %d", len(embeddings)) } + // Ollama all-minilm uses 384 dimensions if len(embeddings[0]) != 384 { t.Errorf("Expected dimension 384, got %d", len(embeddings[0])) } @@ -88,19 +55,12 @@ func TestManagerWithPlaceholder(t *testing.T) { texts := []string{"text 1", "text 2", "text 3"} embeddings, err = manager.GenerateEmbedding(texts) if err != nil { - t.Fatalf("Failed to generate batch embeddings: %v", err) + // Model might not be pulled - skip gracefully + t.Skipf("Skipping test: Failed to generate batch embeddings. Error: %v. Run 'ollama pull nomic-embed-text'", err) + return } if len(embeddings) != 3 { t.Errorf("Expected 3 embeddings, got %d", len(embeddings)) } - - // Verify embeddings are deterministic - embeddings2, _ := manager.GenerateEmbedding([]string{"text 1"}) - for i := range embeddings[0] { - if embeddings[0][i] != embeddings2[0][i] { - t.Error("Embeddings should be deterministic") - break - } - } } diff --git a/pkg/optimizer/embeddings/openai_compatible_test.go b/pkg/optimizer/embeddings/openai_compatible_test.go index 916ad0cb8f..e829d2d6ac 100644 --- a/pkg/optimizer/embeddings/openai_compatible_test.go +++ b/pkg/optimizer/embeddings/openai_compatible_test.go @@ -206,30 +206,18 @@ func TestManagerWithUnified(t *testing.T) { func TestManagerFallbackBehavior(t *testing.T) { t.Parallel() - // Test that invalid vLLM backend falls back to placeholder + // Test that invalid vLLM backend fails gracefully during initialization + // (No fallback behavior is currently implemented) config := &Config{ BackendType: "vllm", - BaseURL: "http://invalid-host-that-does-not-exist:99999", + BaseURL: "http://invalid-host-that-does-not-exist:9999", Model: "test-model", Dimension: 384, } - manager, err := NewManager(config) - if err != nil { - t.Fatalf("Failed to create manager: %v", err) - } - defer manager.Close() - - // Should still work with placeholder fallback - embeddings, err := manager.GenerateEmbedding([]string{"test"}) - if err != nil { - t.Fatalf("Failed to generate embeddings with fallback: %v", err) - } - - if len(embeddings) != 1 { - t.Errorf("Expected 1 embedding, got %d", len(embeddings)) - } - if len(embeddings[0]) != 384 { - t.Errorf("Expected dimension 384, got %d", len(embeddings[0])) + _, err := NewManager(config) + if err == nil { + t.Error("Expected error when creating manager with invalid backend URL") } + // Test passes if error is returned (no fallback behavior) } diff --git a/pkg/optimizer/ingestion/service.go b/pkg/optimizer/ingestion/service.go index 821f970d6f..9b63e01289 100644 --- a/pkg/optimizer/ingestion/service.go +++ b/pkg/optimizer/ingestion/service.go @@ -65,6 +65,11 @@ func NewService(config *Config) (*Service, error) { return nil, fmt.Errorf("failed to initialize database: %w", err) } + // Clear database on startup to ensure fresh embeddings + // This is important when the embedding model changes or for consistency + database.Reset() + logger.Info("Cleared optimizer database on startup") + // Initialize embedding manager embeddingManager, err := embeddings.NewManager(config.EmbeddingConfig) if err != nil { @@ -124,7 +129,7 @@ func (s *Service) IngestServer( description *string, tools []mcp.Tool, ) error { - logger.Infof("Ingesting server: %s (%d tools)", serverName, len(tools)) + logger.Infof("Ingesting server: %s (%d tools) [serverID=%s]", serverName, len(tools), serverID) // Create backend server record (simplified - vMCP manages lifecycle) // chromem-go will generate embeddings automatically from the content @@ -155,6 +160,7 @@ func (s *Service) IngestServer( // syncBackendTools synchronizes tools for a backend server func (s *Service) syncBackendTools(ctx context.Context, serverID string, serverName string, tools []mcp.Tool) (int, error) { + logger.Debugf("syncBackendTools: server=%s, serverID=%s, tool_count=%d", serverName, serverID, len(tools)) // Delete existing tools if err := s.backendToolOps.DeleteByServer(ctx, serverID); err != nil { return 0, fmt.Errorf("failed to delete existing tools: %w", err) @@ -195,6 +201,33 @@ func (s *Service) syncBackendTools(ctx context.Context, serverID string, serverN return len(tools), nil } +// GetEmbeddingManager returns the embedding manager for this service +func (s *Service) GetEmbeddingManager() *embeddings.Manager { + return s.embeddingManager +} + +// GetBackendToolOps returns the backend tool operations for search and retrieval +func (s *Service) GetBackendToolOps() *db.BackendToolOps { + return s.backendToolOps +} + +// GetTotalToolTokens returns the total token count across all tools in the database +func (s *Service) GetTotalToolTokens(ctx context.Context) int { + // Use FTS database to efficiently count all tool tokens + if s.database.GetFTSDB() != nil { + totalTokens, err := s.database.GetFTSDB().GetTotalToolTokens(ctx) + if err != nil { + logger.Warnw("Failed to get total tool tokens from FTS", "error", err) + return 0 + } + return totalTokens + } + + // Fallback: query all tools (less efficient but works) + logger.Warn("FTS database not available, using fallback for token counting") + return 0 +} + // Close releases resources func (s *Service) Close() error { var errs []error diff --git a/pkg/optimizer/ingestion/service_test.go b/pkg/optimizer/ingestion/service_test.go index 51c73767b8..acc5b18754 100644 --- a/pkg/optimizer/ingestion/service_test.go +++ b/pkg/optimizer/ingestion/service_test.go @@ -25,14 +25,31 @@ func TestServiceCreationAndIngestion(t *testing.T) { // Create temporary directory for persistence (optional) tmpDir := t.TempDir() - // Initialize service with placeholder embeddings (no dependencies) + // Try to use Ollama if available, otherwise skip test + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull all-minilm'", err) + return + } + _ = embeddingManager.Close() + + // Initialize service with Ollama embeddings config := &Config{ DBConfig: &db.Config{ PersistPath: filepath.Join(tmpDir, "test-db"), }, EmbeddingConfig: &embeddings.Config{ - BackendType: "placeholder", // Use placeholder for testing - Dimension: 384, + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "nomic-embed-text", + Dimension: 768, }, } @@ -78,11 +95,11 @@ func TestServiceCreationAndIngestion(t *testing.T) { require.NoError(t, err) require.NotEmpty(t, results, "Should find at least one similar tool") - // With placeholder embeddings (hash-based), semantic similarity isn't guaranteed - // Just verify we got results back - require.Len(t, results, 2, "Should return both tools") + require.NotEmpty(t, results, "Should return at least one result") - // Verify both tools are present (order doesn't matter with placeholder embeddings) + // Weather tool should be most similar to weather query + require.Equal(t, "get_weather", results[0].ToolName, + "Weather tool should be most similar to weather query") toolNamesFound := make(map[string]bool) for _, result := range results { toolNamesFound[result.ToolName] = true @@ -142,7 +159,6 @@ func TestServiceWithOllama(t *testing.T) { require.NoError(t, err) require.NotEmpty(t, results) - // With real embeddings, weather tool should be most similar require.Equal(t, "get_weather", results[0].ToolName, "Weather tool should be most similar to weather query") } diff --git a/pkg/optimizer/ingestion/service_test_coverage.go b/pkg/optimizer/ingestion/service_test_coverage.go new file mode 100644 index 0000000000..2328db7120 --- /dev/null +++ b/pkg/optimizer/ingestion/service_test_coverage.go @@ -0,0 +1,282 @@ +package ingestion + +import ( + "context" + "path/filepath" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/optimizer/db" + "github.com/stacklok/toolhive/pkg/optimizer/embeddings" +) + +// TestService_GetTotalToolTokens tests token counting +func TestService_GetTotalToolTokens(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + config := &Config{ + DBConfig: &db.Config{ + PersistPath: filepath.Join(tmpDir, "test-db"), + }, + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + svc, err := NewService(config) + require.NoError(t, err) + defer func() { _ = svc.Close() }() + + // Ingest some tools + tools := []mcp.Tool{ + { + Name: "tool1", + Description: "Tool 1", + }, + { + Name: "tool2", + Description: "Tool 2", + }, + } + + err = svc.IngestServer(ctx, "server-1", "TestServer", nil, tools) + require.NoError(t, err) + + // Get total tokens + totalTokens := svc.GetTotalToolTokens(ctx) + assert.GreaterOrEqual(t, totalTokens, 0, "Total tokens should be non-negative") +} + +// TestService_GetTotalToolTokens_NoFTS tests token counting without FTS +func TestService_GetTotalToolTokens_NoFTS(t *testing.T) { + t.Parallel() + ctx := context.Background() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + config := &Config{ + DBConfig: &db.Config{ + PersistPath: "", // In-memory + FTSDBPath: "", // Will default to :memory: + }, + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + svc, err := NewService(config) + require.NoError(t, err) + defer func() { _ = svc.Close() }() + + // Get total tokens (should use FTS if available, fallback otherwise) + totalTokens := svc.GetTotalToolTokens(ctx) + assert.GreaterOrEqual(t, totalTokens, 0, "Total tokens should be non-negative") +} + +// TestService_GetBackendToolOps tests backend tool ops accessor +func TestService_GetBackendToolOps(t *testing.T) { + t.Parallel() + tmpDir := t.TempDir() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + config := &Config{ + DBConfig: &db.Config{ + PersistPath: filepath.Join(tmpDir, "test-db"), + }, + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + svc, err := NewService(config) + require.NoError(t, err) + defer func() { _ = svc.Close() }() + + toolOps := svc.GetBackendToolOps() + require.NotNil(t, toolOps) +} + +// TestService_GetEmbeddingManager tests embedding manager accessor +func TestService_GetEmbeddingManager(t *testing.T) { + t.Parallel() + tmpDir := t.TempDir() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + config := &Config{ + DBConfig: &db.Config{ + PersistPath: filepath.Join(tmpDir, "test-db"), + }, + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + svc, err := NewService(config) + require.NoError(t, err) + defer func() { _ = svc.Close() }() + + manager := svc.GetEmbeddingManager() + require.NotNil(t, manager) +} + +// TestService_IngestServer_ErrorHandling tests error handling during ingestion +func TestService_IngestServer_ErrorHandling(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + config := &Config{ + DBConfig: &db.Config{ + PersistPath: filepath.Join(tmpDir, "test-db"), + }, + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + svc, err := NewService(config) + require.NoError(t, err) + defer func() { _ = svc.Close() }() + + // Test with empty tools list + err = svc.IngestServer(ctx, "server-1", "TestServer", nil, []mcp.Tool{}) + require.NoError(t, err, "Should handle empty tools list gracefully") + + // Test with nil description + err = svc.IngestServer(ctx, "server-2", "TestServer2", nil, []mcp.Tool{ + { + Name: "tool1", + Description: "Tool 1", + }, + }) + require.NoError(t, err, "Should handle nil description gracefully") +} + +// TestService_Close_ErrorHandling tests error handling during close +func TestService_Close_ErrorHandling(t *testing.T) { + t.Parallel() + tmpDir := t.TempDir() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + config := &Config{ + DBConfig: &db.Config{ + PersistPath: filepath.Join(tmpDir, "test-db"), + }, + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + svc, err := NewService(config) + require.NoError(t, err) + + // Close should succeed + err = svc.Close() + require.NoError(t, err) + + // Multiple closes should be safe + err = svc.Close() + require.NoError(t, err) +} diff --git a/pkg/vmcp/client/client.go b/pkg/vmcp/client/client.go index 09d69a2fda..e99533a83a 100644 --- a/pkg/vmcp/client/client.go +++ b/pkg/vmcp/client/client.go @@ -123,8 +123,6 @@ func (a *authRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) return nil, fmt.Errorf("authentication failed for backend %s: %w", a.target.WorkloadID, err) } - logger.Debugf("Applied authentication strategy %q to backend %s", a.authStrategy.Name(), a.target.WorkloadID) - return a.base.RoundTrip(reqClone) } diff --git a/pkg/vmcp/discovery/manager.go b/pkg/vmcp/discovery/manager.go index 9bdfdc1d39..86c6b82482 100644 --- a/pkg/vmcp/discovery/manager.go +++ b/pkg/vmcp/discovery/manager.go @@ -15,6 +15,8 @@ import ( "sync" "time" + "golang.org/x/sync/singleflight" + "github.com/stacklok/toolhive/pkg/auth" "github.com/stacklok/toolhive/pkg/logger" "github.com/stacklok/toolhive/pkg/vmcp" @@ -65,6 +67,9 @@ type DefaultManager struct { stopCh chan struct{} stopOnce sync.Once wg sync.WaitGroup + // singleFlight ensures only one aggregation happens per cache key at a time + // This prevents concurrent requests from all triggering aggregation + singleFlight singleflight.Group } // NewManager creates a new discovery manager with the given aggregator. @@ -128,6 +133,9 @@ func NewManagerWithRegistry(agg aggregator.Aggregator, registry vmcp.DynamicRegi // // The context must contain an authenticated user identity (set by auth middleware). // Returns ErrNoIdentity if user identity is not found in context. +// +// This method uses singleflight to ensure that concurrent requests for the same +// cache key only trigger one aggregation, preventing duplicate work. func (m *DefaultManager) Discover(ctx context.Context, backends []vmcp.Backend) (*aggregator.AggregatedCapabilities, error) { // Validate user identity is present (set by auth middleware) // This ensures discovery happens with proper user authentication context @@ -139,7 +147,7 @@ func (m *DefaultManager) Discover(ctx context.Context, backends []vmcp.Backend) // Generate cache key from user identity and backend set cacheKey := m.generateCacheKey(identity.Subject, backends) - // Check cache first + // Check cache first (with read lock) if caps := m.getCachedCapabilities(cacheKey); caps != nil { logger.Debugf("Cache hit for user %s (key: %s)", identity.Subject, cacheKey) return caps, nil @@ -147,16 +155,33 @@ func (m *DefaultManager) Discover(ctx context.Context, backends []vmcp.Backend) logger.Debugf("Cache miss - performing capability discovery for user: %s", identity.Subject) - // Cache miss - perform aggregation - caps, err := m.aggregator.AggregateCapabilities(ctx, backends) + // Use singleflight to ensure only one aggregation happens per cache key + // Even if multiple requests come in concurrently, they'll all wait for the same result + result, err, _ := m.singleFlight.Do(cacheKey, func() (interface{}, error) { + // Double-check cache after acquiring singleflight lock + // Another goroutine might have populated it while we were waiting + if caps := m.getCachedCapabilities(cacheKey); caps != nil { + logger.Debugf("Cache populated while waiting - returning cached result for user %s", identity.Subject) + return caps, nil + } + + // Perform aggregation + caps, err := m.aggregator.AggregateCapabilities(ctx, backends) + if err != nil { + return nil, fmt.Errorf("%w: %w", ErrDiscoveryFailed, err) + } + + // Cache the result (skips caching if at capacity and key doesn't exist) + m.cacheCapabilities(cacheKey, caps) + + return caps, nil + }) + if err != nil { - return nil, fmt.Errorf("%w: %w", ErrDiscoveryFailed, err) + return nil, err } - // Cache the result (skips caching if at capacity and key doesn't exist) - m.cacheCapabilities(cacheKey, caps) - - return caps, nil + return result.(*aggregator.AggregatedCapabilities), nil } // Stop gracefully stops the manager and cleans up resources. diff --git a/pkg/vmcp/discovery/manager_test_coverage.go b/pkg/vmcp/discovery/manager_test_coverage.go new file mode 100644 index 0000000000..2d31a9db56 --- /dev/null +++ b/pkg/vmcp/discovery/manager_test_coverage.go @@ -0,0 +1,173 @@ +package discovery + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/auth" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + aggmocks "github.com/stacklok/toolhive/pkg/vmcp/aggregator/mocks" + vmcpmocks "github.com/stacklok/toolhive/pkg/vmcp/mocks" +) + +// TestDefaultManager_CacheVersionMismatch tests cache invalidation on version mismatch +func TestDefaultManager_CacheVersionMismatch(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockAggregator := aggmocks.NewMockAggregator(ctrl) + mockRegistry := vmcpmocks.NewMockDynamicRegistry(ctrl) + + // First call - version 1 + mockRegistry.EXPECT().Version().Return(uint64(1)).Times(2) + mockAggregator.EXPECT(). + AggregateCapabilities(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil). + Times(1) + + manager, err := NewManagerWithRegistry(mockAggregator, mockRegistry) + require.NoError(t, err) + defer manager.Stop() + + ctx := context.WithValue(context.Background(), auth.IdentityContextKey{}, &auth.Identity{ + Subject: "user-1", + }) + + backends := []vmcp.Backend{ + {ID: "backend-1", Name: "Backend 1"}, + } + + // First discovery - should cache + caps1, err := manager.Discover(ctx, backends) + require.NoError(t, err) + require.NotNil(t, caps1) + + // Second discovery with same version - should use cache + mockRegistry.EXPECT().Version().Return(uint64(1)).Times(1) + caps2, err := manager.Discover(ctx, backends) + require.NoError(t, err) + require.NotNil(t, caps2) + + // Third discovery with different version - should invalidate cache + mockRegistry.EXPECT().Version().Return(uint64(2)).Times(2) + mockAggregator.EXPECT(). + AggregateCapabilities(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil). + Times(1) + + caps3, err := manager.Discover(ctx, backends) + require.NoError(t, err) + require.NotNil(t, caps3) +} + +// TestDefaultManager_CacheAtCapacity tests cache eviction when at capacity +func TestDefaultManager_CacheAtCapacity(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockAggregator := aggmocks.NewMockAggregator(ctrl) + + // Create many different cache keys to fill cache + mockAggregator.EXPECT(). + AggregateCapabilities(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil). + Times(maxCacheSize + 1) // One more than capacity + + manager, err := NewManager(mockAggregator) + require.NoError(t, err) + defer manager.Stop() + + // Fill cache to capacity + for i := 0; i < maxCacheSize; i++ { + ctx := context.WithValue(context.Background(), auth.IdentityContextKey{}, &auth.Identity{ + Subject: "user-" + string(rune(i)), + }) + + backends := []vmcp.Backend{ + {ID: "backend-" + string(rune(i)), Name: "Backend"}, + } + + _, err := manager.Discover(ctx, backends) + require.NoError(t, err) + } + + // Next discovery should not cache (at capacity) + ctx := context.WithValue(context.Background(), auth.IdentityContextKey{}, &auth.Identity{ + Subject: "user-new", + }) + + backends := []vmcp.Backend{ + {ID: "backend-new", Name: "Backend"}, + } + + _, err = manager.Discover(ctx, backends) + require.NoError(t, err) +} + +// TestDefaultManager_CacheAtCapacity_ExistingKey tests cache update when at capacity but key exists +func TestDefaultManager_CacheAtCapacity_ExistingKey(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockAggregator := aggmocks.NewMockAggregator(ctrl) + + // First call + mockAggregator.EXPECT(). + AggregateCapabilities(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil). + Times(1) + + manager, err := NewManager(mockAggregator) + require.NoError(t, err) + defer manager.Stop() + + ctx := context.WithValue(context.Background(), auth.IdentityContextKey{}, &auth.Identity{ + Subject: "user-1", + }) + + backends := []vmcp.Backend{ + {ID: "backend-1", Name: "Backend 1"}, + } + + // First discovery + _, err = manager.Discover(ctx, backends) + require.NoError(t, err) + + // Fill cache to capacity with other keys + for i := 0; i < maxCacheSize-1; i++ { + ctxOther := context.WithValue(context.Background(), auth.IdentityContextKey{}, &auth.Identity{ + Subject: "user-" + string(rune(i+2)), + }) + + backendsOther := []vmcp.Backend{ + {ID: "backend-" + string(rune(i+2)), Name: "Backend"}, + } + + mockAggregator.EXPECT(). + AggregateCapabilities(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil). + Times(1) + + _, err := manager.Discover(ctxOther, backendsOther) + require.NoError(t, err) + } + + // Update existing key should work even at capacity + mockAggregator.EXPECT(). + AggregateCapabilities(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil). + Times(1) + + _, err = manager.Discover(ctx, backends) + require.NoError(t, err) +} diff --git a/pkg/vmcp/health/checker.go b/pkg/vmcp/health/checker.go index 9705a0b788..c84a341fc4 100644 --- a/pkg/vmcp/health/checker.go +++ b/pkg/vmcp/health/checker.go @@ -8,6 +8,8 @@ import ( "context" "errors" "fmt" + "net/url" + "strings" "time" "github.com/stacklok/toolhive/pkg/logger" @@ -26,6 +28,10 @@ type healthChecker struct { // If a health check succeeds but takes longer than this duration, the backend is marked degraded. // Zero means disabled (backends will never be marked degraded based on response time alone). degradedThreshold time.Duration + + // selfURL is the server's own URL. If a health check targets this URL, it's short-circuited. + // This prevents the server from trying to health check itself. + selfURL string } // NewHealthChecker creates a new health checker that uses BackendClient.ListCapabilities @@ -36,13 +42,20 @@ type healthChecker struct { // - client: BackendClient for communicating with backend MCP servers // - timeout: Maximum duration for health check operations (0 = no timeout) // - degradedThreshold: Response time threshold for marking backend as degraded (0 = disabled) +// - selfURL: Optional server's own URL. If provided, health checks targeting this URL are short-circuited. // // Returns a new HealthChecker implementation. -func NewHealthChecker(client vmcp.BackendClient, timeout time.Duration, degradedThreshold time.Duration) vmcp.HealthChecker { +func NewHealthChecker( + client vmcp.BackendClient, + timeout time.Duration, + degradedThreshold time.Duration, + selfURL string, +) vmcp.HealthChecker { return &healthChecker{ client: client, timeout: timeout, degradedThreshold: degradedThreshold, + selfURL: selfURL, } } @@ -59,16 +72,28 @@ func NewHealthChecker(client vmcp.BackendClient, timeout time.Duration, degraded // The error return is informational and provides context about what failed. // The BackendHealthStatus return indicates the categorized health state. func (h *healthChecker) CheckHealth(ctx context.Context, target *vmcp.BackendTarget) (vmcp.BackendHealthStatus, error) { - // Apply timeout if configured - checkCtx := ctx + // Mark context as health check to bypass authentication logging + // Health checks verify backend availability and should not require user credentials + healthCheckCtx := WithHealthCheckMarker(ctx) + + // Apply timeout if configured (after adding health check marker) + checkCtx := healthCheckCtx var cancel context.CancelFunc if h.timeout > 0 { - checkCtx, cancel = context.WithTimeout(ctx, h.timeout) + checkCtx, cancel = context.WithTimeout(healthCheckCtx, h.timeout) defer cancel() } logger.Debugf("Performing health check for backend %s (%s)", target.WorkloadName, target.BaseURL) + // Short-circuit health check if targeting ourselves + // This prevents the server from trying to health check itself, which would work + // but is wasteful and can cause connection issues during startup + if h.selfURL != "" && h.isSelfCheck(target.BaseURL) { + logger.Debugf("Skipping health check for backend %s - this is the server itself", target.WorkloadName) + return vmcp.BackendHealthy, nil + } + // Track response time for degraded detection startTime := time.Now() @@ -134,3 +159,62 @@ func categorizeError(err error) vmcp.BackendHealthStatus { // Default to unhealthy for unknown errors return vmcp.BackendUnhealthy } + +// isSelfCheck checks if a backend URL matches the server's own URL. +// URLs are normalized before comparison to handle variations like: +// - http://127.0.0.1:PORT vs http://localhost:PORT +// - http://HOST:PORT vs http://HOST:PORT/ +func (h *healthChecker) isSelfCheck(backendURL string) bool { + if h.selfURL == "" || backendURL == "" { + return false + } + + // Normalize both URLs for comparison + backendNormalized, err := NormalizeURLForComparison(backendURL) + if err != nil { + return false + } + + selfNormalized, err := NormalizeURLForComparison(h.selfURL) + if err != nil { + return false + } + + return backendNormalized == selfNormalized +} + +// NormalizeURLForComparison normalizes a URL for comparison by: +// - Parsing and reconstructing the URL +// - Converting localhost/127.0.0.1 to a canonical form +// - Comparing only scheme://host:port (ignoring path, query, fragment) +// - Lowercasing scheme and host +// Exported for testing purposes +func NormalizeURLForComparison(rawURL string) (string, error) { + u, err := url.Parse(rawURL) + if err != nil { + return "", err + } + // Validate that we have a scheme and host (basic URL validation) + if u.Scheme == "" || u.Host == "" { + return "", fmt.Errorf("invalid URL: missing scheme or host") + } + + // Normalize host: convert localhost to 127.0.0.1 for consistency + host := strings.ToLower(u.Hostname()) + if host == "localhost" { + host = "127.0.0.1" + } + + // Reconstruct URL with normalized components (scheme://host:port only) + // We ignore path, query, and fragment for comparison + normalized := &url.URL{ + Scheme: strings.ToLower(u.Scheme), + } + if u.Port() != "" { + normalized.Host = host + ":" + u.Port() + } else { + normalized.Host = host + } + + return normalized.String(), nil +} diff --git a/pkg/vmcp/health/checker_selfcheck_test.go b/pkg/vmcp/health/checker_selfcheck_test.go new file mode 100644 index 0000000000..fc42b071f0 --- /dev/null +++ b/pkg/vmcp/health/checker_selfcheck_test.go @@ -0,0 +1,501 @@ +package health + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" +) + +// TestHealthChecker_CheckHealth_SelfCheck tests self-check detection +func TestHealthChecker_CheckHealth_SelfCheck(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + // Should not call ListCapabilities for self-check + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Times(0) + + checker := NewHealthChecker(mockClient, 5*time.Second, 0, "http://127.0.0.1:8080") + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "test-backend", + BaseURL: "http://127.0.0.1:8080", // Same as selfURL + } + + status, err := checker.CheckHealth(context.Background(), target) + assert.NoError(t, err) + assert.Equal(t, vmcp.BackendHealthy, status) +} + +// TestHealthChecker_CheckHealth_SelfCheck_Localhost tests localhost normalization +func TestHealthChecker_CheckHealth_SelfCheck_Localhost(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Times(0) + + checker := NewHealthChecker(mockClient, 5*time.Second, 0, "http://localhost:8080") + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "test-backend", + BaseURL: "http://127.0.0.1:8080", // localhost should match 127.0.0.1 + } + + status, err := checker.CheckHealth(context.Background(), target) + assert.NoError(t, err) + assert.Equal(t, vmcp.BackendHealthy, status) +} + +// TestHealthChecker_CheckHealth_SelfCheck_Reverse tests reverse localhost normalization +func TestHealthChecker_CheckHealth_SelfCheck_Reverse(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Times(0) + + checker := NewHealthChecker(mockClient, 5*time.Second, 0, "http://127.0.0.1:8080") + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "test-backend", + BaseURL: "http://localhost:8080", // 127.0.0.1 should match localhost + } + + status, err := checker.CheckHealth(context.Background(), target) + assert.NoError(t, err) + assert.Equal(t, vmcp.BackendHealthy, status) +} + +// TestHealthChecker_CheckHealth_SelfCheck_DifferentPort tests different ports don't match +func TestHealthChecker_CheckHealth_SelfCheck_DifferentPort(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil). + Times(1) + + checker := NewHealthChecker(mockClient, 5*time.Second, 0, "http://127.0.0.1:8080") + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "test-backend", + BaseURL: "http://127.0.0.1:8081", // Different port + } + + status, err := checker.CheckHealth(context.Background(), target) + assert.NoError(t, err) + assert.Equal(t, vmcp.BackendHealthy, status) +} + +// TestHealthChecker_CheckHealth_SelfCheck_EmptyURL tests empty URLs +func TestHealthChecker_CheckHealth_SelfCheck_EmptyURL(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil). + Times(1) + + checker := NewHealthChecker(mockClient, 5*time.Second, 0, "") + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "test-backend", + BaseURL: "http://127.0.0.1:8080", + } + + status, err := checker.CheckHealth(context.Background(), target) + assert.NoError(t, err) + assert.Equal(t, vmcp.BackendHealthy, status) +} + +// TestHealthChecker_CheckHealth_SelfCheck_InvalidURL tests invalid URLs +func TestHealthChecker_CheckHealth_SelfCheck_InvalidURL(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil). + Times(1) + + checker := NewHealthChecker(mockClient, 5*time.Second, 0, "not-a-valid-url") + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "test-backend", + BaseURL: "http://127.0.0.1:8080", + } + + status, err := checker.CheckHealth(context.Background(), target) + assert.NoError(t, err) + assert.Equal(t, vmcp.BackendHealthy, status) +} + +// TestHealthChecker_CheckHealth_SelfCheck_WithPath tests URLs with paths are normalized +func TestHealthChecker_CheckHealth_SelfCheck_WithPath(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Times(0) + + checker := NewHealthChecker(mockClient, 5*time.Second, 0, "http://127.0.0.1:8080") + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "test-backend", + BaseURL: "http://127.0.0.1:8080/mcp", // Path should be ignored + } + + status, err := checker.CheckHealth(context.Background(), target) + assert.NoError(t, err) + assert.Equal(t, vmcp.BackendHealthy, status) +} + +// TestHealthChecker_CheckHealth_DegradedThreshold tests degraded threshold detection +func TestHealthChecker_CheckHealth_DegradedThreshold(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + // Simulate slow response + time.Sleep(150 * time.Millisecond) + return &vmcp.CapabilityList{}, nil + }). + Times(1) + + // Set degraded threshold to 100ms + checker := NewHealthChecker(mockClient, 5*time.Second, 100*time.Millisecond, "") + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "test-backend", + BaseURL: "http://localhost:8080", + } + + status, err := checker.CheckHealth(context.Background(), target) + assert.NoError(t, err) + assert.Equal(t, vmcp.BackendDegraded, status, "Should mark as degraded when response time exceeds threshold") +} + +// TestHealthChecker_CheckHealth_DegradedThreshold_Disabled tests disabled degraded threshold +func TestHealthChecker_CheckHealth_DegradedThreshold_Disabled(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, _ *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + // Simulate slow response + time.Sleep(150 * time.Millisecond) + return &vmcp.CapabilityList{}, nil + }). + Times(1) + + // Set degraded threshold to 0 (disabled) + checker := NewHealthChecker(mockClient, 5*time.Second, 0, "") + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "test-backend", + BaseURL: "http://localhost:8080", + } + + status, err := checker.CheckHealth(context.Background(), target) + assert.NoError(t, err) + assert.Equal(t, vmcp.BackendHealthy, status, "Should not mark as degraded when threshold is disabled") +} + +// TestHealthChecker_CheckHealth_DegradedThreshold_FastResponse tests fast response doesn't trigger degraded +func TestHealthChecker_CheckHealth_DegradedThreshold_FastResponse(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockBackendClient(ctrl) + mockClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil). + Times(1) + + // Set degraded threshold to 100ms + checker := NewHealthChecker(mockClient, 5*time.Second, 100*time.Millisecond, "") + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "test-backend", + BaseURL: "http://localhost:8080", + } + + status, err := checker.CheckHealth(context.Background(), target) + assert.NoError(t, err) + assert.Equal(t, vmcp.BackendHealthy, status, "Should not mark as degraded when response is fast") +} + +// TestCategorizeError_SentinelErrors tests sentinel error categorization +func TestCategorizeError_SentinelErrors(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + expectedStatus vmcp.BackendHealthStatus + }{ + { + name: "ErrAuthenticationFailed", + err: vmcp.ErrAuthenticationFailed, + expectedStatus: vmcp.BackendUnauthenticated, + }, + { + name: "ErrAuthorizationFailed", + err: vmcp.ErrAuthorizationFailed, + expectedStatus: vmcp.BackendUnauthenticated, + }, + { + name: "ErrTimeout", + err: vmcp.ErrTimeout, + expectedStatus: vmcp.BackendUnhealthy, + }, + { + name: "ErrCancelled", + err: vmcp.ErrCancelled, + expectedStatus: vmcp.BackendUnhealthy, + }, + { + name: "ErrBackendUnavailable", + err: vmcp.ErrBackendUnavailable, + expectedStatus: vmcp.BackendUnhealthy, + }, + { + name: "wrapped ErrAuthenticationFailed", + err: errors.New("wrapped: " + vmcp.ErrAuthenticationFailed.Error()), + expectedStatus: vmcp.BackendUnauthenticated, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + status := categorizeError(tt.err) + assert.Equal(t, tt.expectedStatus, status) + }) + } +} + +// TestNormalizeURLForComparison tests URL normalization +func TestNormalizeURLForComparison(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input string + expected string + wantErr bool + }{ + { + name: "localhost normalized to 127.0.0.1", + input: "http://localhost:8080", + expected: "http://127.0.0.1:8080", + wantErr: false, + }, + { + name: "127.0.0.1 stays as is", + input: "http://127.0.0.1:8080", + expected: "http://127.0.0.1:8080", + wantErr: false, + }, + { + name: "path is ignored", + input: "http://127.0.0.1:8080/mcp", + expected: "http://127.0.0.1:8080", + wantErr: false, + }, + { + name: "query is ignored", + input: "http://127.0.0.1:8080?param=value", + expected: "http://127.0.0.1:8080", + wantErr: false, + }, + { + name: "fragment is ignored", + input: "http://127.0.0.1:8080#fragment", + expected: "http://127.0.0.1:8080", + wantErr: false, + }, + { + name: "scheme is lowercased", + input: "HTTP://127.0.0.1:8080", + expected: "http://127.0.0.1:8080", + wantErr: false, + }, + { + name: "host is lowercased", + input: "http://EXAMPLE.COM:8080", + expected: "http://example.com:8080", + wantErr: false, + }, + { + name: "no port", + input: "http://127.0.0.1", + expected: "http://127.0.0.1", + wantErr: false, + }, + { + name: "invalid URL", + input: "not-a-url", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + result, err := NormalizeURLForComparison(tt.input) + if tt.wantErr { + assert.Error(t, err) + } else { + assert.NoError(t, err) + assert.Equal(t, tt.expected, result) + } + }) + } +} + +// TestIsSelfCheck_EdgeCases tests edge cases for self-check detection +func TestIsSelfCheck_EdgeCases(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + t.Cleanup(func() { ctrl.Finish() }) + + mockClient := mocks.NewMockBackendClient(ctrl) + + tests := []struct { + name string + selfURL string + backendURL string + expected bool + }{ + { + name: "both empty", + selfURL: "", + backendURL: "", + expected: false, + }, + { + name: "selfURL empty", + selfURL: "", + backendURL: "http://127.0.0.1:8080", + expected: false, + }, + { + name: "backendURL empty", + selfURL: "http://127.0.0.1:8080", + backendURL: "", + expected: false, + }, + { + name: "localhost matches 127.0.0.1", + selfURL: "http://localhost:8080", + backendURL: "http://127.0.0.1:8080", + expected: true, + }, + { + name: "127.0.0.1 matches localhost", + selfURL: "http://127.0.0.1:8080", + backendURL: "http://localhost:8080", + expected: true, + }, + { + name: "different ports", + selfURL: "http://127.0.0.1:8080", + backendURL: "http://127.0.0.1:8081", + expected: false, + }, + { + name: "different hosts", + selfURL: "http://127.0.0.1:8080", + backendURL: "http://192.168.1.1:8080", + expected: false, + }, + { + name: "path ignored", + selfURL: "http://127.0.0.1:8080", + backendURL: "http://127.0.0.1:8080/mcp", + expected: true, + }, + { + name: "query ignored", + selfURL: "http://127.0.0.1:8080", + backendURL: "http://127.0.0.1:8080?param=value", + expected: true, + }, + { + name: "invalid selfURL", + selfURL: "not-a-url", + backendURL: "http://127.0.0.1:8080", + expected: false, + }, + { + name: "invalid backendURL", + selfURL: "http://127.0.0.1:8080", + backendURL: "not-a-url", + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + checker := NewHealthChecker(mockClient, 5*time.Second, 0, tt.selfURL) + hc, ok := checker.(*healthChecker) + require.True(t, ok) + + result := hc.isSelfCheck(tt.backendURL) + assert.Equal(t, tt.expected, result) + }) + } +} diff --git a/pkg/vmcp/health/checker_test.go b/pkg/vmcp/health/checker_test.go index a0515cb3c2..818021bd33 100644 --- a/pkg/vmcp/health/checker_test.go +++ b/pkg/vmcp/health/checker_test.go @@ -41,7 +41,7 @@ func TestNewHealthChecker(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - checker := NewHealthChecker(mockClient, tt.timeout, 0) + checker := NewHealthChecker(mockClient, tt.timeout, 0, "") require.NotNil(t, checker) // Type assert to access internals for verification @@ -65,7 +65,7 @@ func TestHealthChecker_CheckHealth_Success(t *testing.T) { Return(&vmcp.CapabilityList{}, nil). Times(1) - checker := NewHealthChecker(mockClient, 5*time.Second, 0) + checker := NewHealthChecker(mockClient, 5*time.Second, 0, "") target := &vmcp.BackendTarget{ WorkloadID: "backend-1", WorkloadName: "test-backend", @@ -92,7 +92,7 @@ func TestHealthChecker_CheckHealth_ContextCancellation(t *testing.T) { }). Times(1) - checker := NewHealthChecker(mockClient, 100*time.Millisecond, 0) + checker := NewHealthChecker(mockClient, 100*time.Millisecond, 0, "") target := &vmcp.BackendTarget{ WorkloadID: "backend-1", WorkloadName: "test-backend", @@ -120,7 +120,7 @@ func TestHealthChecker_CheckHealth_NoTimeout(t *testing.T) { Times(1) // Create checker with no timeout - checker := NewHealthChecker(mockClient, 0, 0) + checker := NewHealthChecker(mockClient, 0, 0, "") target := &vmcp.BackendTarget{ WorkloadID: "backend-1", WorkloadName: "test-backend", @@ -210,7 +210,7 @@ func TestHealthChecker_CheckHealth_ErrorCategorization(t *testing.T) { Return(nil, tt.err). Times(1) - checker := NewHealthChecker(mockClient, 5*time.Second, 0) + checker := NewHealthChecker(mockClient, 5*time.Second, 0, "") target := &vmcp.BackendTarget{ WorkloadID: "backend-1", WorkloadName: "test-backend", @@ -427,7 +427,7 @@ func TestHealthChecker_CheckHealth_Timeout(t *testing.T) { }). Times(1) - checker := NewHealthChecker(mockClient, 100*time.Millisecond, 0) + checker := NewHealthChecker(mockClient, 100*time.Millisecond, 0, "") target := &vmcp.BackendTarget{ WorkloadID: "backend-1", WorkloadName: "test-backend", @@ -464,7 +464,7 @@ func TestHealthChecker_CheckHealth_MultipleBackends(t *testing.T) { }). Times(4) - checker := NewHealthChecker(mockClient, 5*time.Second, 0) + checker := NewHealthChecker(mockClient, 5*time.Second, 0, "") // Test healthy backend status, err := checker.CheckHealth(context.Background(), &vmcp.BackendTarget{ diff --git a/pkg/vmcp/health/monitor.go b/pkg/vmcp/health/monitor.go index aa6a26240f..ee49e5bf70 100644 --- a/pkg/vmcp/health/monitor.go +++ b/pkg/vmcp/health/monitor.go @@ -105,12 +105,14 @@ func DefaultConfig() MonitorConfig { // - client: BackendClient for communicating with backend MCP servers // - backends: List of backends to monitor // - config: Configuration for health monitoring +// - selfURL: Optional server's own URL. If provided, health checks targeting this URL are short-circuited. // // Returns (monitor, error). Error is returned if configuration is invalid. func NewMonitor( client vmcp.BackendClient, backends []vmcp.Backend, config MonitorConfig, + selfURL string, ) (*Monitor, error) { // Validate configuration if config.CheckInterval <= 0 { @@ -120,8 +122,8 @@ func NewMonitor( return nil, fmt.Errorf("unhealthy threshold must be >= 1, got %d", config.UnhealthyThreshold) } - // Create health checker with degraded threshold - checker := NewHealthChecker(client, config.Timeout, config.DegradedThreshold) + // Create health checker with degraded threshold and self URL + checker := NewHealthChecker(client, config.Timeout, config.DegradedThreshold, selfURL) // Create status tracker statusTracker := newStatusTracker(config.UnhealthyThreshold) diff --git a/pkg/vmcp/health/monitor_test.go b/pkg/vmcp/health/monitor_test.go index 0bb74f163f..36defadd04 100644 --- a/pkg/vmcp/health/monitor_test.go +++ b/pkg/vmcp/health/monitor_test.go @@ -63,7 +63,7 @@ func TestNewMonitor_Validation(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - monitor, err := NewMonitor(mockClient, backends, tt.config) + monitor, err := NewMonitor(mockClient, backends, tt.config, "") if tt.expectError { assert.Error(t, err) assert.Nil(t, monitor) @@ -98,7 +98,7 @@ func TestMonitor_StartStop(t *testing.T) { Return(&vmcp.CapabilityList{}, nil). AnyTimes() - monitor, err := NewMonitor(mockClient, backends, config) + monitor, err := NewMonitor(mockClient, backends, config, "") require.NoError(t, err) // Start monitor @@ -175,7 +175,7 @@ func TestMonitor_StartErrors(t *testing.T) { t.Run(tt.name, func(t *testing.T) { t.Parallel() - monitor, err := NewMonitor(mockClient, backends, config) + monitor, err := NewMonitor(mockClient, backends, config, "") require.NoError(t, err) err = tt.setupFunc(monitor) @@ -205,7 +205,7 @@ func TestMonitor_StopWithoutStart(t *testing.T) { Timeout: 50 * time.Millisecond, } - monitor, err := NewMonitor(mockClient, backends, config) + monitor, err := NewMonitor(mockClient, backends, config, "") require.NoError(t, err) // Try to stop without starting @@ -236,7 +236,7 @@ func TestMonitor_PeriodicHealthChecks(t *testing.T) { Return(nil, errors.New("backend unavailable")). MinTimes(2) - monitor, err := NewMonitor(mockClient, backends, config) + monitor, err := NewMonitor(mockClient, backends, config, "") require.NoError(t, err) ctx := context.Background() @@ -286,7 +286,7 @@ func TestMonitor_GetHealthSummary(t *testing.T) { }). AnyTimes() - monitor, err := NewMonitor(mockClient, backends, config) + monitor, err := NewMonitor(mockClient, backends, config, "") require.NoError(t, err) ctx := context.Background() @@ -330,7 +330,7 @@ func TestMonitor_GetBackendStatus(t *testing.T) { Return(&vmcp.CapabilityList{}, nil). AnyTimes() - monitor, err := NewMonitor(mockClient, backends, config) + monitor, err := NewMonitor(mockClient, backends, config, "") require.NoError(t, err) ctx := context.Background() @@ -379,7 +379,7 @@ func TestMonitor_GetBackendState(t *testing.T) { Return(&vmcp.CapabilityList{}, nil). AnyTimes() - monitor, err := NewMonitor(mockClient, backends, config) + monitor, err := NewMonitor(mockClient, backends, config, "") require.NoError(t, err) ctx := context.Background() @@ -430,7 +430,7 @@ func TestMonitor_GetAllBackendStates(t *testing.T) { Return(&vmcp.CapabilityList{}, nil). AnyTimes() - monitor, err := NewMonitor(mockClient, backends, config) + monitor, err := NewMonitor(mockClient, backends, config, "") require.NoError(t, err) ctx := context.Background() @@ -474,7 +474,7 @@ func TestMonitor_ContextCancellation(t *testing.T) { Return(&vmcp.CapabilityList{}, nil). AnyTimes() - monitor, err := NewMonitor(mockClient, backends, config) + monitor, err := NewMonitor(mockClient, backends, config, "") require.NoError(t, err) // Start with cancellable context diff --git a/pkg/vmcp/optimizer/find_tool_semantic_search_test.go b/pkg/vmcp/optimizer/find_tool_semantic_search_test.go new file mode 100644 index 0000000000..a539937fe9 --- /dev/null +++ b/pkg/vmcp/optimizer/find_tool_semantic_search_test.go @@ -0,0 +1,690 @@ +package optimizer + +import ( + "context" + "encoding/json" + "path/filepath" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/optimizer/embeddings" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + "github.com/stacklok/toolhive/pkg/vmcp/discovery" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +const ( + testBackendOllama = "ollama" + testBackendOpenAI = "openai" +) + +// verifyEmbeddingBackendWorking verifies that the embedding backend is actually working by attempting to generate an embedding +// This ensures the service is not just reachable but actually functional +func verifyEmbeddingBackendWorking(t *testing.T, manager *embeddings.Manager, backendType string) { + t.Helper() + _, err := manager.GenerateEmbedding([]string{"test"}) + if err != nil { + if backendType == testBackendOllama { + t.Skipf("Skipping test: Ollama is reachable but embedding generation failed. Error: %v. Ensure 'ollama pull %s' has been executed", err, embeddings.DefaultModelAllMiniLM) + } else { + t.Skipf("Skipping test: Embedding backend is reachable but embedding generation failed. Error: %v", err) + } + } +} + +// TestFindTool_SemanticSearch tests semantic search capabilities +// These tests verify that find_tool can find tools based on semantic meaning, +// not just exact keyword matches +func TestFindTool_SemanticSearch(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Try to use Ollama if available, otherwise skip test + embeddingBackend := testBackendOllama + embeddingConfig := &embeddings.Config{ + BackendType: embeddingBackend, + BaseURL: "http://localhost:11434", + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, // all-MiniLM-L6-v2 dimension + } + + // Test if Ollama is available + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + // Try OpenAI-compatible (might be vLLM or Ollama v1 API) + embeddingConfig.BackendType = testBackendOpenAI + embeddingConfig.BaseURL = "http://localhost:11434" + embeddingConfig.Model = embeddings.DefaultModelAllMiniLM + embeddingConfig.Dimension = 768 + embeddingManager, err = embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping semantic search test: No embedding backend available (Ollama or OpenAI-compatible). Error: %v", err) + return + } + embeddingBackend = testBackendOpenAI + } + t.Cleanup(func() { _ = embeddingManager.Close() }) + + // Verify embedding backend is actually working, not just reachable + verifyEmbeddingBackendWorking(t, embeddingManager, embeddingBackend) + + // Setup optimizer integration with high semantic ratio to favor semantic search + mcpServer := server.NewMCPServer("test-server", "1.0") + mockClient := &mockBackendClient{} + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: embeddingBackend, + BaseURL: embeddingConfig.BaseURL, + Model: embeddingConfig.Model, + Dimension: embeddingConfig.Dimension, + }, + HybridSearchRatio: 0.9, // 90% semantic, 10% BM25 to test semantic search + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer, mockClient, sessionMgr) + require.NoError(t, err) + require.NotNil(t, integration) + t.Cleanup(func() { _ = integration.Close() }) + + // Create tools with diverse descriptions to test semantic understanding + tools := []vmcp.Tool{ + { + Name: "github_pull_request_read", + Description: "Get information on a specific pull request in GitHub repository.", + BackendID: "github", + }, + { + Name: "github_list_pull_requests", + Description: "List pull requests in a GitHub repository.", + BackendID: "github", + }, + { + Name: "github_create_pull_request", + Description: "Create a new pull request in a GitHub repository.", + BackendID: "github", + }, + { + Name: "github_merge_pull_request", + Description: "Merge a pull request in a GitHub repository.", + BackendID: "github", + }, + { + Name: "github_issue_read", + Description: "Get information about a specific issue in a GitHub repository.", + BackendID: "github", + }, + { + Name: "github_list_issues", + Description: "List issues in a GitHub repository.", + BackendID: "github", + }, + { + Name: "github_create_repository", + Description: "Create a new GitHub repository in your account or specified organization", + BackendID: "github", + }, + { + Name: "github_get_commit", + Description: "Get details for a commit from a GitHub repository", + BackendID: "github", + }, + { + Name: "github_get_branch", + Description: "Get information about a branch in a GitHub repository", + BackendID: "github", + }, + { + Name: "fetch_fetch", + Description: "Fetches a URL from the internet and optionally extracts its contents as markdown.", + BackendID: "fetch", + }, + } + + capabilities := &aggregator.AggregatedCapabilities{ + Tools: tools, + RoutingTable: &vmcp.RoutingTable{ + Tools: make(map[string]*vmcp.BackendTarget), + Resources: map[string]*vmcp.BackendTarget{}, + Prompts: map[string]*vmcp.BackendTarget{}, + }, + } + + for _, tool := range tools { + capabilities.RoutingTable.Tools[tool.Name] = &vmcp.BackendTarget{ + WorkloadID: tool.BackendID, + WorkloadName: tool.BackendID, + } + } + + session := &mockSession{sessionID: "test-session"} + err = integration.OnRegisterSession(ctx, session, capabilities) + require.NoError(t, err) + + // Manually ingest tools for testing (OnRegisterSession skips ingestion) + mcpTools := make([]mcp.Tool, len(tools)) + for i, tool := range tools { + mcpTools[i] = mcp.Tool{ + Name: tool.Name, + Description: tool.Description, + } + } + err = integration.IngestToolsForTesting(ctx, "github", "GitHub", nil, mcpTools) + require.NoError(t, err) + + ctxWithCaps := discovery.WithDiscoveredCapabilities(ctx, capabilities) + + // Test cases for semantic search - queries that mean the same thing but use different words + testCases := []struct { + name string + query string + keywords string + expectedTools []string // Tools that should be found semantically + description string + }{ + { + name: "semantic_pr_synonyms", + query: "view code review request", + keywords: "", + expectedTools: []string{"github_pull_request_read", "github_list_pull_requests"}, + description: "Should find PR tools using semantic synonyms (code review = pull request)", + }, + { + name: "semantic_merge_synonyms", + query: "combine code changes", + keywords: "", + expectedTools: []string{"github_merge_pull_request"}, + description: "Should find merge tool using semantic meaning (combine = merge)", + }, + { + name: "semantic_create_synonyms", + query: "make a new code review", + keywords: "", + expectedTools: []string{"github_create_pull_request", "github_list_pull_requests", "github_pull_request_read"}, + description: "Should find PR-related tools using semantic meaning (make = create, code review = PR)", + }, + { + name: "semantic_issue_synonyms", + query: "show bug reports", + keywords: "", + expectedTools: []string{"github_issue_read", "github_list_issues"}, + description: "Should find issue tools using semantic synonyms (bug report = issue)", + }, + { + name: "semantic_repository_synonyms", + query: "start a new project", + keywords: "", + expectedTools: []string{"github_create_repository"}, + description: "Should find repository tool using semantic meaning (project = repository)", + }, + { + name: "semantic_commit_synonyms", + query: "get change details", + keywords: "", + expectedTools: []string{"github_get_commit"}, + description: "Should find commit tool using semantic meaning (change = commit)", + }, + { + name: "semantic_fetch_synonyms", + query: "download web page content", + keywords: "", + expectedTools: []string{"fetch_fetch"}, + description: "Should find fetch tool using semantic synonyms (download = fetch)", + }, + { + name: "semantic_branch_synonyms", + query: "get branch information", + keywords: "", + expectedTools: []string{"github_get_branch"}, + description: "Should find branch tool using semantic meaning", + }, + { + name: "semantic_related_concepts", + query: "code collaboration features", + keywords: "", + expectedTools: []string{"github_pull_request_read", "github_create_pull_request", "github_issue_read"}, + description: "Should find collaboration-related tools (PRs and issues are collaboration features)", + }, + { + name: "semantic_intent_based", + query: "I want to see what code changes were made", + keywords: "", + expectedTools: []string{"github_get_commit", "github_pull_request_read"}, + description: "Should find tools based on user intent (seeing code changes = commits/PRs)", + }, + } + + for _, tc := range testCases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": tc.query, + "tool_keywords": tc.keywords, + "limit": 10, + }, + }, + } + + handler := integration.CreateFindToolHandler() + result, err := handler(ctxWithCaps, request) + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, result.IsError, "Tool call should not return error for query: %s", tc.query) + + // Parse the result + require.NotEmpty(t, result.Content, "Result should have content") + textContent, okText := mcp.AsTextContent(result.Content[0]) + require.True(t, okText, "Result should be text content") + + var response map[string]any + err = json.Unmarshal([]byte(textContent.Text), &response) + require.NoError(t, err, "Result should be valid JSON") + + toolsArray, okArray := response["tools"].([]interface{}) + require.True(t, okArray, "Response should have tools array") + require.NotEmpty(t, toolsArray, "Should return at least one result for semantic query: %s", tc.query) + + // Extract tool names from results + foundTools := make([]string, 0, len(toolsArray)) + for _, toolInterface := range toolsArray { + toolMap, okMap := toolInterface.(map[string]interface{}) + require.True(t, okMap, "Tool should be a map") + toolName, okName := toolMap["name"].(string) + require.True(t, okName, "Tool should have name") + foundTools = append(foundTools, toolName) + + // Verify similarity score exists and is reasonable + similarity, okScore := toolMap["similarity_score"].(float64) + require.True(t, okScore, "Tool should have similarity_score") + assert.Greater(t, similarity, 0.0, "Similarity score should be positive") + } + + // Check that at least one expected tool is found + foundCount := 0 + for _, expectedTool := range tc.expectedTools { + for _, foundTool := range foundTools { + if foundTool == expectedTool { + foundCount++ + break + } + } + } + + assert.GreaterOrEqual(t, foundCount, 1, + "Semantic query '%s' should find at least one expected tool from %v. Found tools: %v (found %d/%d)", + tc.query, tc.expectedTools, foundTools, foundCount, len(tc.expectedTools)) + + // Log results for debugging + if foundCount < len(tc.expectedTools) { + t.Logf("Semantic query '%s': Found %d/%d expected tools. Found: %v, Expected: %v", + tc.query, foundCount, len(tc.expectedTools), foundTools, tc.expectedTools) + } + + // Verify token metrics exist + tokenMetrics, okMetrics := response["token_metrics"].(map[string]interface{}) + require.True(t, okMetrics, "Response should have token_metrics") + assert.Contains(t, tokenMetrics, "baseline_tokens") + assert.Contains(t, tokenMetrics, "returned_tokens") + }) + } +} + +// TestFindTool_SemanticVsKeyword tests that semantic search finds different results than keyword search +func TestFindTool_SemanticVsKeyword(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Try to use Ollama if available + embeddingBackend := "ollama" + embeddingConfig := &embeddings.Config{ + BackendType: embeddingBackend, + BaseURL: "http://localhost:11434", + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + // Try OpenAI-compatible + embeddingConfig.BackendType = testBackendOpenAI + embeddingManager, err = embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: No embedding backend available. Error: %v", err) + return + } + embeddingBackend = testBackendOpenAI + } + + // Verify embedding backend is actually working, not just reachable + verifyEmbeddingBackendWorking(t, embeddingManager, embeddingBackend) + _ = embeddingManager.Close() + + mcpServer := server.NewMCPServer("test-server", "1.0") + mockClient := &mockBackendClient{} + + // Test with high semantic ratio + configSemantic := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db-semantic"), + EmbeddingConfig: &embeddings.Config{ + BackendType: embeddingBackend, + BaseURL: embeddingConfig.BaseURL, + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + }, + HybridSearchRatio: 0.9, // 90% semantic + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integrationSemantic, err := NewIntegration(ctx, configSemantic, mcpServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integrationSemantic.Close() }() + + // Test with low semantic ratio (high BM25) + configKeyword := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db-keyword"), + EmbeddingConfig: &embeddings.Config{ + BackendType: embeddingBackend, + BaseURL: embeddingConfig.BaseURL, + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + }, + HybridSearchRatio: 0.1, // 10% semantic, 90% BM25 + } + + integrationKeyword, err := NewIntegration(ctx, configKeyword, mcpServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integrationKeyword.Close() }() + + tools := []vmcp.Tool{ + { + Name: "github_pull_request_read", + Description: "Get information on a specific pull request in GitHub repository.", + BackendID: "github", + }, + { + Name: "github_create_repository", + Description: "Create a new GitHub repository in your account or specified organization", + BackendID: "github", + }, + } + + capabilities := &aggregator.AggregatedCapabilities{ + Tools: tools, + RoutingTable: &vmcp.RoutingTable{ + Tools: make(map[string]*vmcp.BackendTarget), + Resources: map[string]*vmcp.BackendTarget{}, + Prompts: map[string]*vmcp.BackendTarget{}, + }, + } + + for _, tool := range tools { + capabilities.RoutingTable.Tools[tool.Name] = &vmcp.BackendTarget{ + WorkloadID: tool.BackendID, + WorkloadName: tool.BackendID, + } + } + + session := &mockSession{sessionID: "test-session"} + ctxWithCaps := discovery.WithDiscoveredCapabilities(ctx, capabilities) + + // Register both integrations + err = integrationSemantic.OnRegisterSession(ctx, session, capabilities) + require.NoError(t, err) + + err = integrationKeyword.OnRegisterSession(ctx, session, capabilities) + require.NoError(t, err) + + // Manually ingest tools for testing (OnRegisterSession skips ingestion) + mcpTools := make([]mcp.Tool, len(tools)) + for i, tool := range tools { + mcpTools[i] = mcp.Tool{ + Name: tool.Name, + Description: tool.Description, + } + } + err = integrationSemantic.IngestToolsForTesting(ctx, "github", "GitHub", nil, mcpTools) + require.NoError(t, err) + err = integrationKeyword.IngestToolsForTesting(ctx, "github", "GitHub", nil, mcpTools) + require.NoError(t, err) + + // Query that has semantic meaning but no exact keyword match + query := "view code review" + + // Test semantic search + requestSemantic := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": query, + "tool_keywords": "", + "limit": 10, + }, + }, + } + + handlerSemantic := integrationSemantic.CreateFindToolHandler() + resultSemantic, err := handlerSemantic(ctxWithCaps, requestSemantic) + require.NoError(t, err) + require.False(t, resultSemantic.IsError) + + // Test keyword search + requestKeyword := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": query, + "tool_keywords": "", + "limit": 10, + }, + }, + } + + handlerKeyword := integrationKeyword.CreateFindToolHandler() + resultKeyword, err := handlerKeyword(ctxWithCaps, requestKeyword) + require.NoError(t, err) + require.False(t, resultKeyword.IsError) + + // Parse both results + textSemantic, _ := mcp.AsTextContent(resultSemantic.Content[0]) + var responseSemantic map[string]any + json.Unmarshal([]byte(textSemantic.Text), &responseSemantic) + + textKeyword, _ := mcp.AsTextContent(resultKeyword.Content[0]) + var responseKeyword map[string]any + json.Unmarshal([]byte(textKeyword.Text), &responseKeyword) + + toolsSemantic, _ := responseSemantic["tools"].([]interface{}) + toolsKeyword, _ := responseKeyword["tools"].([]interface{}) + + // Both should find results (semantic should find PR tools, keyword might not) + assert.NotEmpty(t, toolsSemantic, "Semantic search should find results") + assert.NotEmpty(t, toolsKeyword, "Keyword search should find results") + + // Semantic search should find pull request tools even without exact keyword match + foundPRSemantic := false + for _, toolInterface := range toolsSemantic { + toolMap, _ := toolInterface.(map[string]interface{}) + toolName, _ := toolMap["name"].(string) + if toolName == "github_pull_request_read" { + foundPRSemantic = true + break + } + } + + t.Logf("Semantic search (90%% semantic): Found %d tools", len(toolsSemantic)) + t.Logf("Keyword search (10%% semantic): Found %d tools", len(toolsKeyword)) + t.Logf("Semantic search found PR tool: %v", foundPRSemantic) + + // Semantic search should be able to find semantically related tools + // even when keywords don't match exactly + assert.True(t, foundPRSemantic, + "Semantic search should find 'github_pull_request_read' for query 'view code review' even without exact keyword match") +} + +// TestFindTool_SemanticSimilarityScores tests that similarity scores are meaningful +func TestFindTool_SemanticSimilarityScores(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Try to use Ollama if available + embeddingBackend := "ollama" + embeddingConfig := &embeddings.Config{ + BackendType: embeddingBackend, + BaseURL: "http://localhost:11434", + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + // Try OpenAI-compatible + embeddingConfig.BackendType = testBackendOpenAI + embeddingManager, err = embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: No embedding backend available. Error: %v", err) + return + } + embeddingBackend = testBackendOpenAI + } + + // Verify embedding backend is actually working, not just reachable + verifyEmbeddingBackendWorking(t, embeddingManager, embeddingBackend) + _ = embeddingManager.Close() + + mcpServer := server.NewMCPServer("test-server", "1.0") + mockClient := &mockBackendClient{} + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: embeddingBackend, + BaseURL: embeddingConfig.BaseURL, + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + }, + HybridSearchRatio: 0.9, // High semantic ratio + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + tools := []vmcp.Tool{ + { + Name: "github_pull_request_read", + Description: "Get information on a specific pull request in GitHub repository.", + BackendID: "github", + }, + { + Name: "github_create_repository", + Description: "Create a new GitHub repository in your account or specified organization", + BackendID: "github", + }, + { + Name: "fetch_fetch", + Description: "Fetches a URL from the internet and optionally extracts its contents as markdown.", + BackendID: "fetch", + }, + } + + capabilities := &aggregator.AggregatedCapabilities{ + Tools: tools, + RoutingTable: &vmcp.RoutingTable{ + Tools: make(map[string]*vmcp.BackendTarget), + Resources: map[string]*vmcp.BackendTarget{}, + Prompts: map[string]*vmcp.BackendTarget{}, + }, + } + + for _, tool := range tools { + capabilities.RoutingTable.Tools[tool.Name] = &vmcp.BackendTarget{ + WorkloadID: tool.BackendID, + WorkloadName: tool.BackendID, + } + } + + session := &mockSession{sessionID: "test-session"} + err = integration.OnRegisterSession(ctx, session, capabilities) + require.NoError(t, err) + + // Manually ingest tools for testing (OnRegisterSession skips ingestion) + mcpTools := make([]mcp.Tool, len(tools)) + for i, tool := range tools { + mcpTools[i] = mcp.Tool{ + Name: tool.Name, + Description: tool.Description, + } + } + err = integration.IngestToolsForTesting(ctx, "github", "GitHub", nil, mcpTools) + require.NoError(t, err) + + ctxWithCaps := discovery.WithDiscoveredCapabilities(ctx, capabilities) + + // Query for pull request + query := "view pull request" + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": query, + "tool_keywords": "", + "limit": 10, + }, + }, + } + + handler := integration.CreateFindToolHandler() + result, err := handler(ctxWithCaps, request) + require.NoError(t, err) + require.False(t, result.IsError) + + textContent, _ := mcp.AsTextContent(result.Content[0]) + var response map[string]any + json.Unmarshal([]byte(textContent.Text), &response) + + toolsArray, _ := response["tools"].([]interface{}) + require.NotEmpty(t, toolsArray) + + // Check that results are sorted by similarity (highest first) + var similarities []float64 + for _, toolInterface := range toolsArray { + toolMap, _ := toolInterface.(map[string]interface{}) + similarity, _ := toolMap["similarity_score"].(float64) + similarities = append(similarities, similarity) + } + + // Verify results are sorted by similarity (descending) + for i := 1; i < len(similarities); i++ { + assert.GreaterOrEqual(t, similarities[i-1], similarities[i], + "Results should be sorted by similarity score (descending). Scores: %v", similarities) + } + + // The most relevant tool (pull request) should have a higher similarity than unrelated tools + if len(similarities) > 1 { + // First result should have highest similarity + assert.Greater(t, similarities[0], 0.0, "Top result should have positive similarity") + } +} diff --git a/pkg/vmcp/optimizer/find_tool_string_matching_test.go b/pkg/vmcp/optimizer/find_tool_string_matching_test.go new file mode 100644 index 0000000000..b994d7b95d --- /dev/null +++ b/pkg/vmcp/optimizer/find_tool_string_matching_test.go @@ -0,0 +1,696 @@ +package optimizer + +import ( + "context" + "encoding/json" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/optimizer/embeddings" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + "github.com/stacklok/toolhive/pkg/vmcp/discovery" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +// verifyOllamaWorking verifies that Ollama is actually working by attempting to generate an embedding +// This ensures the service is not just reachable but actually functional +func verifyOllamaWorking(t *testing.T, manager *embeddings.Manager) { + t.Helper() + _, err := manager.GenerateEmbedding([]string{"test"}) + if err != nil { + t.Skipf("Skipping test: Ollama is reachable but embedding generation failed. Error: %v. Ensure 'ollama pull %s' has been executed", err, embeddings.DefaultModelAllMiniLM) + } +} + +// getRealToolData returns test data based on actual MCP server tools +// These are real tool descriptions from GitHub and other MCP servers +func getRealToolData() []vmcp.Tool { + return []vmcp.Tool{ + { + Name: "github_pull_request_read", + Description: "Get information on a specific pull request in GitHub repository.", + BackendID: "github", + }, + { + Name: "github_list_pull_requests", + Description: "List pull requests in a GitHub repository. If the user specifies an author, then DO NOT use this tool and use the search_pull_requests tool instead.", + BackendID: "github", + }, + { + Name: "github_search_pull_requests", + Description: "Search for pull requests in GitHub repositories using issues search syntax already scoped to is:pr", + BackendID: "github", + }, + { + Name: "github_create_pull_request", + Description: "Create a new pull request in a GitHub repository.", + BackendID: "github", + }, + { + Name: "github_merge_pull_request", + Description: "Merge a pull request in a GitHub repository.", + BackendID: "github", + }, + { + Name: "github_pull_request_review_write", + Description: "Create and/or submit, delete review of a pull request.", + BackendID: "github", + }, + { + Name: "github_issue_read", + Description: "Get information about a specific issue in a GitHub repository.", + BackendID: "github", + }, + { + Name: "github_list_issues", + Description: "List issues in a GitHub repository. For pagination, use the 'endCursor' from the previous response's 'pageInfo' in the 'after' parameter.", + BackendID: "github", + }, + { + Name: "github_create_repository", + Description: "Create a new GitHub repository in your account or specified organization", + BackendID: "github", + }, + { + Name: "github_get_commit", + Description: "Get details for a commit from a GitHub repository", + BackendID: "github", + }, + { + Name: "fetch_fetch", + Description: "Fetches a URL from the internet and optionally extracts its contents as markdown.", + BackendID: "fetch", + }, + } +} + +// TestFindTool_StringMatching tests that find_tool can match strings correctly +func TestFindTool_StringMatching(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Setup optimizer integration + mcpServer := server.NewMCPServer("test-server", "1.0") + mockClient := &mockBackendClient{} + + // Try to use Ollama if available, otherwise skip test + embeddingConfig := &embeddings.Config{ + BackendType: embeddings.BackendTypeOllama, + BaseURL: "http://localhost:11434", + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull %s'", err, embeddings.DefaultModelAllMiniLM) + return + } + t.Cleanup(func() { _ = embeddingManager.Close() }) + + // Verify Ollama is actually working, not just reachable + verifyOllamaWorking(t, embeddingManager) + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: embeddings.BackendTypeOllama, + BaseURL: "http://localhost:11434", + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + }, + HybridSearchRatio: 0.5, // 50% semantic, 50% BM25 for better string matching + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer, mockClient, sessionMgr) + require.NoError(t, err) + require.NotNil(t, integration) + t.Cleanup(func() { _ = integration.Close() }) + + // Get real tool data + tools := getRealToolData() + + // Create capabilities with real tools + capabilities := &aggregator.AggregatedCapabilities{ + Tools: tools, + RoutingTable: &vmcp.RoutingTable{ + Tools: make(map[string]*vmcp.BackendTarget), + Resources: map[string]*vmcp.BackendTarget{}, + Prompts: map[string]*vmcp.BackendTarget{}, + }, + } + + // Build routing table + for _, tool := range tools { + capabilities.RoutingTable.Tools[tool.Name] = &vmcp.BackendTarget{ + WorkloadID: tool.BackendID, + WorkloadName: tool.BackendID, + } + } + + // Register session and generate embeddings + session := &mockSession{sessionID: "test-session"} + err = integration.OnRegisterSession(ctx, session, capabilities) + require.NoError(t, err) + + // Manually ingest tools for testing (OnRegisterSession skips ingestion) + mcpTools := make([]mcp.Tool, len(tools)) + for i, tool := range tools { + mcpTools[i] = mcp.Tool{ + Name: tool.Name, + Description: tool.Description, + } + } + err = integration.IngestToolsForTesting(ctx, "github", "GitHub", nil, mcpTools) + require.NoError(t, err) + + // Create context with capabilities + ctxWithCaps := discovery.WithDiscoveredCapabilities(ctx, capabilities) + + // Test cases: query -> expected tool names that should be found + testCases := []struct { + name string + query string + keywords string + expectedTools []string // Tools that should definitely be in results + minResults int // Minimum number of results expected + description string + }{ + { + name: "exact_pull_request_match", + query: "pull request", + keywords: "pull request", + expectedTools: []string{"github_pull_request_read", "github_list_pull_requests", "github_create_pull_request"}, + minResults: 3, + description: "Should find tools with exact 'pull request' string match", + }, + { + name: "pull_request_in_name", + query: "pull request", + keywords: "pull_request", + expectedTools: []string{"github_pull_request_read", "github_list_pull_requests"}, + minResults: 2, + description: "Should match tools with 'pull_request' in name", + }, + { + name: "list_pull_requests", + query: "list pull requests", + keywords: "list pull requests", + expectedTools: []string{"github_list_pull_requests"}, + minResults: 1, + description: "Should find list pull requests tool", + }, + { + name: "read_pull_request", + query: "read pull request", + keywords: "read pull request", + expectedTools: []string{"github_pull_request_read"}, + minResults: 1, + description: "Should find read pull request tool", + }, + { + name: "create_pull_request", + query: "create pull request", + keywords: "create pull request", + expectedTools: []string{"github_create_pull_request"}, + minResults: 1, + description: "Should find create pull request tool", + }, + { + name: "merge_pull_request", + query: "merge pull request", + keywords: "merge pull request", + expectedTools: []string{"github_merge_pull_request"}, + minResults: 1, + description: "Should find merge pull request tool", + }, + { + name: "search_pull_requests", + query: "search pull requests", + keywords: "search pull requests", + expectedTools: []string{"github_search_pull_requests"}, + minResults: 1, + description: "Should find search pull requests tool", + }, + { + name: "issue_tools", + query: "issue", + keywords: "issue", + expectedTools: []string{"github_issue_read", "github_list_issues"}, + minResults: 2, + description: "Should find issue-related tools", + }, + { + name: "repository_tool", + query: "create repository", + keywords: "create repository", + expectedTools: []string{"github_create_repository"}, + minResults: 1, + description: "Should find create repository tool", + }, + { + name: "commit_tool", + query: "get commit", + keywords: "commit", + expectedTools: []string{"github_get_commit"}, + minResults: 1, + description: "Should find get commit tool", + }, + { + name: "fetch_tool", + query: "fetch URL", + keywords: "fetch", + expectedTools: []string{"fetch_fetch"}, + minResults: 1, + description: "Should find fetch tool", + }, + } + + for _, tc := range testCases { + tc := tc // capture loop variable + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + // Create the tool call request + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": tc.query, + "tool_keywords": tc.keywords, + "limit": 20, + }, + }, + } + + // Call the handler + handler := integration.CreateFindToolHandler() + result, err := handler(ctxWithCaps, request) + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, result.IsError, "Tool call should not return error") + + // Parse the result + require.NotEmpty(t, result.Content, "Result should have content") + textContent, ok := mcp.AsTextContent(result.Content[0]) + require.True(t, ok, "Result should be text content") + + // Parse JSON response + var response map[string]any + err = json.Unmarshal([]byte(textContent.Text), &response) + require.NoError(t, err, "Result should be valid JSON") + + // Check tools array exists + toolsArray, ok := response["tools"].([]interface{}) + require.True(t, ok, "Response should have tools array") + require.GreaterOrEqual(t, len(toolsArray), tc.minResults, + "Should return at least %d results for query: %s", tc.minResults, tc.query) + + // Extract tool names from results + foundTools := make([]string, 0, len(toolsArray)) + for _, toolInterface := range toolsArray { + toolMap, okMap := toolInterface.(map[string]interface{}) + require.True(t, okMap, "Tool should be a map") + toolName, okName := toolMap["name"].(string) + require.True(t, okName, "Tool should have name") + foundTools = append(foundTools, toolName) + } + + // Check that at least some expected tools are found + // String matching may not be perfect, so we check that at least one expected tool is found + foundCount := 0 + for _, expectedTool := range tc.expectedTools { + for _, foundTool := range foundTools { + if foundTool == expectedTool { + foundCount++ + break + } + } + } + + // We should find at least one expected tool, or at least 50% of expected tools + minExpected := 1 + if len(tc.expectedTools) > 1 { + half := len(tc.expectedTools) / 2 + if half > minExpected { + minExpected = half + } + } + + assert.GreaterOrEqual(t, foundCount, minExpected, + "Query '%s' should find at least %d of expected tools %v. Found tools: %v (found %d/%d)", + tc.query, minExpected, tc.expectedTools, foundTools, foundCount, len(tc.expectedTools)) + + // Log which expected tools were found for debugging + if foundCount < len(tc.expectedTools) { + t.Logf("Query '%s': Found %d/%d expected tools. Found: %v, Expected: %v", + tc.query, foundCount, len(tc.expectedTools), foundTools, tc.expectedTools) + } + + // Verify token metrics exist + tokenMetrics, ok := response["token_metrics"].(map[string]interface{}) + require.True(t, ok, "Response should have token_metrics") + assert.Contains(t, tokenMetrics, "baseline_tokens") + assert.Contains(t, tokenMetrics, "returned_tokens") + assert.Contains(t, tokenMetrics, "tokens_saved") + assert.Contains(t, tokenMetrics, "savings_percentage") + }) + } +} + +// TestFindTool_ExactStringMatch tests that exact string matches work correctly +func TestFindTool_ExactStringMatch(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Setup optimizer integration with higher BM25 ratio for better string matching + mcpServer := server.NewMCPServer("test-server", "1.0") + mockClient := &mockBackendClient{} + + // Try to use Ollama if available, otherwise skip test + embeddingConfig := &embeddings.Config{ + BackendType: embeddings.BackendTypeOllama, + BaseURL: "http://localhost:11434", + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull %s'", err, embeddings.DefaultModelAllMiniLM) + return + } + t.Cleanup(func() { _ = embeddingManager.Close() }) + + // Verify Ollama is actually working, not just reachable + verifyOllamaWorking(t, embeddingManager) + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: embeddings.BackendTypeOllama, + BaseURL: "http://localhost:11434", + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + }, + HybridSearchRatio: 0.3, // 30% semantic, 70% BM25 for better exact string matching + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer, mockClient, sessionMgr) + require.NoError(t, err) + require.NotNil(t, integration) + t.Cleanup(func() { _ = integration.Close() }) + + // Create tools with specific strings to match + tools := []vmcp.Tool{ + { + Name: "test_pull_request_tool", + Description: "This tool handles pull requests in GitHub", + BackendID: "test", + }, + { + Name: "test_issue_tool", + Description: "This tool handles issues in GitHub", + BackendID: "test", + }, + { + Name: "test_repository_tool", + Description: "This tool creates repositories", + BackendID: "test", + }, + } + + capabilities := &aggregator.AggregatedCapabilities{ + Tools: tools, + RoutingTable: &vmcp.RoutingTable{ + Tools: make(map[string]*vmcp.BackendTarget), + Resources: map[string]*vmcp.BackendTarget{}, + Prompts: map[string]*vmcp.BackendTarget{}, + }, + } + + for _, tool := range tools { + capabilities.RoutingTable.Tools[tool.Name] = &vmcp.BackendTarget{ + WorkloadID: tool.BackendID, + WorkloadName: tool.BackendID, + } + } + + session := &mockSession{sessionID: "test-session"} + err = integration.OnRegisterSession(ctx, session, capabilities) + require.NoError(t, err) + + // Manually ingest tools for testing (OnRegisterSession skips ingestion) + mcpTools := make([]mcp.Tool, len(tools)) + for i, tool := range tools { + mcpTools[i] = mcp.Tool{ + Name: tool.Name, + Description: tool.Description, + } + } + err = integration.IngestToolsForTesting(ctx, "test", "test", nil, mcpTools) + require.NoError(t, err) + + ctxWithCaps := discovery.WithDiscoveredCapabilities(ctx, capabilities) + + // Test exact string matching + testCases := []struct { + name string + query string + keywords string + expectedTool string + description string + }{ + { + name: "exact_pull_request_string", + query: "pull request", + keywords: "pull request", + expectedTool: "test_pull_request_tool", + description: "Should match exact 'pull request' string", + }, + { + name: "exact_issue_string", + query: "issue", + keywords: "issue", + expectedTool: "test_issue_tool", + description: "Should match exact 'issue' string", + }, + { + name: "exact_repository_string", + query: "repository", + keywords: "repository", + expectedTool: "test_repository_tool", + description: "Should match exact 'repository' string", + }, + } + + for _, tc := range testCases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": tc.query, + "tool_keywords": tc.keywords, + "limit": 10, + }, + }, + } + + handler := integration.CreateFindToolHandler() + result, err := handler(ctxWithCaps, request) + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, result.IsError) + + textContent, okText := mcp.AsTextContent(result.Content[0]) + require.True(t, okText) + + var response map[string]any + err = json.Unmarshal([]byte(textContent.Text), &response) + require.NoError(t, err) + + toolsArray, okArray := response["tools"].([]interface{}) + require.True(t, okArray) + require.NotEmpty(t, toolsArray, "Should find at least one tool for query: %s", tc.query) + + // Check that the expected tool is in the results + found := false + for _, toolInterface := range toolsArray { + toolMap, okMap := toolInterface.(map[string]interface{}) + require.True(t, okMap) + toolName, okName := toolMap["name"].(string) + require.True(t, okName) + if toolName == tc.expectedTool { + found = true + break + } + } + + assert.True(t, found, + "Expected tool '%s' not found in results for query '%s'. This indicates string matching is not working correctly.", + tc.expectedTool, tc.query) + }) + } +} + +// TestFindTool_CaseInsensitive tests case-insensitive string matching +func TestFindTool_CaseInsensitive(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + mcpServer := server.NewMCPServer("test-server", "1.0") + mockClient := &mockBackendClient{} + + // Try to use Ollama if available, otherwise skip test + embeddingConfig := &embeddings.Config{ + BackendType: embeddings.BackendTypeOllama, + BaseURL: "http://localhost:11434", + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull %s'", err, embeddings.DefaultModelAllMiniLM) + return + } + t.Cleanup(func() { _ = embeddingManager.Close() }) + + // Verify Ollama is actually working, not just reachable + verifyOllamaWorking(t, embeddingManager) + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: embeddings.BackendTypeOllama, + BaseURL: "http://localhost:11434", + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + }, + HybridSearchRatio: 0.3, // Favor BM25 for string matching + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer, mockClient, sessionMgr) + require.NoError(t, err) + require.NotNil(t, integration) + t.Cleanup(func() { _ = integration.Close() }) + + tools := []vmcp.Tool{ + { + Name: "github_pull_request_read", + Description: "Get information on a specific pull request in GitHub repository.", + BackendID: "github", + }, + } + + capabilities := &aggregator.AggregatedCapabilities{ + Tools: tools, + RoutingTable: &vmcp.RoutingTable{ + Tools: map[string]*vmcp.BackendTarget{ + "github_pull_request_read": { + WorkloadID: "github", + WorkloadName: "github", + }, + }, + Resources: map[string]*vmcp.BackendTarget{}, + Prompts: map[string]*vmcp.BackendTarget{}, + }, + } + + session := &mockSession{sessionID: "test-session"} + err = integration.OnRegisterSession(ctx, session, capabilities) + require.NoError(t, err) + + // Manually ingest tools for testing (OnRegisterSession skips ingestion) + mcpTools := make([]mcp.Tool, len(tools)) + for i, tool := range tools { + mcpTools[i] = mcp.Tool{ + Name: tool.Name, + Description: tool.Description, + } + } + err = integration.IngestToolsForTesting(ctx, "github", "GitHub", nil, mcpTools) + require.NoError(t, err) + + ctxWithCaps := discovery.WithDiscoveredCapabilities(ctx, capabilities) + + // Test different case variations + queries := []string{ + "PULL REQUEST", + "Pull Request", + "pull request", + "PuLl ReQuEsT", + } + + for _, query := range queries { + query := query + t.Run("case_"+strings.ToLower(query), func(t *testing.T) { + t.Parallel() + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": query, + "tool_keywords": strings.ToLower(query), + "limit": 10, + }, + }, + } + + handler := integration.CreateFindToolHandler() + result, err := handler(ctxWithCaps, request) + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, result.IsError) + + textContent, okText := mcp.AsTextContent(result.Content[0]) + require.True(t, okText) + + var response map[string]any + err = json.Unmarshal([]byte(textContent.Text), &response) + require.NoError(t, err) + + toolsArray, okArray := response["tools"].([]interface{}) + require.True(t, okArray) + + // Should find the pull request tool regardless of case + found := false + for _, toolInterface := range toolsArray { + toolMap, okMap := toolInterface.(map[string]interface{}) + require.True(t, okMap) + toolName, okName := toolMap["name"].(string) + require.True(t, okName) + if toolName == "github_pull_request_read" { + found = true + break + } + } + + assert.True(t, found, + "Should find pull request tool with case-insensitive query: %s", query) + }) + } +} diff --git a/pkg/vmcp/optimizer/optimizer.go b/pkg/vmcp/optimizer/optimizer.go index 4a24d95576..19553ea2e1 100644 --- a/pkg/vmcp/optimizer/optimizer.go +++ b/pkg/vmcp/optimizer/optimizer.go @@ -13,7 +13,9 @@ package optimizer import ( "context" + "encoding/json" "fmt" + "sync" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" @@ -22,8 +24,11 @@ import ( "github.com/stacklok/toolhive/pkg/optimizer/db" "github.com/stacklok/toolhive/pkg/optimizer/embeddings" "github.com/stacklok/toolhive/pkg/optimizer/ingestion" + "github.com/stacklok/toolhive/pkg/optimizer/models" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + "github.com/stacklok/toolhive/pkg/vmcp/discovery" ) // Config holds optimizer configuration for vMCP integration. @@ -49,10 +54,12 @@ type Config struct { // //nolint:revive // Name is intentional for clarity in external packages type OptimizerIntegration struct { - config *Config - ingestionService *ingestion.Service - mcpServer *server.MCPServer // For registering tools - backendClient vmcp.BackendClient // For querying backends at startup + config *Config + ingestionService *ingestion.Service + mcpServer *server.MCPServer // For registering tools + backendClient vmcp.BackendClient // For querying backends at startup + sessionManager *transportsession.Manager + processedSessions sync.Map // Track sessions that have already been processed } // NewIntegration creates a new optimizer integration. @@ -61,6 +68,7 @@ func NewIntegration( cfg *Config, mcpServer *server.MCPServer, backendClient vmcp.BackendClient, + sessionManager *transportsession.Manager, ) (*OptimizerIntegration, error) { if cfg == nil || !cfg.Enabled { return nil, nil // Optimizer disabled @@ -85,6 +93,7 @@ func NewIntegration( ingestionService: svc, mcpServer: mcpServer, backendClient: backendClient, + sessionManager: sessionManager, }, nil } @@ -96,98 +105,30 @@ func NewIntegration( // 2. Generates embeddings for all tools (parallel per-backend) // 3. Registers optim.find_tool and optim.call_tool as session tools func (o *OptimizerIntegration) OnRegisterSession( - ctx context.Context, + _ context.Context, session server.ClientSession, - capabilities *aggregator.AggregatedCapabilities, + _ *aggregator.AggregatedCapabilities, ) error { if o == nil { return nil // Optimizer not enabled } sessionID := session.SessionID() - logger.Infow("Generating embeddings for session", "session_id", sessionID) - - // Group tools by backend for parallel processing - type backendTools struct { - backendID string - backendName string - backendURL string - transport string - tools []mcp.Tool - } - - backendMap := make(map[string]*backendTools) - - // Extract tools from routing table - if capabilities.RoutingTable != nil { - for toolName, target := range capabilities.RoutingTable.Tools { - // Find the tool definition from capabilities.Tools - var toolDef mcp.Tool - found := false - for i := range capabilities.Tools { - if capabilities.Tools[i].Name == toolName { - // Convert vmcp.Tool to mcp.Tool - // Note: vmcp.Tool.InputSchema is map[string]any, mcp.Tool.InputSchema is ToolInputSchema struct - // For ingestion, we just need the tool name and description - toolDef = mcp.Tool{ - Name: capabilities.Tools[i].Name, - Description: capabilities.Tools[i].Description, - // InputSchema will be empty - we only need name/description for embedding generation - } - found = true - break - } - } - if !found { - logger.Warnw("Tool in routing table but not in capabilities", - "tool_name", toolName, - "backend_id", target.WorkloadID) - continue - } - // Group by backend - if _, exists := backendMap[target.WorkloadID]; !exists { - backendMap[target.WorkloadID] = &backendTools{ - backendID: target.WorkloadID, - backendName: target.WorkloadName, - backendURL: target.BaseURL, - transport: target.TransportType, - tools: []mcp.Tool{}, - } - } - backendMap[target.WorkloadID].tools = append(backendMap[target.WorkloadID].tools, toolDef) - } - } + logger.Debugw("OnRegisterSession called", "session_id", sessionID) - // Ingest each backend's tools (in parallel - TODO: add goroutines) - for _, bt := range backendMap { - logger.Debugw("Ingesting backend for session", - "session_id", sessionID, - "backend_id", bt.backendID, - "backend_name", bt.backendName, - "tool_count", len(bt.tools)) - - // Ingest server with simplified metadata - // Note: URL and transport are not stored - vMCP manages backend lifecycle - err := o.ingestionService.IngestServer( - ctx, - bt.backendID, - bt.backendName, - nil, // description - bt.tools, - ) - if err != nil { - logger.Errorw("Failed to ingest backend", - "session_id", sessionID, - "backend_id", bt.backendID, - "error", err) - // Continue with other backends - } + // Check if this session has already been processed + if _, alreadyProcessed := o.processedSessions.LoadOrStore(sessionID, true); alreadyProcessed { + logger.Debugw("Session already processed, skipping duplicate ingestion", + "session_id", sessionID) + return nil } - logger.Infow("Embeddings generated for session", - "session_id", sessionID, - "backend_count", len(backendMap)) + // Skip ingestion in OnRegisterSession - IngestInitialBackends already handles ingestion at startup + // This prevents duplicate ingestion when sessions are registered + // The optimizer database is populated once at startup, not per-session + logger.Infow("Skipping ingestion in OnRegisterSession (handled by IngestInitialBackends at startup)", + "session_id", sessionID) return nil } @@ -252,7 +193,7 @@ func (o *OptimizerIntegration) RegisterTools(_ context.Context, session server.C Required: []string{"backend_id", "tool_name", "parameters"}, }, }, - Handler: o.createCallToolHandler(), + Handler: o.CreateCallToolHandler(), }, } @@ -265,32 +206,255 @@ func (o *OptimizerIntegration) RegisterTools(_ context.Context, session server.C return nil } -// createFindToolHandler creates the handler for optim.find_tool -func (*OptimizerIntegration) createFindToolHandler() func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { - return func(_ context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { - // TODO: Implement semantic search - // 1. Extract tool_description and tool_keywords from request.Params.Arguments - // 2. Call optimizer search service (hybrid semantic + BM25) - // 3. Return ranked list of tools with scores and token metrics +// CreateFindToolHandler creates the handler for optim.find_tool +// Exported for testing purposes +func (o *OptimizerIntegration) CreateFindToolHandler() func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return o.createFindToolHandler() +} + +// extractFindToolParams extracts and validates parameters from the find_tool request +func extractFindToolParams(args map[string]any) (toolDescription, toolKeywords string, limit int, err *mcp.CallToolResult) { + // Extract tool_description (required) + toolDescription, ok := args["tool_description"].(string) + if !ok || toolDescription == "" { + return "", "", 0, mcp.NewToolResultError("tool_description is required and must be a non-empty string") + } + + // Extract tool_keywords (optional) + toolKeywords, _ = args["tool_keywords"].(string) + + // Extract limit (optional, default: 10) + limit = 10 + if limitVal, ok := args["limit"]; ok { + if limitFloat, ok := limitVal.(float64); ok { + limit = int(limitFloat) + } + } + + return toolDescription, toolKeywords, limit, nil +} +// convertSearchResultsToResponse converts database search results to the response format +func convertSearchResultsToResponse(results []*models.BackendToolWithMetadata) ([]map[string]any, int) { + responseTools := make([]map[string]any, 0, len(results)) + totalReturnedTokens := 0 + + for _, result := range results { + // Unmarshal InputSchema + var inputSchema map[string]any + if len(result.InputSchema) > 0 { + if err := json.Unmarshal(result.InputSchema, &inputSchema); err != nil { + logger.Warnw("Failed to unmarshal input schema", + "tool_id", result.ID, + "tool_name", result.ToolName, + "error", err) + inputSchema = map[string]any{} // Use empty schema on error + } + } + + // Handle nil description + description := "" + if result.Description != nil { + description = *result.Description + } + + tool := map[string]any{ + "name": result.ToolName, + "description": description, + "input_schema": inputSchema, + "backend_id": result.MCPServerID, + "similarity_score": result.Similarity, + "token_count": result.TokenCount, + } + responseTools = append(responseTools, tool) + totalReturnedTokens += result.TokenCount + } + + return responseTools, totalReturnedTokens +} + +// createFindToolHandler creates the handler for optim.find_tool +func (o *OptimizerIntegration) createFindToolHandler() func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { logger.Debugw("optim.find_tool called", "request", request) - return mcp.NewToolResultError("optim.find_tool not yet implemented"), nil + // Extract parameters from request arguments + args, ok := request.Params.Arguments.(map[string]any) + if !ok { + return mcp.NewToolResultError("invalid arguments: expected object"), nil + } + + // Extract and validate parameters + toolDescription, toolKeywords, limit, err := extractFindToolParams(args) + if err != nil { + return err, nil + } + + // Perform hybrid search using database operations + if o.ingestionService == nil { + return mcp.NewToolResultError("backend tool operations not initialized"), nil + } + backendToolOps := o.ingestionService.GetBackendToolOps() + if backendToolOps == nil { + return mcp.NewToolResultError("backend tool operations not initialized"), nil + } + + // Configure hybrid search + hybridConfig := &db.HybridSearchConfig{ + SemanticRatio: o.config.HybridSearchRatio, + Limit: limit, + ServerID: nil, // Search across all servers + } + + // Execute hybrid search + queryText := toolDescription + if toolKeywords != "" { + queryText = toolDescription + " " + toolKeywords + } + results, err2 := backendToolOps.SearchHybrid(ctx, queryText, hybridConfig) + if err2 != nil { + logger.Errorw("Hybrid search failed", + "error", err2, + "tool_description", toolDescription, + "tool_keywords", toolKeywords, + "query_text", queryText) + return mcp.NewToolResultError(fmt.Sprintf("search failed: %v", err2)), nil + } + + // Convert results to response format + responseTools, totalReturnedTokens := convertSearchResultsToResponse(results) + + // Calculate token metrics + baselineTokens := o.ingestionService.GetTotalToolTokens(ctx) + tokensSaved := baselineTokens - totalReturnedTokens + savingsPercentage := 0.0 + if baselineTokens > 0 { + savingsPercentage = (float64(tokensSaved) / float64(baselineTokens)) * 100.0 + } + + tokenMetrics := map[string]any{ + "baseline_tokens": baselineTokens, + "returned_tokens": totalReturnedTokens, + "tokens_saved": tokensSaved, + "savings_percentage": savingsPercentage, + } + + // Build response + response := map[string]any{ + "tools": responseTools, + "token_metrics": tokenMetrics, + } + + // Marshal to JSON for the result + responseJSON, err3 := json.Marshal(response) + if err3 != nil { + logger.Errorw("Failed to marshal response", "error", err3) + return mcp.NewToolResultError(fmt.Sprintf("failed to marshal response: %v", err3)), nil + } + + logger.Infow("optim.find_tool completed", + "query", toolDescription, + "results_count", len(responseTools), + "tokens_saved", tokensSaved, + "savings_percentage", fmt.Sprintf("%.2f%%", savingsPercentage)) + + return mcp.NewToolResultText(string(responseJSON)), nil } } -// createCallToolHandler creates the handler for optim.call_tool -func (*OptimizerIntegration) createCallToolHandler() func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { - return func(_ context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { - // TODO: Implement dynamic tool invocation - // 1. Extract backend_id, tool_name, parameters from request.Params.Arguments - // 2. Validate backend and tool exist - // 3. Route to backend via existing router - // 4. Return result +// CreateCallToolHandler creates the handler for optim.call_tool +// Exported for testing purposes +func (o *OptimizerIntegration) CreateCallToolHandler() func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return o.createCallToolHandler() +} +// createCallToolHandler creates the handler for optim.call_tool +func (o *OptimizerIntegration) createCallToolHandler() func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { logger.Debugw("optim.call_tool called", "request", request) - return mcp.NewToolResultError("optim.call_tool not yet implemented"), nil + // Extract parameters from request arguments + args, ok := request.Params.Arguments.(map[string]any) + if !ok { + return mcp.NewToolResultError("invalid arguments: expected object"), nil + } + + // Extract backend_id (required) + backendID, ok := args["backend_id"].(string) + if !ok || backendID == "" { + return mcp.NewToolResultError("backend_id is required and must be a non-empty string"), nil + } + + // Extract tool_name (required) + toolName, ok := args["tool_name"].(string) + if !ok || toolName == "" { + return mcp.NewToolResultError("tool_name is required and must be a non-empty string"), nil + } + + // Extract parameters (required) + parameters, ok := args["parameters"].(map[string]any) + if !ok { + return mcp.NewToolResultError("parameters is required and must be an object"), nil + } + + // Get routing table from context via discovered capabilities + capabilities, ok := discovery.DiscoveredCapabilitiesFromContext(ctx) + if !ok || capabilities == nil { + return mcp.NewToolResultError("routing information not available in context"), nil + } + + if capabilities.RoutingTable == nil || capabilities.RoutingTable.Tools == nil { + return mcp.NewToolResultError("routing table not initialized"), nil + } + + // Find the tool in the routing table + target, exists := capabilities.RoutingTable.Tools[toolName] + if !exists { + return mcp.NewToolResultError(fmt.Sprintf("tool not found in routing table: %s", toolName)), nil + } + + // Verify the tool belongs to the specified backend + if target.WorkloadID != backendID { + return mcp.NewToolResultError(fmt.Sprintf( + "tool %s belongs to backend %s, not %s", + toolName, + target.WorkloadID, + backendID, + )), nil + } + + // Get the backend capability name (handles renamed tools) + backendToolName := target.GetBackendCapabilityName(toolName) + + logger.Infow("Calling tool via optimizer", + "backend_id", backendID, + "tool_name", toolName, + "backend_tool_name", backendToolName, + "workload_name", target.WorkloadName) + + // Call the tool on the backend using the backend client + result, err := o.backendClient.CallTool(ctx, target, backendToolName, parameters) + if err != nil { + logger.Errorw("Tool call failed", + "error", err, + "backend_id", backendID, + "tool_name", toolName, + "backend_tool_name", backendToolName) + return mcp.NewToolResultError(fmt.Sprintf("tool call failed: %v", err)), nil + } + + // Convert result to JSON + resultJSON, err := json.Marshal(result) + if err != nil { + logger.Errorw("Failed to marshal tool result", "error", err) + return mcp.NewToolResultError(fmt.Sprintf("failed to marshal result: %v", err)), nil + } + + logger.Infow("optim.call_tool completed successfully", + "backend_id", backendID, + "tool_name", toolName) + + return mcp.NewToolResultText(string(resultJSON)), nil } } @@ -362,3 +526,18 @@ func (o *OptimizerIntegration) Close() error { } return o.ingestionService.Close() } + +// IngestToolsForTesting manually ingests tools for testing purposes. +// This is a test helper that bypasses the normal ingestion flow. +func (o *OptimizerIntegration) IngestToolsForTesting( + ctx context.Context, + serverID string, + serverName string, + description *string, + tools []mcp.Tool, +) error { + if o == nil || o.ingestionService == nil { + return fmt.Errorf("optimizer integration not initialized") + } + return o.ingestionService.IngestServer(ctx, serverID, serverName, description, tools) +} diff --git a/pkg/vmcp/optimizer/optimizer_handlers_test.go b/pkg/vmcp/optimizer/optimizer_handlers_test.go new file mode 100644 index 0000000000..3889a47e37 --- /dev/null +++ b/pkg/vmcp/optimizer/optimizer_handlers_test.go @@ -0,0 +1,1026 @@ +package optimizer + +import ( + "context" + "encoding/json" + "path/filepath" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stacklok/toolhive/pkg/optimizer/embeddings" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + "github.com/stacklok/toolhive/pkg/vmcp/discovery" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +// mockMCPServerWithSession implements AddSessionTools for testing +type mockMCPServerWithSession struct { + *server.MCPServer + toolsAdded map[string][]server.ServerTool +} + +func newMockMCPServerWithSession() *mockMCPServerWithSession { + return &mockMCPServerWithSession{ + MCPServer: server.NewMCPServer("test-server", "1.0"), + toolsAdded: make(map[string][]server.ServerTool), + } +} + +func (m *mockMCPServerWithSession) AddSessionTools(sessionID string, tools ...server.ServerTool) error { + m.toolsAdded[sessionID] = tools + return nil +} + +// mockBackendClientWithCallTool implements CallTool for testing +type mockBackendClientWithCallTool struct { + callToolResult map[string]any + callToolError error +} + +func (*mockBackendClientWithCallTool) ListCapabilities(_ context.Context, _ *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + return &vmcp.CapabilityList{}, nil +} + +func (m *mockBackendClientWithCallTool) CallTool(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any) (map[string]any, error) { + if m.callToolError != nil { + return nil, m.callToolError + } + return m.callToolResult, nil +} + +//nolint:revive // Receiver unused in mock implementation +func (m *mockBackendClientWithCallTool) GetPrompt(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any) (string, error) { + return "", nil +} + +//nolint:revive // Receiver unused in mock implementation +func (m *mockBackendClientWithCallTool) ReadResource(_ context.Context, _ *vmcp.BackendTarget, _ string) ([]byte, error) { + return nil, nil +} + +// TestCreateFindToolHandler_InvalidArguments tests error handling for invalid arguments +func TestCreateFindToolHandler_InvalidArguments(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Setup optimizer integration + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClient{} + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + handler := integration.CreateFindToolHandler() + + // Test with invalid arguments type + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: "not a map", + }, + } + + result, err := handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error for invalid arguments") + + // Test with missing tool_description + request = mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "limit": 10, + }, + }, + } + + result, err = handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error for missing tool_description") + + // Test with empty tool_description + request = mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": "", + }, + }, + } + + result, err = handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error for empty tool_description") + + // Test with non-string tool_description + request = mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": 123, + }, + }, + } + + result, err = handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error for non-string tool_description") +} + +// TestCreateFindToolHandler_WithKeywords tests find_tool with keywords +func TestCreateFindToolHandler_WithKeywords(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClient{} + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + // Ingest a tool for testing + tools := []mcp.Tool{ + { + Name: "test_tool", + Description: "A test tool for searching", + }, + } + + err = integration.IngestToolsForTesting(ctx, "server-1", "TestServer", nil, tools) + require.NoError(t, err) + + handler := integration.CreateFindToolHandler() + + // Test with keywords + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": "search tool", + "tool_keywords": "test search", + "limit": 10, + }, + }, + } + + result, err := handler(ctx, request) + require.NoError(t, err) + require.False(t, result.IsError, "Should not return error") + + // Verify response structure + textContent, ok := mcp.AsTextContent(result.Content[0]) + require.True(t, ok) + + var response map[string]any + err = json.Unmarshal([]byte(textContent.Text), &response) + require.NoError(t, err) + + _, ok = response["tools"] + require.True(t, ok, "Response should have tools") + + _, ok = response["token_metrics"] + require.True(t, ok, "Response should have token_metrics") +} + +// TestCreateFindToolHandler_Limit tests limit parameter handling +func TestCreateFindToolHandler_Limit(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClient{} + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + handler := integration.CreateFindToolHandler() + + // Test with custom limit + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": "test", + "limit": 5, + }, + }, + } + + result, err := handler(ctx, request) + require.NoError(t, err) + require.False(t, result.IsError) + + // Test with float64 limit (from JSON) + request = mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": "test", + "limit": float64(3), + }, + }, + } + + result, err = handler(ctx, request) + require.NoError(t, err) + require.False(t, result.IsError) +} + +// TestCreateFindToolHandler_BackendToolOpsNil tests error when backend tool ops is nil +func TestCreateFindToolHandler_BackendToolOpsNil(t *testing.T) { + t.Parallel() + ctx := context.Background() + + // Create integration with nil ingestion service to trigger error path + integration := &OptimizerIntegration{ + config: &Config{Enabled: true}, + ingestionService: nil, // This will cause GetBackendToolOps to return nil + } + + handler := integration.CreateFindToolHandler() + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": "test", + }, + }, + } + + result, err := handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error when backend tool ops is nil") +} + +// TestCreateCallToolHandler_InvalidArguments tests error handling for invalid arguments +func TestCreateCallToolHandler_InvalidArguments(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Check Ollama availability first + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClientWithCallTool{} + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + handler := integration.CreateCallToolHandler() + + // Test with invalid arguments type + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.call_tool", + Arguments: "not a map", + }, + } + + result, err := handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error for invalid arguments") + + // Test with missing backend_id + request = mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.call_tool", + Arguments: map[string]any{ + "tool_name": "test_tool", + "parameters": map[string]any{}, + }, + }, + } + + result, err = handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error for missing backend_id") + + // Test with empty backend_id + request = mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.call_tool", + Arguments: map[string]any{ + "backend_id": "", + "tool_name": "test_tool", + "parameters": map[string]any{}, + }, + }, + } + + result, err = handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error for empty backend_id") + + // Test with missing tool_name + request = mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.call_tool", + Arguments: map[string]any{ + "backend_id": "backend-1", + "parameters": map[string]any{}, + }, + }, + } + + result, err = handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error for missing tool_name") + + // Test with missing parameters + request = mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.call_tool", + Arguments: map[string]any{ + "backend_id": "backend-1", + "tool_name": "test_tool", + }, + }, + } + + result, err = handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error for missing parameters") + + // Test with invalid parameters type + request = mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.call_tool", + Arguments: map[string]any{ + "backend_id": "backend-1", + "tool_name": "test_tool", + "parameters": "not a map", + }, + }, + } + + result, err = handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error for invalid parameters type") +} + +// TestCreateCallToolHandler_NoRoutingTable tests error when routing table is missing +func TestCreateCallToolHandler_NoRoutingTable(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Check Ollama availability first + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClientWithCallTool{} + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + handler := integration.CreateCallToolHandler() + + // Test without routing table in context + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.call_tool", + Arguments: map[string]any{ + "backend_id": "backend-1", + "tool_name": "test_tool", + "parameters": map[string]any{}, + }, + }, + } + + result, err := handler(ctx, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error when routing table is missing") +} + +// TestCreateCallToolHandler_ToolNotFound tests error when tool is not found +func TestCreateCallToolHandler_ToolNotFound(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Check Ollama availability first + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClientWithCallTool{} + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + handler := integration.CreateCallToolHandler() + + // Create context with routing table but tool not found + capabilities := &aggregator.AggregatedCapabilities{ + RoutingTable: &vmcp.RoutingTable{ + Tools: make(map[string]*vmcp.BackendTarget), + Resources: map[string]*vmcp.BackendTarget{}, + Prompts: map[string]*vmcp.BackendTarget{}, + }, + } + + ctxWithCaps := discovery.WithDiscoveredCapabilities(ctx, capabilities) + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.call_tool", + Arguments: map[string]any{ + "backend_id": "backend-1", + "tool_name": "nonexistent_tool", + "parameters": map[string]any{}, + }, + }, + } + + result, err := handler(ctxWithCaps, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error when tool is not found") +} + +// TestCreateCallToolHandler_BackendMismatch tests error when backend doesn't match +func TestCreateCallToolHandler_BackendMismatch(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Check Ollama availability first + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClientWithCallTool{} + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + handler := integration.CreateCallToolHandler() + + // Create context with routing table where tool belongs to different backend + capabilities := &aggregator.AggregatedCapabilities{ + RoutingTable: &vmcp.RoutingTable{ + Tools: map[string]*vmcp.BackendTarget{ + "test_tool": { + WorkloadID: "backend-2", // Different backend + WorkloadName: "Backend 2", + }, + }, + Resources: map[string]*vmcp.BackendTarget{}, + Prompts: map[string]*vmcp.BackendTarget{}, + }, + } + + ctxWithCaps := discovery.WithDiscoveredCapabilities(ctx, capabilities) + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.call_tool", + Arguments: map[string]any{ + "backend_id": "backend-1", // Requesting backend-1 + "tool_name": "test_tool", // But tool belongs to backend-2 + "parameters": map[string]any{}, + }, + }, + } + + result, err := handler(ctxWithCaps, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error when backend doesn't match") +} + +// TestCreateCallToolHandler_Success tests successful tool call +func TestCreateCallToolHandler_Success(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Check Ollama availability first + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClientWithCallTool{ + callToolResult: map[string]any{ + "result": "success", + }, + } + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + handler := integration.CreateCallToolHandler() + + // Create context with routing table + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "Backend 1", + BaseURL: "http://localhost:8000", + } + + capabilities := &aggregator.AggregatedCapabilities{ + RoutingTable: &vmcp.RoutingTable{ + Tools: map[string]*vmcp.BackendTarget{ + "test_tool": target, + }, + Resources: map[string]*vmcp.BackendTarget{}, + Prompts: map[string]*vmcp.BackendTarget{}, + }, + } + + ctxWithCaps := discovery.WithDiscoveredCapabilities(ctx, capabilities) + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.call_tool", + Arguments: map[string]any{ + "backend_id": "backend-1", + "tool_name": "test_tool", + "parameters": map[string]any{ + "param1": "value1", + }, + }, + }, + } + + result, err := handler(ctxWithCaps, request) + require.NoError(t, err) + require.False(t, result.IsError, "Should not return error") + + // Verify response + textContent, ok := mcp.AsTextContent(result.Content[0]) + require.True(t, ok) + + var response map[string]any + err = json.Unmarshal([]byte(textContent.Text), &response) + require.NoError(t, err) + assert.Equal(t, "success", response["result"]) +} + +// TestCreateCallToolHandler_CallToolError tests error handling when CallTool fails +func TestCreateCallToolHandler_CallToolError(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Check Ollama availability first + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClientWithCallTool{ + callToolError: assert.AnError, + } + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + handler := integration.CreateCallToolHandler() + + target := &vmcp.BackendTarget{ + WorkloadID: "backend-1", + WorkloadName: "Backend 1", + BaseURL: "http://localhost:8000", + } + + capabilities := &aggregator.AggregatedCapabilities{ + RoutingTable: &vmcp.RoutingTable{ + Tools: map[string]*vmcp.BackendTarget{ + "test_tool": target, + }, + Resources: map[string]*vmcp.BackendTarget{}, + Prompts: map[string]*vmcp.BackendTarget{}, + }, + } + + ctxWithCaps := discovery.WithDiscoveredCapabilities(ctx, capabilities) + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.call_tool", + Arguments: map[string]any{ + "backend_id": "backend-1", + "tool_name": "test_tool", + "parameters": map[string]any{}, + }, + }, + } + + result, err := handler(ctxWithCaps, request) + require.NoError(t, err) + require.True(t, result.IsError, "Should return error when CallTool fails") +} + +// TestCreateFindToolHandler_InputSchemaUnmarshalError tests error handling for invalid input schema +func TestCreateFindToolHandler_InputSchemaUnmarshalError(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClient{} + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + handler := integration.CreateFindToolHandler() + + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": "test", + }, + }, + } + + // The handler should handle invalid input schema gracefully + result, err := handler(ctx, request) + require.NoError(t, err) + // Should not error even if some tools have invalid schemas + require.False(t, result.IsError) +} + +// TestOnRegisterSession_DuplicateSession tests duplicate session handling +func TestOnRegisterSession_DuplicateSession(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Check Ollama availability first + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClient{} + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + session := &mockSession{sessionID: "test-session"} + capabilities := &aggregator.AggregatedCapabilities{} + + // First call + err = integration.OnRegisterSession(ctx, session, capabilities) + require.NoError(t, err) + + // Second call with same session ID (should be skipped) + err = integration.OnRegisterSession(ctx, session, capabilities) + require.NoError(t, err, "Should handle duplicate session gracefully") +} + +// TestIngestInitialBackends_ErrorHandling tests error handling during ingestion +func TestIngestInitialBackends_ErrorHandling(t *testing.T) { + t.Parallel() + ctx := context.Background() + tmpDir := t.TempDir() + + // Check Ollama availability first + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + mcpServer := newMockMCPServerWithSession() + mockClient := &mockBackendClient{ + err: assert.AnError, // Simulate error when listing capabilities + } + + config := &Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer.MCPServer, mockClient, sessionMgr) + require.NoError(t, err) + defer func() { _ = integration.Close() }() + + backends := []vmcp.Backend{ + { + ID: "backend-1", + Name: "Backend 1", + BaseURL: "http://localhost:8000", + TransportType: "sse", + }, + } + + // Should not fail even if backend query fails + err = integration.IngestInitialBackends(ctx, backends) + require.NoError(t, err, "Should handle backend query errors gracefully") +} + +// TestIngestInitialBackends_NilIntegration tests nil integration handling +func TestIngestInitialBackends_NilIntegration(t *testing.T) { + t.Parallel() + ctx := context.Background() + + var integration *OptimizerIntegration = nil + backends := []vmcp.Backend{} + + err := integration.IngestInitialBackends(ctx, backends) + require.NoError(t, err, "Should handle nil integration gracefully") +} diff --git a/pkg/vmcp/optimizer/optimizer_integration_test.go b/pkg/vmcp/optimizer/optimizer_integration_test.go index 82a51a925a..2fcb912743 100644 --- a/pkg/vmcp/optimizer/optimizer_integration_test.go +++ b/pkg/vmcp/optimizer/optimizer_integration_test.go @@ -4,14 +4,17 @@ import ( "context" "path/filepath" "testing" + "time" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" "github.com/stretchr/testify/require" "github.com/stacklok/toolhive/pkg/optimizer/embeddings" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" ) // mockBackendClient implements vmcp.BackendClient for integration testing @@ -107,18 +110,36 @@ func TestOptimizerIntegration_WithVMCP(t *testing.T) { }, }) + // Try to use Ollama if available, otherwise skip test + embeddingConfig := &embeddings.Config{ + BackendType: embeddings.BackendTypeOllama, + BaseURL: "http://localhost:11434", + Model: embeddings.DefaultModelAllMiniLM, + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull %s'", err, embeddings.DefaultModelAllMiniLM) + return + } + t.Cleanup(func() { _ = embeddingManager.Close() }) + // Configure optimizer optimizerConfig := &Config{ Enabled: true, PersistPath: filepath.Join(tmpDir, "optimizer-db"), EmbeddingConfig: &embeddings.Config{ - BackendType: "placeholder", + BackendType: embeddings.BackendTypeOllama, + BaseURL: "http://localhost:11434", + Model: embeddings.DefaultModelAllMiniLM, Dimension: 384, }, } // Create optimizer integration - integration, err := NewIntegration(ctx, optimizerConfig, mcpServer, mockClient) + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, optimizerConfig, mcpServer, mockClient, sessionMgr) require.NoError(t, err) defer func() { _ = integration.Close() }() diff --git a/pkg/vmcp/optimizer/optimizer_unit_test.go b/pkg/vmcp/optimizer/optimizer_unit_test.go index 794069b851..8b09a99ee8 100644 --- a/pkg/vmcp/optimizer/optimizer_unit_test.go +++ b/pkg/vmcp/optimizer/optimizer_unit_test.go @@ -4,6 +4,7 @@ import ( "context" "path/filepath" "testing" + "time" "github.com/mark3labs/mcp-go/mcp" "github.com/mark3labs/mcp-go/server" @@ -11,8 +12,10 @@ import ( "github.com/stretchr/testify/require" "github.com/stacklok/toolhive/pkg/optimizer/embeddings" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" "github.com/stacklok/toolhive/pkg/vmcp" "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" ) // mockBackendClient implements vmcp.BackendClient for testing @@ -85,13 +88,13 @@ func TestNewIntegration_Disabled(t *testing.T) { ctx := context.Background() // Test with nil config - integration, err := NewIntegration(ctx, nil, nil, nil) + integration, err := NewIntegration(ctx, nil, nil, nil, nil) require.NoError(t, err) assert.Nil(t, integration, "Should return nil when config is nil") // Test with disabled config config := &Config{Enabled: false} - integration, err = NewIntegration(ctx, config, nil, nil) + integration, err = NewIntegration(ctx, config, nil, nil, nil) require.NoError(t, err) assert.Nil(t, integration, "Should return nil when optimizer is disabled") } @@ -102,6 +105,21 @@ func TestNewIntegration_Enabled(t *testing.T) { ctx := context.Background() tmpDir := t.TempDir() + // Try to use Ollama if available, otherwise skip test + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull all-minilm'", err) + return + } + _ = embeddingManager.Close() + mcpServer := server.NewMCPServer("test-server", "1.0") mockClient := &mockBackendClient{} @@ -109,12 +127,15 @@ func TestNewIntegration_Enabled(t *testing.T) { Enabled: true, PersistPath: filepath.Join(tmpDir, "optimizer-db"), EmbeddingConfig: &embeddings.Config{ - BackendType: "placeholder", - Dimension: 384, + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "nomic-embed-text", + Dimension: 768, }, } - integration, err := NewIntegration(ctx, config, mcpServer, mockClient) + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer, mockClient, sessionMgr) require.NoError(t, err) require.NotNil(t, integration) defer func() { _ = integration.Close() }() @@ -129,16 +150,34 @@ func TestOnRegisterSession(t *testing.T) { mcpServer := server.NewMCPServer("test-server", "1.0") mockClient := &mockBackendClient{} + // Try to use Ollama if available, otherwise skip test + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull all-minilm'", err) + return + } + _ = embeddingManager.Close() + config := &Config{ Enabled: true, PersistPath: filepath.Join(tmpDir, "optimizer-db"), EmbeddingConfig: &embeddings.Config{ - BackendType: "placeholder", - Dimension: 384, + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "nomic-embed-text", + Dimension: 768, }, } - integration, err := NewIntegration(ctx, config, mcpServer, mockClient) + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer, mockClient, sessionMgr) require.NoError(t, err) defer func() { _ = integration.Close() }() @@ -189,16 +228,34 @@ func TestRegisterTools(t *testing.T) { mcpServer := server.NewMCPServer("test-server", "1.0") mockClient := &mockBackendClient{} + // Try to use Ollama if available, otherwise skip test + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull all-minilm'", err) + return + } + _ = embeddingManager.Close() + config := &Config{ Enabled: true, PersistPath: filepath.Join(tmpDir, "optimizer-db"), EmbeddingConfig: &embeddings.Config{ - BackendType: "placeholder", - Dimension: 384, + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "nomic-embed-text", + Dimension: 768, }, } - integration, err := NewIntegration(ctx, config, mcpServer, mockClient) + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer, mockClient, sessionMgr) require.NoError(t, err) defer func() { _ = integration.Close() }() @@ -230,16 +287,34 @@ func TestClose(t *testing.T) { mcpServer := server.NewMCPServer("test-server", "1.0") mockClient := &mockBackendClient{} + // Try to use Ollama if available, otherwise skip test + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v. Run 'ollama serve && ollama pull all-minilm'", err) + return + } + _ = embeddingManager.Close() + config := &Config{ Enabled: true, PersistPath: filepath.Join(tmpDir, "optimizer-db"), EmbeddingConfig: &embeddings.Config{ - BackendType: "placeholder", - Dimension: 384, + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "nomic-embed-text", + Dimension: 768, }, } - integration, err := NewIntegration(ctx, config, mcpServer, mockClient) + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := NewIntegration(ctx, config, mcpServer, mockClient, sessionMgr) require.NoError(t, err) err = integration.Close() diff --git a/pkg/vmcp/server/optimizer_test.go b/pkg/vmcp/server/optimizer_test.go new file mode 100644 index 0000000000..0d8cba1ad5 --- /dev/null +++ b/pkg/vmcp/server/optimizer_test.go @@ -0,0 +1,350 @@ +package server + +import ( + "context" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/stacklok/toolhive/pkg/optimizer/embeddings" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + discoveryMocks "github.com/stacklok/toolhive/pkg/vmcp/discovery/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/router" +) + +// TestNew_OptimizerEnabled tests server creation with optimizer enabled +func TestNew_OptimizerEnabled(t *testing.T) { + t.Parallel() + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockBackendClient := mocks.NewMockBackendClient(ctrl) + mockBackendClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil). + AnyTimes() + + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) + mockDiscoveryMgr.EXPECT(). + Discover(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil). + AnyTimes() + mockDiscoveryMgr.EXPECT().Stop().AnyTimes() + + tmpDir := t.TempDir() + + // Try to use Ollama if available + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + cfg := &Config{ + Name: "test-server", + Version: "1.0.0", + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + OptimizerConfig: &OptimizerConfig{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingBackend: "ollama", + EmbeddingURL: "http://localhost:11434", + EmbeddingModel: "all-minilm", + EmbeddingDimension: 384, + HybridSearchRatio: 0.7, + }, + } + + rt := router.NewDefaultRouter() + backends := []vmcp.Backend{ + { + ID: "backend-1", + Name: "Backend 1", + BaseURL: "http://localhost:8000", + TransportType: "sse", + }, + } + + srv, err := New(ctx, cfg, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry(backends), nil) + require.NoError(t, err) + require.NotNil(t, srv) + defer func() { _ = srv.Stop(context.Background()) }() + + // Verify optimizer integration was created + // We can't directly access optimizerIntegration, but we can verify server was created successfully +} + +// TestNew_OptimizerDisabled tests server creation with optimizer disabled +func TestNew_OptimizerDisabled(t *testing.T) { + t.Parallel() + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockBackendClient := mocks.NewMockBackendClient(ctrl) + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) + mockDiscoveryMgr.EXPECT().Stop().AnyTimes() + + cfg := &Config{ + Name: "test-server", + Version: "1.0.0", + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + OptimizerConfig: &OptimizerConfig{ + Enabled: false, // Disabled + }, + } + + rt := router.NewDefaultRouter() + backends := []vmcp.Backend{} + + srv, err := New(ctx, cfg, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry(backends), nil) + require.NoError(t, err) + require.NotNil(t, srv) + defer func() { _ = srv.Stop(context.Background()) }() +} + +// TestNew_OptimizerConfigNil tests server creation with nil optimizer config +func TestNew_OptimizerConfigNil(t *testing.T) { + t.Parallel() + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockBackendClient := mocks.NewMockBackendClient(ctrl) + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) + mockDiscoveryMgr.EXPECT().Stop().AnyTimes() + + cfg := &Config{ + Name: "test-server", + Version: "1.0.0", + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + OptimizerConfig: nil, // Nil config + } + + rt := router.NewDefaultRouter() + backends := []vmcp.Backend{} + + srv, err := New(ctx, cfg, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry(backends), nil) + require.NoError(t, err) + require.NotNil(t, srv) + defer func() { _ = srv.Stop(context.Background()) }() +} + +// TestNew_OptimizerIngestionError tests error handling during optimizer ingestion +func TestNew_OptimizerIngestionError(t *testing.T) { + t.Parallel() + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockBackendClient := mocks.NewMockBackendClient(ctrl) + // Return error when listing capabilities + mockBackendClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(nil, assert.AnError). + AnyTimes() + + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) + mockDiscoveryMgr.EXPECT().Stop().AnyTimes() + + tmpDir := t.TempDir() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + cfg := &Config{ + Name: "test-server", + Version: "1.0.0", + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + OptimizerConfig: &OptimizerConfig{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingBackend: "ollama", + EmbeddingURL: "http://localhost:11434", + EmbeddingModel: "all-minilm", + EmbeddingDimension: 384, + }, + } + + rt := router.NewDefaultRouter() + backends := []vmcp.Backend{ + { + ID: "backend-1", + Name: "Backend 1", + BaseURL: "http://localhost:8000", + TransportType: "sse", + }, + } + + // Should not fail even if ingestion fails + srv, err := New(ctx, cfg, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry(backends), nil) + require.NoError(t, err, "Server should be created even if optimizer ingestion fails") + require.NotNil(t, srv) + defer func() { _ = srv.Stop(context.Background()) }() +} + +// TestNew_OptimizerHybridRatio tests hybrid ratio configuration +func TestNew_OptimizerHybridRatio(t *testing.T) { + t.Parallel() + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockBackendClient := mocks.NewMockBackendClient(ctrl) + mockBackendClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil). + AnyTimes() + + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) + mockDiscoveryMgr.EXPECT(). + Discover(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil). + AnyTimes() + mockDiscoveryMgr.EXPECT().Stop().AnyTimes() + + tmpDir := t.TempDir() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + cfg := &Config{ + Name: "test-server", + Version: "1.0.0", + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + OptimizerConfig: &OptimizerConfig{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingBackend: "ollama", + EmbeddingURL: "http://localhost:11434", + EmbeddingModel: "all-minilm", + EmbeddingDimension: 384, + HybridSearchRatio: 0.5, // Custom ratio + }, + } + + rt := router.NewDefaultRouter() + backends := []vmcp.Backend{} + + srv, err := New(ctx, cfg, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry(backends), nil) + require.NoError(t, err) + require.NotNil(t, srv) + defer func() { _ = srv.Stop(context.Background()) }() +} + +// TestServer_Stop_OptimizerCleanup tests optimizer cleanup on server stop +func TestServer_Stop_OptimizerCleanup(t *testing.T) { + t.Parallel() + ctx := context.Background() + + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockBackendClient := mocks.NewMockBackendClient(ctrl) + mockBackendClient.EXPECT(). + ListCapabilities(gomock.Any(), gomock.Any()). + Return(&vmcp.CapabilityList{}, nil). + AnyTimes() + + mockDiscoveryMgr := discoveryMocks.NewMockManager(ctrl) + mockDiscoveryMgr.EXPECT(). + Discover(gomock.Any(), gomock.Any()). + Return(&aggregator.AggregatedCapabilities{}, nil). + AnyTimes() + mockDiscoveryMgr.EXPECT().Stop().AnyTimes() + + tmpDir := t.TempDir() + + embeddingConfig := &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + } + + embeddingManager, err := embeddings.NewManager(embeddingConfig) + if err != nil { + t.Skipf("Skipping test: Ollama not available. Error: %v", err) + return + } + _ = embeddingManager.Close() + + cfg := &Config{ + Name: "test-server", + Version: "1.0.0", + Host: "127.0.0.1", + Port: 0, + SessionTTL: 5 * time.Minute, + OptimizerConfig: &OptimizerConfig{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingBackend: "ollama", + EmbeddingURL: "http://localhost:11434", + EmbeddingModel: "all-minilm", + EmbeddingDimension: 384, + }, + } + + rt := router.NewDefaultRouter() + backends := []vmcp.Backend{} + + srv, err := New(ctx, cfg, rt, mockBackendClient, mockDiscoveryMgr, vmcp.NewImmutableRegistry(backends), nil) + require.NoError(t, err) + require.NotNil(t, srv) + + // Stop should clean up optimizer + err = srv.Stop(context.Background()) + require.NoError(t, err) +} diff --git a/pkg/vmcp/server/server.go b/pkg/vmcp/server/server.go index 2253f3fd67..0bb302b438 100644 --- a/pkg/vmcp/server/server.go +++ b/pkg/vmcp/server/server.go @@ -367,7 +367,15 @@ func New( if cfg.HealthMonitorConfig != nil { // Get initial backends list from registry for health monitoring setup initialBackends := backendRegistry.List(ctx) - healthMon, err = health.NewMonitor(backendClient, initialBackends, *cfg.HealthMonitorConfig) + + // Construct server's own URL for self-check detection + // Use http:// as default scheme (most common for local development) + var selfURL string + if cfg.Host != "" && cfg.Port > 0 { + selfURL = fmt.Sprintf("http://%s:%d", cfg.Host, cfg.Port) + } + + healthMon, err = health.NewMonitor(backendClient, initialBackends, *cfg.HealthMonitorConfig, selfURL) if err != nil { return nil, fmt.Errorf("failed to create health monitor: %w", err) } @@ -382,40 +390,44 @@ func New( // Initialize optimizer integration if enabled var optimizerInteg OptimizerIntegration - if cfg.OptimizerConfig != nil && cfg.OptimizerConfig.Enabled { - logger.Infow("Initializing optimizer integration (chromem-go)", - "persist_path", cfg.OptimizerConfig.PersistPath, - "embedding_backend", cfg.OptimizerConfig.EmbeddingBackend) - - // Convert server config to optimizer config - hybridRatio := 0.7 // Default - if cfg.OptimizerConfig.HybridSearchRatio != 0 { - hybridRatio = cfg.OptimizerConfig.HybridSearchRatio - } - optimizerCfg := &optimizer.Config{ - Enabled: cfg.OptimizerConfig.Enabled, - PersistPath: cfg.OptimizerConfig.PersistPath, - FTSDBPath: cfg.OptimizerConfig.FTSDBPath, - HybridSearchRatio: hybridRatio, - EmbeddingConfig: &embeddings.Config{ - BackendType: cfg.OptimizerConfig.EmbeddingBackend, - BaseURL: cfg.OptimizerConfig.EmbeddingURL, - Model: cfg.OptimizerConfig.EmbeddingModel, - Dimension: cfg.OptimizerConfig.EmbeddingDimension, - }, - } + if cfg.OptimizerConfig != nil { + if cfg.OptimizerConfig.Enabled { + logger.Infow("Initializing optimizer integration (chromem-go)", + "persist_path", cfg.OptimizerConfig.PersistPath, + "embedding_backend", cfg.OptimizerConfig.EmbeddingBackend) + + // Convert server config to optimizer config + hybridRatio := 0.7 // Default + if cfg.OptimizerConfig.HybridSearchRatio != 0 { + hybridRatio = cfg.OptimizerConfig.HybridSearchRatio + } + optimizerCfg := &optimizer.Config{ + Enabled: cfg.OptimizerConfig.Enabled, + PersistPath: cfg.OptimizerConfig.PersistPath, + FTSDBPath: cfg.OptimizerConfig.FTSDBPath, + HybridSearchRatio: hybridRatio, + EmbeddingConfig: &embeddings.Config{ + BackendType: cfg.OptimizerConfig.EmbeddingBackend, + BaseURL: cfg.OptimizerConfig.EmbeddingURL, + Model: cfg.OptimizerConfig.EmbeddingModel, + Dimension: cfg.OptimizerConfig.EmbeddingDimension, + }, + } - optimizerInteg, err = optimizer.NewIntegration(ctx, optimizerCfg, mcpServer, backendClient) - if err != nil { - return nil, fmt.Errorf("failed to initialize optimizer: %w", err) - } - logger.Info("Optimizer integration initialized successfully") + optimizerInteg, err = optimizer.NewIntegration(ctx, optimizerCfg, mcpServer, backendClient, sessionManager) + if err != nil { + return nil, fmt.Errorf("failed to initialize optimizer: %w", err) + } + logger.Info("Optimizer integration initialized successfully") - // Ingest discovered backends at startup (populate optimizer database) - initialBackends := backendRegistry.List(ctx) - if err := optimizerInteg.IngestInitialBackends(ctx, initialBackends); err != nil { - logger.Warnf("Failed to ingest initial backends: %v", err) - // Don't fail server startup - optimizer can still work with incremental ingestion + // Ingest discovered backends at startup (populate optimizer database) + initialBackends := backendRegistry.List(ctx) + if err := optimizerInteg.IngestInitialBackends(ctx, initialBackends); err != nil { + logger.Warnf("Failed to ingest initial backends: %v", err) + // Don't fail server startup - optimizer can still work with incremental ingestion + } + } else { + logger.Info("Optimizer configuration present but disabled (enabled=false), skipping initialization") } } @@ -509,23 +521,59 @@ func New( "resource_count", len(caps.RoutingTable.Resources), "prompt_count", len(caps.RoutingTable.Prompts)) - // Inject capabilities into SDK session - if err := srv.injectCapabilities(sessionID, caps); err != nil { - logger.Errorw("failed to inject session capabilities", - "error", err, - "session_id", sessionID) - return - } + // When optimizer is enabled, we should NOT inject backend tools directly. + // Instead, only optimizer tools (optim.find_tool, optim.call_tool) will be exposed. + // Backend tools are still discovered and stored for optimizer ingestion, + // but not exposed directly to clients. + if srv.optimizerIntegration == nil { + // Inject capabilities into SDK session (only when optimizer is disabled) + if err := srv.injectCapabilities(sessionID, caps); err != nil { + logger.Errorw("failed to inject session capabilities", + "error", err, + "session_id", sessionID) + return + } - logger.Infow("session capabilities injected", - "session_id", sessionID, - "tool_count", len(caps.Tools), - "resource_count", len(caps.Resources)) + logger.Infow("session capabilities injected", + "session_id", sessionID, + "tool_count", len(caps.Tools), + "resource_count", len(caps.Resources)) + } else { + // Optimizer is enabled - register optimizer tools FIRST so they're available immediately + // Backend tools will be accessible via optim.find_tool and optim.call_tool + if err := srv.optimizerIntegration.RegisterTools(ctx, session); err != nil { + logger.Errorw("failed to register optimizer tools", + "error", err, + "session_id", sessionID) + // Don't fail session initialization - continue without optimizer tools + } else { + logger.Infow("optimizer tools registered", + "session_id", sessionID) + } - // Generate embeddings and register optimizer tools if enabled - if srv.optimizerIntegration != nil { - logger.Debugw("Generating embeddings for optimizer", "session_id", sessionID) + // Inject resources (but not backend tools) + if len(caps.Resources) > 0 { + sdkResources := srv.capabilityAdapter.ToSDKResources(caps.Resources) + if err := srv.mcpServer.AddSessionResources(sessionID, sdkResources...); err != nil { + logger.Errorw("failed to add session resources", + "error", err, + "session_id", sessionID) + return + } + logger.Debugw("added session resources (optimizer mode)", + "session_id", sessionID, + "count", len(sdkResources)) + } + logger.Infow("optimizer mode: backend tools not exposed directly", + "session_id", sessionID, + "backend_tool_count", len(caps.Tools), + "resource_count", len(caps.Resources)) + } + // Generate embeddings for optimizer if enabled + // This happens after tools are registered so tools are available immediately + if srv.optimizerIntegration != nil { + logger.Debugw("Calling OnRegisterSession for optimizer", "session_id", sessionID) // Generate embeddings for all tools in this session if err := srv.optimizerIntegration.OnRegisterSession(ctx, session, caps); err != nil { logger.Errorw("failed to generate embeddings for optimizer", @@ -533,16 +581,7 @@ func New( "session_id", sessionID) // Don't fail session initialization - continue without optimizer } else { - // Register optimizer tools (optim.find_tool, optim.call_tool) - if err := srv.optimizerIntegration.RegisterTools(ctx, session); err != nil { - logger.Errorw("failed to register optimizer tools", - "error", err, - "session_id", sessionID) - // Don't fail session initialization - continue without optimizer tools - } else { - logger.Infow("optimizer tools registered", - "session_id", sessionID) - } + logger.Debugw("OnRegisterSession completed successfully", "session_id", sessionID) } } }) diff --git a/scripts/README.md b/scripts/README.md index 09a382f6b0..fa19fe399d 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -81,7 +81,40 @@ Then open any `.db` file in VSCode to browse tables visually. ## Testing Scripts -### Optimizer Tests +### Optimizer Tool Finding Tests + +These scripts test the `optim.find_tool` functionality in different scenarios: + +#### Test via vMCP Server Connection +```bash +# Test optim.find_tool through a running vMCP server +go run scripts/test-vmcp-find-tool/main.go "read pull requests from GitHub" [server_url] + +# Default server URL: http://localhost:4483/mcp +# Example: +go run scripts/test-vmcp-find-tool/main.go "search the web" http://localhost:4483/mcp +``` +Connects to a running vMCP server and calls `optim.find_tool` via the MCP protocol. Useful for integration testing with a live server. + +#### Call Optimizer Tool Directly +```bash +# Call optim.find_tool via MCP client +go run scripts/call-optim-find-tool/main.go [tool_keywords] [limit] [server_url] + +# Examples: +go run scripts/call-optim-find-tool/main.go "search the web" "web search" 20 +go run scripts/call-optim-find-tool/main.go "read files" "" 10 http://localhost:4483/mcp +``` +A more flexible client for calling `optim.find_tool` with various parameters. Useful for manual testing and debugging. + +#### Test Optimizer Handler Directly +```bash +# Test the optimizer handler directly (unit test style) +go run scripts/test-optim-find-tool/main.go "read pull requests from GitHub" +``` +Tests the optimizer's `find_tool` handler directly without requiring a full vMCP server. Creates a mock environment with test tools and embeddings. Useful for development and debugging the optimizer logic. + +### Other Optimizer Tests ```bash # Test with sqlite-vec extension ./scripts/test-optimizer-with-sqlite-vec.sh diff --git a/scripts/call-optim-find-tool/main.go b/scripts/call-optim-find-tool/main.go new file mode 100644 index 0000000000..3df36a3e86 --- /dev/null +++ b/scripts/call-optim-find-tool/main.go @@ -0,0 +1,137 @@ +//go:build ignore +// +build ignore + +package main + +import ( + "context" + "encoding/json" + "fmt" + "os" + "time" + + "github.com/mark3labs/mcp-go/client" + "github.com/mark3labs/mcp-go/client/transport" + "github.com/mark3labs/mcp-go/mcp" +) + +func main() { + if len(os.Args) < 2 { + fmt.Println("Usage: go run main.go [tool_keywords] [limit] [server_url]") + fmt.Println("Example: go run main.go 'search the web' 'web search' 20") + fmt.Println("Default server URL: http://localhost:4483/mcp") + os.Exit(1) + } + + toolDescription := os.Args[1] + toolKeywords := "" + if len(os.Args) >= 3 { + toolKeywords = os.Args[2] + } + limit := 20 + if len(os.Args) >= 4 { + if l, err := fmt.Sscanf(os.Args[3], "%d", &limit); err != nil || l != 1 { + fmt.Printf("Invalid limit: %s, using default 20\n", os.Args[3]) + limit = 20 + } + } + serverURL := "http://localhost:4483/mcp" + if len(os.Args) >= 5 { + serverURL = os.Args[4] + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + // Create streamable-http client to connect to vmcp server + mcpClient, err := client.NewStreamableHttpClient( + serverURL, + transport.WithHTTPTimeout(30*time.Second), + transport.WithContinuousListening(), + ) + if err != nil { + fmt.Printf("❌ Failed to create MCP client: %v\n", err) + os.Exit(1) + } + defer func() { + if err := mcpClient.Close(); err != nil { + fmt.Printf("⚠️ Error closing client: %v\n", err) + } + }() + + // Start the client connection + if err := mcpClient.Start(ctx); err != nil { + fmt.Printf("❌ Failed to start client connection: %v\n", err) + os.Exit(1) + } + + // Initialize the client + initResult, err := mcpClient.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{ + Name: "optim-find-tool-client", + Version: "1.0.0", + }, + Capabilities: mcp.ClientCapabilities{}, + }, + }) + if err != nil { + fmt.Printf("❌ Failed to initialize client: %v\n", err) + os.Exit(1) + } + fmt.Printf("✅ Connected to: %s %s\n", initResult.ServerInfo.Name, initResult.ServerInfo.Version) + + // Call optim.find_tool + args := map[string]any{ + "tool_description": toolDescription, + "limit": limit, + } + if toolKeywords != "" { + args["tool_keywords"] = toolKeywords + } + + callResult, err := mcpClient.CallTool(ctx, mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: args, + }, + }) + if err != nil { + fmt.Printf("❌ Failed to call optim.find_tool: %v\n", err) + os.Exit(1) + } + + if callResult.IsError { + fmt.Printf("❌ Tool call returned an error\n") + if len(callResult.Content) > 0 { + if textContent, ok := mcp.AsTextContent(callResult.Content[0]); ok { + fmt.Printf("Error: %s\n", textContent.Text) + } + } + os.Exit(1) + } + + // Parse and display the result + if len(callResult.Content) > 0 { + if textContent, ok := mcp.AsTextContent(callResult.Content[0]); ok { + // Try to parse as JSON for pretty printing + var resultData map[string]any + if err := json.Unmarshal([]byte(textContent.Text), &resultData); err == nil { + // Pretty print JSON + prettyJSON, err := json.MarshalIndent(resultData, "", " ") + if err == nil { + fmt.Println(string(prettyJSON)) + } else { + fmt.Println(textContent.Text) + } + } else { + fmt.Println(textContent.Text) + } + } else { + fmt.Printf("%+v\n", callResult.Content) + } + } else { + fmt.Println("(No content returned)") + } +} diff --git a/scripts/inspect-chromem/inspect-chromem.go b/scripts/inspect-chromem/inspect-chromem.go index 672741b5ae..14b5c5e4a0 100644 --- a/scripts/inspect-chromem/inspect-chromem.go +++ b/scripts/inspect-chromem/inspect-chromem.go @@ -35,9 +35,9 @@ func main() { fmt.Println(" - backend_tools") fmt.Println() - // Create a dummy embedding function (we're just inspecting, not querying) + // Create an embedding function for collection access (we're just inspecting, not querying) dummyEmbedding := func(ctx context.Context, text string) ([]float32, error) { - return make([]float32, 384), nil // Placeholder + return make([]float32, 384), nil } // Inspect backend_servers collection diff --git a/scripts/test-optim-find-tool/main.go b/scripts/test-optim-find-tool/main.go new file mode 100644 index 0000000000..e61fc8c9c2 --- /dev/null +++ b/scripts/test-optim-find-tool/main.go @@ -0,0 +1,246 @@ +//go:build ignore +// +build ignore + +package main + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "time" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" + + "github.com/stacklok/toolhive/pkg/optimizer/embeddings" + transportsession "github.com/stacklok/toolhive/pkg/transport/session" + "github.com/stacklok/toolhive/pkg/vmcp" + "github.com/stacklok/toolhive/pkg/vmcp/aggregator" + "github.com/stacklok/toolhive/pkg/vmcp/discovery" + "github.com/stacklok/toolhive/pkg/vmcp/optimizer" + vmcpsession "github.com/stacklok/toolhive/pkg/vmcp/session" +) + +func main() { + if len(os.Args) < 2 { + fmt.Println("Usage: go run main.go ") + fmt.Println("Example: go run main.go 'read pull requests from GitHub'") + os.Exit(1) + } + + query := os.Args[1] + ctx := context.Background() + tmpDir := filepath.Join(os.TempDir(), "optimizer-test") + os.MkdirAll(tmpDir, 0755) + + fmt.Printf("🔍 Testing optim.find_tool with query: %s\n\n", query) + + // Create MCP server + mcpServer := server.NewMCPServer("test-server", "1.0") + + // Create mock backend client + mockClient := &mockBackendClient{} + + // Configure optimizer + optimizerConfig := &optimizer.Config{ + Enabled: true, + PersistPath: filepath.Join(tmpDir, "optimizer-db"), + EmbeddingConfig: &embeddings.Config{ + BackendType: "ollama", + BaseURL: "http://localhost:11434", + Model: "all-minilm", + Dimension: 384, + }, + } + + // Create optimizer integration + sessionMgr := transportsession.NewManager(30*time.Minute, vmcpsession.VMCPSessionFactory()) + integration, err := optimizer.NewIntegration(ctx, optimizerConfig, mcpServer, mockClient, sessionMgr) + if err != nil { + fmt.Printf("❌ Failed to create optimizer integration: %v\n", err) + os.Exit(1) + } + defer func() { _ = integration.Close() }() + + fmt.Println("✅ Optimizer integration created") + + // Ingest some test tools + backends := []vmcp.Backend{ + { + ID: "github", + Name: "GitHub", + BaseURL: "http://localhost:8000", + TransportType: "sse", + }, + } + + err = integration.IngestInitialBackends(ctx, backends) + if err != nil { + fmt.Printf("⚠️ Failed to ingest initial backends: %v (continuing...)\n", err) + } + + // Create a test session + sessionID := "test-session-123" + testSession := &mockSession{sessionID: sessionID} + + // Create capabilities with GitHub tools + capabilities := &aggregator.AggregatedCapabilities{ + Tools: []vmcp.Tool{ + { + Name: "github_pull_request_read", + Description: "Read details of a pull request from GitHub", + BackendID: "github", + }, + { + Name: "github_issue_read", + Description: "Read details of an issue from GitHub", + BackendID: "github", + }, + { + Name: "github_pull_request_list", + Description: "List pull requests in a GitHub repository", + BackendID: "github", + }, + }, + RoutingTable: &vmcp.RoutingTable{ + Tools: map[string]*vmcp.BackendTarget{ + "github_pull_request_read": { + WorkloadID: "github", + WorkloadName: "GitHub", + }, + "github_issue_read": { + WorkloadID: "github", + WorkloadName: "GitHub", + }, + "github_pull_request_list": { + WorkloadID: "github", + WorkloadName: "GitHub", + }, + }, + Resources: map[string]*vmcp.BackendTarget{}, + Prompts: map[string]*vmcp.BackendTarget{}, + }, + } + + // Register session with MCP server first (needed for RegisterTools) + err = mcpServer.RegisterSession(ctx, testSession) + if err != nil { + fmt.Printf("⚠️ Failed to register session: %v\n", err) + } + + // Generate embeddings for session + err = integration.OnRegisterSession(ctx, testSession, capabilities) + if err != nil { + fmt.Printf("❌ Failed to generate embeddings: %v\n", err) + os.Exit(1) + } + fmt.Println("✅ Embeddings generated for session") + + // Skip RegisterTools since we're calling the handler directly + // RegisterTools requires per-session tool support which the mock doesn't have + // err = integration.RegisterTools(ctx, testSession) + // if err != nil { + // fmt.Printf("⚠️ Failed to register optimizer tools: %v (skipping, calling handler directly)\n", err) + // } + fmt.Println("⏭️ Skipping tool registration (testing handler directly)") + + // Now try to call optim.find_tool directly via the handler + fmt.Printf("\n🔍 Calling optim.find_tool handler directly...\n\n") + + // Create a context with capabilities (needed for the handler) + ctxWithCaps := discovery.WithDiscoveredCapabilities(ctx, capabilities) + + // Create the tool call request + request := mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": query, + "tool_keywords": "github pull request", + "limit": 10, + }, + }, + } + + // Call the handler directly using the exported test method + handler := integration.CreateFindToolHandler() + result, err := handler(ctxWithCaps, request) + if err != nil { + fmt.Printf("❌ Failed to call optim.find_tool: %v\n", err) + os.Exit(1) + } + + fmt.Println("\n✅ Successfully called optim.find_tool!") + fmt.Println("\n📊 Results:") + + // Print the result - CallToolResult has Content field which is a slice + resultJSON, err := json.MarshalIndent(result, "", " ") + if err != nil { + fmt.Printf("Error marshaling result: %v\n", err) + fmt.Printf("Raw result: %+v\n", result) + } else { + fmt.Println(string(resultJSON)) + } +} + +type mockBackendClient struct{} + +func (m *mockBackendClient) ListCapabilities(_ context.Context, _ *vmcp.BackendTarget) (*vmcp.CapabilityList, error) { + return &vmcp.CapabilityList{ + Tools: []vmcp.Tool{ + { + Name: "github_pull_request_read", + Description: "Read details of a pull request from GitHub", + }, + { + Name: "github_issue_read", + Description: "Read details of an issue from GitHub", + }, + { + Name: "github_pull_request_list", + Description: "List pull requests in a GitHub repository", + }, + }, + }, nil +} + +func (m *mockBackendClient) CallTool(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any) (map[string]any, error) { + return nil, nil +} + +func (m *mockBackendClient) GetPrompt(_ context.Context, _ *vmcp.BackendTarget, _ string, _ map[string]any) (string, error) { + return "", nil +} + +func (m *mockBackendClient) ReadResource(_ context.Context, _ *vmcp.BackendTarget, _ string) ([]byte, error) { + return nil, nil +} + +type mockSession struct { + sessionID string +} + +func (m *mockSession) SessionID() string { + return m.sessionID +} + +func (m *mockSession) Send(_ interface{}) error { + return nil +} + +func (m *mockSession) Close() error { + return nil +} + +func (m *mockSession) Initialize() {} + +func (m *mockSession) Initialized() bool { + return true +} + +func (m *mockSession) NotificationChannel() chan<- mcp.JSONRPCNotification { + ch := make(chan mcp.JSONRPCNotification, 1) + return ch +} diff --git a/scripts/test-vmcp-find-tool/main.go b/scripts/test-vmcp-find-tool/main.go new file mode 100644 index 0000000000..71861d2508 --- /dev/null +++ b/scripts/test-vmcp-find-tool/main.go @@ -0,0 +1,158 @@ +//go:build ignore +// +build ignore + +package main + +import ( + "context" + "encoding/json" + "fmt" + "os" + "time" + + "github.com/mark3labs/mcp-go/client" + "github.com/mark3labs/mcp-go/client/transport" + "github.com/mark3labs/mcp-go/mcp" +) + +func main() { + if len(os.Args) < 2 { + fmt.Println("Usage: go run main.go [server_url]") + fmt.Println("Example: go run main.go 'read pull requests from GitHub'") + fmt.Println("Default server URL: http://localhost:4483/mcp") + os.Exit(1) + } + + query := os.Args[1] + serverURL := "http://localhost:4483/mcp" + if len(os.Args) >= 3 { + serverURL = os.Args[2] + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + fmt.Printf("🔍 Testing optim.find_tool via vmcp server\n") + fmt.Printf(" Server: %s\n", serverURL) + fmt.Printf(" Query: %s\n\n", query) + + // Create streamable-http client to connect to vmcp server + mcpClient, err := client.NewStreamableHttpClient( + serverURL, + transport.WithHTTPTimeout(30*time.Second), + transport.WithContinuousListening(), + ) + if err != nil { + fmt.Printf("❌ Failed to create MCP client: %v\n", err) + os.Exit(1) + } + defer func() { + if err := mcpClient.Close(); err != nil { + fmt.Printf("⚠️ Error closing client: %v\n", err) + } + }() + + // Start the client connection + if err := mcpClient.Start(ctx); err != nil { + fmt.Printf("❌ Failed to start client connection: %v\n", err) + os.Exit(1) + } + fmt.Println("✅ Connected to vmcp server") + + // Initialize the client + initResult, err := mcpClient.Initialize(ctx, mcp.InitializeRequest{ + Params: mcp.InitializeParams{ + ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION, + ClientInfo: mcp.Implementation{ + Name: "test-vmcp-client", + Version: "1.0.0", + }, + Capabilities: mcp.ClientCapabilities{}, + }, + }) + if err != nil { + fmt.Printf("❌ Failed to initialize client: %v\n", err) + os.Exit(1) + } + fmt.Printf("✅ Initialized - Server: %s %s\n\n", initResult.ServerInfo.Name, initResult.ServerInfo.Version) + + // List available tools to see if optim.find_tool is available + fmt.Println("📋 Listing available tools...") + toolsResult, err := mcpClient.ListTools(ctx, mcp.ListToolsRequest{}) + if err != nil { + fmt.Printf("❌ Failed to list tools: %v\n", err) + os.Exit(1) + } + + fmt.Printf("Found %d tools:\n", len(toolsResult.Tools)) + hasFindTool := false + for _, tool := range toolsResult.Tools { + fmt.Printf(" - %s: %s\n", tool.Name, tool.Description) + if tool.Name == "optim.find_tool" { + hasFindTool = true + } + } + fmt.Println() + + if !hasFindTool { + fmt.Println("⚠️ Warning: optim.find_tool not found in available tools") + fmt.Println(" The optimizer may not be enabled on this vmcp server") + fmt.Println(" Continuing anyway...\n") + } + + // Call optim.find_tool + fmt.Printf("🔍 Calling optim.find_tool with query: %s\n\n", query) + + callResult, err := mcpClient.CallTool(ctx, mcp.CallToolRequest{ + Params: mcp.CallToolParams{ + Name: "optim.find_tool", + Arguments: map[string]any{ + "tool_description": query, + "tool_keywords": "pull request", + "limit": 20, + }, + }, + }) + if err != nil { + fmt.Printf("❌ Failed to call optim.find_tool: %v\n", err) + os.Exit(1) + } + + if callResult.IsError { + fmt.Printf("❌ Tool call returned an error\n") + if len(callResult.Content) > 0 { + if textContent, ok := mcp.AsTextContent(callResult.Content[0]); ok { + fmt.Printf("Error: %s\n", textContent.Text) + } + } + os.Exit(1) + } + + fmt.Println("✅ Successfully called optim.find_tool!") + fmt.Println("\n📊 Results:") + + // Parse and display the result + if len(callResult.Content) > 0 { + if textContent, ok := mcp.AsTextContent(callResult.Content[0]); ok { + // Try to parse as JSON for pretty printing + var resultData map[string]any + if err := json.Unmarshal([]byte(textContent.Text), &resultData); err == nil { + // Pretty print JSON + prettyJSON, err := json.MarshalIndent(resultData, "", " ") + if err == nil { + fmt.Println(string(prettyJSON)) + } else { + fmt.Println(textContent.Text) + } + } else { + // Not JSON, print as-is + fmt.Println(textContent.Text) + } + } else { + // Not text content, print raw + fmt.Printf("%+v\n", callResult.Content) + } + } else { + fmt.Println("(No content returned)") + } +}