Wire optimizer flags into thv vmcp serve (#4940)

Add --optimizer, --optimizer-embedding, --embedding-model, and
--embedding-image flags to the serve subcommand. Extend ServeConfig
with the four corresponding fields. In Serve(), inject a non-nil
OptimizerConfig when either Tier 1 or Tier 2 is active, start the
TEI container via EmbeddingServiceManager when Tier 2 is requested,
and defer Stop() for clean shutdown. --optimizer-embedding implies
--optimizer. Regenerate CLI docs.

Closes #4887

Co-authored-by: taskbot <taskbot@users.noreply.github.com>

Author: yrobla

cmd/thv/app/vmcp.go

@@ -30,11 +30,15 @@ servers from a ToolHive group into a single unified endpoint.`,
 // newVMCPServeCommand returns the "vmcp serve" subcommand.
 func newVMCPServeCommand() *cobra.Command {
 	var (
-		configPath  string
-		group       string
-		host        string
-		port        int
-		enableAudit bool
+		configPath      string
+		group           string
+		host            string
+		port            int
+		enableAudit     bool
+		enableOptimizer bool
+		enableEmbedding bool
+		embeddingModel  string
+		embeddingImage  string
 	)
 	cmd := &cobra.Command{
 		Use:   "serve",
@@ -51,16 +55,28 @@ configuration file is needed for the common case of aggregating a local group.`,
 		Args: cobra.NoArgs,
 		RunE: func(cmd *cobra.Command, _ []string) error {
 			return vmcpcli.Serve(cmd.Context(), vmcpcli.ServeConfig{
-				ConfigPath:  configPath,
-				GroupRef:    group,
-				Host:        host,
-				Port:        port,
-				EnableAudit: enableAudit,
+				ConfigPath:      configPath,
+				GroupRef:        group,
+				Host:            host,
+				Port:            port,
+				EnableAudit:     enableAudit,
+				EnableOptimizer: enableOptimizer,
+				EnableEmbedding: enableEmbedding,
+				EmbeddingModel:  embeddingModel,
+				EmbeddingImage:  embeddingImage,
 			})
 		},
 	}
 	cmd.Flags().StringVarP(&configPath, "config", "c", "", "Path to vMCP configuration file")
 	cmd.Flags().StringVar(&group, "group", "", "ToolHive group name (zero-config quick mode when --config is omitted)")
+	cmd.Flags().BoolVar(&enableOptimizer, "optimizer", false,
+		"Enable FTS5 keyword optimizer (Tier 1): exposes find_tool and call_tool instead of all backend tools")
+	cmd.Flags().BoolVar(&enableEmbedding, "optimizer-embedding", false,
+		"Enable managed TEI semantic optimizer (Tier 2); implies --optimizer")
+	cmd.Flags().StringVar(&embeddingModel, "embedding-model", "BAAI/bge-small-en-v1.5",
+		"HuggingFace model name for semantic search (Tier 2)")
+	cmd.Flags().StringVar(&embeddingImage, "embedding-image",
+		"ghcr.io/huggingface/text-embeddings-inference:cpu-latest", "TEI container image (Tier 2)")
 	cmd.Flags().StringVar(&host, "host", "127.0.0.1", "Host address to bind to")
 	cmd.Flags().IntVar(&port, "port", 4483, "Port to listen on")
 	cmd.Flags().BoolVar(&enableAudit, "enable-audit", false, "Enable audit logging with default configuration")

docs/cli/thv_vmcp_serve.md

