refactor(vmcp): unify composite tools and optimizer as session decorators

Both composite tools and the optimizer now implement the MultiSession
decorator pattern (same as hijackPreventionDecorator) rather than having
bespoke SDK wiring in handleSessionRegistrationImpl.

New session decorators:
- session/compositetools: appends composite tools to Tools(), routes
  their CallTool invocations to per-session workflow executors
- session/optimizerdec: replaces Tools() with [find_tool, call_tool];
  find_tool routes through the optimizer, call_tool delegates to the
  underlying session for normal backend routing

sessionmanager.Manager gains DecorateSession() to swap in a wrapped
session after creation. handleSessionRegistrationImpl becomes a flat
decoration sequence (apply compositetools → apply optimizer → register
tools) with no branching on optimizer vs non-optimizer paths.

adapter.WorkflowExecutor/WorkflowResult become type aliases for the
compositetools package types so the two layers share a single definition.
adapter.CreateOptimizerTools is deleted; its logic lives in optimizerdec.

Author: taskbot

pkg/vmcp/server/adapter/handler_factory.go

@@ -19,6 +19,7 @@ import (
 	"github.com/stacklok/toolhive/pkg/vmcp/conversion"
 	"github.com/stacklok/toolhive/pkg/vmcp/discovery"
 	"github.com/stacklok/toolhive/pkg/vmcp/router"
+	"github.com/stacklok/toolhive/pkg/vmcp/session/compositetools"
 )
 
 //go:generate mockgen -destination=mocks/mock_handler_factory.go -package=mocks github.com/stacklok/toolhive/pkg/vmcp/server/adapter HandlerFactory
@@ -43,20 +44,13 @@ type HandlerFactory interface {
 }
 
 // WorkflowExecutor executes composite tool workflows.
-// This interface abstracts the composer to enable testing without full composer setup.
-type WorkflowExecutor interface {
-	// ExecuteWorkflow executes the workflow with the given parameters.
-	ExecuteWorkflow(ctx context.Context, params map[string]any) (*WorkflowResult, error)
-}
+// Type alias for compositetools.WorkflowExecutor so that adapter consumers and
+// the session decorator share a single interface definition.
+type WorkflowExecutor = compositetools.WorkflowExecutor
 
 // WorkflowResult represents the result of a workflow execution.
-type WorkflowResult struct {
-	// Output contains the workflow output data (typically from the last step).
-	Output map[string]any
-
-	// Error contains error information if the workflow failed.
-	Error error
-}
+// Type alias for compositetools.WorkflowResult.
+type WorkflowResult = compositetools.WorkflowResult
 
 // DefaultHandlerFactory creates MCP request handlers that route to backend workloads.
 type DefaultHandlerFactory struct {

pkg/vmcp/server/adapter/optimizer_adapter.go

@@ -3,124 +3,10 @@
 
 package adapter
 
-import (
-	"context"
-	"encoding/json"
-	"fmt"
-
-	"github.com/mark3labs/mcp-go/mcp"
-	"github.com/mark3labs/mcp-go/server"
-
-	"github.com/stacklok/toolhive/pkg/vmcp/optimizer"
-	"github.com/stacklok/toolhive/pkg/vmcp/schema"
-)
-
-// OptimizerToolNames defines the tool names exposed when optimizer is enabled.
+// OptimizerToolNames defines the tool names exposed when optimizer mode is enabled.
+// These constants are kept here for backwards compatibility with existing tests and
+// callers. The actual tool implementation lives in the optimizerdec session decorator.
 const (
 	FindToolName = "find_tool"
 	CallToolName = "call_tool"
 )
-
-// Pre-generated schemas for optimizer tools.
-// Generated at package init time so any schema errors panic at startup.
-var (
-	findToolInputSchema = mustGenerateSchema[optimizer.FindToolInput]()
-	callToolInputSchema = mustGenerateSchema[optimizer.CallToolInput]()
-)
-
-// CreateOptimizerTools creates the SDK tools for optimizer mode.
-// When optimizer is enabled, only these two tools are exposed to clients
-// instead of all backend tools.
-func CreateOptimizerTools(opt optimizer.Optimizer) []server.ServerTool {
-	return []server.ServerTool{
-		{
-			Tool: mcp.Tool{
-				Name: FindToolName,
-				Description: "Find and return tools that can help accomplish the user's request. " +
-					"This searches available MCP server tools using semantic and keyword-based matching. " +
-					"Use this function when you need to: " +
-					"(1) discover what tools are available for a specific task, " +
-					"(2) find the right tool(s) before attempting to solve a problem, " +
-					"(3) check if required functionality exists in the current environment. " +
-					"Returns matching tools ranked by relevance including their names, descriptions, " +
-					"required parameters and schemas, plus token efficiency metrics showing " +
-					"baseline_tokens, returned_tokens, and savings_percent. " +
-					"Example: for 'Find good restaurants in San Jose', call with " +
-					"tool_description='search the web' and tool_keywords='web search restaurants'. " +
-					"Always call this before call_tool to discover the correct tool name and parameter schema.",
-				RawInputSchema: findToolInputSchema,
-			},
-			Handler: createFindToolHandler(opt),
-		},
-		{
-			Tool: mcp.Tool{
-				Name: CallToolName,
-				Description: "Execute a specific tool with the provided parameters. " +
-					"Use this function to: " +
-					"(1) run a tool after identifying it with find_tool, " +
-					"(2) execute operations that require specific MCP server functionality, " +
-					"(3) perform actions that go beyond your built-in capabilities. " +
-					"Important: always use find_tool first to get the correct tool_name " +
-					"and parameter schema before calling this function. " +
-					"The parameters must match the tool's input schema as returned by find_tool. " +
-					"Returns the tool's execution result which may include success/failure status, " +
-					"result data or content, and error messages if execution failed.",
-				RawInputSchema: callToolInputSchema,
-			},
-			Handler: createCallToolHandler(opt),
-		},
-	}
-}
-
-// createFindToolHandler creates a handler for the find_tool optimizer operation.
-func createFindToolHandler(opt optimizer.Optimizer) func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-		input, err := schema.Translate[optimizer.FindToolInput](request.Params.Arguments)
-		if err != nil {
-			return mcp.NewToolResultError(fmt.Sprintf("invalid arguments: %v", err)), nil
-		}
-
-		output, err := opt.FindTool(ctx, input)
-		if err != nil {
-			return mcp.NewToolResultError(fmt.Sprintf("find_tool failed: %v", err)), nil
-		}
-
-		return mcp.NewToolResultStructuredOnly(output), nil
-	}
-}
-
-// createCallToolHandler creates a handler for the call_tool optimizer operation.
-func createCallToolHandler(opt optimizer.Optimizer) func(context.Context, mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-	return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
-		input, err := schema.Translate[optimizer.CallToolInput](request.Params.Arguments)
-		if err != nil {
-			return mcp.NewToolResultError(fmt.Sprintf("invalid arguments: %v", err)), nil
-		}
-
-		result, err := opt.CallTool(ctx, input)
-		if err != nil {
-			// Exposing the error to the MCP client is important if you want it to correct its behavior.
-			// Without information on the failure, the model is pretty much hopeless in figuring out the problem.
-			return mcp.NewToolResultError(fmt.Sprintf("call_tool failed: %v", err)), nil
-		}
-
-		return result, nil
-	}
-}
-
-// mustMarshalSchema marshals a schema to JSON, panicking on error.
-// This is safe because schemas are generated from known types at startup.
-// This should NOT be called by runtime code.
-func mustGenerateSchema[T any]() json.RawMessage {
-	s, err := schema.GenerateSchema[T]()
-	if err != nil {
-		panic(fmt.Sprintf("failed to generate schema: %v", err))
-	}
-
-	data, err := json.Marshal(s)
-	if err != nil {
-		panic(fmt.Sprintf("failed to marshal schema: %v", err))
-	}
-
-	return data
-}