@@ -31,6 +31,9 @@ const (
 	// defaultTEIImage is the default HuggingFace Text Embeddings Inference image.
 	defaultTEIImage = "ghcr.io/huggingface/text-embeddings-inference:cpu-latest"
 
+	// DefaultEmbeddingModel is the HuggingFace model used when EmbeddingModel is empty.
+	DefaultEmbeddingModel = "BAAI/bge-small-en-v1.5"
+
 	// teiModelCachePath is the path inside the TEI container where models are cached.
 	teiModelCachePath = "/data"
 

pkg/vmcp/cli/embedding_manager.go

@@ -0,0 +1,147 @@
+// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package cli
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config"
+)
+
+// stubEmbeddingManager is a test double for the embeddingManager interface.
+type stubEmbeddingManager struct {
+	startURL  string
+	startErr  error
+	stopErr   error
+	startSeen bool
+	stopSeen  bool
+}
+
+func (s *stubEmbeddingManager) Start(_ context.Context) (string, error) {
+	s.startSeen = true
+	return s.startURL, s.startErr
+}
+
+func (s *stubEmbeddingManager) Stop(_ context.Context) error {
+	s.stopSeen = true
+	return s.stopErr
+}
+
+func TestInjectOptimizerConfig_NeitherTierEnabled(t *testing.T) {
+	t.Parallel()
+
+	vmcpCfg := &vmcpconfig.Config{}
+	cfg := ServeConfig{EnableOptimizer: false, EnableEmbedding: false}
+
+	cleanup, err := injectOptimizerConfig(context.Background(), cfg, vmcpCfg, nil)
+
+	require.NoError(t, err)
+	assert.Nil(t, cleanup)
+	assert.Nil(t, vmcpCfg.Optimizer, "Optimizer must remain nil when neither tier is enabled")
+}
+
+func TestInjectOptimizerConfig_Tier1Only(t *testing.T) {
+	t.Parallel()
+
+	vmcpCfg := &vmcpconfig.Config{}
+	cfg := ServeConfig{EnableOptimizer: true, EnableEmbedding: false}
+
+	cleanup, err := injectOptimizerConfig(context.Background(), cfg, vmcpCfg, nil)
+
+	require.NoError(t, err)
+	assert.Nil(t, cleanup, "Tier 1 does not start TEI — no cleanup needed")
+	require.NotNil(t, vmcpCfg.Optimizer)
+	assert.Empty(t, vmcpCfg.Optimizer.EmbeddingService, "Tier 1 must not set an embedding service URL")
+}
+
+func TestInjectOptimizerConfig_Tier1_PreservesExistingOptimizerConfig(t *testing.T) {
+	t.Parallel()
+
+	existing := &vmcpconfig.OptimizerConfig{MaxToolsToReturn: 5}
+	vmcpCfg := &vmcpconfig.Config{Optimizer: existing}
+	cfg := ServeConfig{EnableOptimizer: true, EnableEmbedding: false}
+
+	_, err := injectOptimizerConfig(context.Background(), cfg, vmcpCfg, nil)
+
+	require.NoError(t, err)
+	assert.Same(t, existing, vmcpCfg.Optimizer, "Existing optimizer config must not be replaced")
+}
+
+func TestInjectOptimizerConfig_Tier2_SetsEmbeddingURL(t *testing.T) {
+	t.Parallel()
+
+	stub := &stubEmbeddingManager{startURL: "http://127.0.0.1:8080"}
+	vmcpCfg := &vmcpconfig.Config{}
+	cfg := ServeConfig{EnableEmbedding: true}
+
+	cleanup, err := injectOptimizerConfig(context.Background(), cfg, vmcpCfg, stub)
+
+	require.NoError(t, err)
+	require.NotNil(t, cleanup, "Tier 2 must return a cleanup func")
+	assert.True(t, stub.startSeen, "Start must be called for Tier 2")
+	require.NotNil(t, vmcpCfg.Optimizer)
+	assert.Equal(t, "http://127.0.0.1:8080", vmcpCfg.Optimizer.EmbeddingService)
+}
+
+func TestInjectOptimizerConfig_Tier2_ImpliesOptimizer(t *testing.T) {
+	t.Parallel()
+
+	stub := &stubEmbeddingManager{startURL: "http://127.0.0.1:8080"}
+	vmcpCfg := &vmcpconfig.Config{}
+	// EnableOptimizer is false — EnableEmbedding alone must activate the optimizer.
+	cfg := ServeConfig{EnableOptimizer: false, EnableEmbedding: true}
+
+	_, err := injectOptimizerConfig(context.Background(), cfg, vmcpCfg, stub)
+
+	require.NoError(t, err)
+	assert.NotNil(t, vmcpCfg.Optimizer, "Tier 2 must activate optimizer even without --optimizer flag")
+}
+
+func TestInjectOptimizerConfig_Tier2_StartError(t *testing.T) {
+	t.Parallel()
+
+	startErr := errors.New("docker daemon unavailable")
+	stub := &stubEmbeddingManager{startErr: startErr}
+	vmcpCfg := &vmcpconfig.Config{}
+	cfg := ServeConfig{EnableEmbedding: true}
+
+	cleanup, err := injectOptimizerConfig(context.Background(), cfg, vmcpCfg, stub)
+
+	require.Error(t, err)
+	assert.ErrorContains(t, err, "docker daemon unavailable")
+	assert.Nil(t, cleanup, "No cleanup func must be returned on Start failure")
+}
+
+func TestInjectOptimizerConfig_Tier2_NilManagerReturnsError(t *testing.T) {
+	t.Parallel()
+
+	vmcpCfg := &vmcpconfig.Config{}
+	cfg := ServeConfig{EnableEmbedding: true}
+
+	cleanup, err := injectOptimizerConfig(context.Background(), cfg, vmcpCfg, nil)
+
+	require.Error(t, err)
+	assert.ErrorContains(t, err, "embedding manager must not be nil")
+	assert.Nil(t, cleanup)
+}
+
+func TestInjectOptimizerConfig_Tier2_CleanupCallsStop(t *testing.T) {
+	t.Parallel()
+
+	stub := &stubEmbeddingManager{startURL: "http://127.0.0.1:9999"}
+	vmcpCfg := &vmcpconfig.Config{}
+	cfg := ServeConfig{EnableEmbedding: true}
+
+	cleanup, err := injectOptimizerConfig(context.Background(), cfg, vmcpCfg, stub)
+	require.NoError(t, err)
+	require.NotNil(t, cleanup)
+
+	cleanup()
+	assert.True(t, stub.stopSeen, "cleanup func must call Stop on the embedding manager")
+}

pkg/vmcp/cli/optimizer_wiring_test.go

@@ -27,6 +27,7 @@ import (
 	authserverconfig "github.com/stacklok/toolhive/pkg/authserver"
 	authserverrunner "github.com/stacklok/toolhive/pkg/authserver/runner"
 	"github.com/stacklok/toolhive/pkg/authserver/server/keys"
+	"github.com/stacklok/toolhive/pkg/container"
 	"github.com/stacklok/toolhive/pkg/container/runtime"
 	"github.com/stacklok/toolhive/pkg/groups"
 	"github.com/stacklok/toolhive/pkg/migration"
@@ -67,6 +68,18 @@ type ServeConfig struct {
 	// EnableAudit enables audit logging with default configuration when
 	// the loaded config does not already define an audit section.
 	EnableAudit bool
+
+	// Optimizer tier selection (Phase 4 — flag-driven).
+	// EnableOptimizer enables Tier 1 FTS5 keyword search (find_tool / call_tool).
+	EnableOptimizer bool
+	// EnableEmbedding enables Tier 2 TEI semantic search; implies EnableOptimizer.
+	EnableEmbedding bool
+	// EmbeddingModel is the HuggingFace model name for the managed TEI container.
+	// Defaults to "BAAI/bge-small-en-v1.5" when empty.
+	EmbeddingModel string
+	// EmbeddingImage is the TEI container image.
+	// Defaults to the CPU TEI image when empty.
+	EmbeddingImage string
 }
 
 // validateQuickModeHost returns an error when the config represents quick mode
@@ -280,6 +293,32 @@ func Serve(ctx context.Context, cfg ServeConfig) error {
 		return fmt.Errorf("failed to create status reporter: %w", err)
 	}
 
+	// Optimizer wiring — Phase 4: flag-driven Tier 1 (FTS5) and Tier 2 (TEI).
+	// Build the embedding manager only when Tier 2 is requested, to avoid
+	// unnecessary Docker / Kubernetes API calls for Tier 0 and Tier 1.
+	var embMgr embeddingManager
+	if cfg.EnableEmbedding {
+		model := cfg.EmbeddingModel
+		if model == "" {
+			model = DefaultEmbeddingModel
+		}
+		m, err := NewEmbeddingServiceManager(container.NewFactory(), EmbeddingServiceManagerConfig{
+			Model: model,
+			Image: cfg.EmbeddingImage,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to create embedding service manager: %w", err)
+		}
+		embMgr = m
+	}
+	teiCleanup, err := injectOptimizerConfig(ctx, cfg, vmcpCfg, embMgr)
+	if err != nil {
+		return err
+	}
+	if teiCleanup != nil {
+		defer teiCleanup()
+	}
+
 	optCfg, err := optimizer.GetAndValidateConfig(vmcpCfg.Optimizer)
 	if err != nil {
 		return fmt.Errorf("failed to validate optimizer config: %w", err)
@@ -371,6 +410,39 @@ func Serve(ctx context.Context, cfg ServeConfig) error {
 	return srv.Start(ctx)
 }
 
+// embeddingManager is the minimal interface over *EmbeddingServiceManager needed
+// by the Serve lifecycle. Defined here to allow stub injection in unit tests;
+// production code passes a *EmbeddingServiceManager.
+type embeddingManager interface {
+	Start(ctx context.Context) (string, error)
+	Stop(ctx context.Context) error
+}
+
+// injectOptimizerConfig ensures vmcpCfg.Optimizer is non-nil when flag-driven
+// optimizer tiers are active, and starts the TEI container when EnableEmbedding
+// is true. Returns a non-nil cleanup func only when a TEI container was started;
+// the caller must defer it. mgr must be non-nil when cfg.EnableEmbedding is true.
+func injectOptimizerConfig(ctx context.Context, cfg ServeConfig, vmcpCfg *config.Config, mgr embeddingManager) (func(), error) {
+	if !cfg.EnableOptimizer && !cfg.EnableEmbedding {
+		return nil, nil
+	}
+	if vmcpCfg.Optimizer == nil {
+		vmcpCfg.Optimizer = &config.OptimizerConfig{}
+	}
+	if !cfg.EnableEmbedding {
+		return nil, nil
+	}
+	if mgr == nil {
+		return nil, fmt.Errorf("embedding manager must not be nil when EnableEmbedding is true")
+	}
+	teiURL, err := mgr.Start(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to start TEI embedding service: %w", err)
+	}
+	vmcpCfg.Optimizer.EmbeddingService = teiURL
+	return func() { _ = mgr.Stop(context.Background()) }, nil
+}
+
 // getStatusReportingInterval extracts the status reporting interval from config.
 // Returns 0 if not configured, which uses the default interval.
 func getStatusReportingInterval(cfg *config.Config) time.Duration {