pkg/vmcp/server/adapter/optimizer_adapter_test.go

@@ -4,103 +4,14 @@
 package adapter
 
 import (
-	"context"
 	"testing"
 
-	"github.com/mark3labs/mcp-go/mcp"
-	"github.com/stretchr/testify/require"
-
-	"github.com/stacklok/toolhive/pkg/vmcp/optimizer"
+	"github.com/stretchr/testify/assert"
 )
 
-// mockOptimizer implements optimizer.Optimizer for testing.
-type mockOptimizer struct {
-	findToolFunc func(ctx context.Context, input optimizer.FindToolInput) (*optimizer.FindToolOutput, error)
-	callToolFunc func(ctx context.Context, input optimizer.CallToolInput) (*mcp.CallToolResult, error)
-}
-
-func (m *mockOptimizer) FindTool(ctx context.Context, input optimizer.FindToolInput) (*optimizer.FindToolOutput, error) {
-	if m.findToolFunc != nil {
-		return m.findToolFunc(ctx, input)
-	}
-	return &optimizer.FindToolOutput{}, nil
-}
-
-func (m *mockOptimizer) CallTool(ctx context.Context, input optimizer.CallToolInput) (*mcp.CallToolResult, error) {
-	if m.callToolFunc != nil {
-		return m.callToolFunc(ctx, input)
-	}
-	return mcp.NewToolResultText("ok"), nil
-}
-
-func TestCreateOptimizerTools(t *testing.T) {
-	t.Parallel()
-
-	opt := &mockOptimizer{}
-	tools := CreateOptimizerTools(opt)
-
-	require.Len(t, tools, 2)
-	require.Equal(t, FindToolName, tools[0].Tool.Name)
-	require.Equal(t, CallToolName, tools[1].Tool.Name)
-}
-
-func TestFindToolHandler(t *testing.T) {
+func TestOptimizerToolNameConstants(t *testing.T) {
 	t.Parallel()
 
-	opt := &mockOptimizer{
-		findToolFunc: func(_ context.Context, input optimizer.FindToolInput) (*optimizer.FindToolOutput, error) {
-			require.Equal(t, "read files", input.ToolDescription)
-			return &optimizer.FindToolOutput{
-				Tools: []mcp.Tool{
-					{
-						Name:        "read_file",
-						Description: "Read a file",
-					},
-				},
-			}, nil
-		},
-	}
-
-	tools := CreateOptimizerTools(opt)
-	handler := tools[0].Handler
-
-	request := mcp.CallToolRequest{}
-	request.Params.Arguments = map[string]any{
-		"tool_description": "read files",
-	}
-
-	result, err := handler(context.Background(), request)
-	require.NoError(t, err)
-	require.NotNil(t, result)
-	require.False(t, result.IsError)
-	require.Len(t, result.Content, 1)
-}
-
-func TestCallToolHandler(t *testing.T) {
-	t.Parallel()
-
-	opt := &mockOptimizer{
-		callToolFunc: func(_ context.Context, input optimizer.CallToolInput) (*mcp.CallToolResult, error) {
-			require.Equal(t, "read_file", input.ToolName)
-			require.Equal(t, "/etc/hosts", input.Parameters["path"])
-			return mcp.NewToolResultText("file contents here"), nil
-		},
-	}
-
-	tools := CreateOptimizerTools(opt)
-	handler := tools[1].Handler
-
-	request := mcp.CallToolRequest{}
-	request.Params.Arguments = map[string]any{
-		"tool_name": "read_file",
-		"parameters": map[string]any{
-			"path": "/etc/hosts",
-		},
-	}
-
-	result, err := handler(context.Background(), request)
-	require.NoError(t, err)
-	require.NotNil(t, result)
-	require.False(t, result.IsError)
-	require.Len(t, result.Content, 1)
+	assert.Equal(t, "find_tool", FindToolName)
+	assert.Equal(t, "call_tool", CallToolName)
 